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

README.md

@@ -183,22 +183,19 @@ on is cold and a cold one pays the lazy 72.16 GiB `PIPE.to("cuda")` inside its f
 
 ## 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

@@ -200,8 +200,7 @@ def get_duration(prompt_embeds, text_token_tags, references, height, width, num_
 PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
-# One `gradio_client.Client` per forwarded token; see `conditioner`.
-CLIENTS: dict[str | None, object] = {}
+CLIENT = None
 
 
 def load_models() -> str | None:
@@ -307,33 +306,21 @@ def _arm_decode_hooks(pipe):
             setattr(module, method, 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 probe(path: str) -> tuple[float | None, float | None]:
     """`(video seconds, audio seconds)` of a media file, either being `None` when the stream is absent."""
@@ -431,7 +418,7 @@ def check(prompt: str, references: list[tuple[str, str]]) -> None:
             )
 
 
-def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False, ip_token=None):
+def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
     """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely.
 
     The references go over with the request: `ref2va`'s presentation puts a vision block in front of the prompt for
@@ -446,7 +433,7 @@ def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False,
     from gradio_client import handle_file
     from safetensors import safe_open
 
-    path, plan = conditioner(ip_token).predict(
+    path, plan = conditioner().predict(
         prompt=prompt,
         media=[handle_file(path) for _, path in references],
         kinds=",".join(kind for kind, _ in references),
@@ -513,7 +500,6 @@ def generate(
     seed=42,
     upsample=False,
     progress=gr.Progress(track_tqdm=True),
-    request: gr.Request | None = None,
 ):
     """One request. `upsample` is appended last and defaults off, so an existing API client is untouched by it."""
     if LOAD_ERROR:
@@ -536,7 +522,7 @@ def generate(
     conditioned = time.time()
     try:
         prompt_embeds, text_token_tags, metadata, plan = encode_remote(
-            prompt, references, canvas, requested, rewrite_prompt=upsample, ip_token=ip_token_of(request)
+            prompt, references, canvas, requested, rewrite_prompt=upsample
         )
     except gr.Error:
         raise