import math from dataclasses import dataclass from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoConfig, AutoModel, PreTrainedModel from transformers.utils import ModelOutput from .config import FineViTConfig class GatedMLP(nn.Module): def __init__( self, input_dim: int, hidden_dim: int, output_dim: int, bias: bool = True, ): super().__init__() self.fc_in = nn.Linear(input_dim, 2 * hidden_dim, bias=bias) self.fc_out = nn.Linear(hidden_dim, output_dim, bias=bias) def forward(self, x): gate, value = self.fc_in(x).chunk(2, dim=-1) return self.fc_out(F.silu(gate) * value) class FourierPositionEncoder(nn.Module): def __init__(self, feat_dim: int, hidden_size: int) -> None: super().__init__() if feat_dim <= 0 or feat_dim % 2 != 0: raise ValueError("feat_dim must be a positive even integer.") self.embed = nn.Linear(2, feat_dim // 2, bias=False) self.proj = nn.Linear(feat_dim, hidden_size, bias=False) self.mlp = GatedMLP(hidden_size, hidden_size * 4, hidden_size) def forward(self, coordinates: torch.Tensor) -> torch.Tensor: features = 2 * math.pi * self.embed(coordinates) features = torch.cat([features.cos(), features.sin()], dim=-1) hidden_states = self.proj(features) return hidden_states + self.mlp(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.RMSNorm(head_dim, eps=1e-6) self.k_norm = nn.RMSNorm(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): batch, seq, _ = hidden_states.shape qkv = self.to_qkv(hidden_states).view( batch, seq, 3, self.num_attention_heads, self.head_dim ) q, k, v = qkv.permute(0, 2, 3, 1, 4).contiguous().unbind(dim=1) attn = ( F.scaled_dot_product_attention( self.q_norm(q), self.k_norm(k), v, 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, ): super().__init__() self.head_dim = head_dim self.num_attention_heads = num_attention_heads self.dropout = dropout inner_dim = head_dim * num_attention_heads self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) self.to_q = nn.Linear(hidden_size, inner_dim, bias=False) self.to_kv = nn.Linear(hidden_size, 2 * inner_dim, bias=False) self.to_out = nn.Linear(inner_dim, hidden_size, bias=False) def forward( self, query_states: torch.Tensor, kv_states: torch.Tensor, return_weights: bool = False, ): batch, q_seq, _ = query_states.shape kv_seq = kv_states.shape[1] q = ( self.to_q(query_states) .view(batch, q_seq, self.num_attention_heads, self.head_dim) .transpose(1, 2) .contiguous() ) kv = self.to_kv(kv_states).view( batch, kv_seq, 2, self.num_attention_heads, self.head_dim ) k, v = kv.permute(2, 0, 3, 1, 4).contiguous().unbind(dim=0) q = self.q_norm(q) k = self.k_norm(k) if return_weights: scale = q.shape[-1] ** -0.5 scores = (q * scale) @ k.transpose(-2, -1) weights = scores.softmax(dim=-1) if self.training and self.dropout > 0.0: attn = F.dropout(weights, p=self.dropout) @ v else: attn = weights @ v else: attn = F.scaled_dot_product_attention( q, k, v, 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) self.attention = Attention(hidden_size, head_dim, num_attention_heads, dropout) self.mlp_norm = nn.LayerNorm(hidden_size) self.mlp = GatedMLP(hidden_size, mlp_hidden_dim, hidden_size) def forward(self, hidden_states: torch.Tensor): hidden_states = hidden_states + self.attention(self.attn_norm(hidden_states)) hidden_states = hidden_states + 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) self.kv_norm = nn.LayerNorm(hidden_size) self.attention = CrossAttention( hidden_size, head_dim, num_attention_heads, dropout ) self.mlp_norm = nn.LayerNorm(hidden_size) self.mlp = GatedMLP(hidden_size, mlp_hidden_dim, hidden_size) def forward( self, query_states: torch.Tensor, kv_states: torch.Tensor, return_weights: bool = False, ): kv_input = ( torch.cat([query_states, kv_states], dim=1) if self.fused_kv else kv_states ) attn_out = self.attention( self.q_norm(query_states), self.kv_norm(kv_input), return_weights=return_weights, ) if return_weights: attn_out, weights = attn_out query_states = query_states + attn_out query_states = query_states + self.mlp(self.mlp_norm(query_states)) return (query_states, weights) if return_weights else query_states def _run_layer(layer, *inputs, training: bool): if training: return torch.utils.checkpoint.checkpoint(layer, *inputs, use_reentrant=False) return layer(*inputs) class LatentEncoder(nn.Module): def __init__( self, num_latents: int, num_layers: int, hidden_size: int, head_dim: int, num_attention_heads: int, mlp_hidden_dim: int, dropout: float = 0.0, ): super().__init__() self.latent_tokens = nn.Parameter(torch.zeros(num_latents, hidden_size)) self.pos_embed = nn.Parameter(torch.zeros(num_latents, hidden_size)) nn.init.trunc_normal_(self.latent_tokens, std=0.02) nn.init.trunc_normal_(self.pos_embed, std=0.02) self.cross_block = CrossAttentionBlock( hidden_size, head_dim, num_attention_heads, mlp_hidden_dim, fused_kv=False, dropout=dropout, ) self.layers = nn.ModuleList( [ TransformerBlock( hidden_size, head_dim, num_attention_heads, mlp_hidden_dim, dropout ) for _ in range(num_layers) ] ) def forward(self, patch_tokens: torch.Tensor, return_weights: bool = False): batch = patch_tokens.shape[0] latents = self.latent_tokens.expand(batch, -1, -1) out = self.cross_block(latents, patch_tokens, return_weights=return_weights) if return_weights: latents, weights = out else: latents, weights = out, None latents = latents + self.pos_embed for layer in self.layers: latents = _run_layer(layer, latents, training=self.training) return (latents, weights) if return_weights else latents class PatchDecoder(nn.Module): def __init__( self, num_layers: int, hidden_size: int, head_dim: int, num_attention_heads: int, mlp_hidden_dim: int, pos_feat_dim: int, dropout: float = 0.0, ): super().__init__() self.position_encoder = FourierPositionEncoder(pos_feat_dim, hidden_size) self.layers = nn.ModuleList( [ CrossAttentionBlock( hidden_size, head_dim, num_attention_heads, mlp_hidden_dim, fused_kv=True, dropout=dropout, ) for _ in range(num_layers) ] ) self.out_norm = nn.LayerNorm(hidden_size, elementwise_affine=False) def _coordinates( self, seq_length: int, device: torch.device, dtype: torch.dtype ) -> torch.Tensor: side = math.isqrt(seq_length) if side * side != seq_length: raise ValueError( f"PatchDecoder expects a square patch sequence, got length {seq_length}" ) positions = torch.linspace(-1.0, 1.0, side, device=device, dtype=dtype) y, x = torch.meshgrid(positions, positions, indexing="ij") return torch.stack([x, y], dim=-1).view(1, seq_length, 2) def forward(self, latents: torch.Tensor, target_seq_length: int) -> torch.Tensor: coordinates = self._coordinates( target_seq_length, latents.device, latents.dtype ) hidden_states = self.position_encoder(coordinates).expand( latents.shape[0], -1, -1 ) for layer in self.layers: hidden_states = _run_layer( layer, hidden_states, latents, training=self.training ) return self.out_norm(hidden_states) @dataclass class FineViTModelOutput(ModelOutput): decoded_patches: Optional[torch.FloatTensor] = None latents: Optional[torch.FloatTensor] = None latent_attention_weights: Optional[torch.FloatTensor] = None cls_register_hidden_states: Optional[torch.FloatTensor] = None patch_tokens: Optional[torch.FloatTensor] = None uncertainty: Optional[torch.FloatTensor] = None class FineViTModel(PreTrainedModel): config_class = FineViTConfig main_input_name = "pixel_values" def __init__(self, config: FineViTConfig): super().__init__(config) self.encoder = self._build_encoder(config) if hasattr(self.encoder, "embeddings"): self.encoder.embeddings.mask_token = None self.num_register_tokens = int( getattr(self.encoder.config, "num_register_tokens", 0) ) hidden_size = int(self.encoder.config.hidden_size) self.uncertainty_head = GatedMLP(hidden_size, hidden_size * 3, 1) self.latent_encoder = LatentEncoder( num_latents=config.num_latents, num_layers=config.latent_encoder_num_layers, hidden_size=hidden_size, head_dim=config.latent_encoder_head_dim, num_attention_heads=config.latent_encoder_num_attention_heads, mlp_hidden_dim=config.latent_encoder_mlp_hidden_dim, dropout=config.dropout, ) self.patch_decoder = PatchDecoder( num_layers=config.patch_decoder_num_layers, hidden_size=hidden_size, head_dim=config.patch_decoder_head_dim, num_attention_heads=config.patch_decoder_num_attention_heads, mlp_hidden_dim=config.patch_decoder_mlp_hidden_dim, pos_feat_dim=hidden_size, dropout=config.dropout, ) self.post_init() def _build_encoder(self, config: FineViTConfig): if getattr(config, "init_backbone_from_pretrained", True): return AutoModel.from_pretrained(config.backbone_model_name) if config.backbone_config is not None: backbone_config_dict = dict(config.backbone_config) model_type = backbone_config_dict.pop("model_type", None) if model_type is not None: backbone_config = AutoConfig.for_model( model_type, **backbone_config_dict ) return AutoModel.from_config(backbone_config) return AutoModel.from_pretrained(config.backbone_model_name) def post_init(self) -> None: for module in ( self.uncertainty_head, self.latent_encoder, self.patch_decoder, ): module.apply(self._init_weights) def _init_weights(self, module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.trunc_normal_(module.weight, std=self.config.initializer_range) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, (nn.LayerNorm, nn.RMSNorm)): if module.weight is not None: nn.init.ones_(module.weight) bias = getattr(module, "bias", None) if bias is not None: nn.init.zeros_(bias) def forward( self, pixel_values: torch.Tensor, return_dict: Optional[bool] = None, only_latents: bool = False, latents: Optional[torch.Tensor] = None, ): return_dict = ( return_dict if return_dict is not None else self.config.return_dict ) patch_offset = 1 + self.num_register_tokens hidden_states = self.encoder(pixel_values).last_hidden_state patch_tokens = hidden_states[:, patch_offset:] cls_register_hidden_states = hidden_states[:, :patch_offset] if latents is None: latents, latent_attention_weights = self.latent_encoder( hidden_states, return_weights=True ) latent_attention_weights = latent_attention_weights[ :, :, :, -patch_tokens.shape[1] : ] else: latent_attention_weights = None if only_latents: return latents decoded_patches = self.patch_decoder( latents, target_seq_length=patch_tokens.shape[1] ) uncertainty = self.uncertainty_head(patch_tokens) if not return_dict: return ( decoded_patches, latents, latent_attention_weights, cls_register_hidden_states, patch_tokens, uncertainty, ) return FineViTModelOutput( decoded_patches=decoded_patches, latents=latents, latent_attention_weights=latent_attention_weights, cls_register_hidden_states=cls_register_hidden_states, patch_tokens=patch_tokens, uncertainty=uncertainty, ) FineViTModel.register_for_auto_class("AutoModel")