refactor(modeling): self-consistent block/function names; drop unused attrs

modeling_motifvae.py

@@ -33,28 +33,28 @@ from diffusers.utils import BaseOutput
 # ---------------------------------------------------------------------------
 #
 # The config encodes residual / mid block choices as strings (e.g.
-# ``"ResnetBlock2D"``). ``resolve_str_to_obj`` maps those strings to the
+# ``"ResnetBlock2D"``). ``block_from_name`` maps those strings to the
 # classes defined in this module. ``"Identity"`` maps to ``nn.Identity`` so a
 # mid stage can be made attention-free while keeping the surrounding key
 # layout intact.
 
-_STR_TO_OBJ_OVERRIDES = {"Identity": nn.Identity}
+_NAME_OVERRIDES = {"Identity": nn.Identity}
 
 
-def resolve_str_to_obj(name: str):
+def block_from_name(name: str):
     """Resolve a block-type string from the config to its class."""
-    if name in _STR_TO_OBJ_OVERRIDES:
-        return _STR_TO_OBJ_OVERRIDES[name]
+    if name in _NAME_OVERRIDES:
+        return _NAME_OVERRIDES[name]
     try:
         return globals()[name]
     except KeyError as exc:
         raise AttributeError(
-            f"resolve_str_to_obj: unknown class name {name!r} "
+            f"block_from_name: unknown class name {name!r} "
             f"(not defined in this module)"
         ) from exc
 
 
-class VideoBaseAE(ModelMixin, ConfigMixin):
+class MotifVAEBase(ModelMixin, ConfigMixin):
     """Base class wiring diffusers' ``ModelMixin`` + ``ConfigMixin`` together."""
 
     config_name = "config.json"
@@ -68,7 +68,7 @@ class VideoBaseAE(ModelMixin, ConfigMixin):
 # ---------------------------------------------------------------------------
 
 
-def video_to_image(func):
+def per_frame_2d(func):
     """Wrap a 2D forward so it also accepts 5D ``(B, C, T, H, W)`` tensors.
 
     Frames are folded into the batch dimension, processed independently, then
@@ -88,12 +88,12 @@ def video_to_image(func):
     return wrapper
 
 
-def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+def silu(x: torch.Tensor) -> torch.Tensor:
     """SiLU / swish activation used throughout the network."""
     return x * torch.sigmoid(x)
 
 
-def cast_tuple(value, length: int = 1):
+def _as_tuple(value, length: int = 1):
     """Broadcast a scalar to a length-``length`` tuple; pass tuples through."""
     return value if isinstance(value, (tuple, list)) else ((value,) * length)
 
@@ -134,7 +134,7 @@ class Conv2d(nn.Conv2d):
             dtype,
         )
 
-    @video_to_image
+    @per_frame_2d
     def forward(self, x):
         return super().forward(x)
 
@@ -156,12 +156,12 @@ class CausalConv3d(nn.Module):
         **kwargs,
     ) -> None:
         super().__init__()
-        self.kernel_size = cast_tuple(kernel_size, 3)
+        self.kernel_size = _as_tuple(kernel_size, 3)
         self.time_kernel_size = self.kernel_size[0]
         self.chan_in = chan_in
         self.chan_out = chan_out
-        stride = cast_tuple(kwargs.pop("stride", 1), 3)
-        padding = list(cast_tuple(kwargs.pop("padding", 0), 3))  # (T, H, W)
+        stride = _as_tuple(kwargs.pop("stride", 1), 3)
+        padding = list(_as_tuple(kwargs.pop("padding", 0), 3))  # (T, H, W)
         self.stride = stride
         self.padding = padding
         self.conv = nn.Conv3d(
@@ -203,7 +203,7 @@ class LayerNorm(nn.Module):
         return x
 
 
-def Normalize(in_channels, num_groups=32, norm_type="groupnorm"):
+def make_norm(in_channels, num_groups=32, norm_type="groupnorm"):
     """Build the normalisation layer selected by ``norm_type``."""
     if norm_type == "groupnorm":
         return torch.nn.GroupNorm(
@@ -317,7 +317,7 @@ class Upsample(nn.Module):
             r = x.repeat_interleave(out_c // in_c, dim=1)
         return F.interpolate(r, scale_factor=2.0, mode="nearest")
 
-    @video_to_image
+    @per_frame_2d
     def forward(self, x):
         h = self.shuffle(self.conv(x))
         if self.residual:
@@ -337,13 +337,13 @@ class Downsample(nn.Module):
             in_channels, out_channels, kernel_size=3, stride=2, padding=0
         )
 
-    @video_to_image
+    @per_frame_2d
     def forward(self, x):
         x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)
         return self.conv(x)
 
 
-class Spatial2xTime2x3DDownsample(nn.Module):
+class STDownsample(nn.Module):
     """Joint 2x spatial + 2x temporal downsample via a stride-2 causal conv."""
 
     def __init__(self, in_channels, out_channels):
@@ -357,7 +357,7 @@ class Spatial2xTime2x3DDownsample(nn.Module):
         return self.conv(x)
 
 
-class Spatial2xTime2x3DUpsample(nn.Module):
+class STUpsample(nn.Module):
     """Joint spatial + temporal 2x upsample.
 
     The temporal axis is stretched with ``F.interpolate`` (cheap and smooth);
