From c2d507691bdaad5612b8b7d801b8ec96432c66fa Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 19 Aug 2026 21:35:07 +0700 Subject: [PATCH] Restore front-zero temporal pad for multi-frame causal conv; two-step F.pad (reflect spatial + constant T) --- src/h3_blackwell_runtime/vae_encoder.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/h3_blackwell_runtime/vae_encoder.py b/src/h3_blackwell_runtime/vae_encoder.py index 9ea30b1..e6dcd5d 100644 --- a/src/h3_blackwell_runtime/vae_encoder.py +++ b/src/h3_blackwell_runtime/vae_encoder.py @@ -58,13 +58,19 @@ def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding): 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) + kernel_5d = weight[:, :, half:half + 1, :, :] return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (spatial_padding, spatial_padding, spatial_padding)) + if spatial_padding > 0: + return _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding) + return F.conv3d(x, weight, bias, stride, (0, 0, 0)) + + +def _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding): + """Two-step pad (reference CausalConv3d): spatial reflect, temporal front-zero.""" + x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect") 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, 0, 0)) def _group_norm_3d(x, weight, bias):