Cache the conditioner client with functools.cache, tidy comments

README.md

@@ -184,18 +184,14 @@ 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 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
 
@@ -207,14 +203,11 @@ 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.
-
-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 (per-modality reference classes, `before_encode` / `after_denoise`
-steps, no `packing` modules), so block names and the shape of the split are exactly what a new head is liable to
-move.
+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, 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.
 
 Two of those are `ref2va`-only and easy to miss. PyAV decodes a reference video or audio file as the reference is
 built, and **`torchaudio`** resamples a soundtrack that is not already at the audio VAE's 32 kHz — a 32 kHz

app.py

@@ -21,6 +21,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.
@@ -200,7 +201,6 @@ def get_duration(prompt_embeds, text_token_tags, references, height, width, num_
 PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
-CLIENT = None
 
 
 def load_models() -> str | None:
@@ -237,8 +237,6 @@ def load_models() -> str | None:
         blocks = MiniMaxH3Ref2VAGeneratorBlocks()
         print(f"[ref2va] 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)
 
         # Pin the two autoencoders to torch SDPA *before* the transformer takes cuDNN, and in that order.
@@ -306,35 +304,13 @@ def _arm_decode_hooks(pipe):
             setattr(module, method, 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 probe(path: str) -> tuple[float | None, float | None]:
-    """`(video seconds, audio seconds)` of a media file, either being `None` when the stream is absent."""
-    import av
-
-    def seconds(stream, container):
-        if stream.duration is not None and stream.time_base is not None:
-            return float(stream.duration * stream.time_base)
-        return None if container.duration is None else container.duration / av.time_base
-
-    with av.open(path) as container:
-        video = seconds(container.streams.video[0], container) if container.streams.video else None
-        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
-    return video, audio
+    return Client(CONDITIONER_SPACE)
 
 
 def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
@@ -356,9 +332,8 @@ def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
 def build_references(references: list[tuple[str, str]]):
     """The `(kind, path)` references of a request as decoded reference dataclasses, in packed order.
 
-    One public class per modality since the blocks were refactored, each decoding its own file through `from_file` —
-    which is also what brings the rates along, a video its own frame rate and its soundtrack, a clip its sample rate.
-    The blocks themselves never open a media file.
+    One class per modality, each decoding its own file through `from_file`, which brings the rates along: a video its
+    own frame rate and its soundtrack, a clip its sample rate. The blocks themselves never open a media file.
     """
     from diffusers.modular_pipelines.minimax_h3 import (
         MiniMaxH3AudioReference,
@@ -425,10 +400,10 @@ def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
     every image and every merged video frame pair, so the conditioner has to see them. It decodes the very same
     files this Space does, which is what keeps the two `setup` runs in agreement.
 
-    `rewrite_prompt` is the conditioner's prompt upsampling: it rewrites the request into MiniMax-H3's trained
-    reference format with its own Qwen3-VL — which is shown the references, so it can name what each one contributes —
-    and encodes that instead, 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 `get_duration` is untouched.
+    `rewrite_prompt` asks the conditioner to rewrite the request into MiniMax-H3's trained reference format with its
+    own Qwen3-VL, which is shown the references so it can name what each one contributes, and encode that instead. It
+    runs on the conditioner's booking, and this call happens before `_generate` books a card here, so `get_duration`
+    is unaffected.
     """
     from gradio_client import handle_file
     from safetensors import safe_open
@@ -628,8 +603,7 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
 
         with gr.Column():
             result = gr.Video(label="Video + soundtrack")
-            # 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)
 
@@ -651,8 +625,7 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
         )
 
     # Same order as `generate`'s signature: the exampled five first, then the remaining image slots. `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.
+    # last and defaults off, so a positional API client that predates it is unaffected.
     request = [prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample]
 
     gr.Examples(
@@ -686,7 +659,6 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
         cache_mode="lazy",
     )
 
-    # The video stays the first output and the upsampled prompt is appended last, so existing consumers are untouched.
     run.click(generate, request, [result, upsampled, upsampled_panel], api_name="generate")
 
 

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