"""Lip Forcing — few-step autoregressive diffusion for real-time lip synchronization.

ZeroGPU Gradio demo for the released 14B student
(https://huggingface.co/JinhyukJang/lipforcing). Given a talking-head reference
video and a driving audio clip, it re-synchronizes the mouth to the audio using
the streaming per-chunk AR pipeline from the official repo
(scripts/inference/inference_streaming.py), reproduced 1:1 here.
"""

import os

# Allocator: the streaming AR loop has transient spikes (VAE encode/decode of
# 512x512 chunks + KV cache). expandable segments avoids fragmentation OOMs.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("ORT_DISABLE_THREAD_AFFINITY", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import spaces  # noqa: E402  — must precede torch / CUDA-touching imports

import sys
import types
import tempfile
import traceback

import numpy as np
import torch
import gradio as gr
from PIL import Image
from huggingface_hub import hf_hub_download, snapshot_download

# The inference scripts import their helpers as top-level modules
# (`from _common import ...`), so make scripts/inference importable that way.
REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, REPO_ROOT)
sys.path.insert(0, os.path.join(REPO_ROOT, "scripts", "inference"))

# ---------------------------------------------------------------------------
# Weights (downloaded once at startup into the HF cache)
# ---------------------------------------------------------------------------
print("Downloading weights ...", flush=True)

CKPT_PATH = hf_hub_download("JinhyukJang/lipforcing", "lipforcing_14b.pth")

WAN_REPO = "Wan-AI/Wan2.1-T2V-14B"
VAE_PATH = hf_hub_download(WAN_REPO, "Wan2.1_VAE.pth")
T5_PATH = hf_hub_download(WAN_REPO, "models_t5_umt5-xxl-enc-bf16.pth")
# UMT5 tokenizer (lives under google/umt5-xxl/ inside the Wan repo)
for _f in (
    "google/umt5-xxl/special_tokens_map.json",
    "google/umt5-xxl/spiece.model",
    "google/umt5-xxl/tokenizer.json",
    "google/umt5-xxl/tokenizer_config.json",
):
    hf_hub_download(WAN_REPO, _f)
# T5_PATH's parent dir now also holds google/umt5-xxl/* (same snapshot dir).

WAV2VEC_DIR = snapshot_download("facebook/wav2vec2-base-960h")

# TAEW tiny streaming decoder + LatentSync mouth mask.
# taew2_1.pth lives in the taehv GitHub repo; mask.png in the LatentSync repo.
import urllib.request  # noqa: E402
_cache = os.path.join(tempfile.gettempdir(), "lipforcing_assets")
os.makedirs(_cache, exist_ok=True)

TAEHV_CKPT = os.path.join(_cache, "taew2_1.pth")
if not os.path.exists(TAEHV_CKPT):
    urllib.request.urlretrieve(
        "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth",
        TAEHV_CKPT,
    )

MASK_PATH = os.path.join(_cache, "mask.png")
if not os.path.exists(MASK_PATH):
    urllib.request.urlretrieve(
        "https://raw.githubusercontent.com/bytedance/LatentSync/main/latentsync/utils/mask.png",
        MASK_PATH,
    )

print("Weights downloaded.", flush=True)

DTYPE = torch.bfloat16
DEVICE = "cuda"

# ---------------------------------------------------------------------------
# Args shim — the loaders/helpers read attributes off an argparse-like object.
# We build one with the released 14B student's default (2-step t769) schedule.
# ---------------------------------------------------------------------------
def _make_args():
    a = types.SimpleNamespace()
    a.ckpt_path = CKPT_PATH
    a.vae_path = VAE_PATH
    a.wav2vec_path = WAV2VEC_DIR
    a.mask_path = MASK_PATH
    a.taehv_ckpt = TAEHV_CKPT
    a.base_model_paths = None
    a.omniavatar_ckpt_path = None
    a.model_size = "14B"
    a.merge_lora_post_load = True
    a.text_embeds_path = None
    a.text_encoder_path = None           # text encoded once at startup (below)
    a.prompt = "a person talking"
    a.streaming_decoder = "streaming_taehv"
    a.t_list = [0.999, 0.769, 0.0]        # released 14B 2-step schedule
    a.chunk_size = 3
    a.num_latent_frames = None
    a.min_latent_frames = 0
    a.context_noise = 0.0
    a.seed = 42
    a.fps = 25.0
    a.dtype = "bf16"
    a.device = DEVICE
    a.local_attn_size = 7
    a.sink_size = 1
    a.use_dynamic_rope = True
    a.skip_preprocessing = False
    a.face_cache_dir = None
    a.composite_full_face = False
    a.streamwise_encode = True
    a.defer_composite = False
    a.compile = False
    a.input_dir = None
    a.output_dir = None
    a.video_path = None
    a.audio_path = None
    a.output_path = None
    return a


