Update app.py

app.py

@@ -1,15 +1,46 @@
 import gradio as gr
 from models import generate_image, MODEL_ID
 from config import APPLE_TENCENT_THEME
+from typing import Optional, Union
+from huggingface_hub import whoami
 
-def generate_image_with_auth(prompt: str, profile: gr.OAuthProfile | None):
+def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
+    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
+    if not token:
+        return False
+    
+    if isinstance(token, gr.OAuthToken):
+        token_str = token.token
+    elif isinstance(token, str):
+        token_str = token
+    else:
+        return False
+    
+    try:
+        user_info = whoami(token=token_str)
+        return (
+            user_info.get("isPro", False) or
+            any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
+        )
+    except Exception as e:
+        print(f"Could not verify user's PRO/Enterprise status: {e}")
+        return False
+
+def generate_image_with_pro_check(prompt: str, profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
     """
-    Wrapper function that checks if user is logged in before generating image.
+    Wrapper function that checks if user is PRO before generating image.
     """
     if profile is None:
-        raise gr.Error("Click Sign in with Hugging Face button to use this app for free")
+        raise gr.Error("Please Sign in with Hugging Face to use this app")
     
-    # User is logged in, proceed with image generation
+    # Check if user has PRO status
+    if not verify_pro_status(oauth_token):
+        raise gr.Error(
+            "This app is exclusively for Hugging Face PRO subscribers. "
+            "Please upgrade at https://huggingface.co/subscribe/pro to access this feature."
+        )
+    
+    # User is PRO, proceed with image generation
     return generate_image(prompt)
 
 def create_ui():
@@ -18,54 +49,101 @@ def create_ui():
             f"<div style='text-align: center; max-width: 700px; margin: 0 auto;'>"
             f"<h1>Tencent {MODEL_ID.split('/')[-1]}</h1>"
             f"<p>Generate images using Tencent's state-of-the-art model hosted by FAL AI.</p>"
-            f"<p style='color: orange;'>⚠️ You must Sign in with Hugging Face using the button to use this app.</p>"
+            f"<p style='color: orange;'>⚠️ This app is exclusively for Hugging Face PRO subscribers.</p>"
+            f"<p><a href='https://huggingface.co/subscribe/pro' target='_blank' style='color: #ff6b6b; font-weight: bold;'>🚀 Subscribe to PRO</a></p>"
             f"Built with <a href='https://huggingface.co/spaces/akhaliq/anycoder' target='_blank'>anycoder</a>"
             f"</div>"
         )
         
         # Add login button - required for OAuth
         gr.LoginButton()
+        
+        # PRO status message area
+        pro_message = gr.Markdown(visible=False)
+        main_interface = gr.Column(visible=False)
 
-        with gr.Row():
-            with gr.Column(scale=1):
-                prompt_input = gr.Textbox(
-                    label="Prompt",
-                    placeholder="e.g., A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves.",
-                    lines=4
+        with main_interface:
+            with gr.Row():
+                with gr.Column(scale=1):
+                    prompt_input = gr.Textbox(
+                        label="Prompt",
+                        placeholder="e.g., A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves.",
+                        lines=4
+                    )
+                    generate_btn = gr.Button("🎨 Generate Image", variant="primary")
+                    
+                with gr.Column(scale=1):
+                    output_image = gr.Image(
+                        label="Generated Image",
+                        height=512,
+                        width=512,
+                        interactive=False,
+                        show_download_button=True
+                    )
+            
+            # Set up the event listener - note the function now takes OAuthProfile and OAuthToken
+            generate_btn.click(
+                fn=generate_image_with_pro_check,
+                inputs=[prompt_input],
+                outputs=[output_image],
+                queue=False,
+                api_name=False,
+                show_api=False,
+            )
+            
+            # Example usage guidance with queue features disabled
+            gr.Examples(
+                examples=[
+                    "A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves.",
+                    "A futuristic city skyline at sunset with flying cars",
+                    "A mystical forest with glowing mushrooms and fireflies",
+                    "An astronaut riding a horse on Mars",
+                ],
+                inputs=prompt_input,
+                outputs=output_image,
+                fn=generate_image_with_pro_check,
+                cache_examples=False,
+                api_name=False,
+                show_api=False,
+            )
+            
+            gr.Markdown("<h3 style='text-align: center'>Thank you for being a PRO subscriber! 🤗</h3>")
+
+        # Control access based on PRO status
+        def control_access(profile: Optional[gr.OAuthProfile] = None, oauth_token: Optional[gr.OAuthToken] = None):
+            if not profile:
+                return (
+                    gr.update(visible=False),  # main_interface
+                    gr.update(visible=True, value=(
+                        "## 👋 Welcome!\n\n"
+                        "Please **Sign in with Hugging Face** using the button above to continue."
+                    ))  # pro_message
                 )
-                generate_btn = gr.Button("🎨 Generate Image", variant="primary")
-                
-            with gr.Column(scale=1):
-                output_image = gr.Image(
-                    label="Generated Image",
-                    height=512,
-                    width=512,
-                    interactive=False,
-                    show_download_button=True
+            
+            if verify_pro_status(oauth_token):
+                return (
+                    gr.update(visible=True),   # main_interface
+                    gr.update(visible=False)   # pro_message
+                )
+            else:
+                message = (
+                    "## ✨ Exclusive Access for PRO Users\n\n"
+                    "Thank you for your interest! This app is available exclusively for our Hugging Face **PRO** members.\n\n"
+                    "With PRO, you'll get:\n"
+                    "- Access to exclusive Spaces like this one\n"
+                    "- Increased GPU quotas\n"
+                    "- Early access to new features\n"
+                    "- Priority support\n\n"
+                    "### [**Become a PRO Today!**](https://huggingface.co/subscribe/pro)\n\n"
+                    "*Already a PRO member? Try logging out and back in if you're having issues.*"
+                )
+                return (
+                    gr.update(visible=False),  # main_interface
+                    gr.update(visible=True, value=message)  # pro_message
                 )
-
-        # Set up the event listener - note the function now takes OAuthProfile
-        generate_btn.click(
-            fn=generate_image_with_auth,
-            inputs=[prompt_input],
-            outputs=[output_image],
-            queue=False,
-            api_name=False,
-            show_api=False,
-        )
         
-        # Example usage guidance with queue features disabled
-        gr.Examples(
-            examples=[
-                "A detailed watercolor painting of a small red fox sleeping on a pile of autumn leaves."
-            ],
-            inputs=prompt_input,
-            outputs=output_image,
-            fn=generate_image,  # Examples use the original function
-            cache_examples=False,
-            api_name=False,
-            show_api=False,
-        )
+        # Load event to check access on page load
+        demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
 
     return demo