"""NVIDIA CMD (Context-Matched Distillation) image-to-video demo.

This app follows the reference inference path of https://github.com/nv-tlabs/cmd
(`inference.py`, `examples/run_examples.sh chunk1-short`) using the released
`chunk1_short_t24_l21.safetensors` checkpoint from https://huggingface.co/nvidia/cmd.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import random
import tempfile
import time
from pathlib import Path
from typing import Optional

import spaces  # noqa: E402  (must be imported before torch)
import torch  # noqa: E402

import gradio as gr  # noqa: E402
import imageio  # noqa: E402
import numpy as np  # noqa: E402
from einops import rearrange  # noqa: E402
from omegaconf import OmegaConf  # noqa: E402
from PIL import Image  # noqa: E402

from pipeline import CausalInferencePipeline  # noqa: E402
from utils.misc import set_seed  # noqa: E402

# --------------------------------------------------------------------------------------
# Released variant: "chunk1-short" from examples/run_examples.sh
# --------------------------------------------------------------------------------------
MODEL_REPO = "nvidia/cmd"
CHECKPOINT_FILE = "chunk1_short_t24_l21.safetensors"
CONFIG_PATH = "configs/cosmos/t24_l21_student_context_distillation.yaml"
DEFAULT_CONFIG_PATH = "configs/cosmos/default_config.yaml"

MAX_LATENT_FRAMES = 24  # t24
MIN_LATENT_FRAMES = 6
NUM_FRAME_PER_BLOCK = 1  # chunk1
LOCAL_ATTN_SIZE = 21  # l21
FPS = 16
HEIGHT, WIDTH = 480, 832
DEFAULT_SEED = 22  # SEED default in examples/run_examples.sh

# The Wan2.1 16-channel VAE that Cosmos-Predict2.5 ships as `tokenizer.pth`.
# Sourced from the ungated Apache-2.0 Wan2.1 release instead; verified to load
# into the vendored `_video_vae` with an exact state-dict match.
VAE_REPO = "Wan-AI/Wan2.1-T2V-1.3B"
VAE_CHECKPOINT_FILE = "Wan2.1_VAE.pth"

torch.set_grad_enabled(False)


def _pixel_frames(latent_frames: int) -> int:
    """Wan2.1 VAE temporal layout: 1 + 4*(n-1) pixel frames per n latent frames."""
    return 1 + (int(latent_frames) - 1) * 4


print("Building the CMD chunk1-short pipeline...", flush=True)
_config = OmegaConf.merge(
    OmegaConf.load(DEFAULT_CONFIG_PATH), OmegaConf.load(CONFIG_PATH)
)
_config.num_frame_per_block = NUM_FRAME_PER_BLOCK
_config.model_kwargs.local_attn_size = LOCAL_ATTN_SIZE
# Build the DiT straight from the released CMD student export rather than
# layering it over the gated Cosmos-Predict2.5-2B base checkpoint.
_config.model_kwargs.model_name = MODEL_REPO
_config.model_kwargs.checkpoint_filename = CHECKPOINT_FILE
_config.vae_model_name = VAE_REPO
_config.vae_checkpoint_filename = VAE_CHECKPOINT_FILE

pipeline = CausalInferencePipeline(_config, device=torch.device("cuda"))
pipeline = pipeline.to(dtype=torch.bfloat16)
pipeline.text_encoder.to("cuda")
pipeline.generator.to("cuda")
pipeline.vae.to("cuda")
print("Pipeline ready.", flush=True)


def _preprocess(image: Image.Image) -> torch.Tensor:
    """Aspect-preserving centre crop to 832x480, then ToTensor + Normalize([0.5],[0.5])."""
    image = image.convert("RGB")
    width, height = image.size
    target = WIDTH / HEIGHT
    if width / height > target:
        crop_w = int(round(height * target))
        left = (width - crop_w) // 2
        image = image.crop((left, 0, left + crop_w, height))
    elif width / height < target:
        crop_h = int(round(width / target))
        top = (height - crop_h) // 2
        image = image.crop((0, top, width, top + crop_h))
    image = image.resize((WIDTH, HEIGHT), Image.LANCZOS)
    array = np.asarray(image, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(array).permute(2, 0, 1)  # [3, H, W]
    return (tensor - 0.5) / 0.5


def _write_mp4(frames: np.ndarray) -> str:
    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    with imageio.get_writer(
        path,
        format="FFMPEG",
        mode="I",
        fps=FPS,
        codec="libx264",
        # pixelformat (not output_params) so ffmpeg receives a single -pix_fmt.
        pixelformat="yuv420p",
        output_params=["-crf", "17", "-movflags", "+faststart"],
    ) as writer:
        for frame in frames:
            writer.append_data(frame)
    return path


def _estimate_duration(
    image=None,
    prompt: str = "",
    num_latent_frames: int = MAX_LATENT_FRAMES,
    seed: int = DEFAULT_SEED,
    randomize_seed: bool = False,
    *args,
    **kwargs,
) -> int:
    # Measured on this Space's ZeroGPU hardware (chunk1-short): 3.9s at t6,
    # 9.7s at t12, 16.3s at t18 and 23.7s at t24, plus ~1s of H.264 encoding.
    # Cost is linear in the latent-frame count; this keeps ~15% headroom over
    # the measured worst case and a small floor for the shortest clips, so the
    # request stays lean on every visitor's ZeroGPU quota.
    frames = int(num_latent_frames or MAX_LATENT_FRAMES)
    return max(15, min(60, int(round(1.36 * frames - 3.0))))


@spaces.GPU(duration=_estimate_duration)
def generate(
    image: Optional[Image.Image],
    prompt: str,
    num_latent_frames: int = MAX_LATENT_FRAMES,
    seed: int = DEFAULT_SEED,
    randomize_seed: bool = False,
) -> tuple:
    """Animate a still image into a short video with NVIDIA CMD.

    Args:
        image: the first frame of the video (centre-cropped to 832x480).
        prompt: a description of the motion and scene to generate.
        num_latent_frames: video length in latent frames; n latents decode to 1+4*(n-1) frames at 16 fps.
        seed: RNG seed for reproducible sampling.
        randomize_seed: draw a fresh random seed instead of using `seed`.

    Returns:
        The generated mp4 path, a run-info string, and the seed that was used.
    """
    if image is None:
        raise gr.Error("Please provide an input image to animate.")
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please provide a text prompt describing the motion.")

    num_latent_frames = int(num_latent_frames)
    if not MIN_LATENT_FRAMES <= num_latent_frames <= MAX_LATENT_FRAMES:
        raise gr.Error(
            f"Video length must be between {MIN_LATENT_FRAMES} and {MAX_LATENT_FRAMES} latent frames."
        )

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

    started = time.perf_counter()
    with torch.no_grad():
        first_frame = (
            _preprocess(image)
            .unsqueeze(0)
            .unsqueeze(2)
            .to(device="cuda", dtype=torch.bfloat16)
        )  # [1, 3, 1, H, W]
        initial_latent = pipeline.vae.encode_to_latent(first_frame).to(
            device="cuda", dtype=torch.bfloat16
        )
        noise = torch.randn(
            [1, num_latent_frames - 1, *_config.image_or_video_shape[2:]],
            device="cuda",
            dtype=torch.bfloat16,
        )
        video, latents = pipeline.inference(
            noise=noise,
            text_prompts=[prompt],
            initial_latent=initial_latent,
            return_latents=True,
        )
        frames = rearrange(video, "b t c h w -> b t h w c")[0].float().cpu()
        frames = (255.0 * frames).round().clamp(0, 255).to(torch.uint8).numpy()
        pipeline.vae.model.clear_cache()
    elapsed = time.perf_counter() - started

    path = _write_mp4(frames)
    info = (
        f"{frames.shape[0]} frames ({frames.shape[0] / FPS:.1f}s) at {WIDTH}x{HEIGHT}, "
        f"{latents.shape[1]} latent frames · seed {seed} · {elapsed:.1f}s on GPU"
    )
    print(info, flush=True)
    return path, info, seed


EXAMPLES_DIR = Path("examples")


def _example(name: str) -> list:
    return [
        str(EXAMPLES_DIR / f"{name}.jpg"),
        (EXAMPLES_DIR / f"{name}.txt").read_text(encoding="utf-8").strip(),
    ]


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

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # NVIDIA CMD — image to video

            Autoregressive image-to-video with
            [**nvidia/cmd**](https://huggingface.co/nvidia/cmd) (*Context-Matched
            Distillation*): a 4-step causal student distilled from
            [Cosmos-Predict2.5-2B](https://huggingface.co/nvidia/Cosmos-Predict2.5-2B),
            generating one latent frame at a time with a rolling KV cache.

            Give it a first frame and a prompt describing the motion. Released
            `chunk1-short` checkpoint · 832×480 · 16 fps · up to 93 frames.

            [Code](https://github.com/nv-tlabs/cmd) ·
            [Model card](https://huggingface.co/nvidia/cmd) ·
            Non-commercial use only (NVIDIA OneWay Noncommercial License).
            """
        )
        with gr.Row():
            with gr.Column():
                image_in = gr.Image(label="First frame", type="pil", height=300)
                prompt_in = gr.Textbox(
                    label="Prompt",
                    lines=4,
                    placeholder="Describe the scene and how it should move…",
                )
                run_btn = gr.Button("Generate video", variant="primary")
            with gr.Column():
                video_out = gr.Video(label="Generated video", autoplay=True, height=300)
                info_out = gr.Markdown()

        with gr.Accordion("Advanced settings", open=False):
            length_in = gr.Slider(
                label="Video length (latent frames)",
                minimum=MIN_LATENT_FRAMES,
                maximum=MAX_LATENT_FRAMES,
                step=1,
                value=MAX_LATENT_FRAMES,
                info="n latent frames decode to 1 + 4·(n−1) video frames at 16 fps "
                "(24 → 93 frames ≈ 5.8 s). Shorter is faster.",
            )
            with gr.Row():
                seed_in = gr.Number(label="Seed", value=DEFAULT_SEED, precision=0)
                randomize_in = gr.Checkbox(label="Randomize seed", value=False)

        gr.Examples(
            examples=[_example("bus_terminal"), _example("robot_welding")],
            inputs=[image_in, prompt_in],
            outputs=[video_out, info_out, seed_in],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    gr.on(
        triggers=[run_btn.click, prompt_in.submit],
        fn=generate,
        inputs=[image_in, prompt_in, length_in, seed_in, randomize_in],
        outputs=[video_out, info_out, seed_in],
        api_name="generate",
    )

demo.launch(mcp_server=True)