ARGS = _make_args()

# ---------------------------------------------------------------------------
# Text embedding: encode the default prompt ONCE on CPU, then free the 11 GB
# UMT5-XXL encoder. This keeps the encoder off the GPU so peak VRAM stays low
# (~37 GB), per the model card's "48 GB cards work with precomputed embeddings".
# ---------------------------------------------------------------------------
def _precompute_text_embeds(prompt: str) -> torch.Tensor:
    from OmniAvatar.models.wan_video_text_encoder import WanTextEncoder
    from OmniAvatar.prompters.wan_prompter import WanPrompter
    from lipforcing import preprocess as pp

    print(f"Encoding text prompt on CPU: {prompt!r} ...", flush=True)
    text_encoder = WanTextEncoder()
    te_state = torch.load(T5_PATH, map_location="cpu", weights_only=False)
    converter = WanTextEncoder.state_dict_converter()
    te_state = converter.from_civitai(te_state)
    text_encoder.load_state_dict(te_state, strict=True)
    text_encoder = text_encoder.to("cpu").eval()

    tokenizer_path = pp._resolve_tokenizer_path(T5_PATH)
    prompter = WanPrompter(tokenizer_path=tokenizer_path, text_len=512)
    prompter.fetch_models(text_encoder=text_encoder)
    with torch.no_grad():
        emb = prompter.encode_prompt(prompt, positive=True, device="cpu")
    if emb.dim() == 2:
        emb = emb.unsqueeze(0)
    emb = emb.to(dtype=DTYPE).contiguous()
    del text_encoder, prompter, te_state
    import gc
    gc.collect()
    print(f"Text embeds: {tuple(emb.shape)}", flush=True)
    return emb


TEXT_EMBEDS_CPU = _precompute_text_embeds(ARGS.prompt)

# ---------------------------------------------------------------------------
# Models — loaded at module scope, .to("cuda") intercepted by ZeroGPU.
# ---------------------------------------------------------------------------
print("Loading diffusion model (14B student) ...", flush=True)
from _loader import load_diffusion_model            # noqa: E402
from _common import (                               # noqa: E402
    TAEHVDecoderWrapper, load_vae, load_wav2vec,
    resolve_audio, compute_generation_length,
    load_image_processor, preprocess_with_latentsync,
)
from inference_streaming import (                    # noqa: E402
    run_streaming_pipeline, build_condition_streamwise,
)

MODEL = load_diffusion_model(ARGS, DEVICE, DTYPE)

print("Loading Wan VAE ...", flush=True)
VAE = load_vae(ARGS.vae_path, DEVICE)

print("Loading TAEHV decoder ...", flush=True)
DECODER_VAE = TAEHVDecoderWrapper(ARGS.taehv_ckpt, DEVICE)

print("Loading Wav2Vec2 ...", flush=True)
WAV2VEC_MODEL, WAV2VEC_EXTRACTOR = load_wav2vec(ARGS.wav2vec_path, DEVICE)

# LatentSync face detector / aligner uses insightface + onnxruntime; those need
# a live GPU context, so it is initialized lazily inside the GPU call.
IMAGE_PROCESSOR = None


def _get_image_processor():
    global IMAGE_PROCESSOR
    if IMAGE_PROCESSOR is None:
        IMAGE_PROCESSOR = load_image_processor(ARGS.mask_path, DEVICE)
    return IMAGE_PROCESSOR


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
MAX_SECONDS = 8.0  # cap driving audio so a single call stays within GPU budget


def _estimate_duration(video_path, audio_path, *a, **k):
    # 14B student: streaming AR + face detect/composite. Budget generously per
    # second of (capped) audio, plus fixed preprocessing/warmup overhead.
    # Measured ~330s for a 4s clip on first (cold) call incl. warmup; scale by
    # audio length with a generous fixed base and cap at the audio limit.
    base = 60.0
    per_sec = 50.0
    secs = MAX_SECONDS
    try:
        import librosa
        secs = min(librosa.get_duration(path=audio_path), MAX_SECONDS)
    except Exception:
        pass
    return int(base + per_sec * secs)


