Update app.py

app.py

@@ -30,8 +30,8 @@ def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
 
 # --- Backend Generation Functions ---
 
-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."""
+def run_single_image_logic(prompt: str, image: Optional[str] = None) -> str:
+    """Handles text-to-image or single image-to-image and returns a single URL string."""
     get_fal_key()
     if image:
         image_url = fal_client.upload_file(image)
@@ -43,26 +43,28 @@ def run_single_image_logic(prompt: str, image: Optional[str] = None) -> List[str
         result = fal_client.run(
             "fal-ai/nano-banana", arguments={"prompt": prompt}
         )
-    return [result["images"][0]["url"]]
+    return result["images"][0]["url"]
 
-def run_multi_image_logic(prompt: str, images: List[str]) -> List[str]:
-    """Handles multi-image editing."""
+def run_multi_image_logic(prompt: str, images: List[str]) -> str:
+    """
+    Handles multi-image editing and returns the URL of the LAST generated image.
+    """
     get_fal_key()
     if not images:
         raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
 
-    output_images = []
+    last_image_url = ""
     for image_path in images:
         image_url = fal_client.upload_file(image_path)
         result = fal_client.run(
             "fal-ai/nano-banana/edit",
             arguments={"prompt": prompt, "image_url": image_url},
         )
-        output_images.append(result["images"][0]["url"])
-    return output_images
+        last_image_url = result["images"][0]["url"]
+    return last_image_url
 
 # --- Gradio App UI ---
-with gr.Blocks(theme=gr.themes.Citrus()) as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# Nano Banana Image Generation")
     gr.Markdown("Generate or edit images with FAL. **Sign in with Hugging Face to begin.**")
 
@@ -94,9 +96,8 @@ with gr.Blocks(theme=gr.themes.Citrus()) as demo:
 
             # 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")
+                output_image = gr.Image(label="Output")
+                use_image_button = gr.Button("♻️ Use this Image for Next Edit")
 
     # --- Event Handlers ---
 
@@ -105,8 +106,8 @@ with gr.Blocks(theme=gr.themes.Citrus()) as demo:
         single_image: Optional[str],
         multi_images: Optional[List[str]],
         active_tab: str,
-        oauth_token: Optional[gr.OAuthToken] = None, # Gradio injects this automatically
-    ) -> List[str]:
+        oauth_token: Optional[gr.OAuthToken] = None,
+    ) -> str:
         if not verify_pro_status(oauth_token):
             raise gr.Error("Access Denied. This service is for PRO users only.")
         if active_tab == "multiple" and multi_images:
@@ -119,26 +120,16 @@ with gr.Blocks(theme=gr.themes.Citrus()) as demo:
 
     generate_button.click(
         unified_generator,
-        # CORRECTED: login_button is removed. Gradio will inject the token via the type hint.
         inputs=[prompt_input, image_input, gallery_input, active_tab_state],
-        outputs=[output_gallery],
+        outputs=[output_image],
     )
 
-    # Handlers for the continuous editing loop
-    def store_selected_image(evt: gr.SelectData):
-        return evt.value['image']['path']
-
-    def reuse_output_image(selected_image_path):
-        if not selected_image_path:
-            gr.Warning("Please select an image from the output gallery first!")
-            return None, gr.update()
-        return selected_image_path, "single"
-
-    output_gallery.select(store_selected_image, None, selected_output_image_state)
+    # Corrected handler for the continuous editing loop.
+    # It takes the output image and directly returns it to be used as the input.
     use_image_button.click(
-        reuse_output_image,
-        inputs=[selected_output_image_state],
-        outputs=[image_input, tabs]
+        lambda img: img, # A simple function that returns its input
+        inputs=[output_image],
+        outputs=[image_input]
     )
 
     # --- Access Control Logic ---
@@ -157,7 +148,7 @@ with gr.Blocks(theme=gr.themes.Citrus()) as demo:
                 "To unlock this and many other benefits, please consider upgrading your account.\n\n"
                 "### [**Become a PRO Member Today!**](https://huggingface.co/pro)"
             )
-            return gr.update(visible=False), gr.update(visible=true, value=message)
+            return gr.update(visible=False), gr.update(visible=True, value=message)
 
     demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])