Upload folder using huggingface_hub

config.json

@@ -0,0 +1,113 @@
+{
+  "architectures": [
+    "FineViTModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.FineViTConfig",
+    "AutoModel": "finevit.FineViTModel"
+  },
+  "backbone_config": {
+    "_name_or_path": "facebook/dinov2-with-registers-base",
+    "add_cross_attention": false,
+    "apply_layernorm": true,
+    "architectures": [
+      "Dinov2WithRegistersModel"
+    ],
+    "attention_probs_dropout_prob": 0.0,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "drop_path_rate": 0.0,
+    "dtype": "float32",
+    "eos_token_id": null,
+    "finetuning_task": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 518,
+    "initializer_range": 0.02,
+    "interpolate_antialias": true,
+    "interpolate_offset": 0.0,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "layerscale_value": 1.0,
+    "mlp_ratio": 4,
+    "model_type": "dinov2_with_registers",
+    "num_attention_heads": 12,
+    "num_channels": 3,
+    "num_hidden_layers": 12,
+    "num_register_tokens": 4,
+    "out_features": [
+      "stage12"
+    ],
+    "out_indices": [
+      12
+    ],
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qkv_bias": true,
+    "reshape_hidden_states": true,
+    "return_dict": true,
+    "sep_token_id": null,
+    "stage_names": [
+      "stem",
+      "stage1",
+      "stage2",
+      "stage3",
+      "stage4",
+      "stage5",
+      "stage6",
+      "stage7",
+      "stage8",
+      "stage9",
+      "stage10",
+      "stage11",
+      "stage12"
+    ],
+    "task_specific_params": null,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "torchscript": false,
+    "transformers_version": "5.8.0",
+    "use_bfloat16": false,
+    "use_swiglu_ffn": false
+  },
+  "backbone_model_name": "facebook/dinov2-with-registers-base",
+  "checkpoint_epoch": 2,
+  "checkpoint_step": 50000,
+  "dropout": 0.0,
+  "dtype": "bfloat16",
+  "image_size": 224,
+  "init_backbone_from_pretrained": false,
+  "initializer_range": 0.02,
+  "latent_encoder_head_dim": 64,
+  "latent_encoder_mlp_hidden_dim": 3072,
+  "latent_encoder_num_attention_heads": 12,
+  "latent_encoder_num_layers": 12,
+  "model_type": "finevit",
+  "num_latents": 64,
+  "patch_decoder_head_dim": 64,
+  "patch_decoder_mlp_hidden_dim": 3072,
+  "patch_decoder_num_attention_heads": 12,
+  "patch_decoder_num_layers": 8,
+  "teacher_encoder_type": "dinov2_with_registers",
+  "train_encoder": true,
+  "transformers_version": "5.8.0"
+}

config.py

@@ -0,0 +1,48 @@
+from transformers import PretrainedConfig
+
+
+class FineViTConfig(PretrainedConfig):
+    model_type = "finevit"
+
+    def __init__(
+        self,
+        backbone_model_name: str = "facebook/dinov2-with-registers-base",
+        teacher_encoder_type: str = "dinov2_with_registers",
+        backbone_config: dict | None = None,
+        num_latents: int = 64,
+        latent_encoder_num_layers: int = 12,
+        latent_encoder_head_dim: int = 64,
+        latent_encoder_num_attention_heads: int = 12,
+        latent_encoder_mlp_hidden_dim: int = 3072,
+        patch_decoder_num_layers: int = 8,
+        patch_decoder_head_dim: int = 64,
+        patch_decoder_num_attention_heads: int = 12,
+        patch_decoder_mlp_hidden_dim: int = 3072,
+        dropout: float = 0.0,
+        image_size: int = 224,
+        train_encoder: bool = False,
+        init_backbone_from_pretrained: bool = True,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.backbone_model_name = backbone_model_name
+        self.teacher_encoder_type = teacher_encoder_type
+        self.backbone_config = backbone_config
+        self.num_latents = num_latents
+        self.latent_encoder_num_layers = latent_encoder_num_layers
+        self.latent_encoder_head_dim = latent_encoder_head_dim
+        self.latent_encoder_num_attention_heads = latent_encoder_num_attention_heads
+        self.latent_encoder_mlp_hidden_dim = latent_encoder_mlp_hidden_dim
+        self.patch_decoder_num_layers = patch_decoder_num_layers
+        self.patch_decoder_head_dim = patch_decoder_head_dim
+        self.patch_decoder_num_attention_heads = patch_decoder_num_attention_heads
+        self.patch_decoder_mlp_hidden_dim = patch_decoder_mlp_hidden_dim
+        self.dropout = dropout
+        self.image_size = image_size
+        self.train_encoder = train_encoder
+        self.init_backbone_from_pretrained = init_backbone_from_pretrained
+        self.initializer_range = initializer_range
+
+
+FineViTConfig.register_for_auto_class("AutoConfig")

finevit.py

@@ -0,0 +1,448 @@
+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 finevit.models.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")

model.safetensors

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

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 224,
+    "width": 224
+  },
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_rescale": true,
+  "do_resize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_processor_type": "BitImageProcessor",
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 256
+  }
+}