clean vision

vision.py

@@ -24,27 +24,16 @@ 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,
     is_flash_attn_2_available,
-    logging,
-    replace_return_docstrings,
-)
+    logging,)
 
 from .configuration_img2html import Img2HTMLVisionConfig
 
 
 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
-# ]
 
 if is_flash_attn_2_available():
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -64,34 +53,6 @@ def _get_unpad_data(attention_mask):
     )
 
 
-# # 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):
@@ -122,75 +83,6 @@ class SiglipVisionModelOutput(ModelOutput):
     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: Img2HTMLVisionConfig):
         super().__init__()
@@ -220,40 +112,6 @@ class SiglipVisionEmbeddings(nn.Module):
         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"""
@@ -618,150 +476,6 @@ class SiglipEncoderLayer(nn.Module):
         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, 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):
     """
@@ -787,35 +501,6 @@ class SiglipEncoder(nn.Module):
         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
@@ -866,140 +551,6 @@ class SiglipEncoder(nn.Module):
         )
 
 
-# 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: Img2HTMLVisionConfig):
         super().__init__()
@@ -1011,8 +562,6 @@ class SiglipVisionTransformer(nn.Module):
         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=Img2HTMLVisionConfig)
     def forward(
         self,
         pixel_values,
@@ -1079,24 +628,13 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):
         return hidden_state[:, 0]
 
 
-# @add_start_docstrings(
-#     """The vision model from SigLIP without any head or projection on top.""",
-#     SIGLIP_START_DOCSTRING,
-# )
 class SiglipVisionModel(nn.Module):
     def __init__(self, config: Img2HTMLVisionConfig):
         super().__init__()
 
+        self.config = 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=Img2HTMLVisionConfig)
     def forward(
         self,
         pixel_values,
@@ -1104,28 +642,6 @@ class SiglipVisionModel(nn.Module):
         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(
@@ -1134,228 +650,3 @@ class SiglipVisionModel(nn.Module):
             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,
-#         )