Let gradio_client forward the caller's ZeroGPU token itself instead of pinning one by hand

README.md

@@ -160,22 +160,19 @@ one-time `PIPE.to("cuda")` is inside the first row's 339 s and does not reappear
 
 ## Whose GPU quota pays
 
-Two cards are booked per request: this Space's denoise loop and the conditioner's forward, and **both are meant to be
-billed to the requesting user**. ZeroGPU attributes a booking to the `X-IP-Token` of the request that triggered it —
-its `/schedule` hands that token to the Spaces API together with the duration and the calling pod's IP, and nothing
-else, so there is no Space identity in the decision and a valid token minted anywhere is honoured. This Space
-therefore forwards the caller's header to the conditioner
-(`gradio_client.Client(..., headers={"X-IP-Token": ...})`, off the `gr.Request` gradio injects on
-the UI path and the `/generate` API path alike) rather than spending a token of its own.
-
-A token ZeroGPU refuses — it answers `401`, which `spaces` surfaces as `Expired ZeroGPU proxy token` — falls back to
-calling the conditioner with no token at all, billed to this Space's pod IP off a small shared quota. That is a safety
-net rather than the intended path, and the log line `conditioner call paid for by ...` records which identity actually
-paid, next to a decoded dump of what the incoming token claimed.
-
-The conditioner is sized so that even the fallback is legal: an unattributed caller may book at most 120 credits at a
-time and an `xlarge` booking costs **twice** its seconds (`_gpu_size_units`), so the conditioner books the encode
-(45 s) and a prompt upsample (60 s) as **two separate calls**, where one combined booking would be refused outright.
+Two cards are booked per request — this Space's denoise loop and the conditioner's forward — and both are billed to the
+**requesting user**, with nothing in this repository arranging it. `gradio_client` attaches the caller's own
+`x-ip-token` to every outgoing call by itself, reading it off gradio's `LocalContext` inside the event listener
+(`Client.send_data` -> `add_zero_gpu_headers`), and ZeroGPU's `/schedule` charges the booking to whatever that token
+identifies. Forwarding the header by hand is not needed and is actively worse: a cached `Client` would pin one stale
+token, which ZeroGPU refuses with `Expired ZeroGPU proxy token`.
+
+A caller with no token to forward — a `gradio_client` script rather than a browser — leaves the conditioner's booking
+attributed to this Space's pod IP and its small shared quota. That path is why the conditioner books small: an
+unattributed caller may book at most 120 credits at a time and an `xlarge` booking costs **twice** its seconds
+(`_gpu_size_units`), so the conditioner books the encode (45 s) and a prompt upsample (60 s) as two separate calls,
+where one combined booking of the old 300 s would be — and was — refused outright with `The requested GPU duration
+(600s) is larger than the maximum allowed`.
 
 ## Secrets
 

app.py

@@ -65,8 +65,7 @@ PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
-# One `gradio_client.Client` per forwarded token; see `conditioner`.
-CLIENTS: dict[str | None, object] = {}
+CLIENT = None
 
 
 def status() -> str:
@@ -171,35 +170,23 @@ def _arm_decode_hooks(pipe):
         module.decode = armed
 
 
-def conditioner(ip_token: str | None = None):
-    """The other half, over the gradio API, billed to the requesting user.
+def conditioner():
+    """The other half, over the gradio API. Cached — building a `Client` costs a round trip to the Space config.
 
-    ZeroGPU attributes a booking to the `X-IP-Token` of the request that triggered it: `/schedule` hands that token to
-    the Spaces API together with the duration and the calling pod's IP and nothing else — there is no Space identity in
-    the decision — so forwarding the caller's header makes the user's own quota pay for both halves of their request,
-    the way it would if this were a single Space.
-
-    Cached per token: building a `Client` costs a round trip to the Space config, and a token is per user session.
+    No token is passed and none has to be: `gradio_client` attaches the caller's own ZeroGPU token itself, per call,
+    by reading the `x-ip-token` of the request being served off gradio's `LocalContext` (`Client.send_data` ->
+    `add_zero_gpu_headers`). So calling this from inside an event listener — which is the only place it is called —
+    bills the conditioner's booking to the user who asked for the video, exactly as this Space's own booking is, and
+    forwarding the header by hand would only pin a stale token onto a cached client.
     """
-    from gradio_client import Client
-
-    if ip_token not in CLIENTS:
-        if len(CLIENTS) >= 32:
-            CLIENTS.pop(next(iter(CLIENTS)))
-        CLIENTS[ip_token] = Client(CONDITIONER_SPACE, headers={"X-IP-Token": ip_token} if ip_token else None)
-    return CLIENTS[ip_token]
-
+    global CLIENT
+    if CLIENT is None:
+        from gradio_client import Client
 
-def ip_token_of(request) -> str | None:
-    """The requesting user's ZeroGPU identity, as the Spaces router put it on this request.
-
-    The UI path and the `/generate` API path both reach this through the `gr.Request` gradio injects for a parameter
-    annotated with it.
-    """
-    headers = getattr(request, "headers", None)
-    return None if headers is None else headers.get("x-ip-token")
+        CLIENT = Client(CONDITIONER_SPACE)
+    return CLIENT
 
-def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=False, ip_token=None):
+def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=False):
     """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely.
 
     `rewrite_prompt` is the conditioner's prompt upsampling: it rewrites the request into MiniMax-H3's trained format
@@ -210,7 +197,7 @@ def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, rewri
     from gradio_client import handle_file
     from safetensors import safe_open
 
-    path, plan = conditioner(ip_token).predict(
+    path, plan = conditioner().predict(
         prompt=prompt,
         image_path=handle_file(image_path) if image_path else None,
         last_image_path=handle_file(last_image_path) if last_image_path else None,
@@ -274,7 +261,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
     return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
 
 
-def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True), request: gr.Request | None = None):
+def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVAS, duration=5, steps=28, seed=42, upsample=False, progress=gr.Progress(track_tqdm=True)):
     """One request. `upsample` is appended last and defaults off, so an existing API client is untouched by it."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
@@ -292,8 +279,7 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
     progress(0.0, desc=f"Upsampling the prompt on {CONDITIONER_SPACE} ..." if upsample else f"Conditioning on {CONDITIONER_SPACE} ...")
     conditioned = time.time()
     prompt_embeds, text_token_tags, metadata, plan = encode_remote(
-        prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=upsample,
-        ip_token=ip_token_of(request),
+        prompt, image_path, last_image_path, canvas, num_frames, rewrite_prompt=upsample
     )
     condition_seconds = time.time() - conditioned
     height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))