Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import fal_client
 import os
 from typing import Optional, List
-
 from huggingface_hub import whoami
 
 # It is recommended to create this as a Secret on your Hugging Face Space
@@ -18,22 +17,37 @@ def get_fal_key():
     if not FAL_KEY:
         raise gr.Error("FAL_KEY is not set. Please add it to your Hugging Face Space secrets.")
 
-def single_image_generation(prompt: str, image: Optional[str] = None) -> str:
-    """Handles text-to-image or single image-to-image."""
+def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
+    """Verifies if the user is a Hugging Face PRO user using their token."""
+    if not token:
+        return False
+    try:
+        user_info = whoami(token=token.token)
+        return user_info.get("isPro", False)
+    except Exception as e:
+        print(f"Could not verify user's PRO status: {e}")
+        return False
+
+# --- 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."""
     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}
         )
-    return result["images"][0]["url"]
+    return [result["images"][0]["url"]]
 
-def multi_image_edit(prompt: str, images: List[str]) -> List[str]:
+def run_multi_image_logic(prompt: str, images: List[str]) -> List[str]:
     """Handles multi-image editing."""
     get_fal_key()
     if not images:
@@ -61,79 +75,89 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with main_interface:
         gr.Markdown("## Welcome, PRO User!")
         with gr.Row():
+            # LEFT COLUMN: Inputs
             with gr.Column(scale=1):
                 prompt_input = gr.Textbox(
                     label="Prompt",
                     placeholder="A delicious looking pizza"
                 )
 
-                with gr.Tabs():
-                    with gr.TabItem("Single Image", id="single"):
+                # 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",
                             label="Input Image (Optional for text-to-image)"
                         )
-                    with gr.TabItem("Multiple Images", id="multiple"):
+                    with gr.TabItem("Multiple Images", id="multiple") as multi_tab:
                         gallery_input = gr.Gallery(
-                            label="Input Images"
+                            label="Input Images", file_types=["image"]
                         )
-
+                
                 generate_button = gr.Button("Generate", variant="primary")
 
+            # RIGHT COLUMN: Outputs
             with gr.Column(scale=1):
-                # A single gallery can handle one or many images
                 output_gallery = gr.Gallery(label="Output")
 
+    # --- Event Handlers ---
+
+    # A single, unified generator function
     def unified_generator(
         prompt: str,
-        current_tab: str,
         single_image: Optional[str],
         multi_images: Optional[List[str]],
+        active_tab: str,
+        oauth_token: Optional[gr.OAuthToken] = None,
     ) -> List[str]:
         """
-        A single handler that routes to the correct function based on the active tab.
+        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.
         """
-        if current_tab == "multiple":
-            return multi_image_edit(prompt, multi_images)
-        else: # Handles both text-to-image and single image-to-image
-            result_url = single_image_generation(prompt, single_image)
-            return [result_url]
-
-    # The `select` event on the tabs gives us the active tab's ID
-    selected_tab = gr.Textbox(value="single", visible=False)
-    for tab in demo.load_queue[0][0].children[1].children: # Hacky way to get tabs
-        if isinstance(tab, gr.Tab):
-            tab.select(lambda: tab.id, None, selected_tab)
+        # 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
+            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, selected_tab, image_input, gallery_input],
+        inputs=[prompt_input, image_input, gallery_input, active_tab_state, login_button],
         outputs=[output_gallery],
     )
 
+    # --- 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 or not oauth_token:
-            return gr.update(visible=False), gr.update(visible=False)
-        try:
-            user_info = whoami(token=oauth_token.token)
-            if user_info.get("isPro", False):
-                return gr.update(visible=True), gr.update(visible=False)
-            else:
-                message = (
-                    "## ✨ Exclusive Access for PRO Users\n\n"
-                    "Thank you for your interest! This feature is available exclusively for our Hugging Face **PRO** members.\n\n"
-                    "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)
-        except Exception as e:
-            gr.Warning(f"Could not verify user status: {e}")
+        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:
+            message = (
+                "## ✨ Exclusive Access for PRO Users\n\n"
+                "Thank you for your interest! This feature is available exclusively for our Hugging Face **PRO** members.\n\n"
+                "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)
+
     demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
 
 if __name__ == "__main__":