Upload 7 files

README.md

@@ -1,13 +1,70 @@
 ---
-title: Bbs231
-emoji: 😻
-colorFrom: gray
+title: Wan-Animate Character Replacement
+emoji: 🎭
+colorFrom: indigo
 colorTo: pink
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.12'
+python_version: "3.12"
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: Replace a person in a video with the character from a photo
+startup_duration_timeout: 1h
+suggested_hardware: zero-a10g
+models:
+  - Wan-AI/Wan2.2-Animate-14B
+  - Wan-AI/Wan2.2-Animate-14B-Diffusers
+  - Kijai/WanVideo_comfy
+  - facebook/sam2.1-hiera-large
+  - PekingU/rtdetr_v2_r50vd
+  - facebook/mask2former-swin-base-coco-panoptic
+  - Qwen/Qwen3.5-4B
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Wan2.2-Animate character replacement on ZeroGPU
+
+Photo + video → the person in the video is replaced by the character from the photo. The new character follows the
+original person's **skeleton** (body proportions, position, every movement) and **face** (expressions, lips, head
+turns) frame by frame; only the person's region is redrawn, the scene and objects stay those of the video.
+
+* **Model:** [Wan2.2-Animate-14B](https://github.com/Wan-Video/Wan2.2) in replacement mode through diffusers'
+  `WanAnimatePipeline`, with the official Relighting LoRA and the Lightx2v distillation LoRA fused at startup. LoRA
+  file and strength (1.2), sampler (DPM++ SDE), 6 steps and shift 5 follow Kijai's Wan-Animate example workflow. The
+  transformer is read from the official bf16 checkpoint (34.5 GB; the diffusers copy is fp32, 69 GB) and renamed with
+  diffusers' own converter (`wan_animate.py`).
+* **Pose and face** (`pose.py`): ViTPose-H whole-body, the checkpoint of the official preprocessing, run in PyTorch on
+  the GPU; skeleton drawing, face crops and keypoint decoding are ported from the official code and checked against
+  it (identical images, keypoints within 0.001 px). The skeleton is not retargeted, as in the official replacement
+  mode, so the proportions stay those of the original person. Face crops come from up to 2x sharper frames.
+* **Masks** (`masking.py`): RT-DETRv2 finds the person on the first frame (or the one you click), SAM 2.1 tracks it,
+  Mask2Former's "person" segmentation cuts off what SAM 2.1 glues to it (chair, wall, room). The redrawn region is
+  that mask grown by 10 px and cut into 32-px blocks (Kijai's GrowMask + BlockifyMask) aligned with the latent grid:
+  a mask edge running through a latent leaves blacked-out pixels in a "kept" latent, which shows as a dark outline.
+* **Reference:** the character is cut out of the photo (or taken from a PNG's transparency) and pasted onto the video's
+  own first frame, blurred and with the person removed; the Relighting LoRA was trained on characters pasted onto
+  other scenes, and a flat black backdrop comes back as a dark outline.
+* **"Keep the original pixels"** pastes the generated person (feathered, around the original person only) over the
+  untouched source frames.
+* **AI scene reading** (`captioner.py`, Qwen3.5-4B): before the replacement it looks at the photo and four frames,
+  names the action ("licks a lollipop"), finds what the person holds or has at the mouth (with boxes, checked against
+  Mask2Former so a box on a hand or a face is dropped) and writes the prompt: the character from the photo doing
+  exactly that. Wan-Animate's skeleton and face crops carry the body and the expression, not a tongue on a lollipop.
+* **Objects to keep:** what the AI finds, and anything clicked on the first frame. SAM 2.1
+  tracks them as separate objects; latents that are mostly object are given to the model as they are, and the
+  original object, plus a few pixels of what touches it (lips, tongue, fingers), is pasted back on top. Otherwise the
+  model has to redraw the object it cannot see, which comes out doubled or bent.
+* **Expression strength** scales what the face adapter adds to the video (Kijai's `face_strength`).
+* **ZeroGPU:** one `xlarge` call (RTX Pro 6000 Blackwell, 96 GB). ≈2.5 min of GPU for 5 s at 480p / 30 fps.
+* All files sit in the repository root, so the Space can be uploaded file by file from the web UI.
+
+## Settings (Space → Settings → Variables)
+
+| Variable | Default | Meaning |
+|---|---|---|
+| `WAN_LIGHTX2V` / `WAN_RELIGHT` | `1.2` / `1.0` | LoRA strengths (`WAN_LIGHTX2V=0` → original 20 steps with DPM++) |
+| `WAN_ATTENTION` | empty (SDPA) | a diffusers attention backend, e.g. `sage_hub` (SageAttention, unverified publisher) |
+| `WAN_PERSON_PARSER` | `facebook/mask2former-swin-base-coco-panoptic` | `none` turns the mask refinement off |
+| `WAN_SCENE_MODEL` | `Qwen/Qwen3.5-4B` | vision-language model of the AI scene reading; `none` turns it off |
+| `WAN_MAX_SECONDS` / `WAN_MAX_GPU_SECONDS` | `15` / `900` | input length / per-call GPU limits |
+| `WAN_TFLOPS_LINEAR` / `WAN_TFLOPS_ATTENTION` / `WAN_STARTUP_SECONDS` | `330` / `260` / `25` | cost model for the requested ZeroGPU duration |

app.py

@@ -0,0 +1,673 @@
+"""Character replacement on Hugging Face ZeroGPU with Wan2.2-Animate.
+
+The person in a video is replaced by the character from a photo. Wan-Animate's replacement mode follows the original
+person's skeleton (body proportions, positions, motion) and face (expressions, lips, head turns) frame by frame and
+regenerates only that person's region; the rest of every frame is the original video.
+"""
+
+import spaces  # ZeroGPU: import before torch
+
+import math
+import os
+import random
+import ssl
+import subprocess
+import tempfile
+import threading
+import time
+import urllib.request
+from pathlib import Path
+
+import certifi
+import gradio as gr
+import imageio_ffmpeg
+import numpy as np
+import torch
+from PIL import Image, ImageDraw, ImageOps
+
+from captioner import SceneReader
+from masking import PersonMasker, adjust, blockify, blocky, feathered, mostly, place_character, scene_backdrop
+from pose import PoseExtractor
+from wan_animate import (
+    DEFAULT_PROMPT,
+    PREV_SEGMENT_FRAMES,
+    SAMPLERS,
+    SEGMENT_FRAMES,
+    load_pipeline,
+    segment_count,
+    set_face_strength,
+    set_sampler,
+)
+
+# Read when a ZeroGPU worker first allocates; limits fragmentation of the large activation buffers.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+# --------------------------------------------------------------------------------------------------------------
+# Configuration (environment variables can be set in the Space settings)
+# --------------------------------------------------------------------------------------------------------------
+
+LORA_STRENGTHS = {
+    "lightx2v": float(os.getenv("WAN_LIGHTX2V", "1.2")),  # step distillation (1.2 as in Kijai's workflow); 0 = off
+    "relight": float(os.getenv("WAN_RELIGHT", "1.0")),  # official replacement LoRA: scene lighting on the character
+}
+DISTILLED = LORA_STRENGTHS["lightx2v"] > 0
+DEFAULT_STEPS = 6 if DISTILLED else 20
+DEFAULT_SHIFT = 5.0  # official sample_shift, also Kijai's
+DEFAULT_SAMPLER = "DPM++ SDE" if DISTILLED else "DPM++"
+DEVICE = os.getenv("WAN_DEVICE", "cuda")
+GPU_SIZE = os.getenv("WAN_GPU_SIZE", "xlarge")  # 52 GB of weights do not fit a `large` (48 GB) slice
+ATTENTION = os.getenv("WAN_ATTENTION", "")  # a diffusers attention backend, e.g. "sage_hub"; empty = PyTorch SDPA
+MAX_SECONDS = float(os.getenv("WAN_MAX_SECONDS", "15"))
+CONCURRENCY = int(os.getenv("WAN_CONCURRENCY", "2"))
+
+# Rough throughput of a full RTX Pro 6000 Blackwell, only used to request a ZeroGPU duration. The log prints the
+# measured times; tune these if requests are cut off or wait too long in the queue.
+TFLOPS_LINEAR = float(os.getenv("WAN_TFLOPS_LINEAR", "330"))
+TFLOPS_ATTENTION = float(os.getenv("WAN_TFLOPS_ATTENTION", "260"))
+VAE_SECONDS_PER_MEGAPIXEL_FRAME = float(os.getenv("WAN_VAE_SEC_PER_MPX_FRAME", "0.15"))
+STARTUP_SECONDS = float(os.getenv("WAN_STARTUP_SECONDS", "25"))  # worker spawn + moving 52 GB to the GPU
+MAX_GPU_SECONDS = int(os.getenv("WAN_MAX_GPU_SECONDS", "900"))
+SCENE_MODEL = os.getenv("WAN_SCENE_MODEL", "Qwen/Qwen3.5-4B")  # "none" turns the AI scene reading off
+SCENE_SECONDS = 25.0  # moving the reader to the GPU and back + one answer
+# Weights applied to every video token per step: the 40 blocks without the cross-attention K/V projections (those
+# only see the 769 text/image tokens) plus the face adapter.
+ACTIVE_PARAMS = 12.4e9
+
+FPS_CHOICES = [16, 24, 30]
+DEFAULT_FPS = 30  # Wan-Animate's frame rate; a 5 s clip takes 2 windows at 16 fps as well as at 30 fps
+# Target areas as (height, width); sides follow the video's aspect ratio and are multiples of 16.
+QUALITIES = {
+    "480p · recommended": (480, 832),
+    "720p · sharper, ~4x slower": (720, 1280),
+}
+DEFAULT_QUALITY = "480p · recommended"
+# Redrawn region: the person grown by MASK_GROWTH, then cut into blocks of this many pixels at 480p (scaled with the
+# resolution); None = the person's whole box (the official default). Kijai's workflow: grow 10 px, 32 px blocks.
+MASK_SHAPES = {
+    "Blocks (recommended)": 32,
+    "Small blocks: less of the background redrawn": 16,
+    "Box: most freedom, redraws the box around the person": None,
+}
+DEFAULT_MASK_SHAPE = next(iter(MASK_SHAPES))
+MASK_GROWTH = 10 / 480
+
+# --------------------------------------------------------------------------------------------------------------
+# Models (loaded once at startup; ZeroGPU moves everything placed on "cuda" to the GPU of each request)
+# --------------------------------------------------------------------------------------------------------------
+
+print("Loading Wan2.2-Animate ...")
+pipe = load_pipeline(LORA_STRENGTHS).to(DEVICE)
+if ATTENTION:
+    try:
+        pipe.transformer.set_attention_backend(ATTENTION)
+        print(f"Attention backend: {ATTENTION}")
+    except Exception as e:  # noqa: BLE001 - SDPA still works
+        print(f"Attention backend {ATTENTION} unavailable, using SDPA: {e!r}")
+masker = PersonMasker().to(DEVICE)
+poser = PoseExtractor().to(DEVICE)
+
+
+def load_reader() -> SceneReader | None:
+    if SCENE_MODEL.lower() == "none":
+        return None
+    try:
+        return SceneReader(SCENE_MODEL)  # stays on the CPU until used
+    except Exception as e:  # noqa: BLE001 - the app still works with clicks and hand-written prompts
+        print(f"Scene reader {SCENE_MODEL} unavailable: {e!r}")
+        return None
+
+
+reader = load_reader()
+print("Models loaded.")
+
+# --------------------------------------------------------------------------------------------------------------
+# Geometry and cost model
+# --------------------------------------------------------------------------------------------------------------
+
+
+def frame_size(video_width: int, video_height: int, quality: str) -> tuple[int, int]:
+    """(height, width): the video's aspect ratio at the target area, both sides multiples of 16."""
+    area_h, area_w = QUALITIES[quality]
+    aspect = video_height / video_width
+    height = max(16, round(math.sqrt(area_h * area_w * aspect) / 16) * 16)
+    width = max(16, round(math.sqrt(area_h * area_w / aspect) / 16) * 16)
+    return height, width
+
+
+def map_click(point, source_size, width: int, height: int) -> list[float]:
+    """Normalised click on the source frame -> pixel position in the frame after `prepare_driving_video`'s
+    cover-resize + centre crop."""
+    src_w, src_h = source_size
+    scale = max(width / src_w, height / src_h)
+    return [point[0] * src_w * scale - (src_w * scale - width) / 2, point[1] * src_h * scale - (src_h * scale - height) / 2]
+
+
+def estimate_gpu_seconds(num_frames: int, height: int, width: int, steps: int, scene: bool = False) -> float:
+    cfg = pipe.transformer.config
+    dim = cfg.num_attention_heads * cfg.attention_head_dim
+    tokens = (1 + (SEGMENT_FRAMES - 1) // 4 + 1) * (height // 16) * (width // 16)  # reference + window latents
+    forward = (
+        2 * ACTIVE_PARAMS * tokens / (TFLOPS_LINEAR * 1e12)
+        + 4 * dim * cfg.num_layers * tokens**2 / (TFLOPS_ATTENTION * 1e12)
+    )
+    vae = VAE_SECONDS_PER_MEGAPIXEL_FRAME * height * width / 1e6 * SEGMENT_FRAMES * 3  # pose + scene in, video out
+    per_segment = steps * forward + vae
+    analysis = 5.0 + 0.12 * num_frames  # SAM 2.1 tracking + Mask2Former + ViTPose per frame
+    if scene and reader is not None:
+        analysis += SCENE_SECONDS
+    return STARTUP_SECONDS + analysis + 5.0 + segment_count(num_frames) * per_segment
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Media helpers (CPU)
+# --------------------------------------------------------------------------------------------------------------
+
+FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()
+
+
+def ffmpeg(*args: str) -> None:
+    subprocess.run([FFMPEG, "-y", "-hide_banner", "-loglevel", "error", *args], check=True, capture_output=True)
+
+
+def temp_path(suffix: str) -> str:
+    return tempfile.NamedTemporaryFile(suffix=suffix, delete=False).name
+
+
+def video_metadata(path: str) -> dict:
+    reader = imageio_ffmpeg.read_frames(path)
+    meta = next(reader)
+    reader.close()
+    return meta
+
+
+def first_frame(path: str) -> np.ndarray:
+    reader = imageio_ffmpeg.read_frames(path)
+    width, height = next(reader)["size"]
+    frame = np.frombuffer(next(reader), np.uint8).reshape(height, width, 3).copy()
+    reader.close()
+    return frame
+
+
+def prepare_image(path: str) -> str:
+    """EXIF-rotated, at most 2048 px, PNG (keeping an alpha channel, which doubles as the character mask)."""
+    image = ImageOps.exif_transpose(Image.open(path))
+    image = image.convert("RGBA" if image.mode in ("RGBA", "LA", "P") else "RGB")
+    image.thumbnail((2048, 2048), Image.Resampling.LANCZOS)
+    out = temp_path(".png")
+    image.save(out)
+    return out
+
+
+def prepare_driving_video(path: str, max_seconds: float, fps: int, width: int, height: int,
+                          audio: bool = True) -> tuple[str, int, str | None]:
+    """Trims, resamples and cover-crops the video to `width` x `height` (on the CPU, before a GPU is held);
+    extracts its audio track if there is one and `audio` is set."""
+    out = temp_path(".mp4")
+    resize = f"scale={width}:{height}:force_original_aspect_ratio=increase:flags=bicubic,crop={width}:{height}"
+    seconds = f"{max_seconds:.3f}"
+    ffmpeg("-i", path, "-t", seconds, "-vf", f"fps={fps},{resize}", "-an",
+           "-c:v", "libx264", "-preset", "veryfast", "-crf", "10", "-pix_fmt", "yuv444p", out)
+    num_frames = imageio_ffmpeg.count_frames_and_secs(out)[0]
+    track = None
+    if audio and video_metadata(path).get("audio_codec"):
+        track = temp_path(".m4a")
+        try:
+            ffmpeg("-i", path, "-t", seconds, "-vn", "-c:a", "aac", "-b:a", "192k", track)
+        except subprocess.CalledProcessError:
+            track = None
+    return out, num_frames, track
+
+
+def read_frames(path: str) -> list[np.ndarray]:
+    reader = imageio_ffmpeg.read_frames(path)
+    width, height = next(reader)["size"]
+    return [np.frombuffer(frame, np.uint8).reshape(height, width, 3) for frame in reader]
+
+
+def write_video(frames: np.ndarray, fps: int, audio_path: str | None = None) -> str:
+    """uint8 [T, H, W, 3] -> mp4 (cropped to even sides for yuv420p), muxed with `audio_path` if given."""
+    height, width = frames.shape[1] // 2 * 2, frames.shape[2] // 2 * 2
+    silent = temp_path(".mp4")
+    writer = imageio_ffmpeg.write_frames(
+        silent, (width, height), fps=fps, codec="libx264", quality=None, macro_block_size=2,
+        output_params=["-crf", "16", "-preset", "veryfast", "-movflags", "+faststart"],
+    )
+    writer.send(None)
+    for frame in frames:
+        writer.send(np.ascontiguousarray(frame[:height, :width]))
+    writer.close()
+    if audio_path is None:
+        return silent
+    final = temp_path(".mp4")
+    try:
+        ffmpeg("-i", silent, "-i", audio_path, "-map", "0:v:0", "-map", "1:a:0",
+               "-c:v", "copy", "-c:a", "copy", "-shortest", "-movflags", "+faststart", final)
+        return final
+    except subprocess.CalledProcessError:
+        return silent
+
+
+def flatten(image: Image.Image) -> Image.Image:
+    """RGB with transparent areas on black."""
+    if image.mode in ("RGBA", "LA", "P"):
+        rgba = image.convert("RGBA")
+        return Image.alpha_composite(Image.new("RGBA", rgba.size, (0, 0, 0, 255)), rgba).convert("RGB")
+    return image.convert("RGB")
+
+
+def near(box, person, slack: float = 0.25) -> bool:
+    """An object box overlapping the person's box grown by `slack` of its size, and not bigger than it."""
+    if person is None:
+        return False
+    w, h = person[2] - person[0], person[3] - person[1]
+    grown = (person[0] - slack * w, person[1] - slack * h, person[2] + slack * w, person[3] + slack * h)
+    overlaps = box[0] < grown[2] and box[2] > grown[0] and box[1] < grown[3] and box[3] > grown[1]
+    return overlaps and (box[2] - box[0]) * (box[3] - box[1]) <= w * h
+
+
+# --------------------------------------------------------------------------------------------------------------
+# GPU work
+# --------------------------------------------------------------------------------------------------------------
+
+
+def gpu_duration(image_path, video_path, face_path, num_frames, fps, prompt, steps, shift, sampler, seed,
+                 keep_background, click, objects, mask_edge, mask_shape, contact, face_strength, use_ai, audio_path,
+                 progress=None):
+    try:
+        width, height = video_metadata(video_path)["size"]
+        seconds = estimate_gpu_seconds(num_frames, height, width, steps, scene=use_ai)
+        return min(MAX_GPU_SECONDS, int(1.2 * seconds) + 10)
+    except Exception:  # noqa: BLE001 - never fail the scheduling call
+        return 300
+
+
+@spaces.GPU(duration=gpu_duration, size=GPU_SIZE)
+def replace_on_gpu(image_path, video_path, face_path, num_frames, fps, prompt, steps, shift, sampler, seed,
+                   keep_background, click, objects, mask_edge, mask_shape, contact, face_strength, use_ai,
+                   audio_path, progress=gr.Progress(track_tqdm=True)):
+    timings = {}
+    started = time.perf_counter()
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()  # workers are reused across requests
+
+    def lap(name):
+        nonlocal started
+        timings[name] = time.perf_counter() - started
+        started = time.perf_counter()
+
+    reference = Image.open(image_path)
+    frames = read_frames(video_path)[:num_frames]  # already at the generation size
+    face_frames = read_frames(face_path)[: len(frames)] if face_path else None  # the same frames, sharper
+    if face_frames is not None and len(face_frames) < len(frames):
+        face_frames = None
+    height, width = frames[0].shape[:2]
+    pil_frames = [Image.fromarray(f) for f in frames]
+    lap("load inputs")
+
+    # The AI looks at the photo and a few frames: the action (for the prompt) and what the person holds or has at
+    # the mouth (kept from the original video like clicked objects, after a check that it is not the person).
+    ai = {"action": "", "objects": [], "prompt": ""}
+    person_box = masker.person_box(pil_frames[0], click)
+    if use_ai and reader is not None and person_box is not None:
+        picks = np.linspace(0, len(frames) - 1, num=min(4, len(frames))).round().astype(int)
+        try:
+            ai = reader(reference, [pil_frames[i] for i in picks], person_box, DEVICE)
+            print("Scene:", {k: v for k, v in ai.items() if k != "raw"})
+        except Exception as e:  # noqa: BLE001 - the replacement works without it
+            print(f"Scene reading failed: {e!r}")
+        lap("scene")
+    found = [{"box": o["box"], "name": o["name"], "check": True}
+             for o in ai["objects"] if near(o["box"], person_box)][:MAX_OBJECTS]
+    clicked = [{"point": xy, "name": "clicked object"} for xy in objects]
+
+    try:
+        # The person's outline on every frame, and the objects to keep (minus from the person).
+        person, things, box, kept = masker.video_masks(pil_frames, click, clicked + found)
+        ref_mask = masker.image_mask(reference)  # uses the alpha channel when the photo is a cut-out
+        # The character alone, on the video's own scene: no trace of the photo's room, and no flat black around
+        # her (which comes back as a dark outline).
+        ref_image = place_character(flatten(reference), ref_mask, scene_backdrop(frames[0], person[0]))
+    except ValueError as e:
+        raise gr.Error(str(e)) from e
+    if person.mean() < 1e-4:
+        raise gr.Error("The person to replace was lost on every frame: click on them on the first frame.")
+    lap("masks")
+
+    skeletons, faces, _ = poser(frames, person, face_frames)  # the original person's pose and face, frame by frame
+    del face_frames
+    lap("pose")
+
+    side = min(height, width)
+    grow = max(2, round(MASK_GROWTH * side) + int(mask_edge))
+    region = adjust(person, grow, masker.device)  # what Wan-Animate redraws; everything else is given to it
+    block = MASK_SHAPES[mask_shape]
+    if block is None:
+        region = blocky(region, 1, 1)
+    else:  # whole 16-px patches only, so no latent is half person, half kept background
+        region = blockify(region, max(16, round(block * side / 480 / 16) * 16), masker.device)
+    if things.any():  # latents that are mostly a clicked object stay as they are: the model draws around it
+        region &= ~mostly(things, 8, 0.5, masker.device)
+    scene = [Image.fromarray(np.where(r[..., None], 0, f)) for f, r in zip(frames, region)]
+    region_images = [Image.fromarray(r.astype(np.uint8) * 255) for r in region]
+
+    set_sampler(pipe, sampler, float(shift))
+    set_face_strength(pipe, face_strength)
+    torch.manual_seed(seed)  # the SDE sampler draws fresh noise at every step
+    try:
+        video = pipe(
+            image=ref_image,
+            pose_video=[Image.fromarray(s) for s in skeletons],
+            face_video=[Image.fromarray(f) for f in faces],
+            background_video=scene,
+            mask_video=region_images,
+            prompt=prompt or ai["prompt"] or DEFAULT_PROMPT,
+            height=height,
+            width=width,
+            segment_frame_length=SEGMENT_FRAMES,
+            prev_segment_conditioning_frames=PREV_SEGMENT_FRAMES,
+            num_inference_steps=steps,
+            guidance_scale=1.0,
+            mode="replace",
+            generator=torch.Generator("cpu").manual_seed(seed),
+            output_type="pt",
+        ).frames[0]  # [T, 3, H, W] in [0, 1]
+    except torch.OutOfMemoryError as e:
+        raise gr.Error("Out of GPU memory: pick 480p or a shorter clip.") from e
+    lap("generate")
+
+    device = video.device
+    original = torch.from_numpy(np.stack(frames)).to(device).permute(0, 3, 1, 2).float() / 255
+    alpha = None
+    if keep_background:  # paste the generated person over the untouched source frames
+        # Around the person only (fading out within the margin the region adds), so the blocks of redrawn
+        # background are not used; with the box, everything up to near its edge.
+        feather = max(2, grow // 2)
+        inner = person if block is not None else adjust(region, -2 * feather, masker.device)
+        alpha = feathered(inner, grow=feather, device=device)
+    if things.any():  # the clicked objects stay the original ones, with what touches them (lips, tongue, fingers)
+        front = feathered(adjust(things, int(contact), masker.device), grow=1, device=device)
+        alpha = 1 - front if alpha is None else alpha * (1 - front)
+        del front
+    if alpha is not None:
+        video = video.float() * alpha + original * (1 - alpha)
+        del alpha
+    result = (video.clamp(0, 1) * 255).round().to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()
+    out_path = write_video(result, fps, audio_path)
+
+    preview = np.stack(frames).copy()
+    preview[region] = preview[region] // 2 + np.array([0, 0, 127], np.uint8)
+    preview[things] = preview[things] // 2 + np.array([0, 127, 0], np.uint8)
+    drawn = np.stack(skeletons)
+    preview = np.where(drawn.any(-1, keepdims=True), drawn, preview)
+    preview_path = write_video(preview, fps)
+    face_path = write_video(np.stack([np.asarray(Image.fromarray(f).resize((256, 256))) for f in faces]), fps)
+    ref_path = temp_path(".png")
+    ref_image.save(ref_path)
+    lap("encode mp4")
+
+    vram = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else 0.0
+    info = (
+        f"**{width // 2 * 2}×{height // 2 * 2}**, {len(frames)} frames @ {fps} fps, "
+        f"{segment_count(len(frames))} window(s), {steps} steps, seed `{seed}`, expression {face_strength:g} · "
+        f"peak VRAM {vram:.1f} GB  \n"
+        + " · ".join(f"{name} {seconds:.1f}s" for name, seconds in timings.items())
+    )
+    if ai["action"]:
+        info += f"  \n**AI:** {ai['action']}"
+    if kept:
+        info += "  \n**Kept from the original:** " + ", ".join(o["name"] for o in kept)
+    info += f"  \n**Prompt:** {prompt or ai['prompt'] or DEFAULT_PROMPT}"
+    print(info.replace("**", "").replace("  \n", " | "), "| person box:", box)
+    return out_path, preview_path, face_path, ref_path, info
+
+
+# --------------------------------------------------------------------------------------------------------------
+# UI
+# --------------------------------------------------------------------------------------------------------------
+
+
+def replace(image_path, video_path, prompt, quality, fps, max_seconds, steps, shift, sampler, seed, keep_background,
+            keep_audio, mask_edge, mask_shape, contact, face_strength, use_ai, clicks,
+            progress=gr.Progress(track_tqdm=True)):
+    if not image_path:
+        raise gr.Error("Upload a photo of the new character.")
+    if not video_path:
+        raise gr.Error("Upload the video.")
+    prompt = (prompt or "").strip()
+    fps = int(fps)
+    image_path = prepare_image(image_path)
+    source_size = video_metadata(video_path)["size"]
+    height, width = frame_size(*source_size, quality)
+    frames_path, num_frames, audio_path = prepare_driving_video(video_path, float(max_seconds), fps, width, height)
+    if num_frames < 2:
+        raise gr.Error("Could not read at least two frames from the video.")
+    # Up to 2x sharper frames for the face crops, never upscaled beyond the source.
+    face_factor = min(2.0, 1 / max(width / source_size[0], height / source_size[1]))
+    face_path = None
+    if face_factor > 1.2:
+        face_w, face_h = round(width * face_factor / 2) * 2, round(height * face_factor / 2) * 2
+        face_path = prepare_driving_video(video_path, float(max_seconds), fps, face_w, face_h, audio=False)[0]
+    clicks = clicks or {}
+    click = None if clicks.get("person") is None else map_click(clicks["person"], source_size, width, height)
+    objects = [map_click(p, source_size, width, height) for p in clicks.get("objects", [])]
+    expected = estimate_gpu_seconds(num_frames, height, width, int(steps), scene=bool(use_ai))
+    if 1.2 * expected + 10 > MAX_GPU_SECONDS:
+        raise gr.Error(
+            f"This needs ≈{expected / 60:.0f} min of GPU in one call (limit {MAX_GPU_SECONDS / 60:.0f} min): "
+            "use fewer seconds, a lower FPS or 480p."
+        )
+    seed = random.randint(0, 2**31 - 1) if seed is None or int(seed) < 0 else int(seed)
+    return replace_on_gpu(
+        image_path, frames_path, face_path, num_frames, fps, prompt, int(steps), float(shift), sampler, seed,
+        bool(keep_background), click, objects, int(mask_edge), mask_shape, int(contact), float(face_strength),
+        bool(use_ai), audio_path if keep_audio else None, progress,
+    )
+
+
+def describe_cost(image_path, video_path, quality, fps, max_seconds, steps, use_ai=True):
+    if not image_path or not video_path:
+        return "Upload a photo and a video to see the expected GPU time."
+    try:
+        meta = video_metadata(video_path)
+        height, width = frame_size(*meta["size"], quality)
+        seconds = min(float(max_seconds), float(meta.get("duration") or max_seconds))
+        num_frames = max(2, int(seconds * int(fps)))
+        gpu = estimate_gpu_seconds(num_frames, height, width, int(steps), scene=bool(use_ai))
+    except Exception as e:  # noqa: BLE001
+        return f"Could not estimate the cost: {e}"
+    quota = gpu * (2 if GPU_SIZE == "xlarge" else 1)
+    return (
+        f"Output **{width}×{height}**, {num_frames} frames ({num_frames / int(fps):.1f} s at {fps} fps) → "
+        f"{segment_count(num_frames)} window(s) of {SEGMENT_FRAMES} frames. Expected GPU time ≈ **{gpu / 60:.1f} min** "
+        f"(≈ {quota / 60:.1f} min of ZeroGPU quota on `{GPU_SIZE}`)."
+    )
+
+
+def show_first_frame(video_path):
+    if not video_path:
+        return None, None, None
+    frame = first_frame(video_path)
+    return frame, frame, None
+
+
+CLICK_MODES = ["The person to replace", "An object to keep (lollipop, cup, phone…)"]
+MAX_OBJECTS = 6
+
+
+def draw_clicks(frame, clicks):
+    """The person's click in red, the objects to keep in green."""
+    if frame is None or not clicks:
+        return frame
+    image = Image.fromarray(frame)
+    draw = ImageDraw.Draw(image)
+    r = max(6, image.width // 80)
+    marks = [(clicks.get("person"), (255, 40, 40))] + [(p, (40, 220, 60)) for p in clicks.get("objects", [])]
+    for point, color in marks:
+        if point is not None:
+            x, y = point[0] * image.width, point[1] * image.height
+            draw.ellipse([x - r, y - r, x + r, y + r], fill=color, outline=(255, 255, 255), width=max(2, r // 3))
+    return np.asarray(image)
+
+
+def select_point(frame, mode, clicks, evt: gr.SelectData):
+    if frame is None:
+        return None, None
+    point = [evt.index[0] / frame.shape[1], evt.index[1] / frame.shape[0]]
+    clicks = {"person": (clicks or {}).get("person"), "objects": list((clicks or {}).get("objects", []))}
+    if mode == CLICK_MODES[0]:
+        clicks["person"] = point
+    else:
+        clicks["objects"] = (clicks["objects"] + [point])[-MAX_OBJECTS:]
+    return clicks, draw_clicks(frame, clicks)
+
+
+def tolerate_empty(slider: gr.Slider) -> gr.Slider:
+    """Lets the slider's number box be cleared: an empty value falls back to the initial one.
+
+    (Patched on the instance: subclassing a Gradio component makes Gradio write an `app.pyi` stub next to app.py.)
+    """
+    preprocess = slider.preprocess
+    slider.preprocess = lambda payload: preprocess(slider.value if payload is None else payload)
+    return slider
+
+
+def fetch_examples(time_budget: float = 20.0) -> list[list[str]]:
+    """Wan-Animate's official replacement example (Apache-2.0), downloaded once. Never holds up the startup for
+    longer than `time_budget`: a stalled download keeps going in the background and the example appears after a
+    restart."""
+    folder = Path(__file__).parent / "examples"
+    base = "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/wan_animate/replace"
+    files = {"wan_replace_image.jpeg": "image.jpeg", "wan_replace_video.mp4": "video.mp4"}
+
+    def download():
+        try:
+            context = ssl.create_default_context(cafile=certifi.where())
+            folder.mkdir(exist_ok=True)
+            for local, remote in files.items():
+                target = folder / local
+                if not target.exists():
+                    with urllib.request.urlopen(f"{base}/{remote}", context=context, timeout=15) as r:
+                        data = r.read()
+                    partial = target.with_name(target.name + ".part")
+                    partial.write_bytes(data)
+                    partial.replace(target)
+        except Exception as e:  # noqa: BLE001 - examples are optional
+            print(f"Examples unavailable: {e!r}")
+
+    worker = threading.Thread(target=download, daemon=True)
+    worker.start()
+    worker.join(time_budget)
+    if all((folder / name).exists() for name in files):
+        return [[str(folder / name) for name in files]]
+    print("Examples skipped: the download did not finish in time")
+    return []
+
+
+HEADER = f"""
+# Character replacement · Wan2.2-Animate
+Upload a **photo of the new character** and a **video**: the person in the video is replaced by the character from the
+photo. The new character takes the original person's **skeleton** (body proportions, position, every movement) and
+**face** (expressions, lips, head turns) frame by frame; only that person's region is redrawn, the rest of the video
+stays as it is. The character is cut out of the photo automatically (a PNG with transparency is used as is).
+An AI (Qwen3.5) looks at the video first: it writes what the person does into the prompt and finds what they hold
+or have at the mouth (a lollipop, a cup, a phone), which then stays the original one, together with the lips or
+fingers touching it. Things it misses can be clicked on the first frame with "An object to keep".
+Model: [Wan2.2-Animate-14B](https://github.com/Wan-Video/Wan2.2) replacement mode + Relighting LoRA{
+" + Lightx2v (6 steps)" if DISTILLED else ""}; masks by SAM 2.1 + Mask2Former, pose by ViTPose-H.
+"""
+
+CSS = """
+#col-container { max-width: 1400px; margin: 0 auto; }
+#run-button { min-height: 52px; font-size: 18px; }
+"""
+
+with gr.Blocks(title="Wan-Animate Character Replacement") as demo:
+    point_state = gr.State(None)
+    frame_state = gr.State(None)
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(HEADER)
+        with gr.Row(equal_height=True):
+            ref_image = gr.Image(label="1 · Photo of the new character", type="filepath", height=420,
+                                 image_mode=None)
+            driving_video = gr.Video(label="2 · Video with the person to replace", height=420, sources=["upload"])
+            output_video = gr.Video(label="Result", height=420, autoplay=True, interactive=False)
+        with gr.Row():
+            with gr.Column(scale=1):
+                picker = gr.Image(label="3 · Optional: click the person to replace (default: the largest one) "
+                                  "and things they hold that must stay as they are", type="numpy",
+                                  interactive=False, height=300)
+                click_mode = gr.Radio(CLICK_MODES, value=CLICK_MODES[0], label="A click marks")
+                use_ai = gr.Checkbox(
+                    value=reader is not None, interactive=reader is not None,
+                    label=f"AI scene reading ({SCENE_MODEL.split('/')[-1]}): writes the action into the prompt and "
+                          "finds what the person holds or has at the mouth",
+                )
+                clear_click = gr.Button("Clear the clicks", size="sm")
+            with gr.Column(scale=2):
+                with gr.Row():
+                    quality = gr.Radio(list(QUALITIES), value=DEFAULT_QUALITY, label="Resolution")
+                    fps = gr.Radio(FPS_CHOICES, value=DEFAULT_FPS, label="FPS (30 = smoothest; lower = longer clips)")
+                max_seconds = tolerate_empty(
+                    gr.Slider(1, MAX_SECONDS, value=5, step=0.5, label="Use the first N seconds of the video")
+                )
+                keep_background = gr.Checkbox(
+                    value=True, label="Keep the original pixels outside the person (scene and objects stay exact)"
+                )
+                mask_edge = tolerate_empty(
+                    gr.Slider(-6, 16, value=0, step=1, label="Mask margin, px: − tighter / + looser (if hair or "
+                              "clothes of the new character get cut)")
+                )
+        with gr.Accordion("Advanced", open=False):
+            mask_shape = gr.Radio(list(MASK_SHAPES), value=DEFAULT_MASK_SHAPE, label="Redrawn region")
+            prompt = gr.Textbox(
+                label="Prompt (optional)", lines=2,
+                placeholder="Empty = written by the AI scene reading (the character from the photo doing what the "
+                            "person in the video does); with it off, the official prompt.",
+            )
+            with gr.Row():
+                sampler = gr.Radio(list(SAMPLERS), value=DEFAULT_SAMPLER, label="Sampler")
+                steps = tolerate_empty(gr.Slider(4, 40, value=DEFAULT_STEPS, step=1, label="Steps"))
+                shift = tolerate_empty(gr.Slider(1.0, 10.0, value=DEFAULT_SHIFT, step=0.5, label="Sampling shift"))
+                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
+                keep_audio = gr.Checkbox(value=True, label="Keep the video's audio")
+            with gr.Row():
+                face_strength = tolerate_empty(gr.Slider(
+                    0.5, 2.0, value=1.0, step=0.05,
+                    label="Expression strength: > 1 = the mouth, eyes and brows follow the video more strongly",
+                ))
+                contact = tolerate_empty(gr.Slider(
+                    0, 24, value=4, step=1,
+                    label="Original pixels kept around the marked objects, px (lips, tongue or fingers on them)",
+                ))
+        cost = gr.Markdown(describe_cost(None, None, DEFAULT_QUALITY, DEFAULT_FPS, 5, DEFAULT_STEPS))
+        run = gr.Button("Replace the person", variant="primary", elem_id="run-button")
+        info = gr.Markdown()
+        with gr.Accordion("What the model got: skeleton and redrawn region, face, cut-out character", open=False):
+            with gr.Row():
+                preview_video = gr.Video(label="Skeleton, redrawn region (blue), kept objects (green)", height=300,
+                                         interactive=False)
+                face_video = gr.Video(label="Face (expressions)", height=300, interactive=False)
+                ref_preview = gr.Image(label="Character cut out of the photo", height=300, interactive=False)
+        examples = fetch_examples()
+        if examples:
+            gr.Examples(examples=examples, inputs=[ref_image, driving_video], label="Wan-Animate example")
+
+    cost_inputs = [ref_image, driving_video, quality, fps, max_seconds, steps, use_ai]
+    for component in cost_inputs:
+        component.change(describe_cost, cost_inputs, cost, show_progress="hidden", api_visibility="private")
+    driving_video.change(show_first_frame, driving_video, [picker, frame_state, point_state], api_visibility="private")
+    picker.select(select_point, [frame_state, click_mode, point_state], [point_state, picker],
+                  api_visibility="private")
+    clear_click.click(lambda f: (None, f), frame_state, [point_state, picker], api_visibility="private")
+    run.click(
+        replace,
+        [ref_image, driving_video, prompt, quality, fps, max_seconds, steps, shift, sampler, seed, keep_background,
+         keep_audio, mask_edge, mask_shape, contact, face_strength, use_ai, point_state],
+        [output_video, preview_video, face_video, ref_preview, info],
+        api_name="replace",
+        concurrency_limit=CONCURRENCY,
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=30).launch(theme=gr.themes.Soft(), css=CSS, ssr_mode=False)

captioner.py

@@ -0,0 +1,148 @@
+"""Reads the source video with a local vision-language model (Qwen3.5-4B) before the replacement:
+
+* what the person to replace is doing and how (e.g. "licking a red lollipop with her tongue, holding its stick"),
+* the things they hold or touch (with boxes on the first frame), which are then tracked and kept from the original
+  video instead of being redrawn,
+* an English prompt for Wan-Animate in which the character from the photo does exactly that.
+
+Wan-Animate's motion comes from the skeleton and the face crops, which carry neither a tongue on a lollipop nor a
+cup at the lips; the prompt and the kept objects fill that gap.
+
+Two separate questions: the photo alone (the new character's looks) and the video frames alone (action, objects,
+scene). The prompt is put together from both answers, so the original person's looks cannot end up in it.
+"""
+
+from __future__ import annotations
+
+import itertools
+import json
+import re
+from contextlib import contextmanager
+
+import torch
+from PIL import Image, ImageDraw
+
+APPEARANCE = """Describe only how the person in this photo looks, in one English sentence of at most 35 words: gender, approximate age, hair (colour, length, style), face and the visible clothing. Do not mention the pose, what they do, the background or the photo. Start with "A" (for example "A young woman with long wavy blonde hair and a pink top"). Output only the sentence."""
+
+VIDEO = """These are frames of a video in chronological order. In the first frame, the person to follow is marked with a red box.
+
+Answer with one JSON object and nothing else:
+{
+  "action": "what the marked person does, in a few words (e.g. licks a lollipop, drinks from a cup, pets a snake)",
+  "motion": "one English sentence on exactly how they do it across the frames: the hands, mouth, tongue, head and body, and the objects involved. Start with a verb ending in -ing, e.g. sitting on a sofa, licking a big red lollipop with the tongue while holding its stick in the right hand",
+  "scene": "one short English sentence about the place, the light and the camera",
+  "objects": [{"name": "...", "box": [x0, y0, x1, y1]}]
+}
+
+"objects": every thing the marked person holds in the hands, has at or in the mouth, or plays with (a lollipop, a cup, a phone, a snake, a guitar), with its box in the FIRST frame as integers from 0 to 1000 relative to that frame's width and height. Do not list clothing, hair, glasses, jewellery, furniture or the background. Use [] if there is none.
+Never describe the marked person's looks (gender, age, face, hair, clothing) and never name them: "motion" only says what is done, e.g. "sitting down, playing the violin with the bow in the right hand"."""
+
+
+@contextmanager
+def visiting(module: torch.nn.Module, device: str):
+    """Move `module` to `device` for the duration of the block, then point it back at its original CPU storage.
+
+    Restoring the original tensors (instead of copying back with `.to("cpu")`) keeps the weights shared with the
+    parent process, so ZeroGPU's forked workers do not each end up with a private CPU copy.
+    """
+    original = {name: t.data for name, t in itertools.chain(module.named_parameters(), module.named_buffers())}
+    module.to(device)
+    try:
+        yield module
+    finally:
+        for name, t in itertools.chain(module.named_parameters(), module.named_buffers()):
+            t.data = original[name]
+
+
+def mark(frame: Image.Image, box, width: int = 4) -> Image.Image:
+    frame = frame.convert("RGB").copy()
+    if box is not None:
+        ImageDraw.Draw(frame).rectangle([round(v) for v in box], outline=(255, 0, 0), width=width)
+    return frame
+
+
+def clean(text: str) -> str:
+    text = re.sub(r"<think>.*?</think>", "", text, flags=re.S)
+    return re.sub(r"\s+", " ", text.replace("*", "").replace("#", "")).strip().strip('"').strip()
+
+
+def parse_video(text: str, width: int, height: int) -> dict:
+    """The JSON answer about the video, with object boxes in pixels of the first frame; empty fields if unusable."""
+    match = re.search(r"\{.*\}", re.sub(r"<think>.*?</think>", "", text, flags=re.S), flags=re.S)
+    try:
+        data = json.loads(match.group(0)) if match else {}
+    except json.JSONDecodeError:
+        data = {}
+    if not isinstance(data, dict):
+        data = {}
+    out = {key: clean(str(data.get(key) or "")) for key in ("action", "motion", "scene")}
+    out["objects"] = []
+    for obj in data.get("objects") or []:
+        try:
+            x0, y0, x1, y1 = (float(v) for v in obj["box"])
+        except (KeyError, TypeError, ValueError):
+            continue
+        x0, x1 = sorted((min(max(x0, 0), 1000), min(max(x1, 0), 1000)))
+        y0, y1 = sorted((min(max(y0, 0), 1000), min(max(y1, 0), 1000)))
+        if x1 - x0 < 2 or y1 - y0 < 2:
+            continue
+        box = [x0 * width / 1000, y0 * height / 1000, x1 * width / 1000, y1 * height / 1000]
+        out["objects"].append({"name": clean(str(obj.get("name") or "object")), "box": box})
+    return out
+
+
+def compose(appearance: str, motion: str, scene: str) -> str:
+    """"A young woman with ..., licking a lollipop .... The room is ..." from the two answers."""
+    appearance = appearance.rstrip(". ")
+    motion = motion.rstrip(". ")
+    if appearance and motion:
+        continues = motion.split()[0].lower().endswith("ing")  # "..., licking a lollipop" reads on
+        text = f"{appearance}, {motion[0].lower() + motion[1:]}." if continues else f"{appearance}. {motion}."
+    elif appearance or motion:
+        text = (appearance or motion)[0].upper() + (appearance or motion)[1:] + "."
+    else:
+        text = ""
+    if scene:
+        text = f"{text} {scene.rstrip('.')}.".strip()
+    return text
+
+
+class SceneReader:
+    def __init__(self, model_id: str):
+        from transformers import AutoModelForImageTextToText, AutoProcessor
+
+        self.processor = AutoProcessor.from_pretrained(model_id)
+        self.model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=torch.bfloat16).eval()
+
+    def _ask(self, images: list[Image.Image], question: str, max_new_tokens: int, device: str) -> str:
+        messages = [{"role": "user", "content": [
+            *({"type": "image", "image": im} for im in images), {"type": "text", "text": question},
+        ]}]
+        inputs = self.processor.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt",
+            enable_thinking=False,
+        ).to(device)
+        output = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False,
+                                     repetition_penalty=1.05)
+        return self.processor.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+
+    @torch.no_grad()
+    def __call__(self, reference: Image.Image, frames: list[Image.Image], person_box, device: str) -> dict:
+        """`frames`: a few video frames in order, the person to replace in `person_box` on the first one.
+        Returns action, motion, scene, objects (boxes in first-frame pixels), prompt and the raw answers."""
+        width, height = frames[0].size
+
+        def fit(image: Image.Image, side: int) -> Image.Image:
+            image = image.convert("RGB").copy()
+            image.thumbnail((side, side), Image.Resampling.LANCZOS)
+            return image
+
+        video = [fit(mark(frames[0], person_box), 1024), *(fit(f, 640) for f in frames[1:])]
+        with visiting(self.model, device):
+            looks = self._ask([fit(reference, 768)], APPEARANCE, 80, device)
+            answer = self._ask(video, VIDEO, 360, device)
+        result = parse_video(answer, width, height)
+        appearance = clean(looks)
+        result["prompt"] = compose(appearance, result["motion"], result["scene"])
+        result["raw"] = {"appearance": looks, "video": answer}
+        return result

masking.py

@@ -0,0 +1,339 @@
+"""Person masks for Wan-Animate replacement.
+
+The official preprocessing prompts SAM 2 with pose keypoints. Here RT-DETRv2 finds the person on the first frame (or
+the one that was clicked), SAM 2.1 tracks it through the video, and a Mask2Former "person" segmentation removes
+whatever SAM 2.1 glued to it (the chair under it, the wall or room behind it). All three are ungated.
+"""
+
+from __future__ import annotations
+
+import os
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageFilter
+from transformers import (
+    AutoImageProcessor,
+    Mask2FormerForUniversalSegmentation,
+    RTDetrV2ForObjectDetection,
+    Sam2VideoModel,
+    Sam2VideoProcessor,
+)
+
+DETECTOR = os.getenv("WAN_DETECTOR", "PekingU/rtdetr_v2_r50vd")
+SEGMENTER = os.getenv("WAN_SEGMENTER", "facebook/sam2.1-hiera-large")
+PERSON_PARSER = os.getenv("WAN_PERSON_PARSER", "facebook/mask2former-swin-base-coco-panoptic")  # "none" = off
+# Below this share of SAM 2.1's mask being "person" pixels, the subject is taken to be non-human (a cartoon, an
+# animal) and the person refinement is skipped rather than erasing it.
+MIN_PERSON_SHARE = 0.4
+
+
+def _area(box):
+    return max(0.0, box[2] - box[0]) * max(0.0, box[3] - box[1])
+
+
+def _person_label(config) -> int:
+    return next(i for i, name in config.id2label.items() if name == "person")
+
+
+class PersonMasker:
+    def __init__(self, detector_id: str = DETECTOR, segmenter_id: str = SEGMENTER, parser_id: str = PERSON_PARSER):
+        self.det_processor = AutoImageProcessor.from_pretrained(detector_id)
+        self.detector = RTDetrV2ForObjectDetection.from_pretrained(detector_id).eval()
+        self.person_label = _person_label(self.detector.config)
+        self.sam_processor = Sam2VideoProcessor.from_pretrained(segmenter_id)
+        self.sam = Sam2VideoModel.from_pretrained(segmenter_id).eval()
+        self.parser = None
+        if parser_id.lower() != "none":
+            self.parser_processor = AutoImageProcessor.from_pretrained(parser_id)
+            self.parser = Mask2FormerForUniversalSegmentation.from_pretrained(parser_id).eval()
+            self.parser_person = _person_label(self.parser.config)
+
+    def to(self, device):
+        for model in (self.detector, self.sam, self.parser):
+            if model is not None:
+                model.to(device)
+        return self
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.sam.parameters()).device
+
+    @torch.no_grad()
+    def find_people(self, image: Image.Image, threshold: float = 0.4) -> list[tuple[list[float], float]]:
+        inputs = self.det_processor(images=image, return_tensors="pt").to(self.device)
+        result = self.det_processor.post_process_object_detection(
+            self.detector(**inputs), target_sizes=[(image.height, image.width)], threshold=threshold
+        )[0]
+        return [
+            (box.tolist(), float(score))
+            for box, score, label in zip(result["boxes"], result["scores"], result["labels"])
+            if int(label) == self.person_label
+        ]
+
+    @staticmethod
+    def pick(people, point=None, size=None):
+        """The clicked person (smallest box containing the click), else the largest confident, fully visible one.
+
+        Each side cut by the frame border halves a candidate's weight: passers-by crossing close to the camera are
+        big but cropped, the subject of the shot usually is not.
+        """
+        if point is not None:
+            inside = [b for b, _ in people if b[0] <= point[0] <= b[2] and b[1] <= point[1] <= b[3]]
+            return min(inside, key=_area) if inside else None
+        if not people:
+            return None
+
+        def weight(person):
+            box, score = person
+            cut = 0
+            if size is not None:
+                cut = sum([box[0] <= 2, box[1] <= 2, box[2] >= size[0] - 2, box[3] >= size[1] - 2])
+            return _area(box) * score * 0.5**cut
+
+        return max(people, key=weight)[0]
+
+    @torch.no_grad()
+    def track(self, frames: list[Image.Image], box=None, point=None, objects=()) -> tuple[np.ndarray, list]:
+        """Masks [T, H, W] of the subject prompted on frame 0 by a box and/or a positive click, and one mask per
+        object of `objects` (dicts with a "point" [x, y] or a "box" [x0, y0, x1, y1] on frame 0), each tracked as a
+        separate SAM 2.1 object."""
+        if box is None and point is None:
+            raise ValueError("Need a box or a point on the first frame")
+        session = self.sam_processor.init_video_session(
+            video=frames, inference_device=self.device, video_storage_device="cpu", dtype=torch.float32
+        )
+        prompt = {}
+        if box is not None:
+            prompt["input_boxes"] = [[list(map(float, box))]]
+        if point is not None:
+            prompt["input_points"] = [[[list(map(float, point))]]]
+            prompt["input_labels"] = [[[1]]]
+        self.sam_processor.add_inputs_to_inference_session(
+            inference_session=session, frame_idx=0, obj_ids=1, **prompt
+        )
+        self.sam(inference_session=session, frame_idx=0)
+        for i, obj in enumerate(objects):  # each prompt is run on frame 0 before the next is added
+            if "box" in obj:
+                hint = {"input_boxes": [[list(map(float, obj["box"]))]]}
+            else:
+                hint = {"input_points": [[[list(map(float, obj["point"]))]]], "input_labels": [[[1]]]}
+            self.sam_processor.add_inputs_to_inference_session(
+                inference_session=session, frame_idx=0, obj_ids=2 + i, **hint
+            )
+            self.sam(inference_session=session, frame_idx=0)
+        height, width = frames[0].height, frames[0].width
+        subject = np.zeros((len(frames), height, width), dtype=bool)
+        things = [np.zeros((len(frames), height, width), dtype=bool) for _ in objects]
+        for output in self.sam.propagate_in_video_iterator(session):
+            frame_masks = self.sam_processor.post_process_masks(
+                [output.pred_masks], original_sizes=[[height, width]], binarize=True
+            )[0].reshape(-1, height, width).cpu().numpy()
+            ids = list(output.object_ids) if output.object_ids is not None else list(range(1, len(frame_masks) + 1))
+            for obj_id, mask in zip(ids, frame_masks):
+                if obj_id == 1:
+                    subject[output.frame_idx] = mask
+                else:
+                    things[obj_id - 2][output.frame_idx] = mask
+        return subject, things
+
+    @torch.no_grad()
+    def person_pixels(self, frames: list[Image.Image], batch: int = 8, short_side: int = 512) -> np.ndarray:
+        """Semantic "person" masks [T, H, W] (Mask2Former COCO panoptic at `short_side`, upsampled)."""
+        height, width = frames[0].height, frames[0].width
+        out = np.zeros((len(frames), height, width), dtype=bool)
+        for start in range(0, len(frames), batch):
+            chunk = [f.convert("RGB") for f in frames[start : start + batch]]
+            size = {"shortest_edge": short_side, "longest_edge": 2 * short_side}
+            inputs = self.parser_processor(images=chunk, size=size, return_tensors="pt")
+            outputs = self.parser(**inputs.to(self.device))
+            maps = self.parser_processor.post_process_semantic_segmentation(
+                outputs, target_sizes=[(height, width)] * len(chunk)
+            )
+            for i, semantic in enumerate(maps):
+                out[start + i] = (semantic == self.parser_person).cpu().numpy()
+        return out
+
+    def refine(self, frames: list[Image.Image], masks: np.ndarray, tolerance: float = 0.012) -> np.ndarray:
+        """Keeps the tracked mask only where Mask2Former also sees a person (dilated a little, so hair and loose
+        clothing at the boundary survive). Skipped for subjects that are mostly not "person" pixels, and on frames
+        where it would remove most of the mask: there the person is usually only partly visible (occluded), which
+        Mask2Former misses, rather than SAM 2.1 having leaked."""
+        if self.parser is None or not masks.any():
+            return masks
+        person = self.person_pixels(frames)
+        tracked = masks.reshape(len(masks), -1).sum(1)
+        visible = tracked > 0
+        if np.median((masks & person).reshape(len(masks), -1).sum(1)[visible] / tracked[visible]) < MIN_PERSON_SHARE:
+            return masks
+        refined = masks & adjust(person, max(1, round(tolerance * min(masks.shape[1:]))), self.device)
+        kept = refined.reshape(len(masks), -1).sum(1) / np.maximum(tracked, 1)
+        refined[kept < 0.5] = masks[kept < 0.5]
+        return refined
+
+    def video_masks(self, frames: list[Image.Image], point=None, objects=()):
+        """(person masks, union of the objects' masks, person box on frame 0, the objects kept). Where an object
+        covers the person (a lollipop at the mouth, a cup in the hand), the pixels belong to the object.
+
+        Objects marked "check" (found by the scene reader rather than clicked) are dropped when their frame-0 mask is
+        mostly "person" pixels for Mask2Former: a box that landed on a hand or a face, not on what it holds."""
+        box = self.pick(self.find_people(frames[0]), point, frames[0].size)
+        if box is None and point is None:
+            raise ValueError("No person found on the first frame of the video: click on the person to replace.")
+        person, masks = self.track(frames, box=box, point=point, objects=objects)
+        body = None
+        if self.parser is not None and any(o.get("check") for o in objects):
+            body = self.person_pixels(frames[:1])[0]
+        things = np.zeros_like(person)
+        kept = []
+        for obj, mask in zip(objects, masks):
+            first = mask[0]
+            if obj.get("check") and body is not None and first.any() and (first & body).sum() > 0.6 * first.sum():
+                continue
+            things |= mask
+            kept.append(obj)
+        return self.refine(frames, person) & ~things, things, box, kept
+
+    def person_box(self, frame: Image.Image, point=None):
+        """The box `video_masks` will track on the first frame (the clicked person, else the main one)."""
+        return self.pick(self.find_people(frame), point, frame.size)
+
+    def image_mask(self, image: Image.Image) -> np.ndarray:
+        """Character mask of the reference: its alpha channel if it has a cut-out, else detection + SAM 2.1."""
+        if image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info):
+            alpha = np.asarray(image.convert("RGBA"))[..., 3] > 127
+            if 0.02 < alpha.mean() < 0.98:
+                return alpha
+        rgb = image.convert("RGB")
+        box = self.pick(self.find_people(rgb), size=rgb.size)
+        if box is None:  # e.g. a stylised character the detector does not see as a person
+            return self.track([rgb], point=[rgb.width / 2, rgb.height / 2])[0][0]
+        return self.refine([rgb], self.track([rgb], box=box)[0])[0]
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Conditioning tensors
+# --------------------------------------------------------------------------------------------------------------
+
+
+def adjust(masks: np.ndarray, pixels: int, device="cpu") -> np.ndarray:
+    """Grows (`pixels` > 0) or shrinks (< 0) bool masks [T, H, W] by about `pixels` px."""
+    if pixels == 0 or not masks.any():
+        return masks
+    m = torch.from_numpy(masks).to(device=device, dtype=torch.float16 if str(device) != "cpu" else torch.float32)
+    k, pad = 2 * abs(pixels) + 1, abs(pixels)
+    out = []
+    for chunk in m.split(32):
+        chunk = chunk[:, None]
+        if pixels > 0:
+            chunk = F.max_pool2d(chunk, k, stride=1, padding=pad)
+        else:
+            chunk = 1 - F.max_pool2d(1 - chunk, k, stride=1, padding=pad)
+        out.append(chunk[:, 0] > 0.5)
+    return torch.cat(out).cpu().numpy()
+
+
+def blocky(masks: np.ndarray, w_len: int, h_len: int) -> np.ndarray:
+    """The official `get_aug_mask`, per frame: the person's box is cut into a w_len x h_len grid and every cell the
+    mask touches is filled. 1 x 1 = the whole box (the official default), larger grids follow the outline closer."""
+    out = masks.copy()
+    for mask in out:
+        ys, xs = np.nonzero(mask)
+        if len(xs) == 0:
+            continue
+        x0, x1, y0, y1 = xs.min(), xs.max(), ys.min(), ys.max()
+        w_step, h_step = max(1, int((x1 - x0) / w_len)), max(1, int((y1 - y0) / h_len))
+        for x in range(x0, x1, w_step):
+            for y in range(y0, y1, h_step):
+                cell = mask[y : min(y + h_step, y1), x : min(x + w_step, x1)]
+                if cell.any():
+                    cell[...] = True
+    return out
+
+
+def blockify(masks: np.ndarray, block: int, device="cpu") -> np.ndarray:
+    """Every block x block cell of the image grid that the mask touches, filled.
+
+    KJNodes' BlockifyMask (32 px in Kijai's Wan-Animate workflow), but on the image grid: with `block` a multiple of
+    16 each cell is a whole number of VAE latents and transformer patches, redrawn or kept entirely. A mask edge that
+    cuts through a latent leaves part of the blacked-out person inside a latent marked as kept, which comes back as
+    a dark outline around the new character."""
+    if not masks.any():
+        return masks
+    height, width = masks.shape[1:]
+    m = torch.from_numpy(masks).to(device=device, dtype=torch.float16 if str(device) != "cpu" else torch.float32)
+    out = []
+    for chunk in m.split(32):
+        cells = F.max_pool2d(chunk[:, None], block, stride=block, ceil_mode=True)
+        out.append(F.interpolate(cells, scale_factor=block, mode="nearest")[:, 0, :height, :width] > 0.5)
+    return torch.cat(out).cpu().numpy()
+
+
+def mostly(masks: np.ndarray, cell: int = 8, share: float = 0.5, device="cpu") -> np.ndarray:
+    """The cell x cell cells of the image grid that the mask covers at least `share` of, filled (8 px = one VAE
+    latent)."""
+    if not masks.any():
+        return masks
+    height, width = masks.shape[1:]
+    m = torch.from_numpy(masks).to(device=device, dtype=torch.float32)
+    out = []
+    for chunk in m.split(32):
+        cover = F.avg_pool2d(chunk[:, None], cell, stride=cell, ceil_mode=True)
+        out.append(F.interpolate(cover, scale_factor=cell, mode="nearest")[:, 0, :height, :width] >= share)
+    return torch.cat(out).cpu().numpy()
+
+
+def scene_backdrop(frame: np.ndarray, person: np.ndarray) -> np.ndarray:
+    """The frame without the person (its area filled in from the surroundings by a normalised Gaussian blur), softly
+    blurred: what the reference character is pasted onto.
+
+    The Relighting LoRA was trained on references cut out of their video and pasted onto other scenes. The photo's
+    own room would show up around the new character, and a flat black backdrop comes back as a dark outline."""
+    height, width = person.shape
+    known = (~adjust(person[None], max(2, round(0.01 * min(height, width))))[0]).astype(np.float32)
+    image = frame.astype(np.float32)
+    mean = (image * known[..., None]).sum((0, 1)) / max(float(known.sum()), 1.0)
+    sigma = 0.08 * max(height, width)
+    num = cv2.GaussianBlur(image * known[..., None], (0, 0), sigma)
+    den = cv2.GaussianBlur(known, (0, 0), sigma)[..., None]
+    eps = 1e-3  # far from any visible background: the average background colour
+    filled = np.where(known[..., None] > 0, image, (num + eps * mean) / (den + eps))
+    return np.clip(cv2.GaussianBlur(filled, (0, 0), 0.015 * max(height, width)), 0, 255).astype(np.uint8)
+
+
+def place_character(image: Image.Image, mask: np.ndarray, backdrop: np.ndarray, margin: float = 0.1) -> Image.Image:
+    """The reference frame: the character cropped from the photo (plus margin), fitted into the frame and pasted onto
+    `backdrop` (uint8 [H, W, 3]) through its mask, so nothing of the photo's own background is left."""
+    height, width = backdrop.shape[:2]
+    ys, xs = np.nonzero(mask)
+    if len(xs) == 0:
+        raise ValueError("Could not find the character on the reference image.")
+    x0, x1, y0, y1 = xs.min(), xs.max() + 1, ys.min(), ys.max() + 1
+    mx, my = int((x1 - x0) * margin), int((y1 - y0) * margin)
+    box = (max(0, x0 - mx), max(0, y0 - my), min(image.width, x1 + mx), min(image.height, y1 + my))
+    crop = image.convert("RGB").crop(box)
+    scale = min(width / crop.width, height / crop.height)
+    size = (max(1, round(crop.width * scale)), max(1, round(crop.height * scale)))
+    offset = ((width - size[0]) // 2, (height - size[1]) // 2)
+    alpha = Image.fromarray(mask.astype(np.uint8) * 255).crop(box).resize(size, Image.Resampling.BILINEAR)
+    alpha = alpha.filter(ImageFilter.MinFilter(3))  # 1 px inwards: no rim of the photo's background
+    canvas = Image.fromarray(backdrop).convert("RGB")
+    canvas.paste(crop.resize(size, Image.Resampling.LANCZOS), offset, alpha)
+    return canvas
+
+
+def feathered(masks: np.ndarray, grow: int, device) -> torch.Tensor:
+    """bool [T, H, W] -> soft [T, 1, H, W]: dilated by `grow` px, then blurred over about `grow` px."""
+    m = torch.from_numpy(masks).to(device=device, dtype=torch.float32)[:, None]
+    if grow > 0:
+        m = F.max_pool2d(m, 2 * grow + 1, stride=1, padding=grow)
+        sigma = grow / 2
+        radius = int(3 * sigma)
+        x = torch.arange(-radius, radius + 1, device=device, dtype=torch.float32)
+        k = torch.exp(-(x**2) / (2 * sigma**2))
+        k = k / k.sum()
+        m = F.conv2d(F.pad(m, (radius, radius, 0, 0), mode="replicate"), k.view(1, 1, 1, -1))
+        m = F.conv2d(F.pad(m, (0, 0, radius, radius), mode="replicate"), k.view(1, 1, -1, 1))
+    return m.clamp_(0, 1)

pose.py

@@ -0,0 +1,387 @@
+"""Wan-Animate's driving signals, extracted from the video the way its official preprocessing does it
+(wan/modules/animate/preprocess in Wan-Video/Wan2.2, Apache-2.0):
+
+* whole-body keypoints (COCO-WholeBody, 133 points) from ViTPose-H, one top-down crop per frame around the person;
+* the skeleton image Wan-Animate was trained on (`draw_aapose_by_meta_new`): body limbs, feet and hands;
+* 512x512 face crops for its motion encoder, which carries expressions, lips, gaze and head turns.
+
+ViTPose-H is the checkpoint the official preprocessing ships (process_checkpoint/pose2d/vitpose_h_wholebody.onnx in
+Wan-AI/Wan2.2-Animate-14B). Its weights are read from the ONNX external-data files into the PyTorch module below, so
+it runs on the ZeroGPU GPU without onnxruntime.
+
+In replacement mode the skeleton is not retargeted (the official default), so the new character gets the original
+person's limb lengths, positions and motion.
+
+One deliberate difference: the official `Pose2d.load_images` converts the already-RGB frames "BGR to RGB", so its
+ViTPose sees channel-swapped images. ViTPose was trained on RGB, which is what it gets here.
+"""
+
+from __future__ import annotations
+
+import colorsys
+import math
+import os
+
+import cv2
+import numpy as np
+import onnx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import snapshot_download
+
+POSE_REPO = os.getenv("WAN_POSE_REPO", "Wan-AI/Wan2.2-Animate-14B")
+POSE_DIR = "process_checkpoint/pose2d/vitpose_h_wholebody.onnx"
+CROP_SIZE = (256, 192)  # (height, width) of ViTPose's input
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406])
+IMAGENET_STD = np.array([0.229, 0.224, 0.225])
+THRESHOLD = 0.5  # keypoint confidence below which a limb or point is not drawn
+FACE_SIZE = 512
+
+# --------------------------------------------------------------------------------------------------------------
+# ViTPose-H (mmpose layout): ViT-H/16 on a 256x192 crop, two deconvolutions, 133 heatmaps at 64x48
+# --------------------------------------------------------------------------------------------------------------
+
+
+class _Block(nn.Module):
+    def __init__(self, dim: int, heads: int, hidden: int):
+        super().__init__()
+        self.heads = heads
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
+        self.fc1 = nn.Linear(dim, hidden)
+        self.fc2 = nn.Linear(hidden, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, n, c = x.shape
+        q, k, v = self.qkv(self.norm1(x)).view(b, n, 3, self.heads, c // self.heads).permute(2, 0, 3, 1, 4)
+        x = x + self.proj(F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(b, n, c))
+        return x + self.fc2(F.gelu(self.fc1(self.norm2(x))))
+
+
+class ViTPose(nn.Module):
+    def __init__(self, dim: int = 1280, depth: int = 32, heads: int = 16, hidden: int = 5120, keypoints: int = 133):
+        super().__init__()
+        self.patch_embed = nn.Conv2d(3, dim, kernel_size=16, stride=16, padding=2)
+        self.pos_embed = nn.Parameter(torch.zeros(1, 16 * 12, dim))  # patch positions + the (unused) cls position
+        self.blocks = nn.ModuleList(_Block(dim, heads, hidden) for _ in range(depth))
+        self.last_norm = nn.LayerNorm(dim, eps=1e-6)
+        self.deconv1 = nn.ConvTranspose2d(dim, 256, 4, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(256)
+        self.deconv2 = nn.ConvTranspose2d(256, 256, 4, stride=2, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(256)
+        self.final = nn.Conv2d(256, keypoints, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patch_embed(x)
+        b, c, h, w = x.shape
+        x = x.flatten(2).transpose(1, 2) + self.pos_embed
+        for block in self.blocks:
+            x = block(x)
+        x = self.last_norm(x).transpose(1, 2).reshape(b, c, h, w)
+        x = F.relu(self.bn1(self.deconv1(x)))
+        x = F.relu(self.bn2(self.deconv2(x)))
+        return self.final(x)
+
+
+def _onnx_weights(folder: str) -> dict[str, np.ndarray]:
+    """Initializers of end2end.onnx by name; anonymous MatMul/Add constants are named after the node using them."""
+    graph = onnx.load(os.path.join(folder, "end2end.onnx"), load_external_data=False).graph
+    users = {}
+    for node in graph.node:
+        for name in node.input:
+            users.setdefault(name, node)
+    weights = {}
+    for init in graph.initializer:
+        if init.data_location == onnx.TensorProto.EXTERNAL:
+            info = {e.key: e.value for e in init.external_data}
+            dtype = onnx.helper.tensor_dtype_to_np_dtype(init.data_type)
+            count = int(info["length"]) // np.dtype(dtype).itemsize if "length" in info else -1
+            array = np.fromfile(os.path.join(folder, info["location"]), dtype=dtype, count=count,
+                                offset=int(info.get("offset", 0))).reshape(tuple(init.dims))
+        else:
+            array = onnx.numpy_helper.to_array(init)
+        name = init.name
+        if name.startswith("onnx::"):
+            node = users[name].name  # e.g. /backbone/blocks.3/mlp/fc1/MatMul, /backbone/Add_1
+            name = node.strip("/").replace("/", ".")
+        weights[name] = array
+    return weights
+
+
+def load_vitpose(repo: str = POSE_REPO) -> ViTPose:
+    return vitpose_from_onnx(os.path.join(snapshot_download(repo, allow_patterns=[f"{POSE_DIR}/*"]), POSE_DIR))
+
+
+def vitpose_from_onnx(folder: str) -> ViTPose:
+    w = _onnx_weights(folder)
+    with torch.device("meta"):  # skip the random init of 632M parameters; every tensor comes from the checkpoint
+        model = ViTPose()
+    state = {
+        "patch_embed.weight": w["backbone.patch_embed.proj.weight"],
+        "patch_embed.bias": w["backbone.patch_embed.proj.bias"],
+        # The export folded x + pos_embed[:, 1:] + pos_embed[:, :1] into two constants.
+        "pos_embed": w["backbone.Add"] + w["backbone.Add_1"],
+        "last_norm.weight": w["backbone.last_norm.weight"],
+        "last_norm.bias": w["backbone.last_norm.bias"],
+        "deconv1.weight": w["keypoint_head.deconv_layers.0.weight"],
+        "deconv2.weight": w["keypoint_head.deconv_layers.3.weight"],
+        "final.weight": w["keypoint_head.final_layer.weight"],
+        "final.bias": w["keypoint_head.final_layer.bias"],
+    }
+    for ours, theirs in (("bn1", "1"), ("bn2", "4")):
+        for p in ("weight", "bias", "running_mean", "running_var"):
+            state[f"{ours}.{p}"] = w[f"keypoint_head.deconv_layers.{theirs}.{p}"]
+    for i in range(len(model.blocks)):
+        src = f"backbone.blocks.{i}"
+        for ours, theirs in (("norm1", "norm1"), ("norm2", "norm2")):
+            state[f"blocks.{i}.{ours}.weight"] = w[f"{src}.{theirs}.weight"]
+            state[f"blocks.{i}.{ours}.bias"] = w[f"{src}.{theirs}.bias"]
+        for ours, theirs in (("qkv", "attn.qkv"), ("proj", "attn.proj"), ("fc1", "mlp.fc1"), ("fc2", "mlp.fc2")):
+            state[f"blocks.{i}.{ours}.weight"] = w[f"{src}.{theirs}.MatMul"].T  # ONNX MatMul stores W^T
+            state[f"blocks.{i}.{ours}.bias"] = w[f"{src}.{theirs}.bias"]
+    state = {k: torch.from_numpy(np.require(v, np.float32, ["C", "W"])) for k, v in state.items()}
+    state["bn1.num_batches_tracked"] = torch.tensor(0)
+    state["bn2.num_batches_tracked"] = torch.tensor(0)
+    model.load_state_dict(state, strict=True, assign=True)
+    return model.eval().requires_grad_(False)
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Top-down crop and heatmap decoding (pose2d.py / pose2d_utils.py)
+# --------------------------------------------------------------------------------------------------------------
+
+
+def _box_center_scale(box, rescale: float = 1.25):
+    """`bbox_from_detector`: centre and scale (in units of 200 px) of a 3:4 crop around the box."""
+    aspect = CROP_SIZE[0] / CROP_SIZE[1]
+    center = np.array([(box[0] + box[2]) / 2.0, (box[1] + box[3]) / 2.0])
+    size = max((box[2] - box[0]) * aspect, box[3] - box[1])
+    return center, np.array([size / aspect, size]) / 200.0 * rescale
+
+
+def _transform(point, center, scale: float, invert: bool):
+    """`transform` with `get_transform` (no rotation), including its integer rounding."""
+    rows, cols = CROP_SIZE
+    h = 200 * scale
+    w = h / (rows / cols)
+    t = np.array([[cols / w, 0, cols * (-center[0] / w + 0.5)], [0, rows / h, rows * (-center[1] / h + 0.5)], [0, 0, 1]])
+    if invert:
+        t = np.linalg.inv(t)
+    p = t @ np.array([point[0] - 1, point[1] - 1, 1.0])
+    return np.array([round(p[0]), round(p[1])], dtype=int) + 1
+
+
+def _crop(image: np.ndarray, center, scale) -> np.ndarray:
+    """`crop`: the crop box (padded with black outside the image) resized to 192x256."""
+    rows, cols = CROP_SIZE
+    ul = _transform([1, 1], center, max(scale), invert=True) - 1
+    br = _transform([cols + 1, rows + 1], center, max(scale), invert=True) - 1
+    out = np.zeros((br[1] - ul[1], br[0] - ul[0], 3), dtype=np.float32)
+    new_x = max(0, -ul[0]), min(br[0], image.shape[1]) - ul[0]
+    new_y = max(0, -ul[1]), min(br[1], image.shape[0]) - ul[1]
+    old_x = max(0, ul[0]), min(image.shape[1], br[0])
+    old_y = max(0, ul[1]), min(image.shape[0], br[1])
+    if new_x[1] > new_x[0] and new_y[1] > new_y[0]:
+        out[new_y[0] : new_y[1], new_x[0] : new_x[1]] = image[old_y[0] : old_y[1], old_x[0] : old_x[1]]
+    return cv2.resize(out, (cols, rows))
+
+
+def _decode(heatmaps: np.ndarray, centers: np.ndarray, scales: np.ndarray, kernel: int = 11) -> np.ndarray:
+    """`keypoints_from_heatmaps(unbiased=True, use_udp=False)` (DARK), vectorised: [N, K, H, W] -> [N, K, 3]."""
+    n, k, h, w = heatmaps.shape
+    flat = heatmaps.reshape(n, k, -1)
+    idx = flat.argmax(-1)
+    maxvals = flat.max(-1)
+    preds = np.stack([idx % w, idx // w], -1).astype(np.float64)
+    preds[maxvals <= 0] = -1
+    # Gaussian modulation (zero borders, peak height kept), then a second-order Taylor step on log-heatmaps.
+    border = (kernel - 1) // 2
+    padded = np.zeros((n * k, h + 2 * border, w + 2 * border), dtype=np.float32)
+    padded[:, border:-border, border:-border] = heatmaps.reshape(n * k, h, w)
+    blurred = np.stack([cv2.GaussianBlur(p, (kernel, kernel), 0) for p in padded])[:, border:-border, border:-border]
+    peak = heatmaps.reshape(n * k, -1).max(-1)
+    blurred = blurred * (peak / np.maximum(blurred.reshape(n * k, -1).max(-1), 1e-12))[:, None, None]
+    logh = np.log(np.maximum(blurred, 1e-10)).astype(np.float64)
+    px, py = preds[..., 0].reshape(-1).astype(int), preds[..., 1].reshape(-1).astype(int)
+    ok = (px > 1) & (px < w - 2) & (py > 1) & (py < h - 2)
+    j = np.arange(n * k)
+    pxc, pyc = np.clip(px, 2, w - 3), np.clip(py, 2, h - 3)
+
+    def at(dy, dx):
+        return logh[j, pyc + dy, pxc + dx]
+
+    dx = 0.5 * (at(0, 1) - at(0, -1))
+    dy = 0.5 * (at(1, 0) - at(-1, 0))
+    dxx = 0.25 * (at(0, 2) - 2 * at(0, 0) + at(0, -2))
+    dxy = 0.25 * (at(1, 1) - at(-1, 1) - at(1, -1) + at(-1, -1))
+    dyy = 0.25 * (at(2, 0) - 2 * at(0, 0) + at(-2, 0))
+    det = dxx * dyy - dxy**2
+    ok &= det != 0
+    safe = np.where(ok, det, 1.0)
+    preds = preds.reshape(-1, 2)
+    preds[ok, 0] -= ((dyy * dx - dxy * dy) / safe)[ok]  # coord -= H^-1 g
+    preds[ok, 1] -= ((dxx * dy - dxy * dx) / safe)[ok]
+    preds = preds.reshape(n, k, 2)
+    # Back to image pixels (`transform_preds`, use_udp=False).
+    size = scales * 200.0
+    xy = preds * (size[:, None, :] / np.array([w, h])) + centers[:, None, :] - size[:, None, :] * 0.5
+    return np.concatenate([xy, maxvals[..., None]], -1)
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Keypoint layout (AAPoseMeta) and drawing (human_visualization.py)
+# --------------------------------------------------------------------------------------------------------------
+
+# 20 body points from the 133 COCO-WholeBody ones: OpenPose-18 (neck = mean of the shoulders) + one toe per foot.
+_BODY_A = [0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]
+_BODY_B = [0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]
+_LIMBS = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], [10, 11], [2, 12], [12, 13], [13, 14],
+          [2, 1], [1, 15], [15, 17], [1, 16], [16, 18], [14, 19], [11, 20]]
+_COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
+           [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
+           [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [200, 200, 0], [100, 100, 0]]
+_HAND_EDGES = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], [10, 11], [11, 12],
+               [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
+_HAND_COLORS = [tuple(c * 255 for c in colorsys.hsv_to_rgb(i / len(_HAND_EDGES), 1.0, 1.0)) for i in range(20)]
+
+
+def split_keypoints(kps: np.ndarray, last_body: np.ndarray | None = None) -> dict:
+    """`load_pose_metas_from_kp2ds_seq` for one frame: [133, 3] pixels -> body (20), hands (21 + 21) and face (69)
+    in pixels. A frame whose body points all fall outside the image keeps the previous frame's body."""
+    body = (kps[_BODY_A] + kps[_BODY_B]) / 2
+    if last_body is not None and body[:, :2].min(axis=1).max() < 0:
+        body = last_body
+    return {"body": body, "left_hand": kps[91:112], "right_hand": kps[112:133], "face": kps[22:91]}
+
+
+def _draw_hand(canvas: np.ndarray, kps: np.ndarray, stick: int) -> None:
+    eps = 0.01
+    for (a, b), color in zip(_HAND_EDGES, _HAND_COLORS):
+        k1, k2 = kps[a], kps[b]
+        if k1[2] < THRESHOLD or k2[2] < THRESHOLD:
+            continue
+        x1, y1, x2, y2 = int(k1[0]), int(k1[1]), int(k2[0]), int(k2[1])
+        if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+            cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=stick)
+    for k in kps:
+        if k[2] >= THRESHOLD and int(k[0]) > eps and int(k[1]) > eps:
+            cv2.circle(canvas, (int(k[0]), int(k[1])), stick, (0, 0, 255), thickness=-1)
+
+
+def draw_skeleton(meta: dict, height: int, width: int) -> np.ndarray:
+    """`draw_aapose_by_meta_new` (stickwidth v2, hands and head drawn): uint8 RGB [H, W, 3] on black."""
+    canvas = np.zeros((height, width, 3), np.uint8)
+    stick = max(int(min(height, width) / 200) - 1, 1)
+    body = meta["body"]
+    for (a, b), color in zip(_LIMBS, _COLORS):
+        k1, k2 = body[a - 1], body[b - 1]
+        if k1[2] < THRESHOLD or k2[2] < THRESHOLD:
+            continue
+        cx, cy = (k1[0] + k2[0]) / 2, (k1[1] + k2[1]) / 2
+        length = math.hypot(k1[0] - k2[0], k1[1] - k2[1])
+        angle = math.degrees(math.atan2(k1[1] - k2[1], k1[0] - k2[0]))
+        polygon = cv2.ellipse2Poly((int(cx), int(cy)), (int(length / 2), stick), int(angle), 0, 360, 1)
+        cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
+    for k, color in zip(body, _COLORS):
+        if k[2] >= THRESHOLD:
+            cv2.circle(canvas, (int(k[0]), int(k[1])), stick, color, thickness=-1)
+    hand_stick = max(max(int(min(height, width) / 200) - 1, 1) // 2, 1)
+    _draw_hand(canvas, meta["left_hand"], hand_stick)
+    _draw_hand(canvas, meta["right_hand"], hand_stick)
+    return canvas
+
+
+def face_box(meta: dict, height: int, width: int, scale: float = 1.3):
+    """`get_face_bboxes` (the version used for replacement): the 68 face landmarks' box, grown to 1.3x its area,
+    mostly upwards (forehead). Returns (x0, x1, y0, y1) or None when the landmarks collapse to a point."""
+    face = meta["face"][1:, :2]
+    min_x, min_y = face.min(axis=0)
+    max_x, max_y = face.max(axis=0)
+    w0, h0 = max_x - min_x, max_y - min_y
+    if w0 < 2 or h0 < 2:
+        return None
+    area = w0 * h0 * scale
+    dw = (math.sqrt(area * w0 / h0) - w0) / 2
+    dh = (math.sqrt(area * h0 / w0) - h0) / 4
+    box = int(max(min_x - dw, 0)), int(min(max_x + dw, width)), int(max(min_y - 3 * dh, 0)), int(min(max_y + dh, height))
+    return box if box[1] - box[0] >= 2 and box[3] - box[2] >= 2 else None
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Whole video
+# --------------------------------------------------------------------------------------------------------------
+
+
+class PoseExtractor:
+    def __init__(self, repo: str = POSE_REPO):
+        self.model = load_vitpose(repo)
+
+    def to(self, device):
+        self.model.to(device)
+        return self
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
+    @torch.no_grad()
+    def keypoints(self, frames: list[np.ndarray], boxes: list, batch: int = 32) -> np.ndarray:
+        """uint8 RGB frames + one person box (x0, y0, x1, y1) per frame -> [T, 133, 3] (x, y in pixels, score)."""
+        crops, centers, scales = [], [], []
+        for frame, box in zip(frames, boxes):
+            if box[2] - box[0] < 10 or box[3] - box[1] < 10:  # as in `ViTPose.preprocess`: then the whole frame
+                box = [0, 0, frame.shape[1], frame.shape[0]]
+            center, scale = _box_center_scale(box)
+            crop = (_crop(frame, center, scale) / 255.0 - IMAGENET_MEAN) / IMAGENET_STD
+            crops.append(crop.transpose(2, 0, 1).astype(np.float32))
+            centers.append(center)
+            scales.append(scale)
+        heatmaps = []
+        for start in range(0, len(crops), batch):
+            x = torch.from_numpy(np.stack(crops[start : start + batch])).to(self.device)
+            heatmaps.append(self.model(x).float().cpu().numpy())
+        return _decode(np.concatenate(heatmaps), np.stack(centers), np.stack(scales))
+
+    def __call__(self, frames: list[np.ndarray], masks: np.ndarray, face_frames: list[np.ndarray] | None = None):
+        """Skeleton images, 512x512 face crops and per-frame keypoint metas for the person covered by `masks`.
+
+        The person box of each frame is its mask's box (the official code uses a YOLO person box); frames where the
+        person is not visible reuse the nearest visible frame's box. `face_frames`, the same frames at a higher
+        resolution, give sharper face crops: the official preprocessing crops faces from 720p frames, and a face of a
+        480p frame blown up to 512x512 carries less of the expression."""
+        height, width = frames[0].shape[:2]
+        boxes = [_mask_box(m) for m in masks]
+        visible = [i for i, b in enumerate(boxes) if b is not None]
+        if not visible:
+            raise ValueError("The person is not visible in any frame.")
+        boxes = [b if b is not None else boxes[min(visible, key=lambda j: abs(j - i))] for i, b in enumerate(boxes)]
+        kps = self.keypoints(frames, boxes)
+        metas, last_body = [], None
+        for k in kps:
+            meta = split_keypoints(k, last_body)
+            last_body = meta["body"]
+            metas.append(meta)
+        skeletons = [draw_skeleton(m, height, width) for m in metas]
+        sources = face_frames if face_frames is not None else frames
+        fy, fx = sources[0].shape[0] / height, sources[0].shape[1] / width
+        faces, last = [], None
+        for source, meta in zip(sources, metas):
+            box = face_box(meta, height, width) or last
+            if box is None:  # no usable face so far: the top of the person box
+                x0, y0, x1, y1 = _mask_box(masks[visible[0]])
+                side = max(2, min(x1 - x0, y1 - y0) // 2)
+                box = (int((x0 + x1 - side) / 2), int((x0 + x1 + side) / 2), int(y0), int(y0 + side))
+            last = box
+            x0, x1, y0, y1 = (int(v * f) for v, f in zip(box, (fx, fx, fy, fy)))
+            faces.append(cv2.resize(source[y0 : max(y1, y0 + 2), x0 : max(x1, x0 + 2)], (FACE_SIZE, FACE_SIZE)))
+        return skeletons, faces, metas
+
+
+def _mask_box(mask: np.ndarray):
+    ys, xs = np.nonzero(mask)
+    if len(xs) == 0:
+        return None
+    return [float(xs.min()), float(ys.min()), float(xs.max() + 1), float(ys.max() + 1)]

requirements.txt

@@ -0,0 +1,18 @@
+torch==2.10.0
+torchvision==0.25.0
+diffusers==0.42.1
+transformers==5.21.1
+accelerate==1.15.0
+huggingface_hub
+safetensors
+sentencepiece
+ftfy
+scipy
+onnx
+opencv-python-headless
+tqdm
+imageio-ffmpeg
+numpy
+pillow
+certifi
+kernels==0.18.0

wan_animate.py

@@ -0,0 +1,246 @@
+"""Wan2.2-Animate-14B (replacement mode) through diffusers' WanAnimatePipeline.
+
+The diffusers repository stores the transformer in fp32 (69 GB). Here the transformer comes from the official bf16
+checkpoint instead (Wan-AI/Wan2.2-Animate-14B, 34.5 GB): the LoRAs are fused into it in the original key layout,
+then its keys are renamed with diffusers' own single-file converter. Everything else (UMT5, CLIP, VAE, tokenizer,
+scheduler) is read from Wan-AI/Wan2.2-Animate-14B-Diffusers.
+
+LoRAs:
+* Relighting LoRA (official, PEFT layout, rank 128): matches the new character's lighting, shadows and colour tone to
+  the scene. The official code enables it for replacement.
+* Lightx2v step/CFG distillation for Wan2.1-I2V-14B (ComfyUI layout from Kijai/WanVideo_comfy): 6 steps instead of
+  20. Wan-Animate is built on Wan2.1-I2V; the file, strength (1.2) and sampler (DPM++ SDE, shift 5) are those of
+  Kijai's Wan-Animate example workflow.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from accelerate import init_empty_weights
+from diffusers import (
+    AutoencoderKLWan,
+    DPMSolverMultistepScheduler,
+    UniPCMultistepScheduler,
+    WanAnimatePipeline,
+    WanAnimateTransformer3DModel,
+)
+from diffusers.loaders.single_file_utils import convert_wan_transformer_to_diffusers
+from huggingface_hub import hf_hub_download, snapshot_download
+from PIL import Image
+from safetensors import safe_open
+from safetensors.torch import load_file
+
+ORIGINAL_REPO = os.getenv("WAN_ORIGINAL_REPO", "Wan-AI/Wan2.2-Animate-14B")
+DIFFUSERS_REPO = os.getenv("WAN_DIFFUSERS_REPO", "Wan-AI/Wan2.2-Animate-14B-Diffusers")
+LORA_FILES = {
+    "relight": (ORIGINAL_REPO, "relighting_lora/adapter_model.safetensors"),
+    "lightx2v": ("Kijai/WanVideo_comfy", "Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank64_bf16.safetensors"),
+}
+# name -> DPMSolver algorithm (None = UniPC, the diffusers default). The official code samples with DPM++, Kijai's
+# workflow with its stochastic variant, which suits the distilled model's few large steps.
+SAMPLERS = {
+    "DPM++ SDE": "sde-dpmsolver++",
+    "DPM++": "dpmsolver++",
+    "UniPC": None,
+}
+# The official default prompt ("the person in the video is doing actions"); Wan-Animate is driven by the pose, face
+# and reference, and its authors advise against custom prompts.
+DEFAULT_PROMPT = "视频中的人在做动作"
+SEGMENT_FRAMES = 77  # official clip length per inference window
+PREV_SEGMENT_FRAMES = 1  # official refert_num: generated frames carried into the next window
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Transformer: official bf16 checkpoint + fused LoRAs -> diffusers WanAnimateTransformer3DModel
+# --------------------------------------------------------------------------------------------------------------
+
+
+def _lora_updates(path: str, adapter_config: str | None = None) -> dict[str, list]:
+    """Target parameter (original key layout) -> what to add to it, for one LoRA file.
+
+    PEFT layout: base_model.model.<module>.lora_A/lora_B.weight, scaled by lora_alpha / r from adapter_config.json.
+    ComfyUI layout: diffusion_model.<module>.lora_down/lora_up.weight (+ optional .alpha), plus full differences
+    `.diff` (norm weights) and `.diff_b` (biases), which ComfyUI adds as they are.
+    """
+    peft_scale = 1.0
+    if adapter_config is not None:
+        with open(adapter_config) as f:
+            cfg = json.load(f)
+        peft_scale = cfg["lora_alpha"] / cfg["r"]
+    updates: dict[str, list] = {}
+    with safe_open(path, framework="pt") as f:
+        keys = set(f.keys())
+    for key in sorted(keys):
+        if key.endswith(".lora_A.weight"):
+            module = key[: -len(".lora_A.weight")].removeprefix("base_model.model.")
+            updates.setdefault(module + ".weight", []).append(("lora", key, key.replace(".lora_A.", ".lora_B."), peft_scale))
+        elif key.endswith(".lora_down.weight"):
+            prefix = key[: -len(".lora_down.weight")]
+            module = prefix.removeprefix("diffusion_model.")
+            alpha = prefix + ".alpha" if prefix + ".alpha" in keys else None
+            updates.setdefault(module + ".weight", []).append(("lora", key, prefix + ".lora_up.weight", alpha))
+        elif key.endswith(".diff_b"):
+            updates.setdefault(key[: -len(".diff_b")].removeprefix("diffusion_model.") + ".bias", []).append(("diff", key))
+        elif key.endswith(".diff"):
+            updates.setdefault(key[: -len(".diff")].removeprefix("diffusion_model.") + ".weight", []).append(("diff", key))
+    return updates
+
+
+@torch.no_grad()
+def fuse_loras(state: dict[str, torch.Tensor], loras: list[tuple[str, str | None, float]]) -> dict[str, int]:
+    """Adds strength * delta of each LoRA to `state` (original key layout) in place.
+
+    All LoRAs touching a weight are summed in fp32 and rounded to the checkpoint's dtype once.
+    `loras`: (safetensors path, PEFT adapter_config.json path or None, strength).
+    """
+    plans = [(path, _lora_updates(path, config), strength) for path, config, strength in loras if strength]
+    handles = [safe_open(path, framework="pt") for path, _, _ in plans]
+    targets = sorted(set().union(*(updates for _, updates, _ in plans))) if plans else []
+    counts = {os.path.basename(path): 0 for path, _, _ in plans}
+    missing = [t for t in targets if t not in state]
+    if missing:
+        raise KeyError(f"LoRA targets missing from the checkpoint: {missing[:5]} ...")
+    for target in targets:
+        weight = state[target]
+        total = weight.float()
+        for (path, updates, strength), f in zip(plans, handles):
+            for update in updates.get(target, []):
+                if update[0] == "lora":
+                    _, down_key, up_key, scale = update
+                    down, up = f.get_tensor(down_key).float(), f.get_tensor(up_key).float()
+                    if isinstance(scale, str):  # ComfyUI alpha key
+                        scale = f.get_tensor(scale).item() / down.shape[0]
+                    total += (strength * (scale if scale is not None else 1.0)) * (up @ down).view_as(total)
+                else:
+                    total += strength * f.get_tensor(update[1]).float().view_as(total)
+                counts[os.path.basename(path)] += 1
+        state[target] = total.to(weight.dtype)
+    return counts
+
+
+def build_transformer(state: dict[str, torch.Tensor], config: dict) -> WanAnimateTransformer3DModel:
+    """Original-layout state dict -> WanAnimateTransformer3DModel, without a second copy of the weights."""
+    return assemble_transformer(convert_wan_transformer_to_diffusers(state), config)
+
+
+def assemble_transformer(state: dict[str, torch.Tensor], config: dict) -> WanAnimateTransformer3DModel:
+    """diffusers-layout state dict -> model. Dtypes follow `from_pretrained(torch_dtype=torch.bfloat16)`: bf16,
+    except the modules diffusers keeps in fp32 (time embedding, modulation tables, norms)."""
+    with init_empty_weights():
+        model = WanAnimateTransformer3DModel.from_config(config)
+    expected = set(model.state_dict())
+    # The original stores the motion encoder's fixed blur kernels; diffusers recomputes them as buffers.
+    extra = sorted(set(state) - expected)
+    unknown = [k for k in extra if not k.endswith(".kernel")]
+    if unknown:
+        raise KeyError(f"Unexpected transformer weights: {unknown[:5]}")
+    keep_fp32 = WanAnimateTransformer3DModel._keep_in_fp32_modules
+    state = {
+        k: v.to(torch.float32 if any(m in k.split(".") for m in keep_fp32) else torch.bfloat16)
+        for k, v in state.items()
+        if k in expected
+    }
+    model.load_state_dict(state, strict=True, assign=True)
+    meta = [n for n, p in model.named_parameters() if p.is_meta]
+    if meta:
+        raise RuntimeError(f"Transformer parameters left uninitialised: {meta[:5]}")
+    return model.eval().requires_grad_(False)
+
+
+def load_transformer(lora_strengths: dict[str, float]) -> WanAnimateTransformer3DModel:
+    folder = snapshot_download(ORIGINAL_REPO, allow_patterns=["diffusion_pytorch_model*.safetensors*"])
+    with open(os.path.join(folder, "diffusion_pytorch_model.safetensors.index.json")) as f:
+        shards = sorted(set(json.load(f)["weight_map"].values()))
+    state: dict[str, torch.Tensor] = {}
+    for shard in shards:
+        state.update(load_file(os.path.join(folder, shard)))
+    loras = []
+    for name, strength in lora_strengths.items():
+        if strength:
+            repo, filename = LORA_FILES[name]
+            config = None
+            if name == "relight":
+                config = hf_hub_download(repo, os.path.join(os.path.dirname(filename), "adapter_config.json"))
+            loras.append((hf_hub_download(repo, filename), config, strength))
+    if loras:
+        print("Fused LoRA tensors:", fuse_loras(state, loras))
+    config = WanAnimateTransformer3DModel.load_config(DIFFUSERS_REPO, subfolder="transformer")
+    return build_transformer(state, config)
+
+
+# --------------------------------------------------------------------------------------------------------------
+# Pipeline
+# --------------------------------------------------------------------------------------------------------------
+
+
+def _wan_clip_encoder(pipe: WanAnimatePipeline):
+    """`encode_image` as in the official code: the whole reference frame (the character letterboxed into the video
+    frame) resized to 224x224, instead of diffusers' resize-then-centre-crop of the raw photo."""
+    mean = torch.tensor(pipe.image_processor.image_mean).view(1, 3, 1, 1)
+    std = torch.tensor(pipe.image_processor.image_std).view(1, 3, 1, 1)
+
+    def encode_image(image: Image.Image, device=None) -> torch.Tensor:
+        encoder = pipe.image_encoder
+        x = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1)[None].float() / 127.5 - 1
+        x = x.to(encoder.device)
+        x = F.interpolate(x, size=(224, 224), mode="bicubic", align_corners=False)
+        x = (x.mul(0.5).add(0.5) - mean.to(x.device)) / std.to(x.device)
+        out = encoder(pixel_values=x.to(encoder.dtype), output_hidden_states=True)
+        return out.hidden_states[-2].to(device or x.device)
+
+    return encode_image
+
+
+def load_pipeline(lora_strengths: dict[str, float]) -> WanAnimatePipeline:
+    transformer = load_transformer(lora_strengths)
+    folder = snapshot_download(
+        DIFFUSERS_REPO,
+        allow_patterns=["model_index.json", "scheduler/*", "tokenizer/*", "text_encoder/*", "image_encoder/*",
+                        "image_processor/*", "vae/*"],
+    )
+    vae = AutoencoderKLWan.from_pretrained(folder, subfolder="vae", torch_dtype=torch.float32)  # fp32 as in the docs
+    pipe = WanAnimatePipeline.from_pretrained(
+        folder, transformer=transformer, vae=vae,
+        torch_dtype={"image_encoder": torch.float16, "default": torch.bfloat16},  # official: CLIP fp16, UMT5 bf16
+    )
+    pipe.encode_image = _wan_clip_encoder(pipe)
+    pipe.set_progress_bar_config(desc="Denoising")
+    pipe.unipc_config = dict(pipe.scheduler.config)
+    return pipe
+
+
+def set_face_strength(pipe: WanAnimatePipeline, strength: float) -> None:
+    """Scales what each face-adapter block adds to the video tokens (Kijai's `face_strength`): above 1 the driving
+    face's expressions (mouth, eyes, brows) come through stronger."""
+    for block in pipe.transformer.face_adapter:
+        if not getattr(block, "scaled_forward", False):
+            def forward(*args, _forward=block.forward, _block=block, **kwargs):
+                out = _forward(*args, **kwargs)
+                return out if _block.face_strength == 1.0 else out * _block.face_strength
+            block.forward = forward
+            block.scaled_forward = True
+        block.face_strength = float(strength)
+
+
+def set_sampler(pipe: WanAnimatePipeline, sampler: str, shift: float) -> None:
+    algorithm = SAMPLERS[sampler]
+    if algorithm is None:
+        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.unipc_config, flow_shift=shift)
+    else:  # = the official FlowDPMSolverMultistepScheduler (solver order 2, midpoint, final sigma 0)
+        pipe.scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=1000, solver_order=2, prediction_type="flow_prediction", use_flow_sigmas=True,
+            flow_shift=shift, algorithm_type=algorithm, solver_type="midpoint", lower_order_final=True,
+            final_sigmas_type="zero",
+        )
+
+
+def segment_count(num_frames: int) -> int:
+    """Inference windows WanAnimatePipeline runs for `num_frames` frames."""
+    step = SEGMENT_FRAMES - PREV_SEGMENT_FRAMES
+    rest = (num_frames - PREV_SEGMENT_FRAMES) % step
+    return (num_frames + (step - rest if rest else 0)) // step