"""JoyAI-Echo — multi-shot text-to-video **with synchronized audio**, on ZeroGPU.

A faithful port of the authors' reference entrypoint
``ltx-distillation/scripts/multishot_inference_dmd.py``
(https://github.com/jd-opensource/JoyAI-Echo) to a Gradio Space.  Every
hyper-parameter below is copied verbatim from the authors'
``configs/inference.yaml``; the pipeline objects, the sigma schedule, the
paired cross-modal memory bank and the decode path are the upstream ones,
imported from the repo's own ``ltx-core`` / ``ltx-pipelines`` /
``ltx-distillation`` packages (vendored under ``vendor/``).

Pipeline (identical structure to the reference):

* **Stage 1** — Gemma-3-12B text encoder + the checkpoint's embeddings
  processor turn each shot prompt into paired video/audio context.
* **Stage 2** — the DMD-distilled LTX-2.3-derived audio-video DiT plus the
  video / audio VAEs.
* **Per shot** — 8-step DMD denoising on the authors' exact sigma schedule.
  Shot 1 runs through ``BidirectionalAVInferencePipeline``; every later shot
  runs through ``BidirectionalMemoryAVInferencePipeline`` conditioned on the
  *paired audio-video memory bank* (video clips + audio latents carried over
  from the previous shots) — that memory bank is what gives JoyAI-Echo its
  story-level face *and* voice consistency.
* Decode video + audio, write one mp4 per shot, concatenate.

Single deviation, forced by hardware: the transformer's linear weights are
loaded with ltx-core's own FP8 cast policy (``QuantizationPolicy.fp8_cast()``
— fp8_e4m3 storage, upcast to bf16 inside each linear forward).  That takes
the 21B audio-video DiT from ~42 GB to ~22 GB so that it, Gemma and the VAEs
can all stay resident on one ZeroGPU xlarge slot, which in turn removes the
reference's per-phase generator<->CPU hot-swap.  Sampling maths, schedule,
memory bank and VAE decode are untouched.
"""

from __future__ import annotations

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")

import sys
from pathlib import Path

_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE / "vendor"))

import gc
import logging
import tempfile
import time
from typing import Any

import spaces  # noqa: E402  — must precede torch / any CUDA-touching import
import gradio as gr
import torch
from huggingface_hub import hf_hub_download, snapshot_download

import examples as EX

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("joyai-echo")

# -----------------------------------------------------------------------------
# Constants — verbatim from configs/inference.yaml
# -----------------------------------------------------------------------------
CKPT_REPO = "jdopensource/JoyAI-Echo"
CKPT_FILE = "JoyAI-Echo-release.safetensors"
GEMMA_REPO = "google/gemma-3-12b-it"

# denoising.sigmas — the 8-step DMD schedule
DENOISING_SIGMAS = [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0]

VIDEO_FPS = 25          # video.fps
DTYPE = torch.bfloat16  # inference.dtype

# memory:
MEMORY_MAX_SIZE = 7
NUM_FIX_FRAMES = 3
MEMORY_DOWNSCALE_FACTOR = 1
MEMORY_POSITION_MODE = "reference"
SAVE_MODE = "random_every_shot_frame"
VIDEO_MEMORY_FRAME_SELECTION_MODE = "center"
VIDEO_MEMORY_CLIP_NUM_FRAMES = 9
# audio_memory:
ENABLE_AUDIO_MEMORY = True
AUDIO_MEMORY_WINDOW_SIZE = 96
AUDIO_MEMORY_WINDOW_SELECTION_MODE = "max_response"
AUDIO_MEMORY_SAMPLE_RATE = 16000
AUDIO_MEMORY_MEL_BINS = 128
AUDIO_MEMORY_MEL_HOP_LENGTH = 160
AUDIO_MEMORY_N_FFT = 1024
AUDIO_MEMORY_DOWNSAMPLE_FACTOR = 4
AUDIO_MEMORY_IS_CAUSAL = True
# inference:
V2A_GRAD_SCALE = 2.0

RESOLUTIONS: dict[str, tuple[int, int]] = {
    "1280 x 736 — native": (736, 1280),
    "832 x 480 — fast": (480, 832),
}
DEFAULT_RESOLUTION = "1280 x 736 — native"
DEFAULT_FRAMES = 121  # 4.8 s per shot at 25 fps (video.num_frames is 241 upstream)
DEFAULT_SEED = 12345  # video.seed

