PyTorch code (#3)

- Add LiT model (ae9d0c031c7a34b756da73f518406e716b55a4c0)
- Fix importw (1943b4b2f19a5f67ff9a4591215c8366a80dc3bb)
- Upload weights (5fd59c5def3a2f5c450bdf63f40708b34420f3a7)
- Handle dictionary in confgw (fb574928240d6acf48aaee63ef759a26a2c1ce4a)

README.md

@@ -20,7 +20,39 @@ AIMv2 pre-training is simple and straightforward to train and to scale effective
 <img src="aimv2_overview_light.png" alt="AIMv2 Overview"/>
 
 ## Usage
-Under construction. Please consider using the models in the [ml-aim](https://github.com/apple/ml-aim) repository.
+
+### PyTorch
+```python
+import requests
+from PIL import Image
+from transformers import AutoProcessor, AutoModel
+
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+text = ["Picture of a dog.", "Picture of a cat.", "Picture of a horse."]
+
+processor = AutoProcessor.from_pretrained(
+    "apple/aimv2-large-patch14-224-lit",
+)
+model = AutoModel.from_pretrained(
+    "apple/aimv2-large-patch14-224-lit",
+    trust_remote_code=True,
+)
+
+inputs = processor(
+    images=image,
+    text=text,
+    add_special_tokens=True,
+    truncation=True,
+    padding=True,
+    return_tensors="pt",
+)
+outputs = model(**inputs)
+probs = outputs.logits_per_image.softmax(dim=-1)
+```
+
+### JAX
+Under construction.
 
 ## Citation
 If you find our work useful, please consider citing us as:

config.json

@@ -0,0 +1,169 @@
+{
+  "architectures": [
+    "AIMv2Model"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_aimv2.AIMv2Config",
+    "AutoModel": "modeling_aimv2.AIMv2Model"
+  },
+  "init_temperature": 0.07,
+  "max_logit_scale": 100.0,
+  "model_type": "aimv2",
+  "projection_dim": 768,
+  "text_config": {
+    "_attn_implementation_autoset": true,
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 49407,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "intermediate_size": 2048,
+    "is_causal": true,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_context_length": 77,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "aimv2",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 6,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "problem_type": null,
+    "projection_dropout": 0.0,
+    "pruned_heads": {},
+    "qkv_bias": false,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_bias": false,
+    "vocab_size": 49408
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "vision_config": {
+    "_attn_implementation_autoset": true,
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 224,
+    "intermediate_size": 2816,
+    "is_causal": false,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "aimv2",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 8,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "num_queries": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "projection_dropout": 0.0,
+    "pruned_heads": {},
+    "qkv_bias": false,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_bias": false
+  }
+}

configuration_aimv2.py

@@ -0,0 +1,174 @@
+from typing import Any, Dict, Optional, Union
+
+from transformers.configuration_utils import PretrainedConfig
+
+__all__ = ["AIMv2VisionConfig", "AIMv2TextConfig", "AIMv2Config"]
+
+
+class AIMv2VisionConfig(PretrainedConfig):
+    """This is the configuration class to store the configuration of an [`AIMv2VisionModel`].
+
+    Instantiating a configuration with the defaults will yield a similar configuration
+    to that of the [apple/aimv2-large-patch14-224-lit](https://huggingface.co/apple/aimv2-large-patch14-224-lit).
+
+    Args:
+        hidden_size: Dimension of the hidden representations.
+        intermediate_size: Dimension of the SwiGLU representations.
+        num_hidden_layers: Number of hidden layers in the Transformer.
+        num_attention_heads: Number of attention heads for each attention layer
+            in the Transformer.
+        num_queries: Number of learnable queries for the attention-pooling head.
+        num_channels: Number of input channels.
+        image_size: Image size.
+        patch_size: Patch size.
+        rms_norm_eps: Epsilon value used for the RMS normalization layer.
+        attention_dropout: Dropout ratio for attention probabilities.
+        projection_dropout: Dropout ratio for the projection layer after the attention.
+        qkv_bias: Whether to add a bias to the queries, keys and values.
+        use_bias: Whether to add a bias in the feed-forward and projection layers.
+        kwargs: Keyword arguments for the [`PretrainedConfig`].
+    """
+
+    model_type: str = "aimv2"
+    base_config_key: str = "vision_config"
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        intermediate_size: int = 2816,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 8,
+        num_queries: int = 1,
+        num_channels: int = 3,
+        image_size: int = 224,
+        patch_size: int = 14,
+        rms_norm_eps: float = 1e-5,
+        attention_dropout: float = 0.0,
+        projection_dropout: float = 0.0,
+        qkv_bias: bool = False,
+        use_bias: bool = False,
+        **kwargs: Any,
+    ):
+        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_queries = num_queries
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.attention_dropout = attention_dropout
+        self.rms_norm_eps = rms_norm_eps
+
+        self.projection_dropout = projection_dropout
+        self.qkv_bias = qkv_bias
+        self.use_bias = use_bias
+        self.is_causal = False
+
+
+class AIMv2TextConfig(PretrainedConfig):
+    """This is the configuration class to store the configuration of an [`AIMv2TextModel`].
+
+    Instantiating a configuration with the defaults will yield a similar configuration
+    to that of the [apple/aimv2-large-patch14-224-lit](https://huggingface.co/apple/aimv2-large-patch14-224-lit).
+
+    Args:
+        vocab_size: Size of the vocabulary.
+        hidden_size: Dimension of the hidden representations.
+        intermediate_size: Dimension of the SwiGLU representations.
+        num_hidden_layers: Number of hidden layers in the Transformer.
+        num_attention_heads: Number of attention heads for each attention layer
+            in the Transformer.
+        rms_norm_eps: Epsilon value used for the RMS normalization layer.
+        attention_dropout: Dropout ratio for attention probabilities.
+        projection_dropout: Dropout ratio for the projection layer after the attention.
+        qkv_bias: Whether to add a bias to the queries, keys and values.
+        use_bias: Whether to add a bias in the feed-forward and projection layers.
+        eos_token_id: End-of-sequence token id.
+        max_context_length: Maximum number of tokens for the context.
+        kwargs: Keyword arguments for the [`PretrainedConfig`].
+    """
+
+    model_type: str = "aimv2"
+    base_config_key: str = "text_config"
+
+    def __init__(
+        self,
+        vocab_size: int = 49408,
+        hidden_size: int = 768,
+        intermediate_size: int = 2048,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 6,
+        rms_norm_eps: float = 1e-5,
+        attention_dropout: float = 0.0,
+        projection_dropout: float = 0.0,
+        qkv_bias: bool = False,
+        use_bias: bool = False,
+        eos_token_id: int = 49407,
+        max_context_length: int = 77,
+        **kwargs: Any,
+    ):
+        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.attention_dropout = attention_dropout
+        self.rms_norm_eps = rms_norm_eps
+
+        self.projection_dropout = projection_dropout
+        self.qkv_bias = qkv_bias
+        self.use_bias = use_bias
+        self.vocab_size = vocab_size
+        self.max_context_length = max_context_length
+        self.eos_token_id = eos_token_id
+        self.is_causal = True
+
+
+class AIMv2Config(PretrainedConfig):
+    """This is the configuration class to store the configuration of an [`AIMv2Model`].
+
+    Instantiating a configuration with the defaults will yield a similar configuration
+    to that of the [apple/aimv2-large-patch14-224-lit](https://huggingface.co/apple/aimv2-large-patch14-224-lit).
+
+    Args:
+        vision_config: Vision config.
+        text_config: Text config.
+        projection_dim: Dimension of the image and text projection layers.
+        kwargs: Keyword arguments for the [`PretrainedConfig`].
+    """
+
+    model_type = "aimv2"
+    is_composition: bool = True
+    sub_configs: Dict[str, PretrainedConfig] = {
+        "vision_config": AIMv2VisionConfig,
+        "text_config": AIMv2TextConfig,
+    }
+
+    def __init__(
+        self,
+        vision_config: Optional[Union[AIMv2VisionConfig, Dict[str, Any]]] = None,
+        text_config: Optional[Union[AIMv2TextConfig, Dict[str, Any]]] = None,
+        projection_dim: int = 768,
+        init_temperature: float = 0.07,
+        max_logit_scale: float = 100.0,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        if vision_config is None:
+            vision_config = AIMv2VisionConfig()
+        elif isinstance(vision_config, dict):
+            vision_config = AIMv2VisionConfig(**vision_config)
+
+        if text_config is None:
+            text_config = AIMv2TextConfig()
+        elif isinstance(text_config, dict):
+            text_config = AIMv2TextConfig(**text_config)
+
+        self.vision_config = vision_config
+        self.text_config = text_config
+        self.projection_dim = projection_dim
+
+        self.init_temperature = init_temperature
+        self.max_logit_scale = max_logit_scale

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c218dc23f746407cacb100506d83a1c5f169db11ff3cf141853ac28771af3222
-size 1746752308
+oid sha256:72c8b9fe2376ac8a5b7d15ea3175325447e819763c2372cf88596260bbcd9583
+size 1746752340

modeling_aimv2.py

@@ -0,0 +1,442 @@
+from typing import Optional, Tuple, Union
+
+import torch
+import dataclasses
+import math
+
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.utils import ModelOutput
+
+from .configuration_aimv2 import AIMv2Config, AIMv2VisionConfig, AIMv2TextConfig
+from torch import nn
+from torch.nn import functional as F
+from transformers.modeling_outputs import BaseModelOutputWithNoAttention
+from transformers.modeling_utils import PreTrainedModel
+
+__all__ = ["AIMv2VisionModel", "AIMv2TextModel", "AIMv2Model"]
+
+AIMv2VisionOrTextConfig = Union[AIMv2VisionConfig, AIMv2TextConfig]
+
+
+@dataclasses.dataclass
+class AIMv2Output(ModelOutput):
+    logits_per_image: torch.Tensor
+    logits_per_text: Optional[torch.Tensor] = None
+    image_features: Optional[torch.Tensor] = None
+    text_features: Optional[torch.Tensor] = None
+    vision_output: Optional[BaseModelOutputWithNoAttention] = None
+    text_output: Optional[BaseModelOutputWithNoAttention] = None
+
+
+class AIMv2TextPreprocessor(nn.Module):
+    def __init__(self, config: AIMv2TextConfig):
+        super().__init__()
+        self.max_context_length = config.max_context_length
+        self.eos_token_id = config.eos_token_id
+
+        self.text_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.positional_embedding = nn.Parameter(
+            torch.zeros(self.max_context_length, config.hidden_size)
+        )
+
+    def forward(self, input_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        _, N = input_ids.shape
+        max_len = min(N, self.max_context_length)
+        eos_token_mask = input_ids == self.eos_token_id
+        tokens = self.text_embedding(input_ids)
+        tokens = tokens[:, :max_len] + self.positional_embedding[:max_len].unsqueeze(0)
+        return tokens, eos_token_mask
+
+
+class AIMv2ExtractEOS(nn.Module):
+    def forward(
+        self, tokens: torch.Tensor, eos_token_mask: torch.Tensor
+    ) -> torch.Tensor:
+        B, _, D = tokens.shape
+        eos_token_mask = torch.argmax(eos_token_mask.float(), dim=-1)
+        assert eos_token_mask.shape == (B,)
+        eos_token_mask = eos_token_mask.reshape(B, 1, 1).expand(B, 1, D)
+        eos_token = torch.gather(tokens, 1, eos_token_mask)
+        eos_token = eos_token.squeeze(1)
+        return eos_token
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = self._norm(x.float()).type_as(x)
+        return output * self.weight
+
+    def extra_repr(self) -> str:
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+
+class AIMv2SwiGLUFFN(nn.Module):
+    def __init__(self, config: AIMv2VisionOrTextConfig):
+        super().__init__()
+        hidden_features = config.intermediate_size
+        in_features = config.hidden_size
+        bias = config.use_bias
+
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.fc2 = nn.Linear(hidden_features, in_features, bias=bias)
+        self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.fc1(x)) * self.fc3(x)
+        x = self.fc2(x)
+        return x
+
+
+class AIMv2PatchEmbed(nn.Module):
+    def __init__(self, config: AIMv2VisionOrTextConfig):
+        super().__init__()
+        self.proj = nn.Conv2d(
+            config.num_channels,
+            config.hidden_size,
+            kernel_size=(config.patch_size, config.patch_size),
+            stride=(config.patch_size, config.patch_size),
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x
+
+
+class AIMv2ViTPreprocessor(nn.Module):
+    def __init__(self, config: AIMv2VisionConfig):
+        super().__init__()
+        num_patches = (config.image_size // config.patch_size) ** 2
+
+        self.patchifier = AIMv2PatchEmbed(config)
+        self.pos_embed = nn.Parameter(torch.zeros((1, num_patches, config.hidden_size)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        tokens = self.patchifier(x)
+        _, N, _ = tokens.shape
+        pos_embed = self.pos_embed.to(tokens.device)
+        tokens = tokens + pos_embed[:, :N]
+        return tokens
+
+
+class AIMv2Attention(nn.Module):
+    def __init__(self, config: AIMv2VisionOrTextConfig):
+        super().__init__()
+        dim = config.hidden_size
+
+        self.num_heads = config.num_attention_heads
+        self.is_causal = config.is_causal
+        self.qkv = nn.Linear(dim, dim * 3, bias=config.qkv_bias)
+        self.attn_drop = nn.Dropout(config.attention_dropout)
+        self.proj = nn.Linear(dim, dim, bias=config.use_bias)
+        self.proj_drop = nn.Dropout(config.projection_dropout)
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)
+
+        if mask is None:
+            x = F.scaled_dot_product_attention(q, k, v, is_causal=self.is_causal)
+        else:
+            mask_converter = AttentionMaskConverter(self.is_causal)
+            mask = mask_converter.to_4d(
+                mask, key_value_length=N, query_length=N, dtype=q.dtype
+            )
+            x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+
+        x = x.transpose(1, 2).contiguous().reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class AIMv2Block(nn.Module):
+    def __init__(self, config: AIMv2VisionOrTextConfig):
+        super().__init__()
+        self.attn = AIMv2Attention(config)
+        self.norm_1 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = AIMv2SwiGLUFFN(config)
+        self.norm_2 = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm_1(x), mask)
+        x = x + self.mlp(self.norm_2(x))
+        return x
+
+
+class AIMv2AttentionPoolingHead(nn.Module):
+    def __init__(self, config: AIMv2VisionConfig):
+        super().__init__()
+        dim = config.hidden_size
+        qkv_bias = config.qkv_bias
+
+        self.num_heads = config.num_attention_heads
+        self.num_queries = config.num_queries
+
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+        self.cls_token = nn.Parameter(torch.randn(1, self.num_queries, dim) * 0.02)
+        self.linear = nn.Linear(dim, dim, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        cls_token = self.cls_token.expand(B, -1, -1)
+
+        q = cls_token.reshape(
+            B, self.num_queries, self.num_heads, C // self.num_heads
+        ).permute(0, 2, 1, 3)
+        k = (
+            self.k(x)
+            .reshape(B, N, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+        v = (
+            self.v(x)
+            .reshape(B, N, self.num_heads, C // self.num_heads)
+            .permute(0, 2, 1, 3)
+        )
+
+        x_cls = F.scaled_dot_product_attention(q, k, v)
+        x_cls = x_cls.transpose(1, 2).reshape(B, self.num_queries, C)
+        x_cls = x_cls.mean(dim=1)
+
+        out = self.linear(x_cls)
+        return out
+
+
+class AIMv2Transformer(nn.Module):
+    def __init__(self, config: AIMv2VisionOrTextConfig):
+        super().__init__()
+        self.blocks = nn.ModuleList(
+            [AIMv2Block(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.post_trunk_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]]]:
+        hidden_states = () if output_hidden_states else None
+        for block in self.blocks:
+            tokens = block(tokens, mask)
+            if output_hidden_states:
+                hidden_states += (tokens,)
+        tokens = self.post_trunk_norm(tokens)
+        return tokens, hidden_states
+
+
+class AIMv2PretrainedModel(PreTrainedModel):
+    base_model_prefix = "aimv2"
+    _supports_sdpa = True
+
+
+class AIMv2VisionModel(AIMv2PretrainedModel):
+    config_class = AIMv2VisionConfig
+    main_input_name = "pixel_values"
+    _no_split_modules = ["AIMv2ViTPreprocessor", "AIMv2Block"]
+
+    def __init__(self, config: AIMv2VisionConfig):
+        super().__init__(config)
+        self.preprocessor = AIMv2ViTPreprocessor(config)
+        self.trunk = AIMv2Transformer(config)
+        self.head = AIMv2AttentionPoolingHead(config)
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[
+        Tuple[torch.Tensor],
+        Tuple[torch.Tensor, Tuple[torch.Tensor, ...]],
+        BaseModelOutputWithNoAttention,
+    ]:
+        if output_hidden_states is None:
+            output_hidden_states = self.config.output_hidden_states
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        x = self.preprocessor(pixel_values)
+        x, hidden_states = self.trunk(
+            x, mask, output_hidden_states=output_hidden_states
+        )
+        x = self.head(x)
+
+        if not return_dict:
+            res = (x,)
+            res += (hidden_states,) if output_hidden_states else ()
+            return res
+
+        return BaseModelOutputWithNoAttention(
+            last_hidden_state=x,
+            hidden_states=hidden_states,
+        )
+
+
+class AIMv2TextModel(AIMv2PretrainedModel):
+    config_class = AIMv2TextConfig
+    main_input_name = "input_ids"
+    _no_split_modules = ["AIMv2TextPreprocessor", "AIMv2Block"]
+
+    def __init__(self, config: AIMv2TextConfig):
+        super().__init__(config)
+        self.preprocessor = AIMv2TextPreprocessor(config)
+        self.trunk = AIMv2Transformer(config)
+        self.head = AIMv2ExtractEOS()
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[
+        Tuple[torch.Tensor],
+        Tuple[torch.Tensor, Tuple[torch.Tensor, ...]],
+        BaseModelOutputWithNoAttention,
+    ]:
+        if output_hidden_states is None:
+            output_hidden_states = self.config.output_hidden_states
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        x, eos_token_mask = self.preprocessor(pixel_values)
+        x, hidden_states = self.trunk(
+            x, mask, output_hidden_states=output_hidden_states
+        )
+        x = self.head(x, eos_token_mask)
+
+        if not return_dict:
+            res = (x,)
+            res += (hidden_states,) if output_hidden_states else ()
+            return res
+
+        return BaseModelOutputWithNoAttention(
+            last_hidden_state=x,
+            hidden_states=hidden_states,
+        )
+
+
+class AIMv2Model(AIMv2PretrainedModel):
+    config_class = AIMv2Config
+    main_input_name = ["input_ids", "pixel_values"]
+    _no_split_modules = ["AIMv2ViTPreprocessor", "AIMv2TextPreprocessor", "AIMv2Block"]
+
+    def __init__(self, config: AIMv2Config):
+        super().__init__(config)
+        self.image_encoder = AIMv2VisionModel(config.vision_config)
+        self.text_encoder = AIMv2TextModel(config.text_config)
+
+        self.image_projector = nn.Linear(
+            config.vision_config.hidden_size, config.projection_dim, bias=False
+        )
+        self.text_projector = nn.Linear(
+            config.text_config.hidden_size, config.projection_dim, bias=False
+        )
+
+        self.log_logit_scale = nn.Parameter(
+            torch.full([], fill_value=math.log(1.0 / config.init_temperature))
+        )
+        self.max_log_logit_scale = math.log(config.max_logit_scale)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        pixel_values: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[
+        Tuple[
+            torch.Tensor,
+            torch.Tensor,
+            torch.Tensor,
+            torch.Tensor,
+            Union[Tuple[torch.Tensor, ...], BaseModelOutputWithNoAttention],
+            Union[Tuple[torch.Tensor, ...], BaseModelOutputWithNoAttention],
+        ],
+        AIMv2Output,
+    ]:
+        if return_dict is None:
+            return_dict = self.config.use_return_dict
+
+        image_out = self.image_encoder(
+            pixel_values,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        image_features = image_out.last_hidden_state if return_dict else image_out[0]
+        image_features = self.image_projector(image_features)
+        image_features = F.normalize(image_features, p=2, dim=-1)
+
+        text_out = self.text_encoder(
+            input_ids,
+            mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        text_features = text_out.last_hidden_state if return_dict else text_out[0]
+        text_features = self.text_projector(text_features)
+        text_features = F.normalize(text_features, p=2, dim=-1)
+
+        logit_scale = self.log_logit_scale.clamp(0.0, self.max_log_logit_scale).exp()
+        logits_per_text = (logit_scale * text_features) @ image_features.t()
+        logits_per_image = logits_per_text.t()
+
+        if not return_dict:
+            return (
+                logits_per_image,
+                logits_per_text,
+                image_features,
+                text_features,
+                image_out,
+                text_out,
+            )
+
+        return AIMv2Output(
+            logits_per_image=logits_per_image,
+            logits_per_text=logits_per_text,
+            image_features=image_features,
+            text_features=text_features,
+            vision_output=image_out,
+            text_output=text_out,
+        )
+
+    def get_image_features(
+        self,
+        input_pixels: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        out = self.image_encoder(input_pixels, mask=attention_mask, return_dict=True)
+        image_features = self.image_projector(out.last_hidden_state)
+        return image_features
+
+    def get_text_features(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        out = self.text_encoder(input_ids, mask=attention_mask, return_dict=True)
+        text_features = self.text_projector(out.last_hidden_state)
+        return text_features

preprocessor_config.json

@@ -0,0 +1,28 @@
+{
+  "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.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_processor_type": "CLIPImageProcessor",
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "processor_class": "CLIPProcessor",
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 224
+  }
+}

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<start_of_text>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<end_of_text>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<end_of_text>",
+  "unk_token": {
+    "content": "<end_of_text>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,29 @@
+{
+  "added_tokens_decoder": {
+    "49406": {
+      "content": "<start_of_text>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "49407": {
+      "content": "<end_of_text>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<start_of_text>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<end_of_text>",
+  "errors": "replace",
+  "model_max_length": 77,
+  "pad_token": "<end_of_text>",
+  "processor_class": "CLIPProcessor",
+  "tokenizer_class": "CLIPTokenizer",
+  "unk_token": "<end_of_text>"
+}