# 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