"""FastH3 8-Step V2 — the 8-forward DMD2 + VSA distillation of MiniMax-H3, text to video + synchronized audio.

`FastVideo/FastVideo-FastH3-8-Step-V2` replaces only the `transformer/` of the MiniMax-H3 release with a data-free
DMD2 student trained under VSA-H3 at 80% sparsity. 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 unchanged.

**The sampling contract**, straight out of the checkpoint's own `fastvideo_inference.json` sidecar:

    num_inference_steps  9        dmd_denoising_steps  [999, 874, 749, 624, 500, 375, 250, 125]
    transformer_forwards 8        video_scheduler_shift 10.0    audio_scheduler_shift 3.0
    guidance_scale       1.0      task                 t2av

`num_inference_steps` counts sigma *grid points*, and `N` points drive `N - 1` transformer forwards — so the
checkpoint's nine-point grid is exactly eight forwards, which is what FastVideo's own
`examples/inference/basic/basic_fasth3_8step.py` pins (`GRID_POINTS = 9`, and it *errors* on any other value).
`linspace(1, 0, 9) * 1000` is `[1000, 875, 750, 625, 500, 375, 250, 125, 0]`, i.e. the trained rung ladder; the
shift that turns those into sigmas is the V2-specific **10.0** (the base model's is 12.0) and ships inside this
checkpoint's `scheduler/scheduler_config.json`, so the schedule is read from the checkpoint rather than set here.
Steps are therefore not a knob.

**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 (80 GB on disk, 72 GiB resident); the 62 GiB
Qwen3-VL conditioner runs in `multimodalart/qwen3vl-conditioner`, which this Space calls over the gradio API for
every request. V2 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 72 GiB load can happen at
# startup rather than on GPU time.
import spaces
import gradio as gr

MODEL_REPO = os.environ.get("H3_MODEL_REPO", "FastVideo/FastVideo-FastH3-8-Step-V2")
BASE_REPO = "MiniMaxAI/MiniMax-H3"
CONDITIONER_SPACE = os.environ.get("H3_CONDITIONER", "multimodalart/qwen3vl-conditioner")
# `lazy` moves all 72 GiB onto the card on the first GPU call and leaves it 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, it
# carries the trained VSA gate parameters alongside the weights diffusers loads) would bust the 150 GB quota.
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. The
# student was *trained* under VSA-H3 block-sparse attention at 80% sparsity, and the checkpoint's sidecar asks for
# `VIDEO_SPARSE_ATTN_H3` with `vsa_kernel: sm100a` — but that kernel is datacenter-Blackwell only and this pool is
# sm120, so the route here is dense. It is the same trade the 4-step sibling Space makes.
ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
# 72 GiB of weights plus activations does not fit a `large` (48 GiB) allocation.
GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")

# The trained ladder, as sigma grid points. 9 points -> 8 transformer forwards on the rungs
# [999, 874, 749, 624, 500, 375, 250, 125], under video shift 10.0 / audio shift 3.0.
SIGMA_GRID_POINTS = 9
NUM_FORWARDS = SIGMA_GRID_POINTS - 1
DMD_RUNGS = [999, 874, 749, 624, 500, 375, 250, 125]

# 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 is 1344x768 / 124 frames / 24 fps (`basic_fasth3.py`'s defaults). Eight
# forwards is twice the 4-step sibling's GPU bill at the same canvas, and a ZeroGPU booking is charged double at
# `xlarge`, so the *default* here is one notch down — a full-quality run at the operating point is one dropdown away.
DEFAULT_CANVAS = "1024x576 · 16:9 fast"
DEFAULT_DURATION = 5
FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
MIN_UI_DURATION, MAX_UI_DURATION = 2, 10


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
VIDEO_SHIFT: float | None = None
AUDIO_SHIFT: float | None = None


def status() -> str:
    if LOAD_ERROR:
        return LOAD_ERROR
    if PIPE is None:
        return f"Loading `{MODEL_REPO}` (transformer + VAEs, 80 GB). Watch the Space logs."
    return (
        f"Ready · distilled transformer + VAEs **bfloat16, unquantized** · {NUM_FORWARDS} transformer forwards "
        f"({SIGMA_GRID_POINTS}-point sigma grid) · shift {VIDEO_SHIFT} video / {AUDIO_SHIFT} audio, read from the "
        f"checkpoint · placement `{PLACEMENT}` · attention `{ATTENTION}` · loaded in {LOADED_IN:.0f}s · "
        f"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
    (and `transformer_ref` in this checkpoint's `modular_model_index.json` points back at the base repo anyway).
    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.

    The distilled `transformer/` carries 50 extra `attn.to_gate_compress` tensors — one per layer, the trained
    VSA-H3 sparse-gate parameters. Dense inference does not read them, so diffusers reports them as unused and drops
    them; the resident module is byte-for-byte the shape of the base release's.

    The two scheduler configs are the V2-specific part of the recipe (video shift 10.0 against the base model's
    12.0), and they come out of the checkpoint, so nothing about the schedule is hardcoded here beyond the grid size.
    """
    global PIPE, MANAGER, LOAD_ERROR, LOADED_IN, VIDEO_SHIFT, AUDIO_SHIFT

    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()
        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-8step-v2")
        pipe.load_components(dtype=torch.bfloat16)
        pipe.transformer.set_attention_backend(ATTENTION)

        VIDEO_SHIFT = float(pipe.scheduler.config.shift)
        AUDIO_SHIFT = float(pipe.audio_scheduler.config.shift)
        print(f"[gen] schedule from checkpoint: video shift {VIDEO_SHIFT}, audio shift {AUDIO_SHIFT}", flush=True)
        if VIDEO_SHIFT != 10.0:
            print(f"[gen] WARNING: expected video shift 10.0 for V2, checkpoint says {VIDEO_SHIFT}", flush=True)

        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. Per forward the cost is linear in
