app.py

"""FastH3 v1 (VSA) — the 4-step DMD2 distillation of MiniMax-H3, text to video + synchronized audio.

`FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree` replaces only the `transformer/` of the MiniMax-H3 release
with a data-free DMD2 student. Everything else in the repo (Qwen3-VL conditioner, both autoencoders, both schedulers)
is an unmodified copy of the base checkpoint, so it runs on the released `diffusers` modular pipeline — the two
differences at inference time are the step count and the **attention backend**.

**The sampling contract.** `num_inference_steps` counts sigma *grid points*, and `N` points drive `N - 1` transformer
forwards. The checkpoint's own `fastvideo_inference.json` states it exactly: `num_inference_steps: 5`,
`transformer_forwards: 4`, `dmd_denoising_steps: [999, 749, 500, 250]`, `guidance_scale: 1.0`. It is fixed here.

**VSA is not optional.** That same file records `attention_backend: VIDEO_SPARSE_ATTN_H3`, `vsa_tile_size: 64`,
`vsa_sparsity: 0.9`. This student was distilled *under* block-sparse attention and ships 50 trained
`attn.to_gate_compress` tensors that only the sparse path reads, so `vsa_h3.py` ports FastVideo's VSA-H3 backend onto
`MiniMaxH3Attention` and runs it, on FastVideo's own Triton kernels (vendored under `vsa_kernel/`). The checkpoint's
`vsa_kernel: sm100a` is the GB200-only fast path for the same mask semantics; this pool is sm120.

**Why the Space is split.** MiniMax-H3 is ~196 GiB in bfloat16 and a ZeroGPU Space is evicted at 150 GB of storage.
This half holds the distilled transformer and the two autoencoders (81 GB); the 62.15 GiB Qwen3-VL conditioner runs in
`multimodalart/qwen3vl-conditioner`, which this Space calls over the gradio API for every request. FastH3 ships the
base release's conditioner verbatim, so that Space encodes this checkpoint exactly. Nothing is quantized anywhere.

**Latency instrumentation.** The `@spaces.GPU` function times the pipeline call alone with `time.perf_counter()` and
returns `gen_s` next to the video, together with per-forward wall times and the peak CUDA allocation, so the
autoresearch loop can see where a request spends its seconds. A per-request id in the report proves each artifact is
freshly generated (a cached replay would repeat the id).
"""

from __future__ import annotations

import functools
import os
import tempfile
import time
import traceback
import uuid
from functools import cache

# Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 75 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", "FastVideo/FastVideo-FastH3-4-step-Preview-v1-VSA-DataFree")
BASE_REPO = "MiniMaxAI/MiniMax-H3"
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `lazy` moves all weights onto the card on the first GPU call and leaves them there; `offload` hands placement to
# `ComponentsManager.enable_auto_cpu_offload`. Packing at startup is not an option here: `spaces` writes every
# startup-resident CUDA tensor to a second on-disk copy, and this checkpoint's transformer (70.1 GB on disk) would
# bust the 150 GB quota.
PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
# `vsa` is the trained route (see the module docstring). `dense` is the escape hatch: it runs the released dense
# `diffusers` path with cuDNN's fused kernel, which is off-distribution for this student but useful to bisect against.
ATTENTION = os.environ.get("H3_ATTENTION", "vsa").lower()
# The checkpoint's own `vsa_sparsity`. Only read when `H3_ATTENTION=vsa`.
VSA_SPARSITY = float(os.environ.get("H3_VSA_SPARSITY", "0.9"))
# 75.7 GiB of weights plus activations does not fit a `large` (48 GiB) allocation.
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# The distilled ladder, as sigma grid points. 5 points -> 4 transformer forwards at t = 1000, 750, 500, 250.
SIGMA_GRID_POINTS = 5
NUM_FORWARDS = SIGMA_GRID_POINTS - 1

# 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),
    "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),
    "1024x1024 · 1:1 max": (1024, 1024),
    # 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),
}
# The distillation's own operating point: 768x1344, 124 frames, 24 fps.
DEFAULT_CANVAS = "1344x768 · 16:9 full"
DEFAULT_DURATION = 5
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
# The ceiling holds for the *snapped* frame count. 75.74 GiB of weights are resident on a 95.0 GiB card, and the
# sparse path's own working set grows with the packed sequence, so this is a memory ceiling, not a policy one.
MIN_UI_DURATION, MAX_UI_DURATION = 2, 8


