idv2v_lib.py

"""
Helpers ported 1:1 from Eyeline-Labs/ID-V2V (github.com/Eyeline-Labs/ID-V2V):

  src/idv2v/inference/pipeline.py   -> center_crop_and_resize, load_frames,
                                        I2V_14B_DIT_CONFIG, VACE_14B_CONFIG,
                                        load_finetuned_dit_vace, DEFAULT_NEGATIVE_PROMPT
  src/idv2v/preprocess/secret_panda.py -> secret_panda (morphology done with OpenCV
                                        instead of scipy.ndimage for speed; identical
                                        Minkowski semantics, border_value=0)
  src/idv2v/preprocess/sam3.py         -> run_sam3_union_masks
  src/idv2v/preprocess/orig_pixel.py   -> foreground_on_gray
"""

import os
import time
from typing import List

import cv2
import imageio.v2 as imageio
import numpy as np
import torch
from PIL import Image
from scipy import ndimage

# Default Wan 2.1 negative prompt (Chinese quality-degradation terms).
DEFAULT_NEGATIVE_PROMPT = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，"
    "JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
    "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
)

# Source of truth: WanModel / VaceWanModel state_dict converters in diffsynth.
I2V_14B_DIT_CONFIG = {
    "has_image_input": True, "patch_size": [1, 2, 2], "in_dim": 36,
    "dim": 5120, "ffn_dim": 13824, "freq_dim": 256, "text_dim": 4096,
    "out_dim": 16, "num_heads": 40, "num_layers": 40, "eps": 1e-6,
}
VACE_14B_CONFIG = {
    "vace_layers": (0, 5, 10, 15, 20, 25, 30, 35), "vace_in_dim": 96,
    "patch_size": (1, 2, 2), "has_image_input": False, "dim": 5120,
    "num_heads": 40, "ffn_dim": 13824, "eps": 1e-6,
}


# --------------------------------------------------------------------------- #
# Video / image I/O
# --------------------------------------------------------------------------- #
def center_crop_and_resize(img: Image.Image, width: int, height: int) -> Image.Image:
    """Center-crop + resize a PIL image to (width, height) with BICUBIC.
    Resizes to the target aspect ratio first, then center-crops. Verbatim from
    src/idv2v/inference/pipeline.py."""
    w, h = img.size
    target_aspect = height / width
    aspect = h / w

    if (h == height) and (w == width):
        return img
    if abs(aspect - target_aspect) < 1e-6:
        return img.resize((width, height), Image.BICUBIC)

    if aspect > target_aspect:  # too tall -> match width, crop height
        new_w = width
        new_h = int(aspect * new_w)
    else:                        # too wide -> match height, crop width
        new_h = height
        new_w = int(new_h / aspect)
    resized = img.resize((new_w, new_h), Image.BICUBIC)

    rw, rh = resized.size
    left = (rw - width) // 2
    top = (rh - height) // 2
    return resized.crop((left, top, left + width, top + height))


def read_video_rgb(path: str):
    """Return (list[PIL.Image RGB], fps)."""
    reader = imageio.get_reader(path, format="ffmpeg")
    fps = float(reader.get_meta_data().get("fps", 25.0) or 25.0)
    frames = []
    try:
        for f in reader:
            arr = np.asarray(f)
            if arr.ndim == 3 and arr.shape[2] > 3:
                arr = arr[:, :, :3]
            frames.append(Image.fromarray(arr.astype(np.uint8), "RGB"))
    finally:
        reader.close()
    if not frames:
        raise ValueError(f"No frames decoded from {path}")
    return frames, fps


def load_source_frames(path: str, width: int, height: int, num_frames: int, stride: int):
    """Read `num_frames` source frames with temporal `stride`, center-cropped and
    resized to (width, height). Pads by repeating the last frame if the source is
    too short. Also returns the source fps."""
    raw, fps = read_video_rgb(path)
    picked = raw[::max(1, int(stride))][:num_frames]
    if len(picked) < num_frames:
        picked = picked + [picked[-1]] * (num_frames - len(picked))
    return [center_crop_and_resize(f, width, height) for f in picked], fps


def save_video(frames: List[Image.Image], path: str, fps: float):
    writer = imageio.get_writer(
        path, fps=max(1.0, float(fps)), codec="libx264",
        quality=8, macro_block_size=None, pixelformat="yuv420p",
    )
    try:
        for f in frames:
            writer.append_data(np.asarray(f.convert("RGB")))
    finally:
        writer.close()
    return path


# --------------------------------------------------------------------------- #
# Secret Panda mask cleanup
# --------------------------------------------------------------------------- #
def _close(mask_u8: np.ndarray, k: int) -> np.ndarray:
    """Binary morphological close with a k x k square, border value 0
    (matches scipy.ndimage defaults, but OpenCV-fast)."""
    kernel = np.ones((k, k), np.uint8)
    d = cv2.dilate(mask_u8, kernel, borderType=cv2.BORDER_CONSTANT, borderValue=0)
    return cv2.erode(d, kernel, borderType=cv2.BORDER_CONSTANT, borderValue=0)


def secret_panda(mask: np.ndarray, fill_holes_first: bool = True,
                 close_kernel: int = 10, bridge_distance: int = 15) -> np.ndarray:
    """Clean a binary mask: hole-fill -> close -> bridge wider gaps -> hole-fill.
    Port of src/idv2v/preprocess/secret_panda.py."""
    result = (np.asarray(mask) > 0).astype(np.uint8)
    if fill_holes_first:
        result = ndimage.binary_fill_holes(result).astype(np.uint8)
    if close_kernel > 0:
        result = _close(result, close_kernel)
    if bridge_distance > 0:
        result = _close(result, bridge_distance)
    return ndimage.binary_fill_holes(result).astype(bool)


