multi-images

modeling_minicpmv.py

@@ -3,6 +3,7 @@ import json
 import torch
 from threading import Thread
 from copy import deepcopy
+from PIL import Image
 from torchvision import transforms
 from transformers import LlamaPreTrainedModel, LlamaForCausalLM, TextIteratorStreamer
 from transformers.models.idefics2.modeling_idefics2 import Idefics2VisionTransformer
@@ -291,17 +292,37 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             msgs = json.loads(msgs)
         copy_msgs = deepcopy(msgs)
 
-        assert len(msgs) > 0, 'msgs is empty'
-        assert sampling or not stream, 'if use stream mode, make sure sampling=True'
+        assert len(msgs) > 0, "msgs is empty"
+        assert sampling or not stream, "if use stream mode, make sure sampling=True"
+
+        if image is not None and isinstance(copy_msgs[0]["content"], str):
+            # copy_msgs[0]['content'] = '(<image>./</image>)\n' + copy_msgs[0]['content']
+            copy_msgs[0]["content"] = [image, copy_msgs[0]["content"]]
+
+        images = []
+        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):
+                    images.append(c)
+                    cur_msgs.append("(<image>./</image>)")
+                elif isinstance(c, str):
+                    cur_msgs.append(c)
+            msg["content"] = "\n".join(cur_msgs)
 
-        if image is not None and isinstance(msgs[0]['content'], str):
-            copy_msgs[0]['content'] = '(<image>./</image>)\n' + copy_msgs[0]['content']
         if system_prompt:
             sys_msg = {'role': 'system', 'content': system_prompt}
             copy_msgs = [sys_msg] + copy_msgs        
 
         prompt = processor.tokenizer.apply_chat_template(copy_msgs, tokenize=False, add_generation_prompt=True)
-        inputs = processor(prompt, [image], return_tensors="pt", max_length=max_inp_length).to(self.device)
+        inputs = processor(prompt, images, return_tensors="pt", max_length=max_inp_length).to(self.device)
 
         if sampling:
             generation_config = {