modeling_motif_vision_encoder.py

# Copyright (c) Motif Technologies.
# Self-contained inference model for Motif Vision Encoder (image + video).
# Auto-assembled from the training repo's inference path; NO training code.
"""Motif Vision Encoder — unified image/video ViT backbone (inference-only).

Usage:
    from transformers import AutoModel
    import torch
    model = AutoModel.from_pretrained("Motif-Technologies/motif-vision-encoder",
                                      trust_remote_code=True).eval()
    # image: (B, 3, H, W)   video: (B, T, 3, H, W)  (H,W multiples of 16)
    out = model(pixel_values=torch.randn(1, 3, 224, 224))
    out.last_hidden_state            # (B, 1+num_register+N, D)
    out.pooler_output                # (B, D)  CLS token
"""
import logging
import math
from functools import partial
from typing import Any, Callable, Literal

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPooling



# ---- utils ----


def cat_keep_shapes(x_list: list[Tensor]) -> tuple[Tensor, list[tuple[int]], list[int]]:
    """Concatenate list of tensors while preserving their shapes for later reconstruction."""
    shapes = [x.shape for x in x_list]
    num_tokens = [x.select(dim=-1, index=0).numel() for x in x_list]
    flattened = torch.cat([x.flatten(0, -2) for x in x_list])
    return flattened, shapes, num_tokens



def uncat_with_shapes(flattened: Tensor, shapes: list[tuple[int]], num_tokens: list[int]) -> list[Tensor]:
    """Reverse of cat_keep_shapes: split and reshape flattened tensor back to original shapes."""
    outputs_splitted = torch.split_with_sizes(flattened, num_tokens, dim=0)
    shapes_adjusted = [shape[:-1] + torch.Size([flattened.shape[-1]]) for shape in shapes]
    outputs_reshaped = [o.reshape(shape) for o, shape in zip(outputs_splitted, shapes_adjusted)]
    return outputs_reshaped



