Update app.py

app.py

@@ -455,6 +455,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
     print("Generating image...")
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
+    final_image = None  # Initialize final_image
     with calculateDuration("Generating image"):
         # Generate image
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -470,7 +471,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, progress)
         ):
             print("Image generated successfully.")  # Debugging statement
             yield img
-            return final_image
+            final_image = img  # Update final_image with the current image
+    return final_image
 
 @spaces.GPU(duration=75)
 def run_lora(prompt, cfg_scale, steps, selected_indices, lora_scale_1, lora_scale_2, lora_scale_3, lora_scale_4, randomize_seed, seed, width, height, loras_state, image_input=None, progress=gr.Progress(track_tqdm=True)):