@@ -469,11 +469,11 @@ class ResnetBlock2D(nn.Module):
         self.out_channels = in_channels if out_channels is None else out_channels
         self.use_conv_shortcut = conv_shortcut
 
-        self.norm1 = Normalize(in_channels, norm_type=norm_type)
+        self.norm1 = make_norm(in_channels, norm_type=norm_type)
         self.conv1 = torch.nn.Conv2d(
             in_channels, out_channels, kernel_size=3, stride=1, padding=1
         )
-        self.norm2 = Normalize(out_channels, norm_type=norm_type)
+        self.norm2 = make_norm(out_channels, norm_type=norm_type)
         self.dropout = torch.nn.Dropout(dropout)
         self.conv2 = torch.nn.Conv2d(
             out_channels, out_channels, kernel_size=3, stride=1, padding=1
@@ -488,13 +488,13 @@ class ResnetBlock2D(nn.Module):
                     in_channels, out_channels, kernel_size=1, stride=1, padding=0
                 )
 
-    @video_to_image
+    @per_frame_2d
     def forward(self, x):
         h = self.norm1(x)
-        h = nonlinearity(h)
+        h = silu(h)
         h = self.conv1(h)
         h = self.norm2(h)
-        h = nonlinearity(h)
+        h = silu(h)
         h = self.dropout(h)
         h = self.conv2(h)
         if self.in_channels != self.out_channels:
@@ -522,9 +522,9 @@ class ResnetBlock3D(nn.Module):
         self.out_channels = in_channels if out_channels is None else out_channels
         self.use_conv_shortcut = conv_shortcut
 
-        self.norm1 = Normalize(in_channels, norm_type=norm_type)
+        self.norm1 = make_norm(in_channels, norm_type=norm_type)
         self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
-        self.norm2 = Normalize(out_channels, norm_type=norm_type)
+        self.norm2 = make_norm(out_channels, norm_type=norm_type)
         self.dropout = torch.nn.Dropout(dropout)
         self.conv2 = CausalConv3d(out_channels, out_channels, 3, padding=1)
         if self.in_channels != self.out_channels:
@@ -535,10 +535,10 @@ class ResnetBlock3D(nn.Module):
 
     def forward(self, x):
         h = self.norm1(x)
-        h = nonlinearity(h)
+        h = silu(h)
         h = self.conv1(h)
         h = self.norm2(h)
-        h = nonlinearity(h)
+        h = silu(h)
         h = self.dropout(h)
         h = self.conv2(h)
         if self.in_channels != self.out_channels:
@@ -559,7 +559,7 @@ class ResnetBlock3D(nn.Module):
 
 
 @dataclass
-class AutoencoderKLOutput(BaseOutput):
+class MotifEncoderOutput(BaseOutput):
     latent_dist: "DeterministicLatent"
     extra_output: Optional[tuple] = None
 
@@ -621,7 +621,7 @@ def build_mid_layer(
     """
     if layer_type is None or layer_type == "Identity":
         return nn.Identity()
-    return resolve_str_to_obj(layer_type)(
+    return block_from_name(layer_type)(
         in_channels=channels,
         out_channels=channels,
         dropout=dropout,
@@ -632,7 +632,7 @@ def build_mid_layer(
 def pad_time_to_even(x: torch.Tensor) -> torch.Tensor:
     """Prepend a repeated leading frame so the temporal axis becomes even.
 
-    This matches the temporal arithmetic of ``Spatial2xTime2x3DDownsample``:
+    This matches the temporal arithmetic of ``STDownsample``:
     folding the time axis by 2 needs an even length, and pre-pending (rather
     than appending) the repeated frame keeps the fold causal.
     """
@@ -686,7 +686,7 @@ class MotifDownBlock(nn.Module):
         )
 
         if down_type == "thw":
-            self.down = Spatial2xTime2x3DDownsample(
+            self.down = STDownsample(
                 in_channels=in_channels, out_channels=in_channels
             )
         elif down_type == "hw":
@@ -761,7 +761,7 @@ class MotifUpBlock(nn.Module):
     """Decoder up stage: residual stack, spatial/temporal upsample, residual.
 
     ``upsample_residual=True`` forwards the parameter-free identity skip into
-    the upsample module (see :class:`Upsample` / :class:`Spatial2xTime2x3DUpsample`).
+    the upsample module (see :class:`Upsample` / :class:`STUpsample`).
     """
 
     def __init__(
@@ -796,7 +796,7 @@ class MotifUpBlock(nn.Module):
         )
 
         if up_type == "thw":
-            self.up = Spatial2xTime2x3DUpsample(
+            self.up = STUpsample(
                 in_channels=in_channels,
                 out_channels=in_channels,
                 t_interpolation=t_interpolation,
@@ -829,7 +829,7 @@ class MotifUpBlock(nn.Module):
 # ---------------------------------------------------------------------------
 
 
-class MotifEncoder(VideoBaseAE):
+class MotifEncoder(MotifVAEBase):
     """Deterministic 3D causal encoder.
 
     For a ``(B, 3, T, H, W)`` clip the stem is a stride-1 conv; five down
@@ -884,7 +884,7 @@ class MotifEncoder(VideoBaseAE):
                     out_channels=base_channels[idx + 1],
                     num_res_blocks=num_resblocks,
                     down_type=down_type,
-                    res_block=resolve_str_to_obj(down_res_type),
+                    res_block=block_from_name(down_res_type),
                     dropout=dropout,
                     norm_type=norm_type,
                 )
