Update app.py

app.py

@@ -55,9 +55,11 @@ pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_ret
 
 # History functions
 def append_to_history(image, prompt, seed, width, height, guidance_scale, steps, history):
+    """Store only the final generated image"""
     if image is None:
         return history
     
+    # Convert final image to bytes
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     img_bytes = buffered.getvalue()
@@ -93,21 +95,29 @@ def create_history_html(history):
 
 
 @spaces.GPU(duration=75)
-def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
+    
+    # Generate all images and keep the final one
+    final_image = None
     for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
-            prompt=prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            output_type="pil",
-            good_vae=good_vae,
-        ):
-        yield img, seed
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+        output_type="pil",
+        good_vae=good_vae,
+    ):
+        final_image = img  # Keep updating until we get the final image
+        yield img, seed  # Live preview
+    
+    # Return the final image explicitly
+    yield final_image, seed
+
 
 def enhance_prompt(user_prompt):
     """Enhances the given prompt using Groq and returns the refined prompt."""
@@ -215,19 +225,20 @@ with gr.Blocks(css=css) as demo:
 
     # Event handling
     generation_event = run_button.click(
-    fn=infer,
-    inputs=[enhanced_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
-    outputs=[result, seed]  # Image is OUTPUT here
-).then(
-    fn=append_to_history,
-    inputs=[result, enhanced_prompt, seed, width, height, guidance_scale, num_inference_steps, history_state],  # Image becomes INPUT here
-    outputs=history_state
-).then(
-    fn=create_history_html,
-    inputs=history_state,
-    outputs=history_display
-)
-
+        fn=infer,
+        inputs=[enhanced_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, seed]
+    )
+        # This will execute AFTER the generator completes
+    generation_event.then(
+        fn=append_to_history,
+        inputs=[result, enhanced_prompt, seed, width, height, guidance_scale, num_inference_steps, history_state],
+        outputs=history_state
+    ).then(
+        fn=create_history_html,
+        inputs=history_state,
+        outputs=history_display
+    )
     enhanced_prompt.submit(
         fn=infer,
         inputs=[enhanced_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],