Fix VAE encoder to load canonical checkpoint key names

This commit is contained in:
Daniel Maddern 2026-08-19 21:04:30 +07:00
parent 807bd64a82
commit 9bec53ac36

View file

@ -6,15 +6,18 @@ upcast to FP32 before mean/std normalization (the reference contract).
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 the frame
is not convolved against zero frames.
(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
@ -28,182 +31,152 @@ 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 _causal_front_padding(t_in: int, kernel: int, stride: int) -> int:
"""Front-only zero padding matching the reference causal 3D conv.
``padding = kernel - 1 - (t_out - 1) * stride`` with ``t_out = ceil(t_in /
stride)`` (the reference computes the output length from the unpadded
input). The result is non-negative: with a stride-s grid and a (2k+1)
kernel, ``ceil(t/stride) >= 1 + (t-1) // stride`` for all t, so
``(t_in - 1) % stride * stride >= (kernel - 1) % (2 * stride)``.
"""
"""Front-only zero padding: ``kernel-1-(ceil(t_in/stride)-1)*stride``."""
t_out = math.ceil(t_in / stride)
return max(0, kernel - 1 - (t_out - 1) * stride)
class _CausalConv3d(nn.Module):
"""3D conv: reflect spatial padding, causal (front-zero) temporal padding."""
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int | tuple[int, int, int] = 1, spatial_padding: int = 0):
super().__init__()
self.kernel_size = kernel_size
self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
self.spatial_padding = spatial_padding
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=self.stride, padding=(0, spatial_padding, spatial_padding))
def forward(self, x: torch.Tensor) -> torch.Tensor:
t = x.shape[2]
if t == 1:
# A single input frame never convolves against zero frames: the
# temporal taps are truncated (Comfy autopad="causal_zero").
if self.spatial_padding > 0:
half = (self.kernel_size - 1) // 2
kernel = self.conv.weight[:, :, half:half + 2 * self.spatial_padding + 1]
return F.conv3d(x, kernel, self.conv.bias, (1, self.stride[1], self.stride[2]), (0, self.spatial_padding, self.spatial_padding))
return self.conv(x)
front = _causal_front_padding(t, self.kernel_size, self.stride[0])
if front > 0:
x = F.pad(x, (0, 0, 0, 0, front, 0))
return self.conv(x)
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding):
"""Reflect-spatial / causal-temporal 3D conv (single-frame -> truncated taps)."""
t = x.shape[2]
if t == 1:
if spatial_padding > 0:
half = (kernel_size - 1) // 2
kernel = weight[:, :, half : half + 2 * spatial_padding + 1]
return F.conv3d(x, kernel, bias, (1, stride[1], stride[2]), (0, spatial_padding, spatial_padding))
return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))
front = _causal_front_padding(t, kernel_size, stride[0])
if front > 0:
x = F.pad(x, (0, 0, 0, 0, front, 0))
return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))
class TemporalIsolatedGroupNorm(nn.GroupNorm):
"""GroupNorm with statistics computed per frame (time merged into batch)."""
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.dim() != 5:
return super().forward(x)
b, c, t, h, w = x.shape
x = x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w)
x = super().forward(x)
return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()
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 group_norm_3d(num_channels: int) -> TemporalIsolatedGroupNorm:
return TemporalIsolatedGroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
def _resnet(x, p):
residual = x if p["nin"] is None else _causal_conv3d(x, p["nin"][0], p["nin"][1], kernel_size=1, stride=(1, 1, 1), spatial_padding=0)
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)
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)
return h.add_(residual)
class Downsample3D(nn.Module):
def __init__(self, in_channels: int, out_channels: int, time_stride: int = 1, space_stride: int = 2):
super().__init__()
self.space_stride = space_stride
self.conv = _CausalConv3d(in_channels, out_channels, kernel_size=3, stride=(time_stride, space_stride, space_stride), spatial_padding=0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
if self.space_stride == 2:
x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
return self.conv(x)
def _downsample(x, p):
if p["space"] == 2:
x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0)
class ResnetBlock3D(nn.Module):
def __init__(self, in_channels: int, out_channels: int | None = None):
super().__init__()
self.in_channels = in_channels
self.out_channels = in_channels if out_channels is None else out_channels
self.norm1 = group_norm_3d(in_channels)
self.norm2 = group_norm_3d(self.out_channels)
self.conv1 = _CausalConv3d(in_channels, self.out_channels, kernel_size=3, spatial_padding=1)
self.conv2 = _CausalConv3d(self.out_channels, self.out_channels, kernel_size=3, spatial_padding=1)
if self.in_channels != self.out_channels:
self.nin_shortcut = _CausalConv3d(self.in_channels, self.out_channels, kernel_size=1, spatial_padding=0)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.conv1(F.silu(self.norm1(x), inplace=True))
h = self.conv2(F.silu(self.norm2(h), inplace=True))
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return h.add_(x)
class EncoderFCN3D(nn.Module):
def __init__(self, ch: int, ch_mult: tuple[int, ...], space_down: tuple[int, ...], time_down: tuple[int, ...], num_res_blocks: int, in_channels: int, z_channels: int, double_z: bool = True):
super().__init__()
self.num_levels = len(ch_mult)
self.num_res_blocks = [num_res_blocks] * self.num_levels
block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
block_in = [block_mid[0]] + block_mid[:-1]
self.conv_in = _CausalConv3d(in_channels, block_in[0], kernel_size=3, spatial_padding=1)
self.down = nn.ModuleList()
for i_level in range(self.num_levels):
down = nn.Module()
down.block = nn.ModuleList(
ResnetBlock3D(block_in[i_level] if i == 0 else block_mid[i_level], block_mid[i_level])
for i in range(self.num_res_blocks[i_level])
)
if space_down[i_level] * time_down[i_level] > 1:
down.downsample = Downsample3D(block_mid[i_level], block_mid[i_level], time_stride=time_down[i_level], space_stride=space_down[i_level])
self.down.append(down)
self.norm_out = group_norm_3d(block_mid[-1])
self.conv_out = _CausalConv3d(block_mid[-1], 2 * z_channels if double_z else z_channels, kernel_size=3, spatial_padding=1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.conv_in(x)
for i_level in range(self.num_levels):
for i_block in range(self.num_res_blocks[i_level]):
h = self.down[i_level].block[i_block](h)
if hasattr(self.down[i_level], "downsample"):
h = self.down[i_level].downsample(h)
h = F.silu(self.norm_out(h), inplace=True)
return self.conv_out(h)
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)
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)
class MiniMaxH3VideoVAEEncoder(nn.Module):
"""Encoder-only H3 VAE. ``encode`` matches the public contract of upstream_vae.py."""
"""Encoder-only H3 VAE. ``encode`` matches the public contract of upstream_vae.py.
def __init__(self, *, device: torch.device | str | None = None, tiling: bool = True):
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):
super().__init__()
self.vae_ratio, self.vae_ratio_t = 16, 4
self.vae_ratio, self.vae_ratio_t = VAE_RATIO, 4
self.clip_length, self.token_drop = 17, 3
self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64
self.encoder = EncoderFCN3D(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, in_channels=3, z_channels=24, double_z=True)
self.quant_conv = nn.Conv3d(48, 48, 1, device=device)
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"]
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) -> "MiniMaxH3VideoVAEEncoder":
model = cls(device="meta", tiling=tiling)
expected = model.state_dict()
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)
available = set()
with fastsafe_open(filenames=[str(path)], nogds=True, device=fast_device) as checkpoint:
available = set(checkpoint.keys())
missing = sorted(set(expected) - available)
weights = {name: checkpoint.get_tensor(name).clone().detach() for name in expected}
model = cls(tiling=tiling)
names = model._required_encoder_names()
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(missing)}")
weights = {name: t.to(device=device, dtype=torch.float32) for name, t in weights.items()}
elif os.getenv("H3_DISABLE_MMAP", "").lower() in {"1", "true", "yes", "on"}:
from safetensors.torch import load
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(sorted(missing)[:16])}")
W = {n: ck.get_tensor(n).to(dtype=torch.float32).to(device) for n in names}
with open(path, "rb") as file:
available_weights = load(file.read())
available = set(available_weights)
missing = sorted(set(expected) - available)
if missing:
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(missing)}")
weights = {name: available_weights[name].to(device=device, dtype=torch.float32) for name in expected}
else:
with safe_open(str(path), framework="pt", device=str(device)) as checkpoint:
available = set(checkpoint.keys())
shape_errors = [(name, tuple(expected[name].shape), tuple(checkpoint.get_slice(name).get_shape())) for name in expected if name in available and tuple(expected[name].shape) != tuple(checkpoint.get_slice(name).get_shape())]
if shape_errors:
raise ValueError(f"incompatible H3 VAE checkpoint; shape mismatch: {shape_errors}")
weights = {name: checkpoint.get_tensor(name).to(dtype=torch.float32) for name in expected}
missing = sorted(set(expected) - available)
if missing:
raise ValueError(f"incompatible H3 VAE checkpoint; missing: {', '.join(missing)}")
model.load_state_dict(weights, strict=True, assign=True)
# 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, torch.float32)
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 self.quant_conv(self.encoder(x))
return F.conv3d(_encoder_run(x.to(torch.float32), self.W), self.quant_conv.weight, self.quant_conv.bias)
def _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor:
if self.tiling:
@ -259,7 +232,7 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
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)
tile = self.blend(row[j - 1], tile, latent_x_overlap[j], dim=-1)
if i < len(rows) - 1:
tile = tile[..., :-latent_y_overlap[i], :]
if j < len(row) - 1:
@ -273,14 +246,14 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
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_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, T, H, W]`` pixels in ``[-1, 1]`` -> normalized latents ``[B, 24, T_lat, H//16, W//16]``."""
"""``[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)
x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))