OpenGVLab/InternVideo2_5_Chat_8B
923.9k
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, Union, Callable9 10import torch11import torch.utils.checkpoint12import transformers13from torch import nn14from torch.nn import CrossEntropyLoss15from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,16 LlamaTokenizer)17from transformers.modeling_outputs import CausalLMOutputWithPast18from transformers.modeling_utils import PreTrainedModel19from transformers.utils import ModelOutput, logging20 21from .configuration_internvl_chat import InternVLChatConfig22from .conversation import get_conv_template23from .modeling_intern_vit import InternVisionModel, has_flash_attn24from .modeling_internlm2 import InternLM2ForCausalLM25 26logger = logging.get_logger(__name__)27 28 29 30def bipartite_soft_matching(31 metric: torch.Tensor,32 r: int, 33) -> Tuple[Callable, Callable]:34 """35 Applies ToMe with a balanced matching set (50%, 50%).36 37 Input size is [batch, tokens, channels].38 r indicates the number of tokens to remove (max 50% of tokens).39 """40 protected = 041 42 t = metric.shape[1]43 r = min(r, (t - protected) // 2)44 45 assert r > 0, r46 47 with torch.no_grad():48 metric = metric / metric.norm(dim=-1, keepdim=True)49 a, b = metric[..., ::2, :], metric[..., 1::2, :]50 scores = a @ b.transpose(-1, -2)51 52 node_max, node_idx = scores.max(dim=-1)53 edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]54 55 unm_idx = edge_idx[..., r:, :] # Unmerged Tokens56 src_idx = edge_idx[..., :r, :] # Merged Tokens57 dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)58 59 def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:60 src, dst = x[..., ::2, :], x[..., 1::2, :]61 n, t1, c = src.shape62 unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c))63 src = src.gather(dim=-2, index=src_idx.expand(n, r, c))64 dst = dst.scatter_add(-2, dst_idx.expand(n, r, c), src) # , reduce=mode)65 66 return torch.cat([unm, dst], dim=1)67 68 def unmerge(x: torch.Tensor) -> torch.Tensor:69 unm_len = unm_idx.shape[1]70 unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]71 n, _, c = unm.shape72 73 src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c))74 75 out = torch.zeros(n, metric.shape[1], c, device=x.device, dtype=x.dtype)76 77 out[..., 1::2, :] = dst78 out.scatter_(dim=-2, index=(2 * unm_idx).expand(n, unm_len, c), src=unm)79 out.scatter_(dim=-2, index=(2 * src_idx).expand(n, r, c), src=src)80 81 return out82 83 return merge, unmerge84 85 86def merge_wavg(87 merge: Callable, x: torch.Tensor, size: torch.Tensor = None88) -> Tuple[torch.Tensor, torch.Tensor]:89 """90 Applies the merge function by taking a weighted average based on token size.91 Returns the merged tensor and the new token sizes.92 """93 if size is None:94 size = torch.ones_like(x[..., 0, None])95 96 x = merge(x * size, mode="sum")97 size = merge(size, mode="sum")98 99 x = x / size100 return x, size101 102 103def version_cmp(v1, v2, op='eq'):104 import operator105 106 from packaging import version107 op_func = getattr(operator, op)108 return op_func(version.parse(v1), version.parse(v2))109 110 111class InternVLChatModel(PreTrainedModel):112 config_class = InternVLChatConfig113 main_input_name = 'pixel_values'114 base_model_prefix = 'language_model'115 _supports_flash_attn_2 = True116 _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'InternLM2DecoderLayer']117 118 def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):119 super().__init__(config)120 121 assert version_cmp(transformers.__version__, '4.36.2', 'ge')122 image_size = config.force_image_size or config.vision_config.image_size123 patch_size = config.vision_config.patch_size124 self.local_num_frames = 4125 self.num_tome_tokens = 64126 self.config = config127 self.patch_size = patch_size128 self.select_layer = config.select_layer129 self.template = config.template130 # self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2)) //131 self.num_image_token = self.num_tome_tokens // self.local_num_frames132 self.downsample_ratio = config.downsample_ratio133 self.ps_version = config.ps_version134 use_flash_attn = use_flash_attn if has_flash_attn else False135 config.vision_config.use_flash_attn = True if use_flash_attn else False136 config.llm_config.attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'137 138 logger.info(f'num_image_token: {self.num_image_token}')139 logger.info(f'ps_version: {self.ps_version}')140 if vision_model is not None:141 self.vision_model = vision_model142 else:143 self.vision_model = InternVisionModel(config.vision_config)144 if language_model is not None:145 self.language_model = language_model146 else:147 if config.llm_config.architectures[0] == 'LlamaForCausalLM':148 self.language_model = LlamaForCausalLM(config.llm_config)149 elif config.llm_config.architectures[0] == 'InternLM2ForCausalLM':150 self.language_model = InternLM2ForCausalLM(config.llm_config)151 else:152 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')153 154 vit_hidden_size = config.vision_config.hidden_size155 llm_hidden_size = config.llm_config.hidden_size156 157 self.mlp1 = nn.Sequential(158 nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),159 nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),160 nn.GELU(),161 nn.Linear(llm_hidden_size, llm_hidden_size)162 )163 164 self.img_context_token_id = None165 self.conv_template = get_conv_template(self.template)166 self.system_message = self.conv_template.system_message167 168 def merge_tokens(self, x, target_num_token):169 r"""170 x = torch.randn(10, 2560, c)171 x = merge_tokens(x, r_merge_list=[1280])172 """173 size = None174 b, p, c = x.shape175 tmp_p = p176 r_merge_list = []177 assert tmp_p > target_num_token, f"{tmp_p} should greater than {target_num_token}"178 while tmp_p != target_num_token:179 if tmp_p - target_num_token <= (tmp_p // 2):180 r_merge_list.append(tmp_p - target_num_token)181 break182 else:183 r_merge_list.append(tmp_p // 2)184 tmp_p = tmp_p - (tmp_p // 2)185 186 187 head = self.config.llm_config.num_attention_heads188 189 dim = c // head190 for r in r_merge_list:191 metric = x.reshape(b, p, head, dim).mean(2) # [b, p, c//head]192 merge, _ = bipartite_soft_matching(193 metric, 194 r195 )196 x, size = merge_wavg(merge, x, size)197 _, p, _ = x.shape198 # x = x.reshape(-1, c) # 300, 1024199 return x200 201 def forward(202 self,203 pixel_values: torch.FloatTensor,204 input_ids: torch.LongTensor = None,205 attention_mask: Optional[torch.Tensor] = None,206 position_ids: Optional[torch.LongTensor] = None,207 image_flags: Optional[torch.LongTensor] = None,208 past_key_values: Optional[List[torch.FloatTensor]] = None,209 labels: Optional[torch.LongTensor] = None,210 use_cache: Optional[bool] = None,211 output_attentions: Optional[bool] = None,212 output_hidden_states: Optional[bool] = None,213 return_dict: Optional[bool] = None,214 ) -> Union[Tuple, CausalLMOutputWithPast]:215 return_dict = return_dict if return_dict is not None else self.config.use_return_dict216 217 image_flags = image_flags.squeeze(-1)218 input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()219 220 vit_embeds = self.extract_feature(pixel_values)221 vit_embeds = vit_embeds[image_flags == 1]222 vit_batch_size = pixel_values.shape[0]223 224 B, N, C = input_embeds.shape225 input_embeds = input_embeds.reshape(B * N, C)226 227 if torch.distributed.get_rank() == 0:228 print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')229 230 input_ids = input_ids.reshape(B * N)231 selected = (input_ids == self.img_context_token_id)232 try:233 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)234 except Exception as e:235 vit_embeds = vit_embeds.reshape(-1, C)236 print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '237 f'vit_embeds.shape={vit_embeds.shape}')238 n_token = selected.sum()239 input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token]240 241 input_embeds = input_embeds.reshape(B, N, C)242 243 outputs = self.language_model(244 inputs_embeds=input_embeds,245 attention_mask=attention_mask,246 position_ids=position_ids,247 past_key_values=past_key_values,248 use_cache=use_cache,249 output_attentions=output_attentions,250 output_hidden_states=output_hidden_states,251 return_dict=return_dict,252 )253 logits = outputs.logits254 255 loss = None256 if labels is not None:257 # Shift so that tokens < n predict n258 shift_logits = logits[..., :-1, :].contiguous()259 shift_labels = labels[..., 1:].contiguous()260 # Flatten the tokens261 loss_fct = CrossEntropyLoss()262 shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)263 shift_labels = shift_labels.view(-1)264 # Enable model parallelism265 shift_labels = shift_labels.to(shift_logits.device)266 loss = loss_fct(shift_logits, shift_labels)267 268 if not return_dict:269 output = (logits,) + outputs[1:]270 return (loss,) + output if loss is not None else output271 272 return CausalLMOutputWithPast(273 loss=loss,274 logits=logits,275 past_key_values=outputs.past_key_values,276 hidden_states=outputs.hidden_states,277 attentions=outputs.attentions,278 )279 280 def pixel_shuffle(self, x, scale_factor=0.5):281 n, w, h, c = x.size()282 # N, W, H, C --> N, W, H * scale, C // scale283 x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))284 # N, W, H * scale, C // scale --> N, H * scale, W, C // scale285 x = x.permute(0, 2, 1, 3).contiguous()286 # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)287 x = x.view(n, int(h * scale_factor), int(w * scale_factor),288 int(c / (scale_factor * scale_factor)))289 if self.ps_version == 'v1':290 warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "291 'which results in a transposed image.')292 else:293 x = x.permute(0, 2, 1, 3).contiguous()294 295 return x296 297 def extract_feature(self, pixel_values):298 if self.select_layer == -1:299 vit_embeds = self.vision_model(300 pixel_values=pixel_values,301 output_hidden_states=False,302 return_dict=True).last_hidden_state303 else:304 vit_embeds = self.vision_model(305 pixel_values=pixel_values,306 output_hidden_states=True,307 return_dict=True).hidden_states[self.select_layer]308 vit_embeds = vit_embeds[:, 1:, :]309 310 h = w = int(vit_embeds.shape[1] ** 0.5)311 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)312 vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)313 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0] // self.local_num_frames, -1, vit_embeds.shape[-1])314 vit_embeds = self.merge_tokens(vit_embeds, self.num_tome_tokens)315 vit_embeds = vit_embeds.reshape(vit_embeds.shape[0] * self.local_num_frames, -1, vit_embeds.shape[-1])316 vit_embeds = self.mlp1(vit_embeds)317 return vit_embeds318 319 def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,320 history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',321 IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):322 if history is not None or return_history:323 print('Now multi-turn chat is not supported in batch_chat.')324 raise NotImplementedError325 326 if image_counts is not None:327 num_patches_list = image_counts328 print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')329 330 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)331 self.img_context_token_id = img_context_token_id332 333 if verbose and pixel_values is not None:334 image_bs = pixel_values.shape[0]335 print(f'dynamic ViT batch size: {image_bs}')336 337 queries = []338 for idx, num_patches in enumerate(num_patches_list):339 question = questions[idx]340 if pixel_values is not None and '<image>' not in question:341 question = '<image>\n' + question342 template = get_conv_template(self.template)343 template.system_message = self.system_message344 template.append_message(template.roles[0], question)345 template.append_message(template.roles[1], None)346 query = template.get_prompt()347 348 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN349 query = query.replace('<image>', image_tokens, 1)350 queries.append(query)351 352 tokenizer.padding_side = 'left'353 model_inputs = tokenizer(queries, return_tensors='pt', padding=True)354 input_ids = model_inputs['input_ids'].to(self.device)355 attention_mask = model_inputs['attention_mask'].to(self.device)356 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())357 generation_config['eos_token_id'] = eos_token_id358 generation_output = self.generate(359 pixel_values=pixel_values,360 input_ids=input_ids,361 attention_mask=attention_mask,362 **generation_config363 )364 responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)365 responses = [response.split(template.sep.strip())[0].strip() for response in responses]366 return responses367 368 def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,369 num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',370 verbose=False):371 372 if history is None and pixel_values is not None and '<image>' not in question:373 question = '<image>\n' + question374 375 if num_patches_list is None:376 num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []377 assert pixel_values is None or len(pixel_values) == sum(num_patches_list)378 379 img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)380 self.img_context_token_id = img_context_token_id381 382 template = get_conv_template(self.template)383 template.system_message = self.system_message384 eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())385 386 history = [] if history is None else history387 for (old_question, old_answer) in history:388 template.append_message(template.roles[0], old_question)389 template.append_message(template.roles[1], old_answer)390 template.append_message(template.roles[0], question)391 template.append_message(template.roles[1], None)392 query = template.get_prompt()393 394 if verbose and pixel_values is not None:395 image_bs = pixel_values.shape[0]396 print(f'dynamic ViT batch size: {image_bs}')397 398 for num_patches in num_patches_list:399 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN400 query = query.replace('<image>', image_tokens, 1)401 402 model_inputs = tokenizer(query, return_tensors='pt')403 input_ids = model_inputs['input_ids'].to(self.device)404 attention_mask = model_inputs['attention_mask'].to(self.device)405 generation_config['eos_token_id'] = eos_token_id406 generation_output = self.generate(407 pixel_values=pixel_values,408 input_ids=input_ids,409 attention_mask=attention_mask,410 **generation_config411 )412 response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]413 response = response.split(template.sep.strip())[0].strip()414 history.append((question, response))415 if return_history:416 return response, history417 else:418 query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')419 query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')420 if verbose:421 print(query_to_print, response)422 return response423 424 @torch.no_grad()425 def generate(426 self,427 pixel_values: Optional[torch.FloatTensor] = None,428 input_ids: Optional[torch.FloatTensor] = None,429 attention_mask: Optional[torch.LongTensor] = None,430 visual_features: Optional[torch.FloatTensor] = None,431 generation_config: Optional[GenerationConfig] = None,432 output_hidden_states: Optional[bool] = None,433 **generate_kwargs,434 ) -> torch.LongTensor:435 436 assert self.img_context_token_id is not None437 if pixel_values is not None:438 if visual_features is not None:439 vit_embeds = visual_features440 else:441 vit_embeds = self.extract_feature(pixel_values)442 input_embeds = self.language_model.get_input_embeddings()(input_ids)443 B, N, C = input_embeds.shape444 input_embeds = input_embeds.reshape(B * N, C)445 446 input_ids = input_ids.reshape(B * N)447 selected = (input_ids == self.img_context_token_id)448 assert selected.sum() != 0449 input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)450 451 input_embeds = input_embeds.reshape(B, N, C)452 else:453 input_embeds = self.language_model.get_input_embeddings()(input_ids)454 455 outputs = self.language_model.generate(456 inputs_embeds=input_embeds,457 attention_mask=attention_mask,458 generation_config=generation_config,459 output_hidden_states=output_hidden_states,460 use_cache=True,461 **generate_kwargs,462 )463 464 return outputs465