diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 82210fa..6715354 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -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]