Price bookings against the AoTI blocks, generate from 2 s again, cross-link the demos

get_duration is refit on the measured AoTI per-step table, AoTI is opt-in behind H3_AOTI with a card and torch check before it loads, and the duration slider starts at 2 s again by lowering the pipeline's own floor.

app.py

@@ -1,18 +1,8 @@
-"""MiniMax-H3 `ref2va`, split deployment — **the denoising half**.
+"""MiniMax-H3 `ref2va`, split deployment — the denoising half.
 
-This Space holds the `transformer_ref` partition of the checkpoint and the two autoencoders, **unquantized
-bfloat16**, and nothing else. The 62.14 GiB Qwen3-VL conditioner lives in its own Space,
-[`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner), which this
-one calls over the gradio API for every request; what comes back is a safetensors file holding the two tensors the
-denoiser needs, `prompt_embeds` and `text_token_tags`.
-
-Why split at all: MiniMax-H3 is 195.9 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage, so an
-unquantized single Space is impossible. Cut at the text-encoder step, this half pulls 77.3 GB (`transformer_ref/`
-61.73 GiB + `vae/` 9.70 + `audio_vae/` 0.56) and the conditioner 66.7 GB, and neither is quantized.
-
-The blockset is the `ref2va` branch of `MiniMaxH3Blocks` with its `text_encoder` step removed — see
-`h3_split_blocks.py`. Only *text* encoding is remote: `reference_encoder` is the `ref2va` branch's own encoder step
-and runs here, next to the two autoencoders it needs.
+This Space holds the `transformer_ref` partition and the two autoencoders, unquantized bfloat16. Text encoding runs in
+[`qwen3vl-conditioner`](https://huggingface.co/spaces/multimodalart/qwen3vl-conditioner), which this one calls over the
+gradio API for every request; `reference_encoder` stays here, next to the autoencoders it runs.
 """
 
 from __future__ import annotations
@@ -23,32 +13,27 @@ 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.
+# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 72 GiB load can happen at
+# startup rather than on GPU time.
 import spaces
 import gradio as gr
 
 MODEL_REPO = os.environ.get("H3_MODEL_REPO", "MiniMaxAI/MiniMax-H3")
 CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
 # `lazy` moves all 72.16 GiB onto the card on the first GPU call and leaves it there; `offload` hands placement to
-# `ComponentsManager.enable_auto_cpu_offload` instead. Neither puts anything on the card at *startup*, which is
-# deliberate — see `load_models`: the 150 GB storage quota, not the 95 GiB card, is what rules that out here.
+# `ComponentsManager.enable_auto_cpu_offload`. Startup placement is not an option here — see `load_models`.
 PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
 # cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.
 # flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
 ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
 GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
-# Bounds on what `get_duration` may ask the pool to reserve. A request that runs out of GPU time is a total loss, so
-# the estimate is deliberately generous — but a *fixed* 900 s ceiling for every request is what makes the account hit
-# "too many ZeroGPU credits allocated to running tasks", because the pool reserves the number it is given.
+# Bounds on what `get_duration` may reserve. The pool reserves whatever number it is given, so a flat ceiling for every
+# request is what makes an account hit "too many ZeroGPU credits allocated to running tasks".
 MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
 MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))
 
-# MiniMax-H3's own canvases, i.e. `resolve_canvas_size` from
-# `diffusers.modular_pipelines.minimax_h3.modular_pipeline` evaluated for the six released aspect ratios. Hardcoded
-# so the UI renders before `diffusers` is importable.
-# Must stay identical to the conditioner's table: this Space forwards the *label* to the conditioner, so a canvas
-# that half does not know is rejected there and surfaces as a failure here.
+# Must stay identical to the conditioner's table: the *label* goes over the wire, so a canvas that half does not know
+# is rejected there and surfaces as a failure here.
 CANVASES = {
     # 16:9
     "960x544 · 16:9 fast": (544, 960),
@@ -73,41 +58,27 @@ CANVASES = {
     "1536x672 · 21:9 full": (672, 1536),
 }
 DEFAULT_CANVAS = "960x544 · 16:9 fast"
-# The examples carry their own canvas, and it is a full one. They are cached lazily and generated once, so what an
-# example is worth is its quality rather than its latency; an interactive request starts on the fast canvas.
 EXAMPLE_CANVAS, EXAMPLE_PORTRAIT_CANVAS = "1344x768 · 16:9 full", "768x1024 · 3:4 full"
 FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
-# 15 s is the checkpoint's ceiling, but it is the *snapped* frame count that has to hold for it: 15 s is 360 frames,
-# which rounds up to 362, i.e. 15.083 s, and is refused. 14 is the last whole second that survives the snap.
+# It is the *snapped* frame count the ceiling has to hold for: 15 s is 360 frames, which rounds up to 362, i.e.
+# 15.083 s, and is refused. 14 is the last whole second that survives the snap.
 MAX_UI_DURATION = 14
-MIN_DURATION = 5
-# A reference video shorter than 2 s gives the model almost no motion to read, and 15 s is the checkpoint's ceiling.
+MIN_DURATION = 2
+# A reference video shorter than 2 s gives the model almost no motion to read.
 MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
-# `MINIMAX_H3_MAX_REFERENCE_IMAGES`, hardcoded so the UI renders before `diffusers` is importable. The slots are all
-# built up front and revealed one at a time, because a demo asking for two subjects should not open with nine boxes.
+# `MINIMAX_H3_MAX_REFERENCE_IMAGES`. The slots are built up front and revealed one at a time, because a demo asking
+# for two subjects should not open with nine boxes.
 MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
 
-# --- What a request costs, for `get_duration` -------------------------------------------------------------------
-#
-# MiniMax-H3 attends over one packed sequence, so the cost of a step is a function of its length `S` alone. Fitted
-# on the `t2va` half (`minimax-h3`, AoTI on, same pool and same silicon) over two canvases at 124 frames:
-#
-#     544x544, S = 10693 -> 2.4 s/step
-#     960x544, S = 18870 -> 4.4 s/step
-#
-# through `s = LINEAR * S + QUADRATIC * S**2` — linear for the matmuls, quadratic for the attention. Checked against
-# two live `ref2va` requests on this Space, which is the regime the reference rows actually put it in:
-#
-#     one 1344x768 image reference,             S ~= 33232 -> 8.3 predicted, ~8.5 measured
-#     that image plus a 2.5 s video reference,  S ~= 54039 -> 14.6 predicted, ~16.1 measured
-#
-# so the fit holds to about 10% three times past the canvas it was taken from, and `SAFETY` covers the rest.
-STEP_LINEAR, STEP_QUADRATIC, SAFETY = 2.13e-4, 1.069e-9, 1.3
-# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call. Measured at ~45 s; every request
-# has to carry it, because nothing on this side knows whether the worker it lands on is cold.
+# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
+# matmuls, quadratic for the attention, against the AoTI block package this Space runs.
+STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
+# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call; every request carries it, because
+# nothing here knows whether the worker it lands on is cold.
 PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
 AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
 REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
+DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
 
 
 def snap_frames(seconds: float) -> int:
@@ -118,6 +89,13 @@ def snap_frames(seconds: float) -> int:
     return frames
 
 
+def lower_duration_floor(seconds: float = MIN_DURATION) -> None:
+    """Let the pipeline generate below its 5 s floor. 56 frames (2.33 s) is fine on the released checkpoint."""
+    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
+
+    MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
+
+
 def video_latent_frames(num_frames: int) -> int:
     """`17 * n + 5` frames become `5 * n + 2` video latents."""
     return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2
@@ -132,10 +110,9 @@ def target_rows(height: int, width: int, num_frames: int) -> int:
 def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
     """The rows the reference blocks add, from metadata alone — no decode.
 
-    This mirrors what `MiniMaxH3Ref2VASetupStep` and the reference encoder will do, closely enough to size a GPU
-    reservation with. An image is resized to a 2048 pixel short edge and encoded as a single frame; a video is put on
-    the canvas *its own* aspect ratio resolves to, truncated to the generated frame count and then snapped **down**
-    to a `17 * n + 5` the VAE encodes without padding; a soundtrack contributes two rows per 1/40 s.
+    An image is resized to a 2048 pixel short edge and encoded as a single frame; a video is put on the canvas *its
+    own* aspect ratio resolves to, truncated to the generated frame count and snapped **down** to a `17 * n + 5` the
+    VAE encodes without padding; a soundtrack contributes two rows per 1/40 s.
     """
     from PIL import Image
 
@@ -161,7 +138,6 @@ def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
                 stream = container.streams.video[0]
                 source_height, source_width = stream.height, stream.width
             canvas_height, canvas_width = resolve_canvas_size(source_width, source_height, CANVAS_MULTIPLE)
-            # Resampled onto 24 fps and capped at the generated length, then snapped down to `17 * n + 5`.
             frames = min(round(video_seconds * FPS), num_frames)
             snapped = max(1, (frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK
             rows += (
@@ -176,22 +152,16 @@ def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
 
 
 def get_duration(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, **_):
-    """Seconds of GPU to reserve for one request, from the packed sequence it is about to build.
-
-    Takes the arguments of the `@spaces.GPU` function it decorates — and tolerates the `gr.Progress` `spaces`
-    injects — so it can price the request rather than reserve a flat ceiling for all of them.
-
-    The text rows are exact: `text_token_tags` is the conditioner's own answer, already on this side. The reference
-    and target rows come from `reference_rows` and `target_rows`.
-    """
+    """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
+    tolerates the `gr.Progress` `spaces` injects."""
     sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
         height, width, num_frames
     )
     denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
-    # The two reference encoders, ahead of the loop, and the two decoders plus the mux after it. Both scale with what
+    # The two reference encoders ahead of the loop, and the two decoders plus the mux after it. Both scale with what
     # they are handed rather than with the step count.
     encode = 5 + reference_rows(references, num_frames) * 1e-3
-    decode = 15 + 25 * (height * width * num_frames) / (960 * 544 * 124)
+    decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
     total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10
     duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
     print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
@@ -204,22 +174,16 @@ LOAD_ERROR: str | None = None
 
 
 def load_models() -> str | None:
-    """Load the denoising half. At **startup**, but *not* onto the card.
-
-    `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, `scheduler`, `audio_scheduler`
-    and `video_processor` as its pretrained components (plus an `image_processor` built from config), so
-    `load_components` fetches exactly those subfolders out of the shared
-    `modular_model_index.json` — `text_encoder/` and the `transformer/` partition are never touched.
-
-    Both autoencoders carry `_keep_in_fp32_modules` over every module, so the `dtype` below is refused for them and
-    they stay float32: a bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
-
-    Nothing is moved onto the card here, which is the one place this Space departs from the ZeroGPU idiom, and the
-    reason is storage rather than memory. `spaces`' startup `torch.pack()` writes every startup-resident CUDA tensor
-    to a **second copy on disk** and only deletes the downloaded originals afterwards; 77.3 GB of weights plus a
-    77.3 GB pack is 154.6 GB against a 150 GB quota, and the Space is evicted mid-pack with `OSError: [Errno 28] No
-    space left on device` out of `os.posix_fallocate`. Placement therefore happens on the first GPU call, where it
-    costs about 10 s of PCIe and then persists across every later request in the same worker.
+    """Load the denoising half at startup, but *not* onto the card.
+
+    `MiniMaxH3Ref2VAGeneratorBlocks` declares `transformer_ref`, `vae`, `audio_vae`, the two schedulers and
+    `video_processor`, so `load_components` fetches exactly those subfolders — `text_encoder/` and the `transformer/`
+    partition are never touched. Both autoencoders carry `_keep_in_fp32_modules` over every module and stay float32: a
+    bfloat16 audio VAE decodes the soundtrack roughly 20 dB too quiet.
+
+    Nothing moves onto the card here, for storage rather than memory: `spaces`' startup `torch.pack()` writes every
+    startup-resident CUDA tensor to a second copy on disk, and 77.3 GB of weights plus its pack busts the 150 GB quota
+    (`OSError: [Errno 28] No space left on device` out of `os.posix_fallocate`, mid-pack).
     """
     global PIPE, MANAGER, LOAD_ERROR
 
@@ -233,37 +197,24 @@ def load_models() -> str | None:
 
         from h3_split_blocks import MiniMaxH3Ref2VAGeneratorBlocks
 
+        lower_duration_floor()
         manager = ComponentsManager()
         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")
         pipe.load_components(dtype=torch.bfloat16)
 
-        # Pin the two autoencoders to torch SDPA *before* the transformer takes cuDNN, and in that order.
-        #
-        # `set_attention_backend` does two things: it stamps the backend onto every attention processor of the model
-        # it is called on, and it sets the registry's **global** active backend, which every processor that was not
-        # stamped then falls through to. Both VAEs carry `AttentionModuleMixin` attention with `_attention_backend =
-        # None`, so stamping only the transformer leaves them inheriting cuDNN — and they are float32, for which
-        # cuDNN has no kernel:
-        #
-        #     RuntimeError: No available kernel. Aborting execution.    # audio_vae pre_block, is_causal=True
-        #
-        # It is `ref2va` that exposes this. The keyframe half only ever *decodes* audio, and the audio VAE's
-        # attention is on its encoder side, so nothing reached it until a reference brought a soundtrack along.
-        # Stamping the VAEs first leaves both explicitly on `native`; the transformer then stamps itself and takes
-        # the global with it, which no longer matters to anyone.
+        # Both VAEs first, and explicitly. `set_attention_backend` also sets the registry's *global* backend, which
+        # every processor that was not stamped falls through to, and the float32 audio VAE has no cuDNN kernel:
+        # `RuntimeError: No available kernel. Aborting execution.` in its causal encoder attention, which only a
+        # reference soundtrack ever reaches.
         pipe.vae.set_attention_backend("native")
         pipe.audio_vae.set_attention_backend("native")
         pipe.transformer_ref.set_attention_backend(ATTENTION)
 
-        # Still startup, still free: an AoTI package carries no weights and opens its compiled archive lazily inside
-        # the GPU worker, so pointing the 50-block stack at it is CPU work. Off unless `H3_AOTI=1`.
-        #
-        # It is the *same* package the `transformer/` partition runs, `bf16/torch2.11/sm120/dynamic`. Nothing about
-        # it is partition-specific: the two `config.json` files are identical field for field, and `LazyAOTIModel`
-        # binds each block's own live `state_dict()` by name on its first forward, so the compiled code carries no
-        # weights of either partition.
+        # Still startup, still free: an AoTI package carries no weights and opens its archive lazily inside the GPU
+        # worker. Off unless `H3_AOTI=1`. It is the *same* package the `transformer/` partition runs — the two configs
+        # are identical field for field and the compiled code carries no weights of either.
         import h3_aoti
 
         h3_aoti.maybe_load(pipe.transformer_ref)
@@ -286,9 +237,8 @@ def load_models() -> str | None:
 def _arm_decode_hooks(pipe):
     """Make the offload hooks fire for the two VAEs.
 
-    `enable_auto_cpu_offload` installs accelerate hooks, which wrap `forward`. The reference-encoder and decode
-    blocks call `components.vae.encode/decode(...)` and `components.audio_vae.encode/decode(...)` directly, so the
-    hook never runs and the VAE is still on the host when the latents arrive on the card.
+    `enable_auto_cpu_offload` wraps `forward`, and the reference-encoder and decode blocks call `vae.encode/decode(...)`
+    directly, so the hook never runs and the VAE is still on the host when the latents arrive on the card.
     """
     for name in ("vae", "audio_vae"):
         module = getattr(pipe, name)
@@ -306,8 +256,8 @@ def _arm_decode_hooks(pipe):
 
 @cache
 def conditioner():
-    """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."""
+    """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, so the
+    conditioner's booking is billed to whoever asked for the video."""
     from gradio_client import Client
 
     return Client(CONDITIONER_SPACE)
@@ -331,10 +281,8 @@ def probe(path: str) -> tuple[float | None, float | None]:
 def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
     """The `(kind, path)` references of a request, **in the order the model reads them**.
 
-    That order is semantic rather than cosmetic: it numbers the labels of MiniMax-H3's prompt presentation and it
-    advances the shared audio/video rotary clock, so the same references in a different order are a different
-    request. Images first, then a standalone audio clip, then the video — the order the tabs are laid out in, so
-    what the UI shows is what the model is handed.
+    That order numbers the labels of MiniMax-H3's prompt presentation and advances the shared audio/video rotary clock,
+    so the same references in a different order are a different request.
     """
     ordered = [("image", path) for path in image_paths if path]
     if audio_path:
@@ -345,11 +293,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 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.
-    """
+    """The `(kind, path)` references of a request as decoded reference dataclasses, in packed order. `from_file` brings
+    the rates along: a video its own frame rate and soundtrack, a clip its sample rate."""
     from diffusers.modular_pipelines.minimax_h3 import (
         MiniMaxH3AudioReference,
         MiniMaxH3ImageReference,
@@ -373,16 +318,13 @@ def audio_bearing(references: list[tuple[str, str]]) -> list[tuple[str, float]]:
 
 
 def duration_controls(audio_path, video_path, match: bool):
-    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out.
-
-    Only the audio and video slots matter here: an image reference never carries a waveform.
-    """
+    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out."""
     try:
         carried = audio_bearing(collect([], audio_path, video_path))
     except Exception:
         carried = []
-    # Exactly one soundtrack, and one long enough to be a duration MiniMax-H3 generates. Anything else and the
-    # request is ambiguous or out of range, so the slider stays and nothing is derived.
+    # Exactly one soundtrack, long enough to be a duration MiniMax-H3 generates; anything else is ambiguous or out of
+    # range and the slider stays.
     derivable = len(carried) == 1 and MIN_DURATION <= snap_frames(carried[0][1]) / FPS <= MAX_REFERENCE_VIDEO
     return gr.update(visible=derivable), gr.update(visible=not (derivable and match))
 
@@ -409,16 +351,11 @@ def check(prompt: str, references: list[tuple[str, str]]) -> 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.
+    """`/encode_ref2va` on the conditioner Space: a safetensors file holding `prompt_embeds` + `text_token_tags`, with
+    the resolved `height` / `width` / `num_frames` in its metadata, plus the plan.
 
-    The references go over with the request: `ref2va`'s presentation puts a vision block in front of the prompt for
-    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` 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.
+    `canvas` is the label. `media` and `kinds` are parallel and ordered, and the references go over because `ref2va`'s
+    presentation puts a vision block in front of the prompt for every image and every merged video frame pair.
     """
     from gradio_client import handle_file
     from safetensors import safe_open
@@ -440,18 +377,13 @@ def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
 def _generate(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed):
     """The only thing on GPU time: the two reference encoders, the packed-sequence denoise loop and the decoders.
 
-    The references are built here rather than handed over already decoded. A `@spaces.GPU` argument crosses a
-    process boundary by pickling, and a 5 s 1344x768 reference video is 370 MB of frames once PyAV has expanded it;
-    the file path is a few bytes and the decode is CPU work either way.
-
-    Only the three generated outputs come back, for the same reason: the full `PipelineState` still holds the packed
-    latents, the rotary grid and the row indices on the card.
+    References cross as paths and are decoded here; only the three generated outputs come back. A `@spaces.GPU`
+    argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is 370 MB of expanded frames, and
+    the full `PipelineState` still holds the packed latents and the rotary grid on the card.
     """
     import torch
 
     if PLACEMENT == "lazy":
-        # 72.16 GiB across PCIe on the first request of a worker, a no-op walk on every one after it. Startup
-        # placement is not an option here — see `load_models` — and this is what buys the offload-free denoise loop.
         PIPE.to("cuda")
 
     state = PIPE(
@@ -468,9 +400,9 @@ def _generate(prompt_embeds, text_token_tags, references, height, width, num_fra
 
 
 def generate(
-    # The first four are the columns `gr.Examples` varies, and they lead the signature for that reason: an example
-    # row is applied to `inputs` positionally, so the exampled components have to be the leading parameters. Every
-    # parameter has a default, which is what lets a four-column row call this at all.
+    # The first four are the columns `gr.Examples` varies, and they lead the signature for that reason: an example row
+    # is applied to `inputs` positionally. Every parameter has a default, which is what lets a four-column row call
+    # this at all, and `upsample` is last so a positional API client that predates it is unaffected.
     prompt,
     image_1=None,
     audio_path=None,
@@ -491,7 +423,7 @@ def generate(
     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."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
     if PIPE is None:
@@ -518,7 +450,7 @@ def generate(
         raise
     except Exception as error:
         # gradio only puts the exception *type* on the wire, so the useful half of a conditioner-side failure is in
-        # that Space's logs. Print what did come back here, name the Space, and say where the rest of it is.
+        # that Space's logs.
         traceback.print_exc()
         raise gr.Error(
             f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
@@ -557,8 +489,8 @@ INTRO = """# MiniMax-H3 Reference
 
 <div align="center">
   <a href="https://huggingface.co/MiniMaxAI/MiniMax-H3"><strong>[ model ]</strong></a> &nbsp;
-  <a href="PAPER_URL_PLACEHOLDER"><strong>[ paper ]</strong></a> &nbsp;
-  <a href="https://www.minimax.io"><strong>[ project ]</strong></a>
+  <a href="https://www.minimax.io/blog/minimax-h3"><strong>[ blog ]</strong></a> &nbsp;
+  <a href="https://huggingface.co/spaces/multimodalart/minimax-h3"><strong>[ text / image to video ]</strong></a>
 </div>
 
 **MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a
@@ -582,20 +514,19 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
                 value="The character walks through a neon-lit street in the rain, humming to themselves",
             )
             upsample = gr.Checkbox(label="Upsample prompt", value=False)
-            # One tab per modality, in the order the model reads them. A reference left in a tab that is not the
-            # open one is still part of the request — the tabs lay the slots out, they do not choose between them.
+            # One tab per modality, in the order the model reads them. A reference left in a tab that is not the open
+            # one is still part of the request.
             with gr.Tabs():
                 with gr.Tab("Images"):
-                    # One `gr.Row`, so gradio splits the width evenly and wraps once the slots hit `min_width`
-                    # rather than leaving a hole where a hidden slot used to be.
+                    # One `gr.Row`, so gradio splits the width evenly and wraps at `min_width` rather than leaving a
+                    # hole where a hidden slot used to be.
                     with gr.Row():
                         images = [
                             gr.Image(
                                 label="Subject, style or scene",
                                 type="filepath",
                                 min_width=180,
-                                # Fixed, so a row that wraps to a single slot stays the same size as a full one
-                                # instead of stretching to the width of the column.
+                                # Fixed, so a row that wraps to a single slot stays the size of a full one.
                                 height=210,
                                 visible=index < OPEN_IMAGE_SLOTS,
                             )
@@ -639,8 +570,7 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
             duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
         )
 
-    # Same order as `generate`'s signature: the exampled five first, then the remaining image slots. `upsample` is
-    # last and defaults off, so a positional API client that predates it is unaffected.
+    # Same order as `generate`'s signature: the exampled five first, then the remaining image slots.
     request = [prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample]
 
     gr.Examples(

h3_aoti.py

@@ -1,59 +1,6 @@
-"""ZeroGPU AoTI for MiniMax-H3: compile the repeated transformer block once, reuse the package forever.
-
-Shared byte-identically by every MiniMax-H3 Space. A Space only ever calls `maybe_load()`; the compile path runs from
-the debug Space's "Compile (AoTI)" tab, or off-Space from `job_bf16_aoti.py` on an `rtx-pro-6000` Job, and pushes its
-artifacts to `multimodalart/minimax-h3-aoti` under `<width>/torch<X.Y>/sm<cc>/<shape>`.
-
-What is measured, so nobody has to guess whether this is worth turning on. Unquantized bfloat16, 124 frames,
-everything resident, one dynamic-sequence package serving every row — on an RTX PRO 6000 Blackwell, torch 2.11,
-cuDNN attention:
-
-    canvas (HxW)   eager s/step   AoTI s/step   saved     faster
-    768x1344          10.20          9.73       0.47 s     +4.6%
-    704x1280           8.59          7.87       0.72 s     +8.4%
-    640x1152           6.46          5.88       0.59 s     +9.1%
-    576x1024           4.74          4.24       0.50 s    +10.5%
-    544x960            4.02          3.58       0.44 s    +11.0%
-
-Read the *absolute* column: AoTI removes a near-constant ~0.5 s/step no matter how big the canvas is. That is exactly
-the shape of what it can remove — 50 blocks' worth of kernel-launch overhead and the norm / rotary / AdaLN-gather
-epilogues around the matmuls. It cannot touch the matmuls themselves, and at S = 37726 one block is ~70 TFLOP of GEMM
-and attention, so the released 768x1344 canvas is compute bound and only 4.6% comes back. The smaller the default
-canvas gets, the better this pays.
-
-The trap that cost a day, recorded here because the symptom is a segfault with no message: a **shallow clone exported
-in torch.export's default non-strict mode duplicates every weight** — once as a named `PARAMETER` and once as an
-anonymous `lifted_tensor_<N>` `CONSTANT_TENSOR` aliasing the same storage — and `LazyAOTIModel` binds constants by
-name, so the anonymous half binds to nothing and the compiled kernel reads pointers nobody set. It is neither
-accelerate's offload hooks nor torchao's tensor subclasses, which were both blamed first; it reproduces in plain
-bfloat16 with no subclass anywhere, and it goes away with `strict=True`. See `export_block`.
-
-Why block level rather than the whole transformer: `MiniMaxH3Transformer3DModel.forward` decides whether the packed
-sequence needs a padding attention mask with `bool(is_pad.any())`, a data-dependent branch `torch.export` cannot
-trace. One `MiniMaxH3TransformerBlock` is where all the time goes anyway (50 of them per step), and the sequence
-length is the only thing that changes between requests, which a single dynamic dimension covers. The block's fifth
-argument, `attention_mask`, is always `None` in practice — `packing.py` never emits a padding row, so `token_tags` is
-never negative — which is what makes one static signature enough.
-
-What `spaces` 0.51.1 actually provides (checked against the installed package, not the klein-era blog post):
-
-    spaces.aoti_capture(module)                       context manager, grabs the args of the next call and aborts it
-    spaces.aoti_compile(exported_program, configs)    in-process compile, returns a ZeroGPUCompiledModel
-    spaces.aoti_compile_and_save(dir, ep, configs, submodule=...)
-                                                      compile and write `<dir>/submodules/<submodule>/package.pt2`
-    spaces.aoti_apply(compiled, module)               in-process apply
-    spaces.aoti_patch(module, LazyAOTIModel)          apply a package to one module, weights stay live
-    spaces.aoti_load_from_package_dir(module, dir)    walk `<dir>/{root,submodules/*}` and patch, ModuleList aware
-    spaces.aoti_load(module, repo_id, ...)            the convenience wrapper — NOT usable here: it hardcodes
-                                                      `snapshot_download(allow_patterns="package/*")` on a *model*
-                                                      repo, and these artifacts are keyed by quant/torch/arch under a
-                                                      dataset, so the download is done here and only the loader
-                                                      (`aoti_load_from_package_dir`) is reused.
-    spaces.aoti_blocks_load(module, repo_id, variant) the `_repeated_blocks` convenience — same repo-layout mismatch.
-
-Weights are *not* baked into the package: `aoti_patch` binds the block's live `state_dict()`, so one package serves all
-50 blocks, and the quantized weights it reads are whatever the block holds. Which is also why quantization has to
-happen *before* the export — the same ordering constraint as fusing a LoRA before AoTI.
+"""ZeroGPU AoTI for MiniMax-H3: one compiled `MiniMaxH3TransformerBlock` package, reused by all 50 blocks.
+
+Shared byte-identically by every MiniMax-H3 Space. A Space only calls `maybe_load()`; the rest is the build path.
 """
 
 from __future__ import annotations
@@ -62,46 +9,22 @@ import os
 from pathlib import Path
 
 AOTI = os.environ.get("H3_AOTI", "0") == "1"
-# 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:
-# a dynamic sequence dimension exports and compiles cleanly (measured on an rtx-pro-6000 Job, torch 2.11). It has to be
-# dynamic to be useful at all — `build_packed_sequence` pads nothing, so
-# `S = num_text_tokens + condition_rows + audio_rows + video_rows` moves with the *prompt* as well as the canvas, and a
-# static package would only ever serve the one prompt length it was captured from.
-#
-# A `HxWxF` value instead pins the artifact to one static shape. That is the fallback for a width whose dynamic export
-# is refused, which is what the NVFP4 attempt hit: export rejected the dimension and offered only the affine
-# `S = 128 * k - 34` it had derived from that one capture — an offset that is a property of one canvas *and* one prompt
-# rather than of the model, so a package built that way serves almost nothing. The 128 is the alignment torchao's
-# scaled matmuls want, though that has not been re-verified since the bfloat16 path was proven, and a static package is
-# only worth building for a width that has been shown to need one.
+# A package is valid for exactly one `<width>/torch<X.Y>/sm<cc>/<shape>`, and a mismatched one segfaults rather than
+# raising, so `maybe_load` refuses anything but this key.
+AOTI_KEY = os.environ.get("H3_AOTI_KEY", "bf16/torch2.11/sm120/dynamic")
+# `dynamic` is the sequence dimension: `build_packed_sequence` pads nothing, so `S` moves with the prompt as well as
+# the canvas and a static package would serve one prompt length.
 AOTI_SHAPE = os.environ.get("H3_AOTI_SHAPE", "dynamic")
 AOTI_DURATION = int(os.environ.get("H3_AOTI_DURATION", "1500"))
 
-# The 50-deep stack that is the whole cost of a step. `MiniMaxH3TokenRefinerBlock` is also in `_repeated_blocks` but
-# runs a handful of text rows a couple of times per step, so it is left eager.
+# Where a step spends its time. `MiniMaxH3TokenRefinerBlock` is also repeated but runs a handful of text rows.
 BLOCK_CONTAINER = "transformer_blocks"
 
-# Height of the AdaLN table baked into the package. `temb` is `(num_distinct_timesteps, time_embed_dim)` and the block
-# gathers from a `3 * num_distinct_timesteps` table with `adaln_indices = timestep_indices * 3 + token_tag`, so the
-# table's height is part of the compiled shape. It is not constant at runtime: at step 0 the video and audio streams
-# share a noise level and `temb` has a single row, and from step 1 their sigma schedules diverge and it grows one.
-# Exporting whatever the first call happened to show bakes in a 3-row table and the later steps then walk off it:
-#
-#     Assertion `index out of bounds: 0 <= tmp22 < 3` failed
-#
-# A dynamic dimension is the wrong tool — `torch.export` specializes size-1 dimensions unconditionally, so a `Dim`
-# taken from a 2-row capture carries a `>= 2` guard that step 0 violates. Instead `temb` is padded to a fixed height
-# on both sides of the compile. Rows past the live ones are never gathered, so the output is unchanged, and the shape
-# becomes a constant. Must match the `H3_AOTI_TEMB_ROWS` the package was compiled with.
-#
-# 4 is what the published bfloat16 packages were built with, and the padding is *validated* rather than assumed: the
-# build job replays a real 1-row (step 0) call and a real 2-row (step 1) call through the compiled block and diffs both
-# against eager. Two streams at two noise levels is the realistic maximum, so 4 is loose on purpose, and the cost is
-# one slightly taller AdaLN projection per block against the block's own 70 TFLOP.
+# Height of the AdaLN table baked into the package. `temb` grows from 1 row (step 0, both streams at one noise level)
+# to 2 (from step 1, sigmas diverged), and the block gathers from `3 * rows`, so the row count is part of the compiled
+# shape and is pinned by padding on both sides of the compile. Must match the package's `H3_AOTI_TEMB_ROWS`.
 TEMB_ROWS = int(os.environ.get("H3_AOTI_TEMB_ROWS", "4"))
 
 _LOADED: set[int] = set()
@@ -123,12 +46,7 @@ def pad_temb(temb, rows: int = TEMB_ROWS):
 
 
 def width() -> str:
-    """Which transformer these artifacts belong to: `bf16`, `fp8`, `nvfp4`, ...
-
-    `H3_WIDTH` wins, so a Space that has no `h3_core` — the unquantized split deployment is two standalone Spaces —
-    can use this module by setting one variable. Otherwise it comes from `h3_core`, which derives it from `H3_QUANT`
-    or from the pre-quantized repository's suffix.
-    """
+    """Which transformer these artifacts belong to: `bf16`, `fp8`, `nvfp4`, ..."""
     if explicit := os.environ.get("H3_WIDTH"):
         return explicit.lower()
     try:
@@ -139,12 +57,15 @@ def width() -> str:
         return "bf16"
 
 
-def artifact_key() -> str:
-    """`<width>/torch<X.Y>/sm<cc>/<shape>` — an AoTI package is valid for exactly one of each."""
-    import torch
+def artifact_key() -> str | None:
+    """`<width>/torch<X.Y>/sm<cc>/<shape>` of the card this process is on, or `None` when there is no CUDA."""
+    try:
+        import torch
 
-    torch_version = ".".join(torch.__version__.split(".")[:2])
-    major, minor = torch.cuda.get_device_capability()
+        torch_version = ".".join(torch.__version__.split(".")[:2])
+        major, minor = torch.cuda.get_device_capability()
+    except Exception:
+        return None
     return f"{width()}/torch{torch_version}/sm{major}{minor}/{AOTI_SHAPE}"
 
 
@@ -159,32 +80,15 @@ def status() -> str:
 def patch_blocks(transformer, package_dir) -> None:
     """Point all 50 blocks at the one compiled package, binding each block's own weights on its first call.
 
-    This is `spaces.aoti_load_from_package_dir` with two changes, both forced by how this Space runs.
-
-    *Weights are read on the first forward, not at patch time.* `spaces.aoti_patch` snapshots `state_dict()` when it
-    patches, and `maybe_load` runs at startup, while the components are still on the host — `place()` only moves them
-    on the first request, because ZeroGPU cannot pack a startup-resident `Float8Tensor` (it packs CUDA tensors with
-    `aten.empty_like(..., pin_memory=True)`, which the subclass does not implement). `Module.to` rebinds `param.data`
-    to a fresh CUDA tensor, so a snapshot taken at startup keeps pointing at the host copies and the compiled block
-    would run against host memory. Reading the state dict on the first call instead picks it up wherever it now is.
-
-    *`temb` is padded on the way in*, to the fixed height the package was exported with — see `TEMB_ROWS`.
-
-    The clone-and-flatten here is `spaces.aoti_patch`'s own preparation, kept because a quantized width needs it: the
-    names it produces are the FQNs the package's constants were derived from. For an unquantized block it is a no-op
-    and the resulting names are exactly `blocks[0].state_dict()`, which is what `export_block` exported — the two
-    sides agree either way. What must *not* be mirrored is the clone on the export side under non-strict tracing; see
-    `export_block` for why that is the difference between a working package and a SIGSEGV.
+    `spaces.aoti_load_from_package_dir` with two changes. Weights are read on the first forward rather than at patch
+    time, because this runs at startup and `Module.to` later rebinds `param.data` to fresh CUDA tensors. And `temb` is
+    padded to the height the package was exported with — see `TEMB_ROWS`.
     """
     from spaces.zero.torch.aoti import LazyAOTIModel, _shallow_clone_module
     from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters
 
-    # `LazyAOTIModel` binds constants by intersecting `state_dict()` with `get_constant_fqns()` and
-    # silently keeps whatever it does not match, so a package whose constants were lifted anonymously
-    # binds nothing and the compiled block then dereferences constants nobody set — a SIGSEGV. This patch
-    # resolves those names through the `constant_aliases.json` the compile side writes, and **raises** a
-    # readable error if it still cannot. Purely protective: with a well-formed package it changes nothing,
-    # which is why a missing sidecar module is a warning rather than a failure.
+    # `LazyAOTIModel` binds constants by name and silently keeps what it cannot match, which is a SIGSEGV rather than
+    # an error. The patch resolves anonymous names through the compile side's sidecar and raises if it still cannot.
     try:
         from spaces_constant_binding_patch import apply_spaces_constant_binding_patch
 
@@ -210,31 +114,41 @@ def patch_blocks(transformer, package_dir) -> None:
 
 
 def maybe_load(transformer) -> None:
-    """Patch the block stack with its compiled package. Once, and safe to call at **startup**.
+    """Patch the block stack with its compiled package, or leave it eager. Safe to call at **startup**.
 
-    Nothing here touches a GPU: the download is CPU work and the `.pt2` archive is not opened until the first forward,
-    which happens inside the `@spaces.GPU` call. Proven on the pool: `bf16/torch2.11/sm120/dynamic` loads at startup,
-    patches all 50 blocks, and generates. The repo is public, so no token is passed for it.
+    Off unless `H3_AOTI=1`, and anything that does not line up — another card, another torch, no `spaces` AoTI
+    helpers, no published package — falls back to eager with one line rather than raising or segfaulting. Nothing here
+    touches a GPU: the download is CPU work and the `.pt2` is not opened until the first forward.
     """
     if not AOTI or id(transformer) in _LOADED:
         return
 
-    import spaces
-    from huggingface_hub import snapshot_download
-
     key = artifact_key()
+    if key is None:
+        print("[h3-aoti] no CUDA device visible; running eager", flush=True)
+        return
+    if key != AOTI_KEY:
+        print(f"[h3-aoti] this card wants `{key}`, only `{AOTI_KEY}` is published; running eager", flush=True)
+        return
+
+    try:
+        from huggingface_hub import snapshot_download
+        from spaces.zero.torch.aoti import LazyAOTIModel  # noqa: F401
+    except Exception as error:
+        print(f"[h3-aoti] no AoTI loader here ({type(error).__name__}: {error}); running eager", flush=True)
+        return
+
     print(f"[h3-aoti] loading {AOTI_REPO}:{key} ...", flush=True)
-    local = snapshot_download(
-        repo_id=AOTI_REPO,
-        repo_type=AOTI_REPO_TYPE,
-        allow_patterns=f"{key}/package/*",
-    )
+    try:
+        local = snapshot_download(repo_id=AOTI_REPO, repo_type=AOTI_REPO_TYPE, allow_patterns=f"{key}/package/*")
+    except Exception as error:
+        print(f"[h3-aoti] {AOTI_REPO}:{key} unreachable ({type(error).__name__}: {error}); running eager", flush=True)
+        return
     package_dir = Path(local) / key / "package"
     if not package_dir.is_dir():
-        raise RuntimeError(
-            f"No AoTI package at `{AOTI_REPO}:{key}/package`. Run the debug Space's Compile (AoTI) tab on this card "
-            f"with this `H3_QUANT`, or set `H3_AOTI=0`."
-        )
+        print(f"[h3-aoti] no package at `{AOTI_REPO}:{key}/package`; running eager", flush=True)
+        return
+
     patch_blocks(transformer, package_dir)
     _LOADED.add(id(transformer))
     print(f"[h3-aoti] compiled blocks in place (temb padded to {TEMB_ROWS} rows)", flush=True)
@@ -243,8 +157,8 @@ def maybe_load(transformer) -> None:
 def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
     """Capture one block call out of a real request and export it with a dynamic sequence dimension.
 
-    Runs on the GPU, after the transformer has been quantized and moved there: the export traces the quantized module,
-    and a package compiled for one quantization mode is meaningless for another.
+    Runs on the GPU, after the transformer has been quantized and moved there: a package compiled for one
+    quantization mode is meaningless for another.
     """
     import torch
     import spaces
@@ -254,9 +168,8 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
     transformer = h3.transformer_of(pipe)
     blocks = getattr(transformer, BLOCK_CONTAINER)
 
-    # Record every call the block receives over a short real run and keep the widest `temb`, rather than
-    # `spaces.aoti_capture`'s first-call-then-abort. The first call is the unrepresentative one: see `TEMB_ROWS`.
-    # Text encoding and packing run either way, which is the point — these are the real inputs.
+    # Keep the widest `temb` over a short real run rather than `spaces.aoti_capture`'s first call, which is the
+    # 1-row one — see `TEMB_ROWS`.
     original_forward = blocks[0].forward
     widest = {"args": (), "kwargs": {}, "rows": -1}
     seen = []
@@ -285,16 +198,9 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
         raise RuntimeError("Nothing was captured — the transformer block was never called.")
     print(f"[h3-aoti] temb rows seen: {sorted(set(seen))}; exporting with {TEMB_ROWS} (padded)", flush=True)
 
-    # `block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask)`:
-    #   hidden_states   (1, S, hidden)
-    #   temb            (num_distinct_timesteps, time_embed_dim)
-    #   adaln_indices   (S,)
-    #   rotary_emb      ((S, dim), (S, dim))
-    #   attention_mask  None for a padless sequence, which is what these pipelines build
-    #
-    # Only the sequence is asked for. `temb`'s row count is held constant by padding instead (see `TEMB_ROWS`), which
-    # is both cheaper to reason about and the only thing that works: `torch.export` specializes size-1 dimensions
-    # unconditionally, so a `Dim` on a dimension that is 1 at step 0 cannot be expressed at all.
+    # `block(hidden_states, temb, adaln_indices, rotary_emb, attention_mask)`, `attention_mask` being `None` for the
+    # padless sequences these pipelines build. Only the sequence is dynamic: `torch.export` specializes size-1
+    # dimensions unconditionally, so a `Dim` on `temb`'s rows cannot be expressed at all.
     if AOTI_SHAPE == "dynamic":
         sequence = torch.export.Dim("sequence", min=2048, max=262144)
         dynamic_shapes = ({1: sequence}, None, {0: sequence}, ({0: sequence}, {0: sequence}), None)
@@ -302,34 +208,12 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
     else:
         dynamic_shapes = None
 
-    # `temb` to its fixed height, so the AdaLN table the package bakes in is the one `maybe_load` will feed it.
     args = (call.args[0], pad_temb(call.args[1]), *call.args[2:])
 
-    # WHICH MODULE, AND WHICH EXPORT MODE. This is the whole difference between a working package and a SIGSEGV.
-    #
-    # `spaces.aoti_patch` prepares the load side by shallow-cloning the module and flattening any tensor subclass, and
-    # the received wisdom is to do the identical thing before exporting so both sides derive the same constant FQNs.
-    # For a subclass that is genuinely necessary: inductor's constant handling wraps a constant back into
-    # `torch.nn.Parameter`, which rejects a non-floating dtype ("Only Tensors of floating point and complex dtype can
-    # require gradients"), so `Float8Tensor` / `NVFP4Tensor` parameters have to be flattened first.
-    #
-    # But a shallow clone exported in `torch.export`'s **default non-strict** mode duplicates every weight: the same
-    # tensor comes out once as a named `PARAMETER` and again as an anonymous `lifted_tensor_<N>` `CONSTANT_TENSOR`,
-    # 12 of each for this block, 1.2 GiB of them, `data_ptr()` proving the two sets alias. `LazyAOTIModel` binds by
-    # name, so the anonymous half binds to nothing and the compiled block dereferences constants nobody set. That is
-    # the crash that stalled this work, blamed first on accelerate's offload hooks and then on torchao's subclasses;
-    # it is neither. Measured on an rtx-pro-6000 Job at full size, torch 2.11, plain bfloat16, no subclass anywhere:
-    #
-    #     live block,    non-strict   12 PARAMETER,  0 CONSTANT_TENSOR      <- what the shipped bf16 package used
-    #     live block,    strict       12 PARAMETER,  0 CONSTANT_TENSOR
-    #     shallow clone, non-strict   12 PARAMETER, 12 CONSTANT_TENSOR      <- the bug
-    #     shallow clone, strict       12 PARAMETER,  0 CONSTANT_TENSOR
-    #
-    # So: export the **live block** whenever it has no subclass parameters to flatten, which is every unquantized
-    # width and, per the fp8 investigation, quantized ones whose weights are still registered parameters. Only fall
-    # back to the clone when flattening is actually needed, and then in `strict` mode, which is also clean. The clone
-    # is safe to skip for the live-block path precisely because there is nothing to unwrap: `state_dict()` names are
-    # then identical on both sides by construction.
+    # Export the **live** block, non-strict. A shallow clone under non-strict tracing lifts every weight twice — once
+    # named, once as an anonymous `CONSTANT_TENSOR` aliasing it — and the loader binds by name, so the compiled block
+    # dereferences constants nobody set. The clone is only for flattening tensor-subclass parameters, which inductor's
+    # constant handling cannot wrap back into a `Parameter`, and it needs `strict=True`.
     from spaces.zero.torch.aoti import _shallow_clone_module
     from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters
 
@@ -344,13 +228,9 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
         strict = False
         print("[h3-aoti] plain parameters: exporting the live block, non-strict", flush=True)
 
-    # `torch.export` only gives a lifted tensor a real FQN when it is a registered parameter or buffer; a tensor
-    # reached through a plain attribute becomes an anonymous constant the loader can never match against
-    # `state_dict()`. Re-registering such tensors as buffers is numerics-preserving — it changes how a tensor is
-    # registered, never the tensor and never the forward. A well-formed block has none, and this returns empty.
-    # Only ever on the clone: `register_loose_tensors` *re-registers* attributes, so running it on the live block would
-    # mutate the model the eager path uses. A `MiniMaxH3TransformerBlock` has no loose tensor attributes, so this is
-    # empty in practice and the live-block export needs nothing; if that ever changes, the warning below catches it.
+    # `torch.export` only gives a lifted tensor a real FQN when it is a registered parameter or buffer; a plain
+    # attribute becomes an anonymous constant the loader can never match. Only ever on the clone, since this
+    # re-registers attributes and the live block is what the eager path runs.
     if block is not blocks[0]:
         try:
             from spaces_constant_binding_patch import register_loose_tensors
@@ -364,8 +244,6 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
     try:
         exported = torch.export.export(block, args, call.kwargs or None, dynamic_shapes=dynamic_shapes, strict=strict)
     except Exception as error:
-        # Dynamo refuses some modules it cannot trace. Non-strict is still worth attempting, with the duplication
-        # reported loudly below rather than left to segfault at load time.
         if not strict:
             raise
         print(f"[h3-aoti] strict export failed ({type(error).__name__}: {error}); retrying non-strict", flush=True)
@@ -377,8 +255,7 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
     if anonymous:
         print(
             f"[h3-aoti] WARNING {len(anonymous)} constants lifted anonymously: {anonymous[:6]}. The loader binds by "
-            f"name, so this package will not bind them; `compile_and_save` writes the alias sidecar and "
-            f"`patch_blocks` raises rather than letting it segfault.",
+            f"name, so `compile_and_save` writes the alias sidecar and `patch_blocks` raises rather than segfaulting.",
             flush=True,
         )
     return exported
@@ -387,9 +264,8 @@ def export_block(pipe, height: int, width: int, num_frames: int, prompt: str):
 def compile_and_save(exported_program, destination: str | os.PathLike[str]) -> Path:
     """Inductor-compile the exported block into `<destination>/package/submodules/transformer_blocks/package.pt2`.
 
-    That layout is what `aoti_load_from_package_dir` walks: it resolves the submodule name to the transformer's
-    `transformer_blocks` `ModuleList` and, because a `ModuleList` is iterable, patches every one of the 50 blocks with
-    this single package.
+    That layout is what `aoti_load_from_package_dir` walks, resolving the submodule name to the transformer's
+    `transformer_blocks` `ModuleList` and patching every block in it with this one package.
     """
     import spaces
 
@@ -397,9 +273,8 @@ def compile_and_save(exported_program, destination: str | os.PathLike[str]) -> P
     print("[h3-aoti] inductor compile (minutes) ...", flush=True)
     spaces.aoti_compile_and_save(package_dir, exported_program, submodule=BLOCK_CONTAINER)
 
-    # The compiled artifact keeps a constant's dtype, shape and slot index but drops its FQN when the
-    # export lifted it anonymously. The `ExportedProgram` still has the real names, so record the
-    # mapping next to the package while it is still available; the loader reads it back.
+    # The compiled artifact drops a constant's FQN when the export lifted it anonymously; the `ExportedProgram` still
+    # has the real names, so record the mapping for the loader while it is available.
     try:
         from spaces_constant_binding_patch import write_constant_aliases
 
@@ -414,7 +289,7 @@ def compile_and_save(exported_program, destination: str | os.PathLike[str]) -> P
 
 
 def upload(package_dir: str | os.PathLike[str], key: str) -> str:
-    """Push the package under its `<quant>/torch<X.Y>/sm<cc>/<shape>` key. CPU work — never inside GPU time."""
+    """Push the package under its `<width>/torch<X.Y>/sm<cc>/<shape>` key. CPU work — never inside GPU time."""
     from huggingface_hub import HfApi
 
     token = os.environ.get("HF_TOKEN")

h3_split_blocks.py

@@ -1,38 +1,15 @@
 """The halves of a **split** MiniMax-H3 deployment, for both of its checkpoint partitions.
 
-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` 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
-partition:
-
-  * `MiniMaxH3ConditionerBlocks` = `[resize, text_encoder]` — loads `text_encoder` / `tokenizer` / `processor` only
-    (plus the `image_processor`, which is built from config and downloads nothing), and emits `prompt_embeds` +
-    `text_token_tags`, which is the whole wire format between the two halves.
-  * `MiniMaxH3GeneratorBlocks` = everything else — loads `transformer` / `vae` / `audio_vae` / the two schedulers
-    only, and takes `prompt_embeds` + `text_token_tags` as *inputs*.
-  * `MiniMaxH3Ref2VAConditionerBlocks` / `MiniMaxH3Ref2VAGeneratorBlocks` are the same cut through the `ref2va`
-    branch, so one conditioner Space serves both partitions out of the weights it already holds.
-
-`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.
+MiniMax-H3 is 195.9 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage, so `MiniMaxH3Blocks` is cut
+at its `text_encoder` step: the 62.14 GiB Qwen3-VL runs in the conditioner Space, everything else in a generator
+Space, and `prompt_embeds` + `text_token_tags` is the whole wire format between them.
+
+`resize` / `setup` run on **both** sides: they own no pretrained component, and each half needs the canvas and the
+prepared keyframes or normalized references. Both conditioner halves also return the resolved `height` / `width` /
+`num_frames`, which the generating half pins rather than re-deriving.
+
+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 on the denoising side of the cut.
 """
 
 from diffusers.modular_pipelines.minimax_h3.before_encoder import MiniMaxH3Ref2VASetupStep
@@ -55,10 +32,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 declared by the `ref2va` half alone: it is the one whose setup step resolves a frame count.
-    """
+    """The wire format of the split. `num_frames` is declared by the `ref2va` half alone, whose setup resolves one."""
     return [
         OutputParam.template("prompt_embeds"),
         OutputParam("text_token_tags", description="The per-row modality tag of every row of `prompt_embeds`."),
@@ -121,12 +95,9 @@ class MiniMaxH3GeneratorBlocks(SequentialPipelineBlocks):
 class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
     """The conditioner half of a split `ref2va`: the resolved plan plus the Qwen3-VL read at its 50th layer.
 
-    Component for component this is `MiniMaxH3ConditionerBlocks` — `text_encoder`, `tokenizer`, `processor` — which
-    is what lets one conditioner Space serve both partitions of the checkpoint out of the weights it already holds.
-    What differs is the presentation the Qwen3-VL is shown: `ref2va` prepends a label per reference, numbered per
-    modality, and a vision block per image and per merged video frame pair, so the references themselves have to
-    reach this half. An audio reference never does — it contributes its `"<Audio j>: "` label and nothing else — but
-    it is still part of the request here, because the setup step normalizes every soundtrack and validates the mix.
+    Component for component this is `MiniMaxH3ConditionerBlocks`, so one conditioner Space serves both partitions.
+    What differs is the presentation: `ref2va` prepends a label per reference and a vision block per image and per
+    merged video frame pair, so the references themselves have to reach this half.
     """
 
     model_name = "minimax-h3"
@@ -149,9 +120,8 @@ class MiniMaxH3Ref2VAConditionerBlocks(SequentialPipelineBlocks):
 class MiniMaxH3Ref2VAGeneratorBlocks(SequentialPipelineBlocks):
     """The denoising half of a split `ref2va`: the `ref2va` branch with its `text_encoder` step removed.
 
-    Only the text-encoder step is dropped. `reference_encoder` is this half's own encoder — it runs the video VAE
-    over the image and video references and the audio VAE over the soundtracks, and its output shapes are where every
-    reference block's geometry in the packed layout comes from — so it stays here, next to the autoencoders.
+    `reference_encoder` stays here, next to the two autoencoders it runs: its output shapes are where every reference
+    block's geometry in the packed layout comes from.
     """
 
     model_name = "minimax-h3"

requirements.txt

@@ -5,27 +5,23 @@
 # head — whenever the PR updates.
 #
 #   665f578278365ea4a3318cb8c9b66ce6c01204b9 = refs/pull/14371/head at the time of this deploy
-#
-# Nothing is quantized in this deployment, so there is no `torchao`.
 --extra-index-url https://download.pytorch.org/whl/cu130
 diffusers @ git+https://github.com/huggingface/diffusers.git@665f578278365ea4a3318cb8c9b66ce6c01204b9
 torch==2.11.0
 torchvision==0.26.0
-# A reference soundtrack that is not already at the audio VAE's 32 kHz is resampled with `torchaudio`, which is the
-# only thing in the `ref2va` path that needs it — and it is easy to miss, because a 32 kHz reference skips the
-# resample entirely. Both halves need it: the conditioner's `setup` step normalizes the same waveforms this one does.
+# A reference soundtrack that is not already at the audio VAE's 32 kHz is resampled with `torchaudio`; a 32 kHz one
+# skips the resample entirely, so this is easy to miss.
 torchaudio==2.11.0
-# Pinned to the version the MiniMax-H3 parity work was verified against: the Qwen3-VL processor decides the vision
-# patch count, so a different minor changes the conditioning.
+# The Qwen3-VL processor decides the vision patch count, so a different minor changes the conditioning.
 transformers==5.8.0
 accelerate==1.14.0
-# diffusers pins <2; 1.24.0 is the version the parity work ran on.
+# diffusers pins <2.
 huggingface-hub==1.24.0
 gradio==6.20.0
 spaces==0.51.1
 # No `kernels` pin on purpose: the Hub attention backends want `kernels>=0.12.3`, and that version breaks
 # transformers 5.8.0 at import.
-# PyAV: decoding the reference media, i.e. `MiniMaxH3VideoReference.from_file` / `MiniMaxH3AudioReference.from_file`.
+# PyAV decodes the reference media (`MiniMaxH3VideoReference.from_file`).
 av
 pillow
 numpy

spaces_constant_binding_patch.py

@@ -1,32 +1,12 @@
 """Bind AoTI constants that `torch.export` lifted anonymously.
 
-Problem
--------
-`spaces.zero.torch.aoti.LazyAOTIModel` binds a compiled package's constants **by name**::
+`spaces.zero.torch.aoti.LazyAOTIModel` binds a package's constants by intersecting the module's `state_dict()` with
+`compiled_model.get_constant_fqns()`, and keeps whatever it cannot match. `torch.export` only gives a lifted tensor a
+real FQN when it was a registered parameter or buffer; anything else is named `_tensor_constant<N>`, which no
+`state_dict()` can contain, so the compiled model runs against constants nobody set — a SIGSEGV rather than an error.
 
-    constant_fqns = compiled_model.get_constant_fqns()
-    constant_map = {name: tensor for name, tensor in weights.items() if name in constant_fqns}
-    compiled_model.load_constants(constant_map, check_full_update=check_full_update, user_managed=True)
-
-`torch.export` only gives a lifted tensor a real FQN when it was a registered parameter or buffer.
-Anything reached through a plain python attribute is classified `CONSTANT_TENSOR` and the compiled
-artifact names it `_tensor_constant<N>` — a name that can never appear in `state_dict()`. The
-intersection above is then empty, the dict comprehension silently drops every weight, and the
-compiled model runs against constants nobody ever set: a SIGSEGV rather than an error.
-
-This module fixes both halves:
-
-  * `write_constant_aliases(...)` — compile side. Records the exact
-    `_tensor_constant<N> -> real.dotted.fqn` mapping, which the `ExportedProgram` knows even when the
-    compiled package does not, into a `constant_aliases.json` sidecar next to `package.pt2`.
-
-  * `apply_spaces_constant_binding_patch()` — load side. Monkeypatches `LazyAOTIModel.__call__` so it
-    (1) uses that sidecar when present, (2) otherwise falls back to matching anonymous constants
-    against the leftover `state_dict()` entries by dtype+shape read out of the package's own
-    `wrapper.cpp`, and (3) **raises** if the binding is not total instead of segfaulting later.
-
-The load-side patch alone is enough to turn the crash into a clear diagnostic; with the sidecar it
-also makes the package work.
+`write_constant_aliases` records the real names on the compile side; `apply_spaces_constant_binding_patch` uses that
+sidecar on the load side, falls back to matching by dtype+shape, and raises if the binding is still not total.
 """
 
 from __future__ import annotations
@@ -56,10 +36,7 @@ _DTYPES = {
 def register_loose_tensors(module: torch.nn.Module, prefix: str = "") -> list[str]:
     """Re-register plain tensor attributes as buffers so `torch.export` gives them real FQNs.
 
-    Model-agnostic and numerics-preserving: it changes how a tensor is *registered*, never the tensor
-    and never the forward. Run it on the shallow clone right after
-    `unwrap_tensor_subclass_parameters`, immediately before `torch.export.export`. Returns the names
-    it re-registered, which is empty for a module that was already well-formed.
+    Run on the shallow clone, right before `torch.export.export`. Returns the names it re-registered.
     """
     registered = []
     for name, value in list(vars(module).items()):
@@ -78,8 +55,8 @@ def register_loose_tensors(module: torch.nn.Module, prefix: str = "") -> list[st
 def constant_aliases_from_exported_program(exported_program) -> dict[str, str]:
     """`{'_tensor_constant<N>': '<real dotted fqn>'}` for every anonymously lifted constant.
 
-    AOT Inductor numbers its `_tensor_constant<N>` slots in the order the `CONSTANT_TENSOR` inputs
-    appear in the export graph signature, and the signature still carries each one's real FQN.
+    AOT Inductor numbers its slots in the order the `CONSTANT_TENSOR` inputs appear in the graph signature, which
+    still carries each one's real FQN.
     """
     targets = [
         spec.target
@@ -137,13 +114,12 @@ def resolve_constant_map(
     aliases=None,
     allow_shape_fallback: bool = False,
 ):
-    """Map every compiled constant FQN onto one of `weights`, or explain why it cannot."""
+    """Map every compiled constant FQN onto one of `weights`, or report what is left over."""
     constant_map = {name: weights[name] for name in constant_fqns if name in weights}
     missing = [name for name in constant_fqns if name not in constant_map]
     if not missing:
         return constant_map, []
 
-    # 1. the exact mapping, if the compile side recorded one
     aliases = aliases or {}
     for name in list(missing):
         target = aliases.get(name)
@@ -153,10 +129,9 @@ def resolve_constant_map(
     if not missing or not allow_shape_fallback:
         return constant_map, missing
 
-    # 2. otherwise match by dtype+shape against the state_dict entries nobody claimed, preserving
-    #    each side's own order inside a (dtype, shape) group. `get_constant_fqns()` returns the
-    #    slots in *lexicographic* order (`_tensor_constant10` before `_tensor_constant2`), so the
-    #    package's own `constants_info_` index is the only correct order to walk them in.
+    # Match by dtype+shape against the unclaimed `state_dict()` entries, preserving each side's own order inside a
+    # (dtype, shape) group. `get_constant_fqns()` returns slots lexicographically (`_tensor_constant10` before
+    # `_tensor_constant2`), so the package's own `constants_info_` index is the only correct order to walk them in.
     info = _package_constants_info(archive_file)
     by_name = {entry.get("name"): entry for entry in info}
     slot_index = {entry.get("name"): index for index, entry in enumerate(info)}