From 6fc7cf0ad6c57fbe01130b7e6d944ca4ac25c6cc Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 19 Aug 2026 21:31:05 +0700 Subject: [PATCH] Fix causal conv: correct F.pad spatial dim order + 1-frame kernel truncation --- src/h3_blackwell_runtime/vae_encoder.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 6715354..be2096e 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -50,8 +50,23 @@ def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding): the same code path -- the causal front-zero is the reference's ``autopad="causal_zero"`` behaviour for the encoder. """ + # Reference semantics: + # x = F.pad(x, (s_p, s_p, s_p, s_p, 0, 0), mode="reflect") -- spatial + # if x.shape[2] == 1: return super().forward(x, autopad="causal_zero") + # x = F.pad(x, (0, 0, 0, 0, (k-1), 0), mode="constant") -- temporal + # return super().forward(x) if spatial_padding > 0: - x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0)) + x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect") + x = x.permute(0, 2, 1, 3, 4) + if x.shape[2] == 1: + # Keyframe path (matches reference's `autopad="causal_zero"`): truncate + # the temporal taps to the center 1 and run as an effective 2D conv. + # We emulate this with F.conv3d by slicing the kernel's T dim to 1 and + # running conv3d with stride_t=1 and no T padding -- output T stays 1. + half = (kernel_size - 1) // 2 + kernel_5d = weight[:, :, half:half + 1, :, :] # (O, I, 1, k, k) + return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (spatial_padding, spatial_padding, spatial_padding)) + x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 0)) return F.conv3d(x, weight, bias, stride, (0, spatial_padding, spatial_padding))