Create app.py

app.py

@@ -0,0 +1,106 @@
+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)