from __future__ import annotations

import copy
import gc
import os
import random
import sys
import tempfile
import time
import uuid
import warnings
import zipfile
from pathlib import Path
from typing import Any

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

import cv2
import gradio as gr
import numpy as np
import spaces
import torch
import torch._dynamo
from huggingface_hub import hf_hub_download
from PIL import Image, ImageOps
from torch.nn import functional as F
from torchao.quantization import (
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig,
    quantize_,
)
from tqdm import tqdm

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepInverseScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    FlowMatchEulerDiscreteScheduler,
    SASolverScheduler,
    UniPCMultistepScheduler,
)
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline
from diffusers.utils.export_utils import export_to_video

warnings.filterwarnings("ignore")

# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------

DEFAULT_MODEL_ID = (
    "TestOrganizationPleaseIgnore/"
    "WAMU-Merge-MotionBoost_WAN2.2_I2V_LIGHTNING"
)
MODEL_ID = os.getenv("REPO_ID", DEFAULT_MODEL_ID).strip() or DEFAULT_MODEL_ID
AOTI_REPO = os.getenv(
    "AOTI_REPO", "cbensimon/WanTransformer3DModel-sm120-cu130-raa"
).strip()
ENABLE_AOTI = os.getenv("ENABLE_AOTI", "true").lower() not in {"0", "false", "no"}
ENABLE_VAE_TILING = os.getenv("ENABLE_VAE_TILING", "false").lower() in {
    "1",
    "true",
    "yes",
}
RUNPOD_URL = os.getenv("RUNPOD_URL", "").strip()
SUPPORT_URL = os.getenv("SUPPORT_URL", "").strip()
MAX_GPU_SECONDS = int(os.getenv("MAX_GPU_SECONDS", "300"))

