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

Fix causal conv: correct F.pad spatial dim order + 1-frame kernel truncation

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 21:31:05 +07:00
parent 3386975326
commit 6fc7cf0ad6
1 changed files with 16 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -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))