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

Thread single_frame flag so keyframe truncation does not apply to 1-frame tiles

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 21:11:43 +07:00
parent 06fd79a8fd
commit 65e80be1cf
1 changed files with 26 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -41,23 +41,22 @@ NIN_LEVELS = frozenset({1, 3, 5})
DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3})
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding):
def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding, single_frame=False):
"""Reflect-spatial / causal-temporal 3D conv (matches upstream_vae.CausalConv3d).
Front-only zero temporal padding of ``(kernel_size - 1)`` (two taps on each
side) whenever ``spatial_padding > 0`` (the reference pads with ``(k//2,
k//2, 0)`` in time), no temporal padding otherwise. A single input frame
has its temporal taps truncated instead of convolving zero frames
(``autopad="causal_zero"``) so a keyframe never leaks zero into the latent.
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.
"""
t = x.shape[2]
if t == 1 and spatial_padding > 0:
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))
return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))
@ -68,28 +67,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):
def _resnet(x, p, single_frame):
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)
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)
return h.add_(residual)
def _downsample(x, p):
def _downsample(x, p, single_frame):
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)
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)
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)
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)
for level in E["down"]:
for blk in level["blocks"]:
h = _resnet(h, blk)
h = _resnet(h, blk, single_frame)
if level["down"] is not None:
h = _downsample(h, level["down"])
h = _downsample(h, level["down"], single_frame)
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)
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)
class MiniMaxH3VideoVAEEncoder(nn.Module):
@ -174,13 +173,13 @@ class MiniMaxH3VideoVAEEncoder(nn.Module):
return model
@torch.inference_mode()
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 _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 _adaptive_encode(self, x: torch.Tensor) -> torch.Tensor:
def _adaptive_encode(self, x: torch.Tensor, single_frame: bool = False) -> torch.Tensor:
if self.tiling:
return self.tiled_encode(x)
return self._encode_moments(x)
return self.tiled_encode(x, single_frame)
return self._encode_moments(x, single_frame)
def split_tiles(self, length: int) -> tuple[list[int], list[int], list[int]]:
if self.tile_size >= length:
@ -217,11 +216,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) -> torch.Tensor:
def tiled_encode(self, x: torch.Tensor, single_frame: bool = False) -> 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)]
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)]
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 = []
@ -258,7 +257,7 @@ 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)[:, :, -1:, :, :]
moments = self._adaptive_encode(x, single_frame=True)[:, :, -1:, :, :]
else:
moments = self.encode_temporal(x)
mean = torch.chunk(moments.float(), 2, dim=1)[0]