# The frozen benchmark request. One fixed prompt, one fixed seed, the trained 4-forward schedule, the distilled
# operating point (1344x768, 124 frames, 24 fps), audio on (the pipeline always emits it). `bench/bench_client.py`
# duplicates this table; keep the two in sync.
BENCH_PROMPT = (
    "integrated_multimodal_description: [Shot 1] Live-action, cinematic, a medium-wide shot frames a baker opening "
    "the shutters of a small street bakery before sunrise. The camera pushes in with small amplitude at slow speed "
    "as the middle-aged baker with a calm, slightly raspy voice (S1) places a fresh loaf on the wooden counter and "
    "says: <d>[English] First batch of the morning.</d> [Shot 2] At 00:05.000, the camera cuts to a close-up of "
    "steam rising from the sliced bread while the baker's final words carry over from the previous shot.\n\n"
    "overall_soundscape: Wooden shutters scrape open over a quiet street as trays clink softly inside the bakery. "
    "The doorbell rings once, followed by light footsteps and the crisp sound of bread being sliced.\n\n"
    "non_diegetic_music: A soft acoustic-guitar pattern at a moderate tempo, joined by sparse upright-bass notes and "
    "a gentle fade at the end."
)
BENCH_CANVAS = "1344x768 · 16:9 full"
BENCH_DURATION = 5
BENCH_SEED = 42
BENCH_UPSAMPLE = False  # the prompt is already in the trained format


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 lower_duration_floor(seconds: float = MIN_UI_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))


PIPE = None
MANAGER = None
LOAD_ERROR: str | None = None
LOADED_IN: float | None = None
VSA_BLOCKS = 0
VSA_GATES = 0
# Wall seconds of the last request's transformer forwards (reset by every `_generate` call).
FORWARD_TIMES: list[float] = []
# Wall seconds of the last request's two VAE decodes, appended by the wrappers `_arm_decode_timers` arms.
DECODE_TIMES: dict[str, float] = {}


def status() -> str:
    if LOAD_ERROR:
        return LOAD_ERROR
    if PIPE is None:
        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 81 GB). Watch the Space logs."
    if ATTENTION == "vsa":
        attention = (
            f"**VSA-H3** block-sparse, tile 64 / sparsity {VSA_SPARSITY:g} on {VSA_BLOCKS} blocks "
            f"({VSA_GATES} trained compression gates live)"
        )
    else:
        attention = f"dense `{ATTENTION}` (off-distribution for this student)"
    return (
        f"Ready · distilled transformer + VAEs **bfloat16, unquantized** · {NUM_FORWARDS} transformer forwards "
        f"({SIGMA_GRID_POINTS}-point sigma grid) · attention {attention} · placement `{PLACEMENT}` · "
        f"loaded in {LOADED_IN:.0f}s · conditioner `{CONDITIONER_SPACE}` · "
        f"`{torch_version()}` / triton `{triton_version()}`"
    )


def torch_version() -> str:
    import torch

    try:
        capability = ".".join(map(str, torch.cuda.get_device_capability(0)))
    except Exception:  # noqa: BLE001 - no GPU visible at status time
        capability = "nocuda"
    return f"torch {torch.__version__} (sm_{capability})"


def triton_version() -> str:
    try:
        import triton

        return triton.__version__
    except Exception:  # noqa: BLE001
        return "?"


