formatting

modeling_siglip.py

@@ -284,7 +284,7 @@ class SiglipVisionEmbeddings(nn.Module):
         )
 
         self.num_patches_per_side = self.image_size // self.patch_size
-        self.num_patches = self.num_patches_per_side ** 2
+        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)
 
@@ -295,16 +295,22 @@ class SiglipVisionEmbeddings(nn.Module):
         embeddings = patch_embeds.flatten(2).transpose(1, 2)
 
         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
-        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)
+        max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size
+        boundaries = torch.arange(1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side)
+        position_ids = torch.full(
+            size=(
+                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-1e-6, 1/nb_patches_h)
-            fractional_coords_w = torch.arange(0, 1-1e-6, 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)
@@ -1095,27 +1101,26 @@ class SiglipVisionTransformer(nn.Module):
         batch_size = pixel_values.size(0)
         if patch_attention_mask is None:
             patch_attention_mask = torch.ones(
-                size=(batch_size, pixel_values.size(2)//self.config.patch_size, pixel_values.size(3)//self.config.patch_size),
+                size=(
+                    batch_size,
+                    pixel_values.size(2) // self.config.patch_size,
+                    pixel_values.size(3) // self.config.patch_size,
+                ),
                 dtype=torch.bool,
                 device=pixel_values.device,
             )
-        # if pixel_attention_mask is None:
-        #     # 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)
-
-        # 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
-        )
+        hidden_states = self.embeddings(pixel_values=pixel_values, 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=_prepare_4d_attention_mask(patch_attention_mask, hidden_states.dtype) if not self.config._flash_attn_2_enabled else patch_attention_mask,
+            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,
@@ -1156,10 +1161,7 @@ class SiglipMultiheadAttentionPoolingHead(nn.Module):
         probe = self.probe.repeat(batch_size, 1, 1)
 
         hidden_state = self.attention(
-            query=probe,
-            key=hidden_state,
-            value=hidden_state,
-            key_padding_mask=~attention_mask
+            query=probe, key=hidden_state, value=hidden_state, key_padding_mask=~attention_mask
         )[0]
 
         residual = hidden_state