content: # Copyright 2024 The HuggingFace Inc. team.
  # SPDX-License-Identifier: Apache-2.0
  
  """PyTorch Siglip model."""
  
  import math
  import warnings
  from dataclasses import dataclass
  from typing import Any, Optional, Tuple, Union
  
  import numpy as np
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from torch.nn.init import _calculate_fan_in_and_fan_out
  
  from transformers.activations import ACT2FN
  from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers.utils import (
      ModelOutput,
      add_start_docstrings,
      add_start_docstrings_to_model_forward,
      is_flash_attn_2_available,
      is_flash_attn_greater_or_equal_2_10,
      logging,
      replace_return_docstrings,
      torch_int,
  )
  from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig
  
  
  if is_flash_attn_2_available():
      from transformers.modeling_flash_attention_utils import _flash_attention_forward
  
  
  logger = logging.get_logger(__name__)
  
  # General docstring
  _CONFIG_FOR_DOC = "SiglipConfig"
  _CHECKPOINT_FOR_DOC = "google/siglip-base-patch16-224"
  
  
  def _trunc_normal_(tensor, mean, std, a, b):
      # Cut & paste from PyTorch official master until it's in a few official releases - RW
      # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
      def norm_cdf(x):
          # Computes standard normal cumulative distribution function
          return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0
  
      if (mean < a - 2 * std) or (mean > b + 2 * std):
          warnings.warn(
              "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
              "The distribution of values may be incorrect.",
              stacklevel=2,
          )
  
      # Values are generated by using a truncated uniform distribution and
      # then using the inverse CDF for the normal distribution.
      # Get upper and lower cdf values
      l = norm_cdf((a - mean) / std)
      u = norm_cdf((b - mean) / std)
  
      # Uniformly fill tensor with values from [l, u], then translate to
      # [2l-1, 2u-1].
      tensor.uniform_(2 * l - 1, 2 * u - 1)
  
      # Use inverse cdf transform for normal distribution to get truncated
      # standard normal
      tensor.erfinv_()
  
      # Transform to proper mean, std
      tensor.mul_(std * math.sqrt(2.0))
      tensor.add_(mean)
  
      # Clamp to ensure it's in the proper range
      tensor.clamp_(min=a, max=b)
  
  
  def trunc_normal_tf_(
      tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
  ) -> torch.Tensor:
      """Fills the input Tensor with values drawn from a truncated
      normal distribution. The values are effectively drawn from the
      normal distribution :math:`\\mathcal{N}(\text{mean}, \text{std}^2)`
      with values outside :math:`[a, b]` redrawn until they are within
      the bounds. The method used for generating the random values works
      best when :math:`a \\leq \text{mean} \\leq b`.
  
      NOTE: this 'tf' variant behaves closer to Tensorflow / JAX impl where the
      bounds [a, b] are applied when sampling the normal distribution with mean=0, std=1.0
      and the result is subsequently scaled and shifted by the mean and std args.
  
      Args:
          tensor: an n-dimensional `torch.Tensor`
          mean: the mean of the normal distribution
          std: the standard deviation of the normal distribution
          a: the minimum cutoff value
          b: the maximum cutoff value
      """
      with torch.no_grad():
          _trunc_normal_(tensor, 0, 1.0, a, b)
          tensor.mul_(std).add_(mean)
  
  
  def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):
      fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
      if mode == "fan_in":
          denom = fan_in
      elif mode == "fan_out":
          denom = fan_out
      elif mode == "fan_avg":
          denom = (fan_in + fan_out) / 2
  
      variance = scale / denom
  
      if distribution == "truncated_normal":
          # constant is stddev of standard normal truncated to (-2, 2)
          trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
      elif distribution == "normal":
          with torch.no_grad():
              tensor.normal_(std=math.sqrt(variance))
      elif distribution == "uniform":
          bound = math.sqrt(3 * variance)
          with torch.no_grad():
              tensor.uniform_(-bound, bound)
      else:
          raise ValueError(f"invalid distribution {distribution}")
  
  
  def lecun_normal_(tensor):
      variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")
  
  
  def default_flax_embed_init(tensor):
      variance_scaling_(tensor, mode="fan_in", distribution="normal")
  
  
  @dataclass
  # Copied from transformers.models.clip.modeling_clip.CLIPVisionModelOutput with CLIP->Siglip
  class SiglipVisionModelOutput(ModelOutput):
      """
      Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
  
      Args:
          image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
              The image embeddings obtained by applying the projection layer to the pooler_output.
          last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
              Sequence of hidden-states at the output of the last layer of the model.
          hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
              Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
              one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
  
              Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
          attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
              Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
              sequence_length)`.
  
              Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
              heads.
      """
  
      image_embeds: Optional[torch.FloatTensor] = None
      last_hidden_state: torch.FloatTensor = None
      hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
      attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
  
  
  @dataclass
  # Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Siglip
  class SiglipTextModelOutput(ModelOutput):
      """
      Base class for text model's outputs that also contains a pooling of the last hidden states.
  
      Args:
          text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
              The text embeddings obtained by applying the projection layer to the pooler_output.
          last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
              Sequence of hidden-states at the output of the last layer of the model.
          hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
              Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
              one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
  
              Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
          attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
              Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
              sequence_length)`.
  
              Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
              heads.
      """
  
      text_embeds: Optional[torch.FloatTensor] = None
      last_hidden_state: torch.FloatTensor = None
      hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
      attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
  
  
  @dataclass
  # Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Siglip
  class SiglipOutput(ModelOutput):
      """
      Args:
          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 [`SiglipTextModel`].
          image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
              The image embeddings obtained by applying the projection layer to the pooled output of [`SiglipVisionModel`].
          text_model_output (`BaseModelOutputWithPooling`):
              The output of the [`SiglipTextModel`].
          vision_model_output (`BaseModelOutputWithPooling`):
              The output of the [`SiglipVisionModel`].
      """
  
      loss: Optional[torch.FloatTensor] = None
      logits_per_image: torch.FloatTensor = None
      logits_per_text: torch.FloatTensor = None
      text_embeds: torch.FloatTensor = None
      image_embeds: torch.FloatTensor = None
      text_model_output: BaseModelOutputWithPooling = None
      vision_model_output: BaseModelOutputWithPooling = None
  
      def to_tuple(self) -> Tuple[Any]:
          return tuple(
              self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
              for k in self.keys()
          )
  
  
  class SiglipVisionEmbeddings(nn.Module):
      def __init__(self, config: SiglipVisionConfig):
          super().__init__()
          self.config = config
          self.embed_dim = config.hidden_size
          self.image_size = config.image_size
          self.patch_size = config.patch_size
  
          self.patch_embedding = nn.Conv2d(
              in_channels=config.num_channels,
              out_channels=self.embed_dim,
              kernel_size=self.patch_size,
              stride=self.patch_size,
              padding="valid",
          )
  
          self.num_patches = (self.image_size // self.patch_size) ** 2
          self.num_positions = self.num_patches
          self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
          self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
  
      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 method is also adapted to support torch.jit tracing and no class embeddings.
  
          Adapted from:
          - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
          - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
          """
  
          num_patches = embeddings.shape[1]
          num_positions = self.position_embedding.weight.shape[0]
  
          # always interpolate when tracing to ensure the exported model works for dynamic input shapes
          if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
              return self.position_embedding(self.position_ids)
  
          patch_pos_embed = self.position_embedding.weight.unsqueeze(0)
  
          dim = embeddings.shape[-1]
  
          new_height = height // self.patch_size
          new_width = width // self.patch_size
  
          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)
  
          patch_pos_embed = nn.functional.interpolate(
              patch_pos_embed,
              size=(new_height, new_width),
              mode="bicubic",
              align_corners=False,
          )
  
          patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
          return patch_pos_embed
  
      def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor:
          _, _, height, width = pixel_values.shape
          patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
          embeddings = patch_embeds.flatten(2).transpose(1, 2)
  
          if interpolate_pos_encoding:
              embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
          else:
              embeddings = embeddings + self.position_embedding(self.position_ids)
          return embeddings
  
  
  # Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->Siglip
  class SiglipTextEmbeddings(nn.Module):
      def __init__(self, config: SiglipTextConfig):
          super().__init__()
          embed_dim = config.hidden_size
  
          self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
          self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
  
          # 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
          )
  
      def forward(
          self,
          input_ids: Optional[torch.LongTensor] = None,
          position_ids: Optional[torch.LongTensor] = None,
          inputs_embeds: Optional[torch.FloatTensor] = None,
      ) -> torch.Tensor:
          seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
  
          if position_ids is None:
              position_ids = self.position_ids[:, :seq_length]
  
          if inputs_embeds is None:
              inputs_embeds = self.token_embedding(input_ids)
  
          position_embeddings = self.position_embedding(position_ids)
          embeddings = inputs_embeds + position_embeddings
  
          return embeddings
  
  
  class SiglipAttention(nn.Module):
      """Multi-headed attention from 'Attention Is All You Need' paper"""
  
      # Copied from transformers.models.clip.modeling_clip.CLIPAttention.__init__
      def __init__(self, config):
          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
          if self.head_dim * self.num_heads != self.embed_dim:
              raise ValueError(
                  f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                  f" {self.num_heads})."
              )
          self.scale = self.head_dim**-0.5
          self.dropout = config.attention_dropout
  
          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: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = False,
      ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
          """Input shape: Batch x Time x Channel"""
  
          batch_size, q_len, _ = hidden_states.size()
  
          query_states = self.q_proj(hidden_states)
          key_states = self.k_proj(hidden_states)
          value_states = self.v_proj(hidden_states)
  
          query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
  
          k_v_seq_len = key_states.shape[-2]
          attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
  
          if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
              raise ValueError(
                  f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
                  f" {attn_weights.size()}"
              )
  
          if attention_mask is not None:
              if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
                  raise ValueError(
                      f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, but is {attention_mask.size()}"
                  )
              attn_weights = attn_weights + attention_mask
  
          # upcast attention to fp32
          attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
          attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
          attn_output = torch.matmul(attn_weights, value_states)
  
          if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
              raise ValueError(
                  f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
                  f" {attn_output.size()}"
              )
  
          attn_output = attn_output.transpose(1, 2).contiguous()
          attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
  
          attn_output = self.out_proj(attn_output)
  
          return attn_output, attn_weights
  
  
  class SiglipFlashAttention2(SiglipAttention):
      """
      SiglipAttention flash attention module. This module inherits from `SiglipAttention` as the weights of the module stays
      untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
      flash attention and deal with padding tokens in case the input contains any of them.
      """
  
      is_causal = False
  
      # Copied from transformers.models.llama.modeling_llama.LlamaFlashAttention2.__init__
      def __init__(self, *args, **kwargs):
          super().__init__(*args, **kwargs)
  
          # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
          # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
          # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
          self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
  
      # Adapted from transformers.models.llama.modeling_llama.LlamaFlashAttention2.forward
      def forward(
          self,
          hidden_states: torch.Tensor,
          attention_mask: Optional[torch.LongTensor] = None,
          output_attentions: bool = False,
      ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
          output_attentions = False
  
          batch_size, q_len, _ = hidden_states.size()
  
          query_states = self.q_proj(hidden_states)
          key_states = self.k_proj(hidden_states)
          value_states = self.v_proj(hidden_states)
  
          # Flash attention requires the input to have the shape
          # batch_size x seq_length x head_dim x hidden_dim
          # therefore we just need to keep the original shape
          query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
  
          # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
          # to be able to avoid many of these transpose/reshape/view.
          query_states = query_states.transpose(1, 2)
          key_states = key_states.transpose(1, 2)
          value_states = value_states.transpose(1, 2)
  
          dropout_rate = self.dropout if self.training else 0.0
  
          # In PEFT, usually we cast the layer norms in float32 for training stability reasons
          # therefore the input hidden states gets silently casted in float32. Hence, we need
          # cast them back in the correct dtype just to be sure everything works as expected.
          # This might slowdown training & inference so it is recommended to not cast the LayerNorms
          # in fp32.
  
          input_dtype = query_states.dtype
          if input_dtype == torch.float32:
              if torch.is_autocast_enabled():
                  target_dtype = torch.get_autocast_gpu_dtype()
              # Handle the case where the model is quantized
              elif hasattr(self.config, "_pre_quantization_dtype"):
                  target_dtype = self.config._pre_quantization_dtype
              else:
                  target_dtype = self.q_proj.weight.dtype
  
              logger.warning_once(
                  f"The input hidden states seems to be silently casted in float32, this might be related to"
                  f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
                  f" {target_dtype}."
              )
  
              query_states = query_states.to(target_dtype)
              key_states = key_states.to(target_dtype)
              value_states = value_states.to(target_dtype)
  
          attn_output = _flash_attention_forward(
              query_states,
              key_states,
              value_states,
              attention_mask,
              q_len,
              dropout=dropout_rate,
              is_causal=self.is_causal,
              use_top_left_mask=self._flash_attn_uses_top_left_mask,
          )
  
          attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim).contiguous()
          attn_output = self.out_proj(attn_output)
  
          if not output_attentions:
              attn_weights = None
  
          return attn_output, attn_weights
  
  
  class SiglipSdpaAttention(SiglipAttention):
      """
      Siglip attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
      `SiglipAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
      SDPA API.
      """
  
      is_causal = False
  
      # Adapted from SiglipAttention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
      def forward(
          self,
          hidden_states: torch.Tensor,
          attention_mask: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = False,
      ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
          if output_attentions:
              # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
              logger.warning_once(
                  "SiglipModel is using SiglipSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
                  'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
              )
              return super().forward(
                  hidden_states=hidden_states,
                  attention_mask=attention_mask,
                  output_attentions=output_attentions,
              )
  
          batch_size, q_len, _ = hidden_states.size()
  
          query_states = self.q_proj(hidden_states)
          key_states = self.k_proj(hidden_states)
          value_states = self.v_proj(hidden_states)
  
          query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
          value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
  
          # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
          # Reference: https://github.com/pytorch/pytorch/issues/112577.
          if query_states.device.type == "cuda" and attention_mask is not None:
              query_states = query_states.contiguous()
              key_states = key_states.contiguous()
              value_states = value_states.contiguous()
  
          # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
          # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
          is_causal = True if self.is_causal and q_len > 1 else False
  
          attn_output = torch.nn.functional.scaled_dot_product_attention(
              query_states,
              key_states,
              value_states,
              attn_mask=attention_mask,
              dropout_p=self.dropout if self.training else 0.0,
              is_causal=is_causal,
          )
  
          attn_output = attn_output.transpose(1, 2).contiguous()
          attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
  
          attn_output = self.out_proj(attn_output)
  
          return attn_output, None
  
  
  SIGLIP_ATTENTION_CLASSES = {
      "eager": SiglipAttention,
      "flash_attention_2": SiglipFlashAttention2,
      "sdpa": SiglipSdpaAttention,
  }
  
  
  # Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->Siglip
  class SiglipMLP(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 SiglipEncoderLayer(nn.Module):
      def __init__(self, config: SiglipConfig):
          super().__init__()
          self.embed_dim = config.hidden_size
          self.self_attn = SIGLIP_ATTENTION_CLASSES[config._attn_implementation](config=config)
          self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
          self.mlp = SiglipMLP(config)
          self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
  
      # Ignore copy
      def forward(
          self,
          hidden_states: torch.Tensor,
          attention_mask: torch.Tensor,
          output_attentions: Optional[bool] = False,
      ) -> Tuple[torch.FloatTensor]:
          """
          Args:
              hidden_states (`torch.FloatTensor`):
                  Input to the layer of shape `(batch, seq_len, embed_dim)`.
              attention_mask (`torch.FloatTensor`):
                  Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very large negative values.
              output_attentions (`bool`, *optional*, defaults to `False`):
                  Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                  returned tensors for more detail.
          """
          residual = hidden_states
  
          hidden_states = self.layer_norm1(hidden_states)
          hidden_states, attn_weights = self.self_attn(
              hidden_states=hidden_states,
              attention_mask=attention_mask,
              output_attentions=output_attentions,
          )
          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
  
          outputs = (hidden_states,)
  
          if output_attentions:
              outputs += (attn_weights,)
  
          return outputs
  
  
  class SiglipPreTrainedModel(PreTrainedModel):
      """
      An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
      models.
      """
  
      config_class = SiglipConfig
      base_model_prefix = "siglip"
      supports_gradient_checkpointing = True
  
      _no_split_modules = [
          "SiglipTextEmbeddings",
          "SiglipEncoderLayer",
          "SiglipVisionEmbeddings",
          "SiglipEncoderLayer",
          "SiglipMultiheadAttentionPoolingHead",
      ]
      _supports_flash_attn_2 = True
      _supports_sdpa = True
  
      def _init_weights(self, module):
          """Initialize the weights"""
          if isinstance(module, SiglipVisionEmbeddings):
              width = (
                  self.config.vision_config.hidden_size
                  if isinstance(self.config, SiglipConfig)
                  else self.config.hidden_size
              )
              nn.init.normal_(module.position_embedding.weight, std=1 / np.sqrt(width))
          elif isinstance(module, nn.Embedding):
              default_flax_embed_init(module.weight)
          elif isinstance(module, SiglipAttention):
              nn.init.xavier_uniform_(module.q_proj.weight)
              nn.init.xavier_uniform_(module.k_proj.weight)
              nn.init.xavier_uniform_(module.v_proj.weight)
              nn.init.xavier_uniform_(module.out_proj.weight)
              nn.init.zeros_(module.q_proj.bias)
              nn.init.zeros_(module.k_proj.bias)
              nn.init.zeros_(module.v_proj.bias)
              nn.init.zeros_(module.out_proj.bias)
          elif isinstance(module, SiglipMLP):
              nn.init.xavier_uniform_(module.fc1.weight)
              nn.init.xavier_uniform_(module.fc2.weight)
              nn.init.normal_(module.fc1.bias, std=1e-6)
              nn.init.normal_(module.fc2.bias, std=1e-6)
          elif isinstance(module, SiglipMultiheadAttentionPoolingHead):
              nn.init.xavier_uniform_(module.probe.data)
              nn.init.xavier_uniform_(module.attention.in_proj_weight.data)
              nn.init.zeros_(module.attention.in_proj_bias.data)
          elif isinstance(module, SiglipModel):
              logit_scale_init = torch.log(torch.tensor(1.0))
              module.logit_scale.data.fill_(logit_scale_init)
              module.logit_bias.data.zero_()
          elif isinstance(module, SiglipForImageClassification):
              nn.init.normal_(
                  module.classifier.weight,
                  std=self.config.vision_config.hidden_size**-0.5 * self.config.initializer_factor,
              )
          elif isinstance(module, (nn.Linear, nn.Conv2d)):
              lecun_normal_(module.weight)
              if module.bias is not None:
                  nn.init.zeros_(module.bias)
          elif isinstance(module, nn.LayerNorm):
              module.bias.data.zero_()
              module.weight.data.fill_(1.0)
  
  
  SIGLIP_START_DOCSTRING = r"""
      This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
      library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
      etc.)
  
      This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
      Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
      and behavior.
  
      Parameters:
          config ([`SiglipConfig`]): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
  