"""MotifVAE — a 3D causal video VAE for HuggingFace ``AutoModel``.

MotifVAE compresses video by 4x in time and 32x in each spatial dimension
into a 128-channel deterministic latent. The
encoder is fully 3D causal and uses a parameter-free Global Skip Connection
(GSC) at every down stage; the decoder mirrors the stage layout but is
deliberately wider (asymmetric, decoder-heavy). Spatial upsampling is
sub-pixel convolution with ICNR initialisation to suppress checkerboard
artifacts.

The whole model lives in this single file so the released checkpoint loads
straight from the Hub via remote code::

    import diffusers
    vae = diffusers.AutoModel.from_pretrained("<repo>", trust_remote_code=True)
"""

from dataclasses import dataclass
from typing import List, Literal, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange

from diffusers import ConfigMixin, ModelMixin
from diffusers.configuration_utils import register_to_config
from diffusers.utils import BaseOutput


# ---------------------------------------------------------------------------
# Block-type resolution
# ---------------------------------------------------------------------------
#
# The config encodes residual / mid block choices as strings (e.g.
# ``"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.

_NAME_OVERRIDES = {"Identity": nn.Identity}


def block_from_name(name: str):
    """Resolve a block-type string from the config to its class."""
    if name in _NAME_OVERRIDES:
        return _NAME_OVERRIDES[name]
    try:
        return globals()[name]
    except KeyError as exc:
        raise AttributeError(
            f"block_from_name: unknown class name {name!r} "
            f"(not defined in this module)"
        ) from exc


class MotifVAEBase(ModelMixin, ConfigMixin):
    """Base class wiring diffusers' ``ModelMixin`` + ``ConfigMixin`` together."""

    config_name = "config.json"

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)


# ---------------------------------------------------------------------------
# Primitive ops
# ---------------------------------------------------------------------------


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
    unfolded — letting purely spatial layers operate transparently on video.
    """

    def wrapper(self, x, *args, **kwargs):
        if x.dim() == 5:
            t = x.shape[2]
            x = rearrange(x, "b c t h w -> (b t) c h w").contiguous()
            x = func(self, x, *args, **kwargs)
            x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
        else:
            x = func(self, x, *args, **kwargs)
        return x

    return wrapper


def silu(x: torch.Tensor) -> torch.Tensor:
    """SiLU / swish activation used throughout the network."""
    return x * torch.sigmoid(x)


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)


# ---------------------------------------------------------------------------
# Convolutions
# ---------------------------------------------------------------------------


class Conv2d(nn.Conv2d):
    """``nn.Conv2d`` that also accepts 5D video tensors (per-frame)."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int]] = 3,
        stride: Union[int, Tuple[int]] = 1,
        padding: Union[str, int, Tuple[int]] = 0,
        dilation: Union[int, Tuple[int]] = 1,
        groups: int = 1,
        bias: bool = True,
        padding_mode: str = "zeros",
        device=None,
        dtype=None,
    ) -> None:
        super().__init__(
            in_channels,
            out_channels,
            kernel_size,
            stride,
            padding,
            dilation,
            groups,
            bias,
            padding_mode,
            device,
            dtype,
        )

    @per_frame_2d
    def forward(self, x):
        return super().forward(x)


class CausalConv3d(nn.Module):
    """A plain 3D convolution wrapped to keep input dtype aligned with weights.

    The wrapper exists so spatial/temporal padding is configured consistently
    across the network and so bf16/fp32 mismatches under AMP are resolved by
    casting the input to the convolution's weight dtype.
    """

    def __init__(
        self,
        chan_in: int,
        chan_out: int,
        kernel_size: Union[int, Tuple[int, int, int]],
        bias: bool = True,
        **kwargs,
    ) -> None:
        super().__init__()
        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 = _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(
            chan_in,
            chan_out,
            self.kernel_size,
            stride=stride,
            padding=padding,
            bias=bias,
        )

    def forward(self, x):
        if x.dtype is not self.conv.weight.dtype:
            x = x.to(dtype=self.conv.weight.dtype)
        return self.conv(x)


# ---------------------------------------------------------------------------
# Normalisation
# ---------------------------------------------------------------------------


