diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 9aa6f9c..d33e1d3 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -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]