Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import fal_client
-from huggingface_hub import whoami
 import os
-from typing import Optional
+from typing import Optional, List
+
+from huggingface_hub import whoami
 
 # It is recommended to create this as a Secret on your Hugging Face Space
 # For example: FAL_KEY = "fal_key_..."
@@ -17,115 +18,123 @@ 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 generate_image(prompt: str, image: Optional[str] = None) -> str:
-    """
-    Generates or edits an image.
-    - If an image filepath is provided, it performs image-to-image editing.
-    - Otherwise, it performs text-to-image generation.
-    """
+def single_image_generation(prompt: str, image: Optional[str] = None) -> str:
+    """Handles text-to-image or single image-to-image."""
     get_fal_key()
-
     if image:
-        # Image-to-Image: Upload the local file to get a public URL
         image_url = fal_client.upload_file(image)
         result = fal_client.run(
             "fal-ai/nano-banana/edit",
-            arguments={
-                "prompt": prompt,
-                "image_url": image_url,
-            },
+            arguments={"prompt": prompt, "image_url": image_url},
         )
     else:
-        # Text-to-Image
         result = fal_client.run(
-            "fal-ai/nano-banana",
-            arguments={
-                "prompt": prompt,
-            },
+            "fal-ai/nano-banana", arguments={"prompt": prompt}
         )
     return result["images"][0]["url"]
 
-def multi_image_edit(prompt: str, images: list[str]) -> list[str]:
-    """Edits multiple images based on a text prompt."""
+def multi_image_edit(prompt: str, images: List[str]) -> List[str]:
+    """Handles multi-image editing."""
     get_fal_key()
-
     if not images:
-        raise gr.Error("Please upload at least one image.")
+        raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
 
     output_images = []
     for image_path in images:
-        # Upload each image and get its URL
         image_url = fal_client.upload_file(image_path)
         result = fal_client.run(
             "fal-ai/nano-banana/edit",
-            arguments={
-                "prompt": prompt,
-                "image_url": image_url,
-            },
+            arguments={"prompt": prompt, "image_url": image_url},
         )
         output_images.append(result["images"][0]["url"])
-
     return output_images
 
 # --- Gradio App UI ---
 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.**")
-    
+
     login_button = gr.LoginButton()
-    
-    # This entire Column will be hidden until the user logs in
+    pro_message = gr.Markdown(visible=False)
     main_interface = gr.Column(visible=False)
 
     with main_interface:
-        gr.Markdown("## Welcome! You are logged in.")
-        with gr.Tabs():
-            with gr.TabItem("Generate"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        prompt_input = gr.Textbox(label="Prompt", placeholder="A delicious looking pizza")
-                        image_input = gr.Image(type="filepath", label="Input Image (Optional)")
-                        generate_button = gr.Button("Generate", variant="primary")
-                    with gr.Column(scale=1):
-                        generate_output = gr.Image(label="Output")
-                
-                generate_button.click(
-                    generate_image,
-                    inputs=[prompt_input, image_input],
-                    outputs=[generate_output],
+        gr.Markdown("## Welcome, PRO User!")
+        with gr.Row():
+            with gr.Column(scale=1):
+                prompt_input = gr.Textbox(
+                    label="Prompt",
+                    placeholder="A delicious looking pizza"
                 )
 
-            with gr.TabItem("Multi-Image Edit"):
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        multi_prompt_input = gr.Textbox(label="Prompt", placeholder="Make it black and white")
-                        multi_image_input = gr.Gallery(label="Input Images", file_types=["image"])
-                        multi_edit_button = gr.Button("Edit Images", variant="primary")
-                    with gr.Column(scale=1):
-                        multi_image_output = gr.Gallery(label="Output Images")
-                
-                multi_edit_button.click(
-                    multi_image_edit,
-                    inputs=[multi_prompt_input, multi_image_input],
-                    outputs=[multi_image_output],
-                )
+                with gr.Tabs():
+                    with gr.TabItem("Single Image", id="single"):
+                        image_input = gr.Image(
+                            type="filepath",
+                            label="Input Image (Optional for text-to-image)"
+                        )
+                    with gr.TabItem("Multiple Images", id="multiple"):
+                        gallery_input = gr.Gallery(
+                            label="Input Images"
+                        )
 
-    def show_app_on_login(profile: Optional[gr.OAuthProfile] = None, oauth_token: gr.OAuthToken = None):
+                generate_button = gr.Button("Generate", variant="primary")
+
+            with gr.Column(scale=1):
+                # A single gallery can handle one or many images
+                output_gallery = gr.Gallery(label="Output")
+
+    def unified_generator(
+        prompt: str,
+        current_tab: str,
+        single_image: Optional[str],
+        multi_images: Optional[List[str]],
+    ) -> List[str]:
         """
-        Controls the visibility of the main UI.
-        Gradio automatically injects `gr.OAuthProfile` if the user is logged in.
+        A single handler that routes to the correct function based on the active tab.
         """
-        
-        if profile:
-            print(profile)
-            user = whoami(token=oauth_token.token)
-            print(user)
-            return gr.update(visible=True)
-        return gr.update(visible=False)
-
-    # When the app loads or the user logs in/out, this function will be called.
-    # It receives the user's profile (or None) and updates the UI visibility.
-    demo.load(show_app_on_login, inputs=None, outputs=main_interface)
+        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)
+
+    generate_button.click(
+        unified_generator,
+        inputs=[prompt_input, selected_tab, image_input, gallery_input],
+        outputs=[output_gallery],
+    )
+
+    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}")
+            return gr.update(visible=False), gr.update(visible=False)
+
+    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
 
 if __name__ == "__main__":
     demo.launch()