class LayerNorm(nn.Module):
    """Channel-last layer normalisation for both 4D and 5D tensors."""

    def __init__(self, num_channels, eps=1e-6, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.norm = torch.nn.LayerNorm(num_channels, eps=eps, elementwise_affine=True)

    def forward(self, x):
        if x.dim() == 5:
            x = rearrange(x, "b c t h w -> b t h w c")
            x = self.norm(x)
            x = rearrange(x, "b t h w c -> b c t h w")
        else:
            x = rearrange(x, "b c h w -> b h w c")
            x = self.norm(x)
            x = rearrange(x, "b h w c -> b c h w")
        return x


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(
            num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True
        )
    elif norm_type == "layernorm":
        return LayerNorm(num_channels=in_channels, eps=1e-6)


# ---------------------------------------------------------------------------
# Up / down sampling
# ---------------------------------------------------------------------------
#
# Spatial upsampling is sub-pixel convolution everywhere: a stride-1 conv that
# expands channels by ``r**2`` followed by a pixel shuffle and ICNR
# initialisation. This is the only spatial upsample mode the model ships.


def spatial_pixel_shuffle_5d(x: torch.Tensor, upscale_factor: int = 2) -> torch.Tensor:
    """Apply a spatial-only pixel shuffle to a 5D ``(B, C*r**2, T, H, W)`` tensor.

    Returns ``(B, C, T, H*r, W*r)`` — equivalent to running
    ``F.pixel_shuffle`` independently on every frame.
    """
    return rearrange(
        x,
        "b (c r1 r2) t h w -> b c t (h r1) (w r2)",
        r1=upscale_factor,
        r2=upscale_factor,
    )


def icnr_init_(weight: torch.Tensor, upscale_factor: int = 2) -> None:
    """Initialise a sub-pixel conv weight with ICNR (Aitken et al., 2017).

    The ``r**2`` output-channel groups that the pixel shuffle redistributes
    into one ``r x r`` block are initialised identically, so the layer behaves
    as nearest-neighbour upsampling at initialisation (no checkerboard).

    Works for 2D and 3D kernels alike: only the leading channel dimension that
    the shuffle reorders is touched, leaving any temporal axis untouched.
    """
    out_full = weight.shape[0]
    r2 = upscale_factor ** 2
    if out_full % r2 != 0:
        raise ValueError(
            f"icnr_init_: out_channels={out_full} must be divisible by "
            f"upscale_factor**2={r2}"
        )
    sub = out_full // r2
    base = torch.empty(sub, *weight.shape[1:], device=weight.device, dtype=weight.dtype)
    nn.init.kaiming_normal_(base)
    with torch.no_grad():
        weight.copy_(base.repeat_interleave(r2, dim=0))


class Upsample(nn.Module):
    """2x spatial upsample (image branch) via sub-pixel convolution.

    ``Conv2d(out=C*4) -> PixelShuffle(2)`` with ICNR init. When ``residual``
    is set, a parameter-free nearest-neighbour identity skip is added on top,
    scaled by a non-learnable ``residual_alpha`` buffer. The buffer is
    non-persistent so a checkpoint saved without it still loads cleanly.
    """

    def __init__(self, in_channels, out_channels, residual: bool = False, **kwargs):
        super().__init__()
        self.upscale_factor = 2
        self.conv = torch.nn.Conv2d(
            in_channels,
            out_channels * (self.upscale_factor ** 2),
            kernel_size=3,
            stride=1,
            padding=1,
        )
        icnr_init_(self.conv.weight, self.upscale_factor)
        if self.conv.bias is not None:
            nn.init.zeros_(self.conv.bias)
        self.shuffle = nn.PixelShuffle(self.upscale_factor)

        self.residual = residual
        self._in_channels = int(in_channels)
        self._out_channels = int(out_channels)
        if residual:
            self.register_buffer("residual_alpha", torch.tensor(1.0), persistent=False)

    def _identity_residual(self, x: torch.Tensor) -> torch.Tensor:
        """Nearest 2x spatial upsample plus channel match to ``out_channels``.

        The canonical configs use ``in_c == out_c`` for every up block, so only
        the passthrough branch is exercised; the channel-average / repeat
        branches are kept for configs whose stages change channel count.
        """
        in_c = x.shape[1]
        out_c = self._out_channels
        if in_c == out_c:
            r = x
        elif in_c > out_c:
            if in_c % out_c != 0:
                raise ValueError(
                    f"residual channel-average failed: in_c={in_c} not divisible "
                    f"by out_c={out_c}; configure matching channels or disable residual."
                )
            r = x.reshape(x.shape[0], out_c, in_c // out_c, *x.shape[2:]).mean(dim=2)
        else:
            if out_c % in_c != 0:
                raise ValueError(
                    f"residual channel-repeat failed: out_c={out_c} not divisible "
                    f"by in_c={in_c}; configure matching channels or disable residual."
                )
            r = x.repeat_interleave(out_c // in_c, dim=1)
        return F.interpolate(r, scale_factor=2.0, mode="nearest")

    @per_frame_2d
    def forward(self, x):
        h = self.shuffle(self.conv(x))
        if self.residual:
            # Cast alpha to the main path's dtype so an AMP-bf16 forward is not
            # silently promoted to fp32 by the float buffer.
            alpha = self.residual_alpha.to(dtype=h.dtype, device=h.device)
            h = h + alpha * self._identity_residual(x).to(dtype=h.dtype)
        return h


class Downsample(nn.Module):
    """2x spatial downsample (encoder): asymmetric zero-pad then stride-2 conv."""

    def __init__(self, in_channels, out_channels, **kwargs):
        super().__init__()
        self.conv = torch.nn.Conv2d(
            in_channels, out_channels, kernel_size=3, stride=2, padding=0
        )

    @per_frame_2d
    def forward(self, x):
        x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)
        return self.conv(x)


class STDownsample(nn.Module):
    """Joint 2x spatial + 2x temporal downsample via a stride-2 causal conv."""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = CausalConv3d(
            in_channels, out_channels, kernel_size=3, padding=0, stride=2
        )

    def forward(self, x):
        x = F.pad(x, (0, 1, 0, 1, 1, 1), mode="constant", value=0)
        return self.conv(x)


class STUpsample(nn.Module):
    """Joint spatial + temporal 2x upsample.

    The temporal axis is stretched with ``F.interpolate`` (cheap and smooth);
    the spatial 2x is sub-pixel convolution — a ``CausalConv3d`` expanding
    channels by 4 followed by a per-frame pixel shuffle — applied after the
    temporal stretch. An optional parameter-free identity residual (see
    :class:`Upsample`) is computed against the post-temporal tensor.
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        t_interpolation="trilinear",
        residual: bool = False,
    ):
        super().__init__()
        self.t_interpolation = t_interpolation
        self.upscale_factor = 2
        # Produce 4x out_channels so the per-frame pixel shuffle recovers
        # out_channels at 2x spatial resolution.
        self.conv = CausalConv3d(
            in_channels,
            out_channels * (self.upscale_factor ** 2),
            kernel_size=3,
            padding=1,
        )
        icnr_init_(self.conv.conv.weight, self.upscale_factor)
        if self.conv.conv.bias is not None:
            nn.init.zeros_(self.conv.conv.bias)

        self.residual = residual
        self._in_channels = int(in_channels)
        self._out_channels = int(out_channels)
        if residual:
            self.register_buffer("residual_alpha", torch.tensor(1.0), persistent=False)

    def _identity_residual_5d(self, x: torch.Tensor) -> torch.Tensor:
        """Nearest 2x spatial upsample on a 5D tensor plus channel match.

        Called on the post-temporal tensor so its ``T`` axis already matches
        the main path. See :meth:`Upsample._identity_residual` for the
        channel-match rule.
        """
        in_c = x.shape[1]
        out_c = self._out_channels
        if in_c == out_c:
            r = x
        elif in_c > out_c:
            if in_c % out_c != 0:
                raise ValueError(
                    f"residual channel-average failed: in_c={in_c} not divisible "
                    f"by out_c={out_c}; configure matching channels or disable residual."
                )
            r = x.reshape(x.shape[0], out_c, in_c // out_c, *x.shape[2:]).mean(dim=2)
        else:
            if out_c % in_c != 0:
                raise ValueError(
                    f"residual channel-repeat failed: out_c={out_c} not divisible "
                    f"by in_c={in_c}; configure matching channels or disable residual."
                )
            r = x.repeat_interleave(out_c // in_c, dim=1)
        return F.interpolate(r, scale_factor=(1, 2, 2), mode="nearest")

    def forward(self, x):
        # Temporal upsample. A single-frame (image) input has nothing to
        # stretch; an even frame count upsamples directly; an odd frame count
        # is split so the leading frame is preserved before stretching.
        if x.size(2) > 1:
            T = x.size(2)
            if T % 2 == 0:
                x = F.interpolate(x, scale_factor=(2, 1, 1), mode=self.t_interpolation)
            else:
                x, x_tail = x[:, :, :1], x[:, :, 1:]
                x_tail = F.interpolate(
                    x_tail, scale_factor=(2, 1, 1), mode=self.t_interpolation
                )
                x = torch.cat([x, x_tail], dim=2)

        # Spatial sub-pixel upsample. The temporal stretch is shared between
        # the main and residual paths, so the residual only mirrors the
        # spatial branch.
        h = spatial_pixel_shuffle_5d(self.conv(x), self.upscale_factor)
        if self.residual:
            alpha = self.residual_alpha.to(dtype=h.dtype, device=h.device)
            h = h + alpha * self._identity_residual_5d(x).to(dtype=h.dtype)
        return h


# ---------------------------------------------------------------------------
# Residual blocks
# ---------------------------------------------------------------------------


class ResnetBlock2D(nn.Module):
    """Pre-activation residual block over 2D (or per-frame 5D) features."""

    def __init__(
        self,
        *,
        in_channels,
        out_channels=None,
        conv_shortcut=False,
        norm_type,
        dropout,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels if out_channels is None else out_channels
        self.use_conv_shortcut = conv_shortcut

        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 = 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
        )
        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                self.conv_shortcut = torch.nn.Conv2d(
                    in_channels, out_channels, kernel_size=3, stride=1, padding=1
                )
            else:
                self.nin_shortcut = torch.nn.Conv2d(
                    in_channels, out_channels, kernel_size=1, stride=1, padding=0
                )

    @per_frame_2d
    def forward(self, x):
        h = self.norm1(x)
        h = silu(h)
        h = self.conv1(h)
        h = self.norm2(h)
        h = silu(h)
        h = self.dropout(h)
        h = self.conv2(h)
        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                x = self.conv_shortcut(x)
            else:
                x = self.nin_shortcut(x)
        return x + h


class ResnetBlock3D(nn.Module):
    """Pre-activation residual block over 3D features (causal convs)."""

    def __init__(
        self,
        *,
        in_channels,
        out_channels=None,
        conv_shortcut=False,
        dropout,
        norm_type,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = in_channels if out_channels is None else out_channels
        self.use_conv_shortcut = conv_shortcut

        self.norm1 = make_norm(in_channels, norm_type=norm_type)
        self.conv1 = CausalConv3d(in_channels, out_channels, 3, padding=1)
        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:
            if self.use_conv_shortcut:
                self.conv_shortcut = CausalConv3d(in_channels, out_channels, 3, padding=1)
            else:
                self.nin_shortcut = CausalConv3d(in_channels, out_channels, 1, padding=0)

    def forward(self, x):
        h = self.norm1(x)
        h = silu(h)
        h = self.conv1(h)
        h = self.norm2(h)
        h = silu(h)
        h = self.dropout(h)
        h = self.conv2(h)
        if self.in_channels != self.out_channels:
            if self.use_conv_shortcut:
                x = self.conv_shortcut(x)
            else:
                x = self.nin_shortcut(x)
        return x + h


# ---------------------------------------------------------------------------
# Output containers
# ---------------------------------------------------------------------------
#
# The encoder is deterministic, so ``latent_dist`` always carries a
# ``DeterministicLatent`` (defined further down). The annotation is a forward
# reference; it has no runtime effect.


@dataclass
class MotifEncoderOutput(BaseOutput):
    latent_dist: "DeterministicLatent"
    extra_output: Optional[tuple] = None


@dataclass
class DecoderOutput(BaseOutput):
    sample: torch.Tensor
    extra_output: Optional[tuple] = None


@dataclass
class ForwardOutput(BaseOutput):
    sample: torch.Tensor
    latent_dist: "DeterministicLatent"
    extra_output: Optional[tuple] = None
    # The exact latent consumed by the decoder.
    sampled_latent: Optional[torch.Tensor] = None


# ---------------------------------------------------------------------------
# Latent wrapper
# ---------------------------------------------------------------------------


class DeterministicLatent:
    """Latent wrapper exposing the posterior API used by ``encode``.

    The encoder emits the latent directly (no mean/log-variance split). This
    wrapper exposes ``sample`` / ``mode`` (both returning the latent verbatim)
    and a ``mean`` alias.
    """

    def __init__(self, x: torch.Tensor):
        self.x = x
        self.mean = x
        self.parameters = x

    def sample(self) -> torch.Tensor:
        return self.x

    def mode(self) -> torch.Tensor:
        return self.x


# ---------------------------------------------------------------------------
# Encoder / decoder building blocks
# ---------------------------------------------------------------------------


def build_mid_layer(
    layer_type: str, channels: int, dropout: float, norm_type: str
) -> nn.Module:
    """Build one entry of a mid stack from its config string.

    ``"Identity"`` (and ``None``) build ``nn.Identity`` so a Res-Attn-Res mid
    can be turned into an attention-free Res-pass-Res mid while preserving the
    surrounding residual-block keys. Any other string resolves to a residual
    block with matching in/out channels.
    """
    if layer_type is None or layer_type == "Identity":
        return nn.Identity()
    return block_from_name(layer_type)(
        in_channels=channels,
        out_channels=channels,
        dropout=dropout,
        norm_type=norm_type,
    )


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 ``STDownsample``:
    folding the time axis by 2 needs an even length, and pre-pending (rather
    than appending) the repeated frame keeps the fold causal.
    """
    if x.shape[2] % 2 == 0:
        return x
    first = x[:, :, :1].expand(-1, -1, 1, -1, -1)
    return torch.cat([first, x], dim=2)


class MotifDownBlock(nn.Module):
    """Encoder down stage with a parameter-free Global Skip Connection.

    The main path runs ``ResBlock(s) -> downsample -> ResBlock`` to reach
    ``out_channels``. In parallel, the block's own input is folded by
    space-to-channel and channel-averaged to ``out_channels``, then added to
    the main path. ``hw`` stages fold space only; ``thw`` stages also fold time
    by 2 (after padding odd ``T``) to match the main path's joint downsample.

    For the canonical config the space-to-channel fold over-produces channels
    by an integer factor, so the alignment is always an exact channel-average;
    a zero-pad fallback covers configs where it does not.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_res_blocks: int = 2,
        down_type: Literal["thw", "hw"] = "thw",
        res_block: nn.Module = ResnetBlock3D,
        dropout: float = 0.0,
        norm_type: str = "layernorm",
    ) -> None:
        super().__init__()
        assert num_res_blocks >= 2, "num_res_blocks too small (need >= 2)"
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.down_type = down_type

        # Pre-downsample residual stack (channel-preserving).
        self.res_block = nn.Sequential(
            *[
                res_block(
                    in_channels=in_channels,
                    out_channels=in_channels,
                    dropout=dropout,
                    norm_type=norm_type,
                )
                for _ in range(num_res_blocks - 1)
            ]
        )

        if down_type == "thw":
            self.down = STDownsample(
                in_channels=in_channels, out_channels=in_channels
            )
        elif down_type == "hw":
            self.down = Downsample(in_channels=in_channels, out_channels=in_channels)
        else:
            raise ValueError(f"unknown down_type={down_type}")

        # Post-downsample residual carrying the main path to out_channels.
        self.out_res_block = res_block(
            in_channels=in_channels,
            out_channels=out_channels,
            dropout=dropout,
            norm_type=norm_type,
        )

    def _space_to_channel(self, raw: torch.Tensor) -> torch.Tensor:
        """Fold spatial (and, for ``thw``, temporal) resolution into channels."""
        if self.down_type == "hw":
            return rearrange(
                raw, "b c t (h p1) (w p2) -> b (c p1 p2) t h w", p1=2, p2=2
            )
        raw = pad_time_to_even(raw)
        return rearrange(
            raw,
            "b c (t p0) (h p1) (w p2) -> b (c p0 p1 p2) t h w",
            p0=2,
            p1=2,
            p2=2,
        )

    def forward(self, h: torch.Tensor):
        # The skip is tapped from the block input, before any transform.
        h_in = h

        h = self.res_block(h)
        h = self.down(h)
        h_main = self.out_res_block(h)

        skip = self._space_to_channel(h_in)
        out_c = h_main.shape[1]
        skip_c = skip.shape[1]
        if skip_c == out_c:
            skip_aligned = skip
        elif skip_c > out_c:
            assert skip_c % out_c == 0, (
                f"GSC channel-average failed: skip_c={skip_c} not divisible "
                f"by out_c={out_c}.  Re-check base_channels and layer_type."
            )
            g = skip_c // out_c
            skip_aligned = skip.reshape(
                skip.shape[0], out_c, g, *skip.shape[2:]
            ).mean(dim=2)
        else:
            # Zero-pad fallback for configs whose fold under-produces channels.
            pad = torch.zeros(
                skip.shape[0],
                out_c - skip_c,
                *skip.shape[2:],
                dtype=skip.dtype,
                device=skip.device,
            )
            skip_aligned = torch.cat([skip, pad], dim=1)

        assert skip_aligned.shape[2:] == h_main.shape[2:], (
            f"GSC skip shape {skip_aligned.shape} does not match main path "
            f"{h_main.shape}."
        )
        return h_main + skip_aligned


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:`STUpsample`).
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        num_res_blocks: int = 2,
        up_type: Literal["thw", "hw"] = "thw",
        res_block: nn.Module = ResnetBlock3D,
        dropout: float = 0.0,
        norm_type: str = "layernorm",
        t_interpolation: str = "trilinear",
        upsample_residual: bool = False,
    ) -> None:
        super().__init__()
        assert num_res_blocks >= 2, "num_res_blocks too small (need >= 2)"
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.up_type = up_type

        # Pre-upsample residual stack (channel-preserving).
        self.res_block = nn.Sequential(
            *[
                res_block(
                    in_channels=in_channels,
                    out_channels=in_channels,
                    dropout=dropout,
                    norm_type=norm_type,
                )
                for _ in range(num_res_blocks - 1)
            ]
        )

        if up_type == "thw":
            self.up = STUpsample(
                in_channels=in_channels,
                out_channels=in_channels,
                t_interpolation=t_interpolation,
                residual=upsample_residual,
            )
        elif up_type == "hw":
            self.up = Upsample(
                in_channels=in_channels,
                out_channels=in_channels,
                residual=upsample_residual,
            )
        else:
            raise ValueError(f"unknown up_type={up_type}")

        self.out_res_block = res_block(
            in_channels=in_channels,
            out_channels=out_channels,
            dropout=dropout,
            norm_type=norm_type,
        )

    def forward(self, x: torch.Tensor):
        x = self.res_block(x)
        x = self.up(x)
        return self.out_res_block(x)


# ---------------------------------------------------------------------------
# Encoder / decoder
# ---------------------------------------------------------------------------


class MotifEncoder(MotifVAEBase):
    """Deterministic 3D causal encoder.

    For a ``(B, 3, T, H, W)`` clip the stem is a stride-1 conv; five down
    stages reduce resolution (3 spatial-only ``hw`` stages then 2 joint
    ``thw`` stages), each adding its Global Skip Connection; a mid stack and a
    final conv emit ``latent_dim`` channels directly at ``(ceil(T/4), H/32,
    W/32)``.
    """

    @register_to_config
    def __init__(
        self,
        latent_dim: int = 128,
        num_resblocks: int = 2,
        dropout: float = 0.0,
        input_type: Literal["video", "image"] = "video",
        norm_type: str = "layernorm",
        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
        mid_layers_type: List[str] = [
            "ResnetBlock3D",
            "Identity",
            "ResnetBlock3D",
        ],
        down_layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
        down_layer_res_type: List[str] = [
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock3D",
            "ResnetBlock3D",
        ],
    ) -> None:
        super().__init__()
        assert len(base_channels) == len(down_layer_type) + 1, (
            f"base_channels must have one more entry than down_layer_type; "
            f"got {len(base_channels)} vs {len(down_layer_type)}+1"
        )
        assert len(down_layer_type) == len(down_layer_res_type), (
            f"down_layer_type and down_layer_res_type length mismatch: "
            f"{len(down_layer_type)} vs {len(down_layer_res_type)}"
        )
        self.input_type = input_type
        self.conv_in = Conv2d(3, base_channels[0], kernel_size=3, stride=1, padding=1)

        self.down_blocks = nn.ModuleList()
        for idx, (down_type, down_res_type) in enumerate(
            zip(down_layer_type, down_layer_res_type)
        ):
            self.down_blocks.append(
                MotifDownBlock(
                    in_channels=base_channels[idx],
                    out_channels=base_channels[idx + 1],
                    num_res_blocks=num_resblocks,
                    down_type=down_type,
                    res_block=block_from_name(down_res_type),
                    dropout=dropout,
                    norm_type=norm_type,
                )
            )

        self.mid = nn.Sequential(
            *[
                build_mid_layer(t, base_channels[-1], dropout, norm_type)
                for t in mid_layers_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(
                base_channels[-1], latent_dim, kernel_size=3, stride=1, padding=1
            )
        else:
            self.conv_out = Conv2d(
                base_channels[-1], latent_dim, kernel_size=3, stride=1, padding=1
            )

    def forward(self, x: torch.Tensor):
        h = self.conv_in(x)
        for down_block in self.down_blocks:
            h = down_block(h)
        h = self.mid(h)
        h = self.norm_out(h)
        h = silu(h)
        return self.conv_out(h)


class MotifDecoder(MotifVAEBase):
    """Asymmetric decoder mirroring the encoder's stage layout.

    The decoder is intentionally wider than the encoder
    (``base_channels=[144, 288, 576, 1152, 1152, 1152]`` in the released
    config). ``upsample_residual=True`` enables the parameter-free identity
    skip in every spatial upsample, scaled by a non-persistent
    ``residual_alpha`` buffer.
    """

    @register_to_config
    def __init__(
        self,
        latent_dim: int = 128,
        num_resblocks: int = 2,
        dropout: float = 0.0,
        input_type: Literal["video", "image"] = "video",
        norm_type: str = "layernorm",
        t_interpolation: str = "trilinear",
        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
        up_layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
        mid_layers_type: List[str] = [
            "ResnetBlock3D",
            "Identity",
            "ResnetBlock3D",
        ],
        up_layer_res_type: List[str] = [
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock3D",
            "ResnetBlock3D",
        ],
        upsample_residual: bool = False,
    ) -> None:
        super().__init__()
        assert len(base_channels) == len(up_layer_type) + 1, (
            f"base_channels must have one more entry than up_layer_type; "
            f"got {len(base_channels)} vs {len(up_layer_type)}+1"
        )
        assert len(up_layer_type) == len(up_layer_res_type), (
            f"up_layer_type and up_layer_res_type length mismatch: "
            f"{len(up_layer_type)} vs {len(up_layer_res_type)}"
        )
        self.input_type = input_type
        if self.input_type == "video":
            self.conv_in = CausalConv3d(
                latent_dim, base_channels[-1], kernel_size=3, stride=1, padding=1
            )
        else:
            self.conv_in = Conv2d(
                latent_dim, base_channels[-1], kernel_size=3, stride=1, padding=1
            )

        self.mid = nn.Sequential(
            *[
                build_mid_layer(t, base_channels[-1], dropout, norm_type)
                for t in mid_layers_type
            ]
        )

        self.up_blocks = nn.ModuleList()
        # Iterate from the deepest stage back up to the shallowest.
        for idx, (up_type, up_res_type) in enumerate(
            zip(reversed(up_layer_type), reversed(up_layer_res_type))
        ):
            idx = len(up_layer_type) - idx
            self.up_blocks.append(
                MotifUpBlock(
                    in_channels=base_channels[idx],
                    out_channels=base_channels[idx - 1],
                    num_res_blocks=num_resblocks,
                    up_type=up_type,
                    res_block=block_from_name(up_res_type),
                    t_interpolation=t_interpolation,
                    dropout=dropout,
                    norm_type=norm_type,
                    upsample_residual=upsample_residual,
                )
            )

        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):
        h = self.conv_in(z)
        h = self.mid(h)
        for up_block in self.up_blocks:
            h = up_block(h)
        h = self.norm_out(h)
        h = silu(h)
        return self.conv_out(h)


# ---------------------------------------------------------------------------
# Top-level model
# ---------------------------------------------------------------------------


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
    parameter-free Global Skip Connection at every stage; the decoder is wider
    (width 144).
    """

    @register_to_config
    def __init__(
        self,
        latent_dim: int = 128,
        base_channels: List[int] = [96, 192, 384, 768, 768, 768],
        decoder_base_channels: Optional[List[int]] = [144, 288, 576, 1152, 1152, 1152],
        decoder_variant: Literal["V4"] = "V4",
        layer_type: List[str] = ["hw", "hw", "hw", "thw", "thw"],
        layer_res_type: List[str] = [
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock2D",
            "ResnetBlock3D",
            "ResnetBlock3D",
        ],
        encoder_num_resblocks: int = 2,
        decoder_num_resblocks: int = 2,
        dropout: float = 0.0,
        norm_type: str = "layernorm",
        t_interpolation: str = "trilinear",
        input_type: Literal["video", "image"] = "video",
        mid_layers_type: List[str] = [
            "ResnetBlock3D",
            "Identity",
            "ResnetBlock3D",
        ],
        scale: Optional[List[float]] = None,
        shift: Optional[List[float]] = None,
        upsample_residual: bool = False,
    ) -> None:
        super().__init__()

        if scale is None:
            scale = [1.0] * latent_dim
        if shift is None:
            shift = [0.0] * latent_dim
        if decoder_base_channels is None:
            decoder_base_channels = [144, 288, 576, 1152, 1152, 1152]

        self.encoder = MotifEncoder(
            latent_dim=latent_dim,
            base_channels=base_channels,
            num_resblocks=encoder_num_resblocks,
            dropout=dropout,
            norm_type=norm_type,
            down_layer_type=layer_type,
            down_layer_res_type=layer_res_type,
            input_type=input_type,
            mid_layers_type=mid_layers_type,
        )

        if decoder_variant != "V4":
            raise ValueError(
                f"unknown decoder_variant={decoder_variant!r}; expected 'V4'"
            )
        self.decoder = MotifDecoder(
            latent_dim=latent_dim,
            base_channels=decoder_base_channels,
            num_resblocks=decoder_num_resblocks,
            dropout=dropout,
            norm_type=norm_type,
            t_interpolation=t_interpolation,
            up_layer_type=layer_type,
            up_layer_res_type=layer_res_type,
            input_type=input_type,
            mid_layers_type=mid_layers_type,
            upsample_residual=upsample_residual,
        )

    def get_encoder(self):
        return [self.encoder]

    def get_decoder(self):
        return [self.decoder]

    def encode(self, x, **kwargs):
        # Deterministic encoder: DeterministicLatent exposes the sample()/mode()
        # posterior API over the latent.
        h = self.encoder(x)
        posterior = DeterministicLatent(h)
        return MotifEncoderOutput(latent_dist=posterior, extra_output=None)

    def decode(self, z, **kwargs):
        dec = self.decoder(z)
        return DecoderOutput(sample=dec, extra_output=None)

    def forward(self, input, sample_posterior=True):
        posterior = self.encode(input).latent_dist
        z = posterior.sample() if sample_posterior else posterior.mode()
        dec = self.decode(z).sample
        return ForwardOutput(
            sample=dec,
            latent_dist=posterior,
            sampled_latent=z,
            extra_output=None,
        )