app.py

"""EverAnimate — minute-scale human animation via latent flow restoration.

Faithful port of `scripts/everanimate_inference.py` + `test.sh` from
https://github.com/vita-epfl/EverAnimate onto ZeroGPU.

Base model : Wan-AI/Wan2.2-Animate-14B (DiffSynth layout)
Adapter    : epfl-vita/everanimate  ckpts/everanimate-v1-lora32/stage2_480p.safetensors
"""

import os
import sys

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CKPT_DIR = os.path.join(BASE_DIR, "ckpts")
os.makedirs(CKPT_DIR, exist_ok=True)

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", CKPT_DIR)
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
os.environ.setdefault("NUMBA_DISABLE_CUDA", "1")
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

sys.path.insert(0, BASE_DIR)

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

import time  # noqa: E402
import tempfile  # noqa: E402
import traceback  # noqa: E402

import numpy as np  # noqa: E402
import torch  # noqa: E402
import gradio as gr  # noqa: E402
import imageio  # noqa: E402
from PIL import Image  # noqa: E402
from huggingface_hub import hf_hub_download, snapshot_download  # noqa: E402

from diffsynth.utils.data import crop_and_resize, save_video  # noqa: E402
from diffsynth.pipelines.wan_video_svi import WanVideoSviPipeline, ModelConfig  # noqa: E402

# --------------------------------------------------------------------------------------
# Fixed configuration — mirrors test.sh exactly
# --------------------------------------------------------------------------------------
BASE_MODEL_ID = "Wan-AI/Wan2.2-Animate-14B"
LORA_REPO = "epfl-vita/everanimate"
LORA_FILE = "ckpts/everanimate-v1-lora32/stage2_480p.safetensors"

HEIGHT = 480
WIDTH = 832
FPS = 25
FRAMES_PER_CLIP = 77
NUM_OVERLAP_FRAME = 4
NUM_VIDEO_ANCHOR_LATENTS = 4
NUM_MOTION_LATENTS = 1
NUM_MOTION_FRAME = 1
RANDOM_ANCHOR_FRAMES = 3
CFG_SCALE = 1.0
FACE_SIZE = 512
DEFAULT_PROMPT = "视频中的人在做动作"
PREPROCESS_MAX_SECONDS = 30

# --------------------------------------------------------------------------------------
# Weights
# --------------------------------------------------------------------------------------
print("[boot] downloading Wan2.2-Animate-14B base weights …", flush=True)
WAN_DIR = os.path.join(CKPT_DIR, "Wan-AI", "Wan2.2-Animate-14B")
snapshot_download(
    repo_id=BASE_MODEL_ID,
    local_dir=WAN_DIR,
    allow_patterns=[
        "diffusion_pytorch_model*.safetensors",
        "models_t5_umt5-xxl-enc-bf16.pth",
        "Wan2.1_VAE.pth",
        "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
        "process_checkpoint/det/yolov10m.onnx",
        "process_checkpoint/pose2d/vitpose_h_wholebody.onnx/**",
    ],
    max_workers=8,
)
_vitpose_manifest = os.path.join(
    WAN_DIR, "process_checkpoint", "pose2d", "vitpose_h_wholebody.onnx", "end2end.onnx"
)
if not os.path.isfile(_vitpose_manifest):
    raise RuntimeError("Failed to download the required Wan ViTPose checkpoint directory.")
print("[boot] downloading umt5-xxl tokenizer …", flush=True)
snapshot_download(
    repo_id="google/umt5-xxl",
    local_dir=os.path.join(CKPT_DIR, "Wan-AI", "Wan2.1-T2V-1.3B", "google", "umt5-xxl"),
    allow_patterns=["tokenizer*", "special_tokens_map.json", "spiece.model"],
)
print("[boot] downloading EverAnimate LoRA …", flush=True)
LORA_PATH = hf_hub_download(repo_id=LORA_REPO, filename=LORA_FILE)

