Update modeling_minicpmv.py

modeling_minicpmv.py

@@ -1,4 +1,5 @@
-import math
+
+import math 
 from typing import List, Optional
 import json
 import torch
@@ -22,7 +23,7 @@ class MiniCPMVPreTrainedModel(LlamaPreTrainedModel):
 class MiniCPMV(MiniCPMVPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
-
+        
         self.llm = LlamaForCausalLM(config)
         self.vpm = self.init_vision_module()
         self.vision_dim = self.vpm.embed_dim
@@ -30,6 +31,13 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         self.resampler = self.init_resampler(self.embed_dim, self.vision_dim)
         self.transform = self.init_transform()
 
+        # Wrap the model with DataParallel
+        if torch.cuda.device_count() > 1:
+            self.llm = nn.DataParallel(self.llm)
+            self.vpm = nn.DataParallel(self.vpm)
+            self.resampler = nn.DataParallel(self.resampler)
+            
+
     def init_vision_module(self):
         # same as HuggingFaceM4/siglip-so400m-14-980-flash-attn2-navit
         model = Idefics2VisionTransformer(self.config.vision_config)
@@ -61,9 +69,12 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
         )
 
     def get_vllm_embedding(self, data):
+        # Similar processing as before but make sure to handle DataParallel wrapping if needed
+        # Make sure all tensors are moved to the correct device if needed
         if 'vision_hidden_states' not in data:
-            dtype = self.vpm.embeddings.position_embedding.weight.dtype
-            device = self.vpm.embeddings.position_embedding.weight.device
+            dtype = self.vpm.module.embeddings.position_embedding.weight.dtype if isinstance(self.vpm, nn.DataParallel) else self.vpm.embeddings.position_embedding.weight.dtype
+            device = self.vpm.module.embeddings.position_embedding.weight.device if isinstance(self.vpm, nn.DataParallel) else self.vpm.embeddings.position_embedding.weight.device
+
             tgt_sizes = data['tgt_sizes']
             pixel_values_list = data['pixel_values']
             vision_hidden_states = []
@@ -126,14 +137,89 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         else:
             vision_hidden_states = data['vision_hidden_states']
+    def chat(self, image, msgs, tokenizer, vision_hidden_states=None, max_new_tokens=1024, sampling=True, max_inp_length=2048, **kwargs):
+        if isinstance(msgs, str):
+            msgs = json.loads(msgs)
+
+        copy_msgs = deepcopy(msgs)
+        assert len(copy_msgs) > 0, 'msgs is empty'
+
+        if image is not None and isinstance(copy_msgs[0]['content'], str):
+            copy_msgs[0]['content'] = [image, copy_msgs[0]['content']]
+
+        images = []
+        tgt_sizes = []
+        for i, msg in enumerate(copy_msgs):
+            role = msg["role"]
+            content = msg["content"]
+            assert role in ["user", "assistant"]
+            if i == 0:
+                assert role == "user", "The role of first msg should be user"
+            if isinstance(content, str):
+                content = [content]
+
+            cur_msgs = []
+            for c in content:
+                if isinstance(c, Image.Image):
+                    image = c
+                    if self.config.slice_mode:
+                        slice_images, image_placeholder = self.get_slice_image_placeholder(image, tokenizer)
+                        cur_msgs.append(image_placeholder)
+                        for slice_image in slice_images:
+                            slice_image = self.transform(slice_image)
+                            H, W = slice_image.shape[1:]
+                            images.append(self.reshape_by_patch(slice_image))
+                            tgt_sizes.append(torch.Tensor([H // self.config.patch_size, W // self.config.patch_size]).type(torch.int32))
+                    else:
+                        images.append(self.transform(image))
+                        cur_msgs.append(tokenizer.im_start + tokenizer.unk_token * self.config.query_num + tokenizer.im_end)
+                elif isinstance(c, str):
+                    cur_msgs.append(c)
+
+            msg['content'] = '\n'.join(cur_msgs)
+        if tgt_sizes:
+            tgt_sizes = torch.vstack(tgt_sizes)
+
+        input_ids = tokenizer.apply_chat_template(copy_msgs, tokenize=True, add_generation_prompt=False)
+
+        if sampling:
+            generation_config = {
+                "top_p": 0.8,
+                "top_k": 100,
+                "temperature": 0.7,
+                "do_sample": True,
+                "repetition_penalty": 1.05
+            }
+        else:
+            generation_config = {
+                "num_beams": 3,
+                "repetition_penalty": 1.2,
+            }
+
+        generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys())
+
+        with torch.inference_mode():
+            res, vision_hidden_states = self.generate(
+                input_id_list=[input_ids],
+                max_inp_length=max_inp_length,
+                img_list=[images],
+                tgt_sizes=[tgt_sizes],
+                tokenizer=tokenizer,
+                max_new_tokens=max_new_tokens,
+                vision_hidden_states=vision_hidden_states,
+                return_vision_hidden_states=True,
+                **generation_config
+            )
+        answer = res[0]
+
+        return answer
 
         if hasattr(self.llm.config, 'scale_emb'):
-            vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
+            vllm_embedding = self.llm.module.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb if isinstance(self.llm, nn.DataParallel) else self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb
         else:
-            vllm_embedding = self.llm.model.embed_tokens(data['input_ids'])
+            vllm_embedding = self.llm.module.model.embed_tokens(data['input_ids']) if isinstance(self.llm, nn.DataParallel) else 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]
+        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):
@@ -142,29 +228,36 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 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]))
+                    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)
         position_ids = data["position_ids"]
         if position_ids.dtype != torch.int64:
             position_ids = position_ids.long()
 
-        return self.llm(
-            input_ids=None,
-            position_ids=position_ids,
-            inputs_embeds=vllm_embedding,
-            **kwargs
-        )
+        # Handle DataParallel model
+        if isinstance(self.llm, nn.DataParallel):
+            outputs = self.llm.module(
+                input_ids=None,
+                position_ids=position_ids,
+                inputs_embeds=vllm_embedding,
+                **kwargs
+            )
+        else:
+            outputs = self.llm(
+                input_ids=None,
+                position_ids=position_ids,
+                inputs_embeds=vllm_embedding,
+                **kwargs
+            )
+
+        return outputs
+
 
     def _convert_to_tensors(
         self, tokenizer, input_ids, max_inp_length: Optional[int] = None