"""LeapTalk — real-time audio-driven talking-head generation on ZeroGPU.

Faithful port of the official reference implementation
(https://github.com/zhangrongxiang/LeapTalk, `inference.py` streaming path):

    SoulX-FlashHead-1_3B (Model_Pro)  +  LeapTalk LoRA (merged)
    + LeapTalk audio projector        +  wav2vec2-base-960h audio encoder
    + Lite TAE (taew2_1) VAE          +  ViBT Brownian-bridge scheduler

Everything (chunking, audio windowing, bridge sampling, motion-frame
round-trip, colour correction) mirrors the authors' `--lite` / `--model_type pro`
/ `--audio_encode_mode stream` defaults from `inf.sh`.
"""

import os

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

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

import math  # noqa: E402
import shutil  # noqa: E402
import subprocess  # noqa: E402
import sys  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402
import wave  # noqa: E402
from collections import deque  # noqa: E402

import gradio as gr  # noqa: E402
import imageio  # noqa: E402
import librosa  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402
from loguru import logger  # noqa: E402
from peft import PeftModel  # noqa: E402

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import flash_head.src.pipeline.flash_head_pipeline as fh_pipe_mod  # noqa: E402

# torch.compile is disabled: the released LoRA was not saved from a compiled base
# (`--compile off` in the reference `inf.sh`), and TAEHV uses Python-level loops.
fh_pipe_mod.COMPILE_MODEL = False
fh_pipe_mod.COMPILE_VAE = False

from flash_head.src.pipeline.flash_head_pipeline import FlashHeadPipeline  # noqa: E402
from leaptalk_inference import (  # noqa: E402
    StreamParams,
    _audio_context_from_embeddings_range,
    _bridge_sample_one_chunk,
    _build_infer_timesteps,
    _decode_to_cthw,
    _encode_motion_prefix_from_decoded,
    _get_inner_flashhead_model,
    _maybe_apply_color_correction,
)
from vibt.scheduler import ViBTScheduler  # noqa: E402

# --------------------------------------------------------------------------------------
# Fixed inference configuration (reference defaults)
# --------------------------------------------------------------------------------------
DEVICE = "cuda"
DTYPE = torch.bfloat16
HEIGHT = WIDTH = 512
FPS = 25
SAMPLE_RATE = 16000
FRAME_NUM = 33
MOTION_FRAMES_LATENT_NUM = 2
CACHED_AUDIO_DURATION = 8
SHIFT_GAMMA = 5.0
NOISE_SCALE = 1.0
COLOR_CORRECTION_STRENGTH = 1.0
MAX_SECONDS_CAP = 20

# --------------------------------------------------------------------------------------
# Weights
# --------------------------------------------------------------------------------------
logger.info("Downloading weights…")
CKPT_DIR = snapshot_download(
    "Soul-AILab/SoulX-FlashHead-1_3B", allow_patterns=["Model_Pro/*"]
)
WAV2VEC_DIR = snapshot_download(
    "facebook/wav2vec2-base-960h",
    allow_patterns=["*.json", "*.txt", "*.safetensors", "pytorch_model.bin"],
)
LEAPTALK_DIR = snapshot_download("z-rx/leaptalk")
LORA_DIR = os.path.join(LEAPTALK_DIR, "lora")
TAE_PATH = os.path.join(LEAPTALK_DIR, "taew2_1.pth")
AUDIO_PROJ_PATH = os.path.join(LEAPTALK_DIR, "audio_proj_step_10400.pt")

# --------------------------------------------------------------------------------------
# Pipeline assembly (module scope, eagerly moved to CUDA)
# --------------------------------------------------------------------------------------
logger.info("Building FlashHead pipeline…")
# Built on CPU first so the LoRA merge / projector load happen on real tensors,
# then the whole stack is moved to CUDA eagerly (ZeroGPU packs it from there).
pipeline = FlashHeadPipeline(
    checkpoint_dir=CKPT_DIR,
    model_type="pro",
    wav2vec_dir=WAV2VEC_DIR,
    device="cpu",
    param_dtype=DTYPE,
    use_usp=False,
    use_tae=True,
    tae_path=TAE_PATH,
    tae_model_type="wan21",
)