MAX_SHOTS = 3

# -----------------------------------------------------------------------------
# Weights
# -----------------------------------------------------------------------------
logger.info("Downloading JoyAI-Echo checkpoint (~46 GB) …")
CHECKPOINT_PATH = hf_hub_download(CKPT_REPO, CKPT_FILE)
logger.info("Downloading Gemma-3-12B-it text encoder (~24 GB) …")
GEMMA_ROOT = snapshot_download(
    GEMMA_REPO,
    allow_patterns=["*.safetensors", "*.json", "tokenizer.model", "tokenizer.json"],
)
logger.info("Weights ready:\n  %s\n  %s", CHECKPOINT_PATH, GEMMA_ROOT)

# -----------------------------------------------------------------------------
# Models — built once, at module scope, straight onto "cuda" so the ZeroGPU
# hijack packs them and streams them into VRAM on the first request.
# -----------------------------------------------------------------------------
# torchvision 0.26 (shipped with torch 2.11) deleted `torchvision.io.write_video`,
# which the vendored `ltx_distillation.utils` imports at module load and uses to
# mux the generated soundtrack into the mp4. Put a PyAV implementation of the old
# API back before that import, so the vendor tree stays byte-identical upstream.
import tv_compat  # noqa: E402

tv_compat.install()

# `ltx_core.loader` must be imported before `ltx_core.quantization`: loader's
# __init__ pulls in fuse_loras, which imports back into quantization.fp8_cast,
# so entering the package through quantization first is a circular import.
# The authors' own entrypoint imports ltx_core.loader first for this reason.
import ltx_core.loader  # noqa: E402,F401
from ltx_core.quantization import QuantizationPolicy  # noqa: E402
from ltx_distillation.inference.bidirectional_pipeline import (  # noqa: E402
    BidirectionalAVInferencePipeline,
)
from ltx_distillation.inference.memory_bidirectional_pipeline import (  # noqa: E402
    BidirectionalMemoryAVInferencePipeline,
)
from ltx_distillation.inference.memory_multishot import (  # noqa: E402
    PairedAudioVideoMemoryBank,
    build_paired_audio_memory_kwargs,
    video_uint8_to_pil_frames,
)
from ltx_distillation.models.ltx_wrapper import LTX2DiffusionWrapper  # noqa: E402
from ltx_distillation.models.text_encoder_wrapper import GemmaTextEncoderWrapper  # noqa: E402
from ltx_distillation.models.vae_wrapper import AudioVAEWrapper, VideoVAEWrapper  # noqa: E402
from ltx_distillation.utils import (  # noqa: E402
    add_noise,
    compute_latent_shapes,
    concat_shot_videos,
    decode_benchmark_sample,
    encode_memory_frames_batch,
    write_benchmark_media,
)
from ltx_pipelines.utils.model_ledger import ModelLedger  # noqa: E402

DEVICE = torch.device("cuda")
CPU = torch.device("cpu")

# The ledger is built on CPU and every module is moved to "cuda" afterwards.
# That is the authors' own pattern — `create_ltx2_wrapper` / `create_vae_wrappers`
# / `create_text_encoder_wrapper` all construct a `device=torch.device("cpu")`
# ledger and then `.to(device)`, because `ModelLedger._target_device()` is handed
# straight to `safetensors.safe_open(..., device=...)`. Loading weights inside
# safetensors bypasses ZeroGPU's `torch.cuda` interception, so asking it for
# "cuda" in the main process dies with "No CUDA GPUs are available". Moving the
# assembled `nn.Module`s instead goes through the hijacked `Module.to`, which is
# exactly what ZeroGPU wants: weights get packed and streamed into VRAM on the
# first `@spaces.GPU` call.
logger.info("Building model ledger (weights load on CPU, then move to cuda) …")
_ledger = ModelLedger(
    dtype=DTYPE,
    device=CPU,
    checkpoint_path=CHECKPOINT_PATH,
    gemma_root_path=GEMMA_ROOT,
    quantization=QuantizationPolicy.fp8_cast(),  # transformer only
)

logger.info("Loading audio-video DiT (fp8 cast) …")
# NOTE: no `dtype=` here, unlike the authors' unquantized `create_ltx2_wrapper`.
# `QuantizationPolicy.fp8_cast()` has already stored the transformer-block linear
# weights as float8_e4m3fn; passing dtype=bfloat16 would upcast them straight back
# and undo the memory saving that lets this model fit at all.
_transformer = _ledger.transformer().to(device=DEVICE)

