import os

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

import spaces  # MUST come before torch / diffusers / transformers

import torch
import torch.nn as nn
import numpy as np
import gradio as gr
from omegaconf import OmegaConf
from PIL import Image
from typing import List, Optional
import tempfile
from huggingface_hub import hf_hub_download
import random

from torchvision import transforms
from decord import VideoReader, cpu
from diffusers.utils.export_utils import export_to_video
from transformers import AutoProcessor, AutoTokenizer

# Add project root to path for src imports
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.models.utils import load_model, load_checkpoint
from src.models.transformers.loomvideo import LoomVideo

# ---------------------------------------------------------------------------
# Config — use HF hub repo IDs instead of local checkpoint paths
# ---------------------------------------------------------------------------
UND_MODEL_PATH = "Qwen/Qwen3-VL-8B-Instruct"
GEN_MODEL_PATH = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
LOOMVIDEO_CKPT = "MSALab/LoomVideo"

# Build config the way the codebase expects it
config = OmegaConf.create({
    "model": {
        "trainable_modules": {"gen_model": "all"},
        "gradient_checkpointing": False,
        "und": {
            "pretrained_model_path": UND_MODEL_PATH,
            "only_last_hidden_states": False,
            "add_source_video": True,
        },
        "gen": {
            "pretrained_model_path": GEN_MODEL_PATH,
            "base_timestep_shift": 1.0,
            "num_attn_token_base_shift": 64,
            "timestep_shift_scale": 0.5,
            "use_source_embedding": True,
        },
    },
    "data": {"und_temporal_downsample_factor": 15},
    "generation": {
        "negative_prompt": "Generate an image: 色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
        "guidance_scale": 5.0,
        "guidance_scale_edit": 2.5,
        "guidance_scale_visual": 1.5,
    },
})

# Preprocessing: normalize to [-1, 1] for VAE/DiT conditioning
IMAGE_TRANSFORM = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])


def resize_and_crop(image: Image.Image, target_height: int, target_width: int) -> Image.Image:
    """Resize and center-crop image to exact target dimensions."""
    width, height = image.size
    scale = max(target_width / width, target_height / height)
    image = transforms.functional.resize(
        image,
        (round(height * scale), round(width * scale)),
        interpolation=transforms.InterpolationMode.BILINEAR,
    )
    image = transforms.functional.center_crop(image, (target_height, target_width))
    return image


def load_image(image_path: str, target_height: int = None, target_width: int = None) -> tuple:
    """Load and preprocess a single image, returning both PIL and tensor."""
    image = Image.open(image_path).convert("RGB")
    if target_height is not None and target_width is not None:
        image = resize_and_crop(image, target_height, target_width)
    return image, IMAGE_TRANSFORM(image)


def load_images(image_paths: List[str], target_height: int = None, target_width: int = None) -> tuple:
    """Load multiple images, returning both PIL images and tensors."""
    pil_images = []
    tensors = []
    for path in image_paths:
        pil_img, tensor = load_image(path, target_height, target_width)
        pil_images.append(pil_img)
        tensors.append(tensor)
    return pil_images, tensors


def load_video(
    video_path: str,
    target_fps: int = 24,
    target_num_frames: int = 121,
    scale_factor_temporal: int = 4,
    target_height: int = None,
    target_width: int = None,
) -> tuple:
    """Load and preprocess a video, returning both PIL frames and normalized tensor."""
    video_reader = VideoReader(video_path, ctx=cpu(0))
    video_fps = video_reader.get_avg_fps()
    num_video_frames = len(video_reader)

    target_duration = (target_num_frames - 1) / target_fps
    num_video_in_range_frames = min(int(video_fps * target_duration) + 1, num_video_frames)

    start_frame_index = 0
    frame_float = float(start_frame_index)
    interval = video_fps / target_fps

    frame_indices = []
    max_loops = target_num_frames * 3
    loop_count = 0

    while True:
        idx = min(round(frame_float), num_video_frames - 1)
        frame_indices.append(idx)
        frame_float += interval

        if (round(frame_float) - start_frame_index) >= num_video_in_range_frames:
            break

        loop_count += 1
        if loop_count > max_loops:
            break

    needed_len = (len(frame_indices) - 1) // scale_factor_temporal * scale_factor_temporal + 1
    frame_indices = frame_indices[:needed_len]

    batch = video_reader.get_batch(frame_indices)
    video_data = batch.asnumpy()

    pil_frames = []
    frame_tensors = []
    for frame in video_data:
        frame_pil = Image.fromarray(frame)
        if target_height is not None and target_width is not None:
            frame_pil = resize_and_crop(frame_pil, target_height, target_width)
        pil_frames.append(frame_pil)
        frame_tensors.append(IMAGE_TRANSFORM(frame_pil))

    video_tensor = torch.stack(frame_tensors, dim=0).permute(1, 0, 2, 3)
    return pil_frames, video_tensor


