"""
ID-V2V — identity-preserving video-to-video (Eyeline Labs).

Faithful port of github.com/Eyeline-Labs/ID-V2V's reference inference path to a
single-GPU ZeroGPU Space:

  1. SAM3 promptable concept segmentation over the source clip ("person"),
     Secret Panda mask cleanup (per object + on the union).
  2. foreground-on-gray condition video  ->  the single VACE control stream.
  3. Wan2.1 I2V-14B DiT + VACE ControlNet (both from Eyeline-Labs/ID-V2V's
     finetuned idv2v.pth), with the stylized first frame as the I2V anchor and
     as the SVI anti-drift reference pad (ref_pad_num = -1).

Pipeline code is the authors' own diffsynth fork (vendored under ./diffsynth).
"""

import os

# Must precede any torch import.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  (before torch)

import tempfile  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from PIL import Image  # noqa: E402

from diffsynth.pipelines.wan_video_new_multiVace_svi import (  # noqa: E402
    ModelConfig,
    WanVideoPipeline,
)
from idv2v_lib import (  # noqa: E402
    DEFAULT_NEGATIVE_PROMPT,
    center_crop_and_resize,
    foreground_on_gray,
    load_finetuned_dit_vace,
    load_source_frames,
    read_video_rgb,
    run_sam3_union_masks,
    save_video,
)

# --------------------------------------------------------------------------- #
# Fixed generation geometry (the repo's own scripts/infer.sh CLI default res).
# --------------------------------------------------------------------------- #
WIDTH, HEIGHT = 832, 480
REF_PAD_NUM = -1          # -1 = full SVI anti-drift padding (repo default)

DEF_FRAMES = 33
DEF_STRIDE = 2
DEF_STEPS = 20
DEF_CFG = 5.0
DEF_VACE = 1.0
DEF_SEED = 123
DEF_SAM = "person"

DTYPE = torch.bfloat16

# --------------------------------------------------------------------------- #
# Weights
# --------------------------------------------------------------------------- #
print("[idv2v] fetching weights...", flush=True)
_t0 = time.time()

T5_PATH = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "models_t5_umt5-xxl-enc-bf16.pth")
VAE_PATH = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", "Wan2.1_VAE.pth")
for _f in ("special_tokens_map.json", "spiece.model", "tokenizer.json", "tokenizer_config.json"):
    _tok = hf_hub_download("Wan-AI/Wan2.1-T2V-14B", f"google/umt5-xxl/{_f}")
TOKENIZER_DIR = os.path.dirname(_tok)
CLIP_PATH = hf_hub_download(
    "Wan-AI/Wan2.1-I2V-14B-480P",
    "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
)
CKPT_PATH = hf_hub_download("Eyeline-Labs/ID-V2V", "idv2v.pth")
print(f"[idv2v] weights on disk in {time.time() - _t0:.0f}s", flush=True)

# --------------------------------------------------------------------------- #
# SAM3 (gated repo -> needs the HF_TOKEN space secret)
# --------------------------------------------------------------------------- #
from transformers import Sam3VideoModel, Sam3VideoProcessor  # noqa: E402

try:  # upstream typo: initializer_range annotated int but defaults to 0.02
    from transformers import Sam3TrackerVideoConfig

    Sam3TrackerVideoConfig.__dataclass_fields__["initializer_range"].type = float
except Exception:  # pragma: no cover
    pass

print("[idv2v] loading SAM3...", flush=True)
sam3_model = Sam3VideoModel.from_pretrained("facebook/sam3", dtype=DTYPE).eval().to("cuda")
sam3_processor = Sam3VideoProcessor.from_pretrained("facebook/sam3")

