save changes to modeling_siglip

modeling_siglip.py

@@ -39,7 +39,7 @@ from transformers.utils import (
     logging,
     replace_return_docstrings,
 )
-from .configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig
+from configuration_siglip import SiglipConfig, SiglipTextConfig, SiglipVisionConfig
 
 
 logger = logging.get_logger(__name__)
@@ -283,16 +283,45 @@ class SiglipVisionEmbeddings(nn.Module):
             padding="valid",
         )
 
-        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_patches_per_side = self.image_size // self.patch_size
+        self.num_patches = self.num_patches_per_side ** 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)
+    def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor:
+        batch_size = pixel_values.size(0)
 
-        embeddings = embeddings + self.position_embedding(self.position_ids)
+        patch_embeds = self.patch_embedding(pixel_values)
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+        patches_to_select = patch_attention_mask.view(batch_size, -1)
+        max_num_patches = patches_to_select.sum(dim=-1).max()
+        embeddings = torch.zeros((batch_size, max_num_patches, patch_embeds.size(2)), device=patch_embeds.device, dtype=patch_embeds.dtype)
+        for b_idx, (p_embeds, p_to_select) in enumerate(zip(patch_embeds, patches_to_select)):
+            sub_p_embds = p_embeds[p_to_select]
+            embeddings[b_idx][:len(sub_p_embds)] = sub_p_embds
+
+        boundaries = torch.arange(1/self.num_patches_per_side, 1., 1/self.num_patches_per_side)
+        max_im_h, max_im_w = pixel_values.size(2), pixel_values.size(3)
+        max_nb_patches_h, max_nb_patches_w = max_im_h//self.patch_size, max_im_w//self.patch_size
+        position_ids = torch.full((batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0)
+
+        for batch_idx, p_attn_mask in enumerate(patch_attention_mask):
+            nb_patches_h = p_attn_mask[0].sum()
+            nb_patches_w = p_attn_mask[:, 0].sum()
+
+            fractional_coords_h = torch.arange(0, 1, 1/nb_patches_h)
+            fractional_coords_w = torch.arange(0, 1, 1/nb_patches_w)
+
+            bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True)
+            bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True)
+
+            pos_ids = (self.num_patches_per_side * bucket_coords_w[:, None] + bucket_coords_h[None, :]).flatten()
+            position_ids[batch_idx][:len(pos_ids)] = pos_ids
+
+        position_ids = position_ids.to(self.position_embedding.weight.device)
+
+        embeddings = embeddings + self.position_embedding(position_ids)
         return embeddings
 
 
@@ -675,7 +704,7 @@ class SiglipPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         """Initialize the weights"""
-        
+
         if isinstance(module, SiglipVisionEmbeddings):
             width = (
                 self.config.vision_config.hidden_size
@@ -1055,6 +1084,7 @@ class SiglipVisionTransformer(nn.Module):
     def forward(
         self,
         pixel_values,
+        pixel_attention_mask: Optional[torch.BoolTensor] = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
@@ -1069,10 +1099,22 @@ class SiglipVisionTransformer(nn.Module):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-        hidden_states = self.embeddings(pixel_values)
+        if pixel_attention_mask is None:
+            #TODO
+            pass
+
+        batch_size = pixel_attention_mask.size(0) # assuming `pixel_attention_mask` is of size bs x h x w
+        subgrids = pixel_attention_mask.unfold(dimension=1, size=self.config.patch_size, step=self.config.patch_size).unfold(dimension=2, size=self.config.patch_size, step=self.config.patch_size)
+        patch_attention_mask = (subgrids.sum(dim=(-1, -2)) > 0).bool()
+
+        hidden_states = self.embeddings(
+            pixel_values=pixel_values,
+            patch_attention_mask=patch_attention_mask
+        )
 
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
+            attention_mask=patch_attention_mask.view(batch_size, -1),
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
@@ -1081,7 +1123,10 @@ class SiglipVisionTransformer(nn.Module):
         last_hidden_state = encoder_outputs[0]
         last_hidden_state = self.post_layernorm(last_hidden_state)
 
-        pooled_output = self.head(last_hidden_state)
+        pooled_output = self.head(
+            hidden_state=last_hidden_state,
+            attention_mask=patch_attention_mask.view(batch_size, -1)
+        )
 
         if not return_dict:
             return (last_hidden_state, pooled_output) + encoder_outputs[1:]
@@ -1105,11 +1150,16 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):
         self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.mlp = SiglipMLP(config)
 
-    def forward(self, hidden_state):
+    def forward(self, hidden_state, attention_mask):
         batch_size = hidden_state.shape[0]
         probe = self.probe.repeat(batch_size, 1, 1)
 
-        hidden_state = self.attention(probe, hidden_state, hidden_state)[0]
+        hidden_state = self.attention(
+            query=probe,
+            key=hidden_state,
+            value=hidden_state,
+            key_padding_mask=~attention_mask
+        )[0]
 
         residual = hidden_state
         hidden_state = self.layernorm(hidden_state)