def temporal_downsample_for_vlm(pil_frames: List[Image.Image], cfg) -> List[Image.Image]:
    """Downsample video frames for VLM input."""
    downsample_factor = OmegaConf.select(cfg, "data.und_temporal_downsample_factor", default=15)
    total_frames = len(pil_frames)
    if total_frames > downsample_factor:
        indices = np.linspace(
            0, total_frames - 1,
            max(total_frames // downsample_factor, 1),
            dtype=int,
        ).tolist()
    else:
        indices = list(range(total_frames))
    return [pil_frames[i] for i in indices]


# ---------------------------------------------------------------------------
# Model loading at module scope (ZeroGPU requirement)
# ---------------------------------------------------------------------------
print("Building LoomVideo model...")
model = load_model(config)
print("Loading LoomVideo checkpoint...")
ckpt_file = hf_hub_download(
    repo_id=LOOMVIDEO_CKPT,
    filename="Stage3/gen_model.pth",
    repo_type="model",
)
load_checkpoint(model, ckpt_file)
model.to(dtype=torch.bfloat16, device="cuda")
model.eval()
print("Model loaded successfully.")

# Cache processor and tokenizer
_processor = AutoProcessor.from_pretrained(UND_MODEL_PATH, trust_remote_code=True)


def _get_negative_prompt(negative_prompt: Optional[str]) -> str:
    if negative_prompt is not None:
        return negative_prompt
    return OmegaConf.select(config, "generation.negative_prompt", default="")


def _get_guidance_scale(guidance_scale: Optional[float], is_edit: bool) -> float:
    if guidance_scale is not None:
        return guidance_scale
    if is_edit:
        return OmegaConf.select(config, "generation.guidance_scale_edit", default=2.5)
    return OmegaConf.select(config, "generation.guidance_scale", default=5.0)


def _get_guidance_scale_visual(guidance_scale_visual: Optional[float]) -> float:
    if guidance_scale_visual is not None:
        return guidance_scale_visual
    return OmegaConf.select(config, "generation.guidance_scale_visual", default=1.5)


def _build_generate_kwargs(inputs, **extra):
    """Build kwargs dict for model.generate() from processor outputs."""
    kwargs = {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
    }
    # Pass through optional keys that the processor may produce
    for key in ["pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw", "mm_token_type_ids"]:
        if key in inputs:
            kwargs[key] = inputs[key]
    kwargs.update(extra)
    return kwargs


# ---------------------------------------------------------------------------
# Inference functions
# ---------------------------------------------------------------------------

def _estimate_duration(task, num_frames, num_steps, *args, **kwargs):
    """Estimate GPU duration based on task and parameters."""
    base = 30
    if task in ("t2v", "mi2v"):
        base = 120
    elif task in ("edit", "ref_edit"):
        base = 120
    # Scale by steps ratio (50 is the default)
    step_factor = num_steps / 50.0
    # Scale by frame count
    frame_factor = num_frames / 81.0
    return min(300, int(base * step_factor * max(frame_factor, 0.5)))


@spaces.GPU(duration=_estimate_duration)
def generate_t2v(
    prompt: str,
    num_frames: int = 81,
    height: int = 480,
    width: int = 832,
    num_inference_steps: int = 50,
    guidance_scale: float = 5.0,
    guidance_scale_visual: float = 1.5,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a video from a text prompt.

    Args:
        prompt: Text description of the video to generate.
        num_frames: Number of output video frames (1 = image, >1 = video).
        height: Output video height in pixels.
        width: Output video width in pixels.
        num_inference_steps: Number of denoising steps.
        guidance_scale: Text CFG scale.
        guidance_scale_visual: Visual CFG scale.
        seed: Random seed (-1 for random).
    """
    is_image = (num_frames == 1)
    prefix = "Generate an image: " if is_image else "Generate a video: "

    messages = [{"role": "user", "content": [{"type": "text", "text": f"{prefix}{prompt}"}]}]
    inputs = _processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    ).to(model.device)

    if seed >= 0:
        generator = torch.Generator(device=model.device).manual_seed(seed)
    else:
        generator = None

    generate_kwargs = _build_generate_kwargs(
        inputs,
        negative_prompt=_get_negative_prompt(None),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=False),
        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
        generator=generator,
    )

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        output = model.generate(**generate_kwargs)

    return _save_output(output, fps=24)


@spaces.GPU(duration=_estimate_duration)
def generate_edit(
    prompt: str,
    source_video,
    num_inference_steps: int = 50,
    guidance_scale: float = 2.5,
    guidance_scale_visual: float = 1.5,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit a video following a text instruction.

    Args:
        prompt: Editing instruction text.
        source_video: Source video file to edit.
        num_inference_steps: Number of denoising steps.
        guidance_scale: Text CFG scale for editing.
        guidance_scale_visual: Visual CFG scale.
        seed: Random seed (-1 for random).
    """
    if source_video is None:
        return None, "Please upload a source video."

    image_extensions = (".jpg", ".jpeg", ".png", ".bmp", ".webp")
    is_source_image = str(source_video).lower().endswith(image_extensions)

    if is_source_image:
        source_image = Image.open(source_video).convert("RGB")
        native_width, native_height = source_image.size
        height, width = native_height, native_width
        source_image = resize_and_crop(source_image, height, width)
        source_tensor = IMAGE_TRANSFORM(source_image).to(model.device)
        num_frames = 1
        content = [
            {"type": "image", "image": source_image},
            {"type": "text", "text": f"Edit this image: {prompt}"},
        ]
    else:
        probe_reader = VideoReader(source_video, ctx=cpu(0))
        first_frame = probe_reader[0].asnumpy()
        native_height, native_width = first_frame.shape[:2]
        del probe_reader
        height, width = native_height, native_width

        pil_frames, source_tensor = load_video(
            source_video, target_height=height, target_width=width
        )
        source_tensor = source_tensor.to(model.device)
        num_source_frames = source_tensor.shape[1]
        num_frames = num_source_frames

        vlm_frames = temporal_downsample_for_vlm(pil_frames, config)
        content = [
            {"type": "video", "video": vlm_frames},
            {"type": "text", "text": f"Edit this video: {prompt}"},
        ]

    messages = [{"role": "user", "content": content}]
    inputs = _processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    ).to(model.device)

    if seed >= 0:
        generator = torch.Generator(device=model.device).manual_seed(seed)
    else:
        generator = None

    generate_kwargs = _build_generate_kwargs(
        inputs,
        negative_prompt=_get_negative_prompt(None),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=True),
        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
        source_pixel_values=[source_tensor],
        generator=generator,
    )

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        output = model.generate(**generate_kwargs)

    return _save_output(output, fps=24)


