Update app.py

app.py

@@ -137,6 +137,17 @@ 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 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 = {
@@ -278,7 +289,10 @@ 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")
         
         with gr.Accordion("Advanced Settings", open=False):
             
@@ -324,11 +338,30 @@ 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
+    )
+
+    # Main generation pipeline with result clearing 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
     )
 
 demo.launch()