Use user OAuth token for inference calls instead of billing to Space owner

app.py

@@ -1,50 +1,23 @@
 import os
 import gradio as gr
 from huggingface_hub import InferenceClient
-from huggingface_hub import get_token as hf_get_token
-from gradio.context import LocalContext
-import contextvars
+from gradio.oauth import OAuthToken
 import tempfile
 import uuid
 
-workflow_token = contextvars.ContextVar("workflow_token", default=None)
 
-
-def get_hf_token() -> str | None:
-    """
-    Retrieves the HF API token from either the workflow context,
-    the user's Gradio OAuth session, or falls back to the system environment.
-    """
-    w_token = workflow_token.get()
-    if w_token:
-        return w_token
-
-    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")
-            if token and token != "mock-oauth-token-for-local-dev":
-                return token
-    try:
-        return hf_get_token()
-    except Exception:
-        return None
-
-
-def generate_prompt(concept: str) -> str:
+def generate_prompt(concept: str, oauth_token: OAuthToken | None = None) -> 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.
     """
     if not concept:
         return "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
+        token = oauth_token.token if oauth_token else None
         client = InferenceClient(
             provider="together",
             api_key=token,
-            bill_to="huggingface",
         )
         system_instruction = (
             "You are an expert prompt engineer for text-to-image models. "
@@ -76,19 +49,19 @@ def generate_prompt(concept: str) -> str:
         return f"A detailed, high-quality, professional commercial product photograph of {concept}"
 
 
-def generate_z_image(prompt: str) -> dict:
+def generate_z_image(prompt: str, oauth_token: OAuthToken | None = None) -> 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.
     Returns a dictionary structure compatible with Gradio's image viewer.
     """
     if not prompt:
         prompt = "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
+        token = oauth_token.token if oauth_token else None
         client = InferenceClient(
             provider="auto",
             api_key=token,
-            bill_to="huggingface",
         )
         image = client.text_to_image(
             prompt,