# --------------------------------------------------------------------------------------
# Pipeline (built on CPU, then moved to CUDA so ZeroGPU can pack the weights)
# --------------------------------------------------------------------------------------
print("[boot] building WanVideoSviPipeline …", flush=True)
_t0 = time.time()
pipe = WanVideoSviPipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cpu",
    model_configs=[
        ModelConfig(
            model_id=BASE_MODEL_ID,
            origin_file_pattern="diffusion_pytorch_model*.safetensors",
            computation_device="cpu",
        ),
        ModelConfig(
            model_id=BASE_MODEL_ID,
            origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth",
            computation_device="cpu",
        ),
        ModelConfig(
            model_id=BASE_MODEL_ID,
            origin_file_pattern="Wan2.1_VAE.pth",
            computation_device="cpu",
        ),
        ModelConfig(
            model_id=BASE_MODEL_ID,
            origin_file_pattern="models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth",
            computation_device="cpu",
        ),
    ],
    redirect_common_files=False,
)
print(f"[boot] models loaded in {time.time() - _t0:.0f}s — fusing LoRA …", flush=True)
pipe.load_lora(pipe.dit, LORA_PATH, alpha=1)

# test.sh flags: --use_image_anchor --num_video_anchor_latents 4 --use_random_frame_anchor
#                --random_anchor_with_user_first --random_anchor_frames 3
#                --num_motion_latents 1 --num_overlap_frame 4 --use_pingpong
pipe.num_video_anchor_latents = NUM_VIDEO_ANCHOR_LATENTS
pipe.use_zero_padding = False
pipe.add_noise_to_motion_latent = False
pipe.motion_latent_shared_noise = False
pipe.bi_sink = False
pipe.enable_anchor_key_focus = False
pipe.remove_pose = False
pipe.mask_anchor_motion = False
pipe.pad_first_clip_with_anchor = False
if pipe.dit is not None:
    pipe.dit.enable_anchor_key_focus = False

print("[boot] moving pipeline to CUDA …", flush=True)
pipe.to("cuda")
pipe.eval()
print(f"[boot] ready in {time.time() - _t0:.0f}s", flush=True)


# --------------------------------------------------------------------------------------
# Helpers (ported from scripts/everanimate_inference.py)
# --------------------------------------------------------------------------------------
def _to_pil_rgb(frame):
    if isinstance(frame, Image.Image):
        return frame.convert("RGB")
    if isinstance(frame, torch.Tensor):
        frame = frame.detach().cpu().numpy()
    if isinstance(frame, np.ndarray):
        arr = frame
        if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
            arr = np.transpose(arr, (1, 2, 0))
        if np.issubdtype(arr.dtype, np.floating):
            max_val = float(np.nanmax(arr)) if arr.size > 0 else 0.0
            arr = np.clip(arr * 255.0, 0, 255) if max_val <= 1.0 else np.clip(arr, 0, 255)
        arr = arr.astype(np.uint8)
        if arr.ndim == 3 and arr.shape[2] == 1:
            arr = arr[:, :, 0]
        return Image.fromarray(arr).convert("RGB")
    raise TypeError(f"Unsupported frame type: {type(frame)}")


def _read_frames(video_path, height, width, max_frames):
    """Decode at most `max_frames` frames, cropped/resized exactly like diffsynth VideoData."""
    frames = []
    reader = imageio.get_reader(video_path)
    try:
        for idx, raw in enumerate(reader):
            if idx >= max_frames:
                break
            frame = Image.fromarray(np.array(raw)).convert("RGB")
            if frame.size != (width, height):
                frame = crop_and_resize(frame, height, width)
            frames.append(frame)
    finally:
        reader.close()
    if not frames:
        raise gr.Error(f"Could not decode any frame from {os.path.basename(video_path)}.")
    return frames


def _pingpong_clip(all_frames, start_frame, num_frames):
    """Ping-pong looping, identical to StreamingAnimateVideoProcessor.load_video_clip."""
    forward = all_frames
    backward = all_frames[-2:0:-1]
    sequence = forward + backward
    n = len(sequence)
    start = start_frame % n
    return [sequence[(start + i) % n] for i in range(num_frames)]


def _encode_single_latent(frame, tiled):
    frame = _to_pil_rgb(frame).resize((WIDTH, HEIGHT))
    tensor = (
        pipe.preprocess_image(frame)
        .to(pipe.device)
        .transpose(0, 1)
        .to(dtype=pipe.torch_dtype, device=pipe.device)
    )
    latent = pipe.vae.encode([tensor], device=pipe.device, tiled=tiled)[0]
    latent = latent.to(dtype=pipe.torch_dtype, device=pipe.device)
    return latent[:, :1]


def _normalize_anchor_latent_count(latent_slices, target_count):
    normalized = list(latent_slices[:target_count])
    while len(normalized) < target_count:
        normalized.append(normalized[-1].clone())
    return normalized


def _build_video_anchor_latent(clip_frames, anchor_image):
    """--use_image_anchor --use_random_frame_anchor --random_anchor_with_user_first branch."""
    user_anchor_latent = _encode_single_latent(anchor_image, tiled=False)
    num_random_slots = max(NUM_VIDEO_ANCHOR_LATENTS - 1, 0)
    total = len(clip_frames)
    if total <= 1 or num_random_slots == 0:
        sampled_latents = []
        selected = []
    else:
        candidate_ids = np.arange(1, total)
        sample_count = min(num_random_slots, len(candidate_ids))
        selected = sorted(
            int(x) for x in np.random.choice(candidate_ids, size=sample_count, replace=False).tolist()
        )
        sampled_latents = [_encode_single_latent(clip_frames[int(i)], tiled=True) for i in selected]
    latent_slices = sampled_latents + [user_anchor_latent.clone()]
    latent_slices = _normalize_anchor_latent_count(latent_slices, NUM_VIDEO_ANCHOR_LATENTS)
    print(f"Random-anchor (random + user first): indices {selected}")
    return torch.concat(latent_slices, dim=1)


def _side_by_side(rgb_frames, pose_frames):
    out = []
    for rgb, pose in zip(rgb_frames, pose_frames):
        rgb = _to_pil_rgb(rgb)
        pose = _to_pil_rgb(pose).resize(rgb.size)
        combined = Image.new("RGB", (rgb.size[0] * 2, rgb.size[1]))
        combined.paste(rgb, (0, 0))
        combined.paste(pose, (rgb.size[0], 0))
        out.append(combined)
    return out


# --------------------------------------------------------------------------------------
# Wan direct pose + face extraction (CPU-only)
# --------------------------------------------------------------------------------------
_wan_preprocessor = None


def _get_wan_preprocessor(progress=None):
    global _wan_preprocessor
    if _wan_preprocessor is None:
        from wan_preprocess import WanDirectPreprocessor

        if progress is not None:
            progress(0.02, desc="Loading Wan YOLO detector (first use may take a moment)")
        _wan_preprocessor = WanDirectPreprocessor(
            detector_path=os.path.join(WAN_DIR, "process_checkpoint", "det", "yolov10m.onnx"),
            vitpose_path=os.path.join(
                WAN_DIR, "process_checkpoint", "pose2d", "vitpose_h_wholebody.onnx"
            ),
            progress=progress,
            prefer_cuda=True,
        )
    return _wan_preprocessor


@spaces.GPU(duration=180, size="xlarge")
def extract_driving_pair(reference_image, source_video, progress=gr.Progress()):
    """Generate a matched Wan pose/face driving pair from one uploaded video."""
    if not reference_image:
        raise gr.Error("Please provide the character image before extracting a driving pair.")
    if not source_video:
        raise gr.Error("Please upload or record a single-person source video to extract its driving pair.")
    try:
        progress(0.01, desc="Preparing Wan pose extractor")
        output_dir = tempfile.mkdtemp(prefix="wan_preprocess_")
        pose_path, face_path = _get_wan_preprocessor(progress=progress).process(
            source_video, reference_image, output_dir, progress=progress
        )
        return pose_path, face_path
    except gr.Error:
        raise
    except (RuntimeError, ValueError) as exc:
        raise gr.Error(str(exc)) from exc
    except Exception as exc:
        traceback.print_exc()
        raise gr.Error(f"Wan preprocessing failed: {exc}") from exc


