"""Open-Sora v2 demo — 11B text-to-video / image-to-video on ZeroGPU.

The checkpoints are not downloaded: hpcai-tech/Open-Sora-v2 (69 GB) and a bf16 T5
encoder are mounted read-only as Space volumes at /ckpts and /t5, which is what the
config paths below point at.
"""

import spaces  # noqa: F401  — must precede torch so the ZeroGPU CUDA hijack installs

import os
import random
import tempfile
import time
from contextlib import contextmanager

import gradio as gr
import torch
from mmengine.config import Config
from PIL import Image

from opensora.datasets import save_sample
from opensora.utils.inference import (
    add_fps_info_to_text,
    add_motion_score_to_text,
    modify_option_to_t2i,
)
from opensora.utils.sampling import (
    SamplingOption,
    prepare_api,
    sanitize_sampling_option,
)

CKPT = "/ckpts"
# T5-XXL ships in the model repo as 19 GB of fp32 pickle shards. This is the same
# encoder as 9.5 GB of bf16 safetensors (identical config and tokenizer), mounted
# separately — half the bytes, memory-mapped, and no fp32->bf16 cast.
T5_CKPT = "/t5"
DTYPE = torch.bfloat16
FPS_SAVE = 24
# 768px takes ~28 min/video on a single GPU — far past any ZeroGPU allocation.
RESOLUTION = "256px"

