Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import spaces
 import torch
 from diffusers import DiffusionPipeline
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
@@ -33,9 +34,14 @@ def inference(
 ):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    generator = torch.Generator().manual_seed(seed)
+    progress(0, "Starting image generation...")
 
+    for i in range(1, steps + 1):
+        if i % (steps // 10) == 0:
+            progress(i / steps * 100, f"Processing step {i} of {steps}...")
+    
     image = pipeline(
         prompt=prompt,
         guidance_scale=guidance_scale,
@@ -43,9 +49,11 @@ def inference(
         width=width,
         height=height,
         generator=generator,
-        lora_scale=lora_scale,
+        joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
 
+    progress(100, "Completed!")
+
     return image, seed
 
 
@@ -118,14 +126,6 @@ with gr.Blocks(css=css) as demo:
                     value=3.5,
                 )
 
-                lora_scale = gr.Slider(
-                    label="LoRA scale",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.1,
-                    value=1.0,
-                )
-
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
@@ -134,7 +134,15 @@ with gr.Blocks(css=css) as demo:
                     value=30,
                 )
 
-        gr.Examples(examples=examples, inputs=[prompt])
+                lora_scale = gr.Slider(
+                    label="LoRA scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=1.0,
+                )
+
+        gr.Examples(examples=examples, inputs=[prompt], outputs=[Image.open("./example.png")])
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -147,8 +155,10 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            lora_scale,
         ],
         outputs=[result, seed],
     )
 
-demo.queue().launch()
+demo.queue()
+demo.launch()