app.py

"""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.
"""

from __future__ import annotations

import os
import tempfile
import time
import traceback

# 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.
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.
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.
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.
CANVASES = {
    # 16:9
    "960x544 · 16:9 fast": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    # 9:16
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    # 1:1
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    # 4:3 / 3:4
    "768x576 · 4:3 fast": (576, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "576x768 · 3:4 fast": (768, 576),
    "768x1024 · 3:4 full": (1024, 768),
    # 21:9
    "1152x512 · 21:9 fast": (512, 1152),
    "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.
MAX_UI_DURATION = 14
MIN_DURATION = 2
# A reference video shorter than 2 s gives the model almost no motion to read, and 15 s is the checkpoint's ceiling.
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.
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.
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


def snap_frames(seconds: float) -> int:
    """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
    frames = max(1, round(float(seconds) * FPS))
    while frames % FRAMES_PER_CHUNK != LATENTS_PER_CHUNK:
        frames += 1
    return frames


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


def target_rows(height: int, width: int, num_frames: int) -> int:
    """The generated rows of the packed sequence: video patched `(1, 2, 2)`, plus two audio rows per latent."""
    video = video_latent_frames(num_frames) * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS


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.
    """
    from PIL import Image

    from diffusers.modular_pipelines.minimax_h3.modular_pipeline import resolve_canvas_size

    rows = 0
    for kind, path in references:
        if kind == "image":
            width, height = Image.open(path).size
            scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
            resolved = [
                max(CANVAS_MULTIPLE, round(edge * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
                for edge in (height, width)
            ]
            rows += (resolved[0] // CANVAS_MULTIPLE) * (resolved[1] // CANVAS_MULTIPLE)
            continue

        video_seconds, audio_seconds = probe(path)
        if kind == "video" and video_seconds is not None:
            import av

            with av.open(path) as container:
                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 += (
                video_latent_frames(snapped)
                * (canvas_height // CANVAS_MULTIPLE)
                * (canvas_width // CANVAS_MULTIPLE)
            )
        if audio_seconds is not None:
            seconds = min(audio_seconds, num_frames / FPS)
            rows += round(seconds * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
    return rows


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`.
    """
    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
    # 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)
    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)
    return duration


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
CLIENTS: dict[str | None, object] = {}


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.
    """
    global PIPE, MANAGER, LOAD_ERROR

    if PIPE is not None or LOAD_ERROR is not None:
        return LOAD_ERROR

    started = time.time()
    try:
        import torch
        from diffusers import ComponentsManager

        from h3_split_blocks import MiniMaxH3Ref2VAGeneratorBlocks

        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")
        # 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.
        #
        # `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.
        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.
        import h3_aoti

        h3_aoti.maybe_load(pipe.transformer_ref)

        if PLACEMENT == "offload":
            manager.enable_auto_cpu_offload(device="cuda")
            _arm_decode_hooks(pipe)

        PIPE, MANAGER = pipe, manager
        print(f"[ref2va] ready in {time.time() - started:.0f}s", flush=True)
    except Exception as error:
        traceback.print_exc()
        LOAD_ERROR = (
            f"**Loading `{MODEL_REPO}` failed** after {time.time() - started:.0f}s: "
            f"`{type(error).__name__}: {error}`"
        )
    return LOAD_ERROR


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.
    """
    for name in ("vae", "audio_vae"):
        module = getattr(pipe, name)
        for method in ("encode", "decode"):
            inner = getattr(module, method)

            def armed(*args, _module=module, _inner=inner, **kwargs):
                hook = getattr(_module, "_hf_hook", None)
                if hook is not None:
                    hook.pre_forward(_module)
                return _inner(*args, **kwargs)

            setattr(module, method, armed)


def conditioner(ip_token: str | None = None):
    """The other half, over the gradio API — booked against *the caller's* ZeroGPU quota, not this org's.

    ZeroGPU attributes a booking to the `X-IP-Token` header of the request that triggered it
    (`spaces/zero/client.py`), which the Spaces router puts on every browser request. That header is what pays for
    this Space's own `@spaces.GPU` call, and forwarding it to the conditioner makes the same identity pay for the
    conditioner's — the two halves of one user's request then bill as one request, the way they would if this were a
    single Space.

    Without it the conditioner falls back to an IP-based quota, whose ceiling is low enough that an `xlarge` booking
    is refused outright ("The requested GPU duration (Ns) is larger than the maximum allowed"), so a call that does
    not forward a token only works because the conditioner keeps its own reservation small.

    Cached per token: building a `Client` costs a round trip to the Space config, and a token is per user session.
    """
    from gradio_client import Client

    if ip_token in CLIENTS:
        return CLIENTS[ip_token]
    # No org token: the request runs on the caller side quota, which is the point of forwarding theirs.
    client = Client(CONDITIONER_SPACE, headers={"X-IP-Token": ip_token} if ip_token else None)
    if len(CLIENTS) >= 32:
        CLIENTS.pop(next(iter(CLIENTS)))
    CLIENTS[ip_token] = client
    return client


def ip_token_of(request) -> str | None:
    """The caller's ZeroGPU identity, as the Spaces router put it on this request.

    Present on a browser request and on an API request the router could attribute; absent for a truly anonymous
    caller, which then falls back to the conditioner's IP-based quota. Both the UI path and the `/generate` API path
    reach this through the same `gr.Request` gradio injects for a parameter annotated with it.
    """
    headers = getattr(request, "headers", None)
    token = None if headers is None else headers.get("x-ip-token")
    print(f"[ref2va] conditioner call {'forwards the caller ZeroGPU token' if token else 'is anonymous (IP quota)'}", flush=True)
    return token


LOG_TAG = "ref2va"


def call_conditioner(ip_token, **arguments):
    """One conditioner call, on the caller's ZeroGPU identity when that is accepted and anonymously when it is not.

    Forwarding is best effort. ZeroGPU's `/schedule` answers `401` for a proxy token it will not honour — which is
    what a token minted for *this* Space looks like when it arrives at the conditioner — and `spaces` surfaces that as
    `Expired ZeroGPU proxy token`. So the forwarded call is tried first, and a rejected token falls back to no token
    at all rather than failing the request. The conditioner's own bookings are sized to fit the unattributed ceiling,
    so the fallback is a working path and not a degraded one.
    """
    api_name = arguments.pop("api_name")
    if ip_token is not None:
        try:
            return conditioner(ip_token).predict(**arguments, api_name=api_name)
        except Exception as error:
            if "proxy token" not in str(error):
                raise
            print(f"[{LOG_TAG}] the forwarded ZeroGPU token was refused ({error}); retrying anonymously", flush=True)
            CLIENTS.pop(ip_token, None)
    return conditioner(None).predict(**arguments, api_name=api_name)


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


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.
    """
    ordered = [("image", path) for path in image_paths if path]
    if audio_path:
        ordered.append(("audio", audio_path))
    if video_path:
        ordered.append(("video", video_path))
    return ordered


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.
    """
    from diffusers.modular_pipelines.minimax_h3 import (
        MiniMaxH3AudioReference,
        MiniMaxH3ImageReference,
        MiniMaxH3VideoReference,
    )

    classes = {"image": MiniMaxH3ImageReference, "video": MiniMaxH3VideoReference, "audio": MiniMaxH3AudioReference}
    return [classes[kind].from_file(path) for kind, path in references]


def audio_bearing(references: list[tuple[str, str]]) -> list[tuple[str, float]]:
    """The references that carry a waveform, and how long it is. A video reference brings its own soundtrack."""
    carried = []
    for kind, path in references:
        if kind == "image":
            continue
        _, audio_seconds = probe(path)
        if audio_seconds is not None:
            carried.append((kind, audio_seconds))
    return carried


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.
    """
    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.
    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))


def check(prompt: str, references: list[tuple[str, str]]) -> None:
    """The model's own rules, before anything is uploaded or a card is allocated."""
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt, references or not.")
    if not references:
        raise gr.Error("Add at least one reference — an image or a video for the model to condition on.")
    if {kind for kind, _ in references} == {"audio"}:
        raise gr.Error("An audio reference needs an image or a video alongside it; it cannot go on its own.")
    for kind, path in references:
        if kind != "video":
            continue
        video_seconds, _ = probe(path)
        if video_seconds is None:
            raise gr.Error("That reference video has no video stream. Drop it in the audio slot instead.")
        if not MIN_REFERENCE_VIDEO <= video_seconds <= MAX_REFERENCE_VIDEO:
            raise gr.Error(
                f"The reference video is {video_seconds:.1f} s. Use a clip between "
                f"{MIN_REFERENCE_VIDEO:g} and {MAX_REFERENCE_VIDEO:g} seconds."
            )


def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False, ip_token=None):
    """Ask the conditioner Space for `prompt_embeds` + `text_token_tags`. Off this Space's GPU time entirely.

    The references go over with the request: `ref2va`'s presentation puts a vision block in front of the prompt for
    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.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = call_conditioner(
        ip_token,
        prompt=prompt,
        media=[handle_file(path) for _, path in references],
        kinds=",".join(kind for kind, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        api_name="/encode_ref2va",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
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.
    """
    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(
        prompt_embeds=prompt_embeds.to("cuda"),
        text_token_tags=text_token_tags,
        references=build_references(references),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        generator=torch.Generator("cpu").manual_seed(int(seed)),
    )
    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")


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.
    prompt,
    image_1=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    match=True,
    duration=5,
    steps=28,
    seed=42,
    upsample=False,
    progress=gr.Progress(track_tqdm=True),
    request: gr.Request | None = None,
):
    """One request. `upsample` is appended last and defaults off, so an existing API client is untouched by it."""
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")

    from diffusers.utils import encode_video

    images = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9]
    references = collect(images, audio_path, video_path)
    check(prompt, references)

    # `0` is "leave it to the references" over the wire, which MiniMax-H3 accepts when exactly one of them carries a
    # soundtrack. The conditioner resolves it either way and this Space pins whatever comes back.
    derivable = len(audio_bearing(references)) == 1
    requested = 0 if (match and derivable) else snap_frames(duration)

    progress(0.0, desc="Upsampling the prompt ..." if upsample else "Reading the prompt and references ...")
    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(
            prompt, references, canvas, requested, rewrite_prompt=upsample, ip_token=ip_token_of(request)
        )
    except gr.Error:
        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.
        traceback.print_exc()
        raise gr.Error(
            f"The conditioner ({CONDITIONER_SPACE}) failed with `{type(error).__name__}: {error}`. "
            "Its logs carry the full traceback."
        ) from error
    condition_seconds = time.time() - conditioned
    height, width, num_frames = (int(metadata[key]) for key in ("height", "width", "num_frames"))
    refined = plan.get("refined_prompt") or ""

    progress(0.1, desc=f"Generating {num_frames / FPS:.1f} s at {width}x{height} ...")
    started = time.time()
    frames, audio, sampling_rate = _generate(
        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed
    )
    generate_seconds = time.time() - started

    directory = os.path.join(tempfile.gettempdir(), "h3-outputs")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"h3-ref2va-{int(time.time() * 1000)}.mp4")
    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)

    print(
        f"[ref2va] {[kind for kind, _ in references]} · `{width}x{height}`, {num_frames} frames "
        f"({num_frames / FPS:.3f} s), {int(steps)} steps · conditioner {condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens{', upsampled' if refined else ''}) · "
        f"denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}",
        flush=True,
    )
    return path, refined, gr.update(visible=bool(refined))


load_models()

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>
</div>

**MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a
fully synchronized soundtrack (ambience, foley, speech). Bring your own subject, voice or camera move as a
reference.
"""

CSS = """
.main.fillable { max-width: 1250px !important; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="MiniMax-H3 Reference") as demo:
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=3,
                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.
            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.
                    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.
                                height=210,
                                visible=index < OPEN_IMAGE_SLOTS,
                            )
                            for index in range(MAX_IMAGE_SLOTS)
                        ]
                    add_image = gr.Button("+ Add another image", size="sm", variant="secondary")
                with gr.Tab("Audio"):
                    audio = gr.Audio(label="A voice or a piece of music", type="filepath")
                with gr.Tab("Video"):
                    video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                match = gr.Checkbox(label="Match the reference soundtrack", value=True, visible=False)
                duration = gr.Slider(
                    label="Duration (s)", minimum=MIN_DURATION, maximum=MAX_UI_DURATION, step=1, value=5
                )
                steps = gr.Slider(label="Steps", minimum=10, maximum=40, step=1, value=28)
                seed = gr.Number(label="Seed", value=42, precision=0)

        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.
            with gr.Accordion("Upsampled prompt", open=False, visible=False) as upsampled_panel:
                upsampled = gr.Textbox(show_label=False, lines=8, interactive=False)

    open_slots = gr.State(OPEN_IMAGE_SLOTS)

    def reveal_image_slot(open_count):
        open_count = min(open_count + 1, MAX_IMAGE_SLOTS)
        return [
            open_count,
            *[gr.update(visible=index < open_count) for index in range(MAX_IMAGE_SLOTS)],
            gr.update(visible=open_count < MAX_IMAGE_SLOTS),
        ]

    add_image.click(reveal_image_slot, open_slots, [open_slots, *images, add_image], api_name=False)

    for control in (audio, video, match):
        control.change(
            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
    # 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.
    request = [prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample]

    gr.Examples(
        examples=[
            [
                "The character walks through a neon-lit street in the rain, humming to themselves",
                "examples/subject.png",
                None,
                None,
                EXAMPLE_CANVAS,
            ],
            [
                "The character speaks to camera in a quiet room, lips matching every word",
                "examples/subject.png",
                "examples/voice.wav",
                None,
                EXAMPLE_PORTRAIT_CANVAS,
            ],
            [
                "The character moves with the same camera push, down a rainy alley at night",
                "examples/subject.png",
                None,
                "examples/motion.mp4",
                EXAMPLE_CANVAS,
            ],
        ],
        inputs=[prompt, images[0], audio, video, canvas],
        outputs=[result, upsampled, upsampled_panel],
        fn=generate,
        cache_examples=True,
        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")


if __name__ == "__main__":
    demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS)