"""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