@@ -897,7 +897,7 @@ class MotifEncoder(VideoBaseAE):
             ]
         )
 
-        self.norm_out = Normalize(base_channels[-1], norm_type=norm_type)
+        self.norm_out = make_norm(base_channels[-1], norm_type=norm_type)
         # conv_out emits latent_dim channels directly (no mean/log-var split).
         if self.input_type == "video":
             self.conv_out = CausalConv3d(
@@ -914,11 +914,11 @@ class MotifEncoder(VideoBaseAE):
             h = down_block(h)
         h = self.mid(h)
         h = self.norm_out(h)
-        h = nonlinearity(h)
+        h = silu(h)
         return self.conv_out(h)
 
 
-class MotifDecoder(VideoBaseAE):
+class MotifDecoder(MotifVAEBase):
     """Asymmetric decoder mirroring the encoder's stage layout.
 
     The decoder is intentionally wider than the encoder
@@ -991,7 +991,7 @@ class MotifDecoder(VideoBaseAE):
                     out_channels=base_channels[idx - 1],
                     num_res_blocks=num_resblocks,
                     up_type=up_type,
-                    res_block=resolve_str_to_obj(up_res_type),
+                    res_block=block_from_name(up_res_type),
                     t_interpolation=t_interpolation,
                     dropout=dropout,
                     norm_type=norm_type,
@@ -999,7 +999,7 @@ class MotifDecoder(VideoBaseAE):
                 )
             )
 
-        self.norm_out = Normalize(base_channels[0], norm_type=norm_type)
+        self.norm_out = make_norm(base_channels[0], norm_type=norm_type)
         self.conv_out = Conv2d(base_channels[0], 3, kernel_size=3, stride=1, padding=1)
 
     def forward(self, z: torch.Tensor):
@@ -1008,7 +1008,7 @@ class MotifDecoder(VideoBaseAE):
         for up_block in self.up_blocks:
             h = up_block(h)
         h = self.norm_out(h)
-        h = nonlinearity(h)
+        h = silu(h)
         return self.conv_out(h)
 
 
@@ -1017,7 +1017,7 @@ class MotifDecoder(VideoBaseAE):
 # ---------------------------------------------------------------------------
 
 
-class MotifVAE(VideoBaseAE):
+class MotifVAE(MotifVAEBase):
     """3D causal video VAE with a 128-channel deterministic latent.
 
     Compression is 4x temporal and 32x spatial. The encoder (width 96) uses a
@@ -1056,13 +1056,9 @@ class MotifVAE(VideoBaseAE):
         upsample_residual: bool = False,
     ) -> None:
         super().__init__()
-        self.use_tiling = False
-        self.t_chunk_enc = 16
-        self.t_chunk_dec = 4
-        self.use_quant_layer = False
 
         if scale is None:
-            scale = [0.18215] * latent_dim
+            scale = [1.0] * latent_dim
         if shift is None:
             shift = [0.0] * latent_dim
         if decoder_base_channels is None:
@@ -1109,7 +1105,7 @@ class MotifVAE(VideoBaseAE):
         # posterior API over the latent.
         h = self.encoder(x)
         posterior = DeterministicLatent(h)
-        return AutoencoderKLOutput(latent_dist=posterior, extra_output=None)
+        return MotifEncoderOutput(latent_dist=posterior, extra_output=None)
 
     def decode(self, z, **kwargs):
         dec = self.decoder(z)
@@ -1125,8 +1121,3 @@ class MotifVAE(VideoBaseAE):
             sampled_latent=z,
             extra_output=None,
         )
-
-    def get_last_layer(self):
-        if hasattr(self.decoder.conv_out, "conv"):
-            return self.decoder.conv_out.conv.weight
-        return self.decoder.conv_out.weight