Fix: read OAuth token from session via LocalContext (call_fn doesn't forward OAuthToken to bound fns)

app.py

@@ -1,12 +1,38 @@
 import os
 import gradio as gr
 from huggingface_hub import InferenceClient
-from gradio.oauth import OAuthToken
+from huggingface_hub import get_token as hf_get_token
+from gradio.context import LocalContext
 import tempfile
 import uuid
 
 
-def generate_prompt(concept: str, oauth_token: OAuthToken | None = None) -> str:
+def _get_user_token() -> str | None:
+    """
+    Get the logged-in user's HF OAuth token from the current request session.
+    On Spaces: returns the real OAuth access_token after HF login.
+    Locally: the OAuth flow mocks the token, so we fall back to hf_get_token().
+    """
+    try:
+        request = LocalContext.request.get(None)
+        if request is not None:
+            session = getattr(request, "session", {})
+            oauth_info = session.get("oauth_info", {})
+            if oauth_info:
+                token = oauth_info.get("access_token")
+                # Skip the local dev mock token — fall through to hf_get_token()
+                if token and token != "mock-oauth-token-for-local-dev":
+                    return token
+    except Exception:
+        pass
+    # Fallback: use the locally saved HF token (hf auth login / HF_TOKEN env var)
+    try:
+        return hf_get_token()
+    except Exception:
+        return None
+
+
+def generate_prompt(concept: str) -> str:
     """
     Expands a simple concept into a detailed image prompt using the NVIDIA Nemotron model.
     Uses the signed-in user's HF OAuth token for inference provider billing.
@@ -14,7 +40,7 @@ def generate_prompt(concept: str, oauth_token: OAuthToken | None = None) -> str:
     if not concept:
         return "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = oauth_token.token if oauth_token else None
+        token = _get_user_token()
         client = InferenceClient(
             provider="together",
             api_key=token,
@@ -49,7 +75,7 @@ def generate_prompt(concept: str, oauth_token: OAuthToken | None = None) -> str:
         return f"A detailed, high-quality, professional commercial product photograph of {concept}"
 
 
-def generate_z_image(prompt: str, oauth_token: OAuthToken | None = None) -> dict:
+def generate_z_image(prompt: str) -> dict:
     """
     Generates an image from a prompt using the Tongyi-MAI/Z-Image-Turbo model.
     Uses the signed-in user's HF OAuth token for inference provider billing.
@@ -58,7 +84,7 @@ def generate_z_image(prompt: str, oauth_token: OAuthToken | None = None) -> dict
     if not prompt:
         prompt = "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = oauth_token.token if oauth_token else None
+        token = _get_user_token()
         client = InferenceClient(
             provider="auto",
             api_key=token,