"""VidDub: two-step replacement-aware video dubbing.

Upload a source video (a talking head) and target speech audio; VidDub
regenerates the lower face so the lips match the new audio in 2 denoising
steps, pasting the result back into the original video.

Reference implementation: https://github.com/LAW1223/VidDub (scripts/infer.py).
Model: AIPeanutman/VidDub on top of Wan2.1-VACE-1.3B + whisper-small.
"""
import json
import math
import os
import sys
import tempfile
import time
from pathlib import Path

import spaces  # MUST be imported before torch / anything CUDA-touching

# The Space repo bundles the VidDub source (viddub/ + data/) at /home/user/app.
APP_ROOT = Path(__file__).resolve().parent
for p in (str(APP_ROOT),):
    if p not in sys.path:
        sys.path.insert(0, p)

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

import decord
import gradio as gr
import librosa
import numpy as np
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from safetensors.torch import load_file

from viddub.causal.model import CausalConfig, CausalWanDub
from viddub.causal.runtime import autocast
from viddub.causal.sampling import rollout, schedule
from viddub.foundation.pipelines.wan_video import ModelConfig, WanVideoPipeline
from viddub.full_context import runtime as full_context_runtime
from viddub.full_context.audio import DubAudioAdapter, WhisperFeatures
from viddub.full_context.inference import dub_video
from viddub.full_context.runtime import FPS, SR, ZeroLatents, build_inputs, write_mp4

from data.s02_make_clips import NUM_FRAMES

DEVICE = "cuda"
VIDDUB_REPO = "AIPeanutman/VidDub"
BASE_REPO = "Wan-AI/Wan2.1-VACE-1.3B"
WHISPER_REPO = "openai/whisper-small"

MAX_FRAMES = 242  # 3 windows x 81 frames (stride 64) + tail; ~10s @ 24fps


# ---------------------------------------------------------------- weight prep
def prepare_weights():
    """Download every checkpoint the runtime needs into one local root."""
    root = Path(os.environ.get("VIDDUB_WEIGHT_ROOT", APP_ROOT / "weights"))
    root.mkdir(parents=True, exist_ok=True)
    vace_dir = root / "Wan2.1-VACE-1.3B"
    need = {
        "diffusion": vace_dir / "diffusion_pytorch_model.safetensors",
        "vae": vace_dir / "Wan2.1_VAE.pth",
    }
    if not all(p.is_file() for p in need.values()):
        snapshot_download(
            BASE_REPO, repo_type="model",
            allow_patterns=["diffusion_pytorch_model.safetensors", "Wan2.1_VAE.pth"],
            local_dir=str(vace_dir))
    whisper_dir = root / "whisper-small"
    if not (whisper_dir / "model.safetensors").is_file():
        snapshot_download(
            WHISPER_REPO, repo_type="model",
            allow_patterns=["*.json", "*.txt", "model.safetensors", "vocab.json",
                            "merges.txt", "preprocessor_config.json",
                            "tokenizer_config.json", "normalizer.json",
                            "special_tokens_map.json", "added_tokens.json"],
            local_dir=str(whisper_dir))
    release_dir = root / "VidDub"
    if not (release_dir / "config.json").is_file():
        snapshot_download(
            VIDDUB_REPO, repo_type="model",
            allow_patterns=["*.json", "dit/*", "vace/*", "audio_adapter/*"],
            local_dir=str(release_dir))
    full_context_runtime.WEIGHTS = root
    full_context_runtime.VACE_DIR = vace_dir
    return release_dir