logger.info("Merging LeapTalk LoRA…")
# `torch_device="cpu"` is required: PEFT otherwise infers "cuda" (ZeroGPU reports a GPU
# as available at import time) and safetensors' loader bypasses the ZeroGPU patching.
pipeline.model = PeftModel.from_pretrained(
    pipeline.model, LORA_DIR, is_trainable=False, torch_device="cpu"
)
pipeline.model = pipeline.model.merge_and_unload()
pipeline.model.eval().requires_grad_(False)

logger.info("Loading LeapTalk audio projector…")
_audio_proj_state = torch.load(AUDIO_PROJ_PATH, map_location="cpu", weights_only=True)
_get_inner_flashhead_model(pipeline.model).audio_proj.load_state_dict(
    _audio_proj_state, strict=True
)
del _audio_proj_state

pipeline.device = DEVICE
pipeline.model.to(DEVICE)
pipeline.vae.device = DEVICE
pipeline.vae.model.to(DEVICE)
pipeline.audio_encoder.to(DEVICE)
pipeline.audio_encoder.eval().requires_grad_(False)

STREAM = StreamParams(
    frame_num=FRAME_NUM,
    motion_frames_latent_num=MOTION_FRAMES_LATENT_NUM,
    tgt_fps=FPS,
    sample_rate=SAMPLE_RATE,
    cached_audio_duration=CACHED_AUDIO_DURATION,
).init_with_stride(int(pipeline.config.vae_stride[0]))
SLICE_SAMPLES = STREAM.slice_len * SAMPLE_RATE // FPS
logger.info(
    f"Ready. frame_num={STREAM.frame_num} motion_frames={STREAM.motion_frames_num} "
    f"slice_len={STREAM.slice_len} ({SLICE_SAMPLES} samples/chunk)"
)


# --------------------------------------------------------------------------------------
# Video helpers
# --------------------------------------------------------------------------------------
def _ffmpeg_exe() -> str:
    exe = shutil.which("ffmpeg")
    if exe:
        return exe
    import imageio_ffmpeg

    return imageio_ffmpeg.get_ffmpeg_exe()


