hhandle batch video generation

examples/{video_chat.py → chat.py}

@@ -1,11 +1,14 @@
 
 import torch
 from transformers import AutoModel, AutoTokenizer
-from utils import load_video
+from utils import load_image, load_video
 
 if __name__ == "__main__":
 
-    dir, rev = 'morpheushoc/InternVL2_5-2B', 'main'
+    # dir, rev = 'morpheushoc/InternVL2_5-2B', 'main'
+    dir, rev = '../', 'main'
+
+    # path = 'OpenGVLab/InternVL2_5-2B'
     model = AutoModel.from_pretrained(dir,
                                       torch_dtype=torch.bfloat16,
                                       load_in_8bit=False,
@@ -16,12 +19,36 @@ if __name__ == "__main__":
     tokenizer = AutoTokenizer.from_pretrained(dir, trust_remote_code=True, use_fast=False)
     generation_config = dict(max_new_tokens=1024, do_sample=False)
 
-    fp, question = 'red-panda.mp4', 'Describe this video in great details'
+    paths = [
+        'image1.jpg',
+        'image1.jpg',
+        'image2.jpg',
+        'red-panda.mp4',
+    ]
+
+    questions = [
+        'describe this image',
+        'describe this image',
+        'describe this image',
+        'describe this video'
+    ]
+
+    for fp, question in zip(paths, questions):
+        if fp.endswith('mp4'):
+            pixel_values, num_patches_list = load_video(fp, num_segments=8, max_num=1)
+            prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches_list))])
+            
+        else:
+            pixel_values = load_image(fp, max_num=12).to(torch.bfloat16).cuda()
+            num_patches_list = [len(pixel_values)]
+            prefix = '<image>\n'
+    
+        question = prefix + question
+        pixel_values = pixel_values.to(torch.bfloat16).cuda()
+        response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                                    num_patches_list=num_patches_list, history=None, return_history=True)
+        print(f'User: {question}\nAssistant: {response}')
 
-    pixel_values, num_patches_list = load_video(fp, num_segments=8, max_num=1)
-    prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches_list))])
-    question = prefix + question
-    pixel_values = pixel_values.to(torch.bfloat16).cuda()
-    response, history = model.chat(tokenizer, pixel_values, question, generation_config,
-                                num_patches_list=num_patches_list, history=None, return_history=True)
-    print(f'User: {question}\nAssistant: {response}')
+        question = 'How many animals ?'
+        response, history = model.chat(tokenizer, pixel_values, question, generation_config, history=history, return_history=True)
+        print(f'User: {question}\nAssistant: {response}')       

examples/{image_chat_batch.py → chat_batch.py}

@@ -1,11 +1,14 @@
 
 import torch
 from transformers import AutoModel, AutoTokenizer
-from utils import load_image
+from utils import load_image, load_video
 
 if __name__ == "__main__":
 
-    dir, rev = 'morpheushoc/InternVL2_5-2B', 'main'
+    # dir, rev = 'morpheushoc/InternVL2_5-2B', 'main'
+    dir, rev = '../', 'main'
+
+    # path = 'OpenGVLab/InternVL2_5-2B'
     model = AutoModel.from_pretrained(dir,
                                       torch_dtype=torch.bfloat16,
                                       load_in_8bit=False,
@@ -15,24 +18,40 @@ if __name__ == "__main__":
                                       revision=rev).eval().cuda()
     tokenizer = AutoTokenizer.from_pretrained(dir, trust_remote_code=True, use_fast=False)
     generation_config = dict(max_new_tokens=1024, do_sample=False)
+
     paths = [
-        'image1.jpg', 
-        'image1.jpg'
-        ]
-    
+        'image1.jpg',
+        'image1.jpg',
+        'image2.jpg',
+        'red-panda.mp4',
+    ]
+
     questions = [
-        'Describe this image in great details',
-        'Describe this image in great details'
-        ]
+        'describe this image',
+        'describe this image',
+        'describe this image',
+        'describe this video'
+    ]
 
-    pixel_values, num_patches_list, l_questions = [], [], []
-    for path, q in zip(paths, questions):
-        pxl_val = load_image(path, max_num=12).to(torch.bfloat16).cuda()
-        pixel_values.append(pxl_val)
-        num_patches_list.append(len(pxl_val))
-        l_questions.append('<image>\n{}'.format(q)) 
-    pixel_values = torch.cat(pixel_values)
+    pixel_values, num_patches_list = [], []
+    for i, fp in enumerate(paths):
+
+        if fp.endswith('mp4'):
+            pxl_val, num_patches = load_video(fp, num_segments=8, max_num=1)
+            prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches))])
+        else:
+            pxl_val = load_image(fp, max_num=12).to(torch.bfloat16)
+            num_patches = [len(pxl_val)]
+            prefix = '<image>\n'
         