def load_release(model_dir: Path):
    """Load the VidDub release exactly as scripts/infer.py does."""
    config = json.loads((model_dir / "config.json").read_text())
    if config.get("model_type") != "viddub":
        raise ValueError("The model directory is not a VidDub release")
    if float(config["inference_audio_cfg"]) != 1.0:
        raise ValueError("Published VidDub checkpoints require inference audio CFG 1")
    if config["stage"] != "bidirectional_sft" and (
        config["stage"] != "dmd"
        or config.get("paper_selected") is not True
        or config.get("causal") is not False
        or config.get("source_run") != "FULL2_DMD_R4"):
        raise ValueError("The released two-step path requires FULL2_DMD_R4")

    # Same T5-skip path as viddub/causal/runtime.py: with a precomputed
    # empty-prompt context, the 11 GB T5 encoder is never loaded.
    empty_context = torch.load(os.environ["VIDDUB_EMPTY_CONTEXT"],
                               map_location="cpu", weights_only=True)
    if not (torch.is_tensor(empty_context) and empty_context.ndim == 2
            and empty_context.shape[1] == 4096 and 1 <= empty_context.shape[0] <= 512
            and torch.isfinite(empty_context).all()):
        raise ValueError("assets/empty_context.pt must be a finite [L,4096] tensor")
    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16, device=DEVICE,
        model_configs=[ModelConfig(path=str(full_context_runtime.VACE_DIR
                                            / "diffusion_pytorch_model.safetensors")),
                       ModelConfig(path=str(full_context_runtime.VACE_DIR
                                            / "Wan2.1_VAE.pth"))],
        tokenizer_config=None)
    for m in (pipe.dit, pipe.vace, pipe.vae):
        m.requires_grad_(False).eval()
    pipe._wan_dub_empty_context = empty_context.to(pipe.device, pipe.torch_dtype)
    whisper = WhisperFeatures(str(full_context_runtime.WEIGHTS
                                  / "whisper-small")).to(DEVICE, torch.bfloat16)
    paths = {name: model_dir / rel for name, rel in config["components"].items()}
    missing = [str(p) for p in paths.values() if not p.is_file()]
    if missing:
        raise FileNotFoundError(f"Missing model files: {missing}")
    # The VidDub release ships bf16 DiT/VACE and fp32 audio adapter; the
    # reference loader loads them straight onto the device with strict=True.
    pipe.dit.load_state_dict(load_file(str(paths["dit"])), strict=True)
    pipe.vace.load_state_dict(load_file(str(paths["vace"])), strict=True)
    adapter = DubAudioAdapter().to(DEVICE)
    adapter.load_state_dict(load_file(str(paths["audio_adapter"])), strict=True)
    adapter.eval()
    return config, pipe, whisper, adapter


def two_step_sampler(model, config):
    """Two-step flow_euler sampler with a full-context (bidirectional) rollout."""
    @torch.no_grad()
    def sample(pipe, whisper, adapter, zero, clip, audio, steps=None,
               audio_cfg=1.0, seed=0, fps=FPS):
        trained_steps = int(config["student_steps"])
        num_frames = len(clip["frames"])
        video = full_context_runtime.to_video_tensor(clip["frames"], pipe)
        keep = torch.from_numpy(~clip["mask"]).to(video)[None, None, None]
        latents = {
            "x0": pipe.vae.encode(video, device=pipe.device)[0],
            "inactive": pipe.vae.encode(video * keep, device=pipe.device)[0],
            "ref": pipe.vae.encode(
                full_context_runtime.to_video_tensor(clip["ref"][None], pipe),
                device=pipe.device)[0],
        }
        x0, vace_context, context = build_inputs(
            latents, clip["mask"], num_frames, zero,
            drop_ref=int(config["num_reference_latents"]) == 0)
        whisper_hs = whisper([audio], pipe.device, torch.bfloat16)
        model.config = CausalConfig(
            int(config["latent_chunk_size"]),
            int(config["num_reference_latents"]), float(fps))
        conditions = {
            "context": context.to(torch.bfloat16),
            "vace_context": vace_context.to(torch.bfloat16),
            "whisper_hs": whisper_hs.to(torch.bfloat16),
            "num_frames": num_frames,
        }
        noise = torch.randn(
            x0.shape, generator=torch.Generator(device=x0.device).manual_seed(seed),
            device=x0.device, dtype=torch.bfloat16)
        with autocast(x0.device, "bf16"):
            output = rollout(
                model, x0.to(torch.bfloat16), conditions,
                schedule(trained_steps, float(config["shift"]), x0.device),
                noise=noise, use_cache=bool(config.get("causal", False)),
                audio_cfg=1.0)
        video = pipe.vae.decode(
            output[:, :, int(config["num_reference_latents"]):].to(pipe.torch_dtype),
            device=pipe.device)
        return ((video[0].permute(1, 2, 3, 0).float().clamp(-1, 1) + 1)
                * 127.5).round().byte().cpu().numpy()

    return sample