# --------------------------------------------------------------------------- #
# SAM3 person segmentation -> per-frame union masks
# --------------------------------------------------------------------------- #
@torch.inference_mode()
def run_sam3_union_masks(model, processor, frames: List[Image.Image], text_prompt: str,
                        device="cuda", dtype=torch.bfloat16,
                        close_kernel: int = 10, bridge_distance: int = 15,
                        mask_bin_threshold: float = 0.5) -> List[np.ndarray]:
    """Promptable Concept Segmentation over the whole clip, then Secret Panda cleanup
    per object + on the union (== sam3.py with --joint_mask_post_proc).
    Returns one bool HxW union mask per input frame."""
    H, W = frames[0].height, frames[0].width

    session = processor.init_video_session(
        video=frames,
        inference_device=device,
        processing_device=device,
        video_storage_device=device,
        dtype=dtype,
    )
    session = processor.add_text_prompt(inference_session=session, text=text_prompt)

    raw = {}
    for model_outputs in model.propagate_in_video_iterator(
        inference_session=session, max_frame_num_to_track=len(frames) - 1
    ):
        out = processor.postprocess_outputs(session, model_outputs)
        masks = out.get("masks", None)
        if masks is None or len(masks) == 0:
            raw[int(model_outputs.frame_idx)] = None
        else:
            raw[int(model_outputs.frame_idx)] = masks.float().cpu().numpy()

    unions = []
    for i in range(len(frames)):
        m_all = raw.get(i, None)
        if m_all is None:
            unions.append(np.zeros((H, W), dtype=bool))
            continue
        union = np.zeros((H, W), dtype=bool)
        for m in m_all:
            m = np.asarray(m)
            if m.shape != (H, W):
                m = cv2.resize(m.astype(np.float32), (W, H), interpolation=cv2.INTER_NEAREST)
            union |= secret_panda(m > mask_bin_threshold,
                                  close_kernel=close_kernel,
                                  bridge_distance=bridge_distance)
        # joint_mask_post_proc: also clean the union (bridges slivers between people)
        unions.append(secret_panda(union, close_kernel=close_kernel,
                                   bridge_distance=bridge_distance))
    return unions


def foreground_on_gray(frames: List[Image.Image], masks: List[np.ndarray],
                       gray_value: int = 127) -> List[Image.Image]:
    """Keep pixels inside the mask, fill the rest with gray 127. This is the single
    VACE condition ID-V2V consumes. Port of src/idv2v/preprocess/orig_pixel.py."""
    out = []
    for frame, mask in zip(frames, masks):
        rgb = np.asarray(frame.convert("RGB"), dtype=np.uint8)
        comp = np.full_like(rgb, gray_value)
        comp[mask] = rgb[mask]
        out.append(Image.fromarray(comp, "RGB"))
    return out


# --------------------------------------------------------------------------- #
# Finetuned DiT + VACE loading
# --------------------------------------------------------------------------- #
def load_finetuned_dit_vace(pipe, checkpoint_path: str, torch_dtype=torch.bfloat16,
                            delete_checkpoint_after: bool = False):
    """Instantiate empty DiT + VACE on meta, then assign the finetuned fp32 weights and
    cast to bf16. Verbatim logic from _load_finetuned_dit_vace in
    src/idv2v/inference/pipeline.py."""
    from diffsynth.models.utils import init_weights_on_device
    from diffsynth.models.wan_video_dit import WanModel
    from diffsynth.models.wan_video_vace import VaceWanModel

    with init_weights_on_device():
        pipe.dit = WanModel(**I2V_14B_DIT_CONFIG)
        pipe.vace = VaceWanModel(**VACE_14B_CONFIG)

    t0 = time.time()
    state_dict = torch.load(checkpoint_path, map_location="cpu",
                            weights_only=True, mmap=True)
    print(f"[idv2v] mmapped checkpoint in {time.time() - t0:.1f}s "
          f"({len(state_dict)} tensors)", flush=True)

    sd_vace = {k: v for k, v in state_dict.items() if "vace" in k}
    sd_dit = {k: v for k, v in state_dict.items() if "vace" not in k}
    pipe.dit.load_state_dict(sd_dit, assign=True)
    pipe.vace.load_state_dict(sd_vace, assign=True)
    del state_dict, sd_vace, sd_dit

    t0 = time.time()
    pipe.dit = pipe.dit.to(dtype=torch_dtype)
    pipe.vace = pipe.vace.to(dtype=torch_dtype)
    print(f"[idv2v] materialized DiT+VACE in {torch_dtype} in {time.time() - t0:.1f}s",
          flush=True)

    if delete_checkpoint_after:
        # The 78 GB fp32 blob is no longer referenced; free the disk before ZeroGPU
        # packs the bf16 weights back out.
        try:
            real = os.path.realpath(checkpoint_path)
            os.remove(real)
            if os.path.islink(checkpoint_path):
                os.remove(checkpoint_path)
            print(f"[idv2v] removed fp32 checkpoint blob {real}", flush=True)
        except OSError as e:
            print(f"[idv2v] could not remove checkpoint: {e}", flush=True)