# -*- coding: utf-8 -*-
"""
EchoMimicV3 - Audio/Text-driven Human Animation
Model: https://huggingface.co/BadToBest/EchoMimicV3
GitHub: https://github.com/antgroup/echomimic_v3
Paper: https://arxiv.org/abs/2507.03905
"""

import os
import sys
import math
import datetime
import subprocess
import random
import gc

# Must be set before importing gradio — HF Spaces may ignore launch(ssr_mode=...).
os.environ["GRADIO_SSR_MODE"] = "0"

from huggingface_hub import snapshot_download

# ---------------------------------------------------------------------------
# Source code + model layout (matches official app_mm.py / infer_preview.py)
#
#   ./echomimic_v3/          # GitHub source (provides src.*)
#   ./models/
#     Wan2.1-Fun-V1.1-1.3B-InP/   # base: VAE / T5 / CLIP / config
#     transformer/                 # EchoMimicV3 fine-tuned weights
#     wav2vec2-base-960h/
# ---------------------------------------------------------------------------

ROOT = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.join(ROOT, "echomimic_v3")
MODELS_DIR = os.path.join(ROOT, "models")
WAN_MODEL_DIR = os.path.join(MODELS_DIR, "Wan2.1-Fun-V1.1-1.3B-InP")
TRANSFORMER_DIR = os.path.join(MODELS_DIR, "transformer")
WAV2VEC_DIR = os.path.join(MODELS_DIR, "wav2vec2-base-960h")
CONFIG_PATH = os.path.join(SRC_DIR, "config", "config.yaml")
TRANSFORMER_WEIGHTS = os.path.join(TRANSFORMER_DIR, "diffusion_pytorch_model.safetensors")


def ensure_source():
    """Clone EchoMimicV3 source and always put it on sys.path."""
    if not os.path.isdir(SRC_DIR):
        print("Cloning EchoMimicV3 source...")
        subprocess.run(
            ["git", "clone", "--depth", "1", "https://github.com/antgroup/echomimic_v3.git", SRC_DIR],
            check=True,
        )
    if SRC_DIR not in sys.path:
        sys.path.insert(0, SRC_DIR)


def ensure_models():
    """Download base Wan2.1 + EchoMimic transformer + wav2vec into local models/."""
    os.makedirs(MODELS_DIR, exist_ok=True)

    if not os.path.isfile(os.path.join(WAN_MODEL_DIR, "config.json")):
        print("Downloading Wan2.1-Fun-V1.1-1.3B-InP base model...")
        snapshot_download(
            repo_id="alibaba-pai/Wan2.1-Fun-V1.1-1.3B-InP",
            local_dir=WAN_MODEL_DIR,
        )

    if not os.path.isfile(TRANSFORMER_WEIGHTS):
        print("Downloading EchoMimicV3 transformer weights...")
        snapshot_download(
            repo_id="BadToBest/EchoMimicV3",
            allow_patterns=["transformer/*"],
            local_dir=MODELS_DIR,
        )

    if not os.path.isdir(WAV2VEC_DIR) or not os.listdir(WAV2VEC_DIR):
        print("Downloading wav2vec2-base-960h...")
        snapshot_download(
            repo_id="facebook/wav2vec2-base-960h",
            local_dir=WAV2VEC_DIR,
        )


ensure_source()

import numpy as np
import torch
from PIL import Image
from omegaconf import OmegaConf
from transformers import AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor
from moviepy import VideoFileClip, AudioFileClip
import librosa
import gradio as gr
from spaces import GPU

from src.dist import set_multi_gpus_devices
from src.wan_vae import AutoencoderKLWan
from src.wan_image_encoder import CLIPModel
from src.wan_text_encoder import WanT5EncoderModel
from src.wan_transformer3d_audio import WanTransformerAudioMask3DModel
from src.pipeline_wan_fun_inpaint_audio import WanFunInpaintAudioPipeline
from src.utils import filter_kwargs, get_image_to_video_latent3, save_videos_grid
from src.fm_solvers import FlowDPMSolverMultistepScheduler
from src.cache_utils import get_teacache_coefficients
from src.face_detect import get_mask_coord

# Inference defaults (aligned with official app_mm.py)
CONFIG = {
    "model_name": WAN_MODEL_DIR,
    "transformer_path": TRANSFORMER_WEIGHTS,
    "wav2vec_model_dir": WAV2VEC_DIR,
    "config_path": CONFIG_PATH,
    "num_inference_steps": 20,
    "guidance_scale": 4.5,
    "audio_guidance_scale": 2.5,
    "fps": 25,
    "sample_size": [768, 768],
    "partial_video_length": 113,
    "overlap_video_length": 8,
    "teacache_threshold": 0.1,
    "shift": 5.0,
}