# ------------------------------------------------------------ module-scope load
# Dubbing is audio-driven: the only text condition is the empty prompt. The
# bundled assets/empty_context.pt holds the precomputed T5 embedding of "",
# so the 11 GB umt5-xxl encoder is never downloaded or loaded (this is the
# repo's own VIDDUB_SKIP_TEXT_ENCODER path, see viddub/causal/runtime.py).
os.environ["VIDDUB_SKIP_TEXT_ENCODER"] = "1"
os.environ["VIDDUB_EMPTY_CONTEXT"] = str(APP_ROOT / "assets" / "empty_context.pt")

print("Downloading weights ...", flush=True)
RELEASE_DIR = prepare_weights()
print("Loading VidDub runtime ...", flush=True)
CONFIG, PIPE, WHISPER, ADAPTER = load_release(RELEASE_DIR)
ZERO = ZeroLatents(PIPE)
INFERENCE_MODEL = CausalWanDub(
    PIPE.dit, PIPE.vace, ADAPTER,
    CausalConfig(int(CONFIG["latent_chunk_size"]),
                 int(CONFIG["num_reference_latents"]), float(CONFIG["fps"])),
    causal=False).eval().requires_grad_(False)
SAMPLE_FN = two_step_sampler(INFERENCE_MODEL, CONFIG)
RUNTIME = {"models": (PIPE, WHISPER, ADAPTER, ZERO), "config": CONFIG,
           "steps": int(CONFIG["student_steps"]), "sample_fn": SAMPLE_FN}
print("VidDub runtime ready.", flush=True)


# ----------------------------------------------------------------- inference
def _probe_duration(*args, **kwargs) -> float:
    """Estimate the GPU duration from the input video's frame count.

    ZeroGPU passes the decorated function's arguments to a callable duration
    estimator; the video path is the first one.
    """
    video_path = args[0] if args else kwargs.get("video")
    try:
        n = len(decord.VideoReader(video_path, num_threads=1))
    except Exception:
        n = NUM_FRAMES
    n = max(n, NUM_FRAMES)
    n_windows = max(1, math.ceil((n - NUM_FRAMES) / 64) + 1)
    # Measured on the ZeroGPU RTX PRO 6000: 113 frames (2 windows) took 53.2 s
    # end-to-end. Linear fit ~15 s fixed (preprocess + VAE + mux) + ~20 s per
    # window; ~20% headroom on top so max-length clips are not killed mid-run.
    return min(420, 15 + 25 * n_windows)


@spaces.GPU(duration=_probe_duration)
def dub(video: str, audio: str, seed: int = 0, harmonize: bool = True,
        feather: int = 15, max_frames: int = MAX_FRAMES,
        progress=gr.Progress(track_tqdm=False)):
    """Re-dub a talking-head video with new speech audio.

    Args:
        video: source video with a mostly front-facing face (mp4/webm).
        audio: target speech (any format; resampled to 16 kHz mono).
        seed: noise seed for the two-step sampler.
        harmonize: apply the paper's lower-face color harmonization.
        feather: mask feather radius in crop pixels.
        max_frames: cap on processed frames (longer videos are truncated).

    Returns:
        The dubbed video with the original resolution/fps and the new speech.
    """
    if not video or not audio:
        raise gr.Error("Please provide both a source video and target audio.")
    t0 = time.perf_counter()
    video = str(video)
    audio = str(audio)
    reader = decord.VideoReader(video, num_threads=2)
    fps = float(reader.get_avg_fps())
    frames = reader[:].asnumpy()
    if len(frames) > max_frames:
        frames = frames[:max_frames]
    if len(frames) < NUM_FRAMES:
        raise gr.Error(
            f"Source video needs at least {NUM_FRAMES} frames "
            f"({NUM_FRAMES / FPS:.1f} s); got {len(frames)}.")
    waveform = librosa.load(audio, sr=SR, mono=True)[0]
    audio_samples = round(len(frames) / fps * SR)
    waveform = np.pad(
        waveform[:audio_samples], (0, max(0, audio_samples - len(waveform))))
    progress(0.1, desc="Dubbing ...")
    generated = dub_video(
        RUNTIME["models"], frames, waveform,
        steps=RUNTIME["steps"], audio_cfg=1.0, feather=feather, seed=seed,
        fps=fps, overlap=17, mode="auto", harmonize=harmonize,
        sample_fn=RUNTIME["sample_fn"],
        log=lambda m: print(m, flush=True))
    out_path = Path(tempfile.mkdtemp()) / "dubbed.mp4"
    write_mp4(out_path, generated, waveform, fps=fps)
    print(f"dub finished in {time.perf_counter() - t0:.1f}s", flush=True)
    return str(out_path)


