Update app.py

app.py

@@ -16,13 +16,24 @@ if not GOOGLE_API_KEY:
 client = genai.Client(api_key=os.environ.get("GOOGLE_API_KEY"))
 GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'
 
-def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
+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.token)
-        return user_info.get("isPro", False) or any(org.get("isEnterprise", False) for org in user_info.get("orgs", []))
+        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