Spaces:
Running on Zero
Running on Zero
File size: 3,510 Bytes
0c9f968 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """
I2V VAE 未知帧像素填充策略（仅推理/实验 opt-in）
Wan 官方 image2video.py 对未知帧使用 ``torch.zeros``，张量值在 **[-1, 1]**
归一化空间里为 **0**，对应 RGB 约 **128 的中灰**，不是 uint8 意义的黑 (0→-1)。
本模块供 ``WanVideoUnit_ImageEmbedderVAE`` 在 ``pipe.i2v_vae_unknown_frame_fill``
被设置时选用；未设置时行为与上游完全一致。
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
# ImageNet 均值 (RGB, [0,1])，映射到 preprocess_image 的 [-1, 1]
_IMAGENET_MEAN_01: Tuple[float, float, float] = (0.485, 0.456, 0.406)
FILL_MODE_ALIASES = {
None: "official_zero",
"official": "official_zero",
"zero": "official_zero",
"official_zero": "official_zero",
"black": "black",
"gray": "official_zero",
"mid_gray": "official_zero",
"first_frame_mean": "first_frame_mean",
"imagenet_mean": "imagenet_mean",
}
FILL_MODE_DESCRIPTIONS = {
"official_zero": "官方 Wan: torch.zeros → 归一化空间 0 ≈ RGB128 中灰",
"black": "归一化空间 -1 ≈ RGB0 真黑（常被误称为 zero pixel）",
"first_frame_mean": "首帧逐通道均值铺满未知帧",
"imagenet_mean": "ImageNet RGB 均值映射到 [-1,1]",
}
def normalize_fill_mode(mode: Optional[str]) -> str:
if mode not in FILL_MODE_ALIASES:
known = sorted({k for k in FILL_MODE_ALIASES if k is not None})
raise ValueError(f"Unknown i2v_vae_unknown_frame_fill={mode!r}. Known: {known}")
return FILL_MODE_ALIASES[mode]
def resolve_padding_mode(pipe) -> Optional[str]:
return getattr(pipe, "i2v_vae_unknown_frame_fill", None)
def _rgb01_to_normalized(v: float) -> float:
return v * 2.0 - 1.0
def make_vae_unknown_frames(
num_pad_frames: int,
height: int,
width: int,
*,
device: torch.device,
dtype: torch.dtype,
mode: Optional[str] = None,
first_frame_chw: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""
构造 VAE encode 用的未知帧像素块 ``[3, T_pad, H, W]``（与首帧同 dtype/device）。
Args:
num_pad_frames: 待填充帧数 (通常 num_frames-1 或 num_frames-2)
first_frame_chw: 已 preprocess 的首帧 ``[3, H, W]``，``first_frame_mean`` 需要
"""
if num_pad_frames <= 0:
return torch.empty(3, 0, height, width, device=device, dtype=dtype)
key = normalize_fill_mode(mode)
shape = (3, num_pad_frames, height, width)
if key == "official_zero":
return torch.zeros(shape, device=device, dtype=dtype)
if key == "black":
return torch.full(shape, -1.0, device=device, dtype=dtype)
if key == "first_frame_mean":
if first_frame_chw is None:
raise ValueError("first_frame_mean requires first_frame_chw")
fm = first_frame_chw
if fm.dim() != 3 or fm.shape[0] != 3:
raise ValueError(f"first_frame_chw must be [3,H,W], got {tuple(fm.shape)}")
mean = fm.reshape(3, -1).mean(dim=1).view(3, 1, 1, 1)
return mean.expand(shape).contiguous()
if key == "imagenet_mean":
vals = torch.tensor(
[_rgb01_to_normalized(v) for v in _IMAGENET_MEAN_01],
device=device,
dtype=dtype,
)
return vals.view(3, 1, 1, 1).expand(shape).contiguous()
raise RuntimeError(f"Unhandled fill mode: {key}")
|