Forward the caller's ZeroGPU token to the conditioner so one request bills as one request

README.md

@@ -181,6 +181,19 @@ on is cold and a cold one pays the lazy 72.16 GiB `PIPE.to("cuda")` inside its f
 | `H3_PLACEMENT_ALLOWANCE` | `90` | Seconds of the reservation set aside for a cold worker's placement. |
 | `H3_GPU_SIZE` | `xlarge` | ZeroGPU allocation size. `large` does not fit. |
 
+## Whose GPU quota pays
+
+Two cards are booked per request: this Space's denoise loop and the conditioner's forward. ZeroGPU attributes a
+booking to the `X-IP-Token` header of the request that triggered it, so this Space forwards the caller's token to the
+conditioner (`gradio_client.Client(..., headers={"X-IP-Token": ...})`, from the `gr.Request` gradio injects — the UI
+path and the `/generate` API path alike). One user's request then bills as one request across both halves, the way it
+would if this were a single Space, and no org token is ever spent on it.
+
+A caller the router cannot attribute — an unauthenticated API call — falls back to the conditioner's IP-based quota,
+whose ceiling is 120 credits. An `xlarge` booking costs **twice** its seconds there, so the conditioner keeps its
+reservation at 60 s (120 credits) for an encode; asking it to upsample a prompt books 120 s (240 credits) and needs a
+forwarded token.
+
 ## Secrets
 
 Nothing this Space loads is private any more: the weights are the public

app.py

@@ -200,7 +200,7 @@ def get_duration(prompt_embeds, text_token_tags, references, height, width, num_
 PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
-CLIENT = None
+CLIENTS: dict[str | None, object] = {}
 
 
 def load_models() -> str | None:
@@ -306,14 +306,44 @@ def _arm_decode_hooks(pipe):
             setattr(module, method, armed)
 
 
-def conditioner():
-    """The other half, over the gradio API. Cached — building a `Client` costs a round trip to the Space config."""
-    global CLIENT
-    if CLIENT is None:
-        from gradio_client import Client
+def conditioner(ip_token: str | None = None):
+    """The other half, over the gradio API — booked against *the caller's* ZeroGPU quota, not this org's.
 
-        CLIENT = Client(CONDITIONER_SPACE)  # public Space, no org token: the request runs on the caller side quota
-    return CLIENT
+    ZeroGPU attributes a booking to the `X-IP-Token` header of the request that triggered it
+    (`spaces/zero/client.py`), which the Spaces router puts on every browser request. That header is what pays for
+    this Space's own `@spaces.GPU` call, and forwarding it to the conditioner makes the same identity pay for the
+    conditioner's — the two halves of one user's request then bill as one request, the way they would if this were a
+    single Space.
+
+    Without it the conditioner falls back to an IP-based quota, whose ceiling is low enough that an `xlarge` booking
+    is refused outright ("The requested GPU duration (Ns) is larger than the maximum allowed"), so a call that does
+    not forward a token only works because the conditioner keeps its own reservation small.
+
+    Cached per token: building a `Client` costs a round trip to the Space config, and a token is per user session.
+    """
+    from gradio_client import Client
+
+    if ip_token in CLIENTS:
+        return CLIENTS[ip_token]
+    # No org token: the request runs on the caller side quota, which is the point of forwarding theirs.
+    client = Client(CONDITIONER_SPACE, headers={"X-IP-Token": ip_token} if ip_token else None)
+    if len(CLIENTS) >= 32:
+        CLIENTS.pop(next(iter(CLIENTS)))
+    CLIENTS[ip_token] = client
+    return client
+
+
+def ip_token_of(request) -> str | None:
+    """The caller's ZeroGPU identity, as the Spaces router put it on this request.
+
+    Present on a browser request and on an API request the router could attribute; absent for a truly anonymous
+    caller, which then falls back to the conditioner's IP-based quota. Both the UI path and the `/generate` API path
+    reach this through the same `gr.Request` gradio injects for a parameter annotated with it.
+    """
+    headers = getattr(request, "headers", None)
+    token = None if headers is None else headers.get("x-ip-token")
+    print(f"[ref2va] conditioner call {'forwards the caller ZeroGPU token' if token else 'is anonymous (IP quota)'}", flush=True)
+    return token
 
 
 def probe(path: str) -> tuple[float | None, float | None]:
@@ -412,7 +442,7 @@ def check(prompt: str, references: list[tuple[str, str]]) -> None:
             )
 
 
-def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
+def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False, ip_token=None):
     """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
@@ -427,7 +457,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().predict(
+    path, plan = conditioner(ip_token).predict(
         prompt=prompt,
         media=[handle_file(path) for _, path in references],
         kinds=",".join(kind for kind, _ in references),
@@ -494,6 +524,7 @@ 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:
@@ -516,7 +547,7 @@ def generate(
     conditioned = time.time()
     try:
         prompt_embeds, text_token_tags, metadata, plan = encode_remote(
-            prompt, references, canvas, requested, rewrite_prompt=upsample
+            prompt, references, canvas, requested, rewrite_prompt=upsample, ip_token=ip_token_of(request)
         )
     except gr.Error:
         raise