Update modeling_minicpmv.py

modeling_minicpmv.py

@@ -42,13 +42,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         return model
 
-    def init_resampler(self, embed_dim, vision_dim):
+    def init_resampler(self, embed_dim, vision_dim,):
         return Resampler(
             num_queries=self.config.query_num,
             embed_dim=embed_dim,
             num_heads=embed_dim // 128,
             kv_dim=vision_dim,
-            adaptive=True
+            adaptive=True,
         )
 
     def init_transform(self):
@@ -60,13 +60,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 ),
             ]
         )
-
+    
     def get_input_embeddings(self):
         return self.llm.get_input_embeddings()
 
     def set_input_embeddings(self, value):
         self.llm.embed_tokens = value
-
+        
     def get_vllm_embedding(self, data):
         if 'vision_hidden_states' not in data:
             dtype = self.vpm.embeddings.position_embedding.weight.dtype
@@ -152,16 +152,105 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                     image_indices = torch.stack(
                         [torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound]
                     ).to(vllm_embedding.device)
-
                     cur_vllm_emb.scatter_(0, image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]),
                                           cur_vs_hs.view(-1, cur_vs_hs.shape[-1]))
                 elif self.training:
                     cur_vllm_emb += cur_vs_hs[0].mean() * 0
 
         return vllm_embedding, vision_hidden_states
-
+        
     def forward(self, data, **kwargs):
-        vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data)
+
+        if 'vision_hidden_states' not in data:
+            dtype = self.llm.lm_head.weight.dtype
+            device = self.llm.lm_head.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:
+                tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
+
+                if self.config.batch_vision_input:
+                    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, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+                    vision_embedding = self.vpm(all_pixel_values.type(dtype), patch_attention_mask=patch_attn_mask).last_hidden_state
+                    vision_embedding = self.resampler(vision_embedding, tgt_sizes)
+                else:
+                    # get vision_embedding foreach
+                    vision_embedding = []
+                    for single_tgt_size, single_pixel_values in zip(tgt_sizes, all_pixel_values):
+                        single_pixel_values = single_pixel_values.unsqueeze(0)
+                        B, L, _ = single_pixel_values.shape
+                        single_pixel_values = single_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+                        single_vision_embedding = self.vpm(single_pixel_values.type(dtype)).last_hidden_state
+                        single_vision_embedding = self.resampler(single_vision_embedding, single_tgt_size.unsqueeze(0))
+                        vision_embedding.append(single_vision_embedding)
+                    vision_embedding = torch.vstack(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
+                if self.training:
+                    dummy_image = torch.zeros(
+                        (1, 3, 224, 224),
+                        device=device, dtype=dtype
+                    )
+                    tgt_sizes = torch.Tensor([[(224 // self.config.patch_size), math.ceil(224 / self.config.patch_size)]]).type(torch.int32)
+                    dummy_feature = self.resampler(self.vpm(dummy_image).last_hidden_state, tgt_sizes)
+                else:
+                    dummy_feature = []
+                for _ in range(len(pixel_values_list)):
+                    vision_hidden_states.append(dummy_feature)
+
+        else:
+            vision_hidden_states = data['vision_hidden_states']
+
+        if hasattr(self.llm.config, 'scale_emb'):
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
+        else:
+            vllm_embedding = self.llm.model.embed_tokens(data['input_ids'])
+
+        vision_hidden_states = [i.type(vllm_embedding.dtype) if isinstance(
+            i, torch.Tensor) else i for i in vision_hidden_states]
+
+        bs = len(data['input_ids'])
+        for i in range(bs):
+            cur_vs_hs = vision_hidden_states[i]
+            if len(cur_vs_hs) > 0:
+                cur_vllm_emb = vllm_embedding[i]
+                cur_image_bound = data['image_bound'][i]
+                if len(cur_image_bound) > 0:
+                    image_indices = torch.stack(
+                        [torch.arange(r[0], r[1], dtype=torch.long) for r in cur_image_bound]
+                    ).to(vllm_embedding.device)
+                    cur_vllm_emb.scatter_(0, image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]),
+                                          cur_vs_hs.view(-1, cur_vs_hs.shape[-1]))
+                elif self.training:
+                    cur_vllm_emb += cur_vs_hs[0].mean() * 0
+
+        # vllm_embedding, vision_hidden_states = self.get_vllm_embedding(data)
         position_ids = data["position_ids"]
         if position_ids.dtype != torch.int64:
             position_ids = position_ids.long()