# --------------------------------------------------------------------------- #
# Wan base models (T5 / VAE / CLIP) then the finetuned DiT + VACE
# --------------------------------------------------------------------------- #
print("[idv2v] building WanVideoPipeline...", flush=True)
pipe = WanVideoPipeline.from_pretrained(
    torch_dtype=DTYPE,
    device="cuda",
    model_configs=[
        ModelConfig(path=T5_PATH, offload_device="cpu", offload_dtype=DTYPE),
        ModelConfig(path=VAE_PATH, offload_device="cpu", offload_dtype=DTYPE),
        ModelConfig(path=CLIP_PATH, offload_device="cpu", offload_dtype=DTYPE),
    ],
    tokenizer_config=ModelConfig(path=TOKENIZER_DIR),
    skip_download=True,
    redirect_common_files=False,
)

print("[idv2v] loading finetuned DiT + VACE from idv2v.pth...", flush=True)
load_finetuned_dit_vace(pipe, CKPT_PATH, torch_dtype=DTYPE, delete_checkpoint_after=True)
assert pipe.dit is not None and pipe.dit.has_image_input

# Everything resident on the GPU — no vram-management / CPU offload.
for _name in ("text_encoder", "vae", "image_encoder", "dit", "vace"):
    _m = getattr(pipe, _name, None)
    if _m is not None:
        _m.to("cuda")
        _m.eval()
pipe.device = "cuda"
print("[idv2v] pipeline ready", flush=True)


