Update modeling_internvl_chat.py

modeling_internvl_chat.py

@@ -16,6 +16,8 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import ModelOutput, logging
 
+from einops import rearrange
+
 from .configuration_internvl_chat import InternVLChatConfig
 from .conversation import get_conv_template
 from .modeling_intern_vit import InternVisionModel
@@ -280,13 +282,6 @@ class InternVLChatModel(PreTrainedModel):
         model_inputs = tokenizer(query, return_tensors='pt')
         input_ids = model_inputs['input_ids'].cuda()
         attention_mask = model_inputs['attention_mask'].cuda()
-
-        if verbose:
-            print(f"hehe: {self.num_image_token * num_patches}")
-            print(f"hehe: {query}")
-            print(f"hehe: {input_ids.shape}")
-            print(f"hehe: {attention_mask.shape}")        
-        
         generation_config['eos_token_id'] = eos_token_id
         generation_output = self.generate(
             pixel_values=pixel_values,
@@ -324,8 +319,18 @@ class InternVLChatModel(PreTrainedModel):
             if visual_features is not None:
                 vit_embeds = visual_features
             else:
-                vit_embeds = self.extract_feature(pixel_values)
-
+                vit_embeds = []
+                num_chunks = 2
+                pixel_values_splitted = pixel_values.chunk(num_chunks)
+
+                for pixel_values_ in pixel_values_splitted:     
+                    pixel_values_ = pixel_values_.flatten(0, 1)
+                    v_feats = self.extract_feature(pixel_values_) # examples: (18, 256, 896) = 2 * (9, 256, 896)
+                    v_feats = rearrange(v_feats, "(b n) t c -> b n t c", n=num_chunks) # b: batch_size, n: num_patches, t: num_tokens, c: hidden_size
+                    vit_embeds.append(v_feats)
+                    
+                vit_embeds = torch.cat(vit_embeds) # 
+            
             print(vit_embeds.shape)
                 
             input_embeds = self.language_model.get_input_embeddings()(input_ids)