vit.py

# MIT License
# 
# Copyright (c) 2026 audio-embeddings contributors
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from collections.abc import Mapping
from typing import Any, Optional, Tuple

import torch
import torch.nn as nn
from timm.layers import build_sincos2d_pos_embed

from .rope import RotaryEmbedding1D
from .rope import RotaryEmbedding2D
from .transformer import build_norm_layer
from .transformer import FullAttentionResidual
from .transformer import RoPEBlock


def vit_config_with_patch_geometry(
    config: Mapping[str, Any],
    *,
    img_size: tuple[int, int],
    patch_size: tuple[int, int],
) -> dict[str, Any]:
    """Return a ViT config aligned with the patch embedding's actual geometry."""
    resolved = dict(config)
    resolved["img_size"] = tuple(img_size)
    resolved["patch_size"] = tuple(patch_size)
    resolved.setdefault("rope_mode", "auto")
    return resolved


class ViT(nn.Module):
    """
    Vision Transformer with support for RoPE and 2D positional embeddings.

    Args:
        embed_dim (int): Embedding dimension.
        depth (int): Number of transformer blocks.
        num_heads (int): Number of attention heads.
        mlp_ratio (float): Ratio of MLP hidden dim to embedding dim.
        qkv_bias (bool): Enable bias for QKV projections.
        drop_rate (float): Dropout rate.
        attn_drop_rate (float): Attention dropout rate.
        drop_path_rate (float): Stochastic depth rate.
        norm_layer (nn.Module): Normalization layer.
        norm_eps (float | None): Explicit normalization epsilon; None preserves
            the normalization implementation's default.
        act_layer (nn.Module): Activation layer.
        num_patches (int): Total number of patches (used for learnable/sincos pos embed).
        img_size (tuple[int, int]): Input image size (H, W).
        patch_size (tuple[int, int]): Patch size (H, W).
        pos_embed_type (str): Type of positional embedding ("rope", "sincos", "learnable").
        rope_mode (str | None): RoPE axes ("2d" or temporal "1d"). ``None``
            or ``"auto"`` selects 1D when the patch height equals the image height,
            otherwise 2D.
    """

    def __init__(
        self,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        mlp_type: str = "gelu_mlp",
        qkv_bias: bool = True,
        proj_bias: bool = True,
        mlp_bias: bool = True,
        qk_norm: bool = False,
        qk_norm_type: str = "layernorm",
        drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        drop_path_rate: float = 0.0,
        norm_layer: nn.Module | None = None,
        norm_type: str = "layernorm",
        act_layer: nn.Module = nn.GELU,
        num_patches: int = 128,
        img_size: tuple[int, int] = (128, 256),
        patch_size: tuple[int, int] = (16, 16),
        pos_embed_type: str = "rope",
        rope_mode: str | None = "2d",
        residual_type: str = "standard",
        norm_eps: float | None = None,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_patches = num_patches
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.pos_embed_type = pos_embed_type
        requested_rope_mode = (
            "auto" if rope_mode is None else rope_mode.strip().lower().replace("-", "_")
        )
        if requested_rope_mode == "auto":
            self.rope_mode = "1d" if patch_size[0] == img_size[0] else "2d"
        else:
            self.rope_mode = requested_rope_mode
        self.norm_type = norm_type
        self.mlp_type = mlp_type
        self.residual_type = residual_type.strip().lower().replace("-", "_")
        if self.residual_type not in {"standard", "full_attnres"}:
            raise ValueError(
                f"Unknown residual_type={residual_type!r}; expected 'standard' "
                "or 'full_attnres'"
            )

        # Positional Embeddings
        if pos_embed_type == "rope":
            head_dim = embed_dim // num_heads
            if self.rope_mode == "2d":
                self.rope = RotaryEmbedding2D(dim=head_dim, max_res=self.grid_size)
            elif self.rope_mode == "1d":
                self.rope = RotaryEmbedding1D(
                    dim=head_dim,
                    max_seq_len=self.grid_size[1],
                )
            else:
                raise ValueError(
                    f"Unknown rope_mode: {rope_mode!r}; expected 'auto', '1d', or '2d'"
                )
            self.pos_embed = None
        elif pos_embed_type == "sincos":
            self.rope = None
            # build_sincos2d_pos_embed(feat_shape, dim, ...)
            # We assume grid_size matches num_patches
            pos_embed = build_sincos2d_pos_embed(self.grid_size, embed_dim)
            self.register_buffer("pos_embed", pos_embed.unsqueeze(0))  # [1, N, D]
        elif pos_embed_type == "learnable":
            self.rope = None
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
            nn.init.trunc_normal_(self.pos_embed, std=0.02)
        else:
            raise ValueError(f"Unknown pos_embed_type: {pos_embed_type}")

        # Stochastic Depth
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]

        self.blocks = nn.ModuleList(
            [
                RoPEBlock(
                    dim=embed_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    mlp_type=mlp_type,
                    qkv_bias=qkv_bias,
                    proj_bias=proj_bias,
                    mlp_bias=mlp_bias,
                    qk_norm=qk_norm,
                    qk_norm_type=qk_norm_type,
                    proj_drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[i],
                    norm_type=norm_type,
                    norm_layer=norm_layer,
                    norm_eps=norm_eps,
                    act_layer=act_layer,
                    rope=self.rope,
                    residual_type=self.residual_type,
                )
                for i in range(depth)
            ]
        )

        self.norm = build_norm_layer(
            dim=embed_dim,
            norm_type=norm_type,
            norm_layer=norm_layer,
            norm_eps=norm_eps,
        )
        self.output_residual = (
            FullAttentionResidual(embed_dim)
            if self.residual_type == "full_attnres"
            else None
        )

        self.apply(self._init_weights)

    def _init_weights(self, m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.RMSNorm):
            nn.init.constant_(m.weight, 1.0)

    def forward(
        self,
        x: torch.Tensor,
        pos_ids: Optional[torch.Tensor] = None,
        add_pos_embed: bool = True,
        grid_size: Optional[Tuple[int, int]] = None,
    ) -> torch.Tensor:
        """
        Forward pass.

        Args:
            x (torch.Tensor): Input tensor [B, N, D].
            pos_ids (Optional[torch.Tensor]): Positional indices [B, N] or [N].
            add_pos_embed (bool): Whether to add positional embeddings (for non-RoPE).
            grid_size (Optional[Tuple[int, int]]): Grid size for RoPE/PosEmbed.

        Returns:
            torch.Tensor: Output tensor [B, N, D].
        """
        # Determine grid size
        if grid_size is None:
            if pos_ids is None:
                # Infer from x assuming full sequence
                B, N, D = x.shape
                H_grid = self.grid_size[0]
                W_grid = N // H_grid
                current_grid_size = (H_grid, W_grid)
            else:
                # Cannot infer, use default (might be wrong if variable length)
                current_grid_size = self.grid_size
        else:
            current_grid_size = grid_size

        if self.pos_embed_type != "rope" and add_pos_embed:
            if pos_ids is not None:
                # Select positional embeddings
                if pos_ids.ndim == 1:
                    # Shared pos_ids across batch
                    pos_embed = self.pos_embed[:, pos_ids, :]  # [1, N_subset, D]
                else:
                    # Different pos_ids per sample
                    pos_embed = self.pos_embed.expand(x.shape[0], -1, -1)
                    pos_embed = torch.gather(
                        pos_embed,
                        1,
                        pos_ids.unsqueeze(-1).expand(-1, -1, self.embed_dim),
                    )
                x = x + pos_embed
            else:
                # Assume full sequence
                if x.shape[1] == self.num_patches:
                    x = x + self.pos_embed
                elif (
                    self.pos_embed is not None and x.shape[1] <= self.pos_embed.shape[1]
                ):
                    x = x + self.pos_embed[:, : x.shape[1], :]

        # For RoPE, we need pos_ids. If not provided, generate them.
        if self.pos_embed_type == "rope" and pos_ids is None:
            device = x.device
            # We need to generate pos_ids for the current grid
            # If we inferred current_grid_size, we should use it.
            # pos_ids should be 0..N-1
            B, N, D = x.shape
            pos_ids = torch.arange(N, device=device)

        if self.residual_type == "full_attnres":
            values = [x]
            for block in self.blocks:
                if block.attention_residual is None or block.mlp_residual is None:
                    raise RuntimeError(
                        "Full AttnRes block is missing depth aggregators"
                    )
                attention_input = block.attention_residual(values)
                values.append(
                    block.attention_output(
                        attention_input,
                        pos_ids=pos_ids,
                        grid_size=current_grid_size,
                    )
                )
                mlp_input = block.mlp_residual(values)
                values.append(block.mlp_output(mlp_input))

            if self.output_residual is None:
                raise RuntimeError("Full AttnRes ViT is missing its output aggregator")
            x = self.output_residual(values)
        else:
            for block in self.blocks:
                x = block(x, pos_ids=pos_ids, grid_size=current_grid_size)

        x = self.norm(x)
        return x