Pass workflow OAuth token to InferenceClient via ContextVar

app.py

@@ -4,12 +4,20 @@ 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
+
+workflow_token = contextvars.ContextVar("workflow_token", default=None)
+
 
 def get_hf_token() -> str | None:
     """
-    Retrieves the HF API token from either the user's Gradio OAuth session
-    or falls back to the system environment/CLI cached token.
+    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", {})
@@ -31,8 +39,12 @@ def generate_prompt(concept: str) -> str:
     if not concept:
         return "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = get_hf_token()
-        client = InferenceClient("nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", token=token)
+        token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
+        client = InferenceClient(
+            provider="together",
+            api_key=token,
+            bill_to="huggingface",
+        )
         system_instruction = (
             "You are an expert prompt engineer for text-to-image models. "
             "Your task is to take a simple concept and expand it into a detailed, "
@@ -46,6 +58,7 @@ def generate_prompt(concept: str) -> str:
             {"role": "user", "content": f"Concept: {concept}"}
         ]
         response = client.chat_completion(
+            model="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4",
             messages=messages,
             temperature=0.7,
             max_tokens=256
@@ -70,7 +83,7 @@ def generate_image(prompt: str) -> dict:
     if not prompt:
         prompt = "a ginger cat wearing a tiny wizard hat reading a spellbook"
     try:
-        token = get_hf_token()
+        token = get_hf_token() or os.environ.get("HF_TOKEN") or os.environ.get("HF_API_TOKEN")
         client = InferenceClient(
             provider="auto",
             api_key=token,
@@ -125,6 +138,10 @@ class LockedWorkflow(gr.Workflow):
         bound = self._bound
 
         def call_fn(data, _token=None) -> str:
+            if _token:
+                t_str = _token.token if hasattr(_token, "token") else _token
+                if t_str:
+                    workflow_token.set(str(t_str))
             fn_name = data[0] if data else ""
             try:
                 args_json = data[1] if len(data) > 1 else "[]"