def load_models() -> str | None:
    """Load the denoising half at startup.

    `MiniMaxH3GeneratorBlocks` declares `transformer`, `vae`, `audio_vae`, the two schedulers and `video_processor`,
    so `load_components` fetches exactly those subfolders — `text_encoder/` and `transformer_ref/` 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.

    `add_gate_compress_modules()` has to run *before* the transformer is instantiated. The checkpoint carries 50
    `transformer_blocks.*.attn.to_gate_compress.weight` tensors — the trained VSA compression gate — and stock
    `MiniMaxH3Attention` does not declare the module, so `from_pretrained` would report them as unexpected and drop
    them. Declaring it first is what makes them load.
    """
    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, VSA_BLOCKS, VSA_GATES

    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 MiniMaxH3GeneratorBlocks

        lower_duration_floor()
        if ATTENTION == "vsa":
            import vsa_h3

            vsa_h3.add_gate_compress_modules()

        manager = ComponentsManager()
        blocks = MiniMaxH3GeneratorBlocks()
        print(f"[gen] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
        pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="fasth3")
        pipe.load_components(dtype=torch.bfloat16)

        if ATTENTION == "vsa":
            VSA_BLOCKS, VSA_GATES = vsa_h3.install(pipe.transformer, sparsity=VSA_SPARSITY)
            print(f"[gen] VSA-H3 on {VSA_BLOCKS} blocks, {VSA_GATES} trained gates", flush=True)
            # Resolve the sparse kernel now (CPU-side import only) so the log shows the outcome at startup.
            vsa_h3._resolve_cuda_sparse_op()
        else:
            pipe.transformer.set_attention_backend(ATTENTION)

        _install_forward_timing(pipe)
        _arm_decode_timers(pipe)

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

        PIPE, MANAGER = pipe, manager
        LOADED_IN = time.time() - started
        print(f"[gen] ready in {LOADED_IN:.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_timers(pipe) -> None:
    """Time each VAE's `decode` call, so a request can report the video/audio decode split.

    Instance-attribute wrapper, like `_arm_decode_hooks`; only meaningful with `PLACEMENT=lazy` (no offload hook in
    front). The video VAE decode runs tiled under fp16 autocast over fp32 weights inside the pipeline's decode
    block; this just measures it.
    """
    for name in ("vae", "audio_vae"):
        module = getattr(pipe, name)
        inner = module.decode

        def timed(*args, _name=name, _decode=inner, **kwargs):
            started = time.perf_counter()
            try:
                return _decode(*args, **kwargs)
            finally:
                DECODE_TIMES[_name] = DECODE_TIMES.get(_name, 0.0) + time.perf_counter() - started

        module.decode = timed


def _install_forward_timing(pipe) -> None:
    """Record the wall time of every transformer forward, so each request can report where its seconds went.

    Runs *after* `vsa_h3.install`, so the timing wrapper sits outside the layout-publishing wrapper. `functools.wraps`
    is load-bearing: the denoise block filters its `token_tags` / `position_ids` / ... kwargs by
    `inspect.signature(transformer.forward).parameters`, and an unwrapped `(*args, **kwargs)` signature would empty
    that set and drop the packed-sequence layout from every forward.
    """
    original_forward = pipe.transformer.forward

    @functools.wraps(original_forward)
    def timed_forward(*args, **kwargs):
        started = time.perf_counter()
        try:
            return original_forward(*args, **kwargs)
        finally:
            FORWARD_TIMES.append(time.perf_counter() - started)

    pipe.transformer.forward = timed_forward


def _arm_decode_hooks(pipe):
    """Make the offload hooks fire for the two VAEs.

    `enable_auto_cpu_offload` wraps `forward`, and the decode blocks call `vae.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)
        inner = module.decode

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

        module.decode = armed


@cache
def conditioner():
    """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)


def encode_remote(prompt: str, canvas: str, num_frames: int, rewrite_prompt: bool = False):
    """`/encode` 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. `canvas` is the label."""
    from safetensors import safe_open

    path, plan = conditioner().predict(
        prompt=prompt,
        image_path=None,
        last_image_path=None,
        canvas=canvas,
        num_frames=num_frames,
        rewrite_prompt=bool(rewrite_prompt),
        api_name="/encode",
    )
    with safe_open(path, framework="pt") as handle:
        return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan


# Seconds of GPU one request needs, from the packed rows it is about to denoise. Block-sparse attention keeps a fixed
# 10% of the tiles, so the cost is linear in the rows and the quadratic term a dense path needs is gone. Measured
# warm on this Space: 37 296 rows in 59-72 s and 57 456 rows in 87 s, i.e. 1.51e-3 to 1.93e-3 s/row depending on how
# fast a slice of the pool the request lands on. Fitted to the *slow* end so a slow slice is not aborted mid-video.
_DUR_A, _DUR_BASE = 1.95e-3, 3.0
# The one-time costs a cold worker pays inside its first request: 75.7 GiB across PCIe (~11 s) plus the Triton JIT of
# the vendored block-sparse kernels (~57 s), measured at 74 s against 6 s warm for the same request. Booking that on
# *every* request would burn 70 s of each visitor's quota for nothing, so it is only booked while this process has
# not yet seen a request come back.
_COLD_ALLOWANCE, _WARM_ALLOWANCE = 75, 8
_WARM = False
# Booked over the estimate. Keep it small: an inflated duration burns the visitor's quota and drops queue priority.
_MARGIN = 1.15