DEFAULT_NEG_PROMPT = (
    "Gesture is bad. Gesture is unclear. Strange and twisted hands. "
    "Bad hands. Bad fingers. Unclear and blurry hands. "
    "手部快速摆动, 手指频繁抽搐, 夸张手势, 重复机械性动作."
)

pipeline = None
wav2vec_processor = None
wav2vec_model = None
device = None
weight_dtype = None


def load_models():
    """Load Wan2.1 base components + EchoMimicV3 transformer + wav2vec."""
    global pipeline, wav2vec_processor, wav2vec_model, device, weight_dtype

    if pipeline is not None:
        return

    print("Loading EchoMimicV3 models...")
    if not torch.cuda.is_available():
        raise RuntimeError("EchoMimicV3 requires CUDA GPU")

    # Heavy downloads happen here (inside @GPU) to avoid Space startup timeouts.
    ensure_source()
    ensure_models()

    device = "cuda"
    weight_dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] >= 8 else torch.float16

    set_multi_gpus_devices(1, 1)
    cfg = OmegaConf.load(CONFIG["config_path"])
    model_name = CONFIG["model_name"]

    # Structure/config come from Wan2.1 base (transformer_subpath is "./")
    transformer = WanTransformerAudioMask3DModel.from_pretrained(
        os.path.join(model_name, cfg["transformer_additional_kwargs"].get("transformer_subpath", "transformer")),
        transformer_additional_kwargs=OmegaConf.to_container(cfg["transformer_additional_kwargs"]),
        torch_dtype=weight_dtype,
    )

    from safetensors.torch import load_file
    state_dict = load_file(CONFIG["transformer_path"])
    missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
    print(f"Transformer loaded. Missing keys: {len(missing)}, Unexpected keys: {len(unexpected)}")

    vae = AutoencoderKLWan.from_pretrained(
        os.path.join(model_name, cfg["vae_kwargs"].get("vae_subpath", "vae")),
        additional_kwargs=OmegaConf.to_container(cfg["vae_kwargs"]),
    ).to(weight_dtype)

    tokenizer = AutoTokenizer.from_pretrained(
        os.path.join(model_name, cfg["text_encoder_kwargs"].get("tokenizer_subpath", "tokenizer")),
    )

    text_encoder = WanT5EncoderModel.from_pretrained(
        os.path.join(model_name, cfg["text_encoder_kwargs"].get("text_encoder_subpath", "text_encoder")),
        additional_kwargs=OmegaConf.to_container(cfg["text_encoder_kwargs"]),
        torch_dtype=weight_dtype,
    ).eval()

    clip_image_encoder = CLIPModel.from_pretrained(
        os.path.join(model_name, cfg["image_encoder_kwargs"].get("image_encoder_subpath", "image_encoder")),
    ).to(weight_dtype).eval()

    scheduler = FlowDPMSolverMultistepScheduler(
        **filter_kwargs(FlowDPMSolverMultistepScheduler, OmegaConf.to_container(cfg["scheduler_kwargs"]))
    )

    pipeline = WanFunInpaintAudioPipeline(
        transformer=transformer,
        vae=vae,
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        scheduler=scheduler,
        clip_image_encoder=clip_image_encoder,
    )
    pipeline.to(device)

    coefficients = get_teacache_coefficients(model_name)
    if coefficients is not None:
        pipeline.transformer.enable_teacache(
            coefficients,
            CONFIG["num_inference_steps"],
            CONFIG["teacache_threshold"],
            num_skip_start_steps=5,
            offload=True,
        )

    wav2vec_processor = Wav2Vec2Processor.from_pretrained(CONFIG["wav2vec_model_dir"])
    wav2vec_model = Wav2Vec2Model.from_pretrained(CONFIG["wav2vec_model_dir"]).eval().to(device)
    wav2vec_model.requires_grad_(False)

    print("All models loaded successfully!")


def extract_audio_features(audio_path):
    """Extract audio features using Wav2Vec."""
    sr = 16000
    audio_segment, sample_rate = librosa.load(audio_path, sr=sr)
    input_values = wav2vec_processor(
        audio_segment, sampling_rate=sample_rate, return_tensors="pt"
    ).input_values
    input_values = input_values.to(wav2vec_model.device)
    features = wav2vec_model(input_values).last_hidden_state
    return features.squeeze(0)


