Upload folder using huggingface_hub

config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "FineViTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_finevit.FineViTConfig",
+    "AutoModel": "modeling_finevit.FineViTModel"
+  },
+  "backbone_name": "dinov2_base_reg",
+  "backbone_patch_size": 14,
+  "backbone_pool_size": 2,
+  "backbone_pretrained": true,
+  "decoder_head_dim": 64,
+  "decoder_num_attention_heads": 12,
+  "decoder_num_layers": 8,
+  "decoder_position_grid_size": 16,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "encoder_head_dim": 64,
+  "encoder_num_attention_heads": 12,
+  "encoder_num_layers": 8,
+  "encoder_position_grid_size": 16,
+  "intermediate_size": 3072,
+  "model_type": "finevit",
+  "prefix_alignment_mlp_ratio": 4.0,
+  "prefix_alignment_num_attention_heads": 12,
+  "prefix_alignment_qk_norm": true,
+  "training_mask_ratio": 0.75,
+  "training_noise_gamma": 0.0,
+  "transformers_version": "5.8.1"
+}

configuration_finevit.py

@@ -0,0 +1,72 @@
+
+from transformers import PretrainedConfig
+
+
+class FineViTConfig(PretrainedConfig):
+    model_type = "finevit"
+    has_no_defaults_at_init = True
+
+    def __init__(
+        self,
+        encoder_head_dim: int,
+        encoder_num_attention_heads: int,
+        decoder_head_dim: int,
+        decoder_num_attention_heads: int,
+        prefix_alignment_num_attention_heads: int | None = None,
+        prefix_alignment_qk_norm: bool = True,
+        prefix_alignment_mlp_ratio: float = 4.0,
+        backbone_patch_size: int = 14,
+        encoder_num_layers: int = 12,
+        decoder_num_layers: int = 4,
+        intermediate_size: int = 3072,
+        backbone_name: str = "dinov2_base_reg",
+        backbone_pool_size: int = 2,
+        encoder_position_grid_size: int | None = None,
+        decoder_position_grid_size: int = 16,
+        training_mask_ratio: float = 0.7,
+        training_noise_gamma: float = 3.0,
+        dropout: float = 0.0,
+        backbone_pretrained: bool = False,
+        **kwargs,
+    ):
+        legacy_keys = {
+            "num_encoder_perceiver_tokens",
+            "perceiver_num_iterations",
+        } & set(kwargs)
+        if legacy_keys:
+            keys = ", ".join(sorted(legacy_keys))
+            raise ValueError(f"Perceiver config keys are no longer supported: {keys}.")
+        super().__init__(**kwargs)
+
+        self.backbone_name = backbone_name
+        self.backbone_pool_size = int(backbone_pool_size)
+        self.encoder_position_grid_size = int(
+            encoder_position_grid_size
+            if encoder_position_grid_size is not None
+            else decoder_position_grid_size
+        )
+        self.decoder_position_grid_size = int(decoder_position_grid_size)
+        self.training_mask_ratio = float(training_mask_ratio)
+        self.training_noise_gamma = float(training_noise_gamma)
+        self.backbone_patch_size = int(backbone_patch_size)
+        self.encoder_num_layers = int(encoder_num_layers)
+        self.decoder_num_layers = int(decoder_num_layers)
+        self.encoder_head_dim = encoder_head_dim
+        self.encoder_num_attention_heads = encoder_num_attention_heads
+        self.decoder_head_dim = decoder_head_dim
+        self.decoder_num_attention_heads = decoder_num_attention_heads
+        self.prefix_alignment_num_attention_heads = int(
+            prefix_alignment_num_attention_heads
+            if prefix_alignment_num_attention_heads is not None
+            else encoder_num_attention_heads
+        )
+        self.prefix_alignment_qk_norm = bool(prefix_alignment_qk_norm)
+        self.prefix_alignment_mlp_ratio = float(prefix_alignment_mlp_ratio)
+        self.intermediate_size = int(intermediate_size)
+        self.dropout = float(dropout)
+        self.backbone_pretrained = bool(backbone_pretrained)
+        self.architectures = ["FineViTModel"]
+        self.auto_map = {
+            "AutoConfig": "configuration_finevit.FineViTConfig",
+            "AutoModel": "modeling_finevit.FineViTModel",
+        }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7515ecc536fca2f995fed063dfc72a0e3f0f81b38daf4e82823a9f6a9c055a4b
+size 866112016

modeling_finevit.py