@spaces.GPU(duration=_estimate_duration)
def lip_sync(video_path: str, audio_path: str,
             seed: int = 42) -> str:
    """Lip-sync a talking-head video to a driving audio clip.

    Args:
        video_path: reference talking-head video (any resolution; a single
            clear front-facing face is detected, aligned to 512x512, and the
            mouth region is regenerated to match the audio).
        audio_path: driving speech audio; the output length follows the audio
            (capped to keep a single request within the GPU budget).
        seed: RNG seed for reproducibility.

    Returns:
        Path to the generated lip-synced mp4 (muxed with the driving audio).
    """
    if not video_path:
        raise gr.Error("Please provide a reference talking-head video.")
    if not audio_path:
        raise gr.Error("Please provide a driving audio clip.")

    import imageio_ffmpeg
    import subprocess

    args = _make_args()
    args.seed = int(seed)
    args.video_path = video_path
    args.audio_path = audio_path

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    image_processor = _get_image_processor()

    # Cap audio length so runtime stays bounded.
    ff = imageio_ffmpeg.get_ffmpeg_exe()
    capped_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    subprocess.run(
        [ff, "-y", "-loglevel", "error", "-nostdin", "-i", audio_path,
         "-t", str(MAX_SECONDS), "-ar", "16000", "-ac", "1", capped_audio],
        check=True,
    )

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name

    tmp_audio = None
    try:
        used_audio, tmp_audio = resolve_audio(audio_path=capped_audio)

        num_latent_frames, num_video_frames = compute_generation_length(
            used_audio, args.num_latent_frames, args.chunk_size, args.fps,
            min_latent_frames=args.min_latent_frames,
        )

        print("Face detection + 512x512 alignment ...", flush=True)
        meta = preprocess_with_latentsync(
            args.video_path, image_processor, args.face_cache_dir,
            num_frames=num_video_frames,
        )
        if meta is None:
            raise gr.Error(
                "Face detection failed — please provide a video with a single, "
                "clear, front-facing talking head."
            )

        aligned_faces = meta["aligned_faces"]
        ref_frames_np = np.stack([
            f.permute(1, 2, 0).numpy() if isinstance(f, torch.Tensor) else f
            for f in aligned_faces[:num_video_frames]
        ], axis=0)

        text_embeds = TEXT_EMBEDS_CPU.to(device=DEVICE, dtype=DTYPE)

        condition, video_tensor, masked_video_tensor = build_condition_streamwise(
            VAE, WAV2VEC_MODEL, WAV2VEC_EXTRACTOR,
            ref_frames_np, used_audio, text_embeds, args.mask_path,
            num_video_frames, num_latent_frames, DEVICE, DTYPE,
        )

        print("Running streaming pipeline ...", flush=True)
        run_streaming_pipeline(
            MODEL, DECODER_VAE, VAE, condition,
            num_latent_frames, num_video_frames,
            args, meta, image_processor,
            used_audio, out_path, DEVICE, DTYPE,
            video_tensor=video_tensor,
            masked_video_tensor=masked_video_tensor,
        )
    except gr.Error:
        raise
    except Exception as e:
        traceback.print_exc()
        raise gr.Error(f"Inference failed: {e}")
    finally:
        MODEL.clear_caches()
        torch.cuda.empty_cache()
        if tmp_audio and os.path.exists(tmp_audio):
            os.remove(tmp_audio)
        if os.path.exists(capped_audio):
            os.remove(capped_audio)

    return out_path


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

DESCRIPTION = """
# Lip Forcing 🗣️
**Few-Step Autoregressive Diffusion for Real-time Lip Synchronization**  · 
14B student  · 
[Paper](https://arxiv.org/abs/2606.11180)  · 
[Project](https://cvlab-kaist.github.io/LipForcing/)  · 
[Code](https://github.com/cvlab-kaist/LipForcing)  · 
[Weights](https://huggingface.co/JinhyukJang/lipforcing)

Give it a **talking-head video** and a **driving audio** clip — it detects and aligns
the face, then regenerates the mouth to match the audio with a 2-step causal diffusion
student. Audio is capped to the first few seconds per run.
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(DESCRIPTION)
        with gr.Row():
            with gr.Column():
                video_in = gr.Video(label="Reference talking-head video", height=340)
                audio_in = gr.Audio(label="Driving audio", type="filepath")
                run_btn = gr.Button("Lip-sync", variant="primary")
            with gr.Column():
                video_out = gr.Video(label="Lip-synced result", height=340)
        with gr.Accordion("Advanced settings", open=False):
            seed = gr.Number(label="Seed", value=42, precision=0)

        run_btn.click(
            fn=lip_sync,
            inputs=[video_in, audio_in, seed],
            outputs=video_out,
            api_name="lip_sync",
        )

        gr.Examples(
            examples=[
                ["examples/example1_video.mp4", "examples/example1_audio.wav"],
                ["examples/example2_video.mp4", "examples/example2_audio.wav"],
            ],
            inputs=[video_in, audio_in],
            outputs=video_out,
            fn=lip_sync,
            cache_examples=True,
            cache_mode="lazy",
        )

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