def get_duration(prompt_embeds, text_token_tags, height, width, num_frames, seed, *a, **k):
    height, width, num_frames = int(height), int(width), int(num_frames)
    latent_frames = (num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
    rows = latent_frames * (height // 32) * (width // 32)
    allowance = _WARM_ALLOWANCE if _WARM else _COLD_ALLOWANCE
    return max(60, int((_DUR_A * rows + _DUR_BASE + allowance) * _MARGIN) + 2)


@spaces.GPU(duration=get_duration, size=GPU_SIZE)
def _generate(prompt_embeds, text_token_tags, height: int, width: int, num_frames: int, seed: int):
    """The only thing on GPU time: the four-forward packed-sequence denoise loop and the two decoders.

    `gen_s` wraps the pipeline call alone — the four transformer forwards plus the two decoders — measured inside the
    GPU worker with `time.perf_counter()`, and travels back to the client next to the video. The per-forward wall
    times and the peak CUDA allocation ride along for profiling.

    Only the generated outputs come back — a `@spaces.GPU` return crosses a process boundary by pickling, and the
    full `PipelineState` still holds the packed latents, the rotary grid and the row indices on the card.
    """
    import torch

    if PLACEMENT == "lazy":
        PIPE.to("cuda")

    if ATTENTION == "vsa":
        # One packed-sequence request: let every forward reuse the first forward's VSA geometry instead of
        # re-deriving it from `token_tags` with a host sync each time.
        import vsa_h3

        vsa_h3.begin_request()

    torch.cuda.reset_peak_memory_stats()
    FORWARD_TIMES.clear()
    DECODE_TIMES.clear()
    started = time.perf_counter()
    if os.environ.get("H3_PROFILE") == "1":
        from torch.profiler import ProfilerActivity, profile

        with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
            state = PIPE(
                prompt_embeds=prompt_embeds.to("cuda"),
                text_token_tags=text_token_tags,
                height=height,
                width=width,
                num_frames=num_frames,
                num_inference_steps=SIGMA_GRID_POINTS,
                generator=torch.Generator("cpu").manual_seed(int(seed)),
            )
        print(
            prof.key_averages().table(sort_by="cuda_time_total", row_limit=40, max_name_column_width=60),
            flush=True,
        )
    else:
        state = PIPE(
            prompt_embeds=prompt_embeds.to("cuda"),
            text_token_tags=text_token_tags,
            height=height,
            width=width,
            num_frames=num_frames,
            num_inference_steps=SIGMA_GRID_POINTS,
            generator=torch.Generator("cpu").manual_seed(int(seed)),
        )
    gen_s = time.perf_counter() - started
    peak_gib = torch.cuda.max_memory_allocated() / 1024**3
    forwards = [round(t, 2) for t in FORWARD_TIMES]
    decodes = {name: round(t, 2) for name, t in DECODE_TIMES.items()}
    print(f"[gen] pipeline {gen_s:.2f}s · forwards {forwards} · decodes {decodes} · peak {peak_gib:.2f} GiB", flush=True)
    return (
        state.get("videos")[0],
        state.get("audio")[0].cpu(),
        state.get("sampling_rate"),
        gen_s,
        peak_gib,
        forwards,
        decodes,
    )


@spaces.GPU(duration=300, size=GPU_SIZE)
def aoti_diag() -> str:
    """Phase A of the AOTI experiment: validate the traceable functional reimplementation of the VSA
    sparse-attention soup, and (when `H3_AOTI=1`) export + AOTI-compile it for this GPU and check the
    artifact against eager, with timings.

    Uses the *real bench geometry* (1344x768, 124 frames, 203 text + 414 audio + 37 296 video rows),
    synthesized directly — no generation, no weights.
    """
    import torch

    import aoti_attention

    device = torch.device("cuda")
    heads, dim = 56, 128
    text_rows, audio_rows = 203, 414
    grid_t, grid_h, grid_w = 37, 24, 42
    n_video = grid_t * grid_h * grid_w
    seq_len = text_rows + audio_rows + n_video

    tags = torch.tensor([1] * text_rows + [2] * audio_rows + [0] * n_video, dtype=torch.long)
    position_ids = torch.zeros(seq_len, 3, dtype=torch.float64)
    video_start = text_rows + audio_rows
    frame = torch.cartesian_prod(torch.arange(float(grid_h)), torch.arange(float(grid_w)))
    position_ids[video_start:, 0] = torch.arange(float(grid_t)).repeat_interleave(grid_h * grid_w)
    position_ids[video_start:, 1:] = frame.repeat(grid_t, 1)

    import vsa_h3

    geometry = vsa_h3.geometry_from_layout(tags.to(device), position_ids.to(device), 0.9)
    if geometry is None:
        return "**FAILED**: synthetic bench layout did not yield VSA geometry."

    def timed(fn, *args, repeats=3):
        fn(*args)  # warm (JIT/compile)
        torch.cuda.synchronize()
        import time as _t

        started = _t.perf_counter()
        for _ in range(repeats):
            fn(*args)
        torch.cuda.synchronize()
        return (_t.perf_counter() - started) / repeats

    generator = torch.Generator(device=device).manual_seed(0)
    shape = (1, seq_len, heads, dim)
    q, k, v = (torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) for _ in range(3))
    gate = torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) * 0.05

    lines = [f"`{torch.cuda.get_device_name()}` · seq {seq_len}, padded {geometry.padded_len}, "
             f"tiles {geometry.n_tiles}, topk {geometry.topk}", ""]

    def eager(gate_h):
        # eager sparse_attention takes the gate as [B, S, H, D] and transposes internally
        return vsa_h3.sparse_attention(q, k, v, gate_h, geometry)

    def functional(gate_h):
        gate_t = None if gate_h is None else gate_h.transpose(1, 2)
        return aoti_attention.sparse_attention_functional(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate_t,
            geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
            geometry.topk, geometry.num_prefix_tiles,
        )

    for label, gate_h in (("gate=None", None), ("gate=rand*0.05", gate)):
        ref, got = eager(gate_h), functional(gate_h)
        error = (ref.float() - got.float()).abs().max().item()
        scale = ref.float().abs().max().item()
        cosine = torch.nn.functional.cosine_similarity(
            ref.float().flatten(), got.float().flatten(), dim=0
        ).item()
        t_eager = timed(eager, gate_h)
        t_fn = timed(functional, gate_h)
        lines.append(
            f"| {label} | eager {t_eager * 1000:.0f} ms | functional {t_fn * 1000:.0f} ms | "
            f"rel err {error / scale:.2e} | cosine {cosine:.6f} |"
        )

    report = (
        "VSA soup: eager vs functional (op-for-op reimplementation)\n\n"
        "| case | eager | functional | rel err | cosine |\n|---|---|---|---|---|\n" + "\n".join(lines[2:])
    )

    if os.environ.get("H3_AOTI") == "1":
        try:
            from torch._inductor import aoti_compile_and_package, aoti_load_package

            with torch.no_grad():
                exported = torch.export.export(
                    aoti_attention.sparse_module(geometry.topk, geometry.num_prefix_tiles),
                    args=(
                        q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
                        gate.transpose(1, 2), geometry.untile_index,
                        geometry.variable_block_sizes, geometry.tile_divisor,
                    ),
                )
            package = "/tmp/vsa_sparse_aoti.pt2"
            aoti_compile_and_package(exported, package_path=package)
            compiled = aoti_load_package(package)
            got = compiled(
                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate.transpose(1, 2),
                geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
            )
            ref = eager(gate)
            error = (ref.float() - got.float()).abs().max().item()
            scale = ref.float().abs().max().item()
            cosine = torch.nn.functional.cosine_similarity(
                ref.float().flatten(), got.float().flatten(), dim=0
            ).item()
            t_eager = timed(eager, gate)
            t_compiled = timed(lambda *a: compiled(*a), (
                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate.transpose(1, 2),
                geometry.untile_index, geometry.variable_block_sizes, geometry.tile_divisor,
            ))
            report += (
                f"\n\nAOTI artifact: `{package}`\n\n| eager | compiled | rel err | cosine |\n|---|---|---|---|\n"
                f"| {t_eager * 1000:.0f} ms | {t_compiled * 1000:.0f} ms | {error / scale:.2e} | {cosine:.6f} |"
            )
        except Exception as error:  # noqa: BLE001 - surfaced verbatim for the log
            report += f"\n\n**AOTI FAILED**: `{type(error).__name__}: {error}`"

    return report