def use_extracted_driving_pair(pose_path, face_path):
    """Copy the completed extraction outputs into the animation driving inputs."""
    if not pose_path or not face_path:
        raise gr.Error("Extract a pose and face pair before using it as driving input.")
    return pose_path, face_path


# --------------------------------------------------------------------------------------
# Inference
# --------------------------------------------------------------------------------------
def _estimate_duration(
    reference_image=None,
    pose_video=None,
    face_video=None,
    num_chunks=2,
    num_inference_steps=20,
    seed=42,
    prompt=DEFAULT_PROMPT,
    sigma_shift=5.0,
    show_pose=False,
    *args,
    **kwargs,
):
    # Measured on ZeroGPU xlarge: 1 chunk x 20 steps = 147 s (~6.5 s/step + ~17 s
    # of VAE encode/decode & anchor work). Keep this tight: xlarge costs 2x quota.
    chunks = int(num_chunks or 1)
    steps = int(num_inference_steps or 20)
    raw = 25 + chunks * 4 + chunks * (steps * 6.6 + 20)
    return int(min(1500, raw * 1.12))


@spaces.GPU(duration=_estimate_duration, size="xlarge")
def animate(
    reference_image: str,
    pose_video: str,
    face_video: str,
    num_chunks: int = 2,
    num_inference_steps: int = 20,
    seed: int = 42,
    prompt: str = DEFAULT_PROMPT,
    sigma_shift: float = 5.0,
    show_pose: bool = False,
    progress=gr.Progress(track_tqdm=True),
) -> str:
    """Animate a reference character with a Wan-Animate pose + face driving pair.

    Args:
        reference_image: Path to the still image of the character to animate.
        pose_video: Path to the Wan-Animate skeleton (pose) driving video.
        face_video: Path to the matching 512x512 cropped-face driving video.
        num_chunks: Number of 77-frame chunks to generate (~3 s of video each).
        num_inference_steps: Flow-matching denoising steps per chunk.
        seed: Base seed; chunk i uses seed * i.
        prompt: Text prompt fed to the Wan text encoder.
        sigma_shift: Flow-matching timestep shift.
        show_pose: Also render the driving skeleton next to the result.

    Returns:
        Path to the generated mp4 video.
    """
    if not reference_image:
        raise gr.Error("Please provide a reference character image.")
    if not pose_video or not face_video:
        raise gr.Error("Please provide both a pose video and a matching face video.")

    num_chunks = max(1, int(num_chunks))
    num_inference_steps = max(1, int(num_inference_steps))
    seed = int(seed)
    prompt = (prompt or "").strip() or DEFAULT_PROMPT

    np.random.seed(seed if seed != 0 else 42)
    torch.manual_seed(seed)

    stride = FRAMES_PER_CLIP - NUM_OVERLAP_FRAME
    frames_needed = (num_chunks - 1) * stride + FRAMES_PER_CLIP

    print(f"[run] {num_chunks} chunk(s) x {num_inference_steps} steps, need {frames_needed} driving frames")
    pose_frames_all = _read_frames(pose_video, HEIGHT, WIDTH, frames_needed)
    face_frames_all = _read_frames(face_video, FACE_SIZE, FACE_SIZE, frames_needed)

    input_image = crop_and_resize(Image.open(reference_image).convert("RGB"), HEIGHT, WIDTH)
    anchor_image = input_image

    all_video_frames = []
    all_pose_frames = []
    current_input_image = input_image
    prev_last_latent = None
    video_anchor_latent = None
    prev_clip_latent_trajectory = None

    t_start = time.time()
    for clip_idx in range(num_chunks):
        start_frame = clip_idx * stride
        animate_pose_video = _pingpong_clip(pose_frames_all, start_frame, FRAMES_PER_CLIP)
        animate_face_video = _pingpong_clip(face_frames_all, start_frame, FRAMES_PER_CLIP)
        pose_frames_pil = [_to_pil_rgb(f) for f in animate_pose_video]

        video_anchor_for_call = (
            video_anchor_latent.clone() if (clip_idx > 0 and video_anchor_latent is not None) else None
        )

        print(f"[run] chunk {clip_idx + 1}/{num_chunks} …", flush=True)
        out = pipe(
            prompt=prompt,
            seed=clip_idx * seed,
            tiled=False,
            input_image=current_input_image,
            animate_pose_video=animate_pose_video,
            animate_face_video=animate_face_video,
            anchor=anchor_image,
            num_frames=FRAMES_PER_CLIP + 4 * NUM_VIDEO_ANCHOR_LATENTS,
            height=HEIGHT,
            width=WIDTH,
            num_inference_steps=num_inference_steps,
            sigma_shift=float(sigma_shift),
            cfg_scale=CFG_SCALE,
            prev_last_latent=prev_last_latent,
            video_anchor_latent=video_anchor_for_call,
            num_motion_latents=NUM_MOTION_LATENTS,
            use_face_anchor=False,
            debug_save_latents=False,
            clip_idx=clip_idx,
            num_video_anchor_latents=NUM_VIDEO_ANCHOR_LATENTS,
            prev_clip_latent_trajectory=prev_clip_latent_trajectory,
        )

        video_clip = out["video"]
        prev_last_latent = out["prev_last_latent"]
        prev_clip_latent_trajectory = out.get("latent_trajectory", None)
        video_frames = [_to_pil_rgb(f) for f in video_clip]

        if clip_idx == 0:
            video_anchor_latent = _build_video_anchor_latent(video_frames, anchor_image)
            all_video_frames.extend(video_frames)
            all_pose_frames.extend(pose_frames_pil)
        else:
            all_video_frames.extend(video_frames[NUM_OVERLAP_FRAME:])
            all_pose_frames.extend(pose_frames_pil[NUM_OVERLAP_FRAME:])

        current_input_image = video_frames[-NUM_MOTION_FRAME:]
        print(
            f"[run] chunk {clip_idx + 1} done — {len(all_video_frames)} frames "
            f"({time.time() - t_start:.0f}s elapsed)",
            flush=True,
        )

    frames = _side_by_side(all_video_frames, all_pose_frames) if show_pose else all_video_frames
    out_path = os.path.join(tempfile.mkdtemp(), "everanimate.mp4")
    save_video(frames, out_path, fps=FPS, quality=8)
    print(
        f"[run] finished: {len(all_video_frames)} frames "
        f"({len(all_video_frames) / FPS:.1f}s) in {time.time() - t_start:.0f}s",
        flush=True,
    )
    return out_path


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
EX = os.path.join(BASE_DIR, "examples")

