Cache the conditioner client with functools.cache, tidy comments

README.md

@@ -161,18 +161,14 @@ 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 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`.
+requesting user, with nothing here arranging it: `gradio_client` attaches the caller's own `x-ip-token` to every
+outgoing call, reading it off gradio's `LocalContext` inside the event listener (`Client.send_data` ->
+`add_zero_gpu_headers`), and ZeroGPU charges the booking to whatever that token identifies.
 
 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`.
+attributed to this Space's pod IP and its small shared quota. An unattributed caller may book at most 120 credits at a
+time and an `xlarge` booking costs twice its seconds, so the conditioner books the encode (45 s) and a prompt upsample
+(60 s) as two separate calls, each within that ceiling.
 
 ## Secrets
 
@@ -184,10 +180,8 @@ is a public Space called on the requesting user's own ZeroGPU token, never on an
 ## Where diffusers comes from
 
 MiniMax-H3 is modular-only and not in a released `diffusers`, so `requirements.txt` installs it from the canonical
-pull request, [huggingface/diffusers#14371](https://github.com/huggingface/diffusers/pull/14371), pinned to the
-**commit** `665f5782` (`refs/pull/14371/head` at deploy time) rather than to the moving `minimax-h3-refactor` branch.
+pull request, [huggingface/diffusers#14371](https://github.com/huggingface/diffusers/pull/14371), pinned to the commit
+`665f5782` (`refs/pull/14371/head`) rather than to the moving `minimax-h3-refactor` branch.
 
-That PR is a WIP: it needs **re-pinning whenever it updates**, and `h3_split_blocks.py` — which subclasses its block
-classes to cut the pipeline in two — has to be re-checked against the new head at the same time. The PR refactored
-the blocks into one workflow-selected pipeline, so block names and the shape of the split are exactly what a new head
-is liable to move.
+That PR is a WIP, so it needs re-pinning whenever it updates, and `h3_split_blocks.py` — which subclasses its block
+classes to cut the pipeline in two — has to be re-checked against the new head at the same time.

app.py

@@ -7,6 +7,7 @@ import os
 import tempfile
 import time
 import traceback
+from functools import cache
 
 # First, and at module level. `import spaces` patches `torch.cuda` before any GPU is attached, which is what lets the
 # 72 GiB load happen at **startup** rather than on GPU time; it also has to precede anything that initializes CUDA.
@@ -65,7 +66,6 @@ PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
-CLIENT = None
 
 
 def status() -> str:
@@ -117,8 +117,6 @@ def load_models() -> str | None:
         blocks = MiniMaxH3GeneratorBlocks()
         print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
         pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
-        # Every repository this Space reads is public — the checkpoint, the AoTI packages and the conditioner
-        # Space — so no token is passed anywhere.
         pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)
 
@@ -170,45 +168,13 @@ def _arm_decode_hooks(pipe):
         module.decode = armed
 
 
+@cache
 def conditioner():
-    """The other half, over the gradio API. Cached — building a `Client` costs a round trip to the Space config.
+    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, off
+    gradio's `LocalContext`, so the conditioner's booking is billed to the user who asked for the video."""
+    from gradio_client import Client
 
-    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.
-    """
-    global CLIENT
-    if CLIENT is None:
-        from gradio_client import Client
-
-        CLIENT = Client(CONDITIONER_SPACE)
-    return CLIENT
-
-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
-    with its own Qwen3-VL and encodes *that*, handing the rewrite back under the plan's `refined_prompt`. It runs on
-    the conditioner's GPU booking, and this whole call happens before `_generate` books a card here, so it costs this
-    Space's `get_duration` nothing.
-    """
-    from gradio_client import handle_file
-    from safetensors import safe_open
-
-    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,
-        canvas=canvas,
-        num_frames=num_frames,
-        rewrite_prompt=bool(rewrite_prompt),
-        api_name="/encode",
-    )
-    with safe_open(path, framework="pt") as handle:
-        metadata = handle.metadata()
-        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), metadata, plan
+    return Client(CONDITIONER_SPACE)
 
 
 
@@ -262,7 +228,7 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
 
 
 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."""
+    """One request. `upsample` is last and defaults off, so a positional API client is unaffected by it."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
     if PIPE is None:
