siglip-so400m-14-384 / modeling_siglip.py

Leyo HF staff

take off prints

e4bb96b 8 months ago

No virus

50.2 kB

	# coding=utf-8
	# Copyright 2023 Google AI and The HuggingFace 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.
	""" PyTorch Siglip model."""


	from dataclasses import dataclass
	from typing import Any, Optional, Tuple, Union

	import torch
	import torch.utils.checkpoint
	from torch import nn

	from transformers.activations import ACT2FN
	from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
	from transformers.modeling_utils import PreTrainedModel
	from transformers.utils import (
	ModelOutput,
	add_start_docstrings,
	add_start_docstrings_to_model_forward,
	logging,
	replace_return_docstrings,
	)
	from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig


	logger = logging.get_logger(__name__)

	_CHECKPOINT_FOR_DOC = "google/siglip-base-patch16-224"

	SIGLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [
	"google/siglip-base-patch16-224",
	# See all SigLIP models at https://huggingface.co/models?filter=siglip
	]


	# Copied from transformers.models.bart.modeling_bart._expand_mask
	def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
	"""
	Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
	"""
	bsz, src_len = mask.size()
	tgt_len = tgt_len if tgt_len is not None else src_len

	expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)

	inverted_mask = 1.0 - expanded_mask

	return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)


	# contrastive loss function, adapted from
	# https://sachinruk.github.io/blog/2021-03-07-siglip.html
	def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
	return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))


	# Copied from transformers.models.clip.modeling_clip.clip_loss with clip->siglip
	def siglip_loss(similarity: torch.Tensor) -> torch.Tensor:
	caption_loss = contrastive_loss(similarity)
	image_loss = contrastive_loss(similarity.t())
	return (caption_loss + image_loss) / 2.0


	@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 forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:

	patch_embeds = self.patch_embedding(pixel_values) # shape = [*, width, grid, grid]
	embeddings = patch_embeds.flatten(2).transpose(1, 2)

	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


	# Copied from transformers.models.clip.modeling_clip.CLIPAttention with CLIP->Siglip
	class SiglipAttention(nn.Module):
	"""Multi-headed attention from 'Attention Is All You Need' paper"""

	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 _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
	return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: Optional[torch.Tensor] = None,
	causal_attention_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = False,
	) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
	"""Input shape: Batch x Time x Channel"""

	bsz, tgt_len, embed_dim = hidden_states.size()

	# get query proj
	query_states = self.q_proj(hidden_states) * self.scale
	key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
	value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

	proj_shape = (bsz * self.num_heads, -1, self.head_dim)
	query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
	key_states = key_states.view(*proj_shape)
	value_states = value_states.view(*proj_shape)

	src_len = key_states.size(1)
	attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

	if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
	raise ValueError(
	f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
	f" {attn_weights.size()}"
	)

	# apply the causal_attention_mask first
	if causal_attention_mask is not None:
	if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
	f" {causal_attention_mask.size()}"
	)
	attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	if attention_mask is not None:
	if attention_mask.size() != (bsz, 1, tgt_len, src_len):
	raise ValueError(
	f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
	)
	attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
	attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

	attn_weights = nn.functional.softmax(attn_weights, dim=-1)

	if output_attentions:
	# this operation is a bit akward, but it's required to
	# make sure that attn_weights keeps its gradient.
	# In order to do so, attn_weights have to reshaped
	# twice and have to be reused in the following
	attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
	attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
	else:
	attn_weights_reshaped = None

	attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

	attn_output = torch.bmm(attn_probs, value_states)

	if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
	raise ValueError(
	f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
	f" {attn_output.size()}"
	)

	attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
	attn_output = attn_output.transpose(1, 2)
	attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)

	attn_output = self.out_proj(attn_output)

	return attn_output, attn_weights_reshaped


	# 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


	# Copied from transformers.models.clip.modeling_clip.CLIPEncoderLayer with CLIP->Siglip
	class SiglipEncoderLayer(nn.Module):
	def __init__(self, config: SiglipConfig):
	super().__init__()
	self.embed_dim = config.hidden_size
	self.self_attn = SiglipAttention(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)

	def forward(
	self,
	hidden_states: torch.Tensor,
	attention_mask: torch.Tensor,
	causal_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 size
	`(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
	`(config.encoder_attention_heads,)`.
	output_attentions (`bool`, optional):
	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,
	causal_attention_mask=causal_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

	def _init_weights(self, module):
	"""Initialize the weights"""
	factor = self.config.initializer_factor
	if isinstance(module, SiglipTextEmbeddings):
	module.token_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
	module.position_embedding.weight.data.normal_(mean=0.0, std=factor * 0.02)
	elif isinstance(module, SiglipVisionEmbeddings):
	factor = self.config.initializer_factor
	nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
	nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
	elif isinstance(module, SiglipAttention):
	factor = self.config.initializer_factor
	in_proj_std = (module.embed_dim*-0.5) ((2 * module.config.num_hidden_layers) ** -0.5) * factor
	out_proj_std = (module.embed_dim*-0.5) factor
	nn.init.normal_(module.q_proj.weight, std=in_proj_std)
	nn.init.normal_(module.k_proj.weight, std=in_proj_std)
	nn.init.normal_(module.v_proj.weight, std=in_proj_std)
	nn.init.normal_(module.out_proj.weight, std=out_proj_std)
	elif isinstance(module, SiglipMLP):
	factor = self.config.initializer_factor
	in_proj_std = (
	(module.config.hidden_size*-0.5) ((2 * module.config.num_hidden_layers) ** -0.5) * factor
	)
	fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
	nn.init.normal_(module.fc1.weight, std=fc_std)
	nn.init.normal_(module.fc2.weight, std=in_proj_std)
	if isinstance(module, nn.LayerNorm):
	module.bias.data.zero_()
	module.weight.data.fill_(1.0)
	if isinstance(module, nn.Linear) and module.bias is not None:
	module.bias.data.zero_()

	def _set_gradient_checkpointing(self, module, value=False):
	if isinstance(module, SiglipEncoder):
	module.gradient_checkpointing = value


	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
	configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
	"""

	SIGLIP_TEXT_INPUTS_DOCSTRING = r"""
	Args:
	input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
	Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
	it.

	Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
	[`PreTrainedTokenizer.__call__`] for details.

	[What are input IDs?](../glossary#input-ids)
	attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, optional):
	Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

	- 1 for tokens that are not masked,
	- 0 for tokens that are masked.

	[What are attention masks?](../glossary#attention-mask)
	position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):
	Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
	config.max_position_embeddings - 1]`.

	[What are position IDs?](../glossary#position-ids)
	output_attentions (`bool`, optional):
	Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
	tensors for more detail.
	output_hidden_states (`bool`, optional):
	Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
	more detail.
	return_dict (`bool`, optional):
	Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
	"""

	SIGLIP_VISION_INPUTS_DOCSTRING = r"""
	Args:
	pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
	Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
	[`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
	output_attentions (`bool`, optional):
	Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
	tensors for more detail.
	output_hidden_states (`bool`, optional):
	Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
	more detail.
	return_dict (`bool`, optional):
	Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
	"""

	SIGLIP_INPUTS_DOCSTRING = r"""
	Args:
	input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
	Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
	it.

	Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
	[`PreTrainedTokenizer.__call__`] for details.

	[What are input IDs?](../glossary#input-ids)
	attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, optional):
	Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

	- 1 for tokens that are not masked,
	- 0 for tokens that are masked.

	[What are attention masks?](../glossary#attention-mask)
	position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, optional):
	Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
	config.max_position_embeddings - 1]`.

	[What are position IDs?](../glossary#position-ids)
	pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
	Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
	[`AutoImageProcessor`]. See [`CLIPImageProcessor.__call__`] for details.
	return_loss (`bool`, optional):
	Whether or not to return the contrastive loss.
	output_attentions (`bool`, optional):
	Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
	tensors for more detail.
	output_hidden_states (`bool`, optional):
	Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
	more detail.
	return_dict (`bool`, optional):
	Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
	"""


	# Copied from transformers.models.clip.modeling_clip.CLIPEncoder with CLIP->Siglip
	class SiglipEncoder(nn.Module):
	"""
	Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
	[`SiglipEncoderLayer`].

	Args:
	config: SiglipConfig
	"""

	def __init__(self, config: SiglipConfig):
	super().__init__()
	self.config = config
	self.layers = nn.ModuleList([SiglipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
	self.gradient_checkpointing = False

	def forward(
	self,
	inputs_embeds,
	attention_mask: Optional[torch.Tensor] = None,
	causal_attention_mask: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutput]:
	r"""
	Args:
	inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
	Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
	This is useful if you want more control over how to convert `input_ids` indices into associated vectors
	than the model's internal embedding lookup matrix.
	attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, optional):
	Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

	- 1 for tokens that are not masked,
	- 0 for tokens that are masked.

	[What are attention masks?](../glossary#attention-mask)
	causal_attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, optional):
	Causal mask for the text model. Mask values selected in `[0, 1]`:

	- 1 for tokens that are not masked,
	- 0 for tokens that are masked.

	[What are attention masks?](../glossary#attention-mask)
	output_attentions (`bool`, optional):
	Whether or not to return the attentions tensors of all attention layers. See `attentions` under
	returned tensors for more detail.
	output_hidden_states (`bool`, optional):
	Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
	for more detail.
	return_dict (`bool`, optional):
	Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
	"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	encoder_states = () if output_hidden_states else None
	all_attentions = () if output_attentions else None

	hidden_states = inputs_embeds
	for idx, encoder_layer in enumerate(self.layers):
	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)
	if self.gradient_checkpointing and self.training:

	def create_custom_forward(module):
	def custom_forward(*inputs):
	return module(*inputs, output_attentions)

	return custom_forward

	layer_outputs = torch.utils.checkpoint.checkpoint(
	create_custom_forward(encoder_layer),
	hidden_states,
	attention_mask,
	causal_attention_mask,
	)
	else:
	layer_outputs = encoder_layer(
	hidden_states,
	attention_mask,
	causal_attention_mask,
	output_attentions=output_attentions,
	)

	hidden_states = layer_outputs[0]

	if output_attentions:
	all_attentions = all_attentions + (layer_outputs[1],)

	if output_hidden_states:
	encoder_states = encoder_states + (hidden_states,)

	if not return_dict:
	return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
	return BaseModelOutput(
	last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
	)


	class SiglipTextTransformer(nn.Module):
	def __init__(self, config: SiglipTextConfig):
	super().__init__()
	self.config = config
	embed_dim = config.hidden_size
	self.embeddings = SiglipTextEmbeddings(config)
	self.encoder = SiglipEncoder(config)
	self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

	self.head = nn.Linear(embed_dim, embed_dim)

	@add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipTextConfig)
	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	r"""
	Returns:

	"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	if input_ids is None:
	raise ValueError("You have to specify input_ids")

	input_shape = input_ids.size()
	input_ids = input_ids.view(-1, input_shape[-1])

	hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)

	# note: SigLIP's text model does not use q causal mask, unlike the original CLIP model.
	# expand attention_mask
	if attention_mask is not None:
	# [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
	attention_mask = _expand_mask(attention_mask, hidden_states.dtype)

	encoder_outputs = self.encoder(
	inputs_embeds=hidden_states,
	attention_mask=None,
	causal_attention_mask=None,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	last_hidden_state = encoder_outputs[0]
	last_hidden_state = self.final_layer_norm(last_hidden_state)

	# Assuming "sticky" EOS tokenization, last token is always EOS.
	pooled_output = last_hidden_state[:, -1, :]
	pooled_output = self.head(pooled_output)

	if not return_dict:
	return (last_hidden_state, pooled_output) + encoder_outputs[1:]

	return BaseModelOutputWithPooling(
	last_hidden_state=last_hidden_state,
	pooler_output=pooled_output,
	hidden_states=encoder_outputs.hidden_states,
	attentions=encoder_outputs.attentions,
	)


	@add_start_docstrings(
	"""The text model from SigLIP without any head or projection on top.""",
	SIGLIP_START_DOCSTRING,
	)
	class SiglipTextModel(SiglipPreTrainedModel):
	config_class = SiglipTextConfig

	_no_split_modules = ["SiglipTextEmbeddings", "SiglipEncoderLayer"]

	def __init__(self, config: SiglipTextConfig):
	super().__init__(config)
	self.text_model = SiglipTextTransformer(config)
	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.text_model.embeddings.token_embedding

	def set_input_embeddings(self, value):
	self.text_model.embeddings.token_embedding = value

	@add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipTextConfig)
	def forward(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	r"""
	Returns:

	Examples:

	```python
	>>> from transformers import AutoTokenizer, SiglipTextModel

	>>> model = SiglipTextModel.from_pretrained("google/siglip-base-patch16-224")
	>>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")

	>>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")

	>>> outputs = model(**inputs)
	>>> last_hidden_state = outputs.last_hidden_state
	>>> pooled_output = outputs.pooler_output # pooled (EOS token) states
	```"""
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	return self.text_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)


	class SiglipVisionTransformer(nn.Module):
	def __init__(self, config: SiglipVisionConfig):
	super().__init__()
	self.config = config
	embed_dim = config.hidden_size

	self.embeddings = SiglipVisionEmbeddings(config)
	self.encoder = SiglipEncoder(config)
	self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
	self.head = SiglipMultiheadAttentionPoolingHead(config)

	@add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
	def forward(
	self,
	pixel_values,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	r"""
	Returns:

	"""
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	hidden_states = self.embeddings(pixel_values)

	encoder_outputs = self.encoder(
	inputs_embeds=hidden_states,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	last_hidden_state = encoder_outputs[0]
	last_hidden_state = self.post_layernorm(last_hidden_state)


	pooled_output = self.head(last_hidden_state)

	if not return_dict:
	return (last_hidden_state, pooled_output) + encoder_outputs[1:]

	return BaseModelOutputWithPooling(
	last_hidden_state=last_hidden_state,
	pooler_output=pooled_output,
	hidden_states=encoder_outputs.hidden_states,
	attentions=encoder_outputs.attentions,
	)


	class SiglipMultiheadAttentionPoolingHead(nn.Module):
	"""Multihead Attention Pooling."""

	def __init__(self, config: SiglipVisionConfig):
	super().__init__()

	self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
	self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
	self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
	self.mlp = SiglipMLP(config)

	def forward(self, hidden_state):
	batch_size = hidden_state.shape[0]
	probe = self.probe.repeat(batch_size, 1, 1)

	hidden_state = self.attention(probe, hidden_state, hidden_state)[0]

	residual = hidden_state
	hidden_state = self.layernorm(hidden_state)
	hidden_state = residual + self.mlp(hidden_state)

	return hidden_state[:, 0]


	@add_start_docstrings(
	"""The vision model from SigLIP without any head or projection on top.""",
	SIGLIP_START_DOCSTRING,
	)
	class SiglipVisionModel(SiglipPreTrainedModel):
	config_class = SiglipVisionConfig
	main_input_name = "pixel_values"

	def __init__(self, config: SiglipVisionConfig):
	super().__init__(config)

	self.vision_model = SiglipVisionTransformer(config)

	# Initialize weights and apply final processing
	self.post_init()

	def get_input_embeddings(self) -> nn.Module:
	return self.vision_model.embeddings.patch_embedding

	@add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=SiglipVisionConfig)
	def forward(
	self,
	pixel_values,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, BaseModelOutputWithPooling]:
	r"""
	Returns:

	Examples:

	```python
	>>> from PIL import Image
	>>> import requests
	>>> from transformers import AutoProcessor, SiglipVisionModel

	>>> model = SiglipVisionModel.from_pretrained("google/siglip-base-patch16-224")
	>>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")

	>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
	>>> image = Image.open(requests.get(url, stream=True).raw)

	>>> inputs = processor(images=image, return_tensors="pt")

	>>> outputs = model(**inputs)
	>>> last_hidden_state = outputs.last_hidden_state
	>>> pooled_output = outputs.pooler_output # pooled CLS states
	```"""
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	return self.vision_model(
	pixel_values=pixel_values,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)


	@add_start_docstrings(SIGLIP_START_DOCSTRING)
	class SiglipModel(SiglipPreTrainedModel):
	config_class = SiglipConfig

	def __init__(self, config: SiglipConfig):
	super().__init__(config)

	if not isinstance(config.text_config, SiglipTextConfig):
	raise ValueError(
	"config.text_config is expected to be of type SiglipTextConfig but is of type"
	f" {type(config.text_config)}."
	)

	if not isinstance(config.vision_config, SiglipVisionConfig):
	raise ValueError(
	"config.vision_config is expected to be of type SiglipVisionConfig but is of type"
	f" {type(config.vision_config)}."
	)

	text_config = config.text_config
	vision_config = config.vision_config

	self.text_model = SiglipTextModel(text_config)
	self.vision_model = SiglipVisionModel(vision_config)

	self.temperature = nn.Parameter(
	torch.randn(
	1,
	)
	)
	self.bias = nn.Parameter(
	torch.randn(
	1,
	)
	)

	# Initialize weights and apply final processing
	self.post_init()

	@add_start_docstrings_to_model_forward(SIGLIP_TEXT_INPUTS_DOCSTRING)
	def get_text_features(
	self,
	input_ids: Optional[torch.Tensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.Tensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> torch.FloatTensor:
	r"""
	Returns:
	text_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The text embeddings obtained by
	applying the projection layer to the pooled output of [`SiglipTextModel`].

	Examples:

	```python
	>>> from transformers import AutoTokenizer, SiglipModel

	>>> model = SiglipModel.from_pretrained("google/siglip-base-patch16-224")
	>>> tokenizer = AutoTokenizer.from_pretrained("google/siglip-base-patch16-224")

	>>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
	>>> text_features = model.get_text_features(**inputs)
	```"""
	# Use SigLIP model's config for some fields (if specified) instead of those of vision & text components.
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	text_outputs = self.text_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	pooled_output = text_outputs[1]

	return pooled_output

	@add_start_docstrings_to_model_forward(SIGLIP_VISION_INPUTS_DOCSTRING)
	def get_image_features(
	self,
	pixel_values: Optional[torch.FloatTensor] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> torch.FloatTensor:
	r"""
	Returns:
	image_features (`torch.FloatTensor` of shape `(batch_size, output_dim`): The image embeddings obtained by
	applying the projection layer to the pooled output of [`SiglipVisionModel`].

	Examples:

	```python
	>>> from PIL import Image
	>>> import requests
	>>> from transformers import AutoProcessor, SiglipModel

	>>> model = SiglipModel.from_pretrained("google/siglip-base-patch16-224")
	>>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")

	>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
	>>> image = Image.open(requests.get(url, stream=True).raw)

	>>> inputs = processor(images=image, return_tensors="pt")

	>>> image_features = model.get_image_features(**inputs)
	```"""
	# Use SiglipModel's config for some fields (if specified) instead of those of vision & text components.
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	vision_outputs = self.vision_model(
	pixel_values=pixel_values,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	pooled_output = vision_outputs[1]

	return pooled_output

	@add_start_docstrings_to_model_forward(SIGLIP_INPUTS_DOCSTRING)
	@replace_return_docstrings(output_type=SiglipOutput, config_class=SiglipConfig)
	def forward(
	self,
	input_ids: Optional[torch.LongTensor] = None,
	pixel_values: Optional[torch.FloatTensor] = None,
	attention_mask: Optional[torch.Tensor] = None,
	position_ids: Optional[torch.LongTensor] = None,
	return_loss: Optional[bool] = None,
	output_attentions: Optional[bool] = None,
	output_hidden_states: Optional[bool] = None,
	return_dict: Optional[bool] = None,
	) -> Union[Tuple, SiglipOutput]:
	r"""
	Returns:

	Examples:

	```python
	>>> from PIL import Image
	>>> import requests
	>>> from transformers import AutoProcessor, SiglipModel

	>>> model = SiglipModel.from_pretrained("google/siglip-base-patch16-224")
	>>> processor = AutoProcessor.from_pretrained("google/siglip-base-patch16-224")

	>>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
	>>> image = Image.open(requests.get(url, stream=True).raw)

	>>> inputs = processor(
	... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
	... )

	>>> outputs = model(**inputs)
	>>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
	>>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
	```"""
	# Use SigLIP model's config for some fields (if specified) instead of those of vision & text components.
	output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
	output_hidden_states = (
	output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
	)
	return_dict = return_dict if return_dict is not None else self.config.use_return_dict

	vision_outputs = self.vision_model(
	pixel_values=pixel_values,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	text_outputs = self.text_model(
	input_ids=input_ids,
	attention_mask=attention_mask,
	position_ids=position_ids,
	output_attentions=output_attentions,
	output_hidden_states=output_hidden_states,
	return_dict=return_dict,
	)

	image_embeds = vision_outputs[1]
	text_embeds = text_outputs[1]

	# normalized features
	image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
	text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

	# cosine similarity as logits
	logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * self.temperature.exp() + self.bias
	logits_per_image = logits_per_text.t()

	z = torch.matmul(image_embeds, text_embeds.t()) * self.temperature.exp()

	loss = None
	if return_loss:
	raise NotImplementedError("SigLIP loss to be implemented")

	if not return_dict:
	output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
	return ((loss,) + output) if loss is not None else output

	return SiglipOutput(
	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,
	)