Upload files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+aimv2_overview_light.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,237 @@
----
-license: apple-ascl
----
+---
+library_name: transformers
+license: apple-ascl
+metrics:
+- accuracy
+model-index:
+- name: aimv2-3B-patch14-336
+  results:
+  - dataset:
+      name: imagenet-1k
+      type: imagenet-1k
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 89.2
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: inaturalist-18
+      type: inaturalist-18
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 84.4
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: cifar10
+      type: cifar10
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 99.5
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: cifar100
+      type: cifar100
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 94.4
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: food101
+      type: food101
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 97.2
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: dtd
+      type: dtd
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 89.3
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: oxford-pets
+      type: oxford-pets
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 97.2
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: stanford-cars
+      type: stanford-cars
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 96.6
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: camelyon17
+      type: camelyon17
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 93.2
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: patch-camelyon
+      type: patch-camelyon
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 89.3
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: rxrx1
+      type: rxrx1
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 8.8
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: eurosat
+      type: eurosat
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 99.0
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: fmow
+      type: fmow
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 65.7
+      verified: false
+    task:
+      name: Classification
+      type: classification
+  - dataset:
+      name: domainnet-infographic
+      type: domainnet-infographic
+    metrics:
+    - name: Accuracy
+      type: accuracy
+      value: 74.0
+      verified: false
+    task:
+      name: Classification
+      type: classification
+pipeline_tag: image-feature-extraction
+tags:
+- vision
+- image-feature-extraction
+- mlx
+- pytorch
+---
+# Introduction
+[[`AIMv2 Paper`](#)] [[`BibTeX`](#citation)]
+
+We introduce the AIMv2 family of vision models pre-trained with a multimodal autoregressive objective.
+AIMv2 pre-training is simple and straightforward to train and scale effectively. Some AIMv2 highlights include:
+
+1. Outperforms OAI CLIP and SigLIP on the majority of multimodal understanding benchmarks.
+2. Outperforms DINOv2 on open-vocabulary object detection and referring expression comprehension.
+3. Exhibits strong recognition performance with AIMv2-3B achieving *89.5% on ImageNet using a frozen trunk*.
+
+<img src="aimv2_overview_light.png" alt="AIMv2 Overview"/>
+
+## Usage
+
+### PyTorch
+```python
+import requests
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+processor = AutoImageProcessor.from_pretrained(
+    "apple/aimv2-3B-patch14-336",
+)
+model = AutoModel.from_pretrained(
+    "apple/aimv2-3B-patch14-336",
+    trust_remote_code=True,
+)
+
+inputs = processor(images=image, return_tensors="pt")
+outputs = model(**inputs)
+```
+
+### JAX
+```python
+import requests
+from PIL import Image
+from transformers import AutoImageProcessor, FlaxAutoModel
+
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+
+processor = AutoImageProcessor.from_pretrained(
+    "apple/aimv2-3B-patch14-336",
+)
+model = FlaxAutoModel.from_pretrained(
+    "apple/aimv2-3B-patch14-336",
+    trust_remote_code=True,
+)
+
+inputs = processor(images=image, return_tensors="jax")
+outputs = model(**inputs)
+```
+
+## Citation
+If you find our work useful, please consider citing us as:
+```bibtex
+@misc{fini2024multimodal,
+  title         = {Multimodal Autoregressive Pre-training of Large Vision Encoders},
+  author        = {Enrico Fini and Mustafa Shukor and Xiujun Li and Philipp Dufter and Michal Klein and David Haldimann and Sai Aitharaju and Victor Guilherme Turrisi da Costa and Louis Béthune and Zhe Gan and Alexander T Toshev and Marcin Eichner and Moin Nabi and Yinfei Yang and Joshua M. Susskind and Alaaeldin El-Nouby},
+  year          = {2024},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.CV},
+}
+```
+

config.json

@@ -0,0 +1,25 @@
+{
+  "architectures": [
+    "AIMv2Model"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_aimv2.AIMv2Config",
+    "AutoModel": "modeling_aimv2.AIMv2Model",
+    "FlaxAutoModel": "modeling_flax_aimv2.FlaxAIMv2Model"
+  },
+  "hidden_size": 3072,
+  "image_size": 336,
+  "intermediate_size": 8192,
+  "model_type": "aimv2",
+  "num_attention_heads": 24,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 14,
+  "projection_dropout": 0.0,
+  "qkv_bias": false,
+  "rms_norm_eps": 1e-05,
+  "torch_dtype": "float32",
+  "transformers_version": "4.46.3",
+  "use_bias": false
+}

configuration_aimv2.py

@@ -0,0 +1,62 @@
+from typing import Any
+
+from transformers.configuration_utils import PretrainedConfig
+
+__all__ = ["AIMv2Config"]
+
+
+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](https://huggingface.co/apple/aimv2-large-patch14-224).
+
+    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_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"
+
+    def __init__(
+        self,
+        hidden_size: int = 1024,
+        intermediate_size: int = 2816,
+        num_hidden_layers: int = 24,
+        num_attention_heads: int = 8,
+        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_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

flax_model.msgpack

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

mlx_model.safetensors

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

model.safetensors

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

modeling_aimv2.py

@@ -0,0 +1,191 @@
+from typing import Optional, Tuple, Union
+
+import torch
+from .configuration_aimv2 import AIMv2Config
+from torch import nn
+from torch.nn import functional as F
+from transformers.modeling_outputs import BaseModelOutputWithNoAttention
+from transformers.modeling_utils import PreTrainedModel
+
+__all__ = ["AIMv2Model"]
+
+
+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: AIMv2Config):
+        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: AIMv2Config):
+        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: AIMv2Config):
+        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: AIMv2Config):
+        super().__init__()
+        dim = config.hidden_size
+
+        self.num_heads = config.num_attention_heads
+        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)
+
+        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: AIMv2Config):
+        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 AIMv2Transformer(nn.Module):
+    def __init__(self, config: AIMv2Config):
+        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):
+    config_class = AIMv2Config
+    base_model_prefix = "aimv2"
+    main_input_name = "pixel_values"
+    _supports_sdpa = True
+
+
+class AIMv2Model(AIMv2PretrainedModel):
+    def __init__(self, config: AIMv2Config):
+        super().__init__(config)
+        self.preprocessor = AIMv2ViTPreprocessor(config)
+        self.trunk = AIMv2Transformer(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
+        )
+
+        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,
+        )
+