# --------------------------------------------------------------------------- #
# Duration estimate
# --------------------------------------------------------------------------- #
def _estimate_duration(
    source_video,
    stylized_first_frame,
    prompt,
    num_frames=DEF_FRAMES,
    frame_stride=DEF_STRIDE,
    num_inference_steps=DEF_STEPS,
    cfg_scale=DEF_CFG,
    vace_scale=DEF_VACE,
    seed=DEF_SEED,
    sam_prompt=DEF_SAM,
    progress=None,
):
    n = int(num_frames)
    latent_f = (n - 1) // 4 + 1
    tokens = latent_f * (HEIGHT // 16) * (WIDTH // 16)
    per_forward = 4.0e-4 * tokens
    forwards = int(num_inference_steps) * (2 if float(cfg_scale) != 1.0 else 1)
    overhead = 45.0 + 0.7 * n  # SAM3 + T5 + VAE encode/decode + mp4 write
    return int(min(1500, overhead + per_forward * forwards))


# --------------------------------------------------------------------------- #
# Inference
# --------------------------------------------------------------------------- #
@spaces.GPU(duration=_estimate_duration, size="xlarge")
def generate(
    source_video,
    stylized_first_frame,
    prompt,
    num_frames=DEF_FRAMES,
    frame_stride=DEF_STRIDE,
    num_inference_steps=DEF_STEPS,
    cfg_scale=DEF_CFG,
    vace_scale=DEF_VACE,
    seed=DEF_SEED,
    sam_prompt=DEF_SAM,
    progress=gr.Progress(track_tqdm=True),
):
    """Restyle or relight a video while preserving the identity of the people in it.

    Args:
        source_video (str): Path to the source video; supplies the motion.
        stylized_first_frame (PIL.Image.Image): The target look for frame 0 (e.g. an
            edited/restyled/relit version of the source's first frame).
        prompt (str): Text description of the desired output video.
        num_frames (int): Frames to generate (17, 33 or 49).
        frame_stride (int): Take every Nth source frame; output fps is divided to match.
        num_inference_steps (int): Denoising steps.
        cfg_scale (float): Classifier-free guidance scale (1.0 disables CFG, ~2x faster).
        vace_scale (float): Strength of the VACE motion/identity control stream.
        seed (int): Random seed.
        sam_prompt (str): SAM3 concept prompt used to segment the subject(s).

    Returns:
        tuple[str, str]: paths to the generated mp4 and to the VACE condition mp4.
    """
    if not source_video:
        raise gr.Error("Please provide a source video.")
    if stylized_first_frame is None:
        raise gr.Error("Please provide a stylized first frame.")
    if not (prompt or "").strip():
        raise gr.Error("Please provide a prompt describing the target video.")

    num_frames = int(num_frames)
    frame_stride = max(1, int(frame_stride))
    num_inference_steps = int(num_inference_steps)
    seed = int(seed)

    t_all = time.perf_counter()

    # ---- inputs -> 832x480 -------------------------------------------------
    src_frames, src_fps = load_source_frames(
        source_video, WIDTH, HEIGHT, num_frames, frame_stride
    )
    out_fps = max(1.0, src_fps / frame_stride)

    if isinstance(stylized_first_frame, str):
        ext = os.path.splitext(stylized_first_frame)[1].lower()
        if ext in {".mp4", ".mov", ".webm", ".mkv", ".avi"}:
            first = read_video_rgb(stylized_first_frame)[0][0]
        else:
            first = Image.open(stylized_first_frame).convert("RGB")
    else:
        first = stylized_first_frame.convert("RGB")
    input_image = center_crop_and_resize(first, WIDTH, HEIGHT)

    # ---- SAM3 -> foreground-on-gray VACE condition -------------------------
    t_seg = time.perf_counter()
    masks = run_sam3_union_masks(
        sam3_model,
        sam3_processor,
        src_frames,
        (sam_prompt or DEF_SAM).strip() or DEF_SAM,
        device="cuda",
        dtype=DTYPE,
    )
    condition = foreground_on_gray(src_frames, masks)
    covered = sum(1 for m in masks if m.any())
    print(f"[idv2v] SAM3 + cleanup {time.perf_counter() - t_seg:.1f}s "
          f"({covered}/{len(masks)} frames with a mask)", flush=True)

    white = Image.new("RGB", (WIDTH, HEIGHT), (255, 255, 255))
    mask_frames = [white] * num_frames

    # ---- generate ----------------------------------------------------------
    t_gen = time.perf_counter()
    frames = pipe(
        prompt=prompt.strip(),
        negative_prompt=DEFAULT_NEGATIVE_PROMPT,
        input_image=input_image,          # I2V anchor (frame 0)
        random_ref_frame=input_image,     # SVI anti-drift reference pad
        ref_pad_num=REF_PAD_NUM,
        vace_video=[condition],           # the single VACE control stream
        vace_video_mask=[mask_frames],    # fully reactive
        use_multi_control_vace=True,
        vace_scale=float(vace_scale),
        seed=seed,
        height=HEIGHT,
        width=WIDTH,
        num_frames=num_frames,
        cfg_scale=float(cfg_scale),
        num_inference_steps=num_inference_steps,
        tiled=False,
    )
    gen_s = time.perf_counter() - t_gen

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    cond_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    save_video(frames[:num_frames], out_path, out_fps)
    save_video(condition, cond_path, out_fps)

    print(f"[idv2v] denoise {gen_s:.1f}s | total {time.perf_counter() - t_all:.1f}s "
          f"({num_frames}f, {num_inference_steps} steps, cfg {cfg_scale})", flush=True)
    return out_path, cond_path


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

EXAMPLES = [
    [
        "examples/man_dancing/source.mp4",
        "examples/man_dancing/stylized_first_frame.png",
        "Against a dramatic, high-contrast industrial cityscape at dusk, a young man with "
        "dark hair and a beard, in a black hoodie with red panels, stands on a concrete "
        "ledge. Illuminated by warm orange glows from below and cool ambient light, he "
        "turns from off-camera to face the lens, his light eyes intense. Suddenly, he "
        "flexes into a double-bicep pose, shouting, his body bouncing with raw energy "
        "amidst billowing, orange-lit smoke.",
    ],
    [
        "examples/two_sitting_woman/source.mp4",
        "examples/two_sitting_woman/stylized_first_frame.png",
        "Bathed in cool, ethereal light within an overgrown, abandoned greenhouse, two "
        "young women share a joyful moment. The woman on the left, with dark, voluminous "
        "curls and a crisp white shirt, sits relaxed, initially beaming at the camera. Her "
        "companion, with reddish-brown curls and a wide, mid-laugh smile, gestures gently. "
        "The light filters through broken glass and lush vines, casting dappled shadows.",
    ],
    [
        "examples/music_band/source.mp4",
        "examples/music_band/stylized_first_frame.png",
        "A light-skinned woman with dreadlocks and a black beanie sings passionately into a "
        "microphone, her tattooed arms gesturing expressively as she sways with a joyful "
        "smile. Behind her, a bearded man with a topknot intensely plays a hand drum. The "
        "intimate home studio glows with dramatic, contrasting magenta and cool blue light "
        "illuminating the foreground figures, while warm amber hues emanate from sheer "
        "curtains and vintage amps in the background.",
    ],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# ID-V2V — identity-preserving video-to-video

Restyle / relight a video while keeping the people in it recognisable. Give it a
**source video** (whose motion is followed), a **stylized first frame** (which
defines the target look and identity), and a **prompt**.

SAM3 segments the people, the foreground-on-gray result drives a VACE ControlNet,
and the stylized frame anchors the Wan2.1 I2V-14B DiT (SVI anti-drift padding).

[model](https://huggingface.co/Eyeline-Labs/ID-V2V) ·
[code](https://github.com/Eyeline-Labs/ID-V2V) ·
[paper](https://huggingface.co/papers/2607.22830)
            """
        )

        with gr.Row():
            with gr.Column():
                source_video = gr.Video(label="Source video (motion reference)")
                stylized_first_frame = gr.Image(
                    label="Stylized first frame (target look + identity)",
                    type="pil",
                    height=240,
                )
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=4,
                    placeholder="Describe the target video: subjects, action, lighting, setting…",
                )
                run = gr.Button("Generate", variant="primary")
            with gr.Column():
                out_video = gr.Video(label="Generated video", autoplay=True)
                cond_video = gr.Video(label="VACE condition (foreground-on-gray)")

        with gr.Accordion("Advanced settings", open=False):
            with gr.Row():
                num_frames = gr.Radio(
                    [17, 33, 49],
                    value=DEF_FRAMES,
                    label="Frames to generate",
                    info="Must be 4k+1 (Wan latent stride). More frames = longer clip and longer wait.",
                )
                frame_stride = gr.Slider(
                    1, 3, value=DEF_STRIDE, step=1,
                    label="Source frame stride",
                    info="2 samples every other source frame, so a short clip still covers "
                         "the whole action; output fps is divided to match. 1 = the repo's "
                         "native behaviour.",
                )
            with gr.Row():
                num_inference_steps = gr.Slider(
                    10, 30, value=DEF_STEPS, step=1, label="Inference steps",
                    info="The paper uses 30.",
                )
                cfg_scale = gr.Slider(
                    1.0, 8.0, value=DEF_CFG, step=0.5, label="CFG scale",
                    info="1.0 skips the negative pass and is ~2x faster (lower quality).",
                )
            with gr.Row():
                vace_scale = gr.Slider(
                    0.0, 1.5, value=DEF_VACE, step=0.05, label="VACE scale",
                    info="How strongly the segmented source drives motion.",
                )
                seed = gr.Number(value=DEF_SEED, precision=0, label="Seed")
            sam_prompt = gr.Textbox(
                value=DEF_SAM, label="SAM3 segmentation prompt",
                info="What to keep as foreground: 'person' (default), 'head', 'dog', …",
            )

        gr.Markdown(
            "Rendered at 832×480 (the repo's CLI default). The released checkpoint is a "
            "720p model, so 1280×720 is sharper but ~2.5× slower than fits in a single "
            "ZeroGPU slot; the authors' multi-clip chaining for longer videos is likewise "
            "out of scope here."
        )

        gr.Examples(
            examples=EXAMPLES,
            inputs=[source_video, stylized_first_frame, prompt],
            outputs=[out_video, cond_video],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        fn=generate,
        inputs=[
            source_video, stylized_first_frame, prompt,
            num_frames, frame_stride, num_inference_steps,
            cfg_scale, vace_scale, seed, sam_prompt,
        ],
        outputs=[out_video, cond_video],
    )

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