#!/usr/bin/env python
"""Run the official Wan2.2-Animate-14B example against this quantized repo.

Usage (from the repository root):
    python examples/run_example.py --backend native
    python examples/run_example.py --backend sage_v2 --steps 20
"""

from __future__ import annotations

import argparse
from pathlib import Path

import imageio.v3 as iio
import numpy as np
import torch
from diffusers import AutoencoderKLWan, WanAnimatePipeline
from diffusers.models.attention_dispatch import (
    AttentionBackendName,
    _AttentionBackendRegistry,
    _check_device_cuda_atleast_smXY,
    _check_shape,
)
from diffusers.utils import load_image, load_video

MODEL_DIR = Path(__file__).resolve().parents[1]
EXAMPLE_DIR = Path(__file__).resolve().parent / "animate"
PROMPT = "People in the video are doing actions."


def register_sage_backends() -> None:
    from sageattention import sageattn, sageattn_qk_int8_pv_fp16_triton
    from sageattn3 import sageattn3_blackwell

    def sage_v1(query, key, value, attn_mask=None, is_causal=False, scale=None, **kw):
        if attn_mask is not None:
            raise ValueError("SAGE_ATTN_1 does not support attn_mask")
        return sageattn_qk_int8_pv_fp16_triton(
            query, key, value, tensor_layout="NHD", is_causal=is_causal, sm_scale=scale, smooth_k=True
        )

    def sage_v2(query, key, value, attn_mask=None, is_causal=False, scale=None, **kw):
        if attn_mask is not None:
            raise ValueError("SAGE_ATTN does not support attn_mask")
        return sageattn(query, key, value, tensor_layout="NHD", is_causal=is_causal, sm_scale=scale, smooth_k=True)

    def sage_v3(query, key, value, attn_mask=None, is_causal=False, **kw):
        if attn_mask is not None:
            raise ValueError("SAGE_ATTN_3 does not support attn_mask")
        return sageattn3_blackwell(
            query.transpose(1, 2).contiguous(),
            key.transpose(1, 2).contiguous(),
            value.transpose(1, 2).contiguous(),
            is_causal=is_causal,
        ).transpose(1, 2).contiguous()

    _AttentionBackendRegistry.register(
        AttentionBackendName._SAGE_QK_INT8_PV_FP16_TRITON,
        constraints=[_check_device_cuda_atleast_smXY(8, 0), _check_shape],
    )(sage_v1)
    _AttentionBackendRegistry.register(
        AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA,
        constraints=[_check_device_cuda_atleast_smXY(8, 0), _check_shape],
    )(sage_v2)
    _AttentionBackendRegistry.register(
        AttentionBackendName._SAGE_QK_INT8_PV_FP8_CUDA_SM90,
        constraints=[_check_device_cuda_atleast_smXY(10, 0), _check_shape],
    )(sage_v3)


BACKENDS = {
    "native": "native",
    "sage_v1": "_sage_qk_int8_pv_fp16_triton",
    "sage_v2": "_sage_qk_int8_pv_fp8_cuda",
    "sage_v3": "_sage_qk_int8_pv_fp8_cuda_sm90",
}


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--backend", default="native", choices=sorted(BACKENDS))
    p.add_argument("--steps", type=int, default=20)
    p.add_argument("--segment-frame-length", type=int, default=77)
    p.add_argument("--limit-frames", type=int, default=0)
    p.add_argument("--output", default="")
    p.add_argument("--no-offload", action="store_true", help="Keep all components on GPU (needs ~68 GB).")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    backend = BACKENDS[args.backend]
    if args.backend.startswith("sage"):
        register_sage_backends()

    vae = AutoencoderKLWan.from_pretrained(MODEL_DIR, subfolder="vae", torch_dtype=torch.float32)
    pipe = WanAnimatePipeline.from_pretrained(
        MODEL_DIR,
        torch_dtype=torch.bfloat16,
        vae=vae,
        device_map="balanced",
    )
    if args.no_offload:
        pipe.to("cuda")
    else:
        pipe.reset_device_map()
        pipe.enable_model_cpu_offload(device="cuda")
    pipe.transformer.set_attention_backend(backend)

    image = load_image(str(EXAMPLE_DIR / "src_ref.png"))
    pose = load_video(str(EXAMPLE_DIR / "src_pose.mp4"))
    face = load_video(str(EXAMPLE_DIR / "src_face.mp4"))
    if args.limit_frames:
        pose = pose[: args.limit_frames]
        face = face[: args.limit_frames]

    result = pipe(
        image=image,
        pose_video=pose,
        face_video=face,
        prompt=PROMPT,
        height=720,
        width=1280,
        segment_frame_length=args.segment_frame_length,
        prev_segment_conditioning_frames=1,
        num_inference_steps=args.steps,
        mode="animate",
        guidance_scale=1.0,
        generator=torch.Generator(device="cuda").manual_seed(42),
        output_type="np",
    )
    frames = np.stack([np.asarray(f) for f in result.frames[0]], axis=0)
    output = Path(args.output or f"example_{args.backend}.mp4")
    iio.imwrite(output, frames, fps=30, codec="libx264", quality=8)
    print(f"saved {output} ({frames.shape[0]} frames, {frames.shape[2]}x{frames.shape[1]})")


if __name__ == "__main__":
    main()