# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/tipsv2/modular_tipsv2.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_tipsv2.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # Copyright 2026 Google LLC and the HuggingFace Inc. team. All rights reserved. # # 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. import collections.abc import math from collections.abc import Callable from dataclasses import dataclass from typing import Any import torch from torch import nn from ... import initialization as init from ...activations import ACT2FN from ...backbone_utils import BackboneMixin, filter_output_hidden_states from ...masking_utils import create_bidirectional_mask from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BackboneOutput, BaseModelOutput, BaseModelOutputWithPooling from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_int from ...utils.generic import can_return_tuple, merge_with_config_defaults from ...utils.output_capturing import capture_outputs from .configuration_tipsv2 import Tipsv2Config, Tipsv2TextConfig, Tipsv2VisionConfig @auto_docstring @dataclass class Tipsv2Output(ModelOutput): r""" loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`): Contrastive loss for image-text similarity. logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`): The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text similarity scores. logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`): The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image similarity scores. text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by applying the projection layer to the pooled output of [`Tipsv2TextModel`]. image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by applying the projection layer to the pooled output of [`Tipsv2VisionModel`]. text_model_output (`BaseModelOutputWithPooling`): The output of the [`Tipsv2TextModel`]. vision_model_output (`BaseModelOutputWithPooling`): The output of the [`Tipsv2VisionModel`]. """ loss: torch.FloatTensor | None = None logits_per_image: torch.FloatTensor | None = None logits_per_text: torch.FloatTensor | None = None text_embeds: torch.FloatTensor | None = None image_embeds: torch.FloatTensor | None = None text_model_output: BaseModelOutputWithPooling = None vision_model_output: BaseModelOutputWithPooling = None def to_tuple(self) -> tuple[Any]: return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values()) class Tipsv2VisionPatchEmbeddings(nn.Module): """ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a Transformer. """ def __init__(self, config): super().__init__() image_size, patch_size = config.image_size, config.patch_size num_channels, hidden_size = config.num_channels, config.hidden_size image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size) patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size) num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0]) self.image_size = image_size self.patch_size = patch_size self.num_channels = num_channels self.num_patches = num_patches self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: num_channels = pixel_values.shape[1] if num_channels != self.num_channels: raise ValueError( "Make sure that the channel dimension of the pixel values match with the one set in the configuration." f" Expected {self.num_channels} but got {num_channels}." ) embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) return embeddings class Tipsv2VisionEmbeddings(nn.Module): """ Construct the CLS token, mask token, register tokens, position and patch embeddings. """ def __init__(self, config: Tipsv2VisionConfig) -> None: super().__init__() self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size)) self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size)) self.register_tokens = nn.Parameter(torch.zeros(1, config.num_register_tokens, config.hidden_size)) self.patch_embeddings = Tipsv2VisionPatchEmbeddings(config) num_patches = self.patch_embeddings.num_patches self.position_embeddings = nn.Parameter(torch.randn(1, num_patches + 1, config.hidden_size)) self.dropout = nn.Dropout(config.hidden_dropout_prob) self.patch_size = config.patch_size self.config = config def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor: """ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution images. This implementation supports torch.jit tracing while maintaining backwards compatibility with the original implementation. Adapted from: - https://github.com/facebookresearch/dino/blob/main/vision_transformer.py - https://github.com/facebookresearch/dinov2/blob/main/dinov2/models/vision_transformer.py """ num_patches = embeddings.shape[1] - 1 num_positions = self.position_embeddings.shape[1] - 1 # Skip interpolation for matching dimensions (unless tracing) if not torch.jit.is_tracing() and num_patches == num_positions and height == width: return self.position_embeddings # Handle class token and patch embeddings separately class_pos_embed = self.position_embeddings[:, 0] patch_pos_embed = self.position_embeddings[:, 1:] dim = embeddings.shape[-1] # Calculate new dimensions height = height // self.config.patch_size width = width // self.config.patch_size # Reshape for interpolation sqrt_num_positions = torch_int(num_positions**0.5) patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim) patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) # Store original dtype for restoration after interpolation target_dtype = patch_pos_embed.dtype # Interpolate at float32 precision patch_pos_embed = nn.functional.interpolate( patch_pos_embed.to(dtype=torch.float32), size=(torch_int(height), torch_int(width)), mode="bilinear", # Different from Dinov2 which uses bicubic interpolation align_corners=False, antialias=True, ).to(dtype=target_dtype) # Validate output dimensions if not tracing if not torch.jit.is_tracing(): if int(height) != patch_pos_embed.shape[-2] or int(width) != patch_pos_embed.shape[-1]: raise ValueError("Width or height does not match with the interpolated position embeddings") # Reshape back to original format patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) # Combine class and patch embeddings return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None) -> torch.Tensor: batch_size, _, height, width = pixel_values.shape target_dtype = self.patch_embeddings.projection.weight.dtype embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype)) if bool_masked_pos is not None: embeddings = torch.where( bool_masked_pos.unsqueeze(-1), self.mask_token.to(embeddings.dtype).unsqueeze(0), embeddings ) # add the [CLS] token to the embedded patch tokens cls_tokens = self.cls_token.expand(batch_size, -1, -1) embeddings = torch.cat((cls_tokens, embeddings), dim=1) # add positional encoding to each token embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width) # add register tokens embeddings = torch.cat( (embeddings[:, :1], self.register_tokens.expand(embeddings.shape[0], -1, -1), embeddings[:, 1:]), dim=1 ) embeddings = self.dropout(embeddings) return embeddings def eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float | None = None, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): if scaling is None: scaling = query.size(-1) ** -0.5 # Take the dot product between "query" and "key" to get the raw attention scores. attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights # Todo - Refactor as part of vision refactor. Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Tipsv2Vision class Tipsv2VisionSelfAttention(nn.Module): def __init__(self, config: Tipsv2VisionConfig): super().__init__() if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"): raise ValueError( f"The hidden size {config.hidden_size} is not a multiple of the number of attention " f"heads {config.num_attention_heads}." ) self.config = config self.num_attention_heads = config.num_attention_heads self.attention_head_size = int(config.hidden_size / config.num_attention_heads) self.all_head_size = self.num_attention_heads * self.attention_head_size self.dropout_prob = config.attention_probs_dropout_prob self.scaling = self.attention_head_size**-0.5 self.is_causal = False self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias) def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor]: batch_size = hidden_states.shape[0] new_shape = batch_size, -1, self.num_attention_heads, self.attention_head_size key_layer = self.key(hidden_states).view(*new_shape).transpose(1, 2) value_layer = self.value(hidden_states).view(*new_shape).transpose(1, 2) query_layer = self.query(hidden_states).view(*new_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, eager_attention_forward ) context_layer, attention_probs = attention_interface( self, query_layer, key_layer, value_layer, None, is_causal=self.is_causal, scaling=self.scaling, dropout=0.0 if not self.training else self.dropout_prob, **kwargs, ) new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) context_layer = context_layer.reshape(new_context_layer_shape) return context_layer, attention_probs # Todo - Refactor as part of vision refactor. Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Tipsv2Vision class Tipsv2VisionSelfOutput(nn.Module): """ The residual connection is defined in Tipsv2VisionLayer instead of here (as is the case with other models), due to the layernorm applied before each block. """ def __init__(self, config: Tipsv2VisionConfig): super().__init__() self.dense = nn.Linear(config.hidden_size, config.hidden_size) self.dropout = nn.Dropout(config.hidden_dropout_prob) def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: hidden_states = self.dense(hidden_states) hidden_states = self.dropout(hidden_states) return hidden_states # Todo - Refactor as part of vision refactor. Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Tipsv2Vision class Tipsv2VisionAttention(nn.Module): def __init__(self, config: Tipsv2VisionConfig): super().__init__() self.attention = Tipsv2VisionSelfAttention(config) self.output = Tipsv2VisionSelfOutput(config) def forward( self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: self_attn_output, _ = self.attention(hidden_states, **kwargs) output = self.output(self_attn_output, hidden_states) return output class Tipsv2VisionLayerScale(nn.Module): def __init__(self, config) -> None: super().__init__() self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size)) def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: return hidden_state * self.lambda1 class Tipsv2VisionMLP(nn.Module): def __init__(self, config) -> None: super().__init__() in_features = out_features = config.hidden_size hidden_features = int(config.hidden_size * config.mlp_ratio) self.fc1 = nn.Linear(in_features, hidden_features, bias=True) if isinstance(config.hidden_act, str): self.activation = ACT2FN[config.hidden_act] else: self.activation = config.hidden_act self.fc2 = nn.Linear(hidden_features, out_features, bias=True) def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: hidden_state = self.fc1(hidden_state) hidden_state = self.activation(hidden_state) hidden_state = self.fc2(hidden_state) return hidden_state class Tipsv2VisionSwiGLUFFN(nn.Module): def __init__(self, config) -> None: super().__init__() in_features = out_features = config.hidden_size hidden_features = int(config.hidden_size * config.mlp_ratio) hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True) self.weights_out = nn.Linear(hidden_features, out_features, bias=True) def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: hidden_state = self.weights_in(hidden_state) x1, x2 = hidden_state.chunk(2, dim=-1) hidden = nn.functional.silu(x1) * x2 return self.weights_out(hidden) class Tipsv2VisionDropPath(nn.Module): """Stochastic depth (DropPath) per sample, for residual blocks. Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth `_. """ def __init__(self, drop_prob: float = 0.0) -> None: super().__init__() self.drop_prob = drop_prob def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.drop_prob == 0.0 or not self.training: return hidden_states keep_prob = 1 - self.drop_prob shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1) random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device) random_tensor = torch.floor(random_tensor + keep_prob) return hidden_states.div(keep_prob) * random_tensor def extra_repr(self) -> str: return f"p={self.drop_prob}" class Tipsv2VisionLayer(GradientCheckpointingLayer): """This corresponds to the Block class in the original implementation.""" def __init__(self, config: Tipsv2VisionConfig) -> None: super().__init__() self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.attention = Tipsv2VisionAttention(config) self.layer_scale1 = Tipsv2VisionLayerScale(config) self.drop_path = Tipsv2VisionDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity() self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) if config.use_swiglu_ffn: self.mlp = Tipsv2VisionSwiGLUFFN(config) else: self.mlp = Tipsv2VisionMLP(config) self.layer_scale2 = Tipsv2VisionLayerScale(config) def forward( self, hidden_states: torch.Tensor, ) -> torch.Tensor: hidden_states_norm = self.norm1(hidden_states) self_attention_output = self.attention(hidden_states_norm) self_attention_output = self.layer_scale1(self_attention_output) # first residual connection hidden_states = self.drop_path(self_attention_output) + hidden_states # in Tipsv2Vision, layernorm is also applied after self-attention layer_output = self.norm2(hidden_states) layer_output = self.mlp(layer_output) layer_output = self.layer_scale2(layer_output) # second residual connection layer_output = self.drop_path(layer_output) + hidden_states return layer_output @auto_docstring class Tipsv2VisionPreTrainedModel(PreTrainedModel): config: Tipsv2VisionConfig base_model_prefix = "vision_model" main_input_name = "pixel_values" input_modalities = ("image",) supports_gradient_checkpointing = True _no_split_modules = ["Tipsv2VisionLayer"] _supports_sdpa = True _supports_flash_attn = True _supports_flex_attn = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": Tipsv2VisionLayer, "attentions": Tipsv2VisionSelfAttention, } _keys_to_ignore_on_load_unexpected = {"text_encoder"} @torch.no_grad() def _init_weights(self, module: nn.Linear | nn.Conv2d | nn.LayerNorm) -> None: """Initialize the weights""" super()._init_weights(module) if isinstance(module, (nn.Linear, nn.Conv2d)): init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.bias is not None: init.zeros_(module.bias) elif isinstance(module, Tipsv2VisionEmbeddings): init.trunc_normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range) init.trunc_normal_(module.cls_token, mean=0.0, std=self.config.initializer_range) init.zeros_(module.mask_token) init.zeros_(module.register_tokens) elif isinstance(module, Tipsv2VisionLayerScale): # noqa: F821 init.constant_(module.lambda1, self.config.layerscale_value) class Tipsv2VisionEncoder(Tipsv2VisionPreTrainedModel): def __init__(self, config: Tipsv2VisionConfig): super().__init__(config) self.layer = nn.ModuleList([Tipsv2VisionLayer(config) for _ in range(config.num_hidden_layers)]) self.post_init() @merge_with_config_defaults @capture_outputs(tie_last_hidden_states=False) def forward(self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BaseModelOutput: for layer_module in self.layer: hidden_states = layer_module(hidden_states) return BaseModelOutput(last_hidden_state=hidden_states) @auto_docstring class Tipsv2VisionModel(Tipsv2VisionPreTrainedModel): def __init__(self, config: Tipsv2VisionConfig): super().__init__(config) self.config = config self.embeddings = Tipsv2VisionEmbeddings(config) self.encoder = Tipsv2VisionEncoder(config) self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> Tipsv2VisionPatchEmbeddings: return self.embeddings.patch_embeddings @can_return_tuple @auto_docstring def forward( self, pixel_values: torch.Tensor | None = None, bool_masked_pos: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for pre-training. Example: ```python >>> import torch >>> from transformers import AutoConfig, AutoImageProcessor, AutoModel >>> from transformers.image_utils import load_image >>> model_id = "google/tipsv2-b14" >>> config = AutoConfig.from_pretrained(model_id) >>> model = AutoModel.from_pretrained(model_id, config=config.vision_config, device_map="auto") >>> image_processor = AutoImageProcessor.from_pretrained(model_id) >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg") >>> inputs = image_processor(images=image, return_tensors="pt").to(model.device) >>> with torch.no_grad(): ... outputs = model(**inputs) >>> # Tipsv2 repurposes the register token from DINOv2-with-registers as a secondary class token >>> sequence = outputs.last_hidden_state # (batch_size, 1 + num_register_tokens + num_patches, hidden_size) >>> cls_token_1 = sequence[:, 0] # supervised by web alt-text captions >>> cls_token_2 = sequence[:, 1 : 1 + model.config.num_register_tokens] # supervised by synthetic captions ```""" if pixel_values is None: raise ValueError("You have to specify pixel_values") embedding_output = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos) encoder_outputs: BaseModelOutput = self.encoder(embedding_output, **kwargs) sequence_output = encoder_outputs.last_hidden_state sequence_output = self.layernorm(sequence_output) pooled_output = sequence_output[:, 0, :] return BaseModelOutputWithPooling( last_hidden_state=sequence_output, pooler_output=pooled_output, hidden_states=encoder_outputs.hidden_states, attentions=encoder_outputs.attentions, ) @auto_docstring( custom_intro=""" Tipsv2Vision backbone, to be used with frameworks like DETR and MaskFormer. """ ) class Tipsv2VisionBackbone(BackboneMixin, Tipsv2VisionPreTrainedModel): def __init__(self, config: "Tipsv2Config | Tipsv2VisionConfig"): # reassign config to vision_config for compatibility with AutoBackbone if isinstance(config, Tipsv2Config): config = config.vision_config super().__init__(config) else: super().__init__(config) self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)] self.embeddings = Tipsv2VisionEmbeddings(config) self.encoder = Tipsv2VisionEncoder(config) self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.num_register_tokens = config.num_register_tokens # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self) -> Tipsv2VisionPatchEmbeddings: return self.embeddings.patch_embeddings @can_return_tuple @filter_output_hidden_states @auto_docstring def forward( self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> BackboneOutput: r""" Examples: ```python >>> import torch >>> from transformers import AutoBackbone, Tipsv2VisionConfig >>> config = Tipsv2VisionConfig(out_features=["stage3", "stage6", "stage9", "stage12"]) >>> model = AutoBackbone.from_config(config) >>> pixel_values = torch.randn(1, 3, 448, 448) >>> outputs = model(pixel_values) >>> feature_maps = outputs.feature_maps >>> list(feature_maps[-1].shape) [1, 768, 32, 32] ```""" kwargs["output_hidden_states"] = True # required to extract layers for the stages embedding_output = self.embeddings(pixel_values) output: BaseModelOutput = self.encoder(embedding_output, **kwargs) hidden_states = output.hidden_states feature_maps = [] for stage, hidden_state in zip(self.stage_names, hidden_states): if stage in self.out_features: if self.config.apply_layernorm: hidden_state = self.layernorm(hidden_state) if self.config.reshape_hidden_states: hidden_state = hidden_state[:, 1 + self.num_register_tokens :] # this was actually a bug in the original implementation that we copied here, # cause normally the order is height, width batch_size, _, height, width = pixel_values.shape patch_size = self.config.patch_size hidden_state = hidden_state.reshape(batch_size, height // patch_size, width // patch_size, -1) hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() feature_maps.append(hidden_state) return BackboneOutput( feature_maps=tuple(feature_maps), hidden_states=hidden_states, attentions=output.attentions, ) class Tipsv2SinusoidalPositionalEmbedding(nn.Module): """This module produces sinusoidal positional embeddings of any length.""" def __init__(self, config: Tipsv2TextConfig): super().__init__() self.make_weights( num_embeddings=config.max_position_embeddings, embedding_dim=config.hidden_size, padding_idx=None ) def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: int | None = None): emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx) if hasattr(self, "weights"): # in forward put the weights on the correct dtype and device of the param emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device) self.register_buffer("weights", emb_weights, persistent=False) @staticmethod def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: int | None = None): """ Build sinusoidal embeddings. This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of "Attention Is All You Need". """ half_dim = embedding_dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb) emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0) emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1) if embedding_dim % 2 == 1: # zero pad emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1) if padding_idx is not None: emb[padding_idx, :] = 0 return emb.to(torch.get_default_dtype()) @torch.no_grad() def forward(self, position_ids: torch.LongTensor) -> torch.Tensor: return self.weights[position_ids] class Tipsv2TextEmbeddings(nn.Module): def __init__(self, config: Tipsv2TextConfig): super().__init__() embed_dim = config.hidden_size self.token_embedding = nn.Embedding(config.vocab_size, embed_dim) self.position_embedding = Tipsv2SinusoidalPositionalEmbedding(config) # position_ids (1, len position emb) is contiguous in memory and exported when serialized self.register_buffer( "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False ) self.embed_scale = math.sqrt(config.hidden_size) if config.scale_sqrt_depth else 1.0 def forward( self, input_ids: torch.LongTensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, ) -> torch.Tensor: seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2] max_position_embedding = self.position_ids.shape[-1] if seq_length > max_position_embedding: raise ValueError( f"Sequence length must be less than max_position_embeddings (got `sequence length`: " f"{seq_length} and max_position_embeddings: {max_position_embedding}" ) if position_ids is None: position_ids = self.position_ids[:, :seq_length] if inputs_embeds is None: inputs_embeds = self.token_embedding(input_ids) # Additional scale compared to CLIP embeddings inputs_embeds = inputs_embeds * self.embed_scale position_embeddings = self.position_embedding(position_ids).to(dtype=inputs_embeds.dtype) return inputs_embeds + position_embeddings # Identical to CLIP's implementation but couldn't import it because the vision model # already requires eager_attention_forward from DINOv2 which results in modular # conflicts. def text_eager_attention_forward( module: nn.Module, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor | None, scaling: float, dropout: float = 0.0, **kwargs: Unpack[TransformersKwargs], ): attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling if attention_mask is not None: attn_weights = attn_weights + attention_mask attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) attn_output = torch.matmul(attn_weights, value) attn_output = attn_output.transpose(1, 2).contiguous() return attn_output, attn_weights class Tipsv2TextAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: Tipsv2TextConfig): super().__init__() self.config = config self.embed_dim = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.embed_dim // self.num_heads self.scale = self.head_dim**-0.5 self.dropout = config.attention_dropout self.is_causal = False self.k_proj = nn.Linear(self.embed_dim, self.embed_dim) self.v_proj = nn.Linear(self.embed_dim, self.embed_dim) self.q_proj = nn.Linear(self.embed_dim, self.embed_dim) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple[torch.Tensor, torch.Tensor | None]: """Input shape: Batch x Time x Channel""" input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) queries = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) keys = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) values = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( self.config._attn_implementation, text_eager_attention_forward ) attn_output, attn_weights = attention_interface( self, queries, keys, values, attention_mask, scaling=self.scale, dropout=0.0 if not self.training else self.dropout, is_causal=self.is_causal, **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() attn_output = self.out_proj(attn_output) return attn_output, attn_weights class Tipsv2MLP(nn.Module): def __init__(self, config): super().__init__() self.config = config self.activation_fn = ACT2FN[config.hidden_act] self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = self.fc1(hidden_states) hidden_states = self.activation_fn(hidden_states) hidden_states = self.fc2(hidden_states) return hidden_states class Tipsv2TextEncoderLayer(GradientCheckpointingLayer): def __init__(self, config: Tipsv2TextConfig): super().__init__() self.embed_dim = config.hidden_size self.self_attn = Tipsv2TextAttention(config) self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) self.mlp = Tipsv2MLP(config) self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, **kwargs: Unpack[TransformersKwargs], ) -> torch.FloatTensor: residual = hidden_states hidden_states = self.layer_norm1(hidden_states) hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.layer_norm2(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class Tipsv2TextEncoder(nn.Module): """ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a [`Tipsv2TextEncoderLayer`]. Args: config: Tipsv2TextConfig """ def __init__(self, config: Tipsv2TextConfig): super().__init__() self.config = config self.layers = nn.ModuleList([Tipsv2TextEncoderLayer(config) for _ in range(config.num_hidden_layers)]) self.gradient_checkpointing = False def forward( self, inputs_embeds, attention_mask: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutput: hidden_states = inputs_embeds for encoder_layer in self.layers: hidden_states = encoder_layer( hidden_states, attention_mask, **kwargs, ) return BaseModelOutput( last_hidden_state=hidden_states, ) @auto_docstring class Tipsv2TextPreTrainedModel(PreTrainedModel): config: Tipsv2TextConfig base_model_prefix = "text_model" main_input_name = "input_ids" input_modalities = ["text"] supports_gradient_checkpointing = True _keys_to_ignore_on_load_unexpected = {"vision_encoder"} _no_split_modules = ["Tipsv2TextEmbeddings", "Tipsv2TextEncoderLayer"] _supports_sdpa = True _supports_flash_attn = True _supports_flex_attn = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": Tipsv2TextEncoderLayer, "attentions": Tipsv2TextAttention, } @torch.no_grad() def _init_weights(self, module): super()._init_weights(module) if isinstance(module, Tipsv2TextEmbeddings): init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) elif isinstance(module, Tipsv2SinusoidalPositionalEmbedding): num_embeddings, embedding_dim = module.weights.shape embedding_weights = module.get_embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, padding_idx=None, ) init.copy_(module.weights, embedding_weights) @auto_docstring( custom_intro=""" The TIPSv2 text tower without any projection head on top. """ ) class Tipsv2TextModel(Tipsv2TextPreTrainedModel): _input_embed_layer = "token_embedding" def __init__(self, config: Tipsv2TextConfig): super().__init__(config) self.config = config self.embeddings = Tipsv2TextEmbeddings(config) self.encoder = Tipsv2TextEncoder(config) self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.post_init() @merge_with_config_defaults @capture_outputs(tie_last_hidden_states=False) @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPooling: r""" Example: ```python >>> import torch >>> from transformers import AutoConfig, Tipsv2TextModel, AutoProcessor >>> model_id = "google/tipsv2-b14" >>> config = AutoConfig.from_pretrained(model_id) >>> model = Tipsv2TextModel.from_pretrained(model_id, config=config.text_config, device_map="auto") >>> processor = AutoProcessor.from_pretrained(model_id) >>> candidate_labels = ["a photo of a cat", "a photo of a dog", "a photo of a car"] >>> inputs = processor(text=candidate_labels, return_tensors="pt").to(model.device) >>> with torch.no_grad(): ... outputs = model(**inputs) >>> text_embeds = outputs.pooler_output # (batch_size, hidden_size) ```""" if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds) pooling_mask = attention_mask attention_mask = create_bidirectional_mask( config=self.config, inputs_embeds=hidden_states, attention_mask=attention_mask, ) encoder_outputs: BaseModelOutput = self.encoder( inputs_embeds=hidden_states, attention_mask=attention_mask, **kwargs, ) last_hidden_state = encoder_outputs.last_hidden_state last_hidden_state = self.final_layer_norm(last_hidden_state) if pooling_mask is not None: masked_hidden_state = last_hidden_state * pooling_mask[..., None] pooled_output = masked_hidden_state.sum(dim=1) / ( pooling_mask.sum(dim=1, keepdim=True) + self.config.pooling_epsilon ) else: pooled_output = last_hidden_state.mean(dim=1) return BaseModelOutputWithPooling( last_hidden_state=last_hidden_state, pooler_output=pooled_output, ) @auto_docstring class Tipsv2PreTrainedModel(PreTrainedModel): config: Tipsv2Config base_model_prefix = "model" input_modalities = ["image", "text"] supports_gradient_checkpointing = True _no_split_modules = [ "Tipsv2TextEmbeddings", "Tipsv2TextEncoderLayer", "Tipsv2VisionEmbeddings", "Tipsv2VisionLayer", ] _supports_sdpa = True _supports_flash_attn = True _supports_flex_attn = True _supports_attention_backend = True def _init_weights(self, module): super()._init_weights(module) if isinstance(module, Tipsv2Model): init.constant_(module.temperature, module.config.temperature_init_value) # contrastive loss function, adapted from # https://sachinruk.github.io/blog/2021-03-07-tipsv2.html def contrastive_loss(logits: torch.Tensor) -> torch.Tensor: return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device)) def image_text_contrastive_loss(similarity: torch.Tensor) -> torch.Tensor: caption_loss = contrastive_loss(similarity) image_loss = contrastive_loss(similarity.T) return (caption_loss + image_loss) / 2.0 def _get_vector_norm(tensor: torch.Tensor) -> torch.Tensor: """ This method is equivalent to tensor.norm(p=2, dim=-1, keepdim=True) and used to make model `executorch` exportable. See issue https://github.com/pytorch/executorch/issues/3566 """ square_tensor = torch.pow(tensor, 2) sum_tensor = torch.sum(square_tensor, dim=-1, keepdim=True) normed_tensor = torch.pow(sum_tensor, 0.5) return normed_tensor @auto_docstring class Tipsv2Model(Tipsv2PreTrainedModel): _keys_to_ignore_on_load_missing = ["temperature"] # learnable but not in published checkpoints def __init__(self, config: Tipsv2Config): super().__init__(config) self.text_model = Tipsv2TextModel._from_config(config.text_config) self.vision_model = Tipsv2VisionModel._from_config(config.vision_config) self.temperature = nn.Parameter(torch.tensor(config.temperature_init_value)) self.post_init() @can_return_tuple @auto_docstring def get_text_features( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: return self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, **kwargs, ) @can_return_tuple @auto_docstring def get_image_features( self, pixel_values: torch.FloatTensor, bool_masked_pos: torch.Tensor | None = None, **kwargs: Unpack[TransformersKwargs], ) -> tuple | BaseModelOutputWithPooling: r""" bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for pre-training. """ return self.vision_model( pixel_values=pixel_values, bool_masked_pos=bool_masked_pos, **kwargs, ) def get_input_embeddings(self) -> nn.Module: return self.text_model.get_input_embeddings() def set_input_embeddings(self, value: nn.Module): self.text_model.set_input_embeddings(value) @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor | None = None, pixel_values: torch.FloatTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, inputs_embeds: torch.FloatTensor | None = None, bool_masked_pos: torch.Tensor | None = None, return_loss: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> Tipsv2Output: r""" bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for pre-training. return_loss (`bool`, *optional*): Whether or not to return the contrastive loss when both image and text inputs are provided. Example: ```python >>> import torch >>> from transformers import AutoModel, AutoProcessor >>> from transformers.image_utils import load_image >>> model_id = "google/tipsv2-b14" >>> model = AutoModel.from_pretrained(model_id, device_map="auto") >>> processor = AutoProcessor.from_pretrained(model_id) >>> image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg") >>> candidate_labels = ["a photo of a cat", "a photo of a dog", "a photo of a car"] >>> inputs = processor(text=candidate_labels, images=image, return_tensors="pt").to(model.device) >>> with torch.no_grad(): ... outputs = model(**inputs) >>> probs = outputs.logits_per_image.softmax(dim=1) >>> most_likely_idx = probs.argmax(dim=1).item() >>> most_likely_label = candidate_labels[most_likely_idx] >>> print(f"Most likely label: '{most_likely_label}' with probability: {probs[0][most_likely_idx].item():.3f}") Most likely label: 'a photo of a cat' with probability: 0.975 ``` """ if pixel_values is None and input_ids is None and inputs_embeds is None: raise ValueError("You have to specify pixel_values, input_ids, or inputs_embeds") vision_outputs = None image_embeds = None if pixel_values is not None: vision_outputs = self.vision_model( pixel_values=pixel_values, bool_masked_pos=bool_masked_pos, return_dict=True, **kwargs, ) image_embeds = vision_outputs.pooler_output image_embeds = image_embeds / _get_vector_norm(image_embeds) text_outputs = None text_embeds = None if input_ids is not None or inputs_embeds is not None: text_outputs = self.text_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, inputs_embeds=inputs_embeds, return_dict=True, **kwargs, ) text_embeds = text_outputs.pooler_output text_embeds = text_embeds / _get_vector_norm(text_embeds) logits_per_text = None logits_per_image = None loss = None if image_embeds is not None and text_embeds is not None: logits_per_text = torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) logits_per_text = logits_per_text / self.temperature logits_per_image = logits_per_text.t() if return_loss: loss = image_text_contrastive_loss(logits_per_text) return Tipsv2Output( loss=loss, logits_per_image=logits_per_image, logits_per_text=logits_per_text, text_embeds=text_embeds, image_embeds=image_embeds, text_model_output=text_outputs, vision_model_output=vision_outputs, ) __all__ = [ "Tipsv2Model", "Tipsv2PreTrainedModel", "Tipsv2TextModel", "Tipsv2TextPreTrainedModel", "Tipsv2VisionBackbone", "Tipsv2VisionModel", "Tipsv2VisionPreTrainedModel", ]