working version

modeling_siglip.py

@@ -292,32 +292,25 @@ class SiglipVisionEmbeddings(nn.Module):
         batch_size = pixel_values.size(0)
 
         patch_embeds = self.patch_embedding(pixel_values)
-        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        embeddings = 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)
+        boundaries = torch.arange(1/self.num_patches_per_side, 1., 1/self.num_patches_per_side)
+        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()
+            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)
+            fractional_coords_h = torch.arange(0, 1-1e-6, 1/nb_patches_h)
+            fractional_coords_w = torch.arange(0, 1-1e-6, 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[batch_idx][p_attn_mask.view(-1).cpu()] = pos_ids
 
         position_ids = position_ids.to(self.position_embedding.weight.device)
 
@@ -1099,11 +1092,11 @@ class SiglipVisionTransformer(nn.Module):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        batch_size = pixel_values.size(0)
         if pixel_attention_mask is None:
-            #TODO
-            pass
+            # assuming `pixel_attention_mask` is of size bs x h x w
+            pixel_attention_mask = torch.ones(size=(batch_size, pixel_values.size(2), pixel_values.size(3)), dtype=torch.bool, device=pixel_values.device)
 
-        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()
 
@@ -1112,9 +1105,11 @@ class SiglipVisionTransformer(nn.Module):
             patch_attention_mask=patch_attention_mask
         )
 
+        patch_attention_mask = patch_attention_mask.view(batch_size, -1)
+
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
-            attention_mask=patch_attention_mask.view(batch_size, -1),
+            attention_mask=_prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype) if not self.config._flash_attn_2_enabled else patch_attention_mask,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
@@ -1125,7 +1120,7 @@ class SiglipVisionTransformer(nn.Module):
 
         pooled_output = self.head(
             hidden_state=last_hidden_state,
-            attention_mask=patch_attention_mask.view(batch_size, -1)
+            attention_mask=patch_attention_mask,
         )
 
         if not return_dict: