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

Fix F.pad dim order (spatial reflect + no T pad) for 1-frame path

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 21:36:40 +07:00
parent c2d507691b
commit 136297f4ae
1 changed files with 7 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -56,11 +56,15 @@ def _causal_conv3d(x, weight, bias, *, kernel_size, stride, spatial_padding):
# x = F.pad(x, (0, 0, 0, 0, (k-1), 0), mode="constant") -- temporal # x = F.pad(x, (0, 0, 0, 0, (k-1), 0), mode="constant") -- temporal
# return super().forward(x) # return super().forward(x)
if x.shape[2] == 1: if x.shape[2] == 1:
# Keyframe path (matches reference's `autopad="causal_zero"`): truncate # Keyframe path (matches reference's `autopad="causal_zero"`): apply
# the temporal taps to the center 1 and run as an effective 2D conv. # spatial-reflect pad, then truncate the temporal taps to the center 1
# so the conv acts as a 2D op on the single frame.
if spatial_padding > 0:
# F.pad for 5D input (B,C,T,H,W) uses (W_l, W_r, H_l, H_r, T_l, T_r).
x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
half = (kernel_size - 1) // 2 half = (kernel_size - 1) // 2
kernel_5d = weight[:, :, half:half + 1, :, :] kernel_5d = weight[:, :, half:half + 1, :, :]
return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (spatial_padding, spatial_padding, spatial_padding)) return F.conv3d(x, kernel_5d, bias, (1, stride[1], stride[2]), (0, 0, 0))
if spatial_padding > 0: if spatial_padding > 0:
return _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding) return _reflect_and_causal(x, weight, bias, kernel_size, stride, spatial_padding)
return F.conv3d(x, weight, bias, stride, (0, 0, 0)) return F.conv3d(x, weight, bias, stride, (0, 0, 0))