@@ -0,0 +1,929 @@
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers.attention_pool import AttentionPoolLatent
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration_finevit import FineViTConfig
+from .vision_encoder import build_encoder
+
+
+def _run_layer(layer, *inputs, training: bool):
+    # if training:
+    #     return checkpoint(layer, *inputs, use_reentrant=False)
+    return layer(*inputs)
+
+
+def _square_grid_side(length: int, name: str) -> int:
+    side = math.isqrt(length)
+    if side * side != length:
+        raise ValueError(f"{name} length must be a square grid, got {length}.")
+    return side
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate adjacent feature pairs for rotary position embedding."""
+    if x.shape[-1] % 2 != 0:
+        raise ValueError(f"RoPE requires an even head dimension, got {x.shape[-1]}.")
+    x = x.view(*x.shape[:-1], -1, 2)
+    x1, x2 = x.unbind(dim=-1)
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+    """Apply rotary position embedding to attention states shaped [B, H, S, D]."""
+    if freqs_cis.shape[-1] != x.shape[-1] * 2:
+        raise ValueError(
+            "RoPE tensor width must be twice the attention head dimension, got "
+            f"{freqs_cis.shape[-1]} and {x.shape[-1]}."
+        )
+    freqs_cos, freqs_sin = freqs_cis.to(device=x.device, dtype=x.dtype).chunk(2, dim=-1)
+    freqs_cos = freqs_cos.view(1, 1, freqs_cos.shape[0], freqs_cos.shape[1])
+    freqs_sin = freqs_sin.view(1, 1, freqs_sin.shape[0], freqs_sin.shape[1])
+    return x * freqs_cos + rotate_half(x) * freqs_sin
+
+
+def get_rope_tensor(
+    dim: int, seq_h: int, seq_w: int, max_freq: float = 7.0, min_freq: float = 7e-4
+) -> torch.Tensor:
+    """Generate rotary position embedding tensor for 2D sequences."""
+    if dim % 4 != 0:
+        raise ValueError(f"2D RoPE head dimension must be divisible by 4, got {dim}.")
+    freqs_1d = max_freq * (max_freq / min_freq) ** torch.linspace(0, -1, dim // 4)
+    freqs_1d = torch.cat([freqs_1d, freqs_1d])
+    freqs_2d = torch.zeros(2, dim)
+    freqs_2d[0, : dim // 2] = freqs_1d
+    freqs_2d[1, -dim // 2 :] = freqs_1d
+    freqs_2d = freqs_2d * 2 * torch.pi
+    coord_x = torch.linspace(0, 1, seq_h)
+    coord_y = torch.linspace(0, 1, seq_w)
+    coords_all = torch.cartesian_prod(coord_x, coord_y)
+    angle = coords_all @ freqs_2d
+    return torch.cat([angle.cos(), angle.sin()], dim=-1)
+
+
+class GatedMLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        bias: bool = False,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.fc_in = nn.Linear(input_dim, 2 * hidden_dim, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+        self.fc_out = nn.Linear(hidden_dim, output_dim, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, value = self.fc_in(hidden_states).chunk(2, dim=-1)
+        hidden_states = F.silu(gate) * value
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.fc_out(hidden_states)
+        return hidden_states
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        output_dim: int,
+        bias: bool = False,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.fc_in = nn.Linear(input_dim, hidden_dim, bias=bias)
+        self.dropout = nn.Dropout(dropout)
+        self.fc_out = nn.Linear(hidden_dim, output_dim, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.silu(self.fc_in(hidden_states))
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.fc_out(hidden_states)
+        return hidden_states
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        head_dim: int,
+        num_attention_heads: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.num_attention_heads = num_attention_heads
+        self.dropout = dropout
+        self.q_norm = nn.LayerNorm(head_dim, eps=1e-6)
+        self.k_norm = nn.LayerNorm(head_dim, eps=1e-6)
+        self.to_qkv = nn.Linear(
+            hidden_size, 3 * head_dim * num_attention_heads, bias=False
+        )
+        self.to_out = nn.Linear(head_dim * num_attention_heads, hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope_tensor: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch, seq, _ = hidden_states.shape
+        qkv = self.to_qkv(hidden_states).view(
+            batch, seq, 3, self.num_attention_heads, self.head_dim
+        )
+        query, key, value = qkv.permute(0, 2, 3, 1, 4).contiguous().unbind(dim=1)
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        if rope_tensor is not None:
+            if rope_tensor.shape[0] != seq:
+                raise ValueError(
+                    "RoPE sequence length must match hidden_states sequence length, "
+                    f"got {rope_tensor.shape[0]} and {seq}."
+                )
+            query = apply_rotary_emb(query, rope_tensor)
+            key = apply_rotary_emb(key, rope_tensor)
+        attn = (
+            F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+            .transpose(1, 2)
+            .contiguous()
+            .view(batch, seq, -1)
+        )
+        return self.to_out(attn)
+
+
+class CrossAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        head_dim: int,
+        num_attention_heads: int,
+        dropout: float = 0.0,
+        attention_scale: float | None = None,
+    ):
+        super().__init__()
+        self.head_dim = head_dim
+        self.num_attention_heads = num_attention_heads
+        self.dropout = dropout
+        self.attention_scale = attention_scale
+        inner_dim = head_dim * num_attention_heads
+        self.q_norm = nn.LayerNorm(head_dim, eps=1e-6)
+        self.k_norm = nn.LayerNorm(head_dim, eps=1e-6)
+        self.to_q = nn.Linear(hidden_size, inner_dim, bias=False)
+        self.to_k = nn.Linear(hidden_size, inner_dim, bias=False)
+        self.to_v = nn.Linear(hidden_size, inner_dim, bias=False)
+        self.to_out = nn.Linear(inner_dim, hidden_size, bias=False)
+
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        return_weights: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        batch, q_seq, _ = query_states.shape
+        key_seq = key_states.shape[1]
+        value_seq = value_states.shape[1]
+        if key_seq != value_seq:
+            raise ValueError(
+                "key_states and value_states sequence lengths must match, got "
+                f"{key_seq} and {value_seq}."
+            )
+
+        query = (
+            self.to_q(query_states)
+            .view(batch, q_seq, self.num_attention_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+        key = (
+            self.to_k(key_states)
+            .view(batch, key_seq, self.num_attention_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+        value = (
+            self.to_v(value_states)
+            .view(batch, value_seq, self.num_attention_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+
+        if return_weights:
+            scale = self.attention_scale or query.shape[-1] ** -0.5
+            scores = (query * scale) @ key.transpose(-2, -1)
+            weights = scores.softmax(dim=-1)
+            if self.training and self.dropout > 0.0:
+                attn = F.dropout(weights, p=self.dropout) @ value
+            else:
+                attn = weights @ value
+        else:
+            attn = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                scale=self.attention_scale,
+                dropout_p=self.dropout if self.training else 0.0,
+            )
+            weights = None
+
+        attn = attn.transpose(1, 2).contiguous().view(batch, q_seq, -1)
+        out = self.to_out(attn)
+        return (out, weights) if return_weights else out
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        head_dim: int,
+        num_attention_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.attn_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.attention = Attention(hidden_size, head_dim, num_attention_heads, dropout)
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.mlp_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.mlp = MLP(
+            hidden_size,
+            mlp_hidden_dim,
+            hidden_size,
+            dropout=dropout,
+        )
+        self.mlp_dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        rope_tensor: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn_dropout(
+            self.attention(self.attn_norm(hidden_states), rope_tensor)
+        )
+        hidden_states = hidden_states + self.mlp_dropout(
+            self.mlp(self.mlp_norm(hidden_states))
+        )
+        return hidden_states
+
+
+class CrossAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        head_dim: int,
+        num_attention_heads: int,
+        mlp_hidden_dim: int,
+        fused_kv: bool = False,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.fused_kv = fused_kv
+
+        self.q_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.kv_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.attention = CrossAttention(
+            hidden_size, head_dim, num_attention_heads, dropout
+        )
+        self.attn_dropout = nn.Dropout(dropout)
+
+        self.mlp_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.mlp = MLP(
+            hidden_size,
+            mlp_hidden_dim,
+            hidden_size,
+            dropout=dropout,
+        )
+        self.mlp_dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        query_states: torch.Tensor,
+        kv_states: torch.Tensor,
+        return_weights: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        kv_input = (
+            torch.cat([query_states, kv_states], dim=1) if self.fused_kv else kv_states
+        )
+        kv_input = self.kv_norm(kv_input)
+
+        attn_out = self.attention(
+            self.q_norm(query_states),
+            kv_input,
+            kv_input,
+            return_weights=return_weights,
+        )
+
+        weights = None
+        if return_weights:
+            attn_out, weights = attn_out
+
+        query_states = query_states + self.attn_dropout(attn_out)
+        query_states = query_states + self.mlp_dropout(
+            self.mlp(self.mlp_norm(query_states))
+        )
+
+        return (query_states, weights) if return_weights else query_states
+
+
+class PatchPooler(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        pool_size: int,
+        head_dim: int,
+        num_attention_heads: int,
+        mlp_hidden_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.pool_size = int(pool_size)
+        self.num_attention_heads = int(num_attention_heads)
+        if self.pool_size <= 0:
+            raise ValueError(f"pool_size must be positive, got {self.pool_size}.")
+
+        self.pool_attention = CrossAttention(
+            hidden_size=hidden_size,
+            head_dim=head_dim,
+            num_attention_heads=num_attention_heads,
+            dropout=dropout,
+        )
+
+    def forward(
+        self,
+        features: torch.Tensor,
+        return_attention_weights: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        batch, length, dim = features.shape
+        grid_side = _square_grid_side(length, "backbone feature")
+        if grid_side % self.pool_size != 0:
+            raise ValueError(
+                "backbone feature grid side must be divisible by backbone_pool_size, "
+                f"got grid_side={grid_side} and pool_size={self.pool_size}."
+            )
+
+        pooled_side = grid_side // self.pool_size
+        num_windows = pooled_side * pooled_side
+        pool = self.pool_size
+        windows = (
+            features.view(batch, pooled_side, pool, pooled_side, pool, dim)
+            .permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .flatten(3, 4)
+        ).view(batch * num_windows, pool * pool, dim)
+
+        query = windows.mean(dim=1, keepdim=True)
+        if not return_attention_weights:
+            pooled = self.pool_attention(query, windows, windows)
+            return pooled.view(batch, num_windows, dim), None
+
+        pooled, weights = self.pool_attention(
+            query, windows, windows, return_weights=True
+        )
+        weights = weights.squeeze(2).view(
+            batch,
+            num_windows,
+            self.num_attention_heads,
+            pool * pool,
+        )
+        return pooled.view(batch, num_windows, dim), weights.permute(
+            0, 2, 1, 3
+        ).contiguous()
+
+
+class PatchUnpooler(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        pool_size: int,
+        intermediate_size: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.pool_size = int(pool_size)
+        if self.pool_size <= 0:
+            raise ValueError(f"pool_size must be positive, got {self.pool_size}.")
+
+        self.up = nn.Linear(
+            hidden_size, hidden_size * self.pool_size * self.pool_size, bias=False
+        )
+
+    def forward(
+        self,
+        pooled_tokens: torch.Tensor,
+        output_grid_side: int,
+    ) -> torch.Tensor:
+        batch, pooled_length, dim = pooled_tokens.shape
+        if dim != self.hidden_size:
+            raise ValueError(
+                "pooled_tokens hidden size must match unpooler hidden size, got "
+                f"{dim} and {self.hidden_size}."
+            )
+        if output_grid_side <= 0:
+            raise ValueError(
+                f"output_grid_side must be positive, got {output_grid_side}."
+            )
+        if output_grid_side % self.pool_size != 0:
+            raise ValueError(
+                "output_grid_side must be divisible by pool_size, got "
+                f"{output_grid_side} and {self.pool_size}."
+            )
+
+        pooled_side = output_grid_side // self.pool_size
+        if pooled_side * pooled_side != pooled_length:
+            raise ValueError(
+                "pooled_tokens length must match output_grid_side and pool_size, got "
+                f"length={pooled_length}, output_grid_side={output_grid_side}, "
+                f"pool_size={self.pool_size}."
+            )
+
+        pool = self.pool_size
+        patches = self.up(pooled_tokens).view(batch, pooled_length, pool, pool, dim)
+        patches = patches + pooled_tokens[:, :, None, None, :]
+        patches = patches.view(batch, pooled_side, pooled_side, pool, pool, dim)
+        return (
+            patches.permute(0, 1, 3, 2, 4, 5)
+            .contiguous()
+            .view(batch, output_grid_side * output_grid_side, dim)
+        )
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, grid_size: int, hidden_size: int) -> None:
+        super().__init__()
+        self.grid_size = int(grid_size)
+        self.hidden_size = hidden_size
+        if self.grid_size <= 0:
+            raise ValueError(f"grid_size must be positive, got {self.grid_size}.")
+
+        self.weight = nn.Parameter(
+            torch.empty(1, self.grid_size * self.grid_size, hidden_size)
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.trunc_normal_(self.weight, std=0.02)
+
+    def forward(
+        self,
+        grid_side: int,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        if grid_side <= 0:
+            raise ValueError(f"grid_side must be positive, got {grid_side}.")
+
+        position = self.weight
+        if grid_side != self.grid_size:
+            position = position.view(
+                1,
+                self.grid_size,
+                self.grid_size,
+                self.hidden_size,
+            ).permute(0, 3, 1, 2)
+            position = F.interpolate(
+                position.contiguous(),
+                size=(grid_side, grid_side),
+                mode="bicubic",
+                align_corners=False,
+            )
+            position = (
+                position.permute(0, 2, 3, 1)
+                .contiguous()
+                .view(1, grid_side * grid_side, self.hidden_size)
+            )
+        return position.to(device=device, dtype=dtype)
+
+
+class FineViTBackbone(nn.Module):
+    def __init__(self, config: FineViTConfig):
+        super().__init__()
+        self.encoder = build_encoder(
+            config.backbone_name,
+            pretrained=config.backbone_pretrained,
+        )
+        self.hidden_size = self.encoder.hidden_size
+        self.patch_size = self.encoder.patch_size
+
+        if self.patch_size != config.backbone_patch_size:
+            raise ValueError(
+                "Configured backbone_patch_size does not match encoder patch size: "
+                f"{config.backbone_patch_size} != {self.patch_size}."
+            )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        return_prefix_tokens: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        return self.encoder(
+            pixel_values,
+            return_prefix_tokens=return_prefix_tokens,
+        )
+
+
+class FineViTEncoder(nn.Module):
+    def __init__(self, config: FineViTConfig, hidden_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = hidden_size
+        self.num_prefix_tokens = 0
+        self.training_mask_ratio = float(config.training_mask_ratio)
+        if config.encoder_head_dim % 4 != 0:
+            raise ValueError(
+                "encoder_head_dim must be divisible by 4 for 2D RoPE, got "
+                f"{config.encoder_head_dim}."
+            )
+
+        self.patch_positions = PositionEmbedding(
+            config.encoder_position_grid_size,
+            hidden_size,
+        )
+        self.input_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.layers = nn.ModuleList(
+            TransformerBlock(
+                hidden_size=hidden_size,
+                head_dim=config.encoder_head_dim,
+                num_attention_heads=config.encoder_num_attention_heads,
+                mlp_hidden_dim=config.intermediate_size,
+                dropout=config.dropout,
+            )
+            for _ in range(config.encoder_num_layers)
+        )
+
+        self.pooler_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.patch_pooler = PatchPooler(
+            hidden_size=hidden_size,
+            pool_size=config.backbone_pool_size,
+            head_dim=config.encoder_head_dim,
+            num_attention_heads=config.encoder_num_attention_heads,
+            mlp_hidden_dim=config.intermediate_size,
+            dropout=config.dropout,
+        )
+        self.head_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.head_proj = MLP(hidden_size, config.intermediate_size, hidden_size)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+        elif isinstance(module, nn.Embedding):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+
+    def _sample_mask_ratio(self, device: torch.device) -> torch.Tensor:
+        return (
+            torch.empty((), device=device)
+            .uniform_(-0.1, self.training_mask_ratio)
+            .clamp_min(0.0)
+        )
+
+    def _apply_masking(
+        self,
+        backbone_features: torch.Tensor,
+        mask_token: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if not self.training:
+            return backbone_features
+        if mask_token is None:
+            raise ValueError(
+                "mask_token is required for encoder masking during training."
+            )
+
+        mask_token = mask_token.to(
+            device=backbone_features.device,
+            dtype=backbone_features.dtype,
+        )
+        if mask_token.shape[-1] != backbone_features.shape[-1]:
+            raise ValueError(
+                "mask_token hidden size must match backbone feature hidden size, got "
+                f"{mask_token.shape[-1]} and {backbone_features.shape[-1]}."
+            )
+
+        mask_ratio = self._sample_mask_ratio(backbone_features.device)
+        anchor = mask_token.sum() * 0.0
+        token_mask = torch.rand(
+            backbone_features.shape[0],
+            backbone_features.shape[1],
+            1,
+            device=backbone_features.device,
+        ).lt(mask_ratio)
+        return (
+            torch.where(
+                token_mask,
+                mask_token.expand_as(backbone_features),
+                backbone_features,
+            )
+            + anchor
+        )
+
+    def forward(
+        self,
+        backbone_features: torch.Tensor,
+        mask_token: torch.Tensor | None = None,
+        return_attention_weights: bool = False,
+    ) -> tuple[
+        torch.Tensor | None,
+        torch.Tensor,
+    ]:
+        grid_side = _square_grid_side(backbone_features.shape[1], "backbone feature")
+        hidden_states = backbone_features
+        hidden_states = hidden_states + self.patch_positions(
+            grid_side,
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        hidden_states = self.input_norm(hidden_states)
+        rope_tensor = get_rope_tensor(
+            self.config.encoder_head_dim,
+            grid_side,
+            grid_side,
+        ).to(device=hidden_states.device, dtype=hidden_states.dtype)
+        for layer in self.layers:
+            hidden_states = _run_layer(
+                layer,
+                hidden_states,
+                rope_tensor,
+                training=self.training,
+            )
+        pooled_tokens, pool_attention_weights = self.patch_pooler(
+            self.pooler_norm(hidden_states),
+            return_attention_weights=return_attention_weights,
+        )
+        encoder_hidden_states = self._apply_masking(pooled_tokens, mask_token)
+        encoder_hidden_states = self.head_proj(self.head_norm(encoder_hidden_states))
+
+        return (
+            pool_attention_weights,
+            encoder_hidden_states,
+        )
+
+
+class FineViTFeatureDecoder(nn.Module):
+    def __init__(self, config: FineViTConfig, hidden_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = hidden_size
+        if config.decoder_head_dim % 4 != 0:
+            raise ValueError(
+                "decoder_head_dim must be divisible by 4 for 2D RoPE, got "
+                f"{config.decoder_head_dim}."
+            )
+
+        self.input_norm_1 = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.input_norm_2 = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+
+        self.patch_positions = PositionEmbedding(
+            config.decoder_position_grid_size,
+            hidden_size,
+        )
+        self.patch_unpooler = PatchUnpooler(
+            hidden_size=hidden_size,
+            pool_size=config.backbone_pool_size,
+            intermediate_size=config.intermediate_size,
+            dropout=config.dropout,
+        )
+        self.layers = nn.ModuleList(
+            TransformerBlock(
+                hidden_size=hidden_size,
+                head_dim=config.decoder_head_dim,
+                num_attention_heads=config.decoder_num_attention_heads,
+                mlp_hidden_dim=config.intermediate_size,
+                dropout=config.dropout,
+            )
+            for _ in range(config.decoder_num_layers)
+        )
+
+        self.head_norm = nn.LayerNorm(hidden_size, eps=1e-6, bias=False)
+        self.head_proj = nn.Linear(hidden_size, hidden_size)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+        elif isinstance(module, nn.Embedding):
+            nn.init.trunc_normal_(module.weight, std=0.02)
+
+    def forward(
+        self,
+        encoder_hidden_states: torch.Tensor,
+        backbone_grid_side: int,
+    ) -> torch.Tensor:
+        if encoder_hidden_states.shape[-1] != self.hidden_size:
+            raise ValueError(
+                "encoder_hidden_states hidden size must match decoder hidden size, got "
+                f"{encoder_hidden_states.shape[-1]} and {self.hidden_size}."
+            )
+
+        encoder_hidden_states = self.input_norm_1(encoder_hidden_states)
+
+        hidden = self.patch_unpooler(
+            encoder_hidden_states,
+            output_grid_side=backbone_grid_side,
+        )
+        hidden = hidden + self.patch_positions(
+            backbone_grid_side,
+            device=encoder_hidden_states.device,
+            dtype=encoder_hidden_states.dtype,
+        )
+        hidden = self.input_norm_2(hidden)
+        rope_tensor = get_rope_tensor(
+            self.config.decoder_head_dim,
+            backbone_grid_side,
+            backbone_grid_side,
+        ).to(device=hidden.device, dtype=hidden.dtype)
+
+        for layer in self.layers:
+            hidden = _run_layer(
+                layer,
+                hidden,
+                rope_tensor,
+                training=self.training,
+            )
+
+        return self.head_proj(self.head_norm(hidden))
+
+
+@dataclass
+class FineViTOutput(ModelOutput):
+    encoder_hidden_states: torch.Tensor | None = None
+    backbone_features: torch.Tensor | None = None
+    backbone_prefix_tokens: torch.Tensor | None = None
+    aligned_prefix_tokens: torch.Tensor | None = None
+    decoded_patches: torch.Tensor | None = None
+    pool_attention_weights: torch.Tensor | None = None
+
+
+class FineViTModel(PreTrainedModel):
+    config_class = FineViTConfig
+    base_model_prefix = "finevit"
+    main_input_name = "pixel_values"
+    _no_split_modules = ["TransformerBlock", "CrossAttentionBlock"]
+    MODEL_FILENAME = "model.safetensors"
+    CONFIG_FILENAME = "config.json"
+    all_tied_weights_keys = {}
+
+    def __init__(self, config: FineViTConfig):
+        super().__init__(config)
+        self.backbone = FineViTBackbone(config)
+        hidden_size = self.backbone.hidden_size
+        self.encoder = FineViTEncoder(config, hidden_size)
+        self.decoder = FineViTFeatureDecoder(config, hidden_size)
+        self.mask_token = nn.Parameter(torch.empty(1, 1, hidden_size))
+        nn.init.trunc_normal_(self.mask_token, std=0.02)
+        self.backbone_num_prefix_tokens = int(self.backbone.encoder.num_prefix_tokens)
+        self.encoder_num_prefix_tokens = int(self.encoder.num_prefix_tokens)
+        self.num_prefix_tokens = self.encoder_num_prefix_tokens
+        self.prefix_aligner = AttentionPoolLatent(
+            in_features=hidden_size,
+            out_features=hidden_size,
+            embed_dim=hidden_size,
+            num_heads=config.prefix_alignment_num_attention_heads,
+            latent_len=self.backbone_num_prefix_tokens,
+            qk_norm=config.prefix_alignment_qk_norm,
+            mlp_ratio=config.prefix_alignment_mlp_ratio,
+            pool_type="",
+            norm_layer=nn.LayerNorm,
+            drop=config.dropout,
+        )
+
+    def _prepare_encoder_hidden_states(self, encoded: FineViTOutput) -> torch.Tensor:
+        if encoded.encoder_hidden_states is None:
+            raise ValueError("encoded must contain encoder_hidden_states.")
+
+        if self.training:
+            encoder_hidden_states = encoded.encoder_hidden_states
+            gamma = float(self.config.training_noise_gamma)
+            if gamma < 0.0:
+                raise ValueError(
+                    f"training_noise_gamma must be non-negative, got {gamma}."
+                )
+            if gamma > 0.0:
+                tau = torch.rand(
+                    encoder_hidden_states.shape[0],
+                    1,
+                    1,
+                    device=encoder_hidden_states.device,
+                    dtype=encoder_hidden_states.dtype,
+                )
+                encoder_hidden_states = torch.lerp(
+                    encoder_hidden_states,
+                    torch.randn_like(encoder_hidden_states) * gamma,
+                    tau,
+                )
+            return encoder_hidden_states
+
+        return encoded.encoder_hidden_states
+
+    def encode(
+        self,
+        pixel_values: torch.Tensor,
+        return_attention_weights: bool = False,
+        output_backbone_prefix_tokens: bool = False,
+        output_aligned_prefix_tokens: bool = False,
+    ) -> FineViTOutput:
+        backbone_prefix_tokens = None
+        backbone_output = self.backbone(
+            pixel_values,
+            return_prefix_tokens=output_backbone_prefix_tokens,
+        )
+        if output_backbone_prefix_tokens:
+            backbone_features, backbone_prefix_tokens = backbone_output
+        else:
+            backbone_features = backbone_output
+        (
+            pool_attention_weights,
+            encoder_hidden_states,
+        ) = self.encoder(
+            backbone_features,
+            mask_token=self.mask_token,
+            return_attention_weights=return_attention_weights,
+        )
+        aligned_prefix_tokens = (
+            self.prefix_aligner(encoder_hidden_states)
+            if output_aligned_prefix_tokens
+            else None
+        )
+        return FineViTOutput(
+            encoder_hidden_states=encoder_hidden_states,
+            backbone_features=backbone_features,
+            backbone_prefix_tokens=backbone_prefix_tokens,
+            aligned_prefix_tokens=aligned_prefix_tokens,
+            pool_attention_weights=pool_attention_weights,
+        )
+
+    def decode_encoder_hidden_states(
+        self,
+        encoder_hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        pooled_grid_side = _square_grid_side(
+            encoder_hidden_states.shape[1],
+            "encoder hidden states",
+        )
+        backbone_grid_side = pooled_grid_side * int(self.config.backbone_pool_size)
+        return self.decoder(encoder_hidden_states, backbone_grid_side)
+
+    def decode(
+        self,
+        encoded: FineViTOutput,
+    ) -> torch.Tensor:
+        return self.decode_encoder_hidden_states(
+            self._prepare_encoder_hidden_states(encoded)
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_decoded_patches: bool = True,
+        output_backbone_prefix_tokens: bool = False,
+        output_aligned_prefix_tokens: bool = False,
+        return_attention_weights: bool = True,
+    ) -> FineViTOutput:
+        output = self.encode(
+            pixel_values,
+            return_attention_weights=return_attention_weights,
+            output_backbone_prefix_tokens=output_backbone_prefix_tokens,
+            output_aligned_prefix_tokens=output_aligned_prefix_tokens,
+        )
+        if output_decoded_patches:
+            output.decoded_patches = self.decode(output)
+        return output
+
+
+FineViT = FineViTModel
+
+FineViTConfig.register_for_auto_class()
+FineViTModel.register_for_auto_class("AutoModel")

vision_encoder.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+import timm
+
+
+class TimmVisionEncoder(nn.Module):
+    def __init__(self, pretrained_encoder_name: str, load_pretrained: bool = True):
+        super().__init__()
+        self.model = timm.create_model(
+            pretrained_encoder_name,
+            pretrained=load_pretrained,
+            dynamic_img_size=True,
+        )
+        self.hidden_size = self.model.embed_dim
+        self.model.norm = nn.LayerNorm(self.hidden_size, elementwise_affine=False)
+        self.num_prefix_tokens = self.model.num_prefix_tokens
+        patch_size = self.model.patch_embed.patch_size
+        self.patch_size = int(
+            patch_size[0] if isinstance(patch_size, tuple) else patch_size
+        )
+        if self.patch_size == 14:
+            input_size = 224
+        else:
+            input_size = 256
+        self.model.set_input_size(input_size, patch_size)
+
+        data_config = timm.data.resolve_model_data_config(self.model)
+        self.register_buffer(
+            "pixel_mean",
+            torch.tensor(data_config["mean"])[None, :, None, None],
+            persistent=True,
+        )
+        self.register_buffer(
+            "pixel_std",
+            torch.tensor(data_config["std"])[None, :, None, None],
+            persistent=True,
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        return_prefix_tokens: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        x = pixel_values.float() / 255.0
+        weight_dtype = self.model.patch_embed.proj.weight.dtype
+        x = ((x - self.pixel_mean) / self.pixel_std).to(dtype=weight_dtype)
+        output = self.model.forward_features(x)
+        if isinstance(output, dict):
+            output = output["x"]
+        prefix_tokens = output[:, : self.num_prefix_tokens, :]
+        patch_tokens = output[:, self.num_prefix_tokens :, :]
+        if return_prefix_tokens:
+            return patch_tokens, prefix_tokens
+        return patch_tokens
+
+
+class DinoV2Encoder(TimmVisionEncoder):
+    pass
+
+
+class DinoV3Encoder(TimmVisionEncoder):
+    pass
+
+
+ALL_ENCODERS = {
+    "dinov3_small": (DinoV3Encoder, "vit_small_patch16_dinov3.lvd1689m"),
+    "dinov3_base": (DinoV3Encoder, "vit_base_patch16_dinov3.lvd1689m"),
+    "dinov3_large": (DinoV3Encoder, "vit_large_patch16_dinov3.lvd1689m"),
+    "dinov2_small_reg": (DinoV2Encoder, "vit_small_patch14_reg4_dinov2.lvd142m"),
+    "dinov2_base_reg": (DinoV2Encoder, "vit_base_patch14_reg4_dinov2.lvd142m"),
+    "dinov2_large_reg": (DinoV2Encoder, "vit_large_patch14_reg4_dinov2.lvd142m"),
+}
+
+
+def build_encoder(encoder_name: str, pretrained: bool = False):
+    if encoder_name not in ALL_ENCODERS:
+        raise ValueError(
+            f"Unknown encoder {encoder_name!r}. Available: {list(ALL_ENCODERS)}"
+        )
+    model_cls, model_id = ALL_ENCODERS[encoder_name]
+    return model_cls(model_id, load_pretrained=pretrained)
+
+
+__all__ = [
+    "ALL_ENCODERS",
+    "DinoV2Encoder",
+    "DinoV3Encoder",
+    "TimmVisionEncoder",
+    "build_encoder",
+]