"""Direct, decoder-only MiniMax H3 video VAE implementation.""" 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) def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor: result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype) return result if weight is None else result * weight.to(dtype=x.dtype) def create_token_ids(patch_dims: tuple[int, int, int], device: torch.device, dtype: torch.dtype) -> torch.Tensor: coords = [2.0 * (torch.arange(0.5, size, dtype=dtype, device=device) / size) - 1.0 for size in patch_dims] return torch.stack(torch.meshgrid(*coords, indexing="ij"), dim=-1).flatten(0, 2).unsqueeze(0) class RotaryEmbeddingND(nn.Module): def __init__(self, dim: int, rotary_base: float = 100.0, n_dim: int = 3, *, device=None): super().__init__() self.rotary_base = rotary_base self.step = 2 * n_dim / dim inv_freq = 1 / rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=device) self.register_buffer("inv_freq", inv_freq, persistent=False) self.angle_scale = 2.0 * math.pi def forward(self, img_ids: torch.Tensor) -> torch.Tensor: inv_freq = self.inv_freq if inv_freq.device.type == "meta": inv_freq = 1 / self.rotary_base ** torch.arange(0, 1, self.step, dtype=torch.float32, device=img_ids.device) else: inv_freq = inv_freq.to(img_ids.device) angles = self.angle_scale * img_ids[:, :, :, None].float() * inv_freq[None, None, None, :] angles = angles.flatten(2, 3) cos, sin = torch.cos(angles), torch.sin(angles) return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(*angles.shape[:2], 1, -1, 2, 2).to(img_ids.dtype) class RMSNorm(nn.Module): def __init__(self, dim: int, eps: float, affine: bool, *, device=None): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.empty(dim, device=device)) if affine else None def forward(self, x: torch.Tensor) -> torch.Tensor: return _rms_norm(x, self.weight, self.eps) class FeedForward(nn.Module): def __init__(self, dim: int, bias: bool = True, *, device=None): super().__init__() self.w1 = nn.Linear(dim, dim * 8, bias=bias, device=device) self.w2 = nn.Linear(dim * 4, dim, bias=bias, device=device) def forward(self, x: torch.Tensor) -> torch.Tensor: gate, value = self.w1(x).chunk(2, dim=-1) return self.w2(F.silu(gate) * value) def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor: """Apply the reference split-half RoPE layout to leading rotary channels.""" pairs = table.shape[-3] rot = pairs * 2 first, second = x[..., :pairs], x[..., pairs:rot] cos, neg_sin, sin = table[..., 0, 0], table[..., 0, 1], table[..., 1, 0] rotated = torch.cat((first * cos + second * neg_sin, first * sin + second * cos), dim=-1) return torch.cat((rotated, x[..., rot:]), dim=-1) class Attention(nn.Module): def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None): super().__init__() dim = heads * dim_head self.heads, self.dim_head = heads, dim_head self.norm_q = RMSNorm(dim_head, eps, False, device=device) self.norm_k = RMSNorm(dim_head, eps, False, device=device) self.to_qkv = nn.Linear(dim, dim * 3, bias=bias, device=device) self.to_out = nn.Linear(dim, dim, bias=bias, device=device) def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor: batch, sequence, _ = x.shape qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head) query, key, value = qkv.chunk(3, dim=-1) query, key = self.norm_q(query), self.norm_k(key) query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb) output = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)) return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).nan_to_num_(0.0)) class TransformerBlock(nn.Module): def __init__(self, heads: int, dim_head: int, bias: bool = True, eps: float = 1e-5, *, device=None): super().__init__() dim = heads * dim_head self.norm1 = RMSNorm(dim, eps, True, device=device) self.attn = Attention(heads, dim_head, bias, eps, device=device) self.scale1 = nn.Parameter(torch.empty(dim, device=device)) self.norm2 = RMSNorm(dim, eps, True, device=device) self.ff = FeedForward(dim, bias, device=device) self.scale2 = nn.Parameter(torch.empty(dim, device=device)) def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor: x = x + self.attn(self.norm1(x), rotary_pos_emb) * self.scale1.to(x.dtype) return x + self.ff(self.norm2(x)) * self.scale2.to(x.dtype) class ViT3DDecoder(nn.Module): def __init__(self, patch_size: int = 16, patch_size_t: int = 4, in_channels: int = 24, out_channels: int = 3, num_layers: int = 36, heads: int = 32, dim_head: int = 64, rope_theta: float = 100.0, rope_dim_ratio: float = 0.75, bias: bool = True, eps: float = 1e-5, num_register_tokens: int = 4, *, device=None): super().__init__() dim = heads * dim_head self.patch_size, self.patch_size_t, self.out_channels = patch_size, patch_size_t, out_channels self.num_register_tokens = num_register_tokens self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, device=device) self.x_embedder = nn.Linear(in_channels, dim, device=device) self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim, device=device)) self.register_buffer("mask_token", torch.empty(1, 1, dim, device=device)) self.transformer_blocks = nn.ModuleList([TransformerBlock(heads, dim_head, bias, eps, device=device) for _ in range(num_layers)]) self.norm_out = nn.LayerNorm(dim, eps=eps, device=device) self.proj_out = nn.Linear(dim, out_channels * patch_size_t * patch_size * patch_size, device=device) def forward(self, x: torch.Tensor) -> torch.Tensor: batch, _, latent_t, latent_h, latent_w = x.shape h = self.x_embedder(x.flatten(2).transpose(1, 2)) patches = h.shape[1] h = torch.cat((h, self.register_tokens.to(h).expand(batch, -1, -1), torch.zeros_like(h[:, :1])), dim=1) ids = create_token_ids((latent_t, latent_h, latent_w), x.device, x.dtype).expand(batch, -1, -1) ids = torch.cat((ids, torch.zeros(batch, 1 + self.num_register_tokens, 3, device=x.device, dtype=x.dtype)), dim=1) rope = self.pos_embed(ids) for block in self.transformer_blocks: h = block(h, rope) output = self.proj_out(self.norm_out(h))[:, :patches] output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size) return output.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size) class MiniMaxH3VideoVAE(nn.Module): """Decoder-only H3 VAE. The public ``decode`` contract matches upstream_vae.py.""" def __init__(self, *, device=None, tiling: bool = True): super().__init__() self.vae_ratio, self.vae_ratio_t = 16, 4 self.clip_length, self.token_drop = 17, 3 self.tokens_chunk_size, self.token_overlap = 5, 3 self.frame_pre_padding, self.frame_overlap = 3, 9 self.tiling, self.tile_size, self.tile_overlap_min = tiling, 256, 64 self.post_quant_conv = nn.Conv3d(24, 24, 1, device=device) self.decoder = ViT3DDecoder(device=device) self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN, device=device)) self.register_buffer("latents_std", torch.tensor(LATENTS_STD, device=device)) self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN, device=device).view(1, 3, 1, 1, 1), persistent=False) self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD, device=device).view(1, 3, 1, 1, 1), persistent=False) @classmethod def from_safetensors(cls, path: str | Path, *, device: str | torch.device = "cuda", tiling: bool = True) -> "MiniMaxH3VideoVAE": model = cls(device="meta", tiling=tiling) expected = model.state_dict() if 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 = sorted(set(expected) - available) shape_errors = [(name, tuple(expected[name].shape), tuple(available_weights[name].shape)) for name in expected if name in available and tuple(expected[name].shape) != tuple(available_weights[name].shape)] if missing or shape_errors: details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details)) weights = {name: available_weights[name].to(device) for name in expected} else: with safe_open(str(path), framework="pt", device=str(device)) as checkpoint: available = set(checkpoint.keys()) missing = sorted(set(expected) - available) 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 missing or shape_errors: details = ([f"missing: {', '.join(missing)}"] if missing else []) + ([f"shape mismatch: {shape_errors}"] if shape_errors else []) raise ValueError("incompatible H3 VAE checkpoint; " + "; ".join(details)) weights = {name: checkpoint.get_tensor(name) for name in expected} model.load_state_dict(weights, strict=True, assign=True) return model def _decode_pixels(self, z: torch.Tensor) -> torch.Tensor: return self.decoder(self.post_quant_conv(z)) 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) shape = [1] * a.ndim shape[dim] = extent position = torch.arange(extent, device=b.device, dtype=b.dtype).view(shape) blended = a.narrow(dim, a.shape[dim] - extent, extent) * (1 - position / extent) + b.narrow(dim, 0, extent) * (position / extent) return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim) def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio y_starts, y_lengths, y_overlaps = self.split_tiles(height) x_starts, x_lengths, x_overlaps = self.split_tiles(width) canvas = None row_tails = [] output_y = 0 for row_index, (y, tile_height) in enumerate(zip(y_starts, y_lengths)): next_row_tails = [] left_tail = None output_x = 0 for column_index, (x, tile_width) in enumerate(zip(x_starts, x_lengths)): tile = self._decode_pixels(z[..., y // self.vae_ratio:(y + tile_height) // self.vae_ratio, x // self.vae_ratio:(x + tile_width) // self.vae_ratio]) if row_index < len(y_starts) - 1: next_row_tails.append(tile[..., -y_overlaps[row_index]:, :].clone()) next_left_tail = tile[..., :, -x_overlaps[column_index]:].clone() if column_index < len(x_starts) - 1 else None if row_index: tile = self.blend(row_tails[column_index], tile, y_overlaps[row_index - 1], -2) if column_index: tile = self.blend(left_tail, tile, x_overlaps[column_index - 1], -1) left_tail = next_left_tail if row_index < len(y_starts) - 1: tile = tile[..., :-y_overlaps[row_index], :] if column_index < len(x_starts) - 1: tile = tile[..., :, :-x_overlaps[column_index]] if canvas is None: canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device) canvas[..., output_y:output_y + tile.shape[-2], output_x:output_x + tile.shape[-1]].copy_(tile) output_x += tile.shape[-1] row_tails = next_row_tails output_y += tile.shape[-2] return canvas def _decode_temporal_pad_frames(self, z_len: int, pad_tokens: int) -> int: if pad_tokens <= 0: return 0 return sum(1 if (z_len - pad_tokens + index) % self.tokens_chunk_size == 0 else self.vae_ratio_t for index in range(pad_tokens)) def _decode_temporal_frame_plan(self, z_len: int, chunks: int, pad_tokens: int) -> int: total, final_overlap = 0, 0 for index in range(chunks): tokens = max(0, min(index * 5 + 8, z_len) - min(index * 5, z_len)) frames = tokens * self.vae_ratio_t for split in range(2): part = max(0, min((split + 1) * 20, frames) - split * 20 - self.frame_pre_padding) if split == 0: total += part else: final_overlap = part return total + final_overlap - self._decode_temporal_pad_frames(z_len, pad_tokens) def decode_temporal(self, z: torch.Tensor) -> torch.Tensor: pseudo_tokens = z.shape[2] + self.token_drop pad_tokens = (-pseudo_tokens) % self.tokens_chunk_size pseudo_tokens += pad_tokens chunks = pseudo_tokens // self.tokens_chunk_size - 1 if chunks < 1: pad_tokens += self.tokens_chunk_size chunks += 1 if pad_tokens: z = torch.cat((z, z[:, :, -1:].expand(-1, -1, pad_tokens, -1, -1)), dim=2) output_frames = self._decode_temporal_frame_plan(z.shape[2], chunks, pad_tokens) output, overlap = [], None for index in range(chunks): clip = self._adaptive_decode(z[:, :, index * 5:index * 5 + 8]) first = clip[:, :, :20, :, :][:, :, self.frame_pre_padding:] tail = clip[:, :, 20:40, :, :][:, :, self.frame_pre_padding:] if overlap is not None: first = self.blend(overlap, first, self.frame_overlap, -3) output.append(first) overlap = tail if overlap is not None: output.append(overlap) return torch.cat(output, dim=2)[:, :, :output_frames] def _adaptive_decode(self, z: torch.Tensor) -> torch.Tensor: return self.tiled_decode(z) if self.tiling else self._decode_pixels(z) def decode(self, z: torch.Tensor) -> torch.Tensor: z = z * self.latents_std.view(1, -1, 1, 1, 1).to(z) + self.latents_mean.view(1, -1, 1, 1, 1).to(z) decoded = self._adaptive_decode(z) if z.shape[2] == 1 else self.decode_temporal(z) if z.shape[2] == 1: decoded = decoded[:, :, -1:] pixel_std = torch.tensor(IMAGENET_STD, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1) pixel_mean = torch.tensor(IMAGENET_MEAN, device=decoded.device, dtype=decoded.dtype).view(1, 3, 1, 1, 1) return (decoded.float() * pixel_std.float() + pixel_mean.float()).clamp_(0, 1).mul_(2).sub_(1)