@spaces.GPU(duration=_estimate_duration)
def generate_ref_edit(
    prompt: str,
    source_video,
    ref_image,
    num_inference_steps: int = 50,
    guidance_scale: float = 2.5,
    guidance_scale_visual: float = 1.5,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit a source video with guidance from a reference image and text instruction.

    Args:
        prompt: Editing instruction text.
        source_video: Source video file to edit.
        ref_image: Reference image for guidance.
        num_inference_steps: Number of denoising steps.
        guidance_scale: Text CFG scale for editing.
        guidance_scale_visual: Visual CFG scale.
        seed: Random seed (-1 for random).
    """
    if source_video is None or ref_image is None:
        return None, "Please upload both a source video and a reference image."

    probe_reader = VideoReader(source_video, ctx=cpu(0))
    first_frame = probe_reader[0].asnumpy()
    height, width = first_frame.shape[0], first_frame.shape[1]
    del probe_reader

    source_pil_frames, source_video_tensor = load_video(
        source_video, target_height=height, target_width=width
    )
    source_video_tensor = source_video_tensor.to(model.device)
    num_source_frames = source_video_tensor.shape[1]
    num_frames = num_source_frames

    ref_pil_images, ref_tensors = load_images(
        [ref_image], target_height=height, target_width=width
    )
    ref_image_tensors = [t.to(model.device) for t in ref_tensors]

    vlm_frames = temporal_downsample_for_vlm(source_pil_frames, config)

    content = []
    for ref_pil in ref_pil_images:
        content.append({"type": "image", "image": ref_pil})
    content.append({"type": "video", "video": vlm_frames})
    content.append({"type": "text", "text": f"Edit this video with reference images: {prompt}"})

    messages = [{"role": "user", "content": content}]
    inputs = _processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    ).to(model.device)

    if seed >= 0:
        generator = torch.Generator(device=model.device).manual_seed(seed)
    else:
        generator = None

    generate_kwargs = _build_generate_kwargs(
        inputs,
        negative_prompt=_get_negative_prompt(None),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=True),
        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
        source_pixel_values=[source_video_tensor],
        ref_pixel_values=ref_image_tensors,
        generator=generator,
    )

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        output = model.generate(**generate_kwargs)

    return _save_output(output, fps=24)


@spaces.GPU(duration=_estimate_duration)
def generate_mi2v(
    prompt: str,
    ref_images,
    num_frames: int = 81,
    height: int = None,
    width: int = None,
    num_inference_steps: int = 50,
    guidance_scale: float = 5.0,
    guidance_scale_visual: float = 1.5,
    seed: int = 0,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a video conditioned on multiple reference images and a text prompt.

    Args:
        prompt: Text prompt describing the video (use @Image N to reference images).
        ref_images: List of reference image files.
        num_frames: Number of output video frames.
        height: Output video height in pixels.
        width: Output video width in pixels.
        num_inference_steps: Number of denoising steps.
        guidance_scale: Text CFG scale.
        guidance_scale_visual: Visual CFG scale.
        seed: Random seed (-1 for random).
    """
    if ref_images is None or len(ref_images) == 0:
        return None, "Please upload at least one reference image."

    # Handle single image or list
    if isinstance(ref_images, str):
        ref_paths = [ref_images]
    elif isinstance(ref_images, list):
        ref_paths = list(ref_images)
    else:
        ref_paths = [ref_images]

    ref_image = Image.open(ref_paths[0]).convert("RGB")
    native_width, native_height = ref_image.size
    if height is None or height == 0:
        height = native_height
    if width is None or width == 0:
        width = native_width
    num_frames = num_frames if num_frames and num_frames > 0 else 81

    ref_pil_images, ref_tensors = load_images(
        ref_paths, target_height=height, target_width=width
    )
    ref_image_tensors = [t.to(model.device) for t in ref_tensors]

    content = []
    for ref_pil in ref_pil_images:
        content.append({"type": "image", "image": ref_pil})
    content.append({"type": "text", "text": f"Generate a video with reference images: {prompt}"})

    messages = [{"role": "user", "content": content}]
    inputs = _processor.apply_chat_template(
        messages, tokenize=True, add_generation_prompt=True,
        return_dict=True, return_tensors="pt",
    ).to(model.device)

    if seed >= 0:
        generator = torch.Generator(device=model.device).manual_seed(seed)
    else:
        generator = None

    generate_kwargs = _build_generate_kwargs(
        inputs,
        negative_prompt=_get_negative_prompt(None),
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=num_inference_steps,
        guidance_scale=_get_guidance_scale(guidance_scale, is_edit=False),
        guidance_scale_visual=_get_guidance_scale_visual(guidance_scale_visual),
        ref_pixel_values=ref_image_tensors,
        generator=generator,
    )

    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
        output = model.generate(**generate_kwargs)

    return _save_output(output, fps=24)


def _save_output(output_np: np.ndarray, fps: int = 24) -> str:
    """Save generation output as image (T=1) or video (T>1) to a temp file."""
    num_frames = output_np.shape[0]

    if num_frames == 1:
        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        tmp.close()
        frame = (output_np[0] * 255).clip(0, 255).astype(np.uint8)
        Image.fromarray(frame).save(tmp.name)
        return tmp.name

    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    tmp.close()
    export_to_video(output_np, tmp.name, fps=fps)
    return tmp.name


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

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:
    gr.Markdown("# 🎬 LoomVideo: Unified Multimodal Video Generation and Editing")
    gr.Markdown(
        "LoomVideo is a compact 5B-parameter unified architecture for video generation and editing. "
        "It supports text-to-video, instruction editing, reference-guided editing, and multi-image-to-video. "
        "[Paper](https://arxiv.org/abs/2606.06042) · [Model](https://huggingface.co/MSALab/LoomVideo)"
    )

    with gr.Tabs():
        # --- Text-to-Video tab ---
        with gr.TabItem("Text-to-Video"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown("### Generate a video from a text prompt")
                with gr.Row():
                    t2v_prompt = gr.Textbox(
                        label="Prompt", show_label=False,
                        placeholder="Describe the video you want to generate...",
                        container=False, scale=4, lines=3,
                    )
                    t2v_btn = gr.Button("Generate", variant="primary", scale=1)
                t2v_output = gr.Video(label="Generated Video")

                with gr.Accordion("Advanced settings", open=False):
                    t2v_frames = gr.Slider(1, 121, value=81, step=1, label="Number of frames (1 = image)")
                    t2v_height = gr.Slider(256, 1080, value=480, step=8, label="Height")
                    t2v_width = gr.Slider(256, 1920, value=832, step=8, label="Width")
                    t2v_steps = gr.Slider(1, 100, value=50, step=1, label="Inference steps")
                    t2v_cfg = gr.Slider(1.0, 15.0, value=5.0, step=0.5, label="Guidance scale")
                    t2v_cfg_visual = gr.Slider(1.0, 5.0, value=1.5, step=0.1, label="Visual guidance scale")
                    t2v_seed = gr.Number(label="Seed", value=0, precision=0, info="-1 for random")

                gr.Examples(
                    examples=[
                        ["Snow rocky mountains peaks canyon. Snow blanketed rocky mountains surround and shadow deep canyons. The canyons twist and bend through the high elevated mountain peaks."],
                        ["Vampire makeup face of beautiful girl, red contact lenses."],
                    ],
                    inputs=[t2v_prompt],
                    outputs=[t2v_output],
                    fn=generate_t2v,
                    cache_examples=False,
                    run_on_click=True,
                )

            t2v_btn.click(
                fn=generate_t2v,
                inputs=[t2v_prompt, t2v_frames, t2v_height, t2v_width, t2v_steps, t2v_cfg, t2v_cfg_visual, t2v_seed],
                outputs=[t2v_output],
                api_name="t2v",
            )

        # --- Instruction Editing tab ---
        with gr.TabItem("Instruction Editing"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown("### Edit a video with a text instruction")
                with gr.Row():
                    edit_prompt = gr.Textbox(
                        label="Instruction", show_label=False,
                        placeholder="Describe the edit to apply...",
                        container=False, scale=4, lines=3,
                    )
                    edit_btn = gr.Button("Edit", variant="primary", scale=1)
                with gr.Row():
                    edit_source = gr.Video(label="Source video")
                    edit_output = gr.Video(label="Edited video")

                with gr.Accordion("Advanced settings", open=False):
                    edit_steps = gr.Slider(1, 100, value=50, step=1, label="Inference steps")
                    edit_cfg = gr.Slider(1.0, 15.0, value=2.5, step=0.5, label="Guidance scale")
                    edit_cfg_visual = gr.Slider(1.0, 5.0, value=1.5, step=0.1, label="Visual guidance scale")
                    edit_seed = gr.Number(label="Seed", value=0, precision=0, info="-1 for random")

                gr.Examples(
                    examples=[
                        ["Apply the Impressionist aesthetic to this video, ensuring seamless temporal consistency across all frames. The result should emulate the fluid brushstroke techniques and atmospheric focus of 19th-century Impressionist art, with each frame retaining the original motion, character actions, and camera movements.", "assets/edit_input.mp4"],
                    ],
                    inputs=[edit_prompt, edit_source],
                    outputs=[edit_output],
                    fn=generate_edit,
                    cache_examples=False,
                    run_on_click=True,
                )

            edit_btn.click(
                fn=generate_edit,
                inputs=[edit_prompt, edit_source, edit_steps, edit_cfg, edit_cfg_visual, edit_seed],
                outputs=[edit_output],
                api_name="edit",
            )

        # --- Reference-Guided Editing tab ---
        with gr.TabItem("Reference-Guided Editing"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown("### Edit a video using a reference image and text instruction")
                with gr.Row():
                    ref_edit_prompt = gr.Textbox(
                        label="Instruction", show_label=False,
                        placeholder="Describe the edit with reference image guidance...",
                        container=False, scale=4, lines=3,
                    )
                    ref_edit_btn = gr.Button("Edit", variant="primary", scale=1)
                with gr.Row():
                    ref_edit_source = gr.Video(label="Source video")
                    ref_edit_ref = gr.Image(label="Reference image", type="filepath")
                    ref_edit_output = gr.Video(label="Edited video")

                with gr.Accordion("Advanced settings", open=False):
                    ref_edit_steps = gr.Slider(1, 100, value=50, step=1, label="Inference steps")
                    ref_edit_cfg = gr.Slider(1.0, 15.0, value=2.5, step=0.5, label="Guidance scale")
                    ref_edit_cfg_visual = gr.Slider(1.0, 5.0, value=1.5, step=0.1, label="Visual guidance scale")
                    ref_edit_seed = gr.Number(label="Seed", value=0, precision=0, info="-1 for random")

                gr.Examples(
                    examples=[
                        ["Replace the green t-shirt of the man with the suit in the image.", "assets/ref_edit_input.mp4", "assets/ref_edit_reference.jpg"],
                    ],
                    inputs=[ref_edit_prompt, ref_edit_source, ref_edit_ref],
                    outputs=[ref_edit_output],
                    fn=generate_ref_edit,
                    cache_examples=False,
                    run_on_click=True,
                )

            ref_edit_btn.click(
                fn=generate_ref_edit,
                inputs=[ref_edit_prompt, ref_edit_source, ref_edit_ref, ref_edit_steps, ref_edit_cfg, ref_edit_cfg_visual, ref_edit_seed],
                outputs=[ref_edit_output],
                api_name="ref_edit",
            )

        # --- Multi-Image-to-Video tab ---
        with gr.TabItem("Multi-Image-to-Video"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown("### Generate a video from multiple reference images. Use `@Image N` in the prompt to reference specific images.")
                with gr.Row():
                    mi2v_prompt = gr.Textbox(
                        label="Prompt", show_label=False,
                        placeholder="Describe the video using @Image N references...",
                        container=False, scale=4, lines=3,
                    )
                    mi2v_btn = gr.Button("Generate", variant="primary", scale=1)
                with gr.Row():
                    mi2v_refs = gr.File(label="Reference images", file_count="multiple", file_types=["image"])
                    mi2v_output = gr.Video(label="Generated video")

                with gr.Accordion("Advanced settings", open=False):
                    mi2v_frames = gr.Slider(1, 121, value=81, step=1, label="Number of frames")
                    mi2v_height = gr.Slider(0, 1080, value=0, step=8, label="Height (0 = auto from reference)")
                    mi2v_width = gr.Slider(0, 1920, value=0, step=8, label="Width (0 = auto from reference)")
                    mi2v_steps = gr.Slider(1, 100, value=50, step=1, label="Inference steps")
                    mi2v_cfg = gr.Slider(1.0, 15.0, value=5.0, step=0.5, label="Guidance scale")
                    mi2v_cfg_visual = gr.Slider(1.0, 5.0, value=1.5, step=0.1, label="Visual guidance scale")
                    mi2v_seed = gr.Number(label="Seed", value=0, precision=0, info="-1 for random")

            mi2v_btn.click(
                fn=generate_mi2v,
                inputs=[mi2v_prompt, mi2v_refs, mi2v_frames, mi2v_height, mi2v_width, mi2v_steps, mi2v_cfg, mi2v_cfg_visual, mi2v_seed],
                outputs=[mi2v_output],
                api_name="mi2v",
            )

demo.launch(mcp_server=True)