Update app.py

app.py

@@ -1,4 +1,5 @@
 import torch
+
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 import gradio as gr
@@ -25,12 +26,17 @@ pipe = FluxWithCFGPipeline.from_pretrained(
 )
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 pipe.to("cuda")
-pipe.load_lora_weights('hugovntr/flux-schnell-realism', weight_name='schnell-realism_v2.3.safetensors', adapter_name="better")
+pipe.load_lora_weights(
+    "hugovntr/flux-schnell-realism",
+    weight_name="schnell-realism_v2.3.safetensors",
+    adapter_name="better",
+)
 pipe.set_adapters(["better"], adapter_weights=[1.0])
 pipe.fuse_lora(adapter_name=["better"], lora_scale=1.0)
 pipe.unload_lora_weights()
 
-pipe.unet.to(memory_format=torch.channels_last)
+# Corrected: Access 'transformer' instead of 'unet'
+pipe.transformer.to(memory_format=torch.channels_last)
 pipe.vae.to(memory_format=torch.channels_last)
 
 pipe.enable_xformers_memory_efficient_attention()
@@ -39,7 +45,15 @@ torch.cuda.empty_cache()
 
 # Inference function
 @spaces.GPU(duration=25)
-def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
+def generate_image(
+    prompt,
+    seed=24,
+    width=DEFAULT_WIDTH,
+    height=DEFAULT_HEIGHT,
+    randomize_seed=False,
+    num_inference_steps=2,
+    progress=gr.Progress(track_tqdm=True),
+):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(int(float(seed)))
@@ -47,9 +61,15 @@ def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT,
     start_time = time.time()
 
     # Initialize static inputs for CUDA graph
