Update app.py

app.py

@@ -4,6 +4,7 @@ import spaces
 import torch
 import random
 import os
+import time
 
 # from diffusers import QwenImageEditInpaintPipeline
 from optimization import optimize_pipeline_
@@ -137,6 +138,37 @@ Return only the rewritten instruction text directly, without JSON formatting or
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+# --- Helper functions for reuse feature ---
+def clear_result():
+    """Clears the result image."""
+    return gr.update(value=None)
+
+def update_history(new_image, history):
+    """Updates the history gallery with the new image."""
+    time.sleep(0.5)  # Small delay to ensure image is ready
+    if history is None:
+        history = []
+    if new_image is not None:
+        # Convert to list if needed (Gradio sometimes returns tuples)
+        if not isinstance(history, list):
+            history = list(history) if history else []
+        history.insert(0, new_image)
+    # Keep only the last 20 images in history
+    history = history[:20]
+    return history
+
+def use_history_as_input(evt: gr.SelectData, history):
+    """Sets the selected history image as the new input image."""
+    if history and evt.index < len(history):
+        return gr.update(value=history[evt.index][0])
+    return gr.update()
+
+def use_output_as_input(output_image):
+    """Sets the generated output as the new input image."""
+    if output_image is not None:
+        return gr.update(value=output_image)
+    return gr.update()
+
 # Initialize Qwen Image Edit pipeline
 # Scheduler configuration for Lightning
 scheduler_config = {
@@ -230,6 +262,9 @@ css = """
     width: 400px;
 }
 #edit_text{margin-top: -62px !important}
+.gallery-container {
+    margin-top: 20px;
+}
 """
 
 
@@ -278,7 +313,27 @@ with gr.Blocks(css=css) as demo:
                 )
                 run_button = gr.Button("Run")
                 
-            result = gr.ImageSlider(label="Result", show_label=False)
+            with gr.Column():
+                result = gr.ImageSlider(label="Result", show_label=False)
+                
+                use_as_input_button = gr.Button("🔄 Use as Input Image", visible=False, variant="secondary")
+                
+                gr.Markdown("---")
+                
+                with gr.Row():
+                    gr.Markdown("### 📜 History")
+                    clear_history_button = gr.Button("🗑️ Clear History", size="sm", variant="stop")
+                
+                history_gallery = gr.Gallery(
+                    label="Click any image to use as input", 
+                    columns=4, 
+                    rows=2,
+                    object_fit="contain", 
+                    height="auto",
+                    interactive=False,
+                    show_label=True,
+                    elem_classes=["gallery-container"]
+                )
         
         with gr.Accordion("Advanced Settings", open=False):
             
@@ -324,11 +379,49 @@ with gr.Blocks(css=css) as demo:
                 value=True
             )
 
+    # Event handlers for reuse functionality
+    use_as_input_button.click(
+        fn=use_output_as_input,
+        inputs=[result],
+        outputs=[edit_image],
+        show_api=False
+    )
+    
+    history_gallery.select(
+        fn=use_history_as_input,
+        inputs=[history_gallery],
+        outputs=[edit_image],
+        show_api=False
+    )
+    
+    clear_history_button.click(
+        fn=lambda: [],
+        inputs=None,
+        outputs=history_gallery,
+        show_api=False
+    )
+
+    # Main generation pipeline with result clearing, history update, and button visibility
     gr.on(
         triggers=[run_button.click, prompt.submit],
+        fn=clear_result,
+        inputs=None,
+        outputs=result,
+        show_api=False
+    ).then(
         fn = infer,
         inputs = [edit_image, prompt, negative_prompt, seed, randomize_seed, strength, num_inference_steps, true_cfg_scale, rewrite_prompt],
         outputs = [result, seed]
+    ).then(
+        fn=lambda: gr.update(visible=True),
+        inputs=None,
+        outputs=use_as_input_button,
+        show_api=False
+    ).then(
+        fn=update_history,
+        inputs=[result, history_gallery],
+        outputs=history_gallery,
+        show_api=False
     )
 
 demo.launch()