Update app.py

app.py

@@ -34,14 +34,12 @@ def run_single_image_logic(prompt: str, image: Optional[str] = None) -> List[str
     """Handles text-to-image or single image-to-image and returns a list."""
     get_fal_key()
     if image:
-        # Single Image-to-Image
         image_url = fal_client.upload_file(image)
         result = fal_client.run(
             "fal-ai/nano-banana/edit",
             arguments={"prompt": prompt, "image_url": image_url},
         )
     else:
-        # Text-to-Image
         result = fal_client.run(
             "fal-ai/nano-banana", arguments={"prompt": prompt}
         )
@@ -81,12 +79,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     label="Prompt",
                     placeholder="A delicious looking pizza"
                 )
-
-                # Use gr.State to track the active tab
                 active_tab_state = gr.State(value="single")
-
                 with gr.Tabs() as tabs:
-                    # When a tab is selected, its 'select' event updates the active_tab_state
                     with gr.TabItem("Single Image", id="single") as single_tab:
                         image_input = gr.Image(
                             type="filepath",
@@ -96,16 +90,16 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         gallery_input = gr.Gallery(
                             label="Input Images", file_types=["image"]
                         )
-                
                 generate_button = gr.Button("Generate", variant="primary")
 
             # RIGHT COLUMN: Outputs
             with gr.Column(scale=1):
                 output_gallery = gr.Gallery(label="Output")
+                selected_output_image_state = gr.State()
+                use_image_button = gr.Button("♻️ Use Selected Image for Next Edit")
 
     # --- Event Handlers ---
 
-    # A single, unified generator function
     def unified_generator(
         prompt: str,
         single_image: Optional[str],
@@ -113,40 +107,53 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         active_tab: str,
         oauth_token: Optional[gr.OAuthToken] = None,
     ) -> List[str]:
-        """
-        A single handler that verifies PRO status and routes to the correct function.
-        The `active_tab` state is the primary source of truth for user intent.
-        """
-        # Server-side security check before any inference
         if not verify_pro_status(oauth_token):
             raise gr.Error("Access Denied. This service is for PRO users only.")
-
-        # Logic based on the active tab and whether inputs exist
         if active_tab == "multiple" and multi_images:
             return run_multi_image_logic(prompt, multi_images)
-        else: # Covers single image, or text-to-image if single_image is None
+        else:
             return run_single_image_logic(prompt, single_image)
 
-    # Cleanly update the state when a tab is selected
     single_tab.select(lambda: "single", None, active_tab_state)
     multi_tab.select(lambda: "multiple", None, active_tab_state)
-
-    # Wire the button to the unified generator
+    
     generate_button.click(
         unified_generator,
         inputs=[prompt_input, image_input, gallery_input, active_tab_state, login_button],
         outputs=[output_gallery],
     )
 
+    # New handlers for the continuous editing loop
+    def store_selected_image(evt: gr.SelectData):
+        """When an image is selected in the output gallery, store its path in state."""
+        return evt.value['image']
+
+    def reuse_output_image(selected_image_path):
+        """
+        Takes the path from state and sends it to the single image input.
+        Also forces the UI to switch to the "Single Image" tab.
+        """
+        if not selected_image_path:
+            gr.Warning("Please select an image from the output gallery first!")
+            return None, gr.update()
+        # Output 1: The image path for the gr.Image component
+        # Output 2: The ID of the tab to select for the gr.Tabs component
+        return selected_image_path, "single"
+
+    output_gallery.select(store_selected_image, None, selected_output_image_state)
+    use_image_button.click(
+        reuse_output_image,
+        inputs=[selected_output_image_state],
+        outputs=[image_input, tabs]
+    )
+
     # --- Access Control Logic ---
     def control_access(
         profile: Optional[gr.OAuthProfile] = None,
         oauth_token: Optional[gr.OAuthToken] = None
     ):
-        """Controls UI visibility based on user's PRO status."""
         if not profile:
             return gr.update(visible=False), gr.update(visible=False)
-
         if verify_pro_status(oauth_token):
             return gr.update(visible=True), gr.update(visible=False)
         else: