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,
[`minimax-h3-conditioner`](https://huggingface.co/spaces/diffusers-internal-dev/minimax-h3-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 `MiniMaxH3Ref2VABlocks` with its `text_encoder` step removed — see `h3_split_blocks.py`. Only *text*
encoding is remote: `reference_encoder` is the `ref2va` blockset'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", "diffusers-internal-dev/MiniMax-H3")
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "diffusers-internal-dev/minimax-h3-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_DURATION = int(os.environ.get("H3_GPU_DURATION", "900"))
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# MiniMax-H3's own canvases, i.e. `resolve_canvas_size` from `diffusers.modular_pipelines.minimax_h3.packing`
# 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 = {
    "960x544 · 16:9 fast (default)": (544, 960),
    "1024x576 · 16:9 fast": (576, 1024),
    "1152x640 · 16:9": (640, 1152),
    "1280x704 · 16:9": (704, 1280),
    "1344x768 · 16:9 full": (768, 1344),
    "544x960 · 9:16 fast": (960, 544),
    "640x1152 · 9:16": (1152, 640),
    "768x1344 · 9:16 full": (1344, 768),
    "544x544 · 1:1 fast": (544, 544),
    "768x768 · 1:1 full": (768, 768),
    "1024x768 · 4:3 full": (768, 1024),
    "768x1024 · 3:4 full": (1024, 768),
    "1536x672 · 21:9 full": (672, 1536),
}
DEFAULT_CANVAS = "960x544 · 16:9 fast (default)"
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 = 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_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


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


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
CLIENT = 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`, 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

    token = os.environ.get("HF_TOKEN")
    if not token:
        LOAD_ERROR = f"**`HF_TOKEN` secret is missing** and `{MODEL_REPO}` is private. Add it and restart."
        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")
        pipe.load_components(dtype=torch.bfloat16, token=token)

        # 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():
    """The other half, over the gradio API. Cached — building a `Client` costs a round trip to the Space config."""
    global CLIENT
    if CLIENT is None:
        from gradio_client import Client

        CLIENT = Client(CONDITIONER_SPACE, token=os.environ.get("HF_TOKEN"))
    return CLIENT


def probe(path: str) -> tuple[float | None, float | None]:
    """`(video seconds, audio seconds)` of a media file, either being `None` when the stream is absent."""
    import av

    def seconds(stream, container):
        if stream.duration is not None and stream.time_base is not None:
            return float(stream.duration * stream.time_base)
        return None if container.duration is None else container.duration / av.time_base

    with av.open(path) as container:
        video = seconds(container.streams.video[0], container) if container.streams.video else None
        audio = seconds(container.streams.audio[0], container) if container.streams.audio else None
    return video, audio


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 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):
    """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.
    """
    from gradio_client import handle_file
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        media=[handle_file(path) for _, path in references],
        kinds=",".join(kind for kind, _ in references),
        canvas=canvas,
        num_frames=num_frames,
        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=GPU_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

    from diffusers.modular_pipelines.minimax_h3 import MiniMaxH3Reference

    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=[MiniMaxH3Reference(**{kind: path}) for kind, path in 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(
    prompt,
    image_1=None,
    image_2=None,
    image_3=None,
    image_4=None,
    image_5=None,
    image_6=None,
    image_7=None,
    image_8=None,
    image_9=None,
    audio_path=None,
    video_path=None,
    canvas=DEFAULT_CANVAS,
    match=True,
    duration=5,
    steps=28,
    seed=42,
    progress=gr.Progress(track_tqdm=True),
):
    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="Reading the prompt and references ...")
    conditioned = time.time()
    try:
        prompt_embeds, text_token_tags, metadata, plan = encode_remote(prompt, references, canvas, requested)
    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"))

    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) · denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}",
        flush=True,
    )
    return path


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",
            )
            # 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")

    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
        )

    run.click(
        generate,
        [prompt, *images, audio, video, canvas, match, duration, steps, seed],
        result,
        api_name="generate",
    )


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