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

Decouple spatial_padding from temporal_causal in causal conv

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

54
src/h3_blackwell_runtime/vae_encoder.py
View file

@ -41,39 +41,29 @@ NIN_LEVELS = frozenset({1, 3, 5})
DOWNSAMPLE_LEVELS = frozenset({0, 1, 2, 3}) 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, temporal_causal):
"""Reflect-spatial / causal-temporal 3D conv (matches upstream_vae.CausalConv3d). """Causal 3D conv (matches upstream_vae.CausalConv3d).
Front-only zero temporal padding of ``(kernel_size - 1)`` whenever - ``spatial_padding > 0``: reflect H and W by ``spatial_padding`` on each side.
``spatial_padding > 0`` (the reference pads with ``(k//2, k//2, 0)`` in - ``temporal_causal``: front-zero T by ``kernel_size - 1`` (the reference's
time), no temporal padding otherwise. Single-frame keyframes go through ``causal_padding[0] * 2`` for ``causal_padding[0]=1``, which is every
the same code path -- the causal front-zero is the reference's 3D-causal conv in the H3 VAE). A single-frame input truncates the
``autopad="causal_zero"`` behaviour for the encoder. temporal taps to the center 1 instead of convolving zero rows.
- If neither applies: no padding at all (reference early-return).
""" """
# 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 x.shape[2] == 1: if x.shape[2] == 1:
# Keyframe path (matches reference's `autopad="causal_zero"`): apply # Keyframe path (matches reference's `autopad="causal_zero"`):
# spatial-reflect pad, then truncate the temporal taps to the center 1 # apply spatial-reflect pad (if any), then run an effective 2D conv
# so the conv acts as a 2D op on the single frame. # by slicing the kernel to its center temporal tap (T_out stays 1).
if spatial_padding > 0: 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") 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]), (0, 0, 0)) 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) x = F.pad(x, (spatial_padding, spatial_padding, spatial_padding, spatial_padding, 0, 0), mode="reflect")
return F.conv3d(x, weight, bias, stride, (0, 0, 0)) if temporal_causal:
x = F.pad(x, (0, 0, 0, 0, kernel_size - 1, 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, 0, 0)) return F.conv3d(x, weight, bias, stride, (0, 0, 0))
@ -85,29 +75,31 @@ def _group_norm_3d(x, weight, bias):
def _resnet(x, p): def _resnet(x, p):
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) # nin_shortcut uses CausalConv3d(k=1, padding=1) in the reference.
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) 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=1, temporal_causal=True)
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, temporal_causal=True)
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, temporal_causal=True)
return h.add_(residual) return h.add_(residual)
def _downsample(x, p): def _downsample(x, p):
if p["space"] == 2: if p["space"] == 2:
# Reference Downsample3D pads H and W by +1 reflect before the conv. # Reference Downsample3D pads H and W by +1 reflect before the conv.
# F.pad for 5D (B,C,T,H,W) uses (W_l, W_r, H_l, H_r, T_l, T_r).
x = F.pad(x, (1, 1, 1, 1, 0, 0), mode="reflect") x = F.pad(x, (1, 1, 1, 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) # Conv uses padding=(1,0,0) -> causal_padding=(1,0,0), so spatial pad=0,
# temporal front-zero is applied.
return _causal_conv3d(x, p["w"], p["b"], kernel_size=3, stride=(p["time"], p["space"], p["space"]), spatial_padding=0, temporal_causal=True)
def _encoder_run(x, E): 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) h = _causal_conv3d(x, E["conv_in"][0], E["conv_in"][1], kernel_size=3, stride=(1, 1, 1), spatial_padding=1, temporal_causal=True)
for level in E["down"]: for level in E["down"]:
for blk in level["blocks"]: for blk in level["blocks"]:
h = _resnet(h, blk) h = _resnet(h, blk)
if level["down"] is not None: if level["down"] is not None:
h = _downsample(h, level["down"]) h = _downsample(h, level["down"])
h = F.silu(_group_norm_3d(h, E["norm_out_w"], E["norm_out_b"])) 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, temporal_causal=True)
class MiniMaxH3VideoVAEEncoder(nn.Module): class MiniMaxH3VideoVAEEncoder(nn.Module):