@@ -285,9 +251,8 @@ def generate(prompt, image_path=None, last_image_path=None, canvas=DEFAULT_CANVA
     height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
     refined = plan.get("refined_prompt") or ""
 
-    # EXIF-transposed and in RGB before the blocks see them, which the resize step that replaced the old setup step no
-    # longer does itself. Both halves have to prepare a keyframe the same way or the conditioning latents encoded here
-    # would not be of the image the conditioner looked at.
+    # EXIF-transposed and in RGB, the same way the conditioner prepares it: the conditioning latents encoded here have
+    # to be of the image the conditioner looked at.
     def keyframe(path):
         return ImageOps.exif_transpose(Image.open(path)).convert("RGB") if path else None
 
@@ -403,8 +368,7 @@ with gr.Blocks(title="MiniMax-H3") as demo:
         with gr.Column():
             video = gr.Video(label="Video + soundtrack")
             report = gr.Markdown(visible=False)
-            # Only shown for a request that actually asked for a rewrite, so a plain request is not left with an
-            # empty panel. The accordion is an output for that reason: its visibility is part of the answer.
+            # An output, so it can be revealed only for a request that asked for a rewrite.
             with gr.Accordion("Upsampled prompt", open=False, visible=False) as upsampled_panel:
                 upsampled = gr.Textbox(show_label=False, lines=8, interactive=False)
 
@@ -425,9 +389,8 @@ with gr.Blocks(title="MiniMax-H3") as demo:
         cache_mode="lazy",
     )
 