modeling_flax_aimv2.py

@@ -0,0 +1,309 @@
+from typing import Any, Optional, Tuple, Union
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+from .configuration_aimv2 import AIMv2Config
+from flax.core import frozen_dict
+from transformers import FlaxPreTrainedModel
+from transformers.modeling_flax_outputs import FlaxBaseModelOutput
+
+__all__ = ["FlaxAIMv2Model"]
+
+
+class FlaxRMSNorm(nn.Module):
+    eps: float = 1e-6
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        dim = x.shape[-1]
+        scale = self.param("scale", nn.initializers.ones_init(), (dim,))
+        output = self._norm(x.astype(jnp.float32)).astype(x.dtype)
+        output = output * scale.astype(x.dtype)
+        return output
+
+    def _norm(self, x: jax.Array) -> jax.Array:
+        return x * jax.lax.rsqrt(jnp.power(x, 2).mean(-1, keepdims=True) + self.eps)
+
+
+class FlaxAIMv2SwiGLUFFN(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        hidden_features = self.config.intermediate_size
+        in_features = self.config.hidden_size
+        bias = self.config.use_bias
+
+        x1 = nn.Dense(hidden_features, use_bias=bias, dtype=self.dtype, name="fc1")(x)
+        x2 = nn.Dense(hidden_features, use_bias=bias, dtype=self.dtype, name="fc3")(x)
+        x = nn.silu(x1) * x2
+        x = nn.Dense(in_features, use_bias=bias, dtype=self.dtype, name="fc2")(x)
+        return x
+
+
+class FlaxAIMv2PatchEmbed(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        patch_size = (self.config.patch_size, self.config.patch_size)
+        x = x.transpose(0, 2, 3, 1)  # (N C H W) -> (N H W C)
+        x = nn.Conv(
+            self.config.hidden_size,
+            kernel_size=patch_size,
+            strides=patch_size,
+            padding=(0, 0),
+            dtype=self.dtype,
+            name="proj",
+        )(x)
+        x = jax.lax.collapse(x, 1, 3)  # (N, H * W, F)
+        x = FlaxRMSNorm(self.config.rms_norm_eps, name="norm")(x)
+        return x
+
+
+class FlaxAIMv2ViTPreprocessor(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(self, x: jax.Array) -> jax.Array:
+        tokens = FlaxAIMv2PatchEmbed(self.config, dtype=self.dtype, name="patchifier")(
+            x
+        )
+        _, N, _ = tokens.shape
+        pos_embed = self.param(
+            "pos_embed",
+            nn.initializers.normal(stddev=0.02),
+            (1, self.num_patches, self.config.hidden_size),
+        )
+        tokens = tokens + pos_embed[:, :N].astype(tokens.dtype)
+        return tokens
+
+    @property
+    def num_patches(self) -> int:
+        return (self.config.image_size // self.config.patch_size) ** 2
+
+
+class FlaxAIMv2Attention(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(
+        self,
+        x: jax.Array,
+        mask: Optional[jax.Array] = None,
+        deterministic: bool = True,
+        output_attentions: bool = False,
+    ) -> Tuple[jax.Array, Optional[jax.Array]]:
+        B, N, C = x.shape
+        dim, num_heads = self.config.hidden_size, self.config.num_attention_heads
+
+        qkv = nn.Dense(
+            dim * 3, use_bias=self.config.qkv_bias, dtype=self.dtype, name="qkv"
+        )(x)
+        qkv = qkv.reshape(B, N, 3, num_heads, C // num_heads).transpose(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn_weights = nn.dot_product_attention_weights(
+            q.swapaxes(-3, -2),  # [B, N, H, C]
+            k.swapaxes(-3, -2),
+            mask=mask,
+            deterministic=deterministic,
+            dtype=self.dtype,
+        )
+        attn_weights = nn.Dropout(
+            self.config.attention_dropout, deterministic=deterministic, name="attn_drop"
+        )(attn_weights)
+
+        x = (attn_weights @ v).swapaxes(1, 2).reshape(B, N, C)
+        x = nn.Dense(dim, use_bias=self.config.use_bias, dtype=self.dtype, name="proj")(
+            x
+        )
+        x = nn.Dropout(
+            self.config.projection_dropout,
+            deterministic=deterministic,
+            name="proj_drop",
+        )(x)
+        return (x, attn_weights) if output_attentions else (x, None)
+
+
+class FlaxAIMv2Block(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    def setup(self):
+        self.attn = FlaxAIMv2Attention(self.config, dtype=self.dtype, name="attn")
+        self.norm_1 = FlaxRMSNorm(self.config.rms_norm_eps, name="norm_1")
+        self.mlp = FlaxAIMv2SwiGLUFFN(self.config, dtype=self.dtype, name="mlp")
+        self.norm_2 = FlaxRMSNorm(self.config.rms_norm_eps, name="norm_2")
+
+    def __call__(
+        self,
+        x: jax.Array,
+        mask: Optional[jax.Array] = None,
+        deterministic: bool = True,
+        output_attentions: bool = False,
+    ) -> Tuple[jax.Array, Optional[jax.Array]]:
+        features, attention = self.attn(
+            self.norm_1(x),
+            mask,
+            deterministic=deterministic,
+            output_attentions=output_attentions,
+        )
+        x = x + features
+        x = x + self.mlp(self.norm_2(x))
+        return x, attention
+
+
+class FlaxAIMv2Transformer(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(
+        self,
+        tokens: jax.Array,
+        mask: Optional[jax.Array] = None,
+        deterministic: bool = True,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+    ) -> Tuple[
+        jax.Array, Optional[Tuple[jax.Array, ...]], Optional[Tuple[jax.Array, ...]]
+    ]:
+        hidden_states = () if output_hidden_states else None
+        attentions = () if output_attentions else None
+        for blk_id, block in enumerate(range(self.config.num_hidden_layers)):
+            tokens, attention = FlaxAIMv2Block(
+                self.config, dtype=self.dtype, name=f"layers_{blk_id}"
+            )(
+                tokens,
+                mask,
+                deterministic=deterministic,
+                output_attentions=output_attentions,
+            )
+            if output_hidden_states:
+                hidden_states += (tokens,)
+            if output_attentions:
+                attentions += (attention,)
+        tokens = FlaxRMSNorm(self.config.rms_norm_eps, name="post_trunk_norm")(tokens)
+        return tokens, hidden_states, attentions
+
+
+class FlaxAIMv2Module(nn.Module):
+    config: AIMv2Config
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(
+        self,
+        x: jax.Array,
+        mask: Optional[jax.Array] = None,
+        deterministic: bool = True,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+    ) -> Tuple[
+        jax.Array, Optional[Tuple[jax.Array, ...]], Optional[Tuple[jax.Array, ...]]
+    ]:
+        x = FlaxAIMv2ViTPreprocessor(
+            self.config, dtype=self.dtype, name="preprocessor"
+        )(x)
+        x, hidden_states, attentions = FlaxAIMv2Transformer(
+            self.config, dtype=self.dtype, name="trunk"
+        )(
+            x,
+            mask,
+            deterministic=deterministic,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+        return x, hidden_states, attentions
+
+
+class FlaxAIMv2PretrainedModel(FlaxPreTrainedModel):
+    config_class = AIMv2Config
+    base_model_prefix = "aimv2"
+    main_input_name = "pixel_values"
+
+    def __init__(
+        self,
+        config: AIMv2Config,
+        input_shape: Optional[Tuple[int, int, int, int]] = None,  # [B, C, H, W]
+        dtype: jnp.dtype = jnp.float32,
+        **kwargs: Any,
+    ):
+        if input_shape is None:
+            input_shape = (1, 3, config.image_size, config.image_size)
+        super().__init__(
+            config,
+            module=FlaxAIMv2Module(config, dtype=dtype),
+            input_shape=input_shape,
+            dtype=dtype,
+            **kwargs,
+        )
+
+    def init_weights(
+        self,
+        rng: jax.Array,
+        input_shape: Tuple[int, ...],
+        params: Optional[frozen_dict.FrozenDict] = None,
+    ) -> frozen_dict.FrozenDict:
+        del params
+        input_pixels = jnp.empty(input_shape)
+        params = self.module.init(rng, input_pixels, deterministic=True)
+        return params["params"]
+
+
+class FlaxAIMv2Model(FlaxAIMv2PretrainedModel):
+    def __call__(
+        self,
+        pixel_values: jax.Array,
+        params: Optional[frozen_dict.FrozenDict] = None,
+        mask: Optional[jax.Array] = None,
+        dropout_rng: Optional[jax.Array] = None,
+        deterministic: bool = True,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[
+        Tuple[jax.Array],
+        Tuple[jax.Array, Tuple[jax.Array, ...]],
+        Tuple[jax.Array, Tuple[jax.Array, ...], Tuple[jax.Array, ...]],
+        FlaxBaseModelOutput,
+    ]:
+        if params is None:
+            params = self.params
+        if output_attentions is None:
+            output_attentions = self.config.output_attentions
+        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
+
+        rngs = None if deterministic else {"dropout": dropout_rng}
+
+        x, hidden_states, attentions = self.module.apply(
+            {"params": params},
+            pixel_values,
+            mask,
+            rngs=rngs,
+            deterministic=deterministic,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        if not return_dict:
+            res = (x,)
+            res += (hidden_states,) if output_hidden_states else ()
+            res += (attentions,) if output_attentions else ()
+            return res
+
+        return FlaxBaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+

preprocessor_config.json

@@ -0,0 +1,27 @@
+{
+  "crop_size": {
+    "height": 336,
+    "width": 336
+  },
+  "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
+  ],
+  "resample": 3,
+  "rescale_factor": 0.00392156862745098,
+  "size": {
+    "shortest_edge": 336
+  }
+}