minimax_h3_video_vae
MiniMax-H3 video VAE for the Apple Silicon MLX runtime.
Faithful MLX port of fastvideo/models/vaes/minimax_h3_video.py (itself
parity-validated against the official diffusers implementation):
- Encoder: causal 3D CNN — causal temporal padding, reflect spatial
padding, per-frame GroupNorm, residual blocks, spatial/temporal downsampling,
then a 1x1x1
quant_convproducing mean/logvar channels. - Decoder: 1x1x1
post_quant_conv, then a 36-layer ViT — per-head Q/K RMSNorm, three-axis rotary embedding (theta 100, rotary width 48 of 64 head dims), register tokens plus a final zero class token, SwiGLU feed-forward, residual scale vectors, FP32 norm accumulation, final LayerNorm, output projection, and channel-major unpatchification (temporal patch 4, spatial patch 16x16) — with the exact clip chunking (clip_length=17,token_drop=3), frame pre-padding, overlap blending, tail trimming, and optional spatial tiling of the released model.
Released weights are FP32; the loader streams shards so peak memory stays bounded, and can optionally store bf16/fp16 after callers have measured the dtype drift against the FP32 acceptance gate.
Production code here never imports PyTorch. Torch parity references live in
the tests under tests/local_tests/minimax_h3/.
fastvideo.mlx_runtime.minimax_h3_video_vae.init_logger
init_logger(name: str) -> _FastvideoLoggerThe main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.
fastvideo.mlx_runtime.minimax_h3_video_vae.QuantizedWeight
QuantizedWeight(weight: Any, scales: Any, biases: Any, mode: str, group_size: int, bits: int, output_scale: float = 1.0)A packed linear weight; output_scale undoes the pre-quantization tensor scale.
fastvideo.mlx_runtime.minimax_h3_video_vae.QuantizedWeight.matmul
matmul(x)fastvideo.mlx_runtime.minimax_h3_video_vae.QuantizedWeight.dequantize
dequantize(dtype = mx.float32)fastvideo.mlx_runtime.minimax_h3_video_vae.is_block_linear_weight
is_block_linear_weight(key: str) -> boolTrue for the weight of one of the 216 decoder transformer-block linears.
fastvideo.mlx_runtime.minimax_h3_video_vae.packed_nvfp4_weight
packed_nvfp4_weight(packed_weight, inv_scale, global_sf) -> QuantizedWeightOne linear of a fastvideo_h3_decoder_nvfp4_packed_v1 file as an MLX NVFP4 weight.
fastvideo.mlx_runtime.minimax_h3_video_vae.quantize_weight
quantize_weight(weight, mode: str) -> QuantizedWeightQuantize a dense (out, in) matrix; mode is a key of QUANT_MODES.
fastvideo.mlx_runtime.minimax_h3_video_vae.MiniMaxH3VideoVAEConfigView
MiniMaxH3VideoVAEConfigView(in_channels: int = 3, out_channels: int = 3, latent_channels: int = 24, block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024), layers_per_block: int = 2, spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1), temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1), norm_num_groups: int = 32, norm_eps: float = 1e-06, decoder_num_layers: int = 36, decoder_num_attention_heads: int = 32, decoder_attention_head_dim: int = 64, decoder_num_register_tokens: int = 4, decoder_ffn_mult: int = 4, decoder_rope_theta: float = 100.0, decoder_rope_dim_ratio: float = 0.75, decoder_norm_eps: float = 1e-05, clip_length: int = 17, token_drop: int = 3, latents_mean: tuple[float, ...] | None = None, latents_std: tuple[float, ...] | None = None)Architecture constants (defaults mirror the released vae/config.json).
Source
@dataclass(frozen=True)
class MiniMaxH3VideoVAEConfigView:
"""Architecture constants (defaults mirror the released vae/config.json)."""
in_channels: int = 3
out_channels: int = 3
latent_channels: int = 24
block_out_channels: tuple[int, ...] = (128, 256, 256, 512, 512, 1024)
layers_per_block: int = 2
spatial_downsample_factors: tuple[int, ...] = (2, 2, 2, 2, 1, 1)
temporal_downsample_factors: tuple[int, ...] = (1, 2, 2, 1, 1, 1)
norm_num_groups: int = 32
norm_eps: float = 1e-6
decoder_num_layers: int = 36
decoder_num_attention_heads: int = 32
decoder_attention_head_dim: int = 64
decoder_num_register_tokens: int = 4
decoder_ffn_mult: int = 4
decoder_rope_theta: float = 100.0
decoder_rope_dim_ratio: float = 0.75
decoder_norm_eps: float = 1e-5
clip_length: int = 17
token_drop: int = 3
latents_mean: tuple[float, ...] | None = None
latents_std: tuple[float, ...] | None = None
@classmethod
def from_vae_dir(cls, vae_dir: str | Path) -> MiniMaxH3VideoVAEConfigView:
return cls.from_dict(json.loads((Path(vae_dir) / "config.json").read_text()))
@classmethod
def from_dict(cls, config: dict[str, Any]) -> MiniMaxH3VideoVAEConfigView:
"""Build from a ``vae/config.json`` dict (also LightVAE's embedded ``teacher_config``)."""
latents_mean = tuple(float(v) for v in config.get("latents_mean", ()))
latents_std = tuple(float(v) for v in config.get("latents_std", ()))
return cls(
in_channels=int(config["in_channels"]),
out_channels=int(config["out_channels"]),
latent_channels=int(config["latent_channels"]),
block_out_channels=tuple(int(v) for v in config["block_out_channels"]),
layers_per_block=int(config["layers_per_block"]),
spatial_downsample_factors=tuple(int(v) for v in config["spatial_downsample_factors"]),
temporal_downsample_factors=tuple(int(v) for v in config["temporal_downsample_factors"]),
norm_num_groups=int(config["norm_num_groups"]),
norm_eps=float(config["norm_eps"]),
decoder_num_layers=int(config["decoder_num_layers"]),
decoder_num_attention_heads=int(config["decoder_num_attention_heads"]),
decoder_attention_head_dim=int(config["decoder_attention_head_dim"]),
decoder_num_register_tokens=int(config["decoder_num_register_tokens"]),
decoder_ffn_mult=int(config.get("decoder_ffn_mult", 4)),
decoder_rope_theta=float(config["decoder_rope_theta"]),
decoder_rope_dim_ratio=float(config["decoder_rope_dim_ratio"]),
decoder_norm_eps=float(config["decoder_norm_eps"]),
clip_length=int(config["clip_length"]),
token_drop=int(config["token_drop"]),
latents_mean=latents_mean or None,
latents_std=latents_std or None,
)
@property
def spatial_compression_ratio(self) -> int:
return math.prod(self.spatial_downsample_factors)
@property
def temporal_compression_ratio(self) -> int:
return math.prod(self.temporal_downsample_factors)
@property
def tokens_chunk_size(self) -> int:
"""Latent frames decoded per clip chunk (ceil(clip_length / ratio))."""
return math.ceil(self.clip_length / self.temporal_compression_ratio)
@property
def token_overlap(self) -> int:
return (-self.token_drop) % self.tokens_chunk_size
@property
def frame_pre_padding(self) -> int:
return (-self.clip_length) % self.temporal_compression_ratio
@property
def frame_overlap(self) -> int:
return max(self.token_overlap * self.temporal_compression_ratio - self.frame_pre_padding, 0)</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_video_vae.MiniMaxH3VideoVAEConfigView.from_vae_dir
from_vae_dir(vae_dir: str | Path) -> MiniMaxH3VideoVAEConfigView
Source
@classmethod
def from_vae_dir(cls, vae_dir: str | Path) -> MiniMaxH3VideoVAEConfigView:
return cls.from_dict(json.loads((Path(vae_dir) / "config.json").read_text()))
fastvideo.mlx_runtime.minimax_h3_video_vae.MiniMaxH3VideoVAEConfigView.from_dict
from_dict(config: dict[str, Any]) -> MiniMaxH3VideoVAEConfigView
Build from a vae/config.json dict (also LightVAE's embedded teacher_config).
Source
@classmethod
def from_dict(cls, config: dict[str, Any]) -> MiniMaxH3VideoVAEConfigView:
"""Build from a ``vae/config.json`` dict (also LightVAE's embedded ``teacher_config``)."""
latents_mean = tuple(float(v) for v in config.get("latents_mean", ()))
latents_std = tuple(float(v) for v in config.get("latents_std", ()))
return cls(
in_channels=int(config["in_channels"]),
out_channels=int(config["out_channels"]),
latent_channels=int(config["latent_channels"]),
block_out_channels=tuple(int(v) for v in config["block_out_channels"]),
layers_per_block=int(config["layers_per_block"]),
spatial_downsample_factors=tuple(int(v) for v in config["spatial_downsample_factors"]),
temporal_downsample_factors=tuple(int(v) for v in config["temporal_downsample_factors"]),
norm_num_groups=int(config["norm_num_groups"]),
norm_eps=float(config["norm_eps"]),
decoder_num_layers=int(config["decoder_num_layers"]),
decoder_num_attention_heads=int(config["decoder_num_attention_heads"]),
decoder_attention_head_dim=int(config["decoder_attention_head_dim"]),
decoder_num_register_tokens=int(config["decoder_num_register_tokens"]),
decoder_ffn_mult=int(config.get("decoder_ffn_mult", 4)),
decoder_rope_theta=float(config["decoder_rope_theta"]),
decoder_rope_dim_ratio=float(config["decoder_rope_dim_ratio"]),
decoder_norm_eps=float(config["decoder_norm_eps"]),
clip_length=int(config["clip_length"]),
token_drop=int(config["token_drop"]),
latents_mean=latents_mean or None,
latents_std=latents_std or None,
)
fastvideo.mlx_runtime.minimax_h3_video_vae.MiniMaxH3VideoVAEConfigView.tokens_chunk_size
tokens_chunk_size: int
Latent frames decoded per clip chunk (ceil(clip_length / ratio)).
Source
def tokens_chunk_size(self) -> int:
"""Latent frames decoded per clip chunk (ceil(clip_length / ratio))."""
return math.ceil(self.clip_length / self.temporal_compression_ratio)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE
MLXMiniMaxH3VideoVAE(weights: dict[str, Any], config: MiniMaxH3VideoVAEConfigView, *, has_encoder: bool = True, has_decoder: bool = True)
Encoder + decoder for the released MiniMax-H3 video VAE.
Source
class MLXMiniMaxH3VideoVAE:
"""Encoder + decoder for the released MiniMax-H3 video VAE."""
def __init__(self,
weights: dict[str, Any],
config: MiniMaxH3VideoVAEConfigView,
*,
has_encoder: bool = True,
has_decoder: bool = True):
self.weights = weights
self.config = config
self.has_encoder = has_encoder
self.has_decoder = has_decoder
self.latent_channels = config.latent_channels
self.spatial_compression_ratio = config.spatial_compression_ratio
self.temporal_compression_ratio = config.temporal_compression_ratio
self.num_heads = config.decoder_num_attention_heads
self.head_dim = config.decoder_attention_head_dim
self.dim = self.num_heads * self.head_dim
self.rotary_dim = int(self.head_dim * config.decoder_rope_dim_ratio)
self._blocks: list[dict[str, Any]] | None = None
arch = config
mean = arch.latents_mean if arch.latents_mean is not None else [0.0] * arch.latent_channels
std = arch.latents_std if arch.latents_std is not None else [1.0] * arch.latent_channels
self._latents_mean = np.asarray(mean, dtype=np.float32).reshape(1, -1, 1, 1, 1)
self._latents_std = np.asarray(std, dtype=np.float32).reshape(1, -1, 1, 1, 1)
self._pixel_mean = np.asarray(PIXEL_MEAN, dtype=np.float32).reshape(1, -1, 1, 1, 1)
self._pixel_std = np.asarray(PIXEL_STD, dtype=np.float32).reshape(1, -1, 1, 1, 1)
# -- normalization -------------------------------------------------
def normalize_latents(self, latents):
return (latents - mx.array(self._latents_mean)) / mx.array(self._latents_std)
def denormalize_latents(self, latents):
return latents * mx.array(self._latents_std) + mx.array(self._latents_mean)
def normalize_pixels(self, pixels):
return (pixels - mx.array(self._pixel_mean)) / mx.array(self._pixel_std)
def denormalize_pixels(self, sample):
return sample * mx.array(self._pixel_std) + mx.array(self._pixel_mean)
# -- encoder ---------------------------------------------------------
def _encode_clip(self, x):
"""(N, 3, T, H, W) pixels -> (N, 2C, T', H', W') moments."""
w = self.weights
cfg = self.config
x = x.astype(w["encoder.conv_in.weight"].dtype) # compute in the storage dtype
x = _causal_conv3d(x,
w["encoder.conv_in.weight"],
w["encoder.conv_in.bias"],
spatial_padding=1,
temporal_padding=2)
for index in range(len(cfg.block_out_channels)):
td = cfg.temporal_downsample_factors[index]
sd = cfg.spatial_downsample_factors[index]
for layer in range(cfg.layers_per_block):
x = _resnet_block(w,
f"encoder.down_blocks.{index}.resnets.{layer}",
x,
num_groups=cfg.norm_num_groups,
eps=cfg.norm_eps)
mx.eval(x) # bound the lazy graph: full-resolution activations are the encoder's peak
if td * sd > 1:
if sd == 2:
x = _reflect_pad_axis(x, 3, 0, 1)
x = _reflect_pad_axis(x, 4, 0, 1)
x = _causal_conv3d(
x,
w[f"encoder.down_blocks.{index}.downsamplers.0.conv.weight"],
w[f"encoder.down_blocks.{index}.downsamplers.0.conv.bias"],
stride=(td, sd, sd),
temporal_padding=2,
)
mx.eval(x)
x = _group_norm_per_frame(x,
w["encoder.norm_out.weight"],
w["encoder.norm_out.bias"],
num_groups=cfg.norm_num_groups,
eps=cfg.norm_eps,
silu=True)
x = _causal_conv3d(x,
w["encoder.conv_out.weight"],
w["encoder.conv_out.bias"],
spatial_padding=1,
temporal_padding=2)
x = _causal_conv3d(x, w["quant_conv.weight"], w["quant_conv.bias"])
return x.astype(mx.float32)
def encode(self, pixels):
"""Normalized pixels (1, 3, T, H, W) -> (mean, logvar) each (1, C, T', H', W').
Pads the clip to ``clip_length`` and drops ``token_drop`` trailing
moment frames exactly like the reference ``_encode``.
"""
if not self.has_encoder:
raise RuntimeError("This MLX H3 video VAE was loaded without encoder weights.")
cfg = self.config
num_frames = pixels.shape[2]
if num_frames % cfg.clip_length != 0:
pad_frames = (-num_frames) % cfg.clip_length
tail = mx.repeat(pixels[:, :, -1:], pad_frames, axis=2)
pixels = mx.concatenate([pixels, tail], axis=2)
moment_chunks: list[Any] = []
for index in range(pixels.shape[2] // cfg.clip_length):
clip = pixels[:, :, index * cfg.clip_length:(index + 1) * cfg.clip_length]
moment_chunks.append(self._encode_clip(clip))
moments = mx.concatenate(moment_chunks, axis=2)
if cfg.token_drop > 0:
moments = moments[:, :, :-cfg.token_drop]
mean, logvar = mx.split(moments, 2, axis=1)
logvar = mx.clip(logvar, -30.0, 20.0)
return mean, logvar
def encode_keyframe(self, pixels):
"""Single-frame conditioning encode without chunk padding."""
if pixels.shape[2] != 1:
raise ValueError(f"encode_keyframe expects exactly one frame, got {pixels.shape}.")
moments = self._encode_clip(pixels)
mean, logvar = mx.split(moments, 2, axis=1)
logvar = mx.clip(logvar, -30.0, 20.0)
return mean, logvar
@staticmethod
def sample_posterior(mean, logvar, noise):
"""Reparameterization with an explicit noise array (deterministic parity)."""
return mean + mx.exp(0.5 * logvar) * noise
# -- decoder ---------------------------------------------------------
def _decoder_blocks(self) -> list[dict[str, Any]]:
"""Split the flat weight dict into per-block dicts once."""
if self._blocks is None:
blocks: list[dict[str, Any]] = [{} for _ in range(self.config.decoder_num_layers)]
prefix = "decoder.transformer_blocks."
for key, value in self.weights.items():
if key.startswith(prefix):
index, sub = key[len(prefix):].split(".", 1)
blocks[int(index)][sub] = value
self._blocks = blocks
return self._blocks
def _decode_clip(self, z):
"""(1, C, T_tokens, H', W') latents -> (1, 3, T_tokens*4, H, W) pixels."""
w = self.weights
cfg = self.config
z = z.astype(w["post_quant_conv.weight"].dtype) # compute in the storage dtype
z = _causal_conv3d(z, w["post_quant_conv.weight"], w["post_quant_conv.bias"])
batch, channels, num_frames, height, width = z.shape
# (B, C, T, H, W) -> (B*T*H*W, C): sequence-major tokens.
tokens = _ct(z, 0, 2, 3, 4, 1).reshape(-1, channels)
tokens = _linear(tokens, w["decoder.proj_in.weight"], w["decoder.proj_in.bias"])
num_patches = tokens.shape[0]
num_reg = self.config.decoder_num_register_tokens
register = mx.broadcast_to(w["decoder.register_tokens"], (batch, num_reg, self.dim)).reshape(-1, self.dim)
cls_token = mx.zeros((batch, self.dim), dtype=tokens.dtype)
hidden = mx.concatenate([tokens, register, cls_token], axis=0)[:, None, :] # (S, B, D)
grids = [2.0 * (mx.arange(0.5, size, dtype=mx.float32) / size) - 1.0 for size in (num_frames, height, width)]
mesh_t, mesh_h, mesh_w = mx.meshgrid(grids[0], grids[1], grids[2], indexing="ij")
position_ids = mx.stack([mesh_t.reshape(-1), mesh_h.reshape(-1), mesh_w.reshape(-1)], axis=-1)
suffix = mx.zeros((batch * (cfg.decoder_num_register_tokens + 1), 3), dtype=position_ids.dtype)
position_ids = mx.concatenate([position_ids, suffix], axis=0) # (S, 3)
cos, sin = _rotary_cos_sin(position_ids, rotary_dim=self.rotary_dim, theta=cfg.decoder_rope_theta)
for block in self._decoder_blocks():
hidden = _transformer_block(
block,
hidden,
block["scale1"],
block["scale2"],
cos,
sin,
num_heads=self.num_heads,
head_dim=self.head_dim,
norm_eps=cfg.decoder_norm_eps,
qk_norm_eps=cfg.decoder_norm_eps,
)
mx.eval(hidden) # per-block sync (36 blocks build a huge lazy graph)
hidden = _layer_norm(hidden, w["decoder.norm_out.weight"], w["decoder.norm_out.bias"], eps=cfg.decoder_norm_eps)
output = _linear(hidden, w["decoder.proj_out.weight"], w["decoder.proj_out.bias"])[:num_patches]
patch_t = cfg.temporal_compression_ratio
patch_h = patch_w = cfg.spatial_compression_ratio
output = output.reshape(num_patches, cfg.out_channels, patch_t, patch_h, patch_w)
output = output.reshape(batch, num_frames, height, width, cfg.out_channels, patch_t, patch_h, patch_w)
output = _ct(output, 0, 4, 1, 5, 2, 6, 3, 7)
return output.reshape(batch, cfg.out_channels, num_frames * patch_t, height * patch_h, width * patch_w)
def _split_tiles(self, length: int, tile_size: int, min_overlap: int):
if tile_size >= length:
return [0], [length], []
num_tiles = math.ceil(length / tile_size)
while tile_size * num_tiles - min_overlap * (num_tiles - 1) - length < 0:
num_tiles += 1
overlaps = [min_overlap] * (num_tiles - 1)
remaining = tile_size * num_tiles - sum(overlaps) - length
for index in range(remaining // self.spatial_compression_ratio):
overlaps[index % (num_tiles - 1)] += self.spatial_compression_ratio
starts = [0]
for index in range(num_tiles - 1):
starts.append(starts[-1] + tile_size - overlaps[index])
return starts, [tile_size] * num_tiles, overlaps
@staticmethod
def _blend(a, b, blend_extent: int, dim: int):
blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
positions = mx.arange(blend_extent, dtype=b.dtype)
shape = [1] * a.ndim
shape[dim] = blend_extent
weight_a = (1 - positions / blend_extent).reshape(shape)
weight_b = (positions / blend_extent).reshape(shape)
slice_a = [slice(None)] * a.ndim
slice_a[dim] = slice(-blend_extent, None)
slice_b = [slice(None)] * b.ndim
slice_b[dim] = slice(0, blend_extent)
blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
if blend_extent == b.shape[dim]:
return blended
rest = [slice(None)] * b.ndim
rest[dim] = slice(blend_extent, None)
return mx.concatenate([blended, b[tuple(rest)]], axis=dim)
def _stitch_tiles(self, tiles, height_overlaps, width_overlaps):
result_rows = []
for row_index, row in enumerate(tiles):
result_row = []
for column_index, tile in enumerate(row):
if row_index > 0:
tile = self._blend(tiles[row_index - 1][column_index], tile, height_overlaps[row_index - 1], -2)
if column_index > 0:
tile = self._blend(row[column_index - 1], tile, width_overlaps[column_index - 1], -1)
if row_index < len(tiles) - 1:
tile = tile[..., :-height_overlaps[row_index], :]
if column_index < len(row) - 1:
tile = tile[..., :, :-width_overlaps[column_index]]
result_row.append(tile)
result_rows.append(mx.concatenate(result_row, axis=-1))
return mx.concatenate(result_rows, axis=-2)
def decode_clip_tiled(self,
z,
tile_sample_min_height: int,
tile_sample_min_width: int,
tile_sample_min_overlap_height: int = 64,
tile_sample_min_overlap_width: int = 64):
"""One clip through the decoder with spatial tiling (memory bounded)."""
height = z.shape[-2] * self.spatial_compression_ratio
width = z.shape[-1] * self.spatial_compression_ratio
y_starts, y_lengths, y_overlaps = self._split_tiles(height, tile_sample_min_height,
tile_sample_min_overlap_height)
x_starts, x_lengths, x_overlaps = self._split_tiles(width, tile_sample_min_width, tile_sample_min_overlap_width)
ratio = self.spatial_compression_ratio
rows = []
for y_start, y_length in zip(y_starts, y_lengths, strict=False):
row = []
for x_start, x_length in zip(x_starts, x_lengths, strict=False):
tile = z[..., y_start // ratio:y_start // ratio + y_length // ratio,
x_start // ratio:x_start // ratio + x_length // ratio]
row.append(self._decode_clip(tile))
rows.append(row)
return self._stitch_tiles(rows, y_overlaps, x_overlaps)
def decode(self,
z,
*,
tiled: bool = True,
tile_sample_min_height: int = 256,
tile_sample_min_width: int = 256,
tile_sample_min_overlap_height: int = 64,
tile_sample_min_overlap_width: int = 64):
"""Chunked decode of normalized latents (1, C, T_lat, H', W') -> (1, 3, T, H, W)."""
if not self.has_decoder:
raise RuntimeError("This MLX H3 video VAE was loaded without decoder weights.")
cfg = self.config
tokens_chunk_size = cfg.tokens_chunk_size
token_drop = cfg.token_drop
temporal_ratio = self.temporal_compression_ratio
chunk_num_frames = tokens_chunk_size * temporal_ratio
num_tokens = z.shape[2] + token_drop
pad_tokens = (-num_tokens) % tokens_chunk_size
num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
if pad_tokens > 0:
tail = mx.repeat(z[:, :, -1:], pad_tokens, axis=2)
z = mx.concatenate([z, tail], axis=2)
def decode_one(chunk):
if tiled:
return self.decode_clip_tiled(chunk, tile_sample_min_height, tile_sample_min_width,
tile_sample_min_overlap_height, tile_sample_min_overlap_width)
return self._decode_clip(chunk)
decoded_chunks = []
overlap = None
for index in range(num_chunks):
start = index * tokens_chunk_size
clip = decode_one(z[:, :, start:start + tokens_chunk_size + cfg.token_overlap])
for overlap_index in range(int(token_drop > 0) + 1):
frame_start = overlap_index * chunk_num_frames
chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
chunk = chunk[:, :, cfg.frame_pre_padding:]
if overlap_index == 0:
if overlap is not None:
chunk = self._blend(overlap, chunk, cfg.frame_overlap, dim=-3)
mx.eval(chunk) # materialize per clip; keeps the lazy graph bounded
decoded_chunks.append(chunk)
else:
overlap = chunk
mx.eval(overlap)
if overlap is not None:
decoded_chunks.append(overlap)
decoded = mx.concatenate(decoded_chunks, axis=2)
if pad_tokens > 0:
intra_tail = cfg.clip_length % temporal_ratio
num_tokens_before_pad = z.shape[2] - pad_tokens
pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
decoded = decoded[:, :, :-pad_frames]
return decoded</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.normalize_latents
normalize_latents(latents)
Source
def normalize_latents(self, latents):
return (latents - mx.array(self._latents_mean)) / mx.array(self._latents_std)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.denormalize_latents
denormalize_latents(latents)
Source
def denormalize_latents(self, latents):
return latents * mx.array(self._latents_std) + mx.array(self._latents_mean)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.normalize_pixels
normalize_pixels(pixels)
Source
def normalize_pixels(self, pixels):
return (pixels - mx.array(self._pixel_mean)) / mx.array(self._pixel_std)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.denormalize_pixels
denormalize_pixels(sample)
Source
def denormalize_pixels(self, sample):
return sample * mx.array(self._pixel_std) + mx.array(self._pixel_mean)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.encode
encode(pixels)
Normalized pixels (1, 3, T, H, W) -> (mean, logvar) each (1, C, T', H', W').
Pads the clip to clip_length and drops token_drop trailing
moment frames exactly like the reference _encode.
Source
def encode(self, pixels):
"""Normalized pixels (1, 3, T, H, W) -> (mean, logvar) each (1, C, T', H', W').
Pads the clip to ``clip_length`` and drops ``token_drop`` trailing
moment frames exactly like the reference ``_encode``.
"""
if not self.has_encoder:
raise RuntimeError("This MLX H3 video VAE was loaded without encoder weights.")
cfg = self.config
num_frames = pixels.shape[2]
if num_frames % cfg.clip_length != 0:
pad_frames = (-num_frames) % cfg.clip_length
tail = mx.repeat(pixels[:, :, -1:], pad_frames, axis=2)
pixels = mx.concatenate([pixels, tail], axis=2)
moment_chunks: list[Any] = []
for index in range(pixels.shape[2] // cfg.clip_length):
clip = pixels[:, :, index * cfg.clip_length:(index + 1) * cfg.clip_length]
moment_chunks.append(self._encode_clip(clip))
moments = mx.concatenate(moment_chunks, axis=2)
if cfg.token_drop > 0:
moments = moments[:, :, :-cfg.token_drop]
mean, logvar = mx.split(moments, 2, axis=1)
logvar = mx.clip(logvar, -30.0, 20.0)
return mean, logvar</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.encode_keyframe
encode_keyframe(pixels)
Single-frame conditioning encode without chunk padding.
Source
def encode_keyframe(self, pixels):
"""Single-frame conditioning encode without chunk padding."""
if pixels.shape[2] != 1:
raise ValueError(f"encode_keyframe expects exactly one frame, got {pixels.shape}.")
moments = self._encode_clip(pixels)
mean, logvar = mx.split(moments, 2, axis=1)
logvar = mx.clip(logvar, -30.0, 20.0)
return mean, logvar
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.sample_posterior
sample_posterior(mean, logvar, noise)
Reparameterization with an explicit noise array (deterministic parity).
Source
@staticmethod
def sample_posterior(mean, logvar, noise):
"""Reparameterization with an explicit noise array (deterministic parity)."""
return mean + mx.exp(0.5 * logvar) * noise
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.decode_clip_tiled
decode_clip_tiled(z, tile_sample_min_height: int, tile_sample_min_width: int, tile_sample_min_overlap_height: int = 64, tile_sample_min_overlap_width: int = 64)
One clip through the decoder with spatial tiling (memory bounded).
Source
def decode_clip_tiled(self,
z,
tile_sample_min_height: int,
tile_sample_min_width: int,
tile_sample_min_overlap_height: int = 64,
tile_sample_min_overlap_width: int = 64):
"""One clip through the decoder with spatial tiling (memory bounded)."""
height = z.shape[-2] * self.spatial_compression_ratio
width = z.shape[-1] * self.spatial_compression_ratio
y_starts, y_lengths, y_overlaps = self._split_tiles(height, tile_sample_min_height,
tile_sample_min_overlap_height)
x_starts, x_lengths, x_overlaps = self._split_tiles(width, tile_sample_min_width, tile_sample_min_overlap_width)
ratio = self.spatial_compression_ratio
rows = []
for y_start, y_length in zip(y_starts, y_lengths, strict=False):
row = []
for x_start, x_length in zip(x_starts, x_lengths, strict=False):
tile = z[..., y_start // ratio:y_start // ratio + y_length // ratio,
x_start // ratio:x_start // ratio + x_length // ratio]
row.append(self._decode_clip(tile))
rows.append(row)
return self._stitch_tiles(rows, y_overlaps, x_overlaps)
fastvideo.mlx_runtime.minimax_h3_video_vae.MLXMiniMaxH3VideoVAE.decode
decode(z, *, tiled: bool = True, tile_sample_min_height: int = 256, tile_sample_min_width: int = 256, tile_sample_min_overlap_height: int = 64, tile_sample_min_overlap_width: int = 64)
Chunked decode of normalized latents (1, C, T_lat, H', W') -> (1, 3, T, H, W).
Source
def decode(self,
z,
*,
tiled: bool = True,
tile_sample_min_height: int = 256,
tile_sample_min_width: int = 256,
tile_sample_min_overlap_height: int = 64,
tile_sample_min_overlap_width: int = 64):
"""Chunked decode of normalized latents (1, C, T_lat, H', W') -> (1, 3, T, H, W)."""
if not self.has_decoder:
raise RuntimeError("This MLX H3 video VAE was loaded without decoder weights.")
cfg = self.config
tokens_chunk_size = cfg.tokens_chunk_size
token_drop = cfg.token_drop
temporal_ratio = self.temporal_compression_ratio
chunk_num_frames = tokens_chunk_size * temporal_ratio
num_tokens = z.shape[2] + token_drop
pad_tokens = (-num_tokens) % tokens_chunk_size
num_chunks = (num_tokens + pad_tokens) // tokens_chunk_size - int(token_drop > 0)
if pad_tokens > 0:
tail = mx.repeat(z[:, :, -1:], pad_tokens, axis=2)
z = mx.concatenate([z, tail], axis=2)
def decode_one(chunk):
if tiled:
return self.decode_clip_tiled(chunk, tile_sample_min_height, tile_sample_min_width,
tile_sample_min_overlap_height, tile_sample_min_overlap_width)
return self._decode_clip(chunk)
decoded_chunks = []
overlap = None
for index in range(num_chunks):
start = index * tokens_chunk_size
clip = decode_one(z[:, :, start:start + tokens_chunk_size + cfg.token_overlap])
for overlap_index in range(int(token_drop > 0) + 1):
frame_start = overlap_index * chunk_num_frames
chunk = clip[:, :, frame_start:frame_start + chunk_num_frames]
chunk = chunk[:, :, cfg.frame_pre_padding:]
if overlap_index == 0:
if overlap is not None:
chunk = self._blend(overlap, chunk, cfg.frame_overlap, dim=-3)
mx.eval(chunk) # materialize per clip; keeps the lazy graph bounded
decoded_chunks.append(chunk)
else:
overlap = chunk
mx.eval(overlap)
if overlap is not None:
decoded_chunks.append(overlap)
decoded = mx.concatenate(decoded_chunks, axis=2)
if pad_tokens > 0:
intra_tail = cfg.clip_length % temporal_ratio
num_tokens_before_pad = z.shape[2] - pad_tokens
pad_frames = sum(intra_tail if intra_tail and (num_tokens_before_pad + offset) %
tokens_chunk_size == 0 else temporal_ratio for offset in range(pad_tokens))
decoded = decoded[:, :, :-pad_frames]
return decoded</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_video_vae.mlx_h3_video_vae_from_files
mlx_h3_video_vae_from_files(*, vae_dir: str | Path | None = None, decoder_path: str | Path | None = None, encoder_path: str | Path | None = None, include_encoder: bool = True, include_decoder: bool = True, storage_dtype: str = 'fp32', config: MiniMaxH3VideoVAEConfigView | None = None, decoder_quantization: str | None = None) -> MLXMiniMaxH3VideoVAE
Load an H3 video VAE from any mix of the released VAE directory and standalone component files.
decoder_path / encoder_path override the matching half of vae_dir (or stand alone when no
vae_dir is given, so the base VAE weights are never read). Standalone files may be a LightVAE
block-pruned component (config embedded in its metadata; decoder_num_layers follows the kept
blocks) or a FastH3 NVFP4 decoder (decoder.safetensors bf16 masters, or decoder-nvfp4.safetensors
pre-packed). A FastH3 decoder carries no latent statistics: pass a LightVAE encoder, a vae_dir
holding config.json, or config.
decoder_quantization ("nvfp4", "mxfp8", "int8", "int4") stores the 216 decoder block linears
quantized as they are read, so the dense decoder never needs to be resident.
Source
def mlx_h3_video_vae_from_files(*,
vae_dir: str | Path | None = None,
decoder_path: str | Path | None = None,
encoder_path: str | Path | None = None,
include_encoder: bool = True,
include_decoder: bool = True,
storage_dtype: str = "fp32",
config: MiniMaxH3VideoVAEConfigView | None = None,
decoder_quantization: str | None = None) -> MLXMiniMaxH3VideoVAE:
"""Load an H3 video VAE from any mix of the released VAE directory and standalone component files.
``decoder_path`` / ``encoder_path`` override the matching half of ``vae_dir`` (or stand alone when no
``vae_dir`` is given, so the base VAE weights are never read). Standalone files may be a LightVAE
block-pruned component (config embedded in its metadata; ``decoder_num_layers`` follows the kept
blocks) or a FastH3 NVFP4 decoder (``decoder.safetensors`` bf16 masters, or ``decoder-nvfp4.safetensors``
pre-packed). A FastH3 decoder carries no latent statistics: pass a LightVAE encoder, a ``vae_dir``
holding ``config.json``, or ``config``.
``decoder_quantization`` ("nvfp4", "mxfp8", "int8", "int4") stores the 216 decoder block linears
quantized as they are read, so the dense decoder never needs to be resident.
"""
cast_dtype = _DTYPE_MAP[storage_dtype]
def component_paths(explicit: str | Path | None, wanted: bool, what: str) -> list[Path]:
if not wanted:
return []
if explicit:
return [Path(explicit)]
if vae_dir is None:
raise ValueError(f"A {what} needs vae_dir or an explicit {what}_path.")
return _shards(Path(vae_dir))
decoder_paths = component_paths(decoder_path, include_decoder, "decoder")
encoder_paths = component_paths(encoder_path, include_encoder, "encoder")
if config is None and vae_dir is not None and (Path(vae_dir) / "config.json").exists():
config = MiniMaxH3VideoVAEConfigView.from_vae_dir(vae_dir)
if config is None:
config = _embedded_config([Path(p) for p in (encoder_path, decoder_path) if p] + encoder_paths + decoder_paths)
if config is None:
raise ValueError("No VAE config: pass vae_dir with config.json, a LightVAE component file, or config=.")
weights: dict[str, Any] = {}
if include_encoder:
weights.update(_read_weights(encoder_paths, _ENCODER_PREFIXES, cast_dtype, None))
_check_required(
weights,
["quant_conv.weight", "encoder.conv_in.weight", "encoder.conv_out.weight", "encoder.norm_out.weight"],
"H3 video VAE encoder")
if include_decoder:
weights.update(_read_weights(decoder_paths, _DECODER_PREFIXES, cast_dtype, decoder_quantization))
_check_required(weights, [
"post_quant_conv.weight", "post_quant_conv.bias", "decoder.proj_in.weight", "decoder.norm_out.weight",
"decoder.proj_out.weight", "decoder.register_tokens"
], "H3 video VAE decoder")
found = {int(key.split(".")[2]) for key in weights if key.startswith("decoder.transformer_blocks.")}
if found != set(range(len(found))) or not found:
raise KeyError(f"H3 video VAE decoder blocks are not contiguous from 0: {sorted(found)}.")
if vae_dir is not None and decoder_path is None and len(found) != config.decoder_num_layers:
raise KeyError(f"H3 video VAE decoder blocks incomplete: found {sorted(found)} "
f"of {config.decoder_num_layers}.")
if len(found) != config.decoder_num_layers: # pruned / depth-cut decoder
config = replace(config, decoder_num_layers=len(found))
return MLXMiniMaxH3VideoVAE(weights, config, has_encoder=include_encoder, has_decoder=include_decoder)</code></pre></details>
fastvideo.mlx_runtime.minimax_h3_video_vae.mlx_h3_video_vae_from_dir
mlx_h3_video_vae_from_dir(vae_dir: str | Path, *, include_encoder: bool = True, include_decoder: bool = True, storage_dtype: str = 'fp32', config: MiniMaxH3VideoVAEConfigView | None = None) -> MLXMiniMaxH3VideoVAE
Load the released H3 video VAE, streaming one shard at a time.
storage_dtype="fp32" keeps the released numerics. bf16/fp16 halve
residency; measure drift against the FP32 acceptance gate before shipping
a reduced-dtype configuration.
Source
def mlx_h3_video_vae_from_dir(vae_dir: str | Path,
*,
include_encoder: bool = True,
include_decoder: bool = True,
storage_dtype: str = "fp32",
config: MiniMaxH3VideoVAEConfigView | None = None) -> MLXMiniMaxH3VideoVAE:
"""Load the released H3 video VAE, streaming one shard at a time.
``storage_dtype="fp32"`` keeps the released numerics. bf16/fp16 halve
residency; measure drift against the FP32 acceptance gate before shipping
a reduced-dtype configuration.
"""
return mlx_h3_video_vae_from_files(vae_dir=vae_dir,
include_encoder=include_encoder,
include_decoder=include_decoder,
storage_dtype=storage_dtype,
config=config)</code></pre></details>