app.py

"""DreamX-Creator 1.0 — native joint audio-video generation on ZeroGPU.

Image + prompt -> a video whose soundtrack is denoised jointly with the frames
by the same model (gated A2V / V2A cross-attention), so speech, foley and
ambience stay in sync with the picture.

The inference path is the authors' own release code (`videox_fun/` +
`dreamx_inference.py`, copied verbatim from the reference Space
`hugging-apps/gd-ml-dreamx-creator` / `AMAP-ML/DreamX-Creator`); this file only
wires it into Gradio and stages the checkpoint download so the container never
holds all 43 GB of fp32 weights on disk at once.

Adapted for the AI Shorts Factory backend (Space `text_amon_API`):

- `image` is optional: an empty value yields a neutral keyframe (Option A), so
  scenes without a first-frame image can still generate.
- Returns ``(mp4, last-frame PNG, seed)`` instead of ``(mp4, seed)``: the last
  frame is what scene 2..N sends back as the first frame for continuity.
- Checkpoint root can be redirected to persistent storage with
  ``DREAMX_CKPT_DIR`` (default: this repo's ``./checkpoints``, the proven
  upstream layout).

Served through Gradio's `/call/generate` protocol (``api_name="generate"``).
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# videox_fun falls back to torch SDPA when flash-attn is absent; the numerics are
# identical here because every sequence in the batch is full-length (no padding).
os.environ.setdefault("VIDEOX_ATTENTION_TYPE", "FLASH_ATTENTION")

import spaces  # noqa: E402  (must precede torch)

import math  # noqa: E402
import random  # noqa: E402
import shutil  # noqa: E402
import subprocess  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
from pathlib import Path  # noqa: E402
from types import SimpleNamespace  # noqa: E402

import numpy as np  # noqa: E402
import torch  # noqa: E402

# Trusted upstream .pth checkpoints (UMT5 encoder, Wan2.2 VAE) are plain tensor
# dicts saved before torch 2.6 flipped `weights_only` to True.
_torch_load = torch.load


def _torch_load_compat(*args, **kwargs):
    kwargs.setdefault("weights_only", False)
    return _torch_load(*args, **kwargs)


torch.load = _torch_load_compat

import gradio as gr  # noqa: E402
from diffusers import FlowMatchEulerDiscreteScheduler  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from PIL import Image  # noqa: E402
from transformers import AutoTokenizer  # noqa: E402

from dreamx_inference import (  # noqa: E402
    DEFAULT_NEGATIVE_PROMPT,
    DirectionalMultimodalCFGAdapter,
    filter_kwargs,
    generate_joint_audio_video,
)
from videox_fun.models import AutoencoderKLWan3_8, WanT5EncoderModel  # noqa: E402
from videox_fun.models.creator_dac_vae import CreatorDACVAE  # noqa: E402
from videox_fun.models.creator_gating import WanCreatorGatingAVModel  # noqa: E402

REPO_ID = "GD-ML/DreamX-Creator"
APP_DIR = Path(__file__).parent.resolve()
CKPT_DIR = Path(os.environ.get("DREAMX_CKPT_DIR", str(APP_DIR / "checkpoints")))
CONFIG_PATH = APP_DIR / "config" / "config.yaml"
WEIGHT_DTYPE = torch.bfloat16
FPS = 24
CACHE_VERSION = 1

# Authors' defaults (audio_video_generation/inference.py + inference.sh).
VIDEO_BRIDGE_GUIDANCE = 3.5
AUDIO_BRIDGE_GUIDANCE = 3.5
VIDEO_SHIFT = 5.0
AUDIO_SHIFT = 5.0

RESOLUTION_CHOICES = [
    ("Fast — ~360p", 220),
    ("Balanced — ~480p", 440),
    ("Sharp — ~600p", 660),
]

# Neutral keyframe (Option A) dimensions — portrait 9:16, shorts-first.
KEYFRAME_WIDTH = int(os.environ.get("DREAMX_KEYFRAME_WIDTH", "720"))
KEYFRAME_HEIGHT = int(os.environ.get("DREAMX_KEYFRAME_HEIGHT", "1280"))


def _log(msg: str) -> None:
    print(f"[dreamx] {msg}", flush=True)


def _disk() -> str:
    total, used, free = shutil.disk_usage("/")
    return f"disk used={used / 2**30:.1f}GB free={free / 2**30:.1f}GB"


# --------------------------------------------------------------------------- #
# Model loading (module scope, eagerly on "cuda" — ZeroGPU packs from here)
# --------------------------------------------------------------------------- #

_cfg = OmegaConf.load(CONFIG_PATH)
_video_kwargs = OmegaConf.to_container(_cfg["video_transformer_additional_kwargs"], resolve=True)
_audio_kwargs = OmegaConf.to_container(_cfg["audio_transformer_additional_kwargs"], resolve=True)
_gating_kwargs = OmegaConf.to_container(_cfg["creator_gating_kwargs"], resolve=True)
_video_vae_kwargs = OmegaConf.to_container(_cfg["video_vae_kwargs"], resolve=True)
_text_encoder_kwargs = OmegaConf.to_container(_cfg["text_encoder_kwargs"], resolve=True)
_scheduler_kwargs = OmegaConf.to_container(_cfg["scheduler_kwargs"], resolve=True)
MAX_SEQUENCE_LENGTH = int(_text_encoder_kwargs.get("text_length", 512))

_log(f"checkpoints in {CKPT_DIR} ({_disk()})")
_log(f"downloading joint AV generator ... ({_disk()})")
snapshot_download(REPO_ID, local_dir=str(CKPT_DIR), allow_patterns=["creator/*", "creator/**/*"])
_log(f"loading joint AV generator ... ({_disk()})")

transformer = WanCreatorGatingAVModel.from_pretrained(
    pretrained_model_path=str(CKPT_DIR / "creator"),
    video_pretrained_model_path=str(CKPT_DIR / "creator" / "video_model"),
    audio_pretrained_model_path=str(CKPT_DIR / "creator" / "audio_model"),
    video_subfolder=_video_kwargs.get("transformer_low_noise_model_subpath", None),
    audio_subfolder=_audio_kwargs.get("transformer_low_noise_model_subpath", None),
    video_kwargs=_video_kwargs,
    audio_kwargs=_audio_kwargs,
    low_cpu_mem_usage=True,
    torch_dtype=WEIGHT_DTYPE,
    use_temporal_rope=_gating_kwargs.get("use_temporal_rope", True),
    audio_fps=_gating_kwargs.get("audio_fps", 48000.0 / 960.0),
    vae_temporal_stride=_gating_kwargs.get("vae_temporal_stride", 4),
    a2v_cross_attn_layers=_gating_kwargs.get("a2v_cross_attn_layers", None),
    v2a_cross_attn_layers=_gating_kwargs.get("v2a_cross_attn_layers", None),
    use_gating=_gating_kwargs.get("use_gating", True),
    zero_init_cross_attn=_gating_kwargs.get("zero_init_cross_attn", False),
    zero_init_gating=_gating_kwargs.get("zero_init_gating", True),
    gate_init_value=_gating_kwargs.get("gate_init_value", 0.0),
    a2v_gate_alphas=_gating_kwargs.get("a2v_gate_alphas", None),
    v2a_gate_alphas=_gating_kwargs.get("v2a_gate_alphas", None),
)
transformer.eval()
# fp32 source shards are no longer needed once the bf16 model is in RAM.
shutil.rmtree(CKPT_DIR / "creator", ignore_errors=True)
_log(f"joint AV generator loaded ({_disk()})")

_log("downloading VAEs + UMT5-xxl text encoder ...")
snapshot_download(
    REPO_ID,
    local_dir=str(CKPT_DIR),
    allow_patterns=["audio_vae/*", "wan2.2_ti2v_5b/*", "wan2.2_ti2v_5b/**/*"],
)
_log(f"loading VAEs + text encoder ... ({_disk()})")

_wan_dir = CKPT_DIR / "wan2.2_ti2v_5b"
_video_vae_path = _wan_dir / _video_vae_kwargs.get("vae_subpath", "Wan2.2_VAE.pth")
video_vae = AutoencoderKLWan3_8.from_pretrained(
    str(_video_vae_path), additional_kwargs=_video_vae_kwargs
).eval()
audio_vae = CreatorDACVAE.from_pretrained(str(CKPT_DIR / "audio_vae"), strict=False).eval()

tokenizer = AutoTokenizer.from_pretrained(
    str(_wan_dir / _text_encoder_kwargs.get("tokenizer_subpath", "google/umt5-xxl"))
)
_text_encoder_path = _wan_dir / _text_encoder_kwargs.get(
    "text_encoder_subpath", "models_t5_umt5-xxl-enc-bf16.pth"
)
text_encoder = WanT5EncoderModel.from_pretrained(
    str(_text_encoder_path),
    additional_kwargs=_text_encoder_kwargs,
    low_cpu_mem_usage=True,
    torch_dtype=WEIGHT_DTYPE,
).eval()

for _stale in (_video_vae_path, _text_encoder_path):
    try:
        os.remove(_stale)
    except OSError:
        pass
shutil.rmtree(CKPT_DIR / "audio_vae", ignore_errors=True)

transformer = DirectionalMultimodalCFGAdapter(
    transformer,
    video_scale=VIDEO_BRIDGE_GUIDANCE,
    audio_scale=AUDIO_BRIDGE_GUIDANCE,
    enable_a2v=True,
    enable_v2a=True,
).eval()

transformer.to("cuda")
text_encoder.to("cuda")
video_vae.to("cuda")
audio_vae.to("cuda")
_log(f"all models resident on cuda ({_disk()})")


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #

def _latent_frames(duration: float) -> int:
    num_frames = int(duration * FPS)
    num_frames = int((num_frames - 1) // 4 * 4) + 1
    return (num_frames - 1) // 4 + 1


DURATION_CAP = 300


def _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) -> float:
    """Predicted GPU seconds, fit to two measured ZeroGPU runs.

    Measured on RTX PRO 6000 (sm_120), bf16, 3-branch directional multimodal CFG:
      2640 video tokens x 10 steps -> 10.77s denoise, 16.0s total
      7560 video tokens x 30 steps -> 95.04s denoise, 104.7s total
    """
    tokens = _latent_frames(float(seconds)) * int(resolution_tokens)
    # per denoising step: linear (FFN/proj) + quadratic (self-attention) terms
    per_step = 4.02e-4 * tokens + 2.256e-9 * tokens * tokens
    denoise = per_step * int(num_inference_steps)
    overhead = 4.5 + 7.0e-4 * tokens  # T5 encode, first-frame VAE encode, decode, mux
    return denoise + overhead


def _estimate_duration(
    image=None,
    prompt="",
    seconds=3.0,
    num_inference_steps=30,
    resolution_tokens=440,
    *args,
    **kwargs,
):
    """GPU seconds to reserve — calibrated against measured runs on ZeroGPU."""
    raw = _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens)
    return int(min(DURATION_CAP, math.ceil(raw * 1.15) + 5))


def _write_mp4(frames: np.ndarray, audio: np.ndarray, sample_rate: int, fps: int) -> str:
    """Mux uint8 RGB frames + mono float audio into a single H.264/AAC mp4."""
    height, width = frames.shape[1], frames.shape[2]
    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    wav_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name

    import soundfile as sf

    sf.write(wav_path, np.clip(audio, -1.0, 1.0), sample_rate)

    base = [
        "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
        "-f", "rawvideo", "-pix_fmt", "rgb24",
        "-s", f"{width}x{height}", "-r", str(fps), "-i", "-",
        "-i", wav_path,
    ]
    for vcodec in ("libx264", "mpeg4"):
        cmd = base + [
            "-c:v", vcodec, "-pix_fmt", "yuv420p", "-crf", "18",
            "-c:a", "aac", "-b:a", "192k", "-shortest", out_path,
        ]
        if vcodec == "mpeg4":
            cmd.remove("-crf")
            cmd.remove("18")
        proc = subprocess.run(cmd, input=frames.tobytes(), capture_output=True)
        if proc.returncode == 0 and os.path.getsize(out_path) > 0:
            os.remove(wav_path)
            return out_path
        _log(f"ffmpeg ({vcodec}) failed: {proc.stderr.decode()[-600:]}")
    raise gr.Error("ffmpeg failed to encode the generated video.")


def _neutral_keyframe(width: int = KEYFRAME_WIDTH, height: int = KEYFRAME_HEIGHT) -> str:
    """Deterministic dark diagonal gradient — the "Option A" neutral first frame.

    By default a 9:16 portrait so generated clips keep a shorts-friendly aspect
    ratio when no real first-frame image is provided.
    """
    y, x = np.mgrid[0:height, 0:width]
    norm = np.sqrt((x / max(width - 1, 1)) ** 2 + (y / max(height - 1, 1)) ** 2)
    tone = (0.06 + 0.16 * norm[..., None] * np.array([0.86, 0.95, 1.0])).clip(0, 1)
    frame = (tone * 255.0).astype(np.uint8)
    path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    Image.fromarray(frame).save(path)
    return path


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #


@spaces.GPU(duration=_estimate_duration)
def generate(
    image: str,
    prompt: str,
    seconds: float = 3.0,
    num_inference_steps: int = 30,
    resolution_tokens: int = 440,
    guidance_scale: float = 5.0,
    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    seed: int = 113,
    randomize_seed: bool = False,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a video with a jointly-denoised soundtrack from a first frame and a prompt.

    Args:
        image: path to the image used as the first frame of the video. Empty
            means "no image": a neutral keyframe is generated in-app instead.
        prompt: description of the action AND the sound to generate; put spoken
            lines in quotes (e.g. `Man says, 'hello there.'`).
        seconds: length of the clip in seconds (24 fps).
        num_inference_steps: number of flow-matching denoising steps.
        resolution_tokens: spatial token budget; higher means higher resolution.
        guidance_scale: classifier-free guidance strength for text.
        negative_prompt: what to avoid in both the video and the audio.
        seed: RNG seed for reproducible sampling.
        randomize_seed: draw a fresh random seed instead of using `seed`.

    Returns:
        A tuple of (path to the generated mp4 with audio, path to the last-frame
        PNG for chaining into the next scene, the seed actually used).
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please provide a prompt describing the motion and the sound.")
    if not image:
        _log("empty first-frame image -> neutral keyframe")
        image = _neutral_keyframe()

    # ZeroGPU kills a task that outruns its reservation, and the reservation is
    # capped, so refuse the few extreme knob combinations that cannot fit.
    if _raw_gpu_seconds(seconds, num_inference_steps, resolution_tokens) * 1.15 + 5 > DURATION_CAP:
        raise gr.Error(
            f"That combination needs more than the {DURATION_CAP}s GPU budget. "
            "Lower the duration, the resolution, or the number of steps."
        )

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)

    device = torch.device("cuda")
    video_scheduler_kwargs = dict(_scheduler_kwargs, shift=VIDEO_SHIFT)
    audio_scheduler_kwargs = dict(_scheduler_kwargs, shift=AUDIO_SHIFT)

    models = {
        "config": _cfg,
        "transformer": transformer,
        "video_vae": video_vae,
        "audio_vae": audio_vae,
        "tokenizer": tokenizer,
        "text_encoder": text_encoder,
        # fresh schedulers per request: they carry mutable step state
        "video_scheduler": FlowMatchEulerDiscreteScheduler(
            **filter_kwargs(FlowMatchEulerDiscreteScheduler, video_scheduler_kwargs)
        ),
        "audio_scheduler": FlowMatchEulerDiscreteScheduler(
            **filter_kwargs(FlowMatchEulerDiscreteScheduler, audio_scheduler_kwargs)
        ),
        "max_sequence_length": MAX_SEQUENCE_LENGTH,
    }

    args = SimpleNamespace(
        config_path=str(CONFIG_PATH),
        image=image,
        output="output.mp4",
        negative_prompt=negative_prompt or DEFAULT_NEGATIVE_PROMPT,
        duration=float(seconds),
        target_spatial_tokens=int(resolution_tokens),
        min_token_ratio=0.95,
        fps=FPS,
        num_inference_steps=int(num_inference_steps),
        guidance_scale=float(guidance_scale),
        cfg_mode="multimodal",
        video_bridge_guidance_scale=VIDEO_BRIDGE_GUIDANCE,
        audio_bridge_guidance_scale=AUDIO_BRIDGE_GUIDANCE,
        seed=seed,
        video_shift=VIDEO_SHIFT,
        audio_shift=AUDIO_SHIFT,
        flow_match_mu=None,
        sampler_name="Flow",
        weight_dtype="bfloat16",
        GPU_memory_mode="model_full_load",
        text_encoder_cpu_offload=False,
        video_vae_cpu_offload=False,
        audio_vae_cpu_offload=False,
        vae_cpu_offload=False,
        use_temporal_rope=True,
        audio_fps=48000.0 / 960.0,
        vae_temporal_stride=4,
        disable_a2v_cross_attn=False,
        disable_v2a_cross_attn=False,
        suppress_aux_writes=True,
        skip_output_decode=False,
        disable_progress=False,
        synchronize_noise=False,
        ulysses_degree=1,
        ring_degree=1,
        fsdp_dit=False,
    )

    item = {
        "prompt": prompt.strip(),
        "video_prompt": prompt.strip(),
        "audio_prompt": prompt.strip(),
        "negative_prompt": args.negative_prompt,
        "audio_negative_prompt": args.negative_prompt,
        "duration": args.duration,
        "guidance_scale": args.guidance_scale,
        "num_inference_steps": args.num_inference_steps,
        "seed": seed,
        "first_frame_path": image,
        "name": "sample",
    }

    started = time.perf_counter()
    video_decoded, audio_decoded, num_frames = generate_joint_audio_video(
        args, models, device, WEIGHT_DTYPE, item
    )
    gpu_seconds = time.perf_counter() - started

    frames = (
        video_decoded[0].permute(1, 2, 3, 0).clamp(0, 1).numpy() * 255.0
    ).astype(np.uint8)
    waveform = audio_decoded.detach().float().cpu()
    while waveform.ndim > 1:
        waveform = waveform[0]
    audio = waveform.numpy()

    out_path = _write_mp4(frames, audio, int(audio_vae.sample_rate), FPS)

    last_frame_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
    Image.fromarray(frames[-1]).save(last_frame_path)

    _log(
        f"done in {gpu_seconds:.1f}s | {num_frames} frames @ {frames.shape[2]}x{frames.shape[1]} "
        f"| steps={args.num_inference_steps} tokens={args.target_spatial_tokens} "
        f"| reserved={_estimate_duration(image, prompt, seconds, num_inference_steps, resolution_tokens)}s"
    )
    return out_path, last_frame_path, seed


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #

EXAMPLES = [
    [
        "examples/case1.jpg",
        "A man in a dark suit and white shirt is seated on a yellow couch, speaking about "
        "the language of Americans. He uses hand gestures to emphasize his points, and the "
        "background shows a cityscape with illuminated buildings, suggesting an urban setting, "
        "possibly a studio with a city view. Man says, 'The thing about Americans that I've "
        "thought about the language is that they speak, they say they speak English.'.",
    ],
    [
        "examples/case4.jpg",
        "The video shows a wolf standing on its hind legs, howling with its mouth wide open, "
        "showing its teeth and tongue. The wolf's ears are perked up, and its eyes are focused "
        "on something in the distance. The background consists of trees with green and yellow "
        "leaves, indicating it is autumn. The wolf's howl is loud and resonant, filling the air "
        "with its powerful voice. The sound of a dog howls and howls can be heard.",
    ],
    [
        "examples/case3.jpg",
        "The video captures a dramatic night scene with a series of lightning strikes "
        "illuminating the dark sky and revealing the city lights below. The clouds move across "
        "the sky, and the lightning strikes again, followed by a thunderclap. The sound of a "
        "thunderstorm and rain falling can be heard.",
    ],
    [
        "examples/case2.jpg",
        "A humanoid robot is cooking in a modern kitchen. The robot, with a white body and blue "
        "eyes, is stirring food in a pan on the stove. The kitchen is equipped with dark cabinets "
        "and various utensils. The robot's movements are smooth and precise as it stirs the food, "
        "causing steam to rise. The ambient sound of cooking can be heard throughout the scene.",
    ],
    [
        "examples/case5.jpg",
        "A person in a red plaid shirt is typing on a white keyboard placed on a wooden table. "
        "The scene is set in a room with a wooden floor and a part of a white blanket visible in "
        "the background. The person's hands are actively moving across the keyboard, indicating "
        "typing activity. The ambient sound is the distinct sound of keys being pressed.",
    ],
    [
        "examples/case6.jpg",
        "A man is playing an acoustic guitar in a modern kitchen setting. He is focused on his "
        "playing, with his hands moving along the strings and fretboard. The background features "
        "a well-lit kitchen with wooden cabinets and hanging lights.",
    ],
]

CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="DreamX-Creator 1.0") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🎬 DreamX-Creator 1.0\n"
            "Turn **one image + one prompt** into a video whose **soundtrack is generated "
            "jointly with the frames** — speech, foley and ambience come out of the same "
            "denoiser as the picture, so they stay in sync.\n\n"
            "Describe the *sound* as well as the action. For speech, quote the line: "
            "`Man says, 'hello there.'`\n\n"
            "[Model](https://huggingface.co/GD-ML/DreamX-Creator) · "
            "[Code](https://github.com/AMAP-ML/DreamX-Creator)"
        )

        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(label="First frame (optional — leave empty for a neutral keyframe)", type="filepath", height=320)
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the motion and the sound you want to hear…",
                    lines=4,
                )
                run = gr.Button("Generate audio + video", variant="primary")
            with gr.Column(scale=1):
                video_out = gr.Video(label="Result (video + generated audio)", height=380)
                last_frame_out = gr.Image(label="Last frame (feeds the next scene)", height=380)

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                seconds = gr.Slider(
                    label="Duration (seconds)", minimum=2.0, maximum=5.0, step=0.5, value=3.0
                )
                num_inference_steps = gr.Slider(
                    label="Denoising steps", minimum=10, maximum=50, step=1, value=30
                )
            with gr.Row():
                resolution_tokens = gr.Dropdown(
                    label="Resolution budget (spatial tokens)",
                    choices=RESOLUTION_CHOICES,
                    value=440,
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0
                )
            negative_prompt = gr.Textbox(
                label="Negative prompt", value=DEFAULT_NEGATIVE_PROMPT, lines=2
            )
            with gr.Row():
                seed = gr.Number(label="Seed", value=113, precision=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=False)

        gr.Markdown(
            "Longer clips, more steps and a bigger resolution budget all cost GPU time "
            "roughly linearly (and attention grows quadratically with tokens × frames)."
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[image, prompt],
            outputs=[video_out, last_frame_out, seed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label=f"Official Verse-Bench cases from the DreamX-Creator repo (v{CACHE_VERSION})",
        )

    run.click(
        fn=generate,
        inputs=[
            image,
            prompt,
            seconds,
            num_inference_steps,
            resolution_tokens,
            guidance_scale,
            negative_prompt,
            seed,
            randomize_seed,
        ],
        outputs=[video_out, last_frame_out, seed],
        api_name="generate",
    )

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