DESCRIPTION = """
# EverAnimate — minute-scale human animation

<a href="https://huggingface.co/epfl-vita/everanimate">model</a> ·
<a href="https://github.com/vita-epfl/EverAnimate">code</a> ·
<a href="https://arxiv.org/abs/2605.15042">paper</a>

**EverAnimate** (EPFL VITA) is a LoRA on top of **Wan2.2-Animate-14B** that keeps chunk-based
human animation stable over long horizons: *Persistent Latent Propagation* carries an identity /
motion memory across chunks, and *Restorative Flow Matching* corrects drifted latent trajectories.

Give it a **character image** plus a **Wan-Animate pose + face driving pair** and it animates your
character chunk by chunk at 480×832 / 25 fps. Short driving clips are ping-pong looped, so a few
seconds of motion can drive a much longer video.
"""

NOTE = """
> **Driving inputs.** The pose video is the Wan-Animate skeleton rendering and the face video is the
> matching 512×512 face crop, both produced by
> [Wan2.2-Animate's `preprocess_data.py`](https://github.com/Wan-Video/Wan2.2). Two ready-made
> driving pairs from the EverAnimate release are provided below — swap in **your own character
> image** and keep a provided pair to see the model transfer that motion onto your character.
"""

SOURCE_VIDEO_CSS = """
#source-motion-video button,
#source-motion-video button svg {
    color: #f8fafc !important;
    fill: currentColor !important;
}
"""