-    static_latents = torch.randn((1, 4, height // 8, width // 8), dtype=dtype, device="cuda")
-    static_prompt_embeds = torch.randn((2, 77, 768), dtype=dtype, device="cuda")  # Adjust dimensions as needed
-    static_pooled_prompt_embeds = torch.randn((2, 768), dtype=dtype, device="cuda")  # Adjust dimensions as needed
+    static_latents = torch.randn(
+        (1, 4, height // 8, width // 8), dtype=dtype, device="cuda"
+    )
+    static_prompt_embeds = torch.randn(
+        (2, 77, 768), dtype=dtype, device="cuda"
+    )  # Adjust dimensions as needed
+    static_pooled_prompt_embeds = torch.randn(
+        (2, 768), dtype=dtype, device="cuda"
+    )  # Adjust dimensions as needed
     static_text_ids = torch.tensor([[[1, 2, 3]]], dtype=torch.int32, device="cuda")
     static_latent_image_ids = torch.tensor([1], dtype=torch.int64, device="cuda")
     static_timestep = torch.tensor([999], dtype=dtype, device="cuda")
@@ -86,11 +106,22 @@ def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT,
             joint_attention_kwargs=pipe.joint_attention_kwargs,
             return_dict=False,
         )[0]
-        static_latents_out = pipe.scheduler.step(static_noise_pred, static_timestep, static_latents, return_dict=False)[0]
-        static_output = pipe._decode_latents_to_image(static_latents_out, height, width, "pil")
+        static_latents_out = pipe.scheduler.step(
+            static_noise_pred, static_timestep, static_latents, return_dict=False
+        )[0]
+        static_output = pipe._decode_latents_to_image(
+            static_latents_out, height, width, "pil"
+        )
 
     # Graph-based generation function
-    def generate_with_graph(latents, prompt_embeds, pooled_prompt_embeds, text_ids, latent_image_ids, timestep):
+    def generate_with_graph(
+        latents,
+        prompt_embeds,
+        pooled_prompt_embeds,
+        text_ids,
+        latent_image_ids,
+        timestep,
+    ):
         static_latents.copy_(latents)
         static_prompt_embeds.copy_(prompt_embeds)
         static_pooled_prompt_embeds.copy_(pooled_prompt_embeds)
@@ -101,15 +132,15 @@ def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT,
         return static_output
 
     # Only generate the last image in the sequence
-    img = pipe.generate_images( 
-            prompt=prompt,
-            width=width,
-            height=height,
-            num_inference_steps=num_inference_steps,
-            generator=generator,
-            generate_with_graph=generate_with_graph
-        )
-    latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
+    img = pipe.generate_images(
+        prompt=prompt,
+        width=width,
+        height=height,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        generate_with_graph=generate_with_graph,
+    )
+    latency = f"Latency: {(time.time()-start_time):.2f} seconds"
     return img, seed, latency
 
 # Example prompts
@@ -127,12 +158,18 @@ examples = [
 with gr.Blocks() as demo:
     with gr.Column(elem_id="app-container"):
         gr.Markdown("# 🎨 Realtime FLUX Image Generator")
-        gr.Markdown("Generate stunning images in real-time with Modified Flux.Schnell pipeline.")
-        gr.Markdown("<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>")
+        gr.Markdown(
+            "Generate stunning images in real-time with Modified Flux.Schnell pipeline."
+        )
+        gr.Markdown(
+            "<span style='color: red;'>Note: Sometimes it stucks or stops generating images (I don't know why). In that situation just refresh the site.</span>"
+        )
 
         with gr.Row():
             with gr.Column(scale=2.5):
-                result = gr.Image(label="Generated Image", show_label=False, interactive=False)
+                result = gr.Image(
+                    label="Generated Image", show_label=False, interactive=False
+                )
             with gr.Column(scale=1):
                 prompt = gr.Text(
                     label="Prompt",
@@ -146,15 +183,39 @@ with gr.Blocks() as demo:
 
                 with gr.Column("Advanced Options"):
                     with gr.Row():
-                        realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=False)
+                        realtime = gr.Checkbox(
+                            label="Realtime Toggler",
+                            info="If TRUE then uses more GPU but create image in realtime.",
+                            value=False,
+                        )
                         latency = gr.Text(label="Latency")
                     with gr.Row():
                         seed = gr.Number(label="Seed", value=42)
-                        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
+                        randomize_seed = gr.Checkbox(
+                            label="Randomize Seed", value=True
+                        )
                     with gr.Row():
-                        width = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
-                        height = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
-                        num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=DEFAULT_WIDTH,
+                        )
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=DEFAULT_HEIGHT,
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Inference Steps",
+                            minimum=1,
+                            maximum=4,
+                            step=1,
+                            value=DEFAULT_INFERENCE_STEPS,
+                        )
 
         with gr.Row():
             gr.Markdown("### 🌟 Inspiration Gallery")
@@ -164,7 +225,7 @@ with gr.Blocks() as demo:
                 fn=generate_image,
                 inputs=[prompt],
                 outputs=[result, seed, latency],
-                cache_examples="lazy" 
+                cache_examples="lazy",
             )
 
     enhanceBtn.click(
@@ -173,7 +234,7 @@ with gr.Blocks() as demo:
         outputs=[result, seed, latency],
         show_progress="full",
         queue=False,
-        concurrency_limit=None
+        concurrency_limit=None,
     )
 
     generateBtn.click(
@@ -182,13 +243,13 @@ with gr.Blocks() as demo:
         outputs=[result, seed, latency],
         show_progress="full",
         api_name="RealtimeFlux",
-        queue=False
+        queue=False,
     )
 
     def update_ui(realtime_enabled):
         return {
             prompt: gr.update(interactive=True),
-            generateBtn: gr.update(visible=not realtime_enabled)
+            generateBtn: gr.update(visible=not realtime_enabled),
         }
 
     realtime.change(
@@ -196,7 +257,7 @@ with gr.Blocks() as demo:
         inputs=[realtime],
         outputs=[prompt, generateBtn],
         queue=False,
-        concurrency_limit=None
+        concurrency_limit=None,
     )
 
     async def realtime_generation(*args):
@@ -211,18 +272,26 @@ with gr.Blocks() as demo:
         outputs=[result, seed, latency],
         show_progress="full",
         queue=False,
-        concurrency_limit=None
+        concurrency_limit=None,
     )
 
     for component in [prompt, width, height, num_inference_steps]:
         component.input(
             fn=realtime_generation,
-            inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
+            inputs=[
+                realtime,
+                prompt,
+                seed,
+                width,
+                height,
+                randomize_seed,
+                num_inference_steps,
+            ],
             outputs=[result, seed, latency],
             show_progress="hidden",
             trigger_mode="always_last",
             queue=True,
-            concurrency_limit=None
+            concurrency_limit=None,
         )
 
 # Launch the app