Update modeling_internvl_chat.py

modeling_internvl_chat.py

@@ -24,6 +24,8 @@ from .modeling_internlm2 import InternLM2ForCausalLM
 
 logger = logging.get_logger(__name__)
 
+import os 
+image_token_num = 0
 
 def version_cmp(v1, v2, op='eq'):
     import operator
@@ -53,6 +55,8 @@ class InternVLChatModel(PreTrainedModel):
         self.downsample_ratio = config.downsample_ratio
         self.ps_version = config.ps_version
         use_flash_attn = use_flash_attn if has_flash_attn else False
+        #use_flash_attn = True
+        #use_flash_attn = False
         config.vision_config.use_flash_attn = True if use_flash_attn else False
         config.llm_config.attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
 
@@ -182,6 +186,7 @@ class InternVLChatModel(PreTrainedModel):
         return x
 
     def extract_feature(self, pixel_values):
+        
         if self.select_layer == -1:
             vit_embeds = self.vision_model(
                 pixel_values=pixel_values,
@@ -193,9 +198,11 @@ class InternVLChatModel(PreTrainedModel):
                 output_hidden_states=True,
                 return_dict=True).hidden_states[self.select_layer]
         vit_embeds = vit_embeds[:, 1:, :]
-
+        
         h = w = int(vit_embeds.shape[1] ** 0.5)
+        os.environ['IMAGE_H'] = str(h)
         vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        # import pdb; pdb.set_trace()
         vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
         vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
         vit_embeds = self.mlp1(vit_embeds)
@@ -233,13 +240,14 @@ class InternVLChatModel(PreTrainedModel):
             image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
             query = query.replace('<image>', image_tokens, 1)
             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)
         eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
         generation_config['eos_token_id'] = eos_token_id
+        
         generation_output = self.generate(
             pixel_values=pixel_values,
             input_ids=input_ids,
@@ -317,14 +325,32 @@ class InternVLChatModel(PreTrainedModel):
             output_hidden_states: Optional[bool] = None,
             **generate_kwargs,
     ) -> torch.LongTensor:
-
+        
         assert self.img_context_token_id is not None
         if pixel_values is not None:
+            
             if visual_features is not None:
                 vit_embeds = visual_features
             else:
-                vit_embeds = self.extract_feature(pixel_values)
+                #vit_embeds = self.extract_feature(pixel_values)
+                # Assuming pixel_values is already defined
+                batch_size = 10
+                num_samples = pixel_values.size(0)  # Total number of samples
+                vit_embeds_list = []
+
+                # Loop through the batches
+                for start_idx in range(0, num_samples, batch_size):
+                    end_idx = min(start_idx + batch_size, num_samples)  # Ensure the end index doesn't exceed the size
+                    batch = pixel_values[start_idx:end_idx]  # Slice the batch
+                    vit_embeds_batch = self.extract_feature(batch)  # Process the batch
+                    vit_embeds_list.append(vit_embeds_batch)  # Collect the results
+
+                # Concatenate the embeddings if requiimport pdb; pdb.set_trace()red
+                vit_embeds = torch.cat(vit_embeds_list, dim=0)
+
+            
             input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            
             B, N, C = input_embeds.shape
             input_embeds = input_embeds.reshape(B * N, C)
 
@@ -332,11 +358,14 @@ class InternVLChatModel(PreTrainedModel):
             selected = (input_ids == self.img_context_token_id)
             assert selected.sum() != 0
             input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
-
+            
+            image_token_num = int(vit_embeds.shape[0] * vit_embeds.shape[1]/B)
+            os.environ['IMAGE_TOKEN_NUM'] = str(image_token_num)
+            
             input_embeds = input_embeds.reshape(B, N, C)
         else:
             input_embeds = self.language_model.get_input_embeddings()(input_ids)
-
+        
         outputs = self.language_model.generate(
             inputs_embeds=input_embeds,
             attention_mask=attention_mask,