@spaces.GPU(duration=120, size=GPU_SIZE)
def selftest() -> str:
    """Check the vendored VSA-H3 kernels against dense attention on this GPU.

    At `sparsity = 0` the block map is all-true, so VSA-H3 has to reproduce full attention exactly (up to the tile
    padding and the fp32 pooled selection, which cannot change an all-true mask). That is the one assertion that
    catches a wrong tile order, a wrong `variable_block_sizes`, or a mis-transposed buffer — all of which would
    otherwise show up only as a subtly wrong video. Runs on random tensors; no weights are touched.

    Returns:
        A markdown report: the dense-equivalence error, and how much of the dense output the 90%-sparse path keeps.
    """
    import torch
    import torch.nn.functional as F

    from diffusers.modular_pipelines.minimax_h3.before_denoise import MiniMaxH3PrepareLayoutStep

    import vsa_h3

    device = torch.device("cuda")
    heads, dim = 8, 128
    # A real packed layout, just a small one: 300 text rows, 5 latent frames of 8x12 video, its soundtrack.
    _, token_tags, *_ = MiniMaxH3PrepareLayoutStep.build_packed_sequence(
        torch.ones(300, dtype=torch.long), 5, 16, 24, 50, (1, 2, 2), 2, 2, 0, ()
    )
    position_ids = torch.zeros(token_tags.numel(), 3, dtype=torch.float64)
    video_start = int((token_tags == 0).nonzero()[0])
    frame = torch.cartesian_prod(torch.arange(8.0), torch.arange(12.0))
    position_ids[video_start:, 0] = torch.arange(5.0).repeat_interleave(96)
    position_ids[video_start:, 1:] = frame.repeat(5, 1)
    token_tags, position_ids = token_tags.to(device), position_ids.to(device)

    lines = []
    generator = torch.Generator(device=device).manual_seed(0)
    shape = (1, token_tags.numel(), heads, dim)
    query, key, value = (
        torch.randn(shape, generator=generator, device=device, dtype=torch.bfloat16) for _ in range(3)
    )
    reference = F.scaled_dot_product_attention(
        query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
    ).transpose(1, 2)

    for sparsity in (0.0, 0.9):
        vsa_h3.reset_tile_buffers()
        geometry = vsa_h3.geometry_from_layout(token_tags, position_ids, sparsity)
        if geometry is None:
            return "**FAILED**: `geometry_from_layout` did not recognize the standard packed layout."
        out = vsa_h3.sparse_attention(query, key, value, None, geometry)
        error = (out.float() - reference.float()).abs().max().item()
        scale = reference.float().abs().max().item()
        similarity = F.cosine_similarity(out.float().flatten(), reference.float().flatten(), dim=0).item()
        lines.append(
            f"| {sparsity:g} | {geometry.topk}/{geometry.num_video_tiles} | {error:.4f} | "
            f"{error / scale:.2e} | {similarity:.6f} |"
        )
        if sparsity == 0.0 and error / scale > 0.02:
            lines.append(f"\n**FAILED**: dense-equivalent VSA differs from SDPA by {error / scale:.3f} relative.")

    return (
        f"VSA-H3 on `{torch.cuda.get_device_name()}`, {token_tags.numel()} packed rows, "
        f"{heads} heads x {dim}.\n\n"
        "| sparsity | tiles kept | max abs err | relative | cosine |\n|---|---|---|---|---|\n" + "\n".join(lines)
    )