with gr.Blocks(title="EverAnimate") as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            reference_image = gr.Image(
                label="Character image (reference)", type="filepath", height=260
            )
            with gr.Row():
                pose_video = gr.Video(label="Pose driving video", height=200)
                face_video = gr.Video(label="Face driving video (512×512)", height=200)
            with gr.Accordion("Create driving pair from a video", open=False):
                gr.Markdown(
                    "Upload a clear **single-person** video (up to "
                    f"{PREPROCESS_MAX_SECONDS} seconds). Wan preprocessing extracts a matched "
                    "skeleton and 512×512 face crop at 25 fps."
                )
                source_video = gr.Video(
                    label="Source motion video",
                    sources=["upload", "webcam"],
                    height=200,
                    format="mp4",
                    elem_id="source-motion-video",
                )
                extract_btn = gr.Button("Extract pose and face", variant="secondary")
                with gr.Row():
                    extracted_pose = gr.Video(label="Extracted pose preview", height=180)
                    extracted_face = gr.Video(label="Extracted face preview", height=180)
                use_extracted_btn = gr.Button("Use extracted driving pair")
            num_chunks = gr.Slider(
                1, 6, value=2, step=1,
                label="Chunks to generate",
                info="Each chunk is 77 frames ≈ 2.9 s of video at 25 fps (~2.5 min of GPU each).",
            )
            run_btn = gr.Button("Animate", variant="primary")
            gr.Markdown(NOTE)
        with gr.Column(scale=1):
            output_video = gr.Video(label="EverAnimate result", height=380, autoplay=True)
            with gr.Accordion("Advanced settings", open=False):
                num_inference_steps = gr.Slider(
                    4, 30, value=20, step=1, label="Denoising steps per chunk"
                )
                seed = gr.Slider(
                    0, 100000, value=42, step=1, label="Seed",
                    info="Faithful to the reference script: chunk i is sampled with seed × i.",
                )
                sigma_shift = gr.Slider(1.0, 12.0, value=5.0, step=0.5, label="Sigma shift")
                prompt = gr.Textbox(
                    label="Prompt", value=DEFAULT_PROMPT,
                    info="Default is the paper's prompt: “the person in the video is performing an action”.",
                )
                show_pose = gr.Checkbox(
                    value=False, label="Show driving skeleton side by side"
                )

    inputs = [
        reference_image, pose_video, face_video, num_chunks,
        num_inference_steps, seed, prompt, sigma_shift, show_pose,
    ]
    run_btn.click(fn=animate, inputs=inputs, outputs=output_video)
    extract_btn.click(
        fn=extract_driving_pair,
        inputs=[reference_image, source_video],
        outputs=[extracted_pose, extracted_face],
    )
    use_extracted_btn.click(
        fn=use_extracted_driving_pair,
        inputs=[extracted_pose, extracted_face],
        outputs=[pose_video, face_video],
    )

    gr.Examples(
        examples=[
            [
                os.path.join(EX, "demo_character.png"),
                os.path.join(EX, "demo_pose.mp4"),
                os.path.join(EX, "demo_face.mp4"),
            ],
            [
                os.path.join(EX, "sample_character.png"),
                os.path.join(EX, "sample_pose.mp4"),
                os.path.join(EX, "sample_face.mp4"),
            ],
        ],
        inputs=[reference_image, pose_video, face_video],
        label="Official EverAnimate demo inputs",
    )

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