File size: 5,322 Bytes
4502b8a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import random
import torch
import numpy as np
from PIL import Image
import gradio as gr
from nodes import NODE_CLASS_MAPPINGS
from totoro_extras import nodes_custom_sampler
from totoro_extras import nodes_flux
# Set device to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Load the necessary models and move them to the GPU
CheckpointLoaderSimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
LoraLoader = NODE_CLASS_MAPPINGS["LoraLoader"]()
FluxGuidance = nodes_flux.NODE_CLASS_MAPPINGS["FluxGuidance"]()
RandomNoise = nodes_custom_sampler.NODE_CLASS_MAPPINGS["RandomNoise"]()
BasicGuider = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicGuider"]()
KSamplerSelect = nodes_custom_sampler.NODE_CLASS_MAPPINGS["KSamplerSelect"]()
BasicScheduler = nodes_custom_sampler.NODE_CLASS_MAPPINGS["BasicScheduler"]()
SamplerCustomAdvanced = nodes_custom_sampler.NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
EmptyLatentImage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
# Load checkpoint and move to GPU
with torch.inference_mode():
unet, clip, vae = CheckpointLoaderSimple.load_checkpoint("flux1-dev-fp8-all-in-one.safetensors")
unet = unet.to(device)
clip = clip.to(device)
vae = vae.to(device)
# Function to find the closest multiple of a number
def closestNumber(n, m):
q = int(n / m)
n1 = m * q
if (n * m) > 0:
n2 = m * (q + 1)
else:
n2 = m * (q - 1)
if abs(n - n1) < abs(n - n2):
return n1
return n2
# Main generation function
@torch.inference_mode()
def generate(positive_prompt, width, height, seed, steps, sampler_name, scheduler, guidance, lora_strength_model, lora_strength_clip):
global unet, clip
if seed == 0:
seed = random.randint(0, 18446744073709551615)
print(seed)
# Load LoRA models and move them to GPU
unet_lora, clip_lora = LoraLoader.load_lora(unet, clip, "flux_realism_lora.safetensors", lora_strength_model, lora_strength_clip)
unet_lora = unet_lora.to(device)
clip_lora = clip_lora.to(device)
# Encode prompt and apply guidance
cond, pooled = clip_lora.encode_from_tokens(clip_lora.tokenize(positive_prompt), return_pooled=True)
cond = [[cond, {"pooled_output": pooled}]]
cond = FluxGuidance.append(cond, guidance)[0]
# Generate noise and move it to the GPU
noise = RandomNoise.get_noise(seed)[0].to(device)
# Setup guider and sampler
guider = BasicGuider.get_guider(unet_lora, cond)[0]
sampler = KSamplerSelect.get_sampler(sampler_name)[0]
# Generate sigmas and latent image
sigmas = BasicScheduler.get_sigmas(unet_lora, scheduler, steps, 1.0)[0]
latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0].to(device)
# Perform sampling
sample, sample_denoised = SamplerCustomAdvanced.sample(noise, guider, sampler, sigmas, latent_image)
# Decode the latent image to a regular image
decoded = VAEDecode.decode(vae, sample)[0].detach().cpu()
# Convert to image and save
output_image = Image.fromarray(np.array(decoded * 255, dtype=np.uint8)[0])
output_image.save("/content/flux.png")
return "/content/flux.png"
# Setup the Gradio interface
with gr.Blocks(analytics_enabled=False) as demo:
with gr.Row():
with gr.Column():
positive_prompt = gr.Textbox(lines=3, interactive=True, value="cute anime girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black dress with a gold leaf pattern and a white apron eating a slice of an apple pie in the kitchen of an old dark victorian mansion with a bright window and very expensive stuff everywhere", label="Prompt")
width = gr.Slider(minimum=256, maximum=2048, value=1024, step=16, label="width")
height = gr.Slider(minimum=256, maximum=2048, value=1024, step=16, label="height")
seed = gr.Slider(minimum=0, maximum=18446744073709551615, value=0, step=1, label="seed (0=random)")
steps = gr.Slider(minimum=4, maximum=50, value=20, step=1, label="steps")
guidance = gr.Slider(minimum=0, maximum=20, value=3.5, step=0.5, label="guidance")
lora_strength_model = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.1, label="lora_strength_model")
lora_strength_clip = gr.Slider(minimum=0, maximum=1, value=1.0, step=0.1, label="lora_strength_clip")
sampler_name = gr.Dropdown(["euler", "heun", "heunpp2", "dpm_2", "lms", "dpmpp_2m", "ipndm", "deis", "ddim", "uni_pc", "uni_pc_bh2"], label="sampler_name", value="euler")
scheduler = gr.Dropdown(["normal", "sgm_uniform", "simple", "ddim_uniform"], label="scheduler", value="simple")
generate_button = gr.Button("Generate")
with gr.Column():
output_image = gr.Image(label="Generated image", interactive=False)
generate_button.click(fn=generate, inputs=[positive_prompt, width, height, seed, steps, sampler_name, scheduler, guidance, lora_strength_model, lora_strength_clip], outputs=output_image)
# Launch the Gradio interface
demo.queue().launch(inline=False, share=True, debug=True)
|