def generate(
    prompt: str,
    canvas: str = DEFAULT_CANVAS,
    duration: float = DEFAULT_DURATION,
    upsample: bool = True,
    seed: int = 42,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a video with a synchronized soundtrack from a text prompt, in four transformer forwards.

    Args:
        prompt: The request. MiniMax-H3 was trained on a structured format
            (`integrated_multimodal_description: ... overall_soundscape: ... non_diegetic_music: ...`); leave
            `upsample` on to have the conditioner rewrite a plain sentence into it first.
        canvas: One of the released canvases, as a `WIDTHxHEIGHT · ratio` label. The distillation's own operating
            point is `1344x768 · 16:9 full`.
        duration: Length in seconds, rounded up to the next frame count the video VAE can decode (`17 * n + 5`).
        upsample: Rewrite the prompt into MiniMax-H3's trained format before encoding it.
        seed: Random seed.

    Returns:
        The path of an mp4 holding h264 video and AAC audio, a one-line report (with a per-request id proving the
        artifact is fresh), the measured `gen_s` of the pipeline call, and the rewritten prompt when there was one.
    """
    if LOAD_ERROR:
        raise gr.Error(LOAD_ERROR)
    if PIPE is None:
        raise gr.Error("The denoiser is still loading.")
    if not prompt or not prompt.strip():
        raise gr.Error("MiniMax-H3 always takes a prompt.")

    from diffusers.utils import encode_video

    num_frames = snap_frames(duration)

    progress(0.0, desc=f"{'Rewriting the prompt' if upsample else 'Conditioning'} on {CONDITIONER_SPACE} ...")
    conditioned = time.time()
    prompt_embeds, text_token_tags, metadata, plan = encode_remote(
        prompt, canvas, num_frames, rewrite_prompt=bool(upsample)
    )
    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.2, desc=f"{NUM_FORWARDS} transformer forwards at {width}x{height}, {num_frames} frames ...")
    frames, audio, sampling_rate, gen_s, peak_gib, forwards, decodes = _generate(
        prompt_embeds, text_token_tags, height, width, num_frames, seed
    )

    directory = os.path.join(tempfile.gettempdir(), "fasth3-outputs")
    os.makedirs(directory, exist_ok=True)
    request_id = uuid.uuid4().hex[:12]
    path = os.path.join(directory, f"fasth3-{request_id}.mp4")
    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)

    forward_total = sum(forwards) if forwards else float("nan")
    forward_detail = "+".join(f"{t:.1f}" for t in forwards) if forwards else "n/a"
    decode_detail = " + ".join(f"{k} {v:.1f}s" for k, v in decodes.items()) if decodes else "n/a"
    report = (
        f"`{width}x{height}`, {num_frames} frames @ {FPS} fps, {NUM_FORWARDS} transformer forwards · "
        f"pipeline **{gen_s:.1f}s** (forwards {forward_detail} = {forward_total:.1f}s, {decode_detail}, "
        f"other {gen_s - forward_total - sum(decodes.values()):.1f}s) · peak {peak_gib:.1f} GiB · "
        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
        f"{', rewritten' if refined else ''}) · seed {int(seed)} · req {request_id}"
    )
    print(f"[gen] {report}", flush=True)
    return path, report, round(gen_s, 2), refined


load_models()

INTRO = f"""# FastH3 v1 (VSA) — MiniMax-H3 in 4 steps, sparse

<div>
  <a href="https://huggingface.co/{MODEL_REPO}" target="_blank" rel="noopener"><strong>[ model ]</strong></a> &nbsp;
  <a href="https://github.com/hao-ai-lab/FastVideo" target="_blank" rel="noopener"><strong>[ FastVideo ]</strong></a> &nbsp;
  <a href="https://huggingface.co/{BASE_REPO}" target="_blank" rel="noopener"><strong>[ base model ]</strong></a>
</div>

[`{MODEL_REPO}`](https://huggingface.co/{MODEL_REPO}) is a **data-free DMD2 distillation** of
[MiniMax-H3](https://huggingface.co/{BASE_REPO}), the 33B dual-modality transformer that generates video **and** a
fully synchronized soundtrack (ambience, foley, speech) in one denoising pass. The base model samples in 50 steps;
this student walks a trained 4-step ladder — `t = 999, 749, 500, 250` — for **{NUM_FORWARDS} transformer forwards**
per video.

It is also distilled **under Video Sparse Attention**: 64-token tiles at 90% sparsity, with a trained per-head
compression gate. This Space runs that sparse path, on FastVideo's own Triton kernels — not a dense substitute.
"""

FORMAT_NOTE = """MiniMax-H3 was trained on a structured prompt, not a caption:

```text
integrated_multimodal_description: [Shot 1] ... <d>[English] spoken line.</d> [Shot 2] At 00:04.500, ...
overall_soundscape: ...
non_diegetic_music: ...
```

**Expand prompt** (on by default) sends a plain sentence through the Qwen3-VL conditioner's own language-model head
first, which writes that format with the same weights that are about to encode it. Turn it off when the prompt is
already written out — as the last two examples below are. See the base model's
[prompt writing guide](https://huggingface.co/MiniMaxAI/MiniMax-H3/blob/main/docs/VIDEO_PROMPT_WRITING_GUIDE_base_en.md).
"""

# Both long examples are the MiniMax-H3 authors' own published T2VA prompts: the first is Case 1 of the prompt writing
# guide, the second is the reproducible 768p T2VA case from the model card
# (`scripts/readme/reproducible-768p-t2va-request.sh`). Both are documentation of the Apache-2.0 base repo.
GUIDE_CASE_1 = BENCH_PROMPT
OFFICIAL_T2VA = (
    "integrated_multimodal_description: [Shot 1] Cinematic, medium wide shot, pushing in slowly. In the cavernous, "
    "dimly lit bridge of a starship, sleek metallic consoles with glowing amber displays flank a massive, curved "
    "observation window. A female captain, in her late 40s with an athletic build and short silver-streaked black "
    "hair, stands in the center midground. She wears a structured, high-collared dark navy military tunic with "
    "silver chest insignias. Her back is to the camera, silhouetted against the cool, ambient starlight pouring "
    "through the thick glass. She stands perfectly still with her hands clasped tightly behind her back. Outside the "
    "window, a massive armada of jagged, dark grey dreadnoughts hovers in tight formation against a deep purple "
    "space nebula. The fleet's massive rear thrusters begin to glow with an intense, escalating bright blue light. "
    "[Shot 2] At 00:04.500, the camera cuts to a close-up of the captain's face and shakes strongly. The brilliant "
    "blue-white light from the fleet's gathering energy reflects vividly in her dark eyes. Suddenly, a blinding "
    "white flash floods through the window, completely washing out the background as the fleet jumps to hyperspace. "
    "The sheer spatial force violently jolts the bridge, causing the captain from Shot 1 to stagger slightly "
    "forward, her shoulders tensing as she visibly braces herself against the physical tremors. As the intense "
    "white light fades abruptly, leaving only the dim, empty expanse of the purple nebula reflected in her starkly "
    "lit skin, her jaw clenches, and she slowly closes her eyes in the newly emptied space.\n"
    "overall_soundscape: A low, resonant hum of the ship's ambient life support systems serves as the baseline, soon "
    "drowned out by an audible, escalating, high-pitched electronic whine as the fleet outside charges its "
    "hyperdrives. A massive, deafening, bass-heavy boom and sharp crackle erupts during the blinding flash, "
    "accompanied by the loud metallic creaking, rattling, and deep thuds of the bridge's bulkheads vibrating under "
    "immense physical stress. The intense roaring impact then cuts abruptly back to a hollow, echoing room tone, "
    "leaving only the faint, steady hum of the isolated bridge.\n"
    "non_diegetic_music: Cinematic space-opera orchestral score, slow tempo, featuring a solitary, mournful French "
    "horn melody over deep, sustained string dissonances that build rapidly in volume and intensity, swelling to a "
    "massive orchestral peak before snapping immediately into silence right after the jump."
)

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

with gr.Blocks(title="FastH3 v1 (VSA)") as demo:
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                lines=5,
                placeholder="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
                value="A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
            )
            upsample = gr.Checkbox(
                label="Expand prompt into MiniMax-H3's trained format",
                value=True,
                info="Runs on the conditioner Space before encoding. Turn off for a prompt already in that format.",
            )
            run = gr.Button("Generate", variant="primary")
            with gr.Accordion("Advanced options", open=False):
                canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                duration = gr.Slider(
                    label="Duration (s)",
                    minimum=MIN_UI_DURATION,
                    maximum=MAX_UI_DURATION,
                    step=1,
                    value=DEFAULT_DURATION,
                )
                seed = gr.Number(label="Seed", value=42, precision=0)
                gr.Markdown(
                    f"Steps are fixed at the trained ladder — a {SIGMA_GRID_POINTS}-point sigma grid, "
                    f"{NUM_FORWARDS} transformer forwards, exactly what the checkpoint's own "
                    "`fastvideo_inference.json` specifies. There is no guidance scale and no negative prompt: the "
                    "base model is guidance-distilled."
                )

        with gr.Column():
            video = gr.Video(label="Video + soundtrack")
            report = gr.Markdown()
            gen_s = gr.Number(label="gen_s — server-side pipeline time (s)", precision=2, interactive=False)
            with gr.Accordion("Expanded prompt", open=False):
                upsampled = gr.Textbox(show_label=False, lines=10, interactive=False)

    with gr.Accordion("Prompt format", open=False):
        gr.Markdown(FORMAT_NOTE)
    banner = gr.Markdown()

    gr.Examples(
        examples=[
            [
                "A red fox trotting through a snowy pine forest at dawn, snow crunching underfoot",
                "960x544 · 16:9 fast",
                5,
                True,
            ],
            [
                "A cellist playing a slow, low melody alone in an empty concert hall",
                "1344x768 · 16:9 full",
                5,
                True,
            ],
            [GUIDE_CASE_1, "1344x768 · 16:9 full", 5, False],
            [OFFICIAL_T2VA, "1344x768 · 16:9 full", 5, False],
        ],
        inputs=[prompt, canvas, duration, upsample],
        outputs=[video, report, gen_s, upsampled],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples — the last two are the MiniMax-H3 authors' own published T2VA prompts",
    )

    run.click(
        generate,
        [prompt, canvas, duration, upsample, seed],
        [video, report, gen_s, upsampled],
        api_name="generate",
    )
    demo.load(status, None, banner, api_name="status")

    # No UI, API only: the sparse-attention equivalence check, so the kernel can be verified on this pool without
    # spending a full generation; and the AOTI Phase A validation.
    diagnose = gr.Button(visible=False)
    diagnose.click(selftest, None, gr.Markdown(visible=False), api_name="selftest")
    diagnose2 = gr.Button(visible=False)
    diagnose2.click(aoti_diag, None, gr.Markdown(visible=False), api_name="aoti_diag")


if __name__ == "__main__":
    # Gradio 6 moved `theme` and `css` off the `Blocks` constructor onto `launch`.
    demo.launch(theme=gr.themes.Citrus(), css=CSS, show_error=True, max_threads=1000, mcp_server=True)