Update app.py

app.py

@@ -3,14 +3,13 @@ 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')
+# Import nodes and custom modules
+import nodes
+from nodes import NODE_CLASS_MAPPINGS
+from totoro_extras import nodes_custom_sampler, nodes_flux
 
-# Load the necessary models and move them to the GPU
+# Initialize necessary components from the nodes
 CheckpointLoaderSimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
 LoraLoader = NODE_CLASS_MAPPINGS["LoraLoader"]()
 FluxGuidance = nodes_flux.NODE_CLASS_MAPPINGS["FluxGuidance"]()
@@ -23,14 +22,10 @@ VAELoader = NODE_CLASS_MAPPINGS["VAELoader"]()
 VAEDecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
 EmptyLatentImage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
 
-# Load checkpoint and move to GPU
+# Load checkpoints and models
 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)
+    unet, clip, vae = CheckpointLoaderSimple.load_checkpoint("models/checkpoints/flux1-dev-fp8-all-in-one.safetensors")
 
-# Function to find the closest multiple of a number
 def closestNumber(n, m):
     q = int(n / m)
     n1 = m * q
@@ -42,51 +37,51 @@ def closestNumber(n, m):
         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)
+    print(f"Seed used: {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)
+    # Load LoRA models
+    unet_lora, clip_lora = LoraLoader.load_lora(unet, clip, "models/loras/flux_realism_lora.safetensors", lora_strength_model, lora_strength_clip)
 
-    # Encode prompt and apply guidance
+    # Encode the prompt
     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
+
+    # Generate noise
+    noise = RandomNoise.get_noise(seed)[0]
+
+    # Get guider and sampler
     guider = BasicGuider.get_guider(unet_lora, cond)[0]
     sampler = KSamplerSelect.get_sampler(sampler_name)[0]
-    
-    # Generate sigmas and latent image
+
+    # Get scheduling sigmas
     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
+
+    # Generate latent image
+    latent_image = EmptyLatentImage.generate(closestNumber(width, 16), closestNumber(height, 16))[0]
+
+    # Sample and decode the image
     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
+    decoded = VAEDecode.decode(vae, sample)[0].detach()
+
+    # Convert to image and return
+    return Image.fromarray(np.array(decoded * 255, dtype=np.uint8)[0])
+
+# Define 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")
+            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)")
@@ -94,13 +89,28 @@ with gr.Blocks(analytics_enabled=False) as demo:
             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")
+            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)
+    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)