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

from __future__ import annotations

import os
import tempfile
import time
import traceback
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


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


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}`"
    )


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)
        else:
            pipe.transformer.set_attention_backend(ATTENTION)

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

    Only the three 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")

    torch.cuda.reset_peak_memory_stats()
    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)),
    )
    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"[gen] peak allocated {peak:.2f} GiB", flush=True)
    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")


@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 of what ran, 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 ...")
    started = time.time()
    frames, audio, sampling_rate = _generate(prompt_embeds, text_token_tags, height, width, num_frames, seed)
    generate_seconds = time.time() - started
    # A request has come back, so the next one does not have to book the cold worker's placement + JIT.
    global _WARM
    _WARM = True

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

    report = (
        f"`{width}x{height}`, {num_frames} frames ({num_frames / FPS:.3f} s), {NUM_FORWARDS} transformer forwards · "
        f"conditioner {condition_seconds:.0f}s ({plan['num_text_tokens']} tokens"
        f"{', rewritten' if refined else ''}) · denoise + decode {generate_seconds:.0f}s "
        f"({generate_seconds / NUM_FORWARDS:.1f} s/forward) · seed {int(seed)}"
    )
    print(f"[gen] {report}", flush=True)
    return path, report, 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 = (
    "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."
)
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 on 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()
            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, 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, 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.
    diagnose = gr.Button(visible=False)
    diagnose.click(selftest, None, gr.Markdown(visible=False), api_name="selftest")


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)