fix

Signed-off-by: guochenxu <[email protected]>

Author: Guo-Chenxu

src/transformers/models/minicpm_o_2_6/modeling_minicpm_o_2_6.py

@@ -723,42 +723,64 @@ def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
 
         return input_lengths_after_cnn, input_lengths_after_pooling
 
-    def get_image_features(self, tgt_sizes, all_pixel_values, dtype, device):
-        tgt_sizes = [tgt_size for tgt_size in tgt_sizes if isinstance(tgt_size, torch.Tensor)]
-        tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
+    def get_image_features(self, pixel_values_list, tgt_sizes, dtype, device):
+        vision_hidden_states = []
+        all_pixel_values = []
+        img_cnt = []
+        for pixel_values in pixel_values_list:
+            img_cnt.append(len(pixel_values))
+            all_pixel_values.extend([i.flatten(end_dim=1).permute(1, 0) for i in pixel_values])
+
+        # exist image
+        if all_pixel_values:
+            tgt_sizes = [tgt_size for tgt_size in tgt_sizes if isinstance(tgt_size, torch.Tensor)]
+            tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
+
+            max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+
+            all_pixel_values = torch.nn.utils.rnn.pad_sequence(
+                all_pixel_values, batch_first=True, padding_value=0.0
+            )
+            B, L, _ = all_pixel_values.shape
+            all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+
+            patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
+            for i in range(B):
+                patch_attn_mask[i, 0, : tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+            vision_batch_size = self.config.vision_batch_size
+            all_pixel_values = all_pixel_values.type(dtype)
+            if B > vision_batch_size:
+                hs = []
+                for i in range(0, B, vision_batch_size):
+                    start_idx = i
+                    end_idx = i + vision_batch_size
+                    tmp_hs = self.vpm(
+                        all_pixel_values[start_idx:end_idx],
+                        patch_attention_mask=patch_attn_mask[start_idx:end_idx],
+                        tgt_sizes=tgt_sizes[start_idx:end_idx],
+                    ).last_hidden_state
+                    hs.append(tmp_hs)
+                vision_embedding = torch.cat(hs, dim=0)
+            else:
+                vision_embedding = self.vpm(
+                    all_pixel_values, patch_attention_mask=patch_attn_mask, tgt_sizes=tgt_sizes
+                ).last_hidden_state
 
-        max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+            vision_embedding = self.resampler(vision_embedding, tgt_sizes)
 
-        all_pixel_values = torch.nn.utils.rnn.pad_sequence(
-            all_pixel_values, batch_first=True, padding_value=0.0
-        )
-        B, L, _ = all_pixel_values.shape
-        all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
-
-        patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
-        for i in range(B):
-            patch_attn_mask[i, 0, : tgt_sizes[i][0] * tgt_sizes[i][1]] = True
-
-        vision_batch_size = self.config.vision_batch_size
-        all_pixel_values = all_pixel_values.type(dtype)
-        if B > vision_batch_size:
-            hs = []
-            for i in range(0, B, vision_batch_size):
-                start_idx = i
-                end_idx = i + vision_batch_size
-                tmp_hs = self.vpm(
-                    all_pixel_values[start_idx:end_idx],
-                    patch_attention_mask=patch_attn_mask[start_idx:end_idx],
-                    tgt_sizes=tgt_sizes[start_idx:end_idx],
-                ).last_hidden_state
-                hs.append(tmp_hs)
-            vision_embedding = torch.cat(hs, dim=0)
-        else:
-            vision_embedding = self.vpm(
-                all_pixel_values, patch_attention_mask=patch_attn_mask, tgt_sizes=tgt_sizes
-            ).last_hidden_state
-        vision_embedding = self.resampler(vision_embedding, tgt_sizes)
-        return vision_embedding
+            start = 0
+            for pixel_values in pixel_values_list:
+                img_cnt = len(pixel_values)
+                if img_cnt > 0:
+                    vision_hidden_states.append(vision_embedding[start : start + img_cnt])
+                    start += img_cnt
+                else:
+                    vision_hidden_states.append([])
+        else:  # no image
+            vision_hidden_states.extend([[]] * len(pixel_values_list))
+
+        return vision_hidden_states
 
     def get_vllm_embedding(self, data):
         """
@@ -773,34 +795,12 @@ def get_vllm_embedding(self, data):
         Returns:
             embedding with vision, vision_hidden_states
         """
+        dtype = self.language_model.embed_tokens.weight.dtype
+        device = self.language_model.embed_tokens.weight.device
         if "vision_hidden_states" not in data:
-            dtype = self.language_model.embed_tokens.weight.dtype
-            device = self.language_model.embed_tokens.weight.device
-            tgt_sizes = data["tgt_sizes"]
-            pixel_values_list = data["pixel_values"]
-            vision_hidden_states = []
-            all_pixel_values = []
-            img_cnt = []
-            for pixel_values in pixel_values_list:
-                img_cnt.append(len(pixel_values))
-                all_pixel_values.extend([i.flatten(end_dim=1).permute(1, 0) for i in pixel_values])
-
-            # exist image
-            if all_pixel_values:
- 
-                vision_embedding = self.get_image_features(tgt_sizes=tgt_sizes, all_pixel_values=all_pixel_values, dtype=dtype, device=device)
-
-                start = 0
-                for pixel_values in pixel_values_list:
-                    img_cnt = len(pixel_values)
-                    if img_cnt > 0:
-                        vision_hidden_states.append(vision_embedding[start : start + img_cnt])
-                        start += img_cnt
-                    else:
-                        vision_hidden_states.append([])
-            else:  # no image
-                vision_hidden_states.extend([[]] * len(pixel_values_list))
-
+            vision_hidden_states = self.get_image_features(
+                data["pixel_values"], data["tgt_sizes"], dtype, device
+            )
         else:
             vision_hidden_states = data["vision_hidden_states"]