# ---- rms_norm ----


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization.

    A simpler alternative to LayerNorm that normalizes by RMS without centering.

    Args:
        dim: Number of features.
        eps: Small constant for numerical stability.
    """

    def __init__(
        self,
        dim: int,
        eps: float = 1e-6,
        device: torch.device | str | None = None,
    ) -> None:
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim, device=device))

    def forward(self, x: Tensor) -> Tensor:
        """Apply RMS normalization."""
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.weight


# ---- layer_scale ----


class LayerScale(nn.Module):
    """Per-channel scaling that allows gradual incorporation of each layer's contribution.

    Initializes to a small value (e.g., 1e-5) so that early in training, each layer's
    contribution is nearly zero, stabilizing deep network training.

    Args:
        dim: Number of channels.
        init_values: Initial value for all channels.
        inplace: Whether to apply scaling in-place.
        device: Device for parameter allocation.
    """

    def __init__(
        self,
        dim: int,
        init_values: float | Tensor = 1e-5,
        inplace: bool = False,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        self.inplace = inplace
        self.gamma = nn.Parameter(torch.empty(dim, device=device))
        self.init_values = init_values

    def forward(self, x: Tensor) -> Tensor:
        """Apply per-channel scaling."""
        return x.mul_(self.gamma) if self.inplace else x * self.gamma


# ---- patch_embed ----


class PatchEmbed(nn.Module):
    """Video (5D) or Image (4D) to Patch Embedding via 3D Convolution.

    Handles both modalities through a single Conv3d projection:
    - Image (B, C, H, W): unsqueeze temporal dim -> (B, C, 1, H, W) -> Conv3d
    - Video (B, T, C, H, W): transpose -> (B, C, T, H, W) -> Conv3d

    Output: (B, N_total, embed_dim)
        N_total = (T // tubelet_size) * (H // patch_size) * (W // patch_size)

    Args:
        img_size: Input image size (used for reference only).
        patch_size: Spatial patch size in pixels.
        in_chans: Number of input channels.
        embed_dim: Output embedding dimension.
        tubelet_size: Temporal patch size (number of frames per temporal token).
        flatten_embedding: Whether to flatten spatial dimensions.
    """

    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        embed_dim: int = 768,
        tubelet_size: int = 1,
        flatten_embedding: bool = True,
    ) -> None:
        super().__init__()
        self.img_size = img_size
        self.patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
        self.tubelet_size = tubelet_size
        self.flatten_embedding = flatten_embedding
        self.in_chans = in_chans

        # 3D Convolution: kernel and stride = (tubelet_size, patch_h, patch_w)
        self.proj = nn.Conv3d(
            in_chans,
            embed_dim,
            kernel_size=(tubelet_size, self.patch_size[0], self.patch_size[1]),
            stride=(tubelet_size, self.patch_size[0], self.patch_size[1]),
        )

    def forward(self, x: Tensor) -> Tensor:
        """Tokenize input images or videos.

        Args:
            x: Input tensor.
               Image: (B, C, H, W) or Video: (B, T, C, H, W)

        Returns:
            Patch tokens of shape (B, N_total, embed_dim).
        """
        if x.ndim == 4:
            # Image: (B, C, H, W) -> (B, C, 1, H, W)
            x = x.unsqueeze(2)
            # If tubelet_size > 1, repeat the single frame to match kernel size
            if self.tubelet_size > 1:
                x = x.expand(-1, -1, self.tubelet_size, -1, -1)
        elif x.ndim == 5:
            # Video: (B, T, C, H, W) -> (B, C, T, H, W)
            x = x.transpose(1, 2)

        # Conv3d Projection -> (B, embed_dim, T', H', W')
        x = self.proj(x)

        if self.flatten_embedding:
            # Flatten spatial+temporal: (B, embed_dim, N) -> (B, N, embed_dim)
            x = x.flatten(2).transpose(1, 2)

        return x


# ---- rope ----


class RopePositionEmbedding3D(nn.Module):
    """Full 3D axial RoPE with independent T/H/W frequency bands.

    Unlike the original Motif implementation which simply repeats 2D spatial angles
    across temporal frames (making temporal positions indistinguishable), this
    implementation partitions the head dimension into three axis groups:

        D_head = D_T + D_H + D_W   (no spare dimensions)

    By default, the split is spatial-heavy for SSL (spatial quality is priority):
        D_T = D_head // 4           (25% temporal)
        D_H = (D_head - D_T) // 2   (37.5% height)
        D_W = D_head - D_T - D_H    (37.5% width)
        e.g., D_head=64 → T=16, H=24, W=24

    This can be overridden via ``fhw_dim=(D_T, D_H, D_W)`` for full control.

    For images (T=1): t=0 for all tokens, making temporal angles constant
    and the output is equivalent to spatial-only RoPE.

    Args:
        embed_dim: Total embedding dimension.
        num_heads: Number of attention heads.
        fhw_dim: Optional explicit (D_T, D_H, D_W) partition. Each must be even.
            If None, uses the spatial-heavy default described above.
        base: Frequency base (100.0 for spatial vision convention).
        min_period: Minimum period (alternative to base).
        max_period: Maximum period (alternative to base).
        normalize_coords: How to normalize coordinates.
        shift_coords: Random shift range during training.
        jitter_coords: Random jitter multiplier during training.
        rescale_coords: Random rescale multiplier during training.
        dtype: Data type for computation.
        device: Device for parameter allocation.
    """

    def __init__(
        self,
        embed_dim: int,
        *,
        num_heads: int,
        fhw_dim: tuple[int, int, int] | None = None,
        base: float | None = 100.0,
        min_period: float | None = None,
        max_period: float | None = None,
        normalize_coords: Literal["min", "max", "separate"] = "separate",
        shift_coords: float | None = None,
        jitter_coords: float | None = None,
        rescale_coords: float | None = None,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ) -> None:
        super().__init__()
        both_periods = min_period is not None and max_period is not None
        if (base is None and not both_periods) or (base is not None and both_periods):
            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")

        D_head = embed_dim // num_heads
        self.base = base
        self.min_period = min_period
        self.max_period = max_period
        self.D_head = D_head
        self.normalize_coords = normalize_coords
        self.shift_coords = shift_coords
        self.jitter_coords = jitter_coords
        self.rescale_coords = rescale_coords

        # Partition head dimension into 3 groups: T, H, W (no spare)
        if fhw_dim is not None:
            self.D_T, self.D_H, self.D_W = fhw_dim
            assert self.D_T + self.D_H + self.D_W == D_head, (
                f"fhw_dim must sum to D_head={D_head}, got {sum(fhw_dim)}"
            )
        else:
            # Default: spatial-heavy split (SSL prioritizes spatial quality)
            self.D_T = D_head // 4                        # 25% temporal
            self.D_H = (D_head - self.D_T) // 2           # 37.5% height
            self.D_W = D_head - self.D_T - self.D_H       # 37.5% width
        assert self.D_T % 2 == 0 and self.D_H % 2 == 0 and self.D_W % 2 == 0, (
            f"All axis dims must be even, got T={self.D_T}, H={self.D_H}, W={self.D_W}"
        )

        self.dtype = dtype
        # Separate period buffers for each axis (n_freqs = D_axis // 2)
        self.register_buffer(
            "periods_t",
            torch.empty(self.D_T // 2, device=device, dtype=dtype),
            persistent=True,
        )
        self.register_buffer(
            "periods_h",
            torch.empty(self.D_H // 2, device=device, dtype=dtype),
            persistent=True,
        )
        self.register_buffer(
            "periods_w",
            torch.empty(self.D_W // 2, device=device, dtype=dtype),
            persistent=True,
        )
        self._init_weights()

    def forward(self, *, T: int = 1, H: int, W: int) -> tuple[Tensor, Tensor]:
        """Compute 3D axial RoPE sin/cos for (T, H, W) grid.

        The head dimension is partitioned as [D_T | D_H | D_W]:
        - D_T: temporal frequency bands (angles vary with t)
        - D_H: height frequency bands (angles vary with h)
        - D_W: width frequency bands (angles vary with w)

        For images (T=1), all tokens get t=0, so temporal angles are constant
        and the output is equivalent to spatial-only RoPE.

        Args:
            T: Number of temporal positions (T_grid = num_frames // tubelet_size).
            H: Height in patches.
            W: Width in patches.

        Returns:
            Tuple of (sin, cos), each of shape (T*H*W, D_head).
        """
        device = self.periods_t.device
        dtype = self.dtype
        dd = {"device": device, "dtype": dtype}

        # 1. Compute normalized coordinates for each axis
        if T > 1:
            coords_t = torch.arange(0.5, T, **dd) / T  # [T]
        else:
            coords_t = torch.tensor([0.5], **dd)  # [1] - constant for images

        coords_h, coords_w = self._compute_spatial_coords(H, W, **dd)

        # Shift to [-1, +1] range
        coords_t = 2.0 * coords_t - 1.0  # [T]
        coords_h = 2.0 * coords_h - 1.0  # [H]
        coords_w = 2.0 * coords_w - 1.0  # [W]

        # Apply training-time augmentations to spatial coords only
        if self.training:
            coords_h, coords_w = self._augment_spatial_coords(coords_h, coords_w, dd)

        # 2. Compute raw angles for each axis (n_freqs = D_axis // 2)
        angles_t = 2 * math.pi * coords_t[:, None] / self.periods_t[None, :]  # [T, D_T//2]
        angles_h = 2 * math.pi * coords_h[:, None] / self.periods_h[None, :]  # [H, D_H//2]
        angles_w = 2 * math.pi * coords_w[:, None] / self.periods_w[None, :]  # [W, D_W//2]

        # 3. Build full 3D grid: create (T*H*W, D_head) angle tensor
        t_idx, h_idx, w_idx = torch.meshgrid(
            torch.arange(T, device=device),
            torch.arange(H, device=device),
            torch.arange(W, device=device),
            indexing="ij",
        )
        t_idx = t_idx.flatten()  # [T*H*W]
        h_idx = h_idx.flatten()  # [T*H*W]
        w_idx = w_idx.flatten()  # [T*H*W]

        # Gather per-token raw angles and concatenate to D_head//2
        token_angles_t = angles_t[t_idx]  # [T*H*W, D_T//2]
        token_angles_h = angles_h[h_idx]  # [T*H*W, D_H//2]
        token_angles_w = angles_w[w_idx]  # [T*H*W, D_W//2]
        angles_half = torch.cat([token_angles_t, token_angles_h, token_angles_w], dim=-1)  # [T*H*W, D_head//2]

        # tile(2) on full concat — matches Motif 2D RoPE pattern
        # This ensures rotate_half pairs (dim i ↔ dim i+D//2) have identical angles,
        # making the rotation orthogonal (preserves dot products in attention).
        angles = angles_half.tile(2)  # [T*H*W, D_head]

        cos = torch.cos(angles)
        sin = torch.sin(angles)

        return (sin, cos)

    def _compute_spatial_coords(self, H: int, W: int, **dd) -> tuple[Tensor, Tensor]:
        """Compute normalized spatial coordinates."""
        if self.normalize_coords == "max":
            max_HW = max(H, W)
            coords_h = torch.arange(0.5, H, **dd) / max_HW
            coords_w = torch.arange(0.5, W, **dd) / max_HW
        elif self.normalize_coords == "min":
            min_HW = min(H, W)
            coords_h = torch.arange(0.5, H, **dd) / min_HW
            coords_w = torch.arange(0.5, W, **dd) / min_HW
        elif self.normalize_coords == "separate":
            coords_h = torch.arange(0.5, H, **dd) / H
            coords_w = torch.arange(0.5, W, **dd) / W
        else:
            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
        return coords_h, coords_w

    def _augment_spatial_coords(
        self,
        coords_h: Tensor,
        coords_w: Tensor,
        dd: dict,
    ) -> tuple[Tensor, Tensor]:
        """Apply training-time coordinate augmentations to spatial coords."""
        if self.shift_coords is not None:
            shift = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)
            coords_h = coords_h + shift[0]
            coords_w = coords_w + shift[1]
        if self.jitter_coords is not None:
            jitter_max = np.log(self.jitter_coords)
            jitter = torch.empty(2, **dd).uniform_(-jitter_max, jitter_max).exp()
            coords_h = coords_h * jitter[0]
            coords_w = coords_w * jitter[1]
        if self.rescale_coords is not None:
            rescale_max = np.log(self.rescale_coords)
            rescale = torch.empty(1, **dd).uniform_(-rescale_max, rescale_max).exp()
            coords_h = coords_h * rescale
            coords_w = coords_w * rescale
        return coords_h, coords_w

    def _compute_periods(self, n_freqs: int, device: torch.device, dtype: torch.dtype | None) -> Tensor:
        """Compute frequency periods for a single axis.

        Args:
            n_freqs: Number of frequency bands (D_axis // 2).
            device: Device for tensor allocation.
            dtype: Data type for computation.

        Returns:
            Tensor of shape (n_freqs,) with logarithmically spaced periods.
        """
        if self.base is not None:
            return self.base ** (
                2 * torch.arange(n_freqs, device=device, dtype=dtype) / (2 * n_freqs)
            )
        else:
            base = self.max_period / self.min_period
            exponents = torch.linspace(0, 1, n_freqs, device=device, dtype=dtype)
            periods = base**exponents
            periods = periods / base
            return periods * self.max_period

    def _init_weights(self) -> None:
        """Initialize frequency periods for all three axes.

        Each axis gets its own frequency schedule based on its dimension size:
            periods[i] = base^(2i / D_axis)

        This produces logarithmically spaced periods from 1.0 to base,
        with more frequencies for axes with more allocated dimensions.
        """
        device = self.periods_t.device
        dtype = self.dtype

        self.periods_t.data = self._compute_periods(self.D_T // 2, device, dtype)
        self.periods_h.data = self._compute_periods(self.D_H // 2, device, dtype)
        self.periods_w.data = self._compute_periods(self.D_W // 2, device, dtype)


# ---- attention ----


def rope_rotate_half(x: Tensor) -> Tensor:
    """Rotate half of the dimensions: [-x2, x1] from [x1, x2].

    Args:
        x: Input tensor of shape (..., D).

    Returns:
        Rotated tensor of shape (..., D).
    """
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
    """Apply rotary position embedding to input tensor.

    Args:
        x: Input tensor of shape (..., D).
        sin: Sine angles of shape (..., D).
        cos: Cosine angles of shape (..., D).

    Returns:
        Rotated tensor of shape (..., D).
    """
    return (x * cos) + (rope_rotate_half(x) * sin)


class LinearKMaskedBias(nn.Linear):
    """Linear layer with masked bias for the K component of QKV.

    Zeroes out the bias for the K component (middle third of output)
    to avoid interference with RoPE positional encoding.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        o = self.out_features
        assert o % 3 == 0
        if self.bias is not None:
            self.register_buffer("bias_mask", torch.full_like(self.bias, fill_value=math.nan))

    def forward(self, input: Tensor) -> Tensor:
        """Forward pass with masked bias."""
        masked_bias = self.bias * self.bias_mask.to(self.bias.dtype) if self.bias is not None else None
        return F.linear(input, self.weight, masked_bias)


class SelfAttention(nn.Module):
    """Multi-head self-attention with RoPE support.

    Uses torch.nn.functional.scaled_dot_product_attention for FlashAttention
    compatibility. RoPE is applied to Q and K on patch tokens only (not CLS/register).

    Args:
        dim: Model dimension.
        num_heads: Number of attention heads.
        qkv_bias: Whether to use bias in QKV projection.
        proj_bias: Whether to use bias in output projection.
        attn_drop: Attention dropout probability.
        proj_drop: Output projection dropout probability.
        mask_k_bias: Whether to mask K bias (for RoPE compatibility).
        device: Device for parameter allocation.
        gated_attention: Gated attention variant. None disables gating,
            "headwise" applies a per-head scalar gate, "elementwise" applies
            a per-element gate. Gate scores are query-dependent (derived from
            input) and applied as sigmoid after SDPA.
            Reference: https://arxiv.org/abs/2505.06708
    """

    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        qkv_bias: bool = False,
        proj_bias: bool = True,
        attn_drop: float = 0.0,
        proj_drop: float = 0.0,
        mask_k_bias: bool = False,
        device: str | None = None,
        gated_attention: str | None = None,
        qk_norm: bool = False,
    ) -> None:
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5

        linear_class = LinearKMaskedBias if mask_k_bias else nn.Linear
        self.qkv = linear_class(dim, dim * 3, bias=qkv_bias, device=device)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim, bias=proj_bias, device=device)
        self.proj_drop = nn.Dropout(proj_drop)

        self.qk_norm = qk_norm
        if qk_norm:
            self.q_norm = RMSNorm(self.head_dim, device=device)
            self.k_norm = RMSNorm(self.head_dim, device=device)

        self.gated_attention = gated_attention
        if gated_attention == "headwise":
            self.gate_proj = nn.Linear(dim, num_heads, bias=True, device=device)
        elif gated_attention == "elementwise":
            self.gate_proj = nn.Linear(dim, dim, bias=True, device=device)
        elif gated_attention is not None:
            raise ValueError(f"Unknown gated_attention mode: {gated_attention!r}. Use 'headwise' or 'elementwise'.")

    def apply_rope(
        self,
        q: Tensor,
        k: Tensor,
        rope: tuple[Tensor, Tensor],
    ) -> tuple[Tensor, Tensor]:
        """Apply RoPE to query and key tensors.

        RoPE is applied only to patch tokens (prefix tokens like CLS and register
        are excluded based on the difference between sequence length and rope length).

        Args:
            q: Query tensor of shape (B, heads, N, D_head).
            k: Key tensor of shape (B, heads, N, D_head).
            rope: Tuple of (sin, cos), each of shape (N_patches, D_head).

        Returns:
            Tuple of rotated (q, k) tensors.
        """
        q_dtype = q.dtype
        k_dtype = k.dtype
        sin, cos = rope
        rope_dtype = sin.dtype
        q = q.to(dtype=rope_dtype)
        k = k.to(dtype=rope_dtype)
        N = q.shape[-2]
        prefix = N - sin.shape[-2]
        assert prefix >= 0
        q_prefix = q[:, :, :prefix, :]
        q = rope_apply(q[:, :, prefix:, :], sin, cos)
        q = torch.cat((q_prefix, q), dim=-2)
        k_prefix = k[:, :, :prefix, :]
        k = rope_apply(k[:, :, prefix:, :], sin, cos)
        k = torch.cat((k_prefix, k), dim=-2)
        q = q.to(dtype=q_dtype)
        k = k.to(dtype=k_dtype)
        return q, k

    def forward(self, x: Tensor, attn_bias: Tensor | None = None, rope: Tensor | None = None) -> Tensor:
        """Forward pass for single tensor input.

        Args:
            x: Input tensor of shape (B, N, D).
            attn_bias: Unused (kept for interface compatibility).
            rope: Optional RoPE (sin, cos) tuple.

        Returns:
            Output tensor of shape (B, N, D).
        """
        gate_score = self._compute_gate(x) if self.gated_attention else None
        qkv = self.qkv(x)
        attn_v = self.compute_attention(qkv=qkv, attn_bias=attn_bias, rope=rope, gate_score=gate_score)
        x = self.proj(attn_v)
        x = self.proj_drop(x)
        return x

    def forward_list(
        self,
        x_list: list[Tensor],
        attn_bias: Tensor | None = None,
        rope_list: list[tuple[Tensor, Tensor]] | None = None,
    ) -> list[Tensor]:
        """Forward pass for list of tensors (multi-crop efficiency).

        Concatenates inputs for a single QKV projection, then splits for per-crop
        attention computation (needed because different crops have different RoPE).

        Args:
            x_list: List of input tensors.
            attn_bias: Unused.
            rope_list: List of RoPE (sin, cos) tuples, one per input.

        Returns:
            List of output tensors.
        """
        assert len(x_list) == len(rope_list)
        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
        qkv_flat = self.qkv(x_flat)
        qkv_list = uncat_with_shapes(qkv_flat, shapes, num_tokens)

        if self.gated_attention:
            gate_flat = self._compute_gate(x_flat)
            gate_list = uncat_with_shapes(gate_flat, shapes, num_tokens)
        else:
            gate_list = [None] * len(x_list)

        att_out = []
        for qkv, _, rope, gate_score in zip(qkv_list, shapes, rope_list, gate_list):
            att_out.append(self.compute_attention(qkv, attn_bias=attn_bias, rope=rope, gate_score=gate_score))
        x_flat, shapes, num_tokens = cat_keep_shapes(att_out)
        x_flat = self.proj(x_flat)
        return uncat_with_shapes(x_flat, shapes, num_tokens)

    def _compute_gate(self, x: Tensor) -> Tensor:
        """Compute raw gate scores from input.

        Returns the raw projection without reshaping so that the output keeps
        the same number of leading dimensions as ``x``.  This is critical for
        ``forward_list`` where ``uncat_with_shapes`` must split a 2-D flat
        tensor back to per-crop 3-D tensors — adding extra dims here would
        break that reshape.  The per-head unflatten happens later inside
        ``compute_attention`` where B and N are known.

        Args:
            x: Input tensor of shape (..., D). Supports both 2D (flat) and 3D (batched).

        Returns:
            Raw gate projection. Headwise: (..., num_heads). Elementwise: (..., D).
        """
        return self.gate_proj(x)

    def compute_attention(
        self,
        qkv: Tensor,
        attn_bias: Tensor | None = None,
        rope: tuple[Tensor, Tensor] | None = None,
        gate_score: Tensor | None = None,
    ) -> Tensor:
        """Compute scaled dot-product attention.

        Args:
            qkv: Combined QKV tensor of shape (B, N, 3*D).
            attn_bias: Unused.
            rope: Optional RoPE (sin, cos) tuple.
            gate_score: Optional gate tensor from _compute_gate.

        Returns:
            Attention output of shape (B, N, D).
        """
        assert attn_bias is None
        B, N, _ = qkv.shape
        C = self.qkv.in_features

        qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
        q, k, v = torch.unbind(qkv, 2)
        q, k, v = [t.transpose(1, 2) for t in [q, k, v]]
        if self.qk_norm:
            q = self.q_norm(q)
            k = self.k_norm(k)
        if rope is not None:
            q, k = self.apply_rope(q, k, rope)
        x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
        x = x.transpose(1, 2)  # (B, N, num_heads, head_dim)
        if gate_score is not None:
            # _compute_gate returns raw projection: (..., num_heads) or (..., D).
            # Reshape to (B, N, num_heads, 1) or (B, N, num_heads, head_dim) here.
            if self.gated_attention == "headwise":
                gate_score = gate_score.unflatten(-1, (self.num_heads, 1))
            else:  # elementwise
                gate_score = gate_score.unflatten(-1, (self.num_heads, self.head_dim))
            x = x * torch.sigmoid(gate_score)
        return x.reshape([B, N, C])


# ---- ffn ----


class ListForwardMixin:
    """Mixin providing forward_list for efficient multi-crop processing."""

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass for a single tensor."""
        raise NotImplementedError

    def forward_list(self, x_list: list[Tensor]) -> list[Tensor]:
        """Forward pass for a list of tensors, concatenated for efficiency."""
        x_flat, shapes, num_tokens = cat_keep_shapes(x_list)
        x_flat = self.forward(x_flat)
        return uncat_with_shapes(x_flat, shapes, num_tokens)


class SwiGLUFFN(nn.Module, ListForwardMixin):
    """SwiGLU Feed-Forward Network: w3(silu(w1(x)) * w2(x)).

    Used for larger ViT models (SO400M+) due to better gradient flow.
    Hidden dimension is aligned to a multiple of `align_to` for GPU efficiency.

    Args:
        in_features: Input dimension.
        hidden_features: Hidden dimension before alignment.
        out_features: Output dimension (default: same as in_features).
        act_layer: Unused (SwiGLU has built-in SiLU activation).
        drop: Unused (no dropout in SwiGLU).
        bias: Whether to use bias in linear layers.
        align_to: Align hidden dimension to this multiple.
        device: Device for parameter allocation.
    """

    def __init__(
        self,
        in_features: int,
        hidden_features: int | None = None,
        out_features: int | None = None,
        act_layer: Callable[..., nn.Module] | None = None,
        drop: float = 0.0,
        bias: bool = True,
        align_to: int = 8,
        device: str | None = None,
    ) -> None:
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        d = int(hidden_features * 2 / 3)
        swiglu_hidden_features = d + (-d % align_to)
        self.w1 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
        self.w2 = nn.Linear(in_features, swiglu_hidden_features, bias=bias, device=device)
        self.w3 = nn.Linear(swiglu_hidden_features, out_features, bias=bias, device=device)

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass: w3(silu(w1(x)) * w2(x))."""
        x1 = self.w1(x)
        x2 = self.w2(x)
        hidden = F.silu(x1) * x2
        return self.w3(hidden)


# ---- block ----


class SelfAttentionBlock(nn.Module):
    """Pre-norm transformer block: Norm -> Attention -> LayerScale -> Residual (x2).

    Supports both single-tensor and list-of-tensors forward for efficient multi-crop
    processing.

    Args:
        dim: Model dimension.
        num_heads: Number of attention heads.
        ffn_ratio: FFN hidden dimension ratio.
        qkv_bias: Whether to use bias in QKV projection.
        proj_bias: Whether to use bias in output projection.
        ffn_bias: Whether to use bias in FFN layers.
        drop: Dropout probability.
        attn_drop: Attention dropout probability.
        init_values: LayerScale initial values (None disables LayerScale).
        drop_path: Stochastic depth drop probability.
        act_layer: Activation function class.
        norm_layer: Normalization layer class.
        attn_class: Attention class.
        ffn_layer: FFN class.
        mask_k_bias: Whether to mask K bias.
        device: Device for parameter allocation.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        ffn_ratio: float = 4.0,
        qkv_bias: bool = False,
        proj_bias: bool = True,
        ffn_bias: bool = True,
        drop: float = 0.0,
        attn_drop: float = 0.0,
        init_values: float | None = None,
        drop_path: float = 0.0,
        act_layer: Callable[..., nn.Module] = nn.GELU,
        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
        attn_class: Callable[..., nn.Module] = SelfAttention,
        ffn_layer: Callable[..., nn.Module] = SwiGLUFFN,
        mask_k_bias: bool = False,
        device: str | None = None,
        gated_attention: str | None = None,
        qk_norm: bool = False,
    ) -> None:
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = attn_class(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            proj_bias=proj_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
            mask_k_bias=mask_k_bias,
            device=device,
            gated_attention=gated_attention,
            qk_norm=qk_norm,
        )
        self.ls1 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()

        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * ffn_ratio)
        self.mlp = ffn_layer(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
            bias=ffn_bias,
            device=device,
        )
        self.ls2 = LayerScale(dim, init_values=init_values, device=device) if init_values else nn.Identity()

    def _forward_list(self, x_list: list[Tensor], rope_list: list | None = None) -> list[Tensor]:
        """Forward pass for a list of tensors (one per crop), each with its own RoPE.

        Pre-norm residual: Norm -> Attention -> LayerScale -> Residual, twice (attn, ffn).
        """
        x_out = []
        for x, rope in zip(x_list, rope_list):
            x_attn = x + self.ls1(self.attn(self.norm1(x), rope=rope))
            x_ffn_item = x_attn + self.ls2(self.mlp(self.norm2(x_attn)))
            x_out.append(x_ffn_item)
        return x_out

    def forward(
        self,
        x_or_x_list: Tensor | list[Tensor],
        rope_or_rope_list: tuple | list | None = None,
    ) -> Tensor | list[Tensor]:
        """Forward pass accepting either a single tensor or list of tensors.

        Args:
            x_or_x_list: Single tensor (B, N, D) or list of tensors.
            rope_or_rope_list: Single RoPE tuple or list of RoPE tuples.

        Returns:
            Output tensor(s) matching input format.
        """
        if isinstance(x_or_x_list, Tensor):
            return self._forward_list([x_or_x_list], rope_list=[rope_or_rope_list])[0]
        elif isinstance(x_or_x_list, list):
            if rope_or_rope_list is None:
                rope_or_rope_list = [None for _ in x_or_x_list]
            return self._forward_list(x_or_x_list, rope_list=rope_or_rope_list)
        else:
            raise AssertionError(f"Unexpected input type: {type(x_or_x_list)}")


# ---- vision_transformer ----


logger = logging.getLogger("motif")

ffn_layer_dict: dict[str, type] = {
    "swiglu": SwiGLUFFN,
    "swiglu32": partial(SwiGLUFFN, align_to=32),
    "swiglu64": partial(SwiGLUFFN, align_to=64),
    "swiglu128": partial(SwiGLUFFN, align_to=128),
}

norm_layer_dict: dict[str, type] = {
    "layernorm": partial(nn.LayerNorm, eps=1e-6),
    "layernormbf16": partial(nn.LayerNorm, eps=1e-5),
    "rmsnorm": RMSNorm,
}

dtype_dict: dict[str, torch.dtype] = {
    "fp32": torch.float32,
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
}


class MotifVisionTransformer(nn.Module):
    """Vision Transformer backbone with 3D RoPE for unified image/video processing.

    Key features:
    - PatchEmbed (Conv3d) for unified image/video tokenization
    - Full 3D axial RoPE positional encoding (T/H/W)
    - CLS token + register (storage) tokens
    - LayerScale
    - MLP or SwiGLU FFN variants

    Token sequence layout: [CLS] + [Register x n_storage_tokens] + [Patch x N_total]

    Args:
        img_size: Input image size.
        patch_size: Spatial patch size.
        in_chans: Number of input channels.
        embed_dim: Embedding dimension.
        depth: Number of transformer blocks.
        num_heads: Number of attention heads.
        ffn_ratio: FFN hidden dimension ratio.
        qkv_bias: Whether to use bias in QKV.
        drop_path_rate: Stochastic depth rate.
        layerscale_init: LayerScale initial value (None to disable).
        norm_layer: Normalization layer name.
        ffn_layer: FFN layer name.
        ffn_bias: Whether to use bias in FFN.
        proj_bias: Whether to use bias in attention output projection.
        n_storage_tokens: Number of register tokens.
        mask_k_bias: Whether to mask K bias in attention.
        untie_cls_and_patch_norms: Use separate norms for CLS and patch tokens.
        untie_global_and_local_cls_norm: Use separate norm for local CLS tokens.
        device: Device for parameter allocation.
        num_frames: Number of input video frames.
        tubelet_size: Temporal patch size for Conv3d.
        pos_embed_rope_base: RoPE frequency base (100.0 for vision).
        gated_attention: Gated attention variant (None, "headwise", "elementwise").
            See https://arxiv.org/abs/2505.06708.
    """

    def __init__(
        self,
        *,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        ffn_ratio: float = 4.0,
        qkv_bias: bool = True,
        drop_path_rate: float = 0.0,
        layerscale_init: float | None = None,
        norm_layer: str = "layernorm",
        ffn_layer: str = "mlp",
        ffn_bias: bool = True,
        proj_bias: bool = True,
        n_storage_tokens: int = 0,
        mask_k_bias: bool = False,
        untie_cls_and_patch_norms: bool = False,
        untie_global_and_local_cls_norm: bool = False,
        device: Any | None = None,
        num_frames: int = 1,
        tubelet_size: int = 1,
        pos_embed_rope_base: float = 100.0,
        pos_embed_rope_rescale_coords: float | None = None,
        pos_embed_rope_shift_coords: float | None = None,
        pos_embed_rope_jitter_coords: float | None = None,
        pos_embed_rope_fhw_dim: tuple[int, int, int] | None = None,
        gated_attention: str | None = None,
        qk_norm: bool = False,
        **ignored_kwargs,
    ) -> None:
        super().__init__()
        if len(ignored_kwargs) > 0:
            logger.warning(f"Ignored kwargs: {ignored_kwargs}")

        norm_layer_cls = norm_layer_dict[norm_layer]

        self.num_features = self.embed_dim = embed_dim
        self.n_blocks = depth
        self.num_heads = num_heads
        self.patch_size = patch_size

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
            tubelet_size=tubelet_size,
            flatten_embedding=True,
        )

        self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim, device=device))
        self.n_storage_tokens = n_storage_tokens
        if self.n_storage_tokens > 0:
            self.storage_tokens = nn.Parameter(torch.empty(1, n_storage_tokens, embed_dim, device=device))

        # Convert 0.0 to None for backward compat (0.0 means disabled)
        _rescale = pos_embed_rope_rescale_coords if pos_embed_rope_rescale_coords else None
        _shift = pos_embed_rope_shift_coords if pos_embed_rope_shift_coords else None
        _jitter = pos_embed_rope_jitter_coords if pos_embed_rope_jitter_coords else None
        self.rope_embed = RopePositionEmbedding3D(
            embed_dim=embed_dim,
            num_heads=num_heads,
            fhw_dim=pos_embed_rope_fhw_dim,
            base=pos_embed_rope_base,
            rescale_coords=_rescale,
            shift_coords=_shift,
            jitter_coords=_jitter,
        )

        logger.info(f"using {ffn_layer} layer as FFN")
        ffn_layer_cls = ffn_layer_dict[ffn_layer]
        ffn_ratio_sequence = [ffn_ratio] * depth

        blocks_list = [
            SelfAttentionBlock(
                dim=embed_dim,
                num_heads=num_heads,
                ffn_ratio=ffn_ratio_sequence[i],
                qkv_bias=qkv_bias,
                proj_bias=proj_bias,
                ffn_bias=ffn_bias,
                drop_path=drop_path_rate,
                norm_layer=norm_layer_cls,
                act_layer=nn.GELU,
                ffn_layer=ffn_layer_cls,
                init_values=layerscale_init,
                mask_k_bias=mask_k_bias,
                device=device,
                gated_attention=gated_attention,
                qk_norm=qk_norm,
            )
            for i in range(depth)
        ]

        self.chunked_blocks = False
        self.blocks = nn.ModuleList(blocks_list)

        self.norm = norm_layer_cls(embed_dim)

        self.untie_cls_and_patch_norms = untie_cls_and_patch_norms
        if untie_cls_and_patch_norms:
            self.cls_norm = norm_layer_cls(embed_dim)
        else:
            self.cls_norm = None

        self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm
        if untie_global_and_local_cls_norm:
            self.local_cls_norm = norm_layer_cls(embed_dim)
        else:
            self.local_cls_norm = None
        self.head = nn.Identity()
        self.mask_token = nn.Parameter(torch.empty(1, embed_dim, device=device))

    def prepare_tokens_with_masks(
        self,
        x: Tensor,
        masks: Tensor | None = None,
    ) -> tuple[Tensor, tuple[int, int, int]]:
        """Tokenize input and assemble token sequence with CLS + register + patches.

        Args:
            x: Input tensor. Image: (B, C, H, W) or Video: (B, T, C, H, W).
            masks: Boolean mask of shape (B, N_spatial) indicating which patches to mask.

        Returns:
            Tuple of:
                - Token sequence: (B, 1 + n_storage + N_total, embed_dim)
                - Grid dimensions: (T_grid, H_grid, W_grid)
        """
        if x.ndim == 5:
            B, T, C, H, W = x.shape
            # Video: Conv3d kernel=stride=tubelet downsamples raw T to T // tubelet.
            T_grid = T // self.patch_embed.tubelet_size
        else:
            B, C, H, W = x.shape
            # Image (4D): PatchEmbed expands raw T=1 to tubelet then Conv3d(stride=tubelet)
            # produces a single temporal token (output T_out = (tubelet - tubelet)/tubelet + 1 = 1).
            # vjepa2 vision_transformer.py:171-173 passes T=1 (no division) for the same reason.
            T_grid = 1

        x = self.patch_embed(x)  # (B, N_total, D)

        # Grid dimensions for RoPE computation
        H_grid = H // self.patch_embed.patch_size[0]
        W_grid = W // self.patch_embed.patch_size[1]

        if masks is not None:
            # Expand spatial mask to spatio-temporal if needed (tube masking)
            if masks.shape[1] != x.shape[1]:
                ratio = x.shape[1] // masks.shape[1]
                masks = masks.unsqueeze(1).repeat(1, ratio, 1).flatten(1)

            # Replace masked positions with mask_token
            x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
            cls_token = self.cls_token
        else:
            # Include mask_token in computation graph even when not masking
            cls_token = self.cls_token + 0 * self.mask_token

        if self.n_storage_tokens > 0:
            storage_tokens = self.storage_tokens
        else:
            storage_tokens = torch.empty(
                1, 0, cls_token.shape[-1],
                dtype=cls_token.dtype, device=cls_token.device,
            )

        x = torch.cat(
            [
                cls_token.expand(B, -1, -1),
                storage_tokens.expand(B, -1, -1),
                x,
            ],
            dim=1,
        )

        return x, (T_grid, H_grid, W_grid)

    def forward_features_list(
        self,
        x_list: list[Tensor],
        masks_list: list[Tensor | None],
    ) -> list[dict[str, Tensor]]:
        """Forward pass for a list of inputs (multi-crop).

        Args:
            x_list: List of input tensors (global crops, local crops).
            masks_list: List of corresponding masks (None for unmasked).

        Returns:
            List of output dictionaries, one per input, containing:
                - x_norm_clstoken: Normalized CLS token (B, D)
                - x_storage_tokens: Normalized register tokens (B, n_storage, D)
                - x_norm_patchtokens: Normalized patch tokens (B, N, D)
                - x_prenorm: Pre-normalization features (B, 1+n_storage+N, D)
                - masks: Original masks
        """
        x = []
        rope_params = []
        for t_x, t_masks in zip(x_list, masks_list):
            t2_x, grid_tuple = self.prepare_tokens_with_masks(t_x, t_masks)
            x.append(t2_x)
            rope_params.append(grid_tuple)

        # Pre-compute RoPE sin/cos once — identical across all blocks.
        # Hoisting this out of the loop avoids breaking FSDP2's forward prefetch
        # chain (rope_embed is part of the outer FSDP unit, calling it between
        # block forwards disrupts the prefetch scheduling).
        if self.rope_embed is not None:
            rope_sincos = [self.rope_embed(T=t, H=h, W=w) for t, h, w in rope_params]
        else:
            rope_sincos = [None for _ in rope_params]

        for _, blk in enumerate(self.blocks):
            x = blk(x, rope_sincos)

        all_x = x
        output = []
        for idx, (x, masks) in enumerate(zip(all_x, masks_list)):
            if self.untie_cls_and_patch_norms or self.untie_global_and_local_cls_norm:
                if self.untie_global_and_local_cls_norm and self.training and idx == 1:
                    x_norm_cls_reg = self.local_cls_norm(x[:, : self.n_storage_tokens + 1])
                elif self.untie_cls_and_patch_norms:
                    x_norm_cls_reg = self.cls_norm(x[:, : self.n_storage_tokens + 1])
                else:
                    x_norm_cls_reg = self.norm(x[:, : self.n_storage_tokens + 1])
                x_norm_patch = self.norm(x[:, self.n_storage_tokens + 1 :])
            else:
                x_norm = self.norm(x)
                x_norm_cls_reg = x_norm[:, : self.n_storage_tokens + 1]
                x_norm_patch = x_norm[:, self.n_storage_tokens + 1 :]
            output.append(
                {
                    "x_norm_clstoken": x_norm_cls_reg[:, 0],
                    "x_storage_tokens": x_norm_cls_reg[:, 1:],
                    "x_norm_patchtokens": x_norm_patch,
                    "x_prenorm": x,
                    "masks": masks,
                }
            )
        return output

    def forward_features(
        self,
        x: Tensor | list[Tensor],
        masks: Tensor | list[Tensor | None] | None = None,
    ) -> dict[str, Tensor] | list[dict[str, Tensor]]:
        """Forward pass for single or multiple inputs.

        Args:
            x: Single tensor or list of tensors.
            masks: Single mask or list of masks.

        Returns:
            Output dict (single input) or list of output dicts (multiple inputs).
        """
        if isinstance(x, torch.Tensor):
            return self.forward_features_list([x], [masks])[0]
        else:
            return self.forward_features_list(x, masks)

    def _get_intermediate_layers_not_chunked(
        self,
        x: Tensor,
        n: int | list[int] = 1,
    ) -> list[Tensor]:
        """Run forward pass and collect intermediate block outputs.

        Args:
            x: Input tensor (B, C, H, W) or (B, T, C, H, W).
            n: If int, return last n layers. If list, return specific layer indices.

        Returns:
            List of intermediate outputs, each (B, 1+n_storage+N, D).
        """
        x, grid_tuple = self.prepare_tokens_with_masks(x, masks=None)
        T, H, W = grid_tuple

        output, total_block_len = [], len(self.blocks)
        blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n

        if self.rope_embed is not None:
            rope_sincos = self.rope_embed(T=T, H=H, W=W)
        else:
            rope_sincos = None

        for i, blk in enumerate(self.blocks):
            x = blk([x], [rope_sincos])[0]
            if i in blocks_to_take:
                output.append(x)

        assert len(output) == len(blocks_to_take), (
            f"only {len(output)} / {len(blocks_to_take)} blocks found"
        )
        return output

    def get_intermediate_layers(
        self,
        x: Tensor,
        n: int | list[int] = 1,
        reshape: bool = False,
        return_class_token: bool = False,
        norm: bool = True,
    ) -> tuple[Tensor, ...]:
        """Extract intermediate layer outputs for downstream evaluation.

        This method is critical for dense prediction tasks (segmentation, depth)
        that need multi-scale features from different transformer blocks.

        Args:
            x: Input image (B, C, H, W) or video (B, T, C, H, W).
            n: If int, return outputs from last n layers.
               If list[int], return outputs from specific layer indices.
            reshape: If True, reshape patch tokens to spatial form (B, D, H_grid, W_grid).
            return_class_token: If True, return (patch_tokens, cls_token) tuples.
            norm: If True, apply final LayerNorm to outputs.

        Returns:
            If return_class_token is False:
                Tuple of patch token tensors, one per requested layer.
                Each tensor is (B, N, D) or (B, D, H_grid, W_grid) if reshape=True.
            If return_class_token is True:
                Tuple of (patch_tokens, cls_token) pairs.
        """
        # Determine spatial dims for reshape
        if x.ndim == 5:
            B, T_in, C, H, W = x.shape
        else:
            B, C, H, W = x.shape
            T_in = 1

        outputs = self._get_intermediate_layers_not_chunked(x, n)

        if norm:
            outputs_normed = []
            for out in outputs:
                if self.untie_cls_and_patch_norms:
                    x_norm_cls_reg = self.cls_norm(out[:, : self.n_storage_tokens + 1])
                    x_norm_patch = self.norm(out[:, self.n_storage_tokens + 1 :])
                    outputs_normed.append(torch.cat((x_norm_cls_reg, x_norm_patch), dim=1))
                else:
                    outputs_normed.append(self.norm(out))
            outputs = outputs_normed

        class_tokens = [out[:, 0] for out in outputs]
        outputs = [out[:, self.n_storage_tokens + 1 :] for out in outputs]

        if reshape:
            # Image (T_in=1): PatchEmbed expands to tubelet then Conv3d(stride=tubelet) → T_out=1.
            # Video: Conv3d downsamples T_in → T_in // tubelet. Matches prepare_tokens_with_masks
            # and vjepa2 vision_transformer.py:171-177.
            if x.ndim == 5:
                T_grid = T_in // self.patch_embed.tubelet_size
            else:
                T_grid = 1
            H_grid = H // self.patch_size
            W_grid = W // self.patch_size
            if T_grid > 1:
                # Video: reshape to (B, D, T_grid, H_grid, W_grid)
                outputs = [
                    out.reshape(B, T_grid, H_grid, W_grid, -1).permute(0, 4, 1, 2, 3).contiguous()
                    for out in outputs
                ]
            else:
                # Image: reshape to (B, D, H_grid, W_grid)
                outputs = [
                    out.reshape(B, H_grid, W_grid, -1).permute(0, 3, 1, 2).contiguous()
                    for out in outputs
                ]

        if return_class_token:
            return tuple(zip(outputs, class_tokens))
        return tuple(outputs)

    def forward(
        self,
        *args,
        is_training: bool = False,
        **kwargs,
    ) -> dict[str, Tensor] | list[dict[str, Tensor]] | Tensor:
        """High-level forward: training returns feature dict, inference returns CLS logits.

        Args:
            is_training: If True, return full feature dictionary.

        Returns:
            Feature dict(s) if training, CLS token logits if inference.
        """
        ret = self.forward_features(*args, **kwargs)
        if is_training:
            return ret
        else:
            return self.head(ret["x_norm_clstoken"])


# ============================================================================
# HuggingFace transformers wrapper (inference-only)
# ============================================================================
class MotifVisionConfig(PretrainedConfig):
    """Config for the Motif Vision Encoder backbone (image + video)."""

    model_type = "motif_vision"

    def __init__(
        self,
        img_size: int = 224,
        patch_size: int = 16,
        in_chans: int = 3,
        embed_dim: int = 4096,
        depth: int = 40,
        num_heads: int = 32,
        ffn_ratio: float = 3.0,
        qkv_bias: bool = False,
        drop_path_rate: float = 0.0,
        layerscale_init: float | None = 1.0e-5,
        norm_layer: str = "layernormbf16",
        ffn_layer: str = "swiglu64",
        ffn_bias: bool = True,
        proj_bias: bool = True,
        n_storage_tokens: int = 4,
        mask_k_bias: bool = True,
        untie_cls_and_patch_norms: bool = False,
        untie_global_and_local_cls_norm: bool = True,
        num_frames: int = 1,
        tubelet_size: int = 2,
        pos_embed_rope_base: float = 100.0,
        pos_embed_rope_rescale_coords: float | None = 2.0,
        gated_attention: str | None = "elementwise",
        qk_norm: bool = True,
        **kwargs,
    ):
        self.img_size = img_size
        self.patch_size = patch_size
        self.in_chans = in_chans
        self.embed_dim = embed_dim
        self.depth = depth
        self.num_heads = num_heads
        self.ffn_ratio = ffn_ratio
        self.qkv_bias = qkv_bias
        self.drop_path_rate = drop_path_rate
        self.layerscale_init = layerscale_init
        self.norm_layer = norm_layer
        self.ffn_layer = ffn_layer
        self.ffn_bias = ffn_bias
        self.proj_bias = proj_bias
        self.n_storage_tokens = n_storage_tokens
        self.mask_k_bias = mask_k_bias
        self.untie_cls_and_patch_norms = untie_cls_and_patch_norms
        self.untie_global_and_local_cls_norm = untie_global_and_local_cls_norm
        self.num_frames = num_frames
        self.tubelet_size = tubelet_size
        self.pos_embed_rope_base = pos_embed_rope_base
        self.pos_embed_rope_rescale_coords = pos_embed_rope_rescale_coords
        self.gated_attention = gated_attention
        self.qk_norm = qk_norm
        super().__init__(**kwargs)


class MotifVisionModel(PreTrainedModel):
    """Motif Vision Encoder for HF `AutoModel` (inference). Returns dense + CLS features."""

    config_class = MotifVisionConfig
    base_model_prefix = "motif"
    main_input_name = "pixel_values"
    _no_split_modules = ["SelfAttentionBlock"]
    supports_gradient_checkpointing = False

    def __init__(self, config: MotifVisionConfig):
        super().__init__(config)
        self.backbone = MotifVisionTransformer(
            img_size=config.img_size,
            patch_size=config.patch_size,
            in_chans=config.in_chans,
            embed_dim=config.embed_dim,
            depth=config.depth,
            num_heads=config.num_heads,
            ffn_ratio=config.ffn_ratio,
            qkv_bias=config.qkv_bias,
            drop_path_rate=config.drop_path_rate,
            layerscale_init=config.layerscale_init,
            norm_layer=config.norm_layer,
            ffn_layer=config.ffn_layer,
            ffn_bias=config.ffn_bias,
            proj_bias=config.proj_bias,
            n_storage_tokens=config.n_storage_tokens,
            mask_k_bias=config.mask_k_bias,
            untie_cls_and_patch_norms=config.untie_cls_and_patch_norms,
            untie_global_and_local_cls_norm=config.untie_global_and_local_cls_norm,
            num_frames=config.num_frames,
            tubelet_size=config.tubelet_size,
            pos_embed_rope_base=config.pos_embed_rope_base,
            pos_embed_rope_rescale_coords=config.pos_embed_rope_rescale_coords,
            gated_attention=config.gated_attention,
            qk_norm=config.qk_norm,
        )
        self.post_init()

    @torch.no_grad()
    def forward(self, pixel_values: Tensor, return_dict: bool = True, **kwargs):
        """pixel_values: image (B,3,H,W) or video (B,T,3,H,W). H,W multiples of patch_size."""
        out = self.backbone.forward_features(pixel_values)
        cls = out["x_norm_clstoken"]
        reg = out["x_storage_tokens"]
        patch = out["x_norm_patchtokens"]
        last_hidden = torch.cat([cls.unsqueeze(1), reg, patch], dim=1)
        if not return_dict:
            return (last_hidden, cls)
        return BaseModelOutputWithPooling(last_hidden_state=last_hidden, pooler_output=cls)


AutoConfig_registered = False
try:
    from transformers import AutoConfig, AutoModel
    AutoConfig.register("motif_vision", MotifVisionConfig)
    AutoModel.register(MotifVisionConfig, MotifVisionModel)
    AutoConfig_registered = True
except Exception:
    pass