logger.info("Loading video VAE …")
_video_vae = VideoVAEWrapper(
    encoder=_ledger.video_encoder().to(device=DEVICE, dtype=DTYPE),
    decoder=_ledger.video_decoder().to(device=DEVICE, dtype=DTYPE),
    device=DEVICE,
    dtype=DTYPE,
).eval()

logger.info("Loading audio VAE + vocoder …")
_audio_vae = AudioVAEWrapper(
    # float32 for the audio encoder, matching `create_vae_wrappers`.
    encoder=_ledger.audio_encoder().to(device=DEVICE, dtype=torch.float32),
    decoder=_ledger.audio_decoder().to(device=DEVICE, dtype=DTYPE),
    vocoder=_ledger.vocoder().to(device=DEVICE, dtype=DTYPE),
    device=DEVICE,
    dtype=DTYPE,
).eval()
AUDIO_SR = _audio_vae.get_output_sample_rate() or 24000

logger.info("Loading Gemma-3-12B text encoder + embeddings processor …")
_text_encoder = GemmaTextEncoderWrapper(
    text_encoder=_ledger.text_encoder().to(device=DEVICE, dtype=DTYPE),
    embeddings_processor=_ledger.gemma_embeddings_processor().to(device=DEVICE, dtype=DTYPE),
    device=DEVICE,
    dtype=DTYPE,
).eval()

gc.collect()
logger.info("All models resident. audio_sample_rate=%d", AUDIO_SR)

_SIGMAS_CPU = torch.tensor(DENOISING_SIGMAS, dtype=torch.float32)


# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------
def _free() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _vram(tag: str) -> None:
    if torch.cuda.is_available():
        free, total = torch.cuda.mem_get_info()
        logger.info(
            "[vram] %-26s alloc=%5.1fGB peak=%5.1fGB free=%5.1fGB/%5.1fGB",
            tag,
            torch.cuda.memory_allocated() / 2**30,
            torch.cuda.max_memory_allocated() / 2**30,
            free / 2**30,
            total / 2**30,
        )


def _collect_shots(*shots: Any) -> list[str]:
    return [s.strip() for s in shots if isinstance(s, str) and s.strip()]