IS_ZERO_GPU = bool(os.getenv("SPACES_ZERO_GPU"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

MAX_DIM = 832
MIN_DIM = 480
SQUARE_DIM = 640
MULTIPLE_OF = 16
BASE_FPS = 16
MIN_MODEL_FRAMES = 9
MAX_MODEL_FRAMES = 161
MIN_DURATION = round((MIN_MODEL_FRAMES - 1) / BASE_FPS, 1)
MAX_DURATION = round((MAX_MODEL_FRAMES - 1) / BASE_FPS, 1)
MAX_SEED = int(np.iinfo(np.int32).max)

ASPECT_DIMENSIONS: dict[str, tuple[int, int] | None] = {
    "Auto (keep source ratio)": None,
    "Landscape 16:9": (832, 480),
    "Portrait 9:16": (480, 832),
    "Square 1:1": (640, 640),
    "Portrait 4:5": (512, 640),
    "Landscape 3:2": (768, 512),
}

SCHEDULER_MAP = {
    "FlowMatchEulerDiscrete": FlowMatchEulerDiscreteScheduler,
    "SASolver": SASolverScheduler,
    "DEISMultistep": DEISMultistepScheduler,
    "DPMSolverMultistepInverse": DPMSolverMultistepInverseScheduler,
    "UniPCMultistep": UniPCMultistepScheduler,
    "DPMSolverMultistep": DPMSolverMultistepScheduler,
    "DPMSolverSinglestep": DPMSolverSinglestepScheduler,
}

QUALITY_PROFILES = {
    "Fast Preview": {"steps": 4, "quality": 5, "flow_shift": 3.0},
    "Balanced": {"steps": 6, "quality": 7, "flow_shift": 3.0},
    "High Quality": {"steps": 8, "quality": 8, "flow_shift": 3.5},
}

MOTION_STYLE_PROMPTS = {
    "Natural": "natural realistic movement, physically plausible motion, stable composition",
    "Cinematic": "cinematic motion, elegant pacing, dramatic depth, premium film look",
    "Portrait": "subtle facial expression, natural blinking, gentle head movement, stable facial identity",
    "Product Ad": "premium product commercial, controlled motion, clean composition, polished advertising look",
    "Dynamic": "energetic movement, stronger motion, dramatic action, coherent fast pacing",
    "Anime": "fluid anime-style animation, clean line consistency, expressive but stable movement",
}

CAMERA_PROMPTS = {
    "Static": "locked-off camera, stable framing",
    "Slow Push In": "slow smooth camera push-in",
    "Slow Pull Back": "slow smooth camera pull-back",
    "Pan Left": "slow cinematic camera pan to the left",
    "Pan Right": "slow cinematic camera pan to the right",
    "Orbit": "smooth controlled camera orbit around the subject",
    "Handheld": "subtle realistic handheld camera movement",
}

MOTION_LEVEL_PROMPTS = {
    "Low": "minimal subtle motion, preserve the original composition",
    "Medium": "moderate natural motion with stable details",
    "High": "strong visible motion while keeping the subject coherent",
}

DEFAULT_PROMPT = "Make this image come alive with smooth, realistic motion."
DEFAULT_NEGATIVE_PROMPT = (
    "overexposed, oversaturated, static frame, blurry details, low quality, worst quality, "
    "jpeg artifacts, text, subtitles, watermark, logo, deformed face, identity drift, "
    "bad anatomy, malformed limbs, extra fingers, fused fingers, duplicated person, "
    "warped body, flicker, jitter, unstable background, abrupt camera movement, reverse walking, "
    "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 最差质量, 低质量, 多余的手指, "
    "画得不好的手部, 画得不好的脸部, 畸形的, 静止不动的画面, 杂乱的背景"
)

IDENTITY_PROMPT = (
    "preserve the exact facial identity, facial structure, hairstyle, skin tone, clothing, "
    "body proportions and distinguishing features throughout the video"
)
IDENTITY_NEGATIVE = (
    "face morphing, changing identity, changing hairstyle, changing clothes, distorted eyes, "
    "asymmetric face, duplicated facial features"
)

# -----------------------------------------------------------------------------
# Model loading: Wan 2.2 + FP8 + AOTI
# -----------------------------------------------------------------------------

if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_float32_matmul_precision("high")


def _load_aoti_component(module: torch.nn.Module, repo_id: str) -> None:
    """Load an AOTI package while remaining compatible with different spaces builds."""
    if hasattr(spaces, "aoti_load"):
        spaces.aoti_load(module=module, repo_id=repo_id)
        return

    from aoti import aoti_blocks_load

    aoti_blocks_load(module=module, repo_id=repo_id)


def load_pipeline() -> WanImageToVideoPipeline:
    print(f"Loading model: {MODEL_ID}")
    pipeline = WanImageToVideoPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.bfloat16,
    ).to("cuda")

    # The selected WAMU repository is already a merged Lightning/MotionBoost model.
    # Do not fuse a second Lightning LoRA here; that can reduce quality and consistency.
    print("Quantizing text encoder to INT8...")
    quantize_(pipeline.text_encoder, Int8WeightOnlyConfig())
    torch._dynamo.reset()

    print("Quantizing high-noise transformer to FP8...")
    quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
    torch._dynamo.reset()

    print("Quantizing low-noise transformer to FP8...")
    quantize_(pipeline.transformer_2, Float8DynamicActivationFloat8WeightConfig())
    torch._dynamo.reset()

    if ENABLE_AOTI:
        print(f"Loading AOTI packages from: {AOTI_REPO}")
        _load_aoti_component(pipeline.transformer, AOTI_REPO)
        _load_aoti_component(pipeline.transformer_2, AOTI_REPO)

    if ENABLE_VAE_TILING:
        try:
            pipeline.vae.enable_slicing()
            pipeline.vae.enable_tiling()
            print("VAE slicing and tiling enabled.")
        except Exception as exc:
            print(f"VAE tiling could not be enabled: {exc}")

    return pipeline


pipe = load_pipeline()
ORIGINAL_SCHEDULER = copy.deepcopy(pipe.scheduler)

# -----------------------------------------------------------------------------
# Optional RIFE interpolation, loaded lazily only when 32/64 FPS is requested
# -----------------------------------------------------------------------------

RIFE_REPO = "r3gm/RIFE"
RIFE_FILENAME = "RIFEv4.26_0921.zip"
RIFE_MODEL: Any | None = None


def ensure_rife_model() -> Any:
    global RIFE_MODEL
    if RIFE_MODEL is not None:
        return RIFE_MODEL

    if not Path("train_log/RIFE_HDv3.py").exists():
        print("Downloading RIFE frame-interpolation model...")
        archive = hf_hub_download(repo_id=RIFE_REPO, filename=RIFE_FILENAME)
        with zipfile.ZipFile(archive, "r") as zip_ref:
            zip_ref.extractall(".")

    if os.getcwd() not in sys.path:
        sys.path.insert(0, os.getcwd())

    from train_log.RIFE_HDv3 import Model

    model = Model()
    model.load_model("train_log", -1)
    model.eval()
    RIFE_MODEL = model
    return RIFE_MODEL


@torch.no_grad()
def interpolate_frames(
    frames_np: np.ndarray | list[np.ndarray],
    multiplier: int = 2,
    scale: float = 1.0,
) -> list[np.ndarray]:
    if multiplier < 2:
        return list(frames_np)

    model = ensure_rife_model()
    model.device()
    model.flownet = model.flownet.half()

    if isinstance(frames_np, list):
        total_frames = len(frames_np)
        height, width, _ = frames_np[0].shape
    else:
        total_frames, height, width, _ = frames_np.shape

    interpolation_count = multiplier - 1
    block = max(128, int(128 / scale))
    padded_height = ((height - 1) // block + 1) * block
    padded_width = ((width - 1) // block + 1) * block
    padding = (0, padded_width - width, 0, padded_height - height)

    def to_tensor(frame: np.ndarray) -> torch.Tensor:
        tensor = torch.from_numpy(frame).to(DEVICE)
        tensor = tensor.permute(2, 0, 1).unsqueeze(0)
        return F.pad(tensor, padding).half()

    def to_numpy(tensor: torch.Tensor) -> np.ndarray:
        tensor = tensor[0, :, :height, :width].permute(1, 2, 0)
        return tensor.float().cpu().numpy()

    def infer_between(first: torch.Tensor, second: torch.Tensor, count: int) -> list[torch.Tensor]:
        if model.version >= 3.9:
            return [
                model.inference(first, second, (index + 1) / (count + 1), scale)
                for index in range(count)
            ]

        middle = model.inference(first, second, scale)
        if count == 1:
            return [middle]
        first_half = infer_between(first, middle, count // 2)
        second_half = infer_between(middle, second, count // 2)
        if count % 2:
            return [*first_half, middle, *second_half]
        return [*first_half, *second_half]

    output: list[np.ndarray] = []
    next_tensor = to_tensor(frames_np[0])

    with tqdm(total=total_frames - 1, desc="Frame interpolation", unit="frame") as bar:
        for index in range(total_frames - 1):
            current_tensor = next_tensor
            output.append(to_numpy(current_tensor))
            next_tensor = to_tensor(frames_np[index + 1])
            for middle_tensor in infer_between(
                current_tensor, next_tensor, interpolation_count
            ):
                output.append(to_numpy(middle_tensor))
            bar.update(1)

        output.append(to_numpy(next_tensor))

    try:
        model.flownet.to("cpu")
    except Exception:
        pass
    clear_vram()
    return output

# -----------------------------------------------------------------------------
# Image, prompt and scheduler helpers
# -----------------------------------------------------------------------------


def clear_vram() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def normalize_image(image: Image.Image) -> Image.Image:
    return ImageOps.exif_transpose(image).convert("RGB")


def center_crop_resize(image: Image.Image, size: tuple[int, int]) -> Image.Image:
    image = normalize_image(image)
    target_width, target_height = size
    width, height = image.size
    scale = max(target_width / width, target_height / height)
    resized_width = max(target_width, round(width * scale))
    resized_height = max(target_height, round(height * scale))
    image = image.resize((resized_width, resized_height), Image.Resampling.LANCZOS)
    left = (resized_width - target_width) // 2
    top = (resized_height - target_height) // 2
    return image.crop((left, top, left + target_width, top + target_height))


def resize_auto(image: Image.Image) -> Image.Image:
    image = normalize_image(image)
    width, height = image.size

    if width == height:
        return image.resize((SQUARE_DIM, SQUARE_DIM), Image.Resampling.LANCZOS)

    aspect_ratio = width / height
    maximum_ratio = MAX_DIM / MIN_DIM
    minimum_ratio = MIN_DIM / MAX_DIM
    source = image

    if aspect_ratio > maximum_ratio:
        target_width, target_height = MAX_DIM, MIN_DIM
        crop_width = round(height * maximum_ratio)
        left = (width - crop_width) // 2
        source = image.crop((left, 0, left + crop_width, height))
    elif aspect_ratio < minimum_ratio:
        target_width, target_height = MIN_DIM, MAX_DIM
        crop_height = round(width / minimum_ratio)
        top = (height - crop_height) // 2
        source = image.crop((0, top, width, top + crop_height))
    elif width > height:
        target_width = MAX_DIM
        target_height = round(target_width / aspect_ratio)
    else:
        target_height = MAX_DIM
        target_width = round(target_height * aspect_ratio)

    final_width = round(target_width / MULTIPLE_OF) * MULTIPLE_OF
    final_height = round(target_height / MULTIPLE_OF) * MULTIPLE_OF
    final_width = max(MIN_DIM, min(MAX_DIM, final_width))
    final_height = max(MIN_DIM, min(MAX_DIM, final_height))
    return source.resize((final_width, final_height), Image.Resampling.LANCZOS)


def prepare_first_frame(image: Image.Image, aspect_ratio: str) -> Image.Image:
    dimensions = ASPECT_DIMENSIONS.get(aspect_ratio)
    if dimensions is None:
        return resize_auto(image)
    return center_crop_resize(image, dimensions)


def prepare_last_frame(image: Image.Image, reference: Image.Image) -> Image.Image:
    return center_crop_resize(image, reference.size)


def get_num_frames(duration_seconds: float) -> int:
    target = int(round(float(duration_seconds) * BASE_FPS)) + 1
    # Wan video frame counts are most reliable as 4n + 1.
    target = 4 * round((target - 1) / 4) + 1
    return int(np.clip(target, MIN_MODEL_FRAMES, MAX_MODEL_FRAMES))


def build_prompts(
    prompt: str,
    motion_style: str,
    camera_motion: str,
    motion_strength: str,
    enhance_prompt: bool,
    preserve_identity: bool,
    negative_prompt: str,
) -> tuple[str, str]:
    prompt = (prompt or "").strip()
    if not prompt:
        prompt = DEFAULT_PROMPT

    additions: list[str] = []
    if enhance_prompt:
        additions.extend(
            [
                MOTION_STYLE_PROMPTS.get(motion_style, ""),
                CAMERA_PROMPTS.get(camera_motion, ""),
                MOTION_LEVEL_PROMPTS.get(motion_strength, ""),
                "smooth temporal consistency, coherent details, natural motion blur, no flicker",
            ]
        )
    if preserve_identity:
        additions.append(IDENTITY_PROMPT)

    final_prompt = ", ".join(part for part in [prompt, *additions] if part)
    final_negative = (negative_prompt or DEFAULT_NEGATIVE_PROMPT).strip()
    if preserve_identity:
        final_negative = f"{final_negative}, {IDENTITY_NEGATIVE}"
    return final_prompt, final_negative


def configure_scheduler(name: str, flow_shift: float) -> None:
    scheduler_class = SCHEDULER_MAP.get(name, UniPCMultistepScheduler)
    current_name = pipe.scheduler.config.get("_class_name", pipe.scheduler.__class__.__name__)
    current_shift = pipe.scheduler.config.get(
        "flow_shift", pipe.scheduler.config.get("shift", None)
    )

    if current_name == scheduler_class.__name__ and current_shift == flow_shift:
        return

    config = copy.deepcopy(ORIGINAL_SCHEDULER.config)
    if scheduler_class is FlowMatchEulerDiscreteScheduler:
        config["shift"] = float(flow_shift)
    else:
        config["flow_shift"] = float(flow_shift)
    pipe.scheduler = scheduler_class.from_config(config)


def profile_settings(profile: str):
    values = QUALITY_PROFILES.get(profile, QUALITY_PROFILES["Balanced"])
    return (
        gr.update(value=values["steps"]),
        gr.update(value=values["quality"]),
        gr.update(value=values["flow_shift"]),
    )


def swap_images(first: Image.Image | None, last: Image.Image | None):
    return last, first


def format_external_links() -> str:
    links: list[str] = []
    if RUNPOD_URL:
        links.append(f"[⚡ Run on Private GPU via RunPod]({RUNPOD_URL})")
    if SUPPORT_URL:
        links.append(f"[❤️ Support This Free Tool]({SUPPORT_URL})")
    return " &nbsp; • &nbsp; ".join(links)

# -----------------------------------------------------------------------------
# GPU generation
# -----------------------------------------------------------------------------


def get_inference_duration(
    first_frame: Image.Image,
    last_frame: Image.Image | None,
    prompt: str,
    negative_prompt: str,
    num_frames: int,
    steps: int,
    guidance_scale: float,
    guidance_scale_2: float,
    seed: int,
    scheduler_name: str,
    flow_shift: float,
    output_fps: int,
    video_quality: int,
    safe_mode: bool,
    progress,
) -> int:
    del last_frame, prompt, negative_prompt, guidance_scale_2, seed
    del scheduler_name, flow_shift, video_quality, progress

    base_pixels = 81 * 832 * 624
    width, height = first_frame.size
    factor = max(0.2, num_frames * width * height / base_pixels)
    seconds_per_step = 5.0 * factor**1.5
    estimated = 15 + int(steps) * seconds_per_step

    if float(guidance_scale) > 1:
        estimated *= 2.4

    interpolation_multiplier = max(1, int(output_fps) // BASE_FPS)
    if interpolation_multiplier > 1:
        extra_frames = num_frames * interpolation_multiplier - num_frames
        estimated += extra_frames * 0.025

    if safe_mode:
        estimated *= 1.25

    return int(max(30, min(MAX_GPU_SECONDS, round(estimated))))


@spaces.GPU(duration=get_inference_duration, size="xlarge")
def run_inference(
    first_frame: Image.Image,
    last_frame: Image.Image | None,
    prompt: str,
    negative_prompt: str,
    num_frames: int,
    steps: int,
    guidance_scale: float,
    guidance_scale_2: float,
    seed: int,
    scheduler_name: str,
    flow_shift: float,
    output_fps: int,
    video_quality: int,
    safe_mode: bool,
    progress=gr.Progress(track_tqdm=True),
) -> tuple[str, str, float]:
    del safe_mode
    clear_vram()
    configure_scheduler(scheduler_name, flow_shift)

    task_id = str(uuid.uuid4())[:8]
    started = time.time()
    print(
        f"Task {task_id}: {num_frames} frames, {first_frame.size}, "
        f"{steps} steps, seed={seed}"
    )

    try:
        result = pipe(
            image=first_frame,
            last_image=last_frame,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=first_frame.height,
            width=first_frame.width,
            num_frames=int(num_frames),
            guidance_scale=float(guidance_scale),
            guidance_scale_2=float(guidance_scale_2),
            num_inference_steps=int(steps),
            generator=torch.Generator(device="cuda").manual_seed(int(seed)),
            output_type="np",
        )

        raw_frames = result.frames[0]
        multiplier = max(1, int(output_fps) // BASE_FPS)
        if multiplier > 1:
            print(f"Task {task_id}: applying RIFE {multiplier}x interpolation")
            final_frames = interpolate_frames(raw_frames, multiplier=multiplier)
        else:
            final_frames = list(raw_frames)

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_file:
            video_path = temp_file.name

        export_to_video(
            final_frames,
            video_path,
            fps=int(output_fps),
            quality=int(video_quality),
        )
        elapsed = time.time() - started
        print(f"Task {task_id}: completed in {elapsed:.1f}s")
        return video_path, task_id, elapsed
    finally:
        pipe.scheduler = ORIGINAL_SCHEDULER
        clear_vram()


def generate_video(
    first_image: Image.Image | None,
    last_image: Image.Image | None,
    generation_mode: str,
    prompt: str,
    motion_style: str,
    camera_motion: str,
    motion_strength: str,
    preserve_identity: bool,
    enhance_prompt: bool,
    aspect_ratio: str,
    duration_seconds: float,
    output_fps: int,
    quality_profile: str,
    steps: int,
    negative_prompt: str,
    video_quality: int,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    guidance_scale_2: float,
    scheduler_name: str,
    flow_shift: float,
    safe_mode: bool,
    show_preview: bool,
    progress=gr.Progress(track_tqdm=True),
):
    if first_image is None:
        raise gr.Error("Please upload a first image.")

    use_last_frame = generation_mode == "First + Last Frame"
    if use_last_frame and last_image is None:
        raise gr.Error("Please upload a last image or switch to Single Image mode.")

    first_frame = prepare_first_frame(first_image, aspect_ratio)
    final_last_frame = (
        prepare_last_frame(last_image, first_frame)
        if use_last_frame and last_image is not None
        else None
    )

    final_prompt, final_negative = build_prompts(
        prompt=prompt,
        motion_style=motion_style,
        camera_motion=camera_motion,
        motion_strength=motion_strength,
        enhance_prompt=enhance_prompt,
        preserve_identity=preserve_identity,
        negative_prompt=negative_prompt,
    )

    num_frames = get_num_frames(duration_seconds)
    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    try:
        video_path, task_id, elapsed = run_inference(
            first_frame,
            final_last_frame,
            final_prompt,
            final_negative,
            num_frames,
            int(steps),
            float(guidance_scale),
            float(guidance_scale_2),
            current_seed,
            scheduler_name,
            float(flow_shift),
            int(output_fps),
            int(video_quality),
            bool(safe_mode),
            progress,
        )
    except gr.Error:
        raise
    except Exception as exc:
        clear_vram()
        message = str(exc).strip() or exc.__class__.__name__
        print(f"Generation failed: {message}")
        raise gr.Error(f"Generation failed: {message}") from exc

    mode_label = "First → Last" if use_last_frame else "Single Image"
    status = (
        f"### ✅ Generation complete\n"
        f"**Task:** `{task_id}` &nbsp; • &nbsp; **Mode:** {mode_label} &nbsp; • &nbsp; "
        f"**Resolution:** {first_frame.width}×{first_frame.height} &nbsp; • &nbsp; "
        f"**Frames:** {num_frames} at {output_fps} FPS &nbsp; • &nbsp; "
        f"**Seed:** `{current_seed}` &nbsp; • &nbsp; **Time:** {elapsed:.1f}s"
    )

    return (
        video_path if show_preview else None,
        video_path,
        current_seed,
        status,
        final_prompt,
    )

# -----------------------------------------------------------------------------
# Frame extraction from generated video
# -----------------------------------------------------------------------------

GET_TIMESTAMP_JS = """
function() {
    const video = document.querySelector('#generated-video video');
    return video ? video.currentTime : 0;
}
"""


def extract_frame(video_path: str | None, timestamp: float) -> np.ndarray | None:
    if not video_path:
        raise gr.Error("Generate or upload a video first.")

    capture = cv2.VideoCapture(video_path)
    if not capture.isOpened():
        raise gr.Error("The video could not be opened.")

    fps = capture.get(cv2.CAP_PROP_FPS) or BASE_FPS
    total_frames = max(1, int(capture.get(cv2.CAP_PROP_FRAME_COUNT)))
    frame_number = int(max(0, float(timestamp)) * fps)
    frame_number = min(frame_number, total_frames - 1)
    capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
    success, frame = capture.read()
    capture.release()

    if not success:
        raise gr.Error("The selected frame could not be extracted.")
    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# -----------------------------------------------------------------------------
# Interface
# -----------------------------------------------------------------------------


def model_heading() -> str:
    name = MODEL_ID.split("/")[-1].replace("_", " ")
    return (
        "<div class='hero'>"
        "<div class='hero-badge'>WAN 2.2 • FP8 • AOTI • Lightning</div>"
        "<h1>Dream Motion Pro</h1>"
        "<p>Turn a single image—or a first and last frame—into a smooth cinematic video.</p>"
        f"<p class='model-line'>Running <a href='https://huggingface.co/{MODEL_ID}' target='_blank'>{name}</a></p>"
        "</div>"
    )


CSS = """
.gradio-container { max-width: 1380px !important; }
.hero {
    padding: 26px 28px;
    margin-bottom: 18px;
    border: 1px solid var(--border-color-primary);
    border-radius: 20px;
    background: linear-gradient(135deg, rgba(79,70,229,.15), rgba(6,182,212,.09));
}
.hero h1 { margin: 8px 0 5px; font-size: 2.15rem; }
.hero p { margin: 4px 0; opacity: .88; }
.hero-badge {
    display: inline-block;
    padding: 6px 11px;
    border-radius: 999px;
    font-weight: 700;
    font-size: .78rem;
    letter-spacing: .04em;
    border: 1px solid var(--border-color-primary);
}
.model-line { font-size: .9rem; }
#generate-button { min-height: 52px; font-weight: 800; font-size: 1.02rem; }
#generated-video video { border-radius: 16px; }
#hidden-timestamp {
    opacity: 0;
    height: 0;
    width: 0;
    margin: 0;
    padding: 0;
    overflow: hidden;
    position: absolute;
    pointer-events: none;
}
.link-row { text-align: center; padding: 8px 0 2px; font-weight: 650; }
.tip-box {
    border: 1px solid var(--border-color-primary);
    border-radius: 14px;
    padding: 12px 14px;
}
"""

with gr.Blocks(
    delete_cache=(3600, 10800),
) as demo:
    gr.HTML(model_heading())

    with gr.Row(equal_height=False):
        with gr.Column(scale=6):
            generation_mode = gr.Radio(
                ["Single Image", "First + Last Frame"],
                value="Single Image",
                label="Generation Mode",
                info="The second mode guides both the beginning and ending of the video.",
            )

            with gr.Row():
                first_image = gr.Image(
                    type="pil",
                    label="First Image",
                    sources=["upload", "clipboard"],
                    height=320,
                )
                last_image = gr.Image(
                    type="pil",
                    label="Last Image (used only in First + Last Frame mode)",
                    sources=["upload", "clipboard"],
                    height=320,
                )

            with gr.Row():
                swap_button = gr.Button("⇄ Swap First / Last", variant="secondary")

            prompt_input = gr.Textbox(
                label="Describe the motion",
                value=DEFAULT_PROMPT,
                lines=3,
                placeholder="Example: The subject looks toward the camera while the camera slowly pushes in.",
            )

            motion_style = gr.Radio(
                choices=list(MOTION_STYLE_PROMPTS.keys()),
                value="Cinematic",
                label="Motion Style",
            )

            with gr.Row():
                quality_profile = gr.Dropdown(
                    choices=list(QUALITY_PROFILES.keys()),
                    value="Balanced",
                    label="Quality Mode",
                )
                aspect_ratio = gr.Dropdown(
                    choices=list(ASPECT_DIMENSIONS.keys()),
                    value="Auto (keep source ratio)",
                    label="Aspect Ratio",
                )

            with gr.Row():
                duration_seconds = gr.Slider(
                    minimum=MIN_DURATION,
                    maximum=MAX_DURATION,
                    step=0.5,
                    value=3.5,
                    label="Duration (seconds)",
                )
                output_fps = gr.Dropdown(
                    choices=[16, 32, 64],
                    value=16,
                    label="Output FPS",
                    info="32/64 FPS uses RIFE interpolation after generation.",
                )

            with gr.Row():
                preserve_identity = gr.Checkbox(
                    value=True,
                    label="Preserve Face & Identity",
                )
                enhance_prompt = gr.Checkbox(
                    value=True,
                    label="Professional Prompt Enhancement",
                )

            with gr.Accordion("Advanced Controls", open=False):
                with gr.Row():
                    camera_motion = gr.Dropdown(
                        choices=list(CAMERA_PROMPTS.keys()),
                        value="Slow Push In",
                        label="Camera Movement",
                    )
                    motion_strength = gr.Radio(
                        choices=list(MOTION_LEVEL_PROMPTS.keys()),
                        value="Medium",
                        label="Motion Strength",
                    )

                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    value=DEFAULT_NEGATIVE_PROMPT,
                    lines=4,
                )

                with gr.Row():
                    steps = gr.Slider(
                        minimum=1,
                        maximum=10,
                        step=1,
                        value=6,
                        label="Inference Steps",
                        info="Lightning models normally perform best around 4–8 steps.",
                    )
                    video_quality = gr.Slider(
                        minimum=1,
                        maximum=10,
                        step=1,
                        value=7,
                        label="MP4 Quality",
                    )

                with gr.Row():
                    seed = gr.Slider(
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                        label="Seed",
                    )
                    randomize_seed = gr.Checkbox(
                        value=True,
                        label="Randomize Seed",
                    )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        minimum=0.0,
                        maximum=6.0,
                        step=0.5,
                        value=1.0,
                        label="High-noise Guidance",
                    )
                    guidance_scale_2 = gr.Slider(
                        minimum=0.0,
                        maximum=6.0,
                        step=0.5,
                        value=1.0,
                        label="Low-noise Guidance",
                    )

                with gr.Row():
                    scheduler_name = gr.Dropdown(
                        choices=list(SCHEDULER_MAP.keys()),
                        value="UniPCMultistep",
                        label="Scheduler",
                    )
                    flow_shift = gr.Slider(
                        minimum=0.5,
                        maximum=12.0,
                        step=0.1,
                        value=3.0,
                        label="Flow Shift",
                    )

                with gr.Row():
                    safe_mode = gr.Checkbox(
                        value=True,
                        label="Extra ZeroGPU Time Buffer",
                    )
                    show_preview = gr.Checkbox(
                        value=True,
                        label="Show Video Preview",
                    )

            generate_button = gr.Button(
                "✨ Generate Video",
                variant="primary",
                elem_id="generate-button",
            )

            gr.Markdown(
                "**Best results:** use a clear image, describe one main action, keep camera movement controlled, "
                "and generate 2–3 variations with different seeds."
            )

        with gr.Column(scale=6):
            video_output = gr.Video(
                label="Generated Video",
                autoplay=True,
                buttons=["download", "share"],
                interactive=False,
                elem_id="generated-video",
            )
            status_output = gr.Markdown("Your generation details will appear here.")

            with gr.Row():
                grab_frame_button = gr.Button(
                    "📸 Use Current Video Frame as First Image",
                    variant="secondary",
                )
                timestamp_box = gr.Number(
                    value=0,
                    visible=True,
                    elem_id="hidden-timestamp",
                )

            file_output = gr.File(label="Download Original MP4")

            with gr.Accordion("Prompt used for this result", open=False):
                final_prompt_output = gr.Textbox(
                    label="Enhanced Prompt",
                    interactive=False,
                    lines=5,
                )

            external_links = format_external_links()
            if external_links:
                gr.Markdown(
                    f"<div class='link-row'>{external_links}</div>",
                    sanitize_html=False,
                )

            gr.Markdown(
                "### Why this Space is fast\n"
                "The Wan 2.2 high-noise and low-noise transformers run with FP8 quantization and AOTI-compiled blocks. "
                "The selected WAMU checkpoint already includes Lightning and MotionBoost behavior."
            )

    quality_profile.change(
        fn=profile_settings,
        inputs=[quality_profile],
        outputs=[steps, video_quality, flow_shift],
    )

    swap_button.click(
        fn=swap_images,
        inputs=[first_image, last_image],
        outputs=[first_image, last_image],
    )

    inputs = [
        first_image,
        last_image,
        generation_mode,
        prompt_input,
        motion_style,
        camera_motion,
        motion_strength,
        preserve_identity,
        enhance_prompt,
        aspect_ratio,
        duration_seconds,
        output_fps,
        quality_profile,
        steps,
        negative_prompt,
        video_quality,
        seed,
        randomize_seed,
        guidance_scale,
        guidance_scale_2,
        scheduler_name,
        flow_shift,
        safe_mode,
        show_preview,
    ]

    generate_button.click(
        fn=generate_video,
        inputs=inputs,
        outputs=[
            video_output,
            file_output,
            seed,
            status_output,
            final_prompt_output,
        ],
    )

    grab_frame_button.click(
        fn=None,
        inputs=None,
        outputs=[timestamp_box],
        js=GET_TIMESTAMP_JS,
    )
    timestamp_box.change(
        fn=extract_frame,
        inputs=[video_output, timestamp_box],
        outputs=[first_image],
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=20).launch(
        mcp_server=True,
        show_error=True,
        css=CSS,
        theme=gr.themes.Soft(),
    )