finevit.py

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")