-    # `upsample` is appended *after* every input that was already here and every existing input keeps its position, so
-    # a positional API client that predates it keeps working and simply takes the default. Same on the way out: the
-    # video and the report stay first and the upsampled prompt is appended last.
+    # `upsample` is last and defaults off, so a positional API client that predates it is unaffected; so is the
+    # output order, with the upsampled prompt appended after the video and the report.
     run.click(
         generate,
         [prompt, image, last_image, canvas, duration, steps, seed, upsample],

h3_aoti.py

@@ -62,9 +62,8 @@ import os
 from pathlib import Path
 
 AOTI = os.environ.get("H3_AOTI", "0") == "1"
-# A public **model** repo. It used to be a private dataset, which is why the repo type is still a variable: the
-# artifacts are keyed by quant/torch/arch under `<width>/torch<X.Y>/sm<cc>/<shape>` rather than laid out the way
-# `spaces.aoti_load` expects, so the download is done by hand either way (see `maybe_load`).
+# The artifacts are keyed by quant/torch/arch under `<width>/torch<X.Y>/sm<cc>/<shape>` rather than laid out the way
+# `spaces.aoti_load` expects, so the download is done by hand (see `maybe_load`).
 AOTI_REPO = os.environ.get("H3_AOTI_REPO", "multimodalart/minimax-h3-aoti")
 AOTI_REPO_TYPE = os.environ.get("H3_AOTI_REPO_TYPE", "model")
 # `dynamic` is the one package that serves every canvas, duration *and prompt*, and for bfloat16 it is what gets built:

h3_split_blocks.py

@@ -1,14 +1,13 @@
 """The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.
 
-MiniMax-H3 is modular-only. Since https://github.com/huggingface/diffusers/pull/14371 the whole model is *one*
-`MiniMaxH3Blocks` sequence whose branches are picked per request (and per `workflow=`) from the inputs:
+MiniMax-H3 is modular-only, and the whole model is one `MiniMaxH3Blocks` sequence whose branches are picked per
+request — and per `workflow=` — from the inputs:
 
     before_encode -> text_encoder -> vae_encoder -> denoise -> after_denoise -> decode
 
-where `before_encode`, `text_encoder`, `vae_encoder` and `denoise` are each an auto-block that switches on
-`references` (the `ref2va` workflow) versus the keyframe inputs (`t2va` / `fl2va`), and `denoise` is a whole
-sub-sequence — `prepare_layout -> prepare_latents -> set_timesteps -> denoise` — against `transformer` or
-`transformer_ref`.
+where `before_encode`, `text_encoder`, `vae_encoder` and `denoise` each switch on `references` (the `ref2va` workflow)
+versus the keyframe inputs (`t2va` / `fl2va`), and `denoise` is itself `prepare_layout -> prepare_latents ->
+set_timesteps -> denoise` against `transformer` or `transformer_ref`.
 
 The conditioner (a 62.14 GiB Qwen3-VL) and the denoiser (a 61.73 GiB transformer plus ~20.5 GiB of float32 VAEs) do
 not fit on one 95 GiB card unquantized, so this module cuts that sequence in two at the `text_encoder` step, once per
@@ -22,25 +21,15 @@ partition:
   * `MiniMaxH3Ref2VAConditionerBlocks` / `MiniMaxH3Ref2VAGeneratorBlocks` are the same cut through the `ref2va`
     branch, so one conditioner Space serves both partitions out of the weights it already holds.
 
-What the refactor moved, and what that means for the cut:
-
-  * the old `MiniMaxH3SetupStep` is gone. Its keyframe half is now `MiniMaxH3ResizeStep` (wrapped in the conditional
-    `MiniMaxH3AutoResizeStep`, skipped entirely for a text-only request) and its geometry half moved *into*
-    `MiniMaxH3PrepareLayoutStep`, which lives on the denoising side. So the `t2va` / `fl2va` conditioner half no
-    longer resolves `num_frames` at all — the caller does, with `align_num_frames`, which is the one line of
-    arithmetic that used to come back in the plan.
-  * `MiniMaxH3Ref2VASetupStep` survived and still resolves the canvas *and* the frame count, but `num_frames` is now
-    required there: a request that leaves the duration to its single audio-bearing reference resolves it caller-side.
-  * unpacking the denoised rows moved out of the decoders into `MiniMaxH3AfterDenoiseStep`, so the generating halves
-    carry that step explicitly.
-  * `model_name` is `"minimax-h3"` on every block now — the `"minimax-h3-ref2va"` pipeline mapping was dropped when
-    the two blocksets became workflows of one pipeline.
-
-`resize` / `setup` run on both sides on purpose. They own no pretrained component (PIL, decoded media and
-arithmetic), they resolve the canvas and prepare the keyframes or normalize the references — which the conditioner
-needs to build its vision blocks and the generator needs to encode with the VAEs. Running them twice over the same
-inputs is deterministic; both conditioner halves return the resolved `height` / `width` / `num_frames` anyway, so the
-caller pins them explicitly on the generating half.
+`resize` / `setup` run on both sides on purpose. They own no pretrained component (PIL, decoded media and arithmetic),
+they resolve the canvas and prepare the keyframes or normalize the references — which the conditioner needs to build
+its vision blocks and the generator needs to encode with the VAEs. Running them twice over the same inputs is
+deterministic; both conditioner halves return the resolved `height` / `width` / `num_frames` anyway, so the caller
+pins them explicitly on the generating half.
+
+Two things the blocks leave to the caller: a keyframe reaches them EXIF-transposed and in RGB, and the `t2va` /
+`fl2va` frame count is aligned to `17 * n + 5` before the call, since that arithmetic lives in the layout step on the
+denoising side of the cut. `ref2va` still resolves its own frame count, but requires one to be passed.
 
 Only *text* encoding is remote. `vae_encoder` / `reference_encoder` stay on the denoising side: they run the two
 autoencoders, which the conditioner Space does not hold.
@@ -68,8 +57,7 @@ from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
 def _wire_outputs(num_frames: bool = True) -> list[OutputParam]:
     """The wire format of the split, plus the plan the caller pins on the generating half.
 
-    `num_frames` is only declared by the `ref2va` half: its setup step is the one that still resolves the frame count,
-    while the keyframe half's own resolution moved into the layout step, on the other side of the cut.
+    `num_frames` is declared by the `ref2va` half alone: it is the one whose setup step resolves a frame count.
     """
     return [
         OutputParam.template("prompt_embeds"),
@@ -96,8 +84,7 @@ class MiniMaxH3ConditionerBlocks(SequentialPipelineBlocks):
         return (
             "The conditioner half of a split MiniMax-H3 deployment: puts the keyframes onto the target canvas and "
             "encodes MiniMax-H3's presentation of the request into the `prompt_embeds` / `text_token_tags` pair the "
-            "denoising half consumes. The frame count is the caller's to align — that arithmetic now lives in the "
-            "layout step, on the denoising side."
+            "denoising half consumes. The frame count is the caller's to align."
         )
 
     @property