# the rows (matmuls) and quadratic (attention); the two decoders add a fixed linear term on top, independent of the
# step count. Fitted to this Space's own measured runs on ZeroGPU `xlarge`, net of placement:
# 8670 rows / 21 s, 21312 / 49 s, 37296 / 89 s. The curve reproduces all three to within a second.
_FWD_A, _FWD_B = 1.921e-4, 1.776e-9
_DECODE_A = 1.853e-4
_DUR_BASE = 5.0
# `lazy` placement: 72 GiB crosses PCIe once per cold worker, measured at ~11 s. Every request has to book it,
# because any request may be the one that lands on a cold one.
_PLACEMENT_ALLOWANCE = 12
# Booked over the estimate. Keep it small: an inflated duration burns the visitor's quota and drops queue priority.
_MARGIN = 1.15
# A ZeroGPU booking is refused outright past the per-call tier cap, so never ask for more than this.
_MAX_BOOKING = 280


def packed_rows(height: int, width: int, num_frames: int) -> int:
    """Rows in the packed video sequence the DiT attends over — the one number both cost terms scale with."""
    latent_frames = (int(num_frames) - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK * LATENTS_PER_CHUNK + 2
    return latent_frames * (int(height) // 32) * (int(width) // 32)


def estimate_seconds(height: int, width: int, num_frames: int) -> float:
    rows = packed_rows(height, width, num_frames)
    return NUM_FORWARDS * (_FWD_A * rows + _FWD_B * rows**2) + _DECODE_A * rows + _DUR_BASE


def get_duration(prompt_embeds, text_token_tags, height, width, num_frames, seed, *a, **k):
    estimate = estimate_seconds(height, width, num_frames) + _PLACEMENT_ALLOWANCE
    return min(_MAX_BOOKING, max(40, int(estimate * _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 eight-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,
        # Nine grid points, eight forwards. The checkpoint's own sidecar pins this and FastVideo's reference script
        # refuses any other value for this student.
        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")


def generate(
    prompt: str,
    canvas: str = DEFAULT_CANVAS,
    duration: float = DEFAULT_DURATION,
    upsample: bool = True,
    seed: int = 1000,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a video with a synchronized soundtrack from a text prompt, in eight 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.")
    if canvas not in CANVASES:
        raise gr.Error(f"Unknown canvas `{canvas}`. One of: {', '.join(CANVASES)}")

    from diffusers.utils import encode_video

    duration = min(max(float(duration), MIN_UI_DURATION), MAX_UI_DURATION)
    num_frames = snap_frames(duration)
    height, width = CANVASES[canvas]
    if estimate_seconds(height, width, num_frames) + _PLACEMENT_ALLOWANCE > _MAX_BOOKING / _MARGIN:
        raise gr.Error(
            f"`{canvas}` at {duration:.0f}s needs more GPU time than one ZeroGPU booking allows for eight "
            f"forwards. Pick a smaller canvas or a shorter 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()
    # A request that lands on a cold ZeroGPU worker can die mid-ladder with `GPU task aborted` (the 72 GiB lazy move
    # races the previous worker's release, or the worker is reset between requests). It is transient: retry once on
    # a fresh booking before surfacing the failure.
    for attempt in (1, 2):
        try:
            frames, audio, sampling_rate = _generate(prompt_embeds, text_token_tags, height, width, num_frames, seed)
            break
        except Exception as error:
            if attempt == 2 or "aborted" not in str(error).lower():
                raise
            print(f"[gen] GPU task aborted on attempt {attempt}, retrying on a fresh worker ...", flush=True)
    generate_seconds = time.time() - started

    directory = os.path.join(tempfile.gettempdir(), "fasth3-outputs")
    os.makedirs(directory, exist_ok=True)
    path = os.path.join(directory, f"fasth3-8step-{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"{packed_rows(height, width, num_frames)} packed rows · conditioner {condition_seconds:.0f}s "
        f"({plan['num_text_tokens']} tokens{', rewritten' if refined else ''}) · denoise + decode "
        f"{generate_seconds:.0f}s ({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 8-Step V2 — MiniMax-H3 in 8 forwards, with sound

<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://haoailab.com/blogs/fasth3-preview/" target="_blank" rel="noopener"><strong>[ blog ]</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 + VSA-H3 distillation** (80% sparsity,
step 1300) of [MiniMax-H3](https://huggingface.co/{BASE_REPO}) — the 33B dual-modality transformer that generates a
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 **{NUM_FORWARDS}-rung ladder** — `{', '.join(str(r) for r in DMD_RUNGS)}`
— for {NUM_FORWARDS} transformer forwards per video, under the V2-specific video scheduler shift of **10.0**.

Its 4-step sibling is faster; this one trades that back for motion and audio detail. FL2VA and Ref2VA were not
distilled, so this is text-to-video+audio only.
"""

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

# Every example prompt below is the authors' own. The first two are FastVideo's, out of the FastH3 reference
# examples in https://github.com/hao-ai-lab/FastVideo: the structured one-liner from
# `examples/inference/basic/basic_fasth3_simplified.sh` and the plain-prose request from
# `examples/inference/basic/basic_fasth3_spark.yaml` (seed 2026, its own). The last two are FastH3's published
# showcase prompts from the FastH3 blog post, both written for the 5 s / 124-frame operating point.
FASTVIDEO_FOX = (
    "integrated_multimodal_description: A red fox runs through fresh snow at dawn. "
    "overall_soundscape: Fast pawsteps in snow, winter wind, and distant birds."
)
FASTVIDEO_MEADOW = (
    "A wide cinematic shot of an alpine meadow at sunrise, pale pink mountain peaks above a blue valley filled "
    "with thin morning mist."
)
BLOG_MOON = (
    "integrated_multimodal_description: [Shot 1] 2D-animated in a whimsical storybook style. A wide shot of the "
    "night sky: a small, glowing moon with a friendly face, big expressive eyes and a gentle smile, surrounded by "
    "soft radiant light, floats gracefully across the dark blue sky above a cozy village below. Twinkling stars dot "
    "the sky and fluffy clouds drift past as the camera follows the moon in a Tracking Shot at slow speed. "
    "[Shot 2] At 00:02.600, the camera cuts to a medium close-up of the little moon as it gently descends toward "
    "the sleepy village, the camera Tilting Down with it at slow speed. Its warm glow spreads over quaint houses "
    "with thatched roofs, a winding river that reflects the moonlight, and a few villagers waking up: one opens "
    "wooden shutters, and another stretches and yawns in a doorway.\n"
    "overall_soundscape: A gentle night wind drifts through the air and crickets chirp softly around the village, "
    "while the winding river laps quietly against its banks. Wooden shutters creak open and a villager yawns.\n"
    "non_diegetic_music: A soft ambient score of sustained string pads and slow piano arpeggios plays at slow "
    "speed, with a light celesta melody entering in the second half. The dynamics stay quiet throughout, swelling "
    "slightly as the moon descends."
)
BLOG_SHIP = (
    "integrated_multimodal_description: [Shot 1] Live-action, cinematic. A sweeping wide shot captures a large "
    "wooden sailing ship navigating the vast open ocean beneath a starlit night sky, the camera performing an Arc "
    "Shot around the vessel at slow speed. The ship's tall sails billow gently in the wind, illuminated by the soft "
    "warm glow of lanterns hanging from the rigging, while moonlight and starlight shimmer across the subtly "
    "undulating waves. Weathered wooden planks, taut ropes, and towering masts are rendered in fine detail as the "
    "hull rocks with the swell, and in the background the deep blue of the sea blends into the dark expanse of a "
    "night sky densely dotted with stars, the lone ship small against the horizon.\n"
    "overall_soundscape: Wind hums steadily through the rigging and fills the canvas sails with soft flapping and "
    "billowing sounds. Wooden timbers and ropes creak rhythmically as the ship rolls, and waves lap and splash "
    "gently against the hull.\n"
    "non_diegetic_music: A slow orchestral score of sustained strings and low brass builds gradually in volume, "
    "layered with a sparse piano melody in a gentle, unhurried rhythm."
)

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

with gr.Blocks(title="FastH3 8-Step V2") 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=1000, precision=0)
                gr.Markdown(
                    f"Steps are fixed at the checkpoint's trained ladder — a {SIGMA_GRID_POINTS}-point sigma grid, "
                    f"{NUM_FORWARDS} transformer forwards on the rungs "
                    f"`{', '.join(str(r) for r in DMD_RUNGS)}`. FastVideo's own reference script for this "
                    "checkpoint refuses any other grid. There is no guidance scale and no negative prompt: the base "
                    "model is guidance-distilled. Larger canvases and longer durations cost proportionally more "
                    "ZeroGPU time; `1344x768 · 16:9 full` is the distillation's operating point."
                )

        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=[
            [FASTVIDEO_FOX, "960x544 · 16:9 fast", 5, False, 1000],
            [FASTVIDEO_MEADOW, "1024x576 · 16:9 fast", 5, True, 2026],
            [BLOG_MOON, "1024x576 · 16:9 fast", 5, False, 1000],
            [BLOG_SHIP, "1024x576 · 16:9 fast", 5, False, 1000],
        ],
        inputs=[prompt, canvas, duration, upsample, seed],
        outputs=[video, report, upsampled],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
        label="Examples — all four are the authors' own prompts (FastVideo's reference examples, then two FastH3 blog showcases)",
    )

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


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)