Modify eval_mm for MiniCPM-o 2.6

eval_mm/vqaeval/models/MiniCPM/minicpmv.py

@@ -33,14 +33,9 @@ class MiniCPM_V:
 
     def generate(self, images, questions, datasetname):
         image = Image.open(images[0]).convert('RGB')
-        try:
-            max_new_tokens = max_token[datasetname]
-        except:
-            max_new_tokens = 1024
-        if (datasetname == 'docVQA') or (datasetname == "docVQATest") :
-            prompt = "Answer the question directly with single word." + "\n" + questions[0]
-        elif (datasetname == 'textVQA') :
-            prompt = "Answer the question directly with single word." + '\n'+ questions[0]
+        max_new_tokens = max_token[datasetname]
+
+        prompt = "Answer the question directly with single word." + '\n' + questions[0]
         
         msgs = [{'role': 'user', 'content': prompt}]
         default_kwargs = dict(
@@ -59,10 +54,7 @@ class MiniCPM_V:
         return [res]
     
     def generate_with_interleaved(self, images, questions, datasetname):
-        try:
-            max_new_tokens = max_token[datasetname]
-        except:
-            max_new_tokens = 1024
+        max_new_tokens = max_token[datasetname]
         
         prompt = "Answer the question directly with single word."
         
@@ -103,11 +95,10 @@ class MiniCPM_V:
 class MiniCPM_V_2_6:
 
     def __init__(self, model_path, ckpt, device=None)->None:
-        seed = 0
-        random.seed(seed)
-        np.random.seed(seed)
-        torch.manual_seed(seed)
-        torch.cuda.manual_seed_all(seed)
+        random.seed(0)
+        np.random.seed(0)
+        torch.manual_seed(0)
+        torch.cuda.manual_seed_all(0)
         
         self.model_path = model_path
         self.ckpt = ckpt
@@ -125,14 +116,17 @@ class MiniCPM_V_2_6:
 
     def generate(self, images, questions, datasetname):
         image = Image.open(images[0]).convert('RGB')
-        try:
-            max_new_tokens = max_token[datasetname]
-        except:
-            max_new_tokens = 1024
-        if (datasetname == 'docVQA') or (datasetname == "docVQATest") :
-            prompt = "Answer the question directly with single word." + "\n" + questions[0]
-        elif (datasetname == 'textVQA') :
-            prompt = "Answer the question directly with single word." + '\n'+ questions[0]
+        img_width, img_height = image.width, image.height
+        if (img_width * img_height) < (1344 * 1344):
+            ratio = math.sqrt((1344 * 1344) / (img_width * img_height))
+            max_img_width = int(img_width * ratio)
+            new_img_width = random.randint(img_width, max_img_width)
+            new_img_height = int(new_img_width / img_width * img_height)
+            image = image.resize((new_img_width, new_img_height))
+        
+        max_new_tokens = max_token[datasetname]
+            
+        prompt = "Answer the question directly with single word." + '\n' + questions[0]
         
         msgs = [{'role': 'user', 'content': prompt}]
         default_kwargs = dict(
@@ -151,10 +145,7 @@ class MiniCPM_V_2_6:
         return [res]
     
     def generate_with_interleaved(self, images, questions, datasetname):
-        try:
-            max_new_tokens = max_token[datasetname]
-        except:
-            max_new_tokens = 1024
+        max_new_tokens = max_token[datasetname]
         
         prompt = "Answer the question directly with single word."
         
@@ -197,5 +188,117 @@ class MiniCPM_V_2_6:
 
         if isinstance(res, tuple) and len(res) > 0:
             res = res[0]
-        print(f"Q: {content}, \nA: {res}")
+
         return [res]
+
+
+class MiniCPM_o_2_6:
+
+    def __init__(self, model_path, ckpt, device=None)->None:
+        random.seed(0)
+        np.random.seed(0)
+        torch.manual_seed(0)
+        torch.cuda.manual_seed_all(0)
+        
+        self.model_path = model_path
+        self.ckpt = ckpt
+        self.model = AutoModel.from_pretrained(
+            self.model_path,
+            trust_remote_code=True,
+            attn_implementation='sdpa',
+            torch_dtype=torch.bfloat16,
+            init_vision=True,
+            init_audio=False,
+            init_tts=False
+        )
+        if self.ckpt is not None:
+            self.ckpt = ckpt
+            self.state_dict = torch.load(self.ckpt, map_location=torch.device('cpu'))
+            self.model.load_state_dict(self.state_dict)
+
+        self.model = self.model.eval().to(device)
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+        torch.cuda.empty_cache()
+
+    def generate(self, images, questions, datasetname):
+        image = Image.open(images[0]).convert('RGB')
+        img_width, img_height = image.width, image.height
+        if (img_width * img_height) < (1344 * 1344):
+            ratio = math.sqrt((1344 * 1344) / (img_width * img_height))
+            max_img_width = int(img_width * ratio)
+            new_img_width = random.randint(img_width, max_img_width)
+            new_img_height = int(new_img_width / img_width * img_height)
+            image = image.resize((new_img_width, new_img_height))     
+
+        max_new_tokens = max_token[datasetname]
+
+        prompt = "Answer the question directly with single word." + '\n' + questions[0]
+        
+        msgs = [{'role': 'user', 'content': prompt}]
+        default_kwargs = dict(
+            max_new_tokens=max_new_tokens,
+            sampling=False,
+            num_beams=3,
+            max_inp_length=8192,
+            use_image_id=True,
+            max_slice_nums=None
+        )
+        res = self.model.chat(
+            image=image,
+            msgs=msgs,
+            context=None,
+            tokenizer=self.tokenizer,
+            **default_kwargs
+        )
+        
+        return [res]
+    
+    def generate_with_interleaved(self, images, questions, datasetname):
+        max_new_tokens = max_token[datasetname]
+        
+        prompt = "Answer the question directly with single word."
+        
+        default_kwargs = dict(
+            max_new_tokens=max_new_tokens,
+            sampling=False,
+            num_beams=3,
+            max_inp_length=8192,
+            use_image_id=True,
+            max_slice_nums=None
+        )
+        
+        content = []
+        message = [
+            {'type': 'text', 'value': prompt},
+            {'type': 'image', 'value': images[0]},
+            {'type': 'text', 'value': questions[0]}
+        ]
+        for x in message:
+            if x['type'] == 'text':
+                content.append(x['value'])
+            elif x['type'] == 'image':
+                image = Image.open(x['value']).convert('RGB')
+                img_width, img_height = image.width, image.height
+                if (img_width * img_height) >= (1344 * 1344):
+                    content.append(image)
+                else:
+                    ratio = math.sqrt((1344 * 1344) / (img_width * img_height))
+                    max_img_width = int(img_width * ratio)
+                    new_img_width = random.randint(img_width, max_img_width)
+                    new_img_height = int(new_img_width / img_width * img_height)
+                    resized_image = image.resize((new_img_width, new_img_height))
+                    content.append(resized_image)
+        msgs = [{'role': 'user', 'content': content}]
+
+        res = self.model.chat(
+            image=None,
+            msgs=msgs,
+            context=None,
+            tokenizer=self.tokenizer,
+            **default_kwargs
+        )
+
+        if isinstance(res, tuple) and len(res) > 0:
+            res = res[0]
+        print(f"Q: {content}, \nA: {res}")
+        return [res]