-    responses = model.batch_chat(tokenizer, pixel_values, num_patches_list=num_patches_list,questions=questions,generation_config=generation_config)
-    for question, response in zip(questions, responses):
-        print(f'User: {question}\nAssistant: {response}')
+        pixel_values.append(pxl_val)
+        num_patches_list.append(num_patches)
+        questions[i] = prefix + questions[i]
+    
+    pixel_values = torch.cat(pixel_values).to(torch.bfloat16).cuda()
+    response = model.batch_chat(tokenizer, pixel_values, questions, generation_config,
+                                num_patches_list=num_patches_list, history=None, return_history=False)
+    
+    for q, r in zip(questions, response):
+        print(f'User: {q}\nAssistant: {r}')
+        print('\n')

examples/image_chat.py

@@ -1,31 +0,0 @@
-
-import torch
-from transformers import AutoModel, AutoTokenizer, AutoConfig
-from utils import load_image
-
-if __name__ == "__main__":
-
-    dir, rev = 'morpheushoc/InternVL2_5-2B', 'main'
-    # path = 'OpenGVLab/InternVL2_5-2B'
-    model = AutoModel.from_pretrained(dir,
-                                      torch_dtype=torch.bfloat16,
-                                      load_in_8bit=False,
-                                      low_cpu_mem_usage=True,
-                                      use_flash_attn=True,
-                                      trust_remote_code=True,
-                                      revision=rev).eval().cuda()
-    tokenizer = AutoTokenizer.from_pretrained(dir, trust_remote_code=True, use_fast=False)
-    generation_config = dict(max_new_tokens=1024, do_sample=False)
-
-    fp, question = 'image1.jpg', 'Describe this image in great details'
-
-    pixel_values = load_image(fp, max_num=12).to(torch.bfloat16).cuda()
-    num_patches_list = [len(pixel_values)]
-    prefix = '<image>\n'
-    
-    question = prefix + question
-    pixel_values = pixel_values.to(torch.bfloat16).cuda()
-    response, history = model.chat(tokenizer, pixel_values, question, generation_config,
-                                num_patches_list=num_patches_list, history=None, return_history=True)
-    print(f'User: {question}\nAssistant: {response}')
-

modeling_internvl_chat.py

@@ -220,24 +220,15 @@ class InternVLChatModel(PreTrainedModel):
             print(f'dynamic ViT batch size: {image_bs}')
 
         queries = []
-        for idx, num_patches in enumerate(num_patches_list):
-            question = questions[idx]
-            if pixel_values is not None and '<image>' not in question:
-                question = '<image>\n' + question
-            template = get_conv_template(self.template)
-            template.system_message = self.system_message
-            template.append_message(template.roles[0], question)
-            template.append_message(template.roles[1], None)
-            query = template.get_prompt()
-
-            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
-            query = query.replace('<image>', image_tokens, 1)
+        for q, num_patches in zip(questions, num_patches_list):
+            query = self.build_query(q, [], num_patches, IMG_START_TOKEN, IMG_END_TOKEN, IMG_CONTEXT_TOKEN)
             queries.append(query)
 
         tokenizer.padding_side = 'left'
         model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
         input_ids = model_inputs['input_ids'].to(self.device)
         attention_mask = model_inputs['attention_mask'].to(self.device)
+        template = get_conv_template(self.template)
         eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
         generation_config['eos_token_id'] = eos_token_id
         generation_output = self.generate(
@@ -264,25 +255,16 @@ class InternVLChatModel(PreTrainedModel):
         img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
         self.img_context_token_id = img_context_token_id
 
-        template = get_conv_template(self.template)
-        template.system_message = self.system_message
-        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
-
         history = [] if history is None else history
-        for (old_question, old_answer) in history:
-            template.append_message(template.roles[0], old_question)
-            template.append_message(template.roles[1], old_answer)
-        template.append_message(template.roles[0], question)
-        template.append_message(template.roles[1], None)
-        query = template.get_prompt()
+        query = self.build_query(question, history, num_patches_list, IMG_START_TOKEN,
+                                 IMG_END_TOKEN, IMG_CONTEXT_TOKEN)
 
         if verbose and pixel_values is not None:
             image_bs = pixel_values.shape[0]
             print(f'dynamic ViT batch size: {image_bs}')
 
-        for num_patches in num_patches_list:
-            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
-            query = query.replace('<image>', image_tokens, 1)
+        template = get_conv_template(self.template)
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
 
         model_inputs = tokenizer(query, return_tensors='pt')
         input_ids = model_inputs['input_ids'].to(self.device)
@@ -306,6 +288,26 @@ class InternVLChatModel(PreTrainedModel):
                 print(query_to_print, response)
             return response
 
+    def build_query(self, question, history, num_patches_list=None, IMG_START_TOKEN='<img>',
+                    IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'):
+
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+
+        for (old_question, old_answer) in history:
+            template.append_message(template.roles[0], old_question)
+            template.append_message(template.roles[1], old_answer)
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+
+        return query
+
+
     @torch.no_grad()
     def generate(
             self,