def _write_wav(path: str, audio: np.ndarray, sample_rate: int = SAMPLE_RATE) -> str:
    pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
    with wave.open(path, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(pcm.tobytes())
    return path


def _mux(video_path: str, audio_path: str, out_path: str) -> str:
    cmd = [
        _ffmpeg_exe(), "-y",
        "-i", video_path,
        "-i", audio_path,
        "-c:v", "copy",
        "-c:a", "aac", "-b:a", "128k",
        "-shortest",
        "-movflags", "+faststart",
        out_path,
    ]
    proc = subprocess.run(cmd, capture_output=True)
    if proc.returncode != 0 or not os.path.exists(out_path):
        logger.warning(f"ffmpeg mux failed: {proc.stderr.decode()[-800:]}")
        shutil.copy(video_path, out_path)
    return out_path


def _num_chunks_for(seconds: float) -> int:
    samples = max(int(seconds * SAMPLE_RATE), FRAME_NUM * SAMPLE_RATE // FPS)
    return max(1, math.ceil(samples / SLICE_SAMPLES))


def _estimate_duration(
    portrait_image=None,
    speech_audio=None,
    max_seconds: float = 9.0,
    num_inference_steps: int = 1,
    guidance_scale: float = 1.0,
    *args,
    **kwargs,
) -> int:
    """ZeroGPU time budget: weight streaming + per-chunk cost."""
    try:
        chunks = _num_chunks_for(float(max_seconds))
        nfe = max(1, int(num_inference_steps)) * (2 if float(guidance_scale) != 1.0 else 1)
    except Exception:
        chunks, nfe = _num_chunks_for(MAX_SECONDS_CAP), 1
    # Measured on ZeroGPU (A100): ~0.54 s/chunk at 1 NFE, ~1.5 s of fixed setup +
    # video encode/mux, plus weight streaming on a cold worker. Kept deliberately tight
    # so the demo does not over-reserve visitors' quota.
    return int(min(90, 12 + chunks * (0.35 + 0.35 * nfe)))


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=_estimate_duration)
def generate(
    portrait_image: str,
    speech_audio: str,
    max_seconds: float = 9.0,
    num_inference_steps: int = 1,
    guidance_scale: float = 1.0,
    seed: int = 42,
    auto_crop_face: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Animate a portrait photo so that it speaks the given audio.

    Args:
        portrait_image: Path to a portrait photo (a single, roughly front-facing face).
        speech_audio: Path to a speech audio file that drives lip and head motion.
        max_seconds: Maximum number of seconds of the audio to animate.
        num_inference_steps: Bridge-sampler steps per chunk. LeapTalk is distilled for 1.
        guidance_scale: Audio classifier-free guidance. 1.0 disables it (2x faster).
        seed: Random seed for the Brownian-bridge noise.
        auto_crop_face: Detect and crop around the face before generating.

    Returns:
        A tuple of (path to the generated talking-head mp4, a short speed report).
    """
    if not portrait_image:
        raise gr.Error("Please provide a portrait image.")
    if not speech_audio:
        raise gr.Error("Please provide a speech audio file.")

    num_inference_steps = max(1, int(num_inference_steps))
    guidance_scale = float(guidance_scale)
    seed = int(seed)
    max_seconds = float(np.clip(max_seconds, 1.0, MAX_SECONDS_CAP))

    workdir = tempfile.mkdtemp(prefix="leaptalk_")
    progress(0.02, desc="Preparing reference portrait…")

    # ---- reference image -> anchor latent X0 (same call as inference.py) --------------
    pipeline.prepare_params(
        cond_image_path_or_dir=portrait_image,
        target_size=(HEIGHT, WIDTH),
        frame_num=STREAM.frame_num,
        motion_frames_num=0,
        sampling_steps=num_inference_steps,
        seed=seed,
        shift=SHIFT_GAMMA,
        color_correction_strength=COLOR_CORRECTION_STRENGTH,
        use_face_crop=bool(auto_crop_face),
    )
    X0 = pipeline.ref_img_latent.to(device=DEVICE, dtype=DTYPE)

    # ---- scheduler -------------------------------------------------------------------
    scheduler = ViBTScheduler(num_train_timesteps=1000)
    scheduler.timesteps = _build_infer_timesteps(
        step_list=None,
        num_inference_steps=num_inference_steps,
        shift_gamma=SHIFT_GAMMA,
        device=DEVICE,
        num_timesteps=1000,
    )
    scheduler.num_inference_steps = int(scheduler.timesteps.numel())
    scheduler.set_parameters(noise_scale=NOISE_SCALE, shift_gamma=SHIFT_GAMMA, seed=seed)

    # ---- audio (streaming ring buffer, exactly as inference.py --audio_encode_mode stream)
    progress(0.06, desc="Loading audio…")
    audio_all, _ = librosa.load(speech_audio, sr=SAMPLE_RATE, mono=True)
    audio_all = audio_all[: int(max_seconds * SAMPLE_RATE)]
    if audio_all.size == 0:
        raise gr.Error("The audio file appears to be empty.")

    frame_window_samples = STREAM.frame_num * SAMPLE_RATE // FPS
    remainder = len(audio_all) % SLICE_SAMPLES
    if remainder > 0:
        audio_all = np.concatenate(
            [audio_all, np.zeros(SLICE_SAMPLES - remainder, dtype=audio_all.dtype)]
        )
    if len(audio_all) < frame_window_samples:
        audio_all = np.concatenate(
            [audio_all, np.zeros(frame_window_samples - len(audio_all), dtype=audio_all.dtype)]
        )
    remainder = len(audio_all) % SLICE_SAMPLES
    if remainder != 0:
        audio_all = np.concatenate(
            [audio_all, np.zeros(SLICE_SAMPLES - remainder, dtype=audio_all.dtype)]
        )

    slices = audio_all.reshape(-1, SLICE_SAMPLES)
    num_chunks = int(slices.shape[0])

    cached_len = SAMPLE_RATE * STREAM.cached_audio_duration
    audio_end_idx = STREAM.cached_audio_duration * FPS
    audio_start_idx = audio_end_idx - STREAM.frame_num
    audio_dq = deque([0.0] * cached_len, maxlen=cached_len)

    latent_motion_frames = X0[:, :1].unsqueeze(0).clone()
    clamp_latent_len = int(latent_motion_frames.shape[2])

    generated: list[np.ndarray] = []
    gen_seconds = 0.0
    gen_frames = 0

    for chunk_idx in range(num_chunks):
        progress(
            0.08 + 0.88 * chunk_idx / num_chunks,
            desc=f"Generating chunk {chunk_idx + 1}/{num_chunks}…",
        )
        torch.cuda.synchronize()
        t0 = time.perf_counter()

        audio_dq.extend(slices[chunk_idx].tolist())
        audio_cache = np.array(audio_dq, dtype=np.float32)
        audio_emb = pipeline.preprocess_audio(audio_cache, sr=SAMPLE_RATE, fps=FPS)
        if audio_emb is None:
            raise gr.Error("Failed to extract audio embeddings.")
        audio_emb = audio_emb.to(device=DEVICE, dtype=DTYPE)
        audio_ctx = _audio_context_from_embeddings_range(
            audio_emb,
            start_idx=audio_start_idx,
            end_idx=audio_end_idx,
            device=DEVICE,
            dtype=DTYPE,
        )

        x_final = _bridge_sample_one_chunk(
            pipeline,
            scheduler=scheduler,
            ref_latent=X0,
            audio_context=audio_ctx,
            guidance_scale=guidance_scale,
            latent_motion_frames=latent_motion_frames,
            clamp_latent_len=clamp_latent_len,
            device=DEVICE,
            dtype=DTYPE,
        )
        decoded_cthw = _decode_to_cthw(pipeline, x_final)
        decoded_cthw = _maybe_apply_color_correction(pipeline, decoded_cthw)

        # SoulX-style VAE round-trip history update (reference default)
        latent_motion_frames = _encode_motion_prefix_from_decoded(
            pipeline,
            decoded_video_cthw=decoded_cthw,
            motion_frames_num=STREAM.motion_frames_num,
            device=DEVICE,
            dtype=DTYPE,
        ).unsqueeze(0)
        clamp_latent_len = int(latent_motion_frames.shape[2])

        decoded_cthw = decoded_cthw[:, STREAM.motion_frames_num:]
        video_thwc = (
            ((decoded_cthw + 1.0) / 2.0)
            .permute(1, 2, 3, 0)
            .clamp(0.0, 1.0)
            .mul(255.0)
            .contiguous()
        )
        torch.cuda.synchronize()
        chunk_seconds = time.perf_counter() - t0

        frames_np = video_thwc.to(torch.float32).cpu().numpy().astype(np.uint8)
        generated.append(frames_np)
        gen_frames += int(frames_np.shape[0])
        gen_seconds += chunk_seconds
        logger.info(
            f"chunk {chunk_idx + 1}/{num_chunks}: {chunk_seconds:.3f}s "
            f"({frames_np.shape[0] / max(chunk_seconds, 1e-6):.1f} FPS)"
        )

    progress(0.97, desc="Encoding video…")
    silent_path = os.path.join(workdir, "silent.mp4")
    with imageio.get_writer(
        silent_path,
        format="mp4",
        mode="I",
        fps=FPS,
        codec="h264",
        pixelformat="yuv420p",
        ffmpeg_params=["-bf", "0"],
    ) as writer:
        for frames_np in generated:
            for frame in frames_np:
                writer.append_data(frame)

    wav_path = _write_wav(os.path.join(workdir, "track.wav"), audio_all)
    out_path = _mux(silent_path, wav_path, os.path.join(workdir, "leaptalk.mp4"))

    video_seconds = gen_frames / FPS
    report = (
        f"**{gen_frames} frames** ({video_seconds:.1f}s of video) in "
        f"**{gen_seconds:.2f}s** of GPU time — "
        f"**{gen_frames / max(gen_seconds, 1e-6):.1f} FPS** generation throughput "
        f"({gen_frames / max(gen_seconds, 1e-6) / FPS:.2f}× real time) over "
        f"{num_chunks} streaming chunks at {num_inference_steps} step"
        f"{'s' if num_inference_steps > 1 else ''}/chunk."
    )
    return out_path, report


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

with gr.Blocks(title="LeapTalk") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
# 🗣️ LeapTalk — real-time talking heads

Animate a **portrait photo** with a **speech clip**. LeapTalk reformulates talking-head
generation as a Brownian-bridge transport (*Bridge Forcing*), which lets it synthesize each
video chunk in a **single sampling step** while keeping identity stable over long rollouts.

[Model](https://huggingface.co/z-rx/leaptalk) · [Paper](https://huggingface.co/papers/2608.00079)
· [Project page](https://zhangrongxiang.github.io/leaptalk-page/)
· [Code](https://github.com/zhangrongxiang/LeapTalk)
· built on [SoulX-FlashHead-1.3B](https://huggingface.co/Soul-AILab/SoulX-FlashHead-1_3B)
            """
        )

        with gr.Row():
            with gr.Column():
                portrait_image = gr.Image(
                    label="Portrait", type="filepath", height=320, sources=["upload", "webcam", "clipboard"]
                )
                speech_audio = gr.Audio(
                    label="Speech audio", type="filepath", sources=["upload", "microphone"]
                )
                run_btn = gr.Button("Generate talking head", variant="primary")
            with gr.Column():
                video_out = gr.Video(
                    label="Result", height=460, autoplay=True
                )
                report_out = gr.Markdown()

        with gr.Accordion("Advanced options", open=False):
            with gr.Row():
                max_seconds = gr.Slider(
                    label="Max audio length (seconds)",
                    minimum=1,
                    maximum=MAX_SECONDS_CAP,
                    step=1,
                    value=9,
                )
                num_inference_steps = gr.Slider(
                    label="Sampling steps per chunk",
                    minimum=1,
                    maximum=4,
                    step=1,
                    value=1,
                    info="LeapTalk is distilled for 1-step (1 NFE) generation.",
                )
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Audio guidance scale",
                    minimum=1.0,
                    maximum=3.0,
                    step=0.1,
                    value=1.0,
                    info="1.0 disables audio CFG; higher strengthens lip motion but doubles compute.",
                )
                seed = gr.Number(label="Seed", value=42, precision=0)
            auto_crop_face = gr.Checkbox(
                label="Auto-crop to face",
                value=True,
                info="Detects the face and crops around it; falls back to a centre crop.",
            )

        gr.Examples(
            examples=[
                ["examples/portrait.jpg", "examples/narration.wav"],
                ["examples/girl.png", "examples/podcast_sichuan.wav"],
            ],
            inputs=[portrait_image, speech_audio],
            outputs=[video_out, report_out],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

        gr.Markdown(
            "Example assets: portrait from the "
            "[LeapTalk](https://github.com/zhangrongxiang/LeapTalk) repository, portrait + podcast "
            "clip from [SoulX-FlashHead](https://github.com/Soul-AILab/SoulX-FlashHead) "
            "(both Apache-2.0). The narration clip is public-domain audiobook narration from "
            "[LibriSpeech](https://www.openslr.org/12) (LibriVox, CC0 / public domain). "
            "Audio clips were trimmed to a few seconds."
        )

    gr.on(
        triggers=[run_btn.click],
        fn=generate,
        inputs=[
            portrait_image,
            speech_audio,
            max_seconds,
            num_inference_steps,
            guidance_scale,
            seed,
            auto_crop_face,
        ],
        outputs=[video_out, report_out],
    )

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