# ----------------------------------------------------------------------- UI
EXAMPLES = [
    ["examples/man_laughing_video.mp4", "examples/man_laughing_audio.wav", 0],
    ["examples/woman_laughing_studio_video.mp4", "examples/woman_laughing_studio_audio.wav", 0],
    ["examples/man_smiling_studio_video.mp4", "examples/man_smiling_studio_audio.wav", 0],
]

with gr.Blocks(theme=gr.themes.Citrus(), title="VidDub") as demo:
    gr.Markdown(
        """
        # 🎬 VidDub — Replacement-Aware Distillation for Two-Step Video Dubbing

        **VidDub** takes a source talking-head video and a new speech clip and
        re-renders the speaker's lower face so the lips match the new audio —
        in **2 denoising steps** on a Wan2.1-VACE-1.3B backbone.

        📄 [Paper](https://huggingface.co/papers/2611.03404) ·
        💻 [GitHub](https://github.com/LAW1223/VidDub) ·
        🤗 [AIPeanutman/VidDub](https://huggingface.co/AIPeanutman/VidDub)

        Upload a short clip with a single, mostly front-facing face and the
        audio you want the face to speak. Everything except the mouth region is
        kept from the source video.
        """)
    with gr.Row():
        with gr.Column():
            video_in = gr.Video(label="Source video (talking head)", sources=["upload"])
            audio_in = gr.Audio(label="Target speech audio", sources=["upload"], type="filepath")
            run_btn = gr.Button("Dub video", variant="primary")
        with gr.Column():
            video_out = gr.Video(label="Dubbed video", autoplay=True)
    with gr.Accordion("Advanced options", open=False):
        seed = gr.Number(label="Seed", value=0, precision=0)
        harmonize = gr.Checkbox(label="Color harmonization (paper default)", value=True)
        feather = gr.Slider(1, 40, value=15, step=1, label="Mask feather (px)")
        max_frames = gr.Slider(NUM_FRAMES, MAX_FRAMES, value=MAX_FRAMES, step=1,
                               label="Max frames (longer videos are truncated)")
    gr.Examples(
        examples=EXAMPLES,
        inputs=[video_in, audio_in, seed],
        fn=dub,
        outputs=[video_out],
        cache_examples=True,
        cache_mode="lazy",
    )
    run_btn.click(
        fn=dub,
        inputs=[video_in, audio_in, seed, harmonize, feather, max_frames],
        outputs=[video_out],
        api_name="dub",
    )
    gr.Markdown(
        """
        ---
        **Notes.** The two-step DMD checkpoint is fixed: 2 denoising steps,
        audio CFG 1, 81-frame windows with 17-frame overlap. Videos need a
        detectable, mostly front-facing face and at least 81 frames
        (~3.4 s). VidDub-Bench is CC BY-NC; the bundled example videos come
        from the CC0 pool at
        [linoyts/repo-to-space-example-videos](https://huggingface.co/datasets/linoyts/repo-to-space-example-videos)
        and the speech from
        [LibriSpeech](https://huggingface.co/datasets/openslr/librispeech_asr)
        (CC BY 4.0).
        """)

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