def get_sample_size(image, default_size):
    """Calculate sample size based on input image dimensions."""
    width, height = image.size
    original_area = width * height
    default_area = default_size[0] * default_size[1]

    if default_area < original_area:
        ratio = math.sqrt(original_area / default_area)
        width = width / ratio // 16 * 16
        height = height / ratio // 16 * 16
    else:
        width = width // 16 * 16
        height = height // 16 * 16

    return int(height), int(width)


def get_ip_mask(coords):
    """Create IP mask for face region."""
    y1, y2, x1, x2, h, w = coords
    Y, X = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
    mask = (Y.unsqueeze(-1) >= y1) & (Y.unsqueeze(-1) < y2) & (X.unsqueeze(-1) >= x1) & (X.unsqueeze(-1) < x2)
    mask = mask.reshape(-1)
    return mask.float()


@GPU
def generate(
    image,
    audio,
    prompt,
    negative_prompt,
    seed_param,
    progress=gr.Progress(),
):
    """Generate animation from image and audio."""
    if image is None:
        raise ValueError("Please upload an image")
    if audio is None:
        raise ValueError("Please upload an audio file")

    progress(0.1, desc="Loading models...")
    load_models()

    if seed_param is None or seed_param < 0:
        seed = random.randint(0, np.iinfo(np.int32).max)
    else:
        seed = int(seed_param)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    save_path = os.path.join(ROOT, "outputs")
    os.makedirs(save_path, exist_ok=True)

    generator = torch.Generator(device=device).manual_seed(seed)

    progress(0.2, desc="Processing image...")
    ref_img = Image.open(image).convert("RGB")
    y1, y2, x1, x2, h_, w_ = get_mask_coord(image)

    progress(0.3, desc="Processing audio...")
    audio_clip = AudioFileClip(audio)
    audio_features = extract_audio_features(audio)
    audio_embeds = audio_features.unsqueeze(0).to(device=device, dtype=weight_dtype)

    video_length = int(audio_clip.duration * CONFIG["fps"])
    video_length = (
        int((video_length - 1) // pipeline.vae.config.temporal_compression_ratio *
            pipeline.vae.config.temporal_compression_ratio) + 1
        if video_length != 1 else 1
    )

    progress(0.4, desc="Preparing generation...")
    sample_height, sample_width = get_sample_size(ref_img, CONFIG["sample_size"])
    downratio = math.sqrt(sample_height * sample_width / h_ / w_)
    coords = (
        int(y1 * downratio // 16), int(y2 * downratio // 16),
        int(x1 * downratio // 16), int(x2 * downratio // 16),
        sample_height // 16, sample_width // 16,
    )
    ip_mask = get_ip_mask(coords).unsqueeze(0)
    ip_mask = torch.cat([ip_mask] * 3).to(device=device, dtype=weight_dtype)

    partial_video_length = int(
        (CONFIG["partial_video_length"] - 1) //
        pipeline.vae.config.temporal_compression_ratio *
        pipeline.vae.config.temporal_compression_ratio
    ) + 1 if video_length != 1 else 1

    _, _, clip_image = get_image_to_video_latent3(
        ref_img, None, video_length=partial_video_length, sample_size=[sample_height, sample_width]
    )

    progress(0.5, desc="Generating video...")
    init_frames = 0
    last_frames = init_frames + partial_video_length
    new_sample = None
    mix_ratio = torch.linspace(0, 1, steps=CONFIG["overlap_video_length"]).view(1, 1, -1, 1, 1)

    total_iterations = (video_length // (partial_video_length - CONFIG["overlap_video_length"])) + 1
    current_iteration = 0

    while init_frames < video_length:
        if last_frames >= video_length:
            partial_video_length = video_length - init_frames
            partial_video_length = (
                int((partial_video_length - 1) // pipeline.vae.config.temporal_compression_ratio *
                    pipeline.vae.config.temporal_compression_ratio) + 1
                if video_length != 1 else 1
            )
            if partial_video_length <= 0:
                break

        input_video, input_video_mask, _ = get_image_to_video_latent3(
            ref_img, None, video_length=partial_video_length, sample_size=[sample_height, sample_width]
        )

        partial_audio_embeds = audio_embeds[:, init_frames * 2 : (init_frames + partial_video_length) * 2]

        with torch.no_grad():
            sample = pipeline(
                prompt,
                num_frames=partial_video_length,
                negative_prompt=negative_prompt or DEFAULT_NEG_PROMPT,
                audio_embeds=partial_audio_embeds,
                audio_scale=1.0,
                ip_mask=ip_mask,
                use_un_ip_mask=False,
                height=sample_height,
                width=sample_width,
                generator=generator,
                neg_scale=1.5,
                neg_steps=2,
                use_dynamic_cfg=True,
                use_dynamic_acfg=True,
                guidance_scale=CONFIG["guidance_scale"],
                audio_guidance_scale=CONFIG["audio_guidance_scale"],
                num_inference_steps=CONFIG["num_inference_steps"],
                video=input_video,
                mask_video=input_video_mask,
                clip_image=clip_image,
                cfg_skip_ratio=0,
                shift=CONFIG["shift"],
            ).videos

        if init_frames != 0:
            new_sample[:, :, -CONFIG["overlap_video_length"]:] = (
                new_sample[:, :, -CONFIG["overlap_video_length"]:] * (1 - mix_ratio) +
                sample[:, :, :CONFIG["overlap_video_length"]] * mix_ratio
            )
            new_sample = torch.cat([new_sample, sample[:, :, CONFIG["overlap_video_length"]:]], dim=2)
            sample = new_sample
        else:
            new_sample = sample

        if last_frames >= video_length:
            break

        ref_img = [
            Image.fromarray(
                (sample[0, :, i].transpose(0, 1).transpose(1, 2) * 255).numpy().astype(np.uint8)
            ) for i in range(-CONFIG["overlap_video_length"], 0)
        ]

        init_frames += partial_video_length - CONFIG["overlap_video_length"]
        last_frames = init_frames + partial_video_length

        current_iteration += 1
        progress(
            0.5 + 0.4 * (current_iteration / max(total_iterations, 1)),
            desc=f"Generating... {current_iteration}/{total_iterations}",
        )

        del input_video, input_video_mask, partial_audio_embeds
        torch.cuda.empty_cache()

    progress(0.95, desc="Saving video...")
    video_path = os.path.join(save_path, f"{timestamp}.mp4")
    video_audio_path = os.path.join(save_path, f"{timestamp}_audio.mp4")

    save_videos_grid(sample[:, :, :video_length], video_path, fps=CONFIG["fps"])

    video_clip = VideoFileClip(video_path)
    audio_clip_sub = audio_clip.subclipped(0, video_length / CONFIG["fps"])
    video_clip = video_clip.with_audio(audio_clip_sub)
    video_clip.write_videofile(video_audio_path, codec="libx264", audio_codec="aac", threads=2)

    gc.collect()
    torch.cuda.empty_cache()

    return video_audio_path, seed


with gr.Blocks(title="EchoMimicV3 - Audio-driven Human Animation") as demo:
    gr.Markdown("""
    # 🎭 EchoMimicV3
    **Audio/Text-driven Human Animation Model**

    Upload a portrait image and an audio file to generate animated talking head video.

    | Parameter | Recommended Range |
    |-----------|-------------------|
    | Audio CFG | 2.0 - 3.0 (higher = better lip sync) |
    | Text CFG | 3.0 - 6.0 (higher = better prompt following) |
    | Steps | 20-25 |

    **Requirements:** NVIDIA GPU with 24GB+ VRAM (A100 or RTX 4090 recommended)
    """)

    with gr.Row():
        with gr.Column():
            image = gr.Image(label="📷 Upload Portrait Image", type="filepath", height=300)
            audio = gr.Audio(label="🎤 Upload Audio", type="filepath")

            with gr.Accordion("⚙️ Advanced Settings", open=False):
                prompt = gr.Textbox(
                    label="Prompt",
                    value="",
                    placeholder="Optional: Describe the animation style...",
                    lines=2,
                )
                negative_prompt = gr.Textbox(
                    label="Negative Prompt",
                    value=DEFAULT_NEG_PROMPT,
                    lines=3,
                )
                seed_param = gr.Number(
                    label="Seed (-1 for random)",
                    value=-1,
                )

            generate_btn = gr.Button("🎬 Generate Animation", variant="primary", size="lg")

        with gr.Column():
            video_output = gr.Video(label="🎥 Generated Animation", interactive=False)
            seed_output = gr.Textbox(label="Seed Used", interactive=False)

    generate_btn.click(
        fn=generate,
        inputs=[image, audio, prompt, negative_prompt, seed_param],
        outputs=[video_output, seed_output],
        show_progress=True,
    )

    gr.Markdown("""
    ---
    **Model:** [EchoMimicV3](https://huggingface.co/BadToBest/EchoMimicV3) by Ant Group

    <details>
    <summary>Technical Details</summary>

    - **Parameters:** 1.3B
    - **Base Model:** Wan2.1-Fun-1.3B-InP
    - **Audio Encoder:** wav2vec2-base-960h
    - **License:** Apache-2.0

    </details>
    """)

if __name__ == "__main__":
    # Disable SSR to avoid Node proxy / asyncio fd cleanup noise on Spaces
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)