daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix VAE: use same multi-frame causal path for keyframes, add frame_pre_padding so 1 keyframe -> 1 latent

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 21:20:32 +07:00
parent a0274a9868
commit 3386975326
1 changed files with 28 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -41,20 +41,15 @@ NIN_LEVELS = frozenset({1, 3, 5})
DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3})
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, single_frame=False):
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding):
"""Reflect-spatial / causal-temporal 3D conv (matches upstream_vae.CausalConv3d).
Front-only zero temporal padding of ``(kernel_size - 1)`` whenever
``spatial_padding > 0``, no temporal padding otherwise. When
``single_frame`` is set (the *top-level* input had exactly one frame, i.e.
a keyframe), the temporal taps are truncated instead of convolving zero
frames (``autopad="causal_zero"``) so a keyframe never leaks zero rows into
the latent.
``spatial_padding > 0`` (the reference pads with ``(k//2, k//2, 0)`` in
time), no temporal padding otherwise. Single-frame keyframes go through
the same code path -- the causal front-zero is the reference's
``autopad="causal_zero"`` behaviour for the encoder.
"""
if single_frame and 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))
if spatial_padding > 0:
x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0))
return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))
@ -67,28 +62,28 @@ def _group_norm_3d(x, weight, bias):
return y.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()
def _resnet(x, p, single_frame):
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, single_frame=single_frame)
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, single_frame=single_frame)
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)
def _downsample(x, p, single_frame):
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, single_frame=single_frame)
return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0)
def _encoder_run(x, E, single_frame):
h = _causal_conv3d(x, E["conv_in"][0], E["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, single_frame=single_frame)
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, single_frame)
h = _resnet(h, blk)
if level["down"] is not None:
h = _downsample(h, level["down"], single_frame)
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, single_frame=single_frame)
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):
@ -102,6 +97,7 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
super().__init__()
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)
@ -173,13 +169,13 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
return model
@torch.inference_mode()
def _encode_moments(self, x: torch.Tensor, single_frame: bool = False) -> torch.Tensor:
return F.conv3d(_encoder_run(x.to(torch.float32), self.W, single_frame), self.quant_conv.weight, self.quant_conv.bias)
def _encode_moments(self, x: torch.Tensor) -> torch.Tensor:
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, single_frame: bool = False) -> torch.Tensor:
def _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor:
if self.tiling:
return self.tiled_encode(x, single_frame)
return self._encode_moments(x, single_frame)
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:
@ -216,11 +212,11 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
return torch.cat((blended, b[tuple(slice_b_rest)]), dim=dim)
return blended
def tiled_encode(self, x: torch.Tensor, single_frame: bool = False) -> torch.Tensor:
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], single_frame) for j_pos, j_len in zip(x_idx, x_len)] for i_pos, i_len in zip(y_idx, y_len)]
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 = []
@ -257,7 +253,11 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
x = (x + 1.0) * 0.5
x = (x - self.pixel_mean.to(x)) / self.pixel_std.to(x)
if x.shape[2] == 1:
moments = self._adaptive_encode(x, single_frame=True)[:, :, -1:, :, :]
# Pad with frame_pre_padding zero frames so the 17-tap stride-2
# temporal downsampling produces exactly one latent frame (1+3-4).
pad = self.frame_pre_padding
x = torch.cat([torch.zeros_like(x[:, :, :pad]), x], dim=2)
moments = self._adaptive_encode(x)
else:
moments = self.encode_temporal(x)
mean = torch.chunk(moments.float(), 2, dim=1)[0]