Upload Mage-ViT: standalone codec-native visual encoder (ViT pre-training only)

README.md

@@ -1,3 +1,132 @@
 ---
 license: apache-2.0
+library_name: transformers
+pipeline_tag: image-feature-extraction
+tags:
+- mage-vl
+- vision-encoder
+- codec-vit
+- video-understanding
 ---
+
+<h1 align="center">Mage-ViT<br><span style="font-size: 0.55em; font-weight: normal;">A Codec-Native Visual Encoder Trained from Scratch</span></h1>
+
+<p align="center">
+  <a href="https://microsoft.github.io/Mage"><img alt="Project Page" src="https://img.shields.io/badge/%F0%9F%8C%90-Project%20Page-blue" height="22" /></a>
+  <a href="https://github.com/microsoft/Mage"><img src="https://img.shields.io/badge/Code-GitHub-181717?logo=github" alt="GitHub" height="22"></a>
+  <a href="https://huggingface.co/microsoft/Mage-VL"><img alt="Mage-VL" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--VL-yellow" height="22" /></a>
+  <a href="https://huggingface.co/microsoft/Mage-ViT"><img alt="Mage-ViT" src="https://img.shields.io/badge/%F0%9F%A4%97-Mage--ViT-yellow" height="22" /></a>
+  <a href="https://www.apache.org/licenses/LICENSE-2.0"><img src="https://img.shields.io/badge/License-Apache%202.0-green" alt="License: Apache 2.0" height="22"></a>
+</p>
+
+---
+
+**Mage-ViT** is the visual encoder at the core of **[Mage-VL](https://huggingface.co/microsoft/Mage-VL)**. It is a *Codec-ViT* built primarily for video, where a single image is simply the degenerate one-frame case. Mage-ViT follows a **codec-aligned sparsity** principle: visual tokens should be spent where a video codec spends bits, since codec bit-allocation is a natural proxy for spatio-temporal importance. On a shared `16×16` patch grid it keeps every anchor (I-frame) patch and only the motion-salient predicted (P-frame) patches, while a shared **3D rotary position encoding** preserves spatio-temporal structure even after large fractions of the grid are dropped.
+
+This repository is the **ViT-pre-trained checkpoint only** — it has not gone through the joint VLM training with the language model. Use it as a data-efficient, codec-native visual encoder, or as a drop-in ViT for your own multimodal training.
+
+## ✨ Highlights
+
+- **Codec-driven patchifier.** Patches are selected by a per-patch importance map derived from the codec — motion vectors + P-frame residual energy for HEVC/H.265, or the learned rate map of the neural codec DCVC-RT. For a 64-frame clip it keeps all I-frame patches plus the top-*k* P-frame patches within a **4096-token budget (~75% token reduction)**. Chunk-wise and collage patchification are also supported.
+- **Codec-agnostic.** The same interface accepts a traditional codec (HEVC/H.265) or a neural codec (DCVC-RT) with no architecture or retraining change.
+- **Trained from scratch.** No billion-scale image-text ViT initialization — Mage-ViT is optimized with a large-scale **cluster-discrimination** objective on **~100M unlabeled images/videos**.
+
+## 🏗️ Architecture
+
+A **24-layer pre-norm Vision Transformer** trunk (hidden size `1024`, `16` attention heads, GELU MLP at 4× expansion) processes the variable-length token sequence produced by the codec-driven patchifier, followed by a multi-head attention pooling head. The standalone encoder in this repo consumes `pixel_values` (and optional `patch_positions`); codec-driven patch **selection** is applied upstream in the Mage-VL data pipeline.
+
+Pre-training is a two-stage, from-scratch recipe in bf16: **(1)** variable-resolution image pre-training (224–448), then **(2)** joint image + video pre-training (video at resolution 256, 64 frames per clip, 4096-token budget).
+
+## 📦 Installation
+
+Mage-ViT follows the Mage-VL runtime:
+
+```bash
+pip install "transformers>=5.7" torch torchvision pillow
+# optional, for the fastest GPU attention path:
+# pip install flash-attn --no-build-isolation
+```
+
+Attention is dispatched internally to `sdpa` (default) → `flash_attention_2` → `eager`, so **flash-attn is optional** and the model runs on CPU or GPU out of the box.
+
+## 🚀 Encoding an image
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoImageProcessor
+
+model = AutoModel.from_pretrained(
+    "microsoft/Mage-ViT", trust_remote_code=True
+).to("cuda").eval()
+processor = AutoImageProcessor.from_pretrained(
+    "microsoft/Mage-ViT", trust_remote_code=True
+)
+
+image = Image.open("your_image.jpg")
+pixel_values = processor(images=image, return_tensors="pt")["pixel_values"].to("cuda")
+
+with torch.no_grad():
+    out = model(pixel_values)              # pixel_values: [B, 3, H, W]
+
+patch_features = out.last_hidden_state     # [B, num_patches, 1024]
+pooled_feature = out.pooler_output         # [B, 1024]
+```
+
+Pass `attn_implementation="flash_attention_2"` (or `"sdpa"` / `"eager"`) to `from_pretrained` to pick the attention backend.
+
+## 🎞️ Encoding a video
+
+Stack the preprocessed frames into `[B, C, T, H, W]` and pass a `patch_positions` tensor of shape `[B, T * tokens_per_frame, 3]` giving the `(t, h, w)` grid coordinate of every patch:
+
+```python
+import torch
+from PIL import Image
+
+PATCH = 16
+
+def build_patch_positions(num_frames, target_frames, grid_h, grid_w, device="cuda"):
+    # temporal index for each frame, spread across the target timeline
+    t = torch.linspace(0, target_frames - 1, num_frames, device=device).long()
+    t = t.repeat_interleave(grid_h * grid_w)                    # [T * H * W]
+    h = torch.arange(grid_h, device=device).repeat_interleave(grid_w).repeat(num_frames)
+    w = torch.arange(grid_w, device=device).repeat(grid_h).repeat(num_frames)
+    return torch.stack([t, h, w], dim=-1).unsqueeze(0)          # [1, T*H*W, 3]
+
+frames = [Image.open(f"frame_{i}.jpg") for i in range(16)]      # your sampled frames
+pv = processor(images=frames, return_tensors="pt")["pixel_values"]   # [T, C, H, W]
+video = pv.unsqueeze(0).permute(0, 2, 1, 3, 4).to("cuda")            # [1, C, T, H, W]
+
+gh, gw = video.shape[-2] // PATCH, video.shape[-1] // PATCH
+patch_positions = build_patch_positions(num_frames=16, target_frames=64,
+                                        grid_h=gh, grid_w=gw)
+
+with torch.no_grad():
+    out = model(video, patch_positions=patch_positions)
+```
+
+## 📤 Outputs
+
+| Field | Shape | Description |
+| ----- | ----- | ----------- |
+| `last_hidden_state` | `[B, num_patches, 1024]` | per-patch features after the final layer norm |
+| `pooler_output`     | `[B, 1024]`              | global feature from the multi-head attention pooling head |
+
+## 📋 Specifications
+
+| | |
+| --- | --- |
+| Architecture | Codec-ViT (pre-norm ViT, SigLIP-style MLP) |
+| Parameters | ~316M |
+| Hidden size / MLP | 1024 / 4096 |
+| Layers / heads | 24 / 16 |
+| Patch size | 16 |
+| Position encoding | shared 3D rotary (4:6:6 split over T:H:W) |
+| Pooling | learned-probe multi-head attention head |
+| Pre-training resolution | images 224–448 (variable), video 256 |
+| Weights dtype | bfloat16 |
+| License | Apache-2.0 |
+
+## 📄 License
+
+Released under the [Apache-2.0 License](https://www.apache.org/licenses/LICENSE-2.0).

__init__.py

@@ -0,0 +1,2 @@
+from .configuration_mage_vit import MageViTConfig
+from .modeling_mage_vit import MageViTModel

config.json

@@ -0,0 +1,27 @@
+{
+  "architectures": [
+    "MageViTModel"
+  ],
+  "attention_dropout": 0.0,
+  "dtype": "bfloat16",
+  "hidden_act": "gelu",
+  "hidden_size": 1024,
+  "image_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-06,
+  "layer_norm_type": "layer_norm",
+  "model_type": "mage_vit",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 16,
+  "rope_theta": 10000.0,
+  "rope_temporal_size": 64,
+  "transformers_version": "5.7.0",
+  "use_head": true,
+  "auto_map": {
+    "AutoConfig": "configuration_mage_vit.MageViTConfig",
+    "AutoModel": "modeling_mage_vit.MageViTModel"
+  }
+}

configuration_mage_vit.py

@@ -0,0 +1,110 @@
+try:  # transformers >= 5
+    from transformers.configuration_utils import PreTrainedConfig
+except ImportError:  # transformers < 5
+    from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
+
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class MageViTConfig(PreTrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`MageViTModel`]. It is used to instantiate a
+    Mage-ViT model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the Mage-ViT architecture.
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+    Args:
+        hidden_size (`int`, *optional*, defaults to 1024):
+            Dimensionality of the encoder layers and the pooler layer.
+        intermediate_size (`int`, *optional*, defaults to 4096):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        num_hidden_layers (`int`, *optional*, defaults to 24):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 16):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_channels (`int`, *optional*, defaults to 3):
+            The number of input channels.
+        image_size (`int`, *optional*, defaults to 256):
+            The size (resolution) of each image.
+        patch_size (`int`, *optional*, defaults to 16):
+            The size (resolution) of each patch.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        layer_norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+            The type of layer normalization to use. Supported values: `"layer_norm"`, `"rms_norm"`.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rope_theta (`float`, *optional*, defaults to 10000.0):
+            The base period of the RoPE embeddings.
+        rope_temporal_size (`int` or `None`, *optional*, defaults to 64):
+            Fixed temporal grid size used to build RoPE positions when a 5D video tensor is passed without explicit
+            `patch_positions`. Set to `None` to use the actual number of frames.
+        use_head (`bool`, *optional*, defaults to `True`):
+            Whether to use the multi-head attention pooling head.
+
+    Example:
+
+    ```python
+    >>> from configuration_mage_vit import MageViTConfig
+    >>> from modeling_mage_vit import MageViTModel
+
+    >>> # Initializing a Mage-ViT configuration
+    >>> configuration = MageViTConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = MageViTModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "mage_vit"
+
+    def __init__(
+        self,
+        hidden_size=1024,
+        intermediate_size=4096,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_channels=3,
+        image_size=256,
+        patch_size=16,
+        hidden_act="gelu",
+        layer_norm_eps=1e-6,
+        layer_norm_type="layer_norm",
+        attention_dropout=0.0,
+        initializer_range=0.02,
+        rope_theta=10000.0,
+        rope_temporal_size=64,
+        use_head=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.hidden_act = hidden_act
+        self.layer_norm_eps = layer_norm_eps
+        self.layer_norm_type = layer_norm_type
+        self.attention_dropout = attention_dropout
+        self.initializer_range = initializer_range
+        self.rope_theta = rope_theta
+        self.rope_temporal_size = rope_temporal_size  # None=use actual frames, int=fixed size (legacy: 64)
+        self.use_head = use_head
+
+
+__all__ = ["MageViTConfig"]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7848ad33832ef5a49c58c0d67a83d9269ace91349585e9a797364f8a0c535799
+size 631420552

modeling_mage_vit.py

@@ -0,0 +1,608 @@
+"""
+Mage-ViT — the standalone vision encoder of the Mage-VL family.
+
+The implementation is adapted from the vision tower of Mage-VL (``modeling_mage_vl.py``):
+
+    * fused ``qkv`` / ``proj`` self-attention,
+    * 3D (T,H,W) rotary position embeddings with a 4:6:6 split (``VisionRotaryEmbedding``),
+    * SigLIP-style MLP blocks and a multi-head attention pooling head.
+
+It is written to be transformers-version agnostic (works with both ``transformers>=5``
+and ``transformers==4.57.x``): attention is dispatched internally across
+``sdpa`` / ``flash_attention_2`` / ``eager`` without relying on any transformers-5-only
+symbols, and it does not require ``flash-attn`` to be installed (``sdpa`` is the default).
+"""
+
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.siglip.modeling_siglip import SiglipMLP
+from transformers.utils import logging
+
+from .configuration_mage_vit import MageViTConfig
+
+
+try:
+    from flash_attn import flash_attn_func
+
+    _flash_attn_available = True
+except ImportError:
+    _flash_attn_available = False
+
+logger = logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helper Functions & Layers
+# ---------------------------------------------------------------------------
+
+
+def get_norm_layer(config):
+    if config.layer_norm_type == "rms_norm":
+        return nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+    else:
+        return nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+
+def rotate_half(x):
+    """
+    Interleaved rotation matching the training-time implementation.
+    (x1, x2, x3, x4) -> (-x2, x1, -x4, x3)
+    """
+    x_even = x[..., ::2]
+    x_odd = x[..., 1::2]
+    return torch.stack((-x_odd, x_even), dim=-1).flatten(-2)
+
+
+def apply_rotary_pos_emb(q, k, freqs):
+    # q, k: (B, H, L, D); freqs: (B, L, D) or (1, L, D)
+    orig_q_dtype = q.dtype
+    orig_k_dtype = k.dtype
+    q, k = q.float(), k.float()
+    cos = freqs.cos().unsqueeze(1).float()  # (B, 1, L, D)
+    sin = freqs.sin().unsqueeze(1).float()
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed.to(orig_q_dtype), k_embed.to(orig_k_dtype)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+):
+    """Eager attention; query/key/value are expected as ``(B, H, L, D)``."""
+    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+    attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = F.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()  # (B, L, H, D)
+    return attn_output, attn_weights
+
+
+class VisionRotaryEmbedding(nn.Module):
+    """
+    3D (T,H,W) Rotary frequency constructor with a 4:6:6 split for T:H:W.
+    """
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        head_dim = config.hidden_size // config.num_attention_heads
+        base = config.rope_theta
+
+        assert head_dim % 2 == 0, "head_dim must be even for rotary."
+        assert head_dim % 16 == 0, "head_dim must be divisible by 16."
+        half = head_dim // 2
+        assert half % 16 == 0, "head_dim//2 must also be divisible by 16 to split into 4:6:6."
+
+        self.head_dim = head_dim
+        self.half = half
+        self.base = base
+
+        unit = half // 16
+        self.t_size = 4 * unit
+        self.h_size = 6 * unit
+        self.w_size = 6 * unit
+
+        self.register_buffer(
+            "inv_freq_t",
+            1.0 / (base ** (torch.arange(self.t_size, dtype=torch.float32) / self.t_size)),
+            persistent=False,
+        )
+        self.register_buffer(
+            "inv_freq_h",
+            1.0 / (base ** (torch.arange(self.h_size, dtype=torch.float32) / self.h_size)),
+            persistent=False,
+        )
+        self.register_buffer(
+            "inv_freq_w",
+            1.0 / (base ** (torch.arange(self.w_size, dtype=torch.float32) / self.w_size)),
+            persistent=False,
+        )
+
+    def forward_with_thw(self, t: int, h: int, w: int, device=None) -> torch.Tensor:
+        """Build RoPE frequencies for a full (t, h, w) patch grid -> [t*h*w, half]."""
+        if device is None:
+            device = self.inv_freq_t.device
+
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        ft = torch.outer(torch.arange(t, device=device, dtype=torch.float32), inv_t)
+        fh = torch.outer(torch.arange(h, device=device, dtype=torch.float32), inv_h)
+        fw = torch.outer(torch.arange(w, device=device, dtype=torch.float32), inv_w)
+
+        t_ids = torch.arange(t, device=device).repeat_interleave(h * w)
+        h_ids = torch.arange(h, device=device).repeat_interleave(w).repeat(t)
+        w_ids = torch.arange(w, device=device).repeat(h).repeat(t)
+
+        return torch.cat([ft[t_ids], fh[h_ids], fw[w_ids]], dim=-1)
+
+    def forward_from_positions(self, patch_positions: torch.Tensor) -> torch.Tensor:
+        """
+        Build RoPE frequencies from explicit patch positions.
+
+        Args:
+            patch_positions: [seq_len, 3] or [batch_size, seq_len, 3] with [t, h, w] per patch.
+
+        Returns:
+            freqs with a leading batch dim: [batch_size, seq_len, half].
+        """
+        if patch_positions.dim() == 2:
+            patch_positions = patch_positions.unsqueeze(0)
+
+        device = patch_positions.device
+        inv_t = self.inv_freq_t.to(device=device)
+        inv_h = self.inv_freq_h.to(device=device)
+        inv_w = self.inv_freq_w.to(device=device)
+
+        t_pos = patch_positions[..., 0].float()  # [B, L]
+        h_pos = patch_positions[..., 1].float()
+        w_pos = patch_positions[..., 2].float()
+
+        ft = torch.einsum("bs,d->bsd", t_pos, inv_t)
+        fh = torch.einsum("bs,d->bsd", h_pos, inv_h)
+        fw = torch.einsum("bs,d->bsd", w_pos, inv_w)
+
+        return torch.cat([ft, fh, fw], dim=-1)
+
+
+class Siglip2MultiheadAttentionPoolingHead(nn.Module):
+    """Multi-Head Attention Pooling with a learned probe (PMA-style)."""
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+        self.attention = nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.mlp = SiglipMLP(config)
+
+    def forward(self, hidden_states):
+        batch_size = hidden_states.shape[0]
+        probe = self.probe.repeat(batch_size, 1, 1)
+
+        attn_output, _ = self.attention(probe, hidden_states, hidden_states)
+
+        residual = attn_output
+        attn_output = self.norm(attn_output)
+        attn_output = residual + self.mlp(attn_output)
+
+        return attn_output[:, 0]
+
+
+# ---------------------------------------------------------------------------
+# Modeling Components
+# ---------------------------------------------------------------------------
+
+
+class MageViTEmbeddings(nn.Module):
+    """Conv2d patch embedding for pixel-value inputs (4D image or 5D video)."""
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=config.num_channels,
+            out_channels=self.embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+            bias=False,
+        )
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        # Handle 4D (B, C, H, W) or 5D (B, C, T, H, W) inputs
+        if pixel_values.dim() == 4:
+            pixel_values = pixel_values.unsqueeze(2)  # (B, C, 1, H, W)
+
+        batch_size, channels, t_frames, height, width = pixel_values.shape
+        target_dtype = self.patch_embedding.weight.dtype
+
+        # Merge time into batch for Conv2d
+        x_2d = pixel_values.permute(0, 2, 1, 3, 4).reshape(batch_size * t_frames, channels, height, width)
+
+        embeddings = self.patch_embedding(x_2d.to(dtype=target_dtype))  # (B*T, C, Hp, Wp)
+        embeddings = embeddings.flatten(2).transpose(1, 2)  # (B*T, L_frame, C)
+
+        total_patches = t_frames * (height // self.patch_size) * (width // self.patch_size)
+        embeddings = embeddings.reshape(batch_size, total_patches, self.embed_dim)
+
+        return embeddings
+
+
+class MageViTAttention(nn.Module):
+    """
+    Multi-headed attention with fused ``qkv`` / ``proj`` projections and RoPE support.
+
+    Attention is dispatched internally across ``sdpa`` (default) / ``flash_attention_2`` /
+    ``eager`` based on ``config._attn_implementation``.
+    """
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and "
+                f"`num_heads`: {self.num_heads})."
+            )
+
+        self.scale = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.qkv = nn.Linear(self.embed_dim, self.embed_dim * 3)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def _attn_impl(self) -> str:
+        impl = getattr(self.config, "_attn_implementation", None) or "sdpa"
+        if impl == "flash_attention_2" and not _flash_attn_available:
+            logger.warning_once("flash-attn is not installed; falling back to `sdpa` attention.")
+            impl = "sdpa"
+        return impl
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        batch_size, q_len, _ = hidden_states.size()
+
+        # (B, L, 3*H*D) -> (B, L, 3, H, D) -> 3 x (B, H, L, D)
+        q, k, v = (
+            self.qkv(hidden_states)
+            .reshape(batch_size, q_len, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 1, 3, 4)
+            .unbind(0)
+        )
+        query_states = q.transpose(1, 2)
+        key_states = k.transpose(1, 2)
+        value_states = v.transpose(1, 2)
+
+        if rotary_pos_emb is not None:
+            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, rotary_pos_emb)
+
+        dropout = self.attention_dropout if self.training else 0.0
+        impl = self._attn_impl()
+        attn_weights = None
+
+        if output_attentions or impl == "eager":
+            attn_output, attn_weights = eager_attention_forward(
+                self, query_states, key_states, value_states, attention_mask, self.scale, dropout
+            )
+        elif impl == "flash_attention_2":
+            # flash-attn expects (B, L, H, D)
+            attn_output = flash_attn_func(
+                query_states.transpose(1, 2),
+                key_states.transpose(1, 2),
+                value_states.transpose(1, 2),
+                dropout_p=dropout,
+                softmax_scale=self.scale,
+                causal=False,
+            )  # (B, L, H, D)
+        else:  # sdpa
+            attn_output = F.scaled_dot_product_attention(
+                query_states,
+                key_states,
+                value_states,
+                attn_mask=attention_mask,
+                dropout_p=dropout,
+                scale=self.scale,
+            )  # (B, H, L, D)
+            attn_output = attn_output.transpose(1, 2)  # (B, L, H, D)
+
+        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
+        attn_output = self.proj(attn_output)
+
+        return attn_output, attn_weights if output_attentions else None
+
+
+class MageViTEncoderLayer(nn.Module):
+    """Vision encoder layer with pre-norm attention and MLP."""
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.self_attn = MageViTAttention(config)
+        self.layer_norm1 = get_norm_layer(config)
+        self.mlp = SiglipMLP(config)
+        self.layer_norm2 = get_norm_layer(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        residual = hidden_states
+        hidden_states = self.layer_norm1(hidden_states)
+        hidden_states, attn_weights = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            rotary_pos_emb=rotary_pos_emb,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.layer_norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states, attn_weights) if output_attentions else (hidden_states,)
+        return outputs
+
+
+class MageViTEncoder(nn.Module):
+    def __init__(self, config: MageViTConfig):
+        super().__init__()
+        self.config = config
+        self.layers = nn.ModuleList([MageViTEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        rotary_pos_emb: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+    ) -> Union[tuple, BaseModelOutput]:
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    rotary_pos_emb,
+                    output_attentions,
+                )
+            else:
+                layer_outputs = layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    rotary_pos_emb=rotary_pos_emb,
+                    output_attentions=output_attentions,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Main Models
+# ---------------------------------------------------------------------------
+
+
+class MageViTPreTrainedModel(PreTrainedModel):
+    config_class = MageViTConfig
+    base_model_prefix = "mage_vit"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["MageViTEncoderLayer"]
+    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        std = self.config.initializer_range
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)):
+            module.weight.data.fill_(1.0)
+            if hasattr(module, "bias") and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, VisionRotaryEmbedding):
+            # inv_freq buffers are registered with persistent=False, so they are not in the
+            # checkpoint. When `from_pretrained` materializes the model from meta tensors these
+            # buffers would otherwise stay uninitialized; re-fill them so RoPE is correct post-load.
+            base = module.base
+            with torch.no_grad():
+                inv_t = 1.0 / (base ** (torch.arange(module.t_size, dtype=torch.float32) / module.t_size))
+                inv_h = 1.0 / (base ** (torch.arange(module.h_size, dtype=torch.float32) / module.h_size))
+                inv_w = 1.0 / (base ** (torch.arange(module.w_size, dtype=torch.float32) / module.w_size))
+                module.inv_freq_t.copy_(inv_t.to(module.inv_freq_t.device))
+                module.inv_freq_h.copy_(inv_h.to(module.inv_freq_h.device))
+                module.inv_freq_w.copy_(inv_w.to(module.inv_freq_w.device))
+
+
+class MageViTModel(MageViTPreTrainedModel):
+    """Mage-ViT vision transformer encoder."""
+
+    def __init__(self, config: MageViTConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = MageViTEmbeddings(config)
+        self.layernorm_pre = get_norm_layer(config)
+        self.encoder = MageViTEncoder(config)
+        self.video_rope = VisionRotaryEmbedding(config)
+
+        if config.use_head:
+            self.layernorm_post = get_norm_layer(config)
+            self.head = Siglip2MultiheadAttentionPoolingHead(config)
+        else:
+            self.layernorm_post = None
+            self.head = None
+
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        visible_indices: Optional[torch.Tensor] = None,
+        patch_positions: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[tuple, BaseModelOutputWithPooling]:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoModel, AutoImageProcessor
+        >>> from PIL import Image
+
+        >>> model = AutoModel.from_pretrained("microsoft/Mage-ViT", trust_remote_code=True)
+        >>> preprocessor = AutoImageProcessor.from_pretrained("microsoft/Mage-ViT", trust_remote_code=True)
+        >>> image = Image.open("path/to/your/image.jpg")
+        >>> pixel_values = preprocessor(images=image, return_tensors="pt")["pixel_values"]
+        >>> outputs = model(pixel_values)
+        >>> last_hidden_states = outputs.last_hidden_state
+        >>> pooled_output = outputs.pooler_output
+        ```
+        """
+        output_attentions = (
+            output_attentions if output_attentions is not None else getattr(self.config, "output_attentions", False)
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else getattr(self.config, "output_hidden_states", False)
+        )
+        return_dict = True if return_dict is None else return_dict
+
+        # Determine grid dimensions for RoPE (pixel_values may be 4D or 5D)
+        if pixel_values.dim() == 5:
+            t_frames = (
+                self.config.rope_temporal_size if self.config.rope_temporal_size is not None else pixel_values.shape[2]
+            )
+            height, width = pixel_values.shape[3], pixel_values.shape[4]
+        else:
+            t_frames = 1
+            height, width = pixel_values.shape[2], pixel_values.shape[3]
+
+        # 1. Embeddings
+        hidden_states = self.embeddings(pixel_values)
+        batch_size, total_patches, _ = hidden_states.shape
+
+        # 2. Visible-index handling (defaults to all patches)
+        if visible_indices is None:
+            visible_indices = torch.arange(total_patches, device=pixel_values.device).unsqueeze(0).expand(
+                batch_size, -1
+            )
+
+        # 3. RoPE construction
+        if patch_positions is not None:
+            freqs_visible = self.video_rope.forward_from_positions(patch_positions)  # (B, L, half)
+        else:
+            freqs_full = self.video_rope.forward_with_thw(
+                t=t_frames,
+                h=height // self.config.patch_size,
+                w=width // self.config.patch_size,
+                device=pixel_values.device,
+            )
+            freqs_visible = freqs_full[visible_indices]  # (B, L, half)
+
+        # Concatenate half + half -> head_dim
+        freqs_visible = torch.cat([freqs_visible, freqs_visible], dim=-1)
+
+        # 4. Pre-norm & encoder
+        hidden_states = self.layernorm_pre(hidden_states)
+
+        # Sparse mode: gather only visible patches to match freqs_visible
+        if visible_indices.shape[1] != total_patches:
+            hidden_states = hidden_states.gather(
+                1, visible_indices.unsqueeze(-1).expand(-1, -1, hidden_states.shape[-1])
+            )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=None,
+            rotary_pos_emb=freqs_visible,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        sequence_output = encoder_outputs.last_hidden_state
+
+        if self.layernorm_post is not None:
+            sequence_output = self.layernorm_post(sequence_output)
+
+        pooled_output = self.head(sequence_output) if self.head is not None else None
+
+        if not return_dict:
+            outputs = (sequence_output, pooled_output)
+            if output_hidden_states:
+                outputs = outputs + (encoder_outputs.hidden_states,)
+            if output_attentions:
+                outputs = outputs + (encoder_outputs.attentions,)
+            return outputs
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+__all__ = ["MageViTModel", "MageViTPreTrainedModel"]

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 256,
+    "width": 256
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 256
+  }
+}