iarchuk/InternVL
021
1# --------------------------------------------------------2# InternVL3# Copyright (c) 2024 OpenGVLab4# Licensed under The MIT License [see LICENSE for details]5# --------------------------------------------------------6 7import warnings8from typing import List, Optional, Tuple, Union9 10import torch.utils.checkpoint11import transformers12from torch import nn13from torch.nn import CrossEntropyLoss14from transformers import GenerationConfig15from transformers.modeling_outputs import CausalLMOutputWithPast16from transformers.modeling_utils import PreTrainedModel17from transformers.utils import logging18from transformers import LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM, Qwen3MoeForCausalLM19 20from .configuration_internvl_chat import InternVLChatConfig21from .conversation import get_conv_template22from .modeling_intern_vit import InternVisionModel, has_flash_attn23 24logger = logging.get_logger(__name__)25 26 27def version_cmp(v1, v2, op='eq'):28 import operator29 30 from packaging import version31 op_func = getattr(operator, op)32 return op_func(version.parse(v1), version.parse(v2))33 34 35class InternVLChatModel(PreTrainedModel):36 config_class = InternVLChatConfig37 main_input_name = 'pixel_values'38 base_model_prefix = 'language_model'39 _supports_flash_attn_2 = True40 supports_gradient_checkpointing = True41 _no_split_modules = [42 "InternVisionModel",43 "Qwen3DecoderLayer",44 ]45 46 # support transformers 4.51.+47 _tp_plan = ''48 49 def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):50 super().__init__(config)51 52 assert version_cmp(transformers.__version__, '4.37.0', 'ge')53 image_size = config.force_image_size or config.vision_config.image_size54 patch_size = config.vision_config.patch_size55 self.patch_size = patch_size56 self.select_layer = config.select_layer57 self.template = config.template58 self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))59 self.downsample_ratio = config.downsample_ratio60 self.ps_version = config.ps_version61 use_flash_attn = use_flash_attn if has_flash_attn else False62 config.vision_config.use_flash_attn = True if use_flash_attn else False63 config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'64 65 logger.info(f'num_image_token: {self.num_image_token}')66 logger.info(f'ps_version: {self.ps_version}')67 if vision_model is not None:68 self.vision_model = vision_model69 else:70 self.vision_model = InternVisionModel(config.vision_config)71 if language_model is not None:72 self.language_model = language_model73 else:74 architecture: str = config.llm_config.architectures[0]75 if architecture == 'LlamaForCausalLM':76 self.language_model = LlamaForCausalLM(config.llm_config)77 elif architecture == 'Qwen2ForCausalLM':78 self.language_model = Qwen2ForCausalLM(config.llm_config)79 elif architecture == 'Qwen3MoeForCausalLM':80 self.language_model = Qwen3MoeForCausalLM(config.llm_config)81 elif architecture == 'Qwen3ForCausalLM':82 self.language_model = Qwen3ForCausalLM(config.llm_config)83 else:84 raise NotImplementedError(f'{architecture} is not implemented.')85 86 vit_hidden_size = config.vision_config.hidden_size87 llm_hidden_size = config.llm_config.hidden_size88 89 self.mlp1 = nn.Sequential(90 nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),91 nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),92 nn.GELU(),93 nn.Linear(llm_hidden_size, llm_hidden_size)94 )95 96 self.img_context_token_id = None97 self.conv_template = get_conv_template(self.template)98 self.system_message = self.conv_template.system_message99 100 def forward(101 self,102 pixel_values: torch.FloatTensor,103 input_ids: torch.LongTensor = None,104 attention_mask: Optional[torch.Tensor] = None,105 position_ids: Optional[torch.LongTensor] = None,106 image_flags: Optional[torch.LongTensor] = None,107 past_key_values: Optional[List[torch.FloatTensor]] = None,108 labels: Optional[torch.LongTensor] = None,109 use_cache: Optional[bool] = None,110 output_attentions: Optional[bool] = None,111 output_hidden_states: Optional[bool] = None,112 return_dict: Optional[bool] = None,113 ) -> Union[Tuple, CausalLMOutputWithPast]:114 return_dict = return_dict if return_dict is not None else self.config.use_return_dict115 116 image_flags = image_flags.squeeze(-1)117 input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()118 119 vit_embeds = self.extract_feature(pixel_values)120 vit_embeds = vit_embeds[image_flags == 1]121 vit_batch_size = pixel_values.shape[0]122 123 B, N, C = input_embeds.shape124 input_embeds = input_embeds.reshape(B * N, C)125 126 # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:127 # print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')128 129 input_ids = input_ids.reshape(B * N)130 selected = (input_ids == self.img_context_token_id)131 try:132 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)133 except Exception as e:134 vit_embeds = vit_embeds.reshape(-1, C)135 print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '136 f'vit_embeds.shape={vit_embeds.shape}')137 n_token = min(selected.sum(), vit_embeds.size(0))138 input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]139 140 input_embeds = input_embeds.reshape(B, N, C)141 142 outputs = self.language_model(143 inputs_embeds=input_embeds,144 attention_mask=attention_mask,145 position_ids=position_ids,146 past_key_values=past_key_values,147 use_cache=use_cache,148 output_attentions=output_attentions,149 output_hidden_states=output_hidden_states,150 return_dict=return_dict,151 )152 logits = outputs.logits153 154 loss = None155 if labels is not None:156 # Shift so that tokens < n predict n157 shift_logits = logits[..., :-1, :].contiguous()158 shift_labels = labels[..., 1:].contiguous()159 # Flatten the tokens160 loss_fct = CrossEntropyLoss()161 shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)162 shift_labels = shift_labels.view(-1)163 # Enable model parallelism164 shift_labels = shift_labels.to(shift_logits.device)165 loss = loss_fct(shift_logits, shift_labels)166 167 if not return_dict:168 output = (logits,) + outputs[1:]169 return (loss,) + output if loss is not None else output170 171 return CausalLMOutputWithPast(172 loss=loss,173 logits=logits,174 past_key_values=outputs.past_key_values,175 hidden_states=outputs.hidden_states,176 attentions=outputs.attentions,177 )178 179 def pixel_shuffle(self, x, scale_factor=0.5):180 n, w, h, c = x.size()181 # N, W, H, C --> N, W, H * scale, C // scale182 x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))183 # N, W, H * scale, C // scale --> N, H * scale, W, C // scale184 x = x.permute(0, 2, 1, 3).contiguous()185 # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)186 x = x.view(n, int(h * scale_factor), int(w * scale_factor),187 int(c / (scale_factor * scale_factor)))188 if self.ps_version == 'v1':189 warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "190 'which results in a transposed image.')191 else:192 x = x.permute(0, 2, 1, 3).contiguous()193 return x194 195 def extract_feature(self, pixel_values):196 if self.select_layer == -1:197 vit_embeds = self.vision_model(198 pixel_values=pixel_values,199 output_hidden_states=False,200 return_dict=True).last_hidden_state201 else:202 vit_embeds = self.vision_model(203 pixel_values=pixel_values,204 output_hidden_states=True,205 return_dict=True).hidden_states[self.select_layer]206 vit_embeds = vit_embeds[:, 1:, :]207 208 h = w = int(vit_embeds.shape[1] ** 0.5)209 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)210 vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)211 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])212 vit_embeds = self.mlp1(vit_embeds)213 return vit_embeds214 215 def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,216 history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',217 IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):218 if history is not None or return_history:219 print('Now multi-turn chat is not supported in batch_chat.')220 raise NotImplementedError221 222 if image_counts is not None:223 num_patches_list = image_counts224 print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')225 226 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)227 self.img_context_token_id = img_context_token_id228 229 if verbose and pixel_values is not None:230 image_bs = pixel_values.shape[0]231 print(f'dynamic ViT batch size: {image_bs}')232 233 queries = []234 for idx, num_patches in enumerate(num_patches_list):235 question = questions[idx]236 if pixel_values is not None and '<image>' not in question:237 question = '<image>\n' + question238 template = get_conv_template(self.template)239 template.system_message = self.system_message240 template.append_message(template.roles[0], question)241 template.append_message(template.roles[1], None)242 query = template.get_prompt()243 244 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN245 query = query.replace('<image>', image_tokens, 1)246 queries.append(query)247 248 tokenizer.padding_side = 'left'249 model_inputs = tokenizer(queries, return_tensors='pt', padding=True)250 input_ids = model_inputs['input_ids'].to(self.device)251 attention_mask = model_inputs['attention_mask'].to(self.device)252 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())253 generation_config['eos_token_id'] = eos_token_id254 generation_output = self.generate(255 pixel_values=pixel_values,256 input_ids=input_ids,257 attention_mask=attention_mask,258 **generation_config259 )260 responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)261 responses = [response.split(template.sep.strip())[0].strip() for response in responses]262 return responses263 264 def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,265 num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',266 verbose=False):267 268 if history is None and pixel_values is not None and '<image>' not in question:269 question = '<image>\n' + question270 271 if num_patches_list is None:272 num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []273 assert pixel_values is None or len(pixel_values) == sum(num_patches_list)274 275 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)276 self.img_context_token_id = img_context_token_id277 278 template = get_conv_template(self.template)279 template.system_message = self.system_message280 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())281 282 history = [] if history is None else history283 for (old_question, old_answer) in history:284 template.append_message(template.roles[0], old_question)285 template.append_message(template.roles[1], old_answer)286 template.append_message(template.roles[0], question)287 template.append_message(template.roles[1], None)288 query = template.get_prompt()289 290 if verbose and pixel_values is not None:291 image_bs = pixel_values.shape[0]292 print(f'dynamic ViT batch size: {image_bs}')293 294 for num_patches in num_patches_list:295 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN296 query = query.replace('<image>', image_tokens, 1)297 298 model_inputs = tokenizer(query, return_tensors='pt')299 input_ids = model_inputs['input_ids'].to(self.device)300 attention_mask = model_inputs['attention_mask'].to(self.device)301 generation_config['eos_token_id'] = eos_token_id302 generation_output = self.generate(303 pixel_values=pixel_values,304 input_ids=input_ids,305 attention_mask=attention_mask,306 **generation_config307 )308 response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]309 response = response.split(template.sep.strip())[0].strip()310 history.append((question, response))311 if return_history:312 return response, history313 else:314 query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')315 query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')316 if verbose:317 print(query_to_print, response)318 return response319 320 @torch.no_grad()321 def generate(322 self,323 pixel_values: Optional[torch.FloatTensor] = None,324 input_ids: Optional[torch.FloatTensor] = None,325 attention_mask: Optional[torch.LongTensor] = None,326 visual_features: Optional[torch.FloatTensor] = None,327 generation_config: Optional[GenerationConfig] = None,328 output_hidden_states: Optional[bool] = None,329 **generate_kwargs,330 ) -> torch.LongTensor:331 332 assert self.img_context_token_id is not None333 if pixel_values is not None:334 if visual_features is not None:335 vit_embeds = visual_features336 else:337 vit_embeds = self.extract_feature(pixel_values)338 input_embeds = self.language_model.get_input_embeddings()(input_ids)339 B, N, C = input_embeds.shape340 input_embeds = input_embeds.reshape(B * N, C)341 342 input_ids = input_ids.reshape(B * N)343 selected = (input_ids == self.img_context_token_id)344 assert selected.sum() != 0345 input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)346 347 input_embeds = input_embeds.reshape(B, N, C)348 else:349 input_embeds = self.language_model.get_input_embeddings()(input_ids)350 351 outputs = self.language_model.generate(352 inputs_embeds=input_embeds,353 attention_mask=attention_mask,354 generation_config=generation_config,355 output_hidden_states=output_hidden_states,356 use_cache=True,357 **generate_kwargs,358 )359 360 return outputs361 362 @property363 def lm_head(self):364 return self.language_model.get_output_embeddings()365 366 def get_output_embeddings(self):367 return self.language_model.get_output_embeddings()368 369 def get_input_embeddings(self):370 return self.language_model.get_input_embeddings()371 372 def set_input_embeddings(self, value):373 return self.language_model.set_input_embeddings(value)374 375 def set_output_embeddings(self, value):376 return self.language_model.set_output_embeddings(value)377