327 lines
17 KiB
Python
327 lines
17 KiB
Python
"""Direct, encoder-only MiniMax H3 video VAE implementation.
|
|
|
|
Mirrors the encoder half of ``upstream_vae.py`` so keyframe/reference images can
|
|
be encoded without ComfyUI. The encoder runs in FP16 like Comfy's keyframe VAE
|
|
path, and latent moments are upcast for mean/std normalization.
|
|
|
|
Causal-conv semantics: spatial padding is reflect; temporal padding is causal
|
|
(front-only zeros) with a stride grid that starts at the first input frame
|
|
(Comfy autopad "same" / ``causal``). For a single input frame the temporal taps
|
|
of the kernel are truncated (Comfy ``autopad="causal_zero"``) so a keyframe is
|
|
never convolved against zero frames.
|
|
|
|
Weights are read straight from the checkpoint's ``encoder.*`` / ``quant_conv.*``
|
|
keys (a direct name-for-name copy into plain tensors) and applied by the
|
|
stateless kernels below.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors import safe_open
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
|
|
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
|
IMAGENET_STD = (0.229, 0.224, 0.225)
|
|
LATENTS_MEAN = (0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075, -0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975, -0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923, -0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543, -0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279, -0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264)
|
|
LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037, 1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875, 0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647, 0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775, 2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265, 3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525)
|
|
|
|
CH = 128
|
|
CH_MULT = (1, 2, 2, 4, 4, 8)
|
|
SPACE_DOWN = (2, 2, 2, 2, 1, 1)
|
|
TIME_DOWN = (1, 2, 2, 1, 1, 1)
|
|
NUM_RES_BLOCKS = 2
|
|
VAE_RATIO = 16
|
|
NIN_LEVELS = frozenset({1, 3, 5})
|
|
DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3})
|
|
|
|
|
|
def _conv3d(x, weight, bias, stride=(1, 1, 1), padding=(0, 0, 0)):
|
|
if x.is_cuda and weight.dtype in (torch.float16, torch.bfloat16):
|
|
output = torch.cudnn_convolution(
|
|
x, weight, padding, stride, (1, 1, 1), 1,
|
|
benchmark=False, deterministic=False, allow_tf32=True,
|
|
)
|
|
if bias is not None:
|
|
output += bias.reshape(1, -1, 1, 1, 1)
|
|
return output
|
|
return F.conv3d(x, weight, bias, stride, padding)
|
|
|
|
|
|
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, temporal_causal):
|
|
"""Causal 3D conv (matches upstream_vae.CausalConv3d).
|
|
|
|
- ``spatial_padding > 0``: reflect H and W by ``spatial_padding`` on each side.
|
|
- ``temporal_causal``: front-zero T by ``kernel_size - 1`` (the reference's
|
|
``causal_padding[0] * 2`` for ``causal_padding[0]=1``, which is every
|
|
3D-causal conv in the H3 VAE). A single-frame input truncates the
|
|
temporal taps to the final input-aligned tap instead of convolving zero rows.
|
|
- If neither applies: no padding at all (reference early-return).
|
|
"""
|
|
if x.shape[2] == 1:
|
|
# Keyframe path (matches reference's `autopad="causal_zero"`):
|
|
# apply spatial-reflect pad (if any), then run an effective 2D conv
|
|
# by slicing the kernel to its final input-aligned temporal tap.
|
|
if spatial_padding > 0:
|
|
x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
|
|
kernel_5d = weight[:, :, -1:, :, :]
|
|
return _conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]))
|
|
if spatial_padding > 0:
|
|
x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
|
|
if temporal_causal:
|
|
x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0))
|
|
return _conv3d(x, weight, bias, stride)
|
|
|
|
|
|
def _group_norm_3d(x, weight, bias):
|
|
"""GroupNorm (32 groups, eps 1e-6) with per-frame statistics."""
|
|
b, c, t, h, w = x.shape
|
|
y = F.group_norm(x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w), 32, weight, bias, 1e-6)
|
|
return y.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()
|
|
|
|
|
|
def _resnet(x, p):
|
|
# nin_shortcut uses CausalConv3d(k=1, padding=1) in the reference.
|
|
residual = x if p["nin"] is None else _conv3d(x, p["nin"][0], p["nin"][1])
|
|
h = _causal_conv3d(F.silu(_group_norm_3d(x, p["norm1_w"], p["norm1_b"])), p["conv1_w"], p["conv1_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
h = _causal_conv3d(F.silu(_group_norm_3d(h, p["norm2_w"], p["norm2_b"])), p["conv2_w"], p["conv2_b"], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
return h.add_(residual)
|
|
|
|
|
|
def _downsample(x, p):
|
|
if p["space"] == 2:
|
|
# Reference Downsample3D pads only the right and bottom edges.
|
|
x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
|
|
# Conv uses padding=(1,0,0) -> causal_padding=(1,0,0), so spatial pad=0,
|
|
# temporal front-zero is applied.
|
|
return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0, temporal_causal=True)
|
|
|
|
|
|
def _encoder_run(x, E):
|
|
h = _causal_conv3d(x, E["conv_in"][0], E["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
for level in E["down"]:
|
|
for blk in level["blocks"]:
|
|
h = _resnet(h, blk)
|
|
if level["down"] is not None:
|
|
h = _downsample(h, level["down"])
|
|
h = F.silu(_group_norm_3d(h, E["norm_out_w"], E["norm_out_b"]))
|
|
return _causal_conv3d(h, E["conv_out"][0], E["conv_out"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
|
|
|
|
|
|
class MiniMaxH3VideoVAEEncoder(nn.Module):
|
|
"""Encoder-only H3 VAE. ``encode`` matches the public contract of upstream_vae.py.
|
|
|
|
Weights are plain tensors loaded from the checkpoint by canonical name into
|
|
``self.W`` (a dict), so no ``nn.Module`` sub-hierarchy is needed.
|
|
"""
|
|
|
|
def __init__(self, *, tiling: bool = True, compute_dtype: torch.dtype = torch.float16):
|
|
super().__init__()
|
|
self.compute_dtype = compute_dtype
|
|
self.vae_ratio, self.vae_ratio_t = VAE_RATIO, 4
|
|
self.clip_length, self.token_drop = 17, 3
|
|
self.frame_pre_padding = (-self.clip_length) % self.vae_ratio_t
|
|
self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
|
|
self.quant_conv = nn.Conv3d(48, 48, 1)
|
|
self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN), persistent=False)
|
|
self.register_buffer("latents_std", torch.tensor(LATENTS_STD), persistent=False)
|
|
self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False)
|
|
self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False)
|
|
|
|
def _required_encoder_names(self) -> list[str]:
|
|
names = [
|
|
"encoder.conv_in.weight", "encoder.conv_in.bias",
|
|
"encoder.norm_out.weight", "encoder.norm_out.bias",
|
|
"encoder.conv_out.weight", "encoder.conv_out.bias",
|
|
"quant_conv.weight", "quant_conv.bias",
|
|
"latents_mean", "latents_std",
|
|
]
|
|
for i in range(len(CH_MULT)):
|
|
for b in range(NUM_RES_BLOCKS):
|
|
base = f"encoder.down.{i}.block.{b}."
|
|
names += [
|
|
base + "conv1.weight", base + "conv1.bias",
|
|
base + "conv2.weight", base + "conv2.bias",
|
|
base + "norm1.weight", base + "norm1.bias",
|
|
base + "norm2.weight", base + "norm2.bias",
|
|
]
|
|
if b == 0 and i in NIN_LEVELS:
|
|
names += [base + "nin_shortcut.weight", base + "nin_shortcut.bias"]
|
|
if i in DOWNSAMPLE_LEVELS:
|
|
names += [f"encoder.down.{i}.downsample.conv.weight", f"encoder.down.{i}.downsample.conv.bias"]
|
|
return names
|
|
|
|
@classmethod
|
|
def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True, dtype: torch.dtype = torch.float16) -> "MiniMaxH3VideoVAEEncoder":
|
|
model = cls(tiling=tiling, compute_dtype=dtype)
|
|
names = model._required_encoder_names()
|
|
if os.getenv("H3_FAST_SAFETENSORS", "").lower() in {"1", "true", "yes", "on"}:
|
|
from fastsafetensors import fastsafe_open
|
|
|
|
fast_device = "cuda:0" if str(device) == "cuda" else str(device)
|
|
with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as ck:
|
|
available = set(ck.keys())
|
|
missing = [n for n in names if n not in available]
|
|
if missing:
|
|
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}")
|
|
W = {n: ck.get_tensor(n).clone().detach().to(dtype=dtype) for n in names}
|
|
elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
|
|
from safetensors.torch import load
|
|
|
|
with open(path, "rb") as file:
|
|
available_weights = load(file.read())
|
|
available = set(available_weights)
|
|
missing = [n for n in names if n not in available]
|
|
if missing:
|
|
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}")
|
|
W = {n: available_weights[n].to(device=device, dtype=dtype) for n in names}
|
|
del available_weights
|
|
else:
|
|
with safe_open(str(path), framework="pt", device=str(device)) as ck:
|
|
available = set(ck.keys())
|
|
missing = [n for n in names if n not in available]
|
|
if missing:
|
|
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}")
|
|
W = {n: ck.get_tensor(n).to(dtype=dtype).to(device) for n in names}
|
|
|
|
# Build the structured params dict.
|
|
down = []
|
|
for i in range(len(CH_MULT)):
|
|
mid = CH * CH_MULT[i]
|
|
blocks = []
|
|
for b in range(NUM_RES_BLOCKS):
|
|
base = f"encoder.down.{i}.block.{b}."
|
|
blk = {
|
|
"conv1_w": W[base + "conv1.weight"], "conv1_b": W[base + "conv1.bias"],
|
|
"conv2_w": W[base + "conv2.weight"], "conv2_b": W[base + "conv2.bias"],
|
|
"norm1_w": W[base + "norm1.weight"], "norm1_b": W[base + "norm1.bias"],
|
|
"norm2_w": W[base + "norm2.weight"], "norm2_b": W[base + "norm2.bias"],
|
|
"nin": None,
|
|
}
|
|
if base + "nin_shortcut.weight" in W:
|
|
blk["nin"] = (W[base + "nin_shortcut.weight"], W[base + "nin_shortcut.bias"])
|
|
blocks.append(blk)
|
|
level_down = None
|
|
if i in DOWNSAMPLE_LEVELS:
|
|
ds = f"encoder.down.{i}.downsample.conv."
|
|
level_down = {"w": W[ds + "weight"], "b": W[ds + "bias"], "time": TIME_DOWN[i], "space": SPACE_DOWN[i]}
|
|
down.append({"blocks": blocks, "down": level_down})
|
|
|
|
E = {
|
|
"conv_in": (W["encoder.conv_in.weight"], W["encoder.conv_in.bias"]),
|
|
"down": down,
|
|
"norm_out_w": W["encoder.norm_out.weight"], "norm_out_b": W["encoder.norm_out.bias"],
|
|
"conv_out": (W["encoder.conv_out.weight"], W["encoder.conv_out.bias"]),
|
|
}
|
|
model.W = E
|
|
model.quant_conv.to(device, dtype)
|
|
with torch.no_grad():
|
|
model.quant_conv.weight.copy_(W["quant_conv.weight"])
|
|
model.quant_conv.bias.copy_(W["quant_conv.bias"])
|
|
model.latents_mean.copy_(W["latents_mean"].float().cpu())
|
|
model.latents_std.copy_(W["latents_std"].float().cpu())
|
|
for b in ("latents_mean", "latents_std", "pixel_mean", "pixel_std"):
|
|
getattr(model, b).to(device)
|
|
return model
|
|
|
|
@torch.inference_mode()
|
|
def _encode_moments(self, x: torch.Tensor) -> torch.Tensor:
|
|
return _conv3d(_encoder_run(x.to(self.compute_dtype), self.W), self.quant_conv.weight, self.quant_conv.bias)
|
|
|
|
def _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self.tiling:
|
|
return self.tiled_encode(x)
|
|
return self._encode_moments(x)
|
|
|
|
def split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]:
|
|
if self.tile_size >= length:
|
|
return [0], [length], []
|
|
count = math.ceil(length / self.tile_size)
|
|
while self.tile_size * count - self.tile_overlap_min * (count - 1) < length:
|
|
count += 1
|
|
overlaps = [self.tile_overlap_min] * (count - 1)
|
|
for index in range((self.tile_size * count - sum(overlaps) - length) // self.vae_ratio):
|
|
overlaps[index % len(overlaps)] += self.vae_ratio
|
|
starts = [0]
|
|
for overlap in overlaps:
|
|
starts.append(starts[-1] + self.tile_size - overlap)
|
|
return starts, [self.tile_size] * count, overlaps
|
|
|
|
@staticmethod
|
|
def blend(a: torch.Tensor, b: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
|
|
extent = min(a.shape[dim], b.shape[dim], extent)
|
|
positions = torch.arange(extent, device=b.device, dtype=b.dtype)
|
|
weight_a = 1 - positions / extent
|
|
weight_b = positions / extent
|
|
shape = [1] * a.ndim
|
|
shape[dim] = extent
|
|
weight_a = weight_a.view(shape)
|
|
weight_b = weight_b.view(shape)
|
|
slice_a = [slice(None)] * a.ndim
|
|
slice_a[dim] = slice(-extent, None)
|
|
slice_b = [slice(None)] * a.ndim
|
|
slice_b[dim] = slice(0, extent)
|
|
blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
|
|
if extent < b.shape[dim]:
|
|
slice_b_rest = [slice(None)] * b.ndim
|
|
slice_b_rest[dim] = slice(extent, None)
|
|
return torch.cat((blended, b[tuple(slice_b_rest)]), dim=dim)
|
|
return blended
|
|
|
|
def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
height, width = x.shape[-2], x.shape[-1]
|
|
y_idx, y_len, y_overlap = self.split_tiles(height)
|
|
x_idx, x_len, x_overlap = self.split_tiles(width)
|
|
rows = [[self._encode_moments(x[..., i_pos:i_pos + i_len, j_pos:j_pos + j_len]) for j_pos, j_len in zip(x_idx, x_len)] for i_pos, i_len in zip(y_idx, y_len)]
|
|
latent_y_overlap = [o // self.vae_ratio for o in y_overlap]
|
|
latent_x_overlap = [o // self.vae_ratio for o in x_overlap]
|
|
result_rows = []
|
|
for i, row in enumerate(rows):
|
|
result_row = []
|
|
for j, tile in enumerate(row):
|
|
if i > 0:
|
|
tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
|
|
if j > 0:
|
|
tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
|
|
if i < len(rows) - 1:
|
|
tile = tile[..., :-latent_y_overlap[i], :]
|
|
if j < len(row) - 1:
|
|
tile = tile[..., :, :-latent_x_overlap[j]]
|
|
result_row.append(tile)
|
|
result_rows.append(torch.cat(result_row, dim=-1))
|
|
return torch.cat(result_rows, dim=-2)
|
|
|
|
def encode_temporal(self, x: torch.Tensor) -> torch.Tensor:
|
|
if x.shape[2] % self.clip_length != 0:
|
|
pad_size = (-x.shape[2]) % self.clip_length
|
|
x = torch.cat([x, x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)], dim=2)
|
|
num_chunks = x.shape[2] // self.clip_length
|
|
z_list = [self._adaptive_encode(x[:, :, i * self.clip_length : (i + 1) * self.clip_length, :, :]) for i in range(num_chunks)]
|
|
z = torch.cat(z_list, dim=2)
|
|
if self.token_drop > 0:
|
|
z = z[:, :, :-self.token_drop]
|
|
return z
|
|
|
|
def encode(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""``[B,3,H,W]`` or ``[B,3,T,H,W]`` pixels in ``[-1, 1]`` -> normalized latents ``[B,24,T_lat,H//16,W//16]``."""
|
|
if x.ndim == 4:
|
|
x = x.unsqueeze(2)
|
|
# Comfy's VAE wrapper casts before entering the model, so image
|
|
# normalization rounds in the VAE compute dtype as well.
|
|
x = x.to(self.compute_dtype)
|
|
x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
|
|
if x.shape[2] == 1:
|
|
moments = self._adaptive_encode(x)
|
|
moments = moments[:, :, -1:, :, :]
|
|
else:
|
|
moments = self.encode_temporal(x)
|
|
mean = torch.chunk(moments.float(), 2, dim=1)[0]
|
|
latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean)
|
|
latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean)
|
|
return (mean - latents_mean) / latents_std
|