OUTPUT_DIR = os.path.join(tempfile.gettempdir(), "open-sora-outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Mirrors configs/diffusion/inference/t2i2v_256px.py, with ./ckpts repointed at the
# mounted volume and the T5 shardformer pass off (it needs colossalai, which pins
# torch<=2.5.1 and so cannot run on Blackwell).
cfg = Config(
    dict(
        model=dict(
            type="flux",
            from_pretrained=f"{CKPT}/Open_Sora_v2.safetensors",
            guidance_embed=False,
            fused_qkv=False,
            # Liger's Triton RoPE autotunes with CUDA events, which fails inside the
            # ZeroGPU worker; EmbedND is the pure-torch equivalent. Same reason as
            # FusedRMSNorm in opensora/models/mmdit/layers.py.
            use_liger_rope=False,
            in_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            cond_embed=True,
        ),
        ae=dict(
            type="hunyuan_vae",
            from_pretrained=f"{CKPT}/hunyuan_vae.safetensors",
            in_channels=3,
            out_channels=3,
            layers_per_block=2,
            latent_channels=16,
            use_spatial_tiling=True,
            use_temporal_tiling=False,
        ),
        t5=dict(
            type="text_embedder",
            from_pretrained=T5_CKPT,
            max_length=512,
            shardformer=False,
        ),
        clip=dict(
            type="text_embedder",
            from_pretrained=f"{CKPT}/openai/clip-vit-large-patch14",
            max_length=77,
        ),
        img_flux=dict(
            type="flux",
            from_pretrained=f"{CKPT}/flux1-dev.safetensors",
            guidance_embed=True,
            in_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            cond_embed=False,
        ),
        img_flux_ae=dict(
            type="autoencoder_2d",
            from_pretrained=f"{CKPT}/flux1-dev-ae.safetensors",
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    )
)

# The volume mounts are init steps outside this process, and when one fails the Space
# dies with an opaque "hf-mount failed" and no logs — so prove they are readable first.
for mount in (CKPT, T5_CKPT):
    if not os.path.isdir(mount):
        raise RuntimeError(
            f"{mount} is not mounted. This Space needs hpcai-tech/Open-Sora-v2 at "
            f"{CKPT} and city96/t5-v1_1-xxl-encoder-bf16 at {T5_CKPT}, both read-only."
        )
    print(f"Volume {mount} mounted:", flush=True)
    for entry in sorted(os.listdir(mount)):
        full = os.path.join(mount, entry)
        size = os.path.getsize(full) / 1e9 if os.path.isfile(full) else 0
        print(f"  {size:6.2f} GB  {entry}", flush=True)


@contextmanager
def _skip_weight_init():
    """Stop nn.Module.__init__ from random-initialising weights it is about to discard.

    `Flux()` builds the model under `with torch.device(device_map)`. Upstream passes
    "cuda", where the kaiming/normal fills are trivially fast. We build on CPU instead
    (constructing under a cuda device context would initialise real CUDA in the main
    process and defeat the ZeroGPU hijack) — and on CPU, initialising 23B parameters
    across Open-Sora and FLUX takes over an hour, which overran the launch timeout.

    Every one of those values is overwritten microseconds later by load_state_dict, so
    the fills are pure waste. Neutralising them turns model construction into little
    more than the allocation. Watch the load warnings for missing keys: a key absent
    from the checkpoint would now keep uninitialised memory.
    """
    names = (
        "uniform_", "normal_", "trunc_normal_",
        "kaiming_uniform_", "kaiming_normal_",
        "xavier_uniform_", "xavier_normal_",
    )
    saved = {n: getattr(torch.nn.init, n) for n in names if hasattr(torch.nn.init, n)}
    for name in saved:
        setattr(torch.nn.init, name, lambda tensor, *a, **k: tensor)
    try:
        yield
    finally:
        for name, fn in saved.items():
            setattr(torch.nn.init, name, fn)


# Replaces prepare_models() so each model — and construct-vs-load within it — is timed
# separately. A single opaque call gave no way to tell which part overran the timeout.
import opensora.utils.ckpt as _ckpt_mod
from opensora.registry import MODELS, build_module

_unpatched_load_checkpoint = _ckpt_mod.load_checkpoint


def _timed_load_checkpoint(model, path, *args, **kwargs):
    """Read the checkpoint sequentially into RAM, then deserialise from memory.

    safetensors.load_file memory-maps the file, so the load_state_dict that follows
    faults it in one tensor at a time — thousands of small synchronous reads against
    the network-backed volume. Measured: ~48 MB/s and 1% CPU, i.e. ~30 min for a 24 GB
    checkpoint. A bulk sequential read of the same file on the same mount runs at
    267 MB/s, and there is ~980 GB of RAM here, so slurping the bytes first turns this
    into roughly a 90-second read plus an in-memory copy.
    """
    if not str(path).endswith(".safetensors"):
        return _unpatched_load_checkpoint(model, path, *args, **kwargs)

    import gc

    from safetensors.torch import load as load_from_bytes

    start = time.perf_counter()
    with open(path, "rb") as fh:
        blob = fh.read()
    read_at = time.perf_counter()

    ckpt = load_from_bytes(blob)
    del blob
    parsed_at = time.perf_counter()

    missing, unexpected = model.load_state_dict(ckpt, strict=kwargs.get("strict", False))
    del ckpt
    gc.collect()

    if missing:
        print(f"    MISSING {len(missing)} keys e.g. {missing[:3]}", flush=True)
    if unexpected:
        print(f"    UNEXPECTED {len(unexpected)} keys e.g. {unexpected[:3]}", flush=True)
    print(
        f"    read={read_at - start:.0f}s parse={parsed_at - read_at:.0f}s "
        f"copy={time.perf_counter() - parsed_at:.0f}s",
        flush=True,
    )
    return model


_ckpt_mod.load_checkpoint = _timed_load_checkpoint
# Flux() imported load_checkpoint by value, so rebind it there too.
import opensora.models.mmdit.model as _mmdit_mod

_mmdit_mod.load_checkpoint = _timed_load_checkpoint


def _build(name: str, section):
    start = time.perf_counter()
    with _skip_weight_init():
        module = build_module(
            section, MODELS, device_map=torch.device("cpu"), torch_dtype=DTYPE
        ).eval()
    print(f"  [load] {name}: {time.perf_counter() - start:.0f}s total", flush=True)
    return module


# A load that makes no progress for 60 min looks identical whether it is CPU-bound or
# thrashing on swap. Report memory up front, then keep reporting it from a watchdog so
# the logs show which one is happening even if nothing else ever prints.
import faulthandler
import sys
import threading

import psutil

# Four theories about this hang have now been wrong. Rather than infer where the load
# is stuck, dump every thread's stack on a timer — that names the exact line.
faulthandler.dump_traceback_later(90, repeat=True, exit=False, file=sys.stderr)

_proc = psutil.Process()
_proc.cpu_percent()  # prime it; the first call always returns 0.0


def _mem() -> str:
    vm = psutil.virtual_memory()
    sw = psutil.swap_memory()
    return (
        f"rss={_proc.memory_info().rss/1e9:.1f}GB "
        f"used={vm.used/1e9:.0f}/{vm.total/1e9:.0f}GB "
        f"avail={vm.available/1e9:.0f}GB "
        f"swap={sw.used/1e9:.0f}/{sw.total/1e9:.0f}GB"
    )


print(f"[mem] at startup: {_mem()}  cpus={psutil.cpu_count()}", flush=True)


def _watchdog():
    while True:
        time.sleep(15)
        print(f"[mem] {_mem()} cpu={_proc.cpu_percent():.0f}%", flush=True)
        sys.stderr.flush()


threading.Thread(target=_watchdog, daemon=True).start()

# NOTE: do not try to speed this up with spaces.zero.torch.unpatch(). It works, and it
# breaks the Space: unpatch() restores the real torch._C._cuda_init, CUDA then gets
# initialised in the main process during loading, and every @spaces.GPU call afterwards
# dies in the forked worker with "Cannot re-initialize CUDA in forked subprocess".
# ZeroGPU requires CUDA to stay untouched in the parent process. The per-tensor
# interception costs ~1000 dispatches, which is minor; the real cost was the mmap I/O,
# handled in _timed_load_checkpoint and _stage_locally above.


def _stage_locally(src: str, name: str) -> str:
    """Bulk-copy an encoder directory to local disk before transformers loads it.

    transformers memory-maps safetensors too, so from_pretrained straight off the
    volume hits the same tensor-by-tensor fault storm as the diffusion checkpoints.
    These two are small enough (~11 GB combined) that a sequential copy at 267 MB/s
    is far cheaper than paying that cost, and there is 1.8 TB of local disk.
    """
    import shutil

    dst = os.path.join(tempfile.gettempdir(), "local-ckpts", name)
    if os.path.isdir(dst):
        return dst
    start = time.perf_counter()
    shutil.copytree(src, dst)
    total = sum(
        os.path.getsize(os.path.join(dst, f))
        for f in os.listdir(dst)
        if os.path.isfile(os.path.join(dst, f))
    )
    elapsed = time.perf_counter() - start
    print(
        f"  [stage] {name}: {total/1e9:.1f} GB in {elapsed:.0f}s "
        f"({total/1e6/max(elapsed, 1e-6):.0f} MB/s)",
        flush=True,
    )
    return dst


print("Staging text encoders to local disk...", flush=True)
cfg.t5.from_pretrained = _stage_locally(T5_CKPT, "t5")
# HFEmbedder picks CLIP vs T5 with `"openai" in from_pretrained`, so the staged path
# has to keep "openai" in it or CLIP gets built as a T5.
cfg.clip.from_pretrained = _stage_locally(
    f"{CKPT}/openai/clip-vit-large-patch14", "openai/clip-vit-large-patch14"
)

print("Loading Open-Sora v2 weights...", flush=True)
_t0 = time.perf_counter()
model = _build("open-sora (11B)", cfg.model)
model_ae = _build("hunyuan vae", cfg.ae)
model_t5 = _build("t5-xxl", cfg.t5)
model_clip = _build("clip-l", cfg.clip)
optional_models = {
    "img_flux": _build("flux (12B)", cfg.img_flux),
    "img_flux_ae": _build("flux vae", cfg.img_flux_ae),
}
print(f"Weights read in {time.perf_counter() - _t0:.0f}s", flush=True)

# Intercepted by the ZeroGPU hijack: registers the fake CUDA tensors and drives the
# startup pack step, so the forked GPU worker can stream the weights into VRAM.
_t0 = time.perf_counter()
for module in (model, model_ae, model_t5, model_clip, *optional_models.values()):
    module.to("cuda")
print(f"Weights staged in {time.perf_counter() - _t0:.0f}s", flush=True)

api_fn = prepare_api(model, model_ae, model_t5, model_clip, optional_models)
api_fn_img = prepare_api(
    optional_models["img_flux"],
    optional_models["img_flux_ae"],
    model_t5,
    model_clip,
    optional_models,
)


@contextmanager
def _phase(name: str):
    """Log wall time and peak VRAM per phase, to size `duration` and large-vs-xlarge."""
    torch.cuda.reset_peak_memory_stats()
    start = time.perf_counter()
    yield
    peak = torch.cuda.max_memory_allocated() / 1e9
    reserved = torch.cuda.max_memory_reserved() / 1e9
    print(
        f"[timing] {name}: {time.perf_counter() - start:.1f}s  "
        f"peak_alloc={peak:.1f}GB reserved={reserved:.1f}GB",
        flush=True,
    )


def _sampling_option(aspect_ratio: str, num_frames: int, num_steps: int, seed: int):
    return sanitize_sampling_option(
        SamplingOption(
            resolution=RESOLUTION,
            aspect_ratio=aspect_ratio,
            num_frames=num_frames,
            num_steps=num_steps,
            shift=True,
            temporal_reduction=4,
            is_causal_vae=True,
            guidance=7.5,
            guidance_img=3.0,
            text_osci=True,
            image_osci=True,
            scale_temporal_osci=True,
            method="i2v",
            seed=seed,
        )
    )


def _budget(num_frames: int, num_steps: int, with_flux: bool) -> int:
    """Worst-case seconds for one generation, used as the ZeroGPU duration."""
    video = 25 + num_steps * (num_frames / 129.0) * 2.2
    return int(video + (55 if with_flux else 0))


def _t2v_duration(prompt, aspect_ratio, num_frames, num_steps, seed, randomize):
    return _budget(num_frames, num_steps, with_flux=True)


def _i2v_duration(image, prompt, aspect_ratio, num_frames, num_steps, seed, randomize):
    return _budget(num_frames, num_steps, with_flux=False)


def _place_models_on_gpu(*modules):
    for module in modules:
        module.to("cuda")


def _repair_inductor_cuda_interface() -> None:
    """Restore real CUDA functions in PyTorch Inductor after ZeroGPU unpatch."""
    try:
        from torch._dynamo.device_interface import CudaInterface

        for attr, cuda_attr in (
            ("synchronize", "synchronize"),
            ("current_device", "current_device"),
            ("device_count", "device_count"),
            ("get_device_properties", "get_device_properties"),
            ("exchange_device", "_exchange_device"),
            ("maybe_exchange_device", "_maybe_exchange_device"),
        ):
            setattr(
                CudaInterface,
                attr,
                staticmethod(getattr(torch.cuda, cuda_attr)),
            )

        print("[ZeroGPU] Inductor CUDA interface repaired.", flush=True)

    except Exception as exc:
        print(f"[ZeroGPU] Inductor repair skipped: {exc}", flush=True)

def _repair_exchange_device_alias() -> None:
    """Work around a bug in spaces 0.51.3 that breaks dynamo on torch 2.13.

    ZeroGPU's forked worker calls its own unpatch() to hand real CUDA back to user
    code, and patching.py:354 restores the wrong original:

        setattr(torch.cuda, '_maybe_exchange_device', _cuda_exchange_device)

    so `_exchange_device` and `_maybe_exchange_device` end up as the same object.
    torch._dynamo keys its handler table by function object and refuses duplicates, so
    the first @torch._disable_dynamo call in the worker dies with
    "AssertionError: Handler already registered for <built-in function
    _cuda_exchangeDevice>". Re-apply the originals the library saved at import.
    """
    from spaces.zero.torch import patching as zero_patching

    original = getattr(zero_patching, "_cuda_maybe_exchange_device", None)
    if original is not None and torch.cuda._maybe_exchange_device is not original:
        torch.cuda._maybe_exchange_device = original
    original_exchange = getattr(zero_patching, "_cuda_exchange_device", None)
    if original_exchange is not None:
        torch.cuda._exchange_device = original_exchange


def _resolve_seed(seed: int, randomize: bool) -> int:
    return random.randint(0, 2**31 - 1) if randomize else int(seed)


def _run_video(prompt: str, ref_path: str, opt: SamplingOption, seed: int) -> str:
    """Sample the video model from a reference image and return the saved .mp4 path."""
    text = add_fps_info_to_text([prompt], fps=FPS_SAVE)
    text = add_motion_score_to_text(text, 4)

    with _phase(f"video {opt.num_frames}f/{opt.num_steps}steps"):
        x = api_fn(
            opt,
            "i2v_head",
            seed=seed,
            text=text,
            ref=[ref_path],
            patch_size=2,
            channel=cfg.model.in_channels,
        ).cpu()

    save_path = os.path.join(OUTPUT_DIR, f"video_{seed}_{random.randint(0, 10**9)}")
    return save_sample(x[0], save_path=save_path, fps=FPS_SAVE, verbose=False)


@spaces.GPU(duration=_t2v_duration, size="xlarge")
def generate_from_text(
    prompt: str,
    aspect_ratio: str = "16:9",
    num_frames: int = 97,
    num_steps: int = 50,
    seed: int = 42,
    randomize_seed: bool = True,
) -> str:
    """Generate a video from a text prompt with Open-Sora v2.

    The prompt is first rendered to a still frame with FLUX.1-dev, then animated by the
    11B Open-Sora video model — the text-to-image-to-video pipeline the authors
    recommend for text-to-video.

    Args:
        prompt: Description of the video to generate.
        aspect_ratio: Frame shape, one of "16:9", "9:16", "1:1", "2.39:1".
        num_frames: Length in frames; must be 4k+1 and at most 129 (24 fps output).
        num_steps: Diffusion sampling steps for the video model.
        seed: Random seed, used when randomize_seed is False.
        randomize_seed: Draw a fresh random seed instead of using `seed`.

    Returns:
        Path to the generated mp4.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    _repair_exchange_device_alias()
    _place_models_on_gpu(
    model,
    model_ae,
    model_t5,
    model_clip,
    )
    _repair_inductor_cuda_interface()
    seed = _resolve_seed(seed, randomize_seed)
    opt = _sampling_option(aspect_ratio, num_frames, num_steps, seed)

    # 1. prompt -> still frame, via FLUX at 768px
    opt_t2i = modify_option_to_t2i(opt, distilled=True, img_resolution="768px")
    with _phase("flux t2i"):
        x_cond = api_fn_img(
            opt_t2i,
            "t2v",
            seed=seed,
            text=[prompt],
            channel=cfg.img_flux.in_channels,
        ).cpu()

    cond_path = os.path.join(OUTPUT_DIR, f"cond_{seed}_{random.randint(0, 10**9)}")
    cond_path = save_sample(x_cond[0], save_path=cond_path, fps=1, verbose=False)

    # 2. still frame -> video
    return _run_video(prompt, cond_path, opt, seed)


@spaces.GPU(duration=_i2v_duration, size="xlarge")
def generate_from_image(
    image: str,
    prompt: str,
    aspect_ratio: str = "16:9",
    num_frames: int = 97,
    num_steps: int = 50,
    seed: int = 42,
    randomize_seed: bool = True,
) -> str:
    """Animate a still image into a video with Open-Sora v2.

    This is the model's native mode: the uploaded image becomes the first frame and the
    prompt describes how the scene should move.

    Args:
        image: Filepath of the reference image to use as the first frame.
        prompt: Description of the motion and the scene.
        aspect_ratio: Frame shape, one of "16:9", "9:16", "1:1", "2.39:1".
        num_frames: Length in frames; must be 4k+1 and at most 129 (24 fps output).
        num_steps: Diffusion sampling steps for the video model.
        seed: Random seed, used when randomize_seed is False.
        randomize_seed: Draw a fresh random seed instead of using `seed`.

    Returns:
        Path to the generated mp4.
    """
    if image is None:
        raise gr.Error("Please upload a reference image.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt describing the motion.")

    _repair_exchange_device_alias()
    _place_models_on_gpu(
    model,
    model_ae,
    model_t5,
    model_clip,
    )
    seed = _resolve_seed(seed, randomize_seed)
    opt = _sampling_option(aspect_ratio, num_frames, num_steps, seed)
    return _run_video(prompt, image, opt, seed)


ASPECT_RATIOS = ["16:9", "9:16", "1:1", "2.39:1"]
FRAME_CHOICES = [33, 49, 65, 81, 97, 113, 129]

T2V_EXAMPLES = [
    ["raining, sea"],
    [
        "A breathtaking aerial shot of a rugged coastline at golden hour, waves crashing "
        "against dark volcanic cliffs, sea spray catching the warm light."
    ],
    [
        "A steaming cup of coffee on a wooden table by a rain-streaked window, soft "
        "morning light, shallow depth of field."
    ],
    [
        "A neon-lit Tokyo alley at night after rain, reflections shimmering in puddles, "
        "steam rising from a ramen stall."
    ],
]

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
"""

with gr.Blocks(css=CSS, title="Open-Sora v2") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🎬 Open-Sora v2

            An 11B video generation model from
            [HPC-AI Tech](https://huggingface.co/hpcai-tech/Open-Sora-v2), running at
            **256px / 24 fps** on ZeroGPU.
            [Paper](https://huggingface.co/papers/2503.09642) ·
            [GitHub](https://github.com/hpcaitech/Open-Sora)

            A generation takes a couple of minutes and uses a sizeable chunk of your
            daily ZeroGPU quota, so start with the shorter frame counts.
            """
        )

        with gr.Tabs():
            with gr.Tab("Text → Video"):
                with gr.Row():
                    with gr.Column():
                        t2v_prompt = gr.Textbox(
                            label="Prompt",
                            placeholder="Describe the video you want",
                            lines=3,
                        )
                        t2v_button = gr.Button("Generate video", variant="primary")
                        gr.Markdown(
                            "Text-to-video runs through FLUX.1-dev first to draw a "
                            "starting frame — the pipeline the authors recommend."
                        )
                        with gr.Accordion("Settings", open=False):
                            t2v_aspect = gr.Radio(
                                ASPECT_RATIOS, value="16:9", label="Aspect ratio"
                            )
                            t2v_frames = gr.Radio(
                                FRAME_CHOICES,
                                value=97,
                                label="Frames",
                                info="97 frames ≈ 4 s at 24 fps",
                            )
                            t2v_steps = gr.Slider(
                                10, 50, value=50, step=1, label="Sampling steps"
                            )
                            t2v_randomize = gr.Checkbox(
                                value=True, label="Randomize seed"
                            )
                            t2v_seed = gr.Slider(
                                0, 2**31 - 1, value=42, step=1, label="Seed"
                            )
                    with gr.Column():
                        t2v_output = gr.Video(label="Result", autoplay=True)

                gr.Examples(
                    examples=T2V_EXAMPLES,
                    inputs=[t2v_prompt],
                    outputs=t2v_output,
                    fn=generate_from_text,
                    cache_examples=True,
                    cache_mode="lazy",
                )

            with gr.Tab("Image → Video"):
                with gr.Row():
                    with gr.Column():
                        i2v_image = gr.Image(
                            label="First frame", type="filepath", height=280
                        )
                        i2v_prompt = gr.Textbox(
                            label="Prompt",
                            placeholder="Describe how the scene should move",
                            lines=3,
                        )
                        i2v_button = gr.Button("Animate image", variant="primary")
                        gr.Markdown(
                            "Image-to-video is what Open-Sora v2 was tuned for, and it "
                            "skips the FLUX step — so it is the faster of the two."
                        )
                        with gr.Accordion("Settings", open=False):
                            i2v_aspect = gr.Radio(
                                ASPECT_RATIOS, value="16:9", label="Aspect ratio"
                            )
                            i2v_frames = gr.Radio(
                                FRAME_CHOICES,
                                value=97,
                                label="Frames",
                                info="97 frames ≈ 4 s at 24 fps",
                            )
                            i2v_steps = gr.Slider(
                                10, 50, value=50, step=1, label="Sampling steps"
                            )
                            i2v_randomize = gr.Checkbox(
                                value=True, label="Randomize seed"
                            )
                            i2v_seed = gr.Slider(
                                0, 2**31 - 1, value=42, step=1, label="Seed"
                            )
                    with gr.Column():
                        i2v_output = gr.Video(label="Result", autoplay=True)

                gr.Examples(
                    examples=[
                        [
                            "assets/texts/i2v.png",
                            "A plump pig wallows in a muddy pond on a rustic farm, its "
                            "pink snout poking out as it snorts contentedly. The camera "
                            "captures the pig's playful splashes, sending ripples "
                            "through the water under the midday sun.",
                        ]
                    ],
                    inputs=[i2v_image, i2v_prompt],
                    outputs=i2v_output,
                    fn=generate_from_image,
                    cache_examples=True,
                    cache_mode="lazy",
                )

    t2v_button.click(
        fn=generate_from_text,
        inputs=[t2v_prompt, t2v_aspect, t2v_frames, t2v_steps, t2v_seed, t2v_randomize],
        outputs=t2v_output,
        api_name="generate_from_text",
    )
    i2v_button.click(
        fn=generate_from_image,
        inputs=[
            i2v_image,
            i2v_prompt,
            i2v_aspect,
            i2v_frames,
            i2v_steps,
            i2v_seed,
            i2v_randomize,
        ],
        outputs=i2v_output,
        api_name="generate_from_image",
    )

demo.queue(max_size=8, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(mcp_server=True)