def _round_frames(num_frames: int) -> int:
    """Snap to the VAE's 1 + 8k temporal grid."""
    num_frames = max(41, int(num_frames))
    return ((num_frames - 1) // 8) * 8 + 1


def _estimate_duration(*args: Any, **kwargs: Any) -> int:
    """ZeroGPU reservation: cold weight streaming + per-shot sampling/decoding.

    Calibrated against measured runs on this Space rather than guessed, so the
    reservation stays lean and visitors don't burn quota they never use:

      * ~7 s to stream the 52 GB pack into VRAM on a cold worker, ~2 s to encode
        the prompts with Gemma, ~1 s for the final ffmpeg concat  ->  BASE.
      * 21.3 s and 22.1 s for the two shots of a 1280x736 / 121-frame run
        (denoise ~18 s + VAE decode ~2.3 s)  ->  25 s per shot with ~15 % margin.

    Both terms scale linearly with frame count; 832x480 is ~2.4x fewer latent
    tokens, so it gets a proportionally smaller per-shot cost.
    """
    n_shots = max(1, len(_collect_shots(*args[:MAX_SHOTS])))
    resolution = args[MAX_SHOTS] if len(args) > MAX_SHOTS else DEFAULT_RESOLUTION
    frames = args[MAX_SHOTS + 1] if len(args) > MAX_SHOTS + 1 else DEFAULT_FRAMES
    try:
        frames = _round_frames(int(frames))
    except (TypeError, ValueError):
        frames = DEFAULT_FRAMES
    height, _ = RESOLUTIONS.get(str(resolution), RESOLUTIONS[DEFAULT_RESOLUTION])
    base = 18.0
    per_shot = (25.0 if height > 480 else 12.0) * (frames / 121.0)
    return int(min(600, base + n_shots * per_shot))


# -----------------------------------------------------------------------------
# Generation — ported 1:1 from multishot_inference_dmd.py::main
# -----------------------------------------------------------------------------
@spaces.GPU(duration=_estimate_duration, size="xlarge")
@torch.inference_mode()
def generate(
    shot_1: str,
    shot_2: str = "",
    shot_3: str = "",
    resolution: str = DEFAULT_RESOLUTION,
    num_frames: int = DEFAULT_FRAMES,
    seed: int = DEFAULT_SEED,
    progress=gr.Progress(track_tqdm=True),
) -> str:
    """Generate a multi-shot video with a synchronized soundtrack from shot prompts.

    Args:
        shot_1: Prompt for the first shot. Describe, in order: who (appearance
            *and* voice timbre), the action plus the spoken line, visual style,
            camera, background, and the sound design.
        shot_2: Optional second shot. It is conditioned on shot 1 through the
            paired audio-video memory bank, so reuse the same character
            description (e.g. `ID_A is ...`) to keep face and voice.
        shot_3: Optional third shot, conditioned on shots 1-2.
        resolution: Output resolution; 1280 x 736 is the training resolution.
        num_frames: Frames per shot at 25 fps, snapped to 1 + 8k.
        seed: Base random seed; shot *i* uses `seed + i`.

    Returns:
        Path to an mp4 holding every generated shot, with audio.
    """
    prompts = _collect_shots(shot_1, shot_2, shot_3)
    if not prompts:
        raise gr.Error("Please write at least one shot prompt.")

    height, width = RESOLUTIONS.get(str(resolution), RESOLUTIONS[DEFAULT_RESOLUTION])
    num_frames = _round_frames(num_frames)
    seed = int(seed)

    torch.cuda.reset_peak_memory_stats()
    _vram("startup")
    t0 = time.perf_counter()

    # ---- Stage 1: encode every prompt once, cache on CPU -----------------
    cached_conds: list[dict[str, Any]] = []
    for prompt in prompts:
        cond = _text_encoder([prompt])
        cached_conds.append(
            {k: (v.detach().cpu() if isinstance(v, torch.Tensor) else v) for k, v in cond.items()}
        )
        del cond
    _free()
    t_encode = time.perf_counter()
    logger.info("[stage 1] encoded %d prompt(s) in %.1fs", len(prompts), t_encode - t0)

    # ---- Stage 2: pipelines ---------------------------------------------
    # LTX2DiffusionWrapper only stores the resolution-dependent RoPE geometry;
    # it wraps (does not copy) the resident transformer.
    generator = LTX2DiffusionWrapper(
        model=_transformer, video_height=height, video_width=width
    ).eval()
    sigmas = _SIGMAS_CPU.to(DEVICE)
    base_pipeline = BidirectionalAVInferencePipeline(
        generator=generator, add_noise_fn=add_noise, denoising_sigmas=sigmas
    )
    memory_pipeline = BidirectionalMemoryAVInferencePipeline(
        generator=generator,
        add_noise_fn=add_noise,
        denoising_sigmas=sigmas,
        memory_downscale_factor=MEMORY_DOWNSCALE_FACTOR,
    )

    video_shape, audio_shape = compute_latent_shapes(
        num_frames=num_frames,
        video_height=height,
        video_width=width,
        batch_size=1,
        video_fps=float(VIDEO_FPS),
    )
    memory_bank = PairedAudioVideoMemoryBank(
        max_size=MEMORY_MAX_SIZE, save_mode=SAVE_MODE, num_fix_frames=NUM_FIX_FRAMES
    )

    out_dir = Path(tempfile.mkdtemp(prefix="joyai-echo-"))
    shot_paths: list[Path] = []

    # ---- Per-shot sampling ----------------------------------------------
    for shot_idx, prompt in enumerate(prompts):
        shot_t0 = time.perf_counter()
        conditional_dict = {
            k: (v.to(DEVICE) if isinstance(v, torch.Tensor) else v)
            for k, v in cached_conds[shot_idx].items()
        }
        shot_seed = seed + shot_idx
        logger.info(
            "[shot %d/%d] memory=%d seed=%d %dx%d frames=%d",
            shot_idx + 1, len(prompts), len(memory_bank), shot_seed, width, height, num_frames,
        )

        with torch.random.fork_rng(devices=[DEVICE]):
            torch.manual_seed(shot_seed)
            torch.cuda.manual_seed(shot_seed)

            if len(memory_bank) > 0:
                memory_video = encode_memory_frames_batch(
                    video_vae=_video_vae,
                    batch_memory_frames=[memory_bank.get_memory_frames()],
                    target_h=height,
                    target_w=width,
                    device=DEVICE,
                    dtype=DTYPE,
                )
                memory_audio_kwargs = build_paired_audio_memory_kwargs(
                    memory_bank,
                    enable_audio_memory=ENABLE_AUDIO_MEMORY,
                    v2a_grad_scale=V2A_GRAD_SCALE,
                    memory_position_mode=MEMORY_POSITION_MODE,
                )
                video_latent, audio_latent = memory_pipeline.generate(
                    video_shape=tuple(video_shape),
                    audio_shape=tuple(audio_shape),
                    conditional_dict=conditional_dict,
                    memory_video=memory_video,
                    seed=shot_seed,
                    **memory_audio_kwargs,
                )
                del memory_video, memory_audio_kwargs
            else:
                video_latent, audio_latent = base_pipeline.generate(
                    video_shape=tuple(video_shape),
                    audio_shape=tuple(audio_shape),
                    conditional_dict=conditional_dict,
                    seed=shot_seed,
                )
        torch.cuda.synchronize()
        t_denoise = time.perf_counter()
        del conditional_dict
        _free()
        _vram(f"shot {shot_idx + 1} denoised")

        audio_memory_latent = (
            audio_latent.detach().cpu().contiguous()
            if (ENABLE_AUDIO_MEMORY and audio_latent is not None)
            else None
        )
        video_uint8, audio_waveform = decode_benchmark_sample(
            _video_vae, _audio_vae, video_latent, audio_latent
        )
        torch.cuda.synchronize()
        t_decode = time.perf_counter()
        _vram(f"shot {shot_idx + 1} decoded")

        # Only pay for memory-bank bookkeeping when another shot follows.
        if audio_memory_latent is not None and shot_idx + 1 < len(prompts):
            memory_bank.save_memory_slot(
                video_uint8_to_pil_frames(video_uint8),
                audio_memory_latent,
                audio_window_size=AUDIO_MEMORY_WINDOW_SIZE,
                video_clip_num_frames=VIDEO_MEMORY_CLIP_NUM_FRAMES,
                audio_waveform=audio_waveform,
                audio_sample_rate=AUDIO_MEMORY_SAMPLE_RATE,
                video_fps=float(VIDEO_FPS),
                audio_window_selection_mode=AUDIO_MEMORY_WINDOW_SELECTION_MODE,
                video_frame_selection_mode=VIDEO_MEMORY_FRAME_SELECTION_MODE,
                audio_memory_mel_bins=AUDIO_MEMORY_MEL_BINS,
                audio_memory_mel_hop_length=AUDIO_MEMORY_MEL_HOP_LENGTH,
                audio_memory_n_fft=AUDIO_MEMORY_N_FFT,
                audio_memory_downsample_factor=AUDIO_MEMORY_DOWNSAMPLE_FACTOR,
                audio_memory_is_causal=AUDIO_MEMORY_IS_CAUSAL,
            )

        shot_path = out_dir / f"shot_{shot_idx:03d}.mp4"
        write_benchmark_media(
            output_path=shot_path,
            video_uint8=video_uint8,
            audio_waveform=audio_waveform,
            fps=VIDEO_FPS,
            audio_sr=int(AUDIO_SR),
        )
        shot_paths.append(shot_path)

        logger.info(
            "[shot %d/%d] denoise=%.1fs decode=%.1fs shot_total=%.1fs",
            shot_idx + 1, len(prompts),
            t_denoise - shot_t0, t_decode - t_denoise, time.perf_counter() - shot_t0,
        )
        del video_latent, audio_latent, video_uint8, audio_waveform, audio_memory_latent
        _free()

    del generator, base_pipeline, memory_pipeline
    _free()

    if len(shot_paths) == 1:
        final_path = shot_paths[0]
    else:
        final_path = out_dir / "combined_shots.mp4"
        concat_shot_videos(shot_paths, final_path)

    logger.info(
        "[done] %d shot(s) %dx%d x%d frames in %.1fs total",
        len(prompts), width, height, num_frames, time.perf_counter() - t0,
    )
    _vram("done")
    return str(final_path)


# -----------------------------------------------------------------------------
# UI
# -----------------------------------------------------------------------------
CSS = """
.fillable { max-width: 1280px !important; }
#shots textarea { font-size: 13px; line-height: 1.45; }
"""

PLACEHOLDER_1 = (
    "ID_A is a young woman in her twenties with shoulder-length dark brown hair … "
    "ID_A's voice is a clear young female voice with a soft mid-high register … "
    "In a soft young female voice, ID_A says, “…” … A stable medium shot frames ID_A "
    "from the waist up … Very soft indoor room tone. No prominent background music."
)

with gr.Blocks(css=CSS, theme=gr.themes.Citrus(), title="JoyAI-Echo") as demo:
    gr.Markdown(
        """
        # 🎬 JoyAI-Echo — multi-shot video **with audio**

        [JoyAI-Echo](https://huggingface.co/jdopensource/JoyAI-Echo) (JD Joy Future Academy) is a
        DMD-distilled, LTX-2.3-derived **joint audio-video** generator: a single 8-step pass
        produces the picture *and* its synchronized soundtrack — dialogue, lip-sync, foley and
        ambience.

        Write one prompt per **shot**. Every shot after the first is conditioned on a **paired
        cross-modal memory bank** (video clips *and* audio latents carried over from the earlier
        shots), which is what keeps a character's face **and** voice timbre stable across a story.
        Reuse the same character description — the authors write `ID_A is …` — in each shot.

        A good shot prompt covers, in order: **who** (appearance + voice timbre) · **action &
        the spoken line** · **visual style** · **camera** · **background** · **sound design**.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            shot_1 = gr.Textbox(
                label="Shot 1", elem_id="shots", lines=8, placeholder=PLACEHOLDER_1
            )
            shot_2 = gr.Textbox(
                label="Shot 2 — optional, continues the story through the memory bank",
                elem_id="shots",
                lines=5,
                placeholder="Leave empty for a single-shot video.",
            )
            shot_3 = gr.Textbox(
                label="Shot 3 — optional",
                elem_id="shots",
                lines=4,
                placeholder="Leave empty.",
            )
            run = gr.Button("Generate", variant="primary", size="lg")

            with gr.Accordion("Advanced options", open=False):
                resolution = gr.Radio(
                    choices=list(RESOLUTIONS),
                    value=DEFAULT_RESOLUTION,
                    label="Resolution",
                    info="1280 x 736 is the model's training resolution; 832 x 480 is ~2x faster.",
                )
                num_frames = gr.Slider(
                    41, 201, value=DEFAULT_FRAMES, step=8,
                    label="Frames per shot (25 fps)",
                    info="121 frames ≈ 4.8 s of video and audio per shot.",
                )
                seed = gr.Number(value=DEFAULT_SEED, precision=0, label="Seed")

        with gr.Column(scale=2):
            video_out = gr.Video(label="Generated video (with audio)", autoplay=True)
            gr.Markdown(
                "⏱️ A cold run streams ~53 GB of weights onto the GPU before sampling, so the "
                "first generation of a session takes noticeably longer than the ones after it. "
                "Each extra shot adds a full 8-step pass plus a VAE decode."
            )

    run.click(
        fn=generate,
        inputs=[shot_1, shot_2, shot_3, resolution, num_frames, seed],
        outputs=video_out,
    )

    gr.Examples(
        examples=[
            [EX.VLOG_1, EX.VLOG_2],
            [EX.RALLY_1, EX.RALLY_2],
            [EX.WAREHOUSE_1, EX.WAREHOUSE_2],
        ],
        inputs=[shot_1, shot_2],
        outputs=video_out,
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
        label="Example stories — the authors' own prompts (prompts/test_*.json)",
        example_labels=[
            "Indoor vlog — soft female voice, two reflective beats",
            "Rally driver — low male baritone at the mountain start line",
            "Warehouse worker — plain low voice, boots and forklift foley",
        ],
    )

    gr.Markdown(
        "Model: [`jdopensource/JoyAI-Echo`](https://huggingface.co/jdopensource/JoyAI-Echo) · "
        "code: [jd-opensource/JoyAI-Echo](https://github.com/jd-opensource/JoyAI-Echo) · "
        "text encoder: [`google/gemma-3-12b-it`](https://huggingface.co/google/gemma-3-12b-it). "
        "Released under the LTX-2 Community License — **non-commercial use**, and generated "
        "content must be disclosed as machine-generated."
    )


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