tips v1 l14

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -0,0 +1,113 @@
+---
+license: apache-2.0
+tags:
+- vision
+- image-text
+- contrastive-learning
+- zero-shot
+- feature-extraction
+- arxiv:2410.16512
+library_name: transformers
+pipeline_tag: zero-shot-image-classification
+---
+
+# TIPS — L/14 (v1)
+
+TIPS (Text-Image Pre-training with Spatial awareness, ICLR 2025) is a family of contrastive vision-language models that produce spatially rich image features aligned with text embeddings. This is the original (v1) L/14 release with 304M vision params and 184M text params, converted from the [official checkpoints](https://github.com/google-deepmind/tips).
+
+| Variant | Vision params | Text params | Embed dim | Resolution |
+|---------|---------------|-------------|-----------|------------|
+| [S/14](https://huggingface.co/google/tipsv1-s14) | 22M | 34M | 384 | 448 |
+| [B/14](https://huggingface.co/google/tipsv1-b14) | 86M | 110M | 768 | 448 |
+| [L/14](https://huggingface.co/google/tipsv1-l14) | 304M | 184M | 1024 | 448 |
+| [So400m/14](https://huggingface.co/google/tipsv1-so400m14) | 413M | 448M | 1152 | 448 |
+| [g/14](https://huggingface.co/google/tipsv1-g14) | 1.1B | 389M | 1536 | 448 |
+| [g/14 low-res](https://huggingface.co/google/tipsv1-g14-lowres) | 1.1B | 389M | 1536 | 224 |
+
+## Usage
+
+```bash
+pip install transformers torch torchvision sentencepiece scikit-learn requests
+```
+
+### Load the model
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("google/tipsv1-l14", trust_remote_code=True)
+model.eval()
+```
+
+### Encode images
+
+Images should be tensors in `[0, 1]` range (just `ToTensor()`, no ImageNet normalization).
+
+```python
+import requests
+from PIL import Image
+from torchvision import transforms
+
+url = "https://huggingface.co/spaces/google/TIPSv2/resolve/main/examples/zeroseg/pascal_context_00049_image.png"
+image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+transform = transforms.Compose([transforms.Resize((448, 448)), transforms.ToTensor()])
+pixel_values = transform(image).unsqueeze(0)
+
+out = model.encode_image(pixel_values)
+print(out.cls_token.shape)     # (1, 1, 1024) — global image embedding
+print(out.patch_tokens.shape)  # (1, 1024, 1024) — per-patch spatial features
+```
+
+The second CLS token (`out.register_tokens`) was trained on synthetic captions; the first (`out.cls_token`) on web alt-text, and is the one aligned with the text tower.
+
+### Encode text
+
+```python
+text_emb = model.encode_text(["a photo of a bus", "a photo of a dog"])
+print(text_emb.shape)  # (2, 1024) — one embedding per query
+```
+
+### Zero-shot classification
+
+```python
+import torch.nn.functional as F
+
+classes = ["bus", "car", "dog", "cat"]
+cls = F.normalize(out.cls_token[:, 0, :], dim=-1)
+text_emb = F.normalize(model.encode_text(classes), dim=-1)
+similarity = cls @ text_emb.T
+print(classes[similarity.argmax()])  # predicted class
+```
+
+### Visualize spatial features
+
+```python
+import numpy as np
+from sklearn.decomposition import PCA
+
+feat = out.patch_tokens[0].detach().cpu().numpy()
+rgb = PCA(n_components=3, whiten=True).fit_transform(feat).reshape(32, 32, 3)
+rgb = 1 / (1 + np.exp(-2.0 * rgb))  # sigmoid for [0, 1] range with good contrast
+```
+
+## Model details
+
+- ViT-L/14 vision encoder (24 layers, patch size 14, two CLS tokens) + 12-layer transformer text encoder
+- Native resolution 448; other patch-multiple resolutions work via positional-embedding interpolation
+- Preprocessing: images to `[0, 1]`, no normalization; SentencePiece tokenizer, lowercased, max 64 tokens
+
+## License
+
+Apache 2.0
+
+## Citation
+
+```bibtex
+@inproceedings{maninis2025tips,
+  title     = {{TIPS: Text-Image Pretraining with Spatial Awareness}},
+  author    = {Maninis, Kevis-Kokitsi and Chen, Kaifeng and Ghosh, Soham and Karpur, Arjun and Chen, Koert and Xia, Ye and Cao, Bingyi and Salz, Daniel and Han, Guangxing and Dlabal, Jan and Gnanapragasam, Dan and Seyedhosseini, Mojtaba and Zhou, Howard and Araujo, Andre},
+  booktitle = {International Conference on Learning Representations (ICLR)},
+  year      = {2025},
+  url       = {https://arxiv.org/abs/2410.16512}
+}
+```

config.json

@@ -0,0 +1,85 @@
+{
+  "architectures": [
+    "TIPSv2Model"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tips.TIPSv2Config",
+    "AutoModel": "modeling_tips.TIPSv2Model"
+  },
+  "model_type": "tipsv2",
+  "temperature_init_value": 0.004205586854368448,
+  "text_config": {
+    "attention_dropout": 0.0,
+    "bos_token_id": null,
+    "eos_token_id": null,
+    "hidden_act": "relu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "layer_norm_eps": 1e-05,
+    "max_position_embeddings": 64,
+    "model_type": "tipsv2_text_model",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 12,
+    "pad_token_id": 0,
+    "pooling_epsilon": 1e-08,
+    "scale_sqrt_depth": true,
+    "vocab_size": 32000
+  },
+  "transformers_version": "5.14.1",
+  "vision_config": {
+    "apply_layernorm": true,
+    "attention_probs_dropout_prob": 0.0,
+    "drop_path_rate": 0.0,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 1024,
+    "image_size": 448,
+    "initializer_range": 0.02,
+    "layer_norm_eps": 1e-06,
+    "layerscale_value": 1.0,
+    "mlp_ratio": 4,
+    "model_type": "tipsv2_vision_model",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 24,
+    "num_register_tokens": 1,
+    "out_features": [
+      "stage24"
+    ],
+    "out_indices": [
+      24
+    ],
+    "patch_size": 14,
+    "qkv_bias": true,
+    "reshape_hidden_states": true,
+    "stage_names": [
+      "stem",
+      "stage1",
+      "stage2",
+      "stage3",
+      "stage4",
+      "stage5",
+      "stage6",
+      "stage7",
+      "stage8",
+      "stage9",
+      "stage10",
+      "stage11",
+      "stage12",
+      "stage13",
+      "stage14",
+      "stage15",
+      "stage16",
+      "stage17",
+      "stage18",
+      "stage19",
+      "stage20",
+      "stage21",
+      "stage22",
+      "stage23",
+      "stage24"
+    ],
+    "use_swiglu_ffn": false
+  }
+}

configuration_tips.py

@@ -0,0 +1,45 @@
+"""TIPSv2 model configuration."""
+
+from transformers import PretrainedConfig
+
+
+_VISION_FN_BY_GEOMETRY = {
+    (384, 12): "vit_small",
+    (768, 12): "vit_base",
+    (1024, 24): "vit_large",
+    (1152, 27): "vit_so400m",
+    (1536, 40): "vit_giant2",
+}
+
+
+class TIPSv2Config(PretrainedConfig):
+    """Configuration for TIPSv2 vision-language model."""
+
+    model_type = "tipsv2"
+
+    def __init__(
+        self,
+        vision_config=None,
+        text_config=None,
+        temperature_init_value=0.01,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        vision_config = vision_config or {}
+        text_config = text_config or {}
+        hidden_size = vision_config.get("hidden_size", 768)
+        num_hidden_layers = vision_config.get("num_hidden_layers", 12)
+        self.vision_fn = _VISION_FN_BY_GEOMETRY[(hidden_size, num_hidden_layers)]
+        self.embed_dim = hidden_size
+        self.patch_size = vision_config.get("patch_size", 14)
+        self.img_size = vision_config.get("image_size", 448)
+        self.ffn_layer = "swiglu" if vision_config.get("use_swiglu_ffn", False) else "mlp"
+        self.init_values = vision_config.get("layerscale_value", 1.0)
+        self.num_register_tokens = vision_config.get("num_register_tokens", 1)
+        self.text_hidden_size = text_config.get("hidden_size", 768)
+        self.text_mlp_dim = text_config.get("intermediate_size", 3072)
+        self.text_num_heads = text_config.get("num_attention_heads", 12)
+        self.text_num_layers = text_config.get("num_hidden_layers", 12)
+        self.vocab_size = text_config.get("vocab_size", 32000)
+        self.max_len = text_config.get("max_position_embeddings", 64)
+        self.temperature = temperature_init_value

image_encoder.py

@@ -0,0 +1,1002 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Vision encoder implementation in PyTorch."""
+
+import functools
+import math
+import os
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+import warnings
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+
+class Mlp(nn.Module):
+  """Transformer MLP, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = nn.GELU,
+      drop: float = 0.0,
+      bias: bool = True,
+  ) -> None:
+    super().__init__()
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+    self.act = act_layer()
+    self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+    self.drop = nn.Dropout(drop)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    x = self.fc1(x)
+    x = self.act(x)
+    x = self.drop(x)
+    x = self.fc2(x)
+    x = self.drop(x)
+    return x
+
+
+def make_2tuple(x):
+  if isinstance(x, tuple):
+    assert len(x) == 2
+    return x
+
+  assert isinstance(x, int)
+  return (x, x)
+
+
+class PatchEmbed(nn.Module):
+  """2D image to patch embedding: (B,C,H,W) -> (B,N,D)."""
+
+  def __init__(
+      self,
+      img_size: Union[int, Tuple[int, int]] = 224,
+      patch_size: Union[int, Tuple[int, int]] = 16,
+      in_chans: int = 3,
+      embed_dim: int = 768,
+      norm_layer: Optional[Callable] = None,  # pylint: disable=g-bare-generic
+      flatten_embedding: bool = True,
+  ) -> None:
+    super().__init__()
+
+    image_hw = make_2tuple(img_size)
+    patch_hw = make_2tuple(patch_size)
+    patch_grid_size = (
+        image_hw[0] // patch_hw[0],
+        image_hw[1] // patch_hw[1],
+    )
+
+    self.img_size = image_hw
+    self.patch_size = patch_hw
+    self.patches_resolution = patch_grid_size
+    self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+    self.in_chans = in_chans
+    self.embed_dim = embed_dim
+
+    self.flatten_embedding = flatten_embedding
+
+    self.proj = nn.Conv2d(
+        in_chans, embed_dim, kernel_size=patch_hw, stride=patch_hw
+    )
+    self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    _, _, h, w = x.shape
+    patch_h, patch_w = self.patch_size
+
+    assert (
+        h % patch_h == 0
+    ), f"Input image height {h} is not a multiple of patch height {patch_h}"
+    assert (
+        w % patch_w == 0
+    ), f"Input image width {w} is not a multiple of patch width: {patch_w}"
+
+    x = self.proj(x)  # B C H W
+    h, w = x.size(2), x.size(3)
+    x = x.flatten(2).transpose(1, 2)  # B HW C
+    x = self.norm(x)
+    if not self.flatten_embedding:
+      x = x.reshape(-1, h, w, self.embed_dim)  # B H W C
+    return x
+
+  def flops(self) -> float:
+    ho, wo = self.patches_resolution
+    flops = (
+        ho
+        * wo
+        * self.embed_dim
+        * self.in_chans
+        * (self.patch_size[0] * self.patch_size[1])
+    )
+    if self.norm is not None:
+      flops += ho * wo * self.embed_dim
+    return flops
+
+
+class SwiGLUFFN(nn.Module):
+  """SwiGLU FFN layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = None,
+      drop: float = 0.0,
+      bias: bool = True,
+  ) -> None:
+    super().__init__()
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+    self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    x12 = self.w12(x)
+    x1, x2 = x12.chunk(2, dim=-1)
+    hidden = F.silu(x1) * x2
+    return self.w3(hidden)
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+  if XFORMERS_ENABLED:
+    from xformers.ops import SwiGLU, memory_efficient_attention, unbind, fmha, scaled_index_add, index_select_cat  # pylint: disable=g-multiple-import, g-import-not-at-top
+
+    XFORMERS_AVAILABLE = True
+    warnings.warn("xFormers is available (SwiGLU)")
+  else:
+    warnings.warn("xFormers is disabled (SwiGLU)")
+    raise ImportError
+except ImportError:
+  SwiGLU = SwiGLUFFN
+  XFORMERS_AVAILABLE = False
+
+  warnings.warn("xFormers is not available (SwiGLU)")
+
+
+class SwiGLUFFNFused(SwiGLU):
+  """SwiGLU FFN layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      in_features: int,
+      hidden_features: Optional[int] = None,
+      out_features: Optional[int] = None,
+      act_layer: Callable[..., nn.Module] = None,  # pylint: disable=unused-argument
+      drop: float = 0.0,  # pylint: disable=unused-argument
+      bias: bool = True,
+  ) -> None:
+    out_features = out_features or in_features
+    hidden_features = hidden_features or in_features
+    hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+    super().__init__(
+        in_features=in_features,
+        hidden_features=hidden_features,
+        out_features=out_features,
+        bias=bias,
+    )
+
+
+class Attention(nn.Module):
+  """Attention layer, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      num_heads: int = 8,
+      qkv_bias: bool = False,
+      proj_bias: bool = True,
+      attn_drop: float = 0.0,
+      proj_drop: float = 0.0,
+  ) -> None:
+    super().__init__()
+    self.num_heads = num_heads
+    head_dim = dim // num_heads
+    self.scale = head_dim**-0.5
+
+    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+    self.attn_drop = nn.Dropout(attn_drop)
+    self.proj = nn.Linear(dim, dim, bias=proj_bias)
+    self.proj_drop = nn.Dropout(proj_drop)
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    b_dim, n_dim, c_dim = x.shape
+    qkv = (
+        self.qkv(x)
+        .reshape(b_dim, n_dim, 3, self.num_heads, c_dim // self.num_heads)
+        .permute(2, 0, 3, 1, 4)
+    )
+
+    q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+    attn = q @ k.transpose(-2, -1)
+
+    attn = attn.softmax(dim=-1)
+    attn = self.attn_drop(attn)
+
+    x = (attn @ v).transpose(1, 2).reshape(b_dim, n_dim, c_dim)
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+class MemEffAttention(Attention):
+  """Memory Efficient Attention layer, following DINOv2 implementation."""
+
+  def forward(self, x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+    if not XFORMERS_AVAILABLE:
+      if attn_bias is not None:
+        raise AssertionError("xFormers is required for using nested tensors")
+      return super().forward(x)
+
+    b_dim, n_dim, c_dim = x.shape
+    qkv = self.qkv(x).reshape(
+        b_dim, n_dim, 3, self.num_heads, c_dim // self.num_heads
+    )
+
+    q, k, v = unbind(qkv, 2)
+
+    x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+    x = x.reshape([b_dim, n_dim, c_dim])
+
+    x = self.proj(x)
+    x = self.proj_drop(x)
+    return x
+
+
+class LayerScale(nn.Module):
+  """Layer scale, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      init_values: Union[float, torch.Tensor] = 1e-5,
+      inplace: bool = False,
+  ) -> None:
+    super().__init__()
+    self.inplace = inplace
+    self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+def drop_path_impl(x, drop_prob: float = 0.0, training: bool = False):
+  if drop_prob == 0.0 or not training:
+    return x
+  keep_prob = 1 - drop_prob
+  shape = (x.shape[0],) + (1,) * (
+      x.ndim - 1
+  )  # work with diff dim tensors, not just 2D ConvNets
+  random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+  if keep_prob > 0.0:
+    random_tensor.div_(keep_prob)
+  output = x * random_tensor
+  return output
+
+
+class DropPath(nn.Module):
+  """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+  def __init__(self, drop_prob=None):
+    super(DropPath, self).__init__()
+    self.drop_prob = drop_prob
+
+  def forward(self, x):
+    return drop_path_impl(x, self.drop_prob, self.training)
+
+
+class Block(nn.Module):
+  """Transformer Block Implementation, following DINOv2 implementation."""
+
+  def __init__(
+      self,
+      dim: int,
+      num_heads: int,
+      mlp_ratio: float = 4.0,
+      qkv_bias: bool = False,
+      proj_bias: bool = True,
+      ffn_bias: bool = True,
+      drop: float = 0.0,
+      attn_drop: float = 0.0,
+      init_values=None,
+      drop_path: float = 0.0,
+      act_layer: Callable[..., nn.Module] = nn.GELU,
+      norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+      attn_class: Callable[..., nn.Module] = Attention,
+      ffn_layer: Callable[..., nn.Module] = Mlp,
+  ) -> None:
+    super().__init__()
+    self.norm1 = norm_layer(dim)
+    self.attn = attn_class(
+        dim,
+        num_heads=num_heads,
+        qkv_bias=qkv_bias,
+        proj_bias=proj_bias,
+        attn_drop=attn_drop,
+        proj_drop=drop,
+    )
+    self.ls1 = (
+        LayerScale(dim, init_values=init_values)
+        if init_values
+        else nn.Identity()
+    )
+    self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    self.norm2 = norm_layer(dim)
+    mlp_hidden_dim = int(dim * mlp_ratio)
+    self.mlp = ffn_layer(
+        in_features=dim,
+        hidden_features=mlp_hidden_dim,
+        act_layer=act_layer,
+        drop=drop,
+        bias=ffn_bias,
+    )
+    self.ls2 = (
+        LayerScale(dim, init_values=init_values)
+        if init_values
+        else nn.Identity()
+    )
+    self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+    self.sample_drop_ratio = drop_path
+
+  def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def attn_residual_func(x: torch.Tensor) -> torch.Tensor:
+      return self.ls1(self.attn(self.norm1(x)))
+
+    def ffn_residual_func(x: torch.Tensor) -> torch.Tensor:
+      return self.ls2(self.mlp(self.norm2(x)))
+
+    if self.training and self.sample_drop_ratio > 0.1:
+      # the overhead is compensated only for a drop path rate larger than 0.1
+      x = drop_add_residual_stochastic_depth(
+          x,
+          residual_func=attn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+      )
+      x = drop_add_residual_stochastic_depth(
+          x,
+          residual_func=ffn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+      )
+    elif self.training and self.sample_drop_ratio > 0.0:
+      x = x + self.drop_path1(attn_residual_func(x))
+      x = x + self.drop_path1(ffn_residual_func(x))
+    else:
+      x = x + attn_residual_func(x)
+      x = x + ffn_residual_func(x)
+    return x
+
+
+def drop_add_residual_stochastic_depth(
+    x: torch.Tensor,
+    residual_func: Callable[[torch.Tensor], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+) -> torch.Tensor:
+  """This function is taken from the original implementation in DINOv2 to implement stochastic depth in the image encoder."""
+  # 1) extract subset using permutation
+  b, _, _ = x.shape
+  sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+  brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+  x_subset = x[brange]
+
+  # 2) apply residual_func to get residual
+  residual = residual_func(x_subset)
+
+  x_flat = x.flatten(1)
+  residual = residual.flatten(1)
+
+  residual_scale_factor = b / sample_subset_size
+
+  # 3) add the residual
+  x_plus_residual = torch.index_add(
+      x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
+  )
+  return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+  b, _, _ = x.shape
+  sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+  brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+  residual_scale_factor = b / sample_subset_size
+  return brange, residual_scale_factor
+
+
+def add_residual(
+    x, brange, residual, residual_scale_factor, scaling_vector=None
+):
+  """Implement residual addition in the image encoder."""
+  if scaling_vector is None:
+    x_flat = x.flatten(1)
+    residual = residual.flatten(1)
+    x_plus_residual = torch.index_add(
+        x_flat,
+        0,
+        brange,
+        residual.to(dtype=x.dtype),
+        alpha=residual_scale_factor,
+    )
+  else:
+    x_plus_residual = scaled_index_add(
+        x,
+        brange,
+        residual.to(dtype=x.dtype),
+        scaling=scaling_vector,
+        alpha=residual_scale_factor,
+    )
+  return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}  # pylint: disable=g-bare-generic
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+  """this will perform the index select, cat the tensors, and provide the attn_bias from cache."""
+  batch_sizes = (
+      [b.shape[0] for b in branges]
+      if branges is not None
+      else [x.shape[0] for x in x_list]
+  )
+  all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+  if all_shapes not in attn_bias_cache.keys():
+    seqlens = []
+    for b, x in zip(batch_sizes, x_list):
+      for _ in range(b):
+        seqlens.append(x.shape[1])
+    attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+    attn_bias._batch_sizes = batch_sizes  # pylint: disable=protected-access
+    attn_bias_cache[all_shapes] = attn_bias
+
+  if branges is not None:
+    cat_tensors = index_select_cat(
+        [x.flatten(1) for x in x_list], branges
+    ).view(1, -1, x_list[0].shape[-1])
+  else:
+    tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+    cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+  return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+    x_list: List[torch.Tensor],
+    residual_func: Callable[[torch.Tensor, Any], torch.Tensor],
+    sample_drop_ratio: float = 0.0,
+    scaling_vector=None,
+) -> torch.Tensor:
+  """Add residual to a list of tensors."""
+  # 1) generate random set of indices for dropping samples in the batch.
+  branges_scales = [
+      get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list
+  ]
+  branges = [s[0] for s in branges_scales]
+  residual_scale_factors = [s[1] for s in branges_scales]
+
+  # 2) get attention bias and index+concat the tensors.
+  attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+  # 3) apply residual_func to get residual, and split the result.
+  residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias))  # type: ignore
+
+  outputs = []
+  for x, brange, residual, residual_scale_factor in zip(
+      x_list, branges, residual_list, residual_scale_factors
+  ):
+    outputs.append(
+        add_residual(
+            x, brange, residual, residual_scale_factor, scaling_vector
+        ).view_as(x)
+    )
+  return outputs
+
+
+class NestedTensorBlock(Block):
+  """Nested tensor block implementation."""
+
+  def forward_nested(self, x_list: List[torch.Tensor]) -> List[torch.Tensor]:
+    """x_list contains a list of tensors to nest together and run."""
+    assert isinstance(self.attn, MemEffAttention)
+
+    if self.training and self.sample_drop_ratio > 0.0:
+
+      def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+      def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        del attn_bias
+        return self.mlp(self.norm2(x))
+
+      x_list = drop_add_residual_stochastic_depth_list(
+          x_list,
+          residual_func=attn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+          scaling_vector=self.ls1.gamma
+          if isinstance(self.ls1, LayerScale)
+          else None,
+      )
+      x_list = drop_add_residual_stochastic_depth_list(
+          x_list,
+          residual_func=ffn_residual_func,
+          sample_drop_ratio=self.sample_drop_ratio,
+          scaling_vector=self.ls2.gamma
+          if isinstance(self.ls1, LayerScale)
+          else None,
+      )
+      return x_list
+    else:
+
+      def attn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+      def ffn_residual_func(x: torch.Tensor, attn_bias=None) -> torch.Tensor:
+        del attn_bias
+        return self.ls2(self.mlp(self.norm2(x)))
+
+      attn_bias, x = get_attn_bias_and_cat(x_list)
+      x = x + attn_residual_func(x, attn_bias=attn_bias)
+      x = x + ffn_residual_func(x)
+      return attn_bias.split(x)
+
+  def forward(self, x):
+    if isinstance(x, torch.Tensor):
+      return super().forward(x)
+    elif isinstance(x, list):
+      if not XFORMERS_AVAILABLE:
+        raise AssertionError("xFormers is required for using nested tensors")
+      return self.forward_nested(x)
+    else:
+      raise AssertionError
+
+
+def named_apply(
+    fn: Callable,  # pylint: disable=g-bare-generic
+    module: nn.Module,
+    name="",
+    depth_first=True,
+    include_root=False,
+) -> nn.Module:
+  """Apply a function to a module and its children."""
+  if not depth_first and include_root:
+    fn(module=module, name=name)
+  for child_name, child_module in module.named_children():
+    child_name = ".".join((name, child_name)) if name else child_name
+    named_apply(
+        fn=fn,
+        module=child_module,
+        name=child_name,
+        depth_first=depth_first,
+        include_root=True,
+    )
+  if depth_first and include_root:
+    fn(module=module, name=name)
+  return module
+
+
+class BlockChunk(nn.ModuleList):
+
+  def forward(self, x):
+    for b in self:
+      x = b(x)
+    return x
+
+
+class VisionTransformer(nn.Module):
+  """Vision Transformer implementation."""
+
+  def __init__(
+      self,
+      img_size=224,
+      patch_size=16,
+      in_chans=3,
+      embed_dim=768,
+      depth=12,
+      num_heads=12,
+      mlp_ratio=4.0,
+      qkv_bias=True,
+      ffn_bias=True,
+      proj_bias=True,
+      drop_path_rate=0.0,
+      drop_path_uniform=False,
+      init_values=None,  # for layerscale: None or 0 => no layerscale
+      embed_layer=PatchEmbed,
+      act_layer=nn.GELU,
+      block_fn=Block,
+      ffn_layer="mlp",
+      block_chunks=1,
+      num_register_tokens=0,
+      interpolate_antialias=False,
+      interpolate_offset=0.1,
+  ):
+    """Defines the Vision Transformer model.
+
+    Args:
+      img_size (int, tuple): input image size
+      patch_size (int, tuple): patch size
+      in_chans (int): number of input channels
+      embed_dim (int): embedding dimension
+      depth (int): depth of transformer
+      num_heads (int): number of attention heads
+      mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+      qkv_bias (bool): enable bias for qkv if True
+      ffn_bias (bool): enable bias for ffn if True
+      proj_bias (bool): enable bias for proj in attn if True
+      drop_path_rate (float): stochastic depth rate
+      drop_path_uniform (bool): apply uniform drop rate across blocks
+      init_values (float): layer-scale init values
+      embed_layer (nn.Module): patch embedding layer
+      act_layer (nn.Module): MLP activation layer
+      block_fn (nn.Module): transformer block class
+      ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+      block_chunks: (int) split block sequence into block_chunks units for FSDP
+        wrap
+      num_register_tokens: (int) number of extra cls tokens (so-called
+        "registers")
+      interpolate_antialias: (str) flag to apply anti-aliasing when
+        interpolating positional embeddings
+      interpolate_offset: (float) work-around offset to apply when interpolating
+        positional embeddings
+    """
+    super().__init__()
+    norm_layer = functools.partial(nn.LayerNorm, eps=1e-6)
+
+    self.num_features = self.embed_dim = (
+        embed_dim  # num_features for consistency with other models
+    )
+    self.num_tokens = 1
+    self.n_blocks = depth
+    self.num_heads = num_heads
+    self.patch_size = patch_size
+    self.num_register_tokens = num_register_tokens
+    self.interpolate_antialias = interpolate_antialias
+    self.interpolate_offset = interpolate_offset
+
+    self.patch_embed = embed_layer(
+        img_size=img_size,
+        patch_size=patch_size,
+        in_chans=in_chans,
+        embed_dim=embed_dim,
+    )
+    num_patches = self.patch_embed.num_patches
+
+    self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+    self.pos_embed = nn.Parameter(
+        torch.zeros(1, num_patches + self.num_tokens, embed_dim)
+    )
+    assert num_register_tokens >= 0
+    self.register_tokens = (
+        nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim))
+        if num_register_tokens
+        else None
+    )
+
+    if drop_path_uniform:
+      dpr = [drop_path_rate] * depth
+    else:
+      dpr = [
+          drop_path_rate * i / max(depth - 1, 1) for i in range(depth)
+      ]  # stochastic depth decay rule
+
+    if ffn_layer == "mlp":
+      ffn_layer = Mlp
+    elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+      ffn_layer = SwiGLUFFNFused
+    else:
+      raise NotImplementedError
+
+    blocks_list = [
+        block_fn(
+            dim=embed_dim,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            ffn_bias=ffn_bias,
+            drop_path=dpr[i],
+            norm_layer=norm_layer,
+            act_layer=act_layer,
+            ffn_layer=ffn_layer,
+            init_values=init_values,
+        )
+        for i in range(depth)
+    ]
+    if block_chunks > 0:
+      self.chunked_blocks = True
+      chunked_blocks = []
+      chunksize = depth // block_chunks
+      for i in range(0, depth, chunksize):
+        # this is to keep the block index consistent if we chunk the block list
+        chunked_blocks.append(
+            [nn.Identity()] * i + blocks_list[i : i + chunksize]
+        )
+      self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+    else:
+      self.chunked_blocks = False
+      self.blocks = nn.ModuleList(blocks_list)
+
+    self.norm = norm_layer(embed_dim)
+    self.head = nn.Identity()
+
+    self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+    self.init_weights()
+
+  def init_weights(self):
+    nn.init.trunc_normal_(self.pos_embed, std=0.02)
+    nn.init.normal_(self.cls_token, std=1e-6)
+    if self.register_tokens is not None:
+      nn.init.normal_(self.register_tokens, std=1e-6)
+    named_apply(init_weights_vit_timm, self)
+
+  def interpolate_pos_encoding(self, x, w, h):
+    previous_dtype = x.dtype
+    npatch = x.shape[1] - 1
+    num_patches = self.pos_embed.shape[1] - 1
+    if npatch == num_patches and w == h:
+      return self.pos_embed
+    pos_embed = self.pos_embed.float()
+    class_pos_embed = pos_embed[:, 0]
+    patch_pos_embed = pos_embed[:, 1:]
+    dim = x.shape[-1]
+    w0 = w // self.patch_size
+    h0 = h // self.patch_size
+    num_patches_dim = int(
+        math.sqrt(num_patches)
+    )  # Recover the number of patches in each dimension
+    assert num_patches == num_patches_dim * num_patches_dim
+    kwargs = {}
+    if self.interpolate_offset:
+      sx = float(w0 + self.interpolate_offset) / num_patches_dim
+      sy = float(h0 + self.interpolate_offset) / num_patches_dim
+      kwargs["scale_factor"] = (sx, sy)
+    else:
+      # Simply specify an output size instead of a scale factor
+      kwargs["size"] = (w0, h0)
+    patch_pos_embed = nn.functional.interpolate(
+        patch_pos_embed.reshape(
+            1, num_patches_dim, num_patches_dim, dim
+        ).permute(0, 3, 1, 2),
+        mode="bilinear",
+        antialias=self.interpolate_antialias,
+        **kwargs,
+    )
+    assert (w0, h0) == patch_pos_embed.shape[-2:]
+    patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+    return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(
+        previous_dtype
+    )
+
+  def prepare_tokens_with_masks(self, x, masks=None):
+    _, _, w, h = x.shape
+    x = self.patch_embed(x)
+    if masks is not None:
+      x = torch.where(
+          masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x
+      )
+
+    x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+    x = x + self.interpolate_pos_encoding(x, w, h)
+
+    if self.register_tokens is not None:
+      x = torch.cat(
+          (
+              x[:, :1],
+              self.register_tokens.expand(x.shape[0], -1, -1),
+              x[:, 1:],
+          ),
+          dim=1,
+      )
+
+    return x
+
+  def forward_features_list(self, x_list, masks_list):
+    x = [
+        self.prepare_tokens_with_masks(x, masks)
+        for x, masks in zip(x_list, masks_list)
+    ]
+    for blk in self.blocks:
+      x = blk(x)
+
+    all_x = x
+    output = []
+    for x, masks in zip(all_x, masks_list):
+      x_norm = self.norm(x)
+      output.append({
+          "x_norm_1st_clstoken": x_norm[:, :1],
+          "x_norm_2nd_clstoken": x_norm[:, 1 : self.num_register_tokens + 1],
+          "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+          "x_prenorm": x,
+          "masks": masks,
+      })
+    return output
+
+  def forward_features(self, x, masks=None):
+    if isinstance(x, list):
+      return self.forward_features_list(x, masks)
+
+    x = self.prepare_tokens_with_masks(x, masks)
+
+    for blk in self.blocks:
+      x = blk(x)
+
+    x_norm = self.norm(x)
+    return {
+        "x_norm_1st_clstoken": x_norm[:, :1],
+        "x_norm_2nd_clstoken": x_norm[:, 1 : self.num_register_tokens + 1],
+        "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+        "x_prenorm": x,
+        "masks": masks,
+    }
+
+  def _get_intermediate_layers_not_chunked(self, x, n=1):
+    x = self.prepare_tokens_with_masks(x)
+    # If n is an int, take the n last blocks. If it's a list, take them
+    output, total_block_len = [], len(self.blocks)
+    blocks_to_take = (
+        range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+    )
+    for i, blk in enumerate(self.blocks):
+      x = blk(x)
+      if i in blocks_to_take:
+        output.append(x)
+    assert len(output) == len(
+        blocks_to_take
+    ), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+    return output
+
+  def _get_intermediate_layers_chunked(self, x, n=1):
+    x = self.prepare_tokens_with_masks(x)
+    output, i, total_block_len = [], 0, len(self.blocks[-1])
+    # If n is an int, take the n last blocks. If it's a list, take them
+    blocks_to_take = (
+        range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+    )
+    for block_chunk in self.blocks:
+      for blk in block_chunk[i:]:  # Passing the nn.Identity()
+        x = blk(x)
+        if i in blocks_to_take:
+          output.append(x)
+        i += 1
+    assert len(output) == len(
+        blocks_to_take
+    ), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+    return output
+
+  def get_intermediate_layers(
+      self,
+      x: torch.torch.Tensor,
+      n: Union[int, Sequence] = 1,  # Layers or n last layers to take  # pylint: disable=g-bare-generic
+      reshape: bool = False,
+      return_class_token: bool = False,
+      norm=True,
+  ) -> Tuple[Union[torch.torch.Tensor, Tuple[torch.torch.Tensor]]]:  # pylint: disable=g-one-element-tuple
+    if self.chunked_blocks:
+      outputs = self._get_intermediate_layers_chunked(x, n)
+    else:
+      outputs = self._get_intermediate_layers_not_chunked(x, n)
+    if norm:
+      outputs = [self.norm(out) for out in outputs]
+    class_tokens = [out[:, 0] for out in outputs]
+    outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+    if reshape:
+      batch_size, _, w, h = x.shape
+      outputs = [
+          out.reshape(
+              batch_size, w // self.patch_size, h // self.patch_size, -1
+          )
+          .permute(0, 3, 1, 2)
+          .contiguous()
+          for out in outputs
+      ]
+    if return_class_token:
+      return tuple(zip(outputs, class_tokens))
+    return tuple(outputs)
+
+  def forward(self, *args, is_training=False, **kwargs):
+    ret = self.forward_features(*args, **kwargs)
+    if is_training:
+      return ret
+    else:
+      return self.head(ret["x_norm_1st_clstoken"]), self.head(
+          ret["x_norm_2nd_clstoken"]
+      ), ret["x_norm_patchtokens"]
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):  # pylint: disable=unused-argument
+  """ViT weight initialization, original timm impl (for reproducibility)."""
+  if isinstance(module, nn.Linear):
+    nn.init.trunc_normal_(module.weight, std=0.02)
+    if module.bias is not None:
+      nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=384,
+      depth=12,
+      num_heads=6,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_base(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=768,
+      depth=12,
+      num_heads=12,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_large(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1024,
+      depth=24,
+      num_heads=16,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_so400m(patch_size=14, **kwargs):
+  """SoViT 400M model (https://arxiv.org/abs/2305.13035)."""
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1152,
+      depth=27,
+      num_heads=16,
+      mlp_ratio=4304 / 1152,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model
+
+
+def vit_giant2(patch_size=14, **kwargs):
+  model = VisionTransformer(
+      patch_size=patch_size,
+      embed_dim=1536,
+      depth=40,
+      num_heads=24,
+      mlp_ratio=4,
+      block_fn=functools.partial(Block, attn_class=MemEffAttention),
+      num_register_tokens=1,
+      **kwargs,
+  )
+  return model

model.safetensors

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

modeling_tips.py

@@ -0,0 +1,141 @@
+"""TIPSv2 model for HuggingFace — wraps vision and text encoders."""
+
+from dataclasses import dataclass
+from typing import List, Optional, Union
+
+import torch
+from transformers import PreTrainedModel
+from transformers.utils import cached_file
+
+from .configuration_tips import TIPSv2Config
+from .image_encoder import vit_base, vit_giant2, vit_large, vit_small, vit_so400m
+from .text_encoder import TextEncoder, Tokenizer
+
+_VISION_FACTORIES = {
+    "vit_small": vit_small,
+    "vit_base": vit_base,
+    "vit_large": vit_large,
+    "vit_so400m": vit_so400m,
+    "vit_giant2": vit_giant2,
+}
+
+
+@dataclass
+class TIPSv2ImageOutput:
+    """Output from the vision encoder."""
+    cls_token: torch.Tensor        # (B, 1, D)
+    register_tokens: torch.Tensor  # (B, R, D)
+    patch_tokens: torch.Tensor     # (B, N, D)
+
+
+@dataclass
+class TIPSv2Output:
+    """Output from the full model."""
+    image_features: Optional[TIPSv2ImageOutput] = None
+    text_embeds: Optional[torch.Tensor] = None
+    temperature: Optional[float] = None
+
+
+class TIPSv2Model(PreTrainedModel):
+    """TIPSv2 vision-language model.
+
+    Usage::
+
+        model = AutoModel.from_pretrained("google/tipsv2-b14", trust_remote_code=True)
+
+        # Image features
+        out = model.encode_image(pixel_values)  # pixel_values in [0, 1]
+        cls = out.cls_token        # (B, 1, D)
+        spatial = out.patch_tokens  # (B, N, D)
+
+        # Text features
+        text_emb = model.encode_text(["a photo of a cat"])  # (B, D)
+    """
+
+    config_class = TIPSv2Config
+    _no_split_modules = []
+    _supports_cache_class = False
+    _tied_weights_keys = []
+
+    @property
+    def all_tied_weights_keys(self):
+        return {}
+
+    def __init__(self, config: TIPSv2Config):
+        super().__init__(config)
+
+        self.vision_encoder = _VISION_FACTORIES[config.vision_fn](
+            img_size=config.img_size,
+            patch_size=config.patch_size,
+            ffn_layer=config.ffn_layer,
+            block_chunks=0,
+            init_values=config.init_values,
+            interpolate_antialias=True,
+            interpolate_offset=0.0,
+        )
+
+        self.text_encoder = TextEncoder(
+            config={
+                "hidden_size": config.text_hidden_size,
+                "mlp_dim": config.text_mlp_dim,
+                "num_heads": config.text_num_heads,
+                "num_layers": config.text_num_layers,
+            },
+            vocab_size=config.vocab_size,
+        )
+
+        self._tokenizer = None
+
+    def _load_tokenizer(self):
+        """Load the SentencePiece tokenizer shipped with the checkpoint."""
+        return Tokenizer(cached_file(self.name_or_path, "tokenizer.model"))
+
+    @torch.no_grad()
+    def encode_image(self, pixel_values: torch.Tensor) -> TIPSv2ImageOutput:
+        """Encode images. pixel_values: (B, 3, H, W) in [0, 1]."""
+        pixel_values = pixel_values.to(self.device)
+        cls_token, register_tokens, patch_tokens = self.vision_encoder(pixel_values)
+        return TIPSv2ImageOutput(
+            cls_token=cls_token,
+            register_tokens=register_tokens,
+            patch_tokens=patch_tokens,
+        )
+
+    @torch.no_grad()
+    def encode_text(
+        self,
+        texts: Union[str, List[str], torch.Tensor],
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode text. Pass strings (auto-tokenized) or pre-tokenized tensors."""
+        if isinstance(texts, (str, list)):
+            if isinstance(texts, str):
+                texts = [texts]
+            if self._tokenizer is None:
+                self._tokenizer = self._load_tokenizer()
+            ids, paddings = self._tokenizer.tokenize(texts, max_len=self.config.max_len)
+            ids = torch.from_numpy(ids).to(self.device)
+            padding_mask = torch.from_numpy(paddings).to(self.device)
+        else:
+            ids = texts.to(self.device)
+            padding_mask = padding_mask.to(self.device)
+        return self.text_encoder(ids, padding_mask)
+
+    def forward(
+        self,
+        pixel_values: Optional[torch.Tensor] = None,
+        input_ids: Optional[torch.Tensor] = None,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> TIPSv2Output:
+        """Forward pass for both or either modality."""
+        image_features = None
+        text_embeds = None
+        if pixel_values is not None:
+            image_features = self.encode_image(pixel_values)
+        if input_ids is not None:
+            text_embeds = self.encode_text(input_ids, padding_mask)
+        return TIPSv2Output(
+            image_features=image_features,
+            text_embeds=text_embeds,
+            temperature=self.config.temperature,
+        )

processor_config.json

@@ -0,0 +1,16 @@
+{
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_processor_type": "Tipsv2ImageProcessor",
+    "resample": 2,
+    "rescale_factor": 0.00392156862745098,
+    "size": {
+      "height": 448,
+      "width": 448
+    }
+  },
+  "processor_class": "Tipsv2Processor"
+}

text_encoder.py

@@ -0,0 +1,342 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Text encoder implementation in PyTorch."""
+
+import typing as t
+
+import numpy as np
+import sentencepiece as spm
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class Tokenizer(object):
+  """A simple tokenizer using SentencePiece."""
+
+  def __init__(self, tokenizer_path: str):
+    self.sp = spm.SentencePieceProcessor(model_file=tokenizer_path)
+    # Explicitly disable BOS/EOS to match the reference Colab implementation.
+    self._add_bos = False
+    self._add_eos = False
+
+  def tokenize(self, input_texts, max_len=64):
+    if isinstance(input_texts, str):
+      input_texts = [input_texts]
+    batch_ids = [
+        self.sp.encode(t.lower(), add_bos=self._add_bos, add_eos=self._add_eos)
+        for t in input_texts
+    ]
+    tokens = np.zeros((len(batch_ids), max_len), dtype=np.int64)
+    for i, ids in enumerate(batch_ids):
+      length = min(len(ids), max_len)
+      tokens[i, :length] = ids[:length]
+    is_padding = (tokens == 0).astype(np.int32)
+    return tokens, is_padding
+
+
+class PositionalEmbedding(nn.Module):
+  """Generates position embedding for a given 1-d sequence.
+
+  Attributes:
+    min_timescale: Start of the geometric index. Determines the periodicity of
+      the added signal.
+    max_timescale: End of the geometric index. Determines the frequency of the
+      added signal.
+    embedding_dim: Dimension of the embedding to be generated.
+  """
+
+  min_timescale: int = 1
+  max_timescale: int = 10_000
+  embedding_dim: int = 0
+
+  def __init__(self, embedding_dim: int):
+    super().__init__()
+    self.embedding_dim = embedding_dim
+
+  def __call__(self, seq_length: int = None, position: torch.tensor = None):
+    """Generates a torch.tensor of sinusoids with different frequencies.
+
+    Args:
+      seq_length: an optional Python int defining the output sequence length.
+        if the `position` argument is specified.
+      position:   [B, seq_length], optional position for each token in the
+        sequence, only required when the sequence is packed.
+
+    Returns:
+      [B, seqlen, D] if `position` is specified, else [1, seqlen, D]
+    """
+    if position is None:
+      assert seq_length is not None
+      # [1, seqlen]
+      position = torch.arange(seq_length, dtype=torch.float32)[None, :]
+    else:
+      assert position.ndim == 2, position.shape
+
+    num_timescales = self.embedding_dim // 2
+    log_timescale_increment = torch.log(
+        torch.tensor(float(self.max_timescale) / float(self.min_timescale))
+    ) / torch.maximum(
+        torch.tensor(num_timescales, dtype=torch.float32) - 1, torch.tensor(1)
+    )
+    inv_timescales = self.min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float32)
+        * -log_timescale_increment
+    )
+    scaled_time = position[:, :, None] * inv_timescales[None, None, :]
+    signal = torch.cat((torch.sin(scaled_time), torch.cos(scaled_time)), dim=2)
+    # Force usage of `np` rather than `jnp` to compute static values at trace
+    # time.
+    signal = F.pad(signal, (0, self.embedding_dim % 2, 0, 0, 0, 0))
+    return signal
+
+
+class MlpBlockWithMask(nn.Module):
+  """Transformer MLP / feed-forward block that supports masking."""
+
+  def __init__(
+      self,
+      mlp_dim: int,
+      d_model: int,
+      use_bias: bool = True,
+      dtype: torch.dtype = torch.float32,
+      activation_fn: nn.Module = nn.GELU,
+  ):
+    super().__init__()
+
+    self.mlp_dim = mlp_dim
+    self.d_model = d_model
+    self.use_bias = use_bias
+    self.dtype = dtype
+    self.activation_fn = activation_fn
+
+    self.c_fc = nn.Linear(
+        in_features=self.d_model,
+        out_features=self.mlp_dim,
+        dtype=self.dtype,
+        bias=self.use_bias,
+    )
+    self.c_proj = nn.Linear(
+        in_features=self.mlp_dim,
+        out_features=self.d_model,
+        dtype=self.dtype,
+        bias=self.use_bias,
+    )
+
+  def __call__(
+      self, inputs: torch.Tensor, mlp_mask: torch.Tensor
+  ) -> torch.Tensor:
+    """Applies Transformer MlpBlock with mask module."""
+    x = self.c_fc(inputs)
+    x = self.activation_fn()(x)
+    x = x * mlp_mask[..., None]  # First masking.
+    x = self.c_proj(x)
+    x = x * mlp_mask[..., None]  # Second masking.
+    return x
+
+
+class ResidualAttentionBlock(nn.Module):
+  """Transformer residual attention block."""
+
+  def __init__(
+      self,
+      d_model: int,
+      n_head: int,
+      mlp_dim: int,
+      dtype: torch.dtype = torch.float32,
+  ):
+    super().__init__()
+    self.d_model = d_model
+    self.n_head = n_head
+    self.mlp_dim = mlp_dim
+    self.dtype = dtype
+
+    self.attn = nn.MultiheadAttention(d_model, n_head, dtype=self.dtype)
+    self.ln_1 = nn.LayerNorm(d_model, dtype=self.dtype)
+    self.mlp = MlpBlockWithMask(
+        self.mlp_dim,
+        d_model,
+        use_bias=True,
+        dtype=self.dtype,
+        activation_fn=nn.ReLU,
+    )
+    self.ln_2 = nn.LayerNorm(d_model, dtype=self.dtype)
+
+  def attention(self, x: torch.Tensor, mask: torch.Tensor):
+    attn_mask = (
+        mask[:, None, None, :]
+        .repeat(1, self.n_head, x.shape[0], 1)
+        .flatten(0, 1)
+    )
+    attn_mask[attn_mask == 0] = float('-inf')
+    attn_mask[attn_mask == 1] = 0
+    return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
+
+  def forward(self, x: torch.Tensor, mask: torch.Tensor):
+    x = x + self.attention(self.ln_1(x), mask.permute(1, 0))
+    x = x + self.mlp(self.ln_2(x), mask)
+    return x, mask
+
+
+class SequentialMultiInput(nn.Sequential):
+  """Sequential module that can take multiple inputs."""
+
+  def forward(self, *inputs):
+    for module in self._modules.values():
+      if isinstance(inputs, tuple):
+        inputs = module(*inputs)
+      else:
+        inputs = module(inputs)
+    return inputs
+
+
+class Transformer(nn.Module):
+  """Transformer implementation."""
+
+  def __init__(
+      self,
+      width: int,
+      layers: int,
+      heads: int,
+      mlp_dim: int,
+      dtype: torch.dtype = torch.float32,
+  ):
+    super().__init__()
+    self.width = width
+    self.layers = layers
+    self.heads = heads
+    self.mlp_dim = mlp_dim
+    self.dtype = dtype
+
+    self.resblocks = SequentialMultiInput(*[
+        ResidualAttentionBlock(self.width, self.heads, self.mlp_dim, self.dtype)
+        for _ in range(self.layers)
+    ])
+
+  def forward(self, x: torch.Tensor, mask: torch.Tensor):
+    return self.resblocks(x, mask)[0]
+
+
+class GlobalAvgPooling(nn.Module):
+  """Performs a simple global pooling over the input with optional paddings.
+
+  Attributes:
+    pooling_dims: A list of dims to perform pooling over.
+    keepdims: If True, keep dimension of inputs after pooling.
+  """
+
+  pooling_dims: t.Sequence[int]
+  epsilon: float = 1e-8
+
+  def __init__(
+      self, pooling_dims: t.Sequence[int], epsilon: float = 1e-8
+  ):
+    super().__init__()
+    self.pooling_dims = pooling_dims
+    self.epsilon = epsilon
+
+    if not all([p_dims >= 0 for p_dims in self.pooling_dims]):
+      raise ValueError('pooling_dims must be non-negative integers.')
+
+  def __call__(
+      self,
+      inputs: torch.tensor,
+      compatible_paddings: torch.tensor,
+  ):
+    """Applies global average spatial pooling to inputs.
+
+    Args:
+      inputs: An input tensor.
+      compatible_paddings: paddings of inputs with shapes compatible with
+        inputs, e.g. compatible_paddings with shape [B, 1] for inputs with shape
+        [B, D].
+
+    Returns:
+      Output tensor with global pooling applied.
+    """
+    padded_value = torch.zeros_like(inputs)
+    padded_value = torch.ones_like(inputs) * padded_value
+    inputs = torch.where(compatible_paddings > 0, padded_value, inputs)
+    valid_inputs = (
+        torch.sum(
+            1.0 - compatible_paddings,
+            self.pooling_dims,
+            keepdims=True,
+            dtype=inputs.dtype,
+        )
+        + self.epsilon
+    )
+    inputs_sum = torch.sum(inputs, self.pooling_dims, keepdims=True)
+    outputs = torch.divide(inputs_sum, valid_inputs).type(inputs.dtype)
+    outputs = torch.squeeze(outputs, axis=self.pooling_dims)
+    return outputs
+
+
+class TextEncoder(nn.Module):
+  """Text encoder implementation."""
+
+  def __init__(
+      self,
+      config: t.Dict[str, int],
+      vocab_size: int,
+      dtype: torch.dtype = torch.float32,
+      scale_sqrt_depth: bool = True,
+  ):
+    super().__init__()
+    self.vocab_size = vocab_size
+    self.dtype = dtype
+    self.scale_sqrt_depth = scale_sqrt_depth
+
+    # The text tower layers are fixed independent of vision tower size.
+    self.transformer_layers = config['num_layers']
+    self.embedding_dim = config['hidden_size']
+    self.transformer_width = config['hidden_size']
+    self.mlp_dim = config['mlp_dim']
+    self.transformer_heads = config['num_heads']
+
+    self.token_embedding = nn.Embedding(
+        self.vocab_size, self.embedding_dim, dtype=self.dtype
+    )
+    self.pos_embedder = PositionalEmbedding(embedding_dim=self.embedding_dim)
+    self.transformer = Transformer(
+        width=self.transformer_width,
+        layers=self.transformer_layers,
+        heads=self.transformer_heads,
+        mlp_dim=self.mlp_dim,
+        dtype=self.dtype,
+    )
+    self.pooling = GlobalAvgPooling(pooling_dims=[1])
+    self.ln_final = nn.LayerNorm(self.transformer_width, dtype=self.dtype)
+
+  def __call__(
+      self,
+      ids: torch.tensor,
+      paddings: torch.tensor,
+  ):
+    """Applies TextEncoder module."""
+    _, seq_length = ids.shape
+    mask = (paddings == 0).type(torch.float32)
+    mask = mask.permute(1, 0)  # NL -> LN
+    x = self.token_embedding(ids)
+    if self.scale_sqrt_depth:
+      x = x * (self.embedding_dim**0.5)
+    x = x + self.pos_embedder(seq_length=seq_length).to(x.device)
+    x = x.permute(1, 0, 2)  # NLD -> LND
+    x = self.transformer(x, mask)
+    x = x.permute(1, 0, 2)  # LND -> NLD
+    x = self.ln_final(x)
+    x = self.pooling(x, compatible_paddings=paddings[:, :, None])
+    return x

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "backend": "tokenizers",
+  "bos_token": null,
+  "do_lower_case": true,
+  "eos_token": null,
+  "model_max_length": 64,
+  "pad_token": "<pad>",
+  "processor_class": "Tipsv2Processor",
+  "token_type_ids_pattern": "all_zeros",
+  "tokenizer_class": "Tipsv2Tokenizer",
+  "unk_token": "<unk>"
+}