OpenGVLab/VideoChat-Flash-Qwen2-7B_res448
131.1k
1# Copyright 20242#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15from abc import ABC, abstractmethod16import re17import torch18import torch.nn as nn19import random20from typing import List, Optional, Tuple, Union, Dict21 22from transformers import AutoConfig, AutoModelForCausalLM23from transformers.modeling_outputs import CausalLMOutputWithPast24from transformers.generation.utils import GenerateOutput25from transformers import Qwen2Config26 27from .vision_tower_builder import build_vision_tower28from .mm_projector_builder import build_vision_projector29 30from .constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_IMAGE_TOKEN31from .conversation import conv_templates, SeparatorStyle32from .mm_utils import tokenizer_image_token, KeywordsStoppingCriteria, get_anyres_image_grid_shape, load_video33from .modeling_qwen2_flash import Qwen2Model_Flash, Qwen2ForCausalLM_Flash34 35 36class LlavaMetaModel:37 38 def __init__(self, config):39 super(LlavaMetaModel, self).__init__(config)40 41 if hasattr(config, "mm_vision_tower"):42 delay_load = getattr(config, "delay_load", False)43 self.vision_tower = build_vision_tower(config, delay_load=delay_load)44 self.mm_projector = build_vision_projector(config, vision_cfg=self.vision_tower.config)45 46 if "unpad" in getattr(config, "mm_patch_merge_type", ""):47 self.image_newline = nn.Parameter(torch.empty(config.hidden_size, dtype=self.dtype))48 if "nopad" in getattr(config, "mm_patch_merge_type", "") and getattr(self.config, "mm_newline_position", "nothing") != "nothing":49 self.frame_newline = nn.Parameter(torch.empty(config.hidden_size, dtype=self.dtype))50 51 def get_vision_tower(self):52 vision_tower = getattr(self, "vision_tower", None)53 if type(vision_tower) is list:54 vision_tower = vision_tower[0]55 return vision_tower56 57 def initialize_vision_modules(self, model_args, fsdp=None):58 vision_tower = model_args.vision_tower59 mm_vision_select_layer = model_args.mm_vision_select_layer60 mm_vision_select_feature = model_args.mm_vision_select_feature61 pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter62 mm_patch_merge_type = model_args.mm_patch_merge_type63 64 self.config.mm_vision_tower = vision_tower65 self.config.vision_tower_pretrained = getattr(model_args, "vision_tower_pretrained", "")66 67 if self.get_vision_tower() is None:68 vision_tower = build_vision_tower(model_args)69 70 if fsdp is not None and len(fsdp) > 0:71 self.vision_tower = [vision_tower]72 else:73 self.vision_tower = vision_tower74 else:75 if fsdp is not None and len(fsdp) > 0:76 vision_tower = self.vision_tower[0]77 else:78 vision_tower = self.vision_tower79 vision_tower.load_model()80 81 82 83 self.config.use_mm_proj = True84 self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")85 self.config.mm_vision_select_layer = mm_vision_select_layer86 self.config.mm_vision_select_feature = mm_vision_select_feature87 self.config.mm_patch_merge_type = mm_patch_merge_type88 89 if getattr(self, "mm_projector", None) is None:90 self.mm_projector = build_vision_projector(self.config, vision_cfg=vision_tower.config)91 92 if "unpad" in mm_patch_merge_type:93 embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))94 self.image_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)95 if "nopad" in getattr(self.config, "mm_patch_merge_type", "") and getattr(self.config, "mm_newline_position", "nothing") != "nothing":96 embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))97 self.frame_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)98 else:99 # In case it is frozen by LoRA100 for p in self.mm_projector.parameters():101 p.requires_grad = True102 103 if pretrain_mm_mlp_adapter is not None:104 mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")105 106 def get_w(weights, keyword):107 return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}108 109 if self.config.mm_projector_type =='lxh_qformer':110 incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"), strict=False)111 else:112 incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))113 print(f"Loaded mm projector weights from {pretrain_mm_mlp_adapter}. Incompatible keys: {incompatible_keys}")114 115 116class LlavaMetaForCausalLM(ABC):117 118 @abstractmethod119 def get_model(self):120 pass121 122 def get_vision_tower(self):123 return self.get_model().get_vision_tower()124 125 126 def encode_video_image(self, images_list, video_idx_in_batch):127 # video encoder编码后按图像的connector处理128 bs = len(images_list)129 130 concat_images = []131 concat_videos = []132 for idx, image in enumerate(images_list):133 if idx in video_idx_in_batch:134 concat_videos.append(image)135 else:136 concat_images.append(image)137 # print(concat_videos[0].shape)138 has_image = len(concat_images) > 0139 has_video = len(concat_videos) > 0140 141 mm_local_num_frames = getattr(self.config, "mm_local_num_frames", -1)142 assert mm_local_num_frames != -1143 if has_image:144 image_split_sizes = [image.shape[0] for image in concat_images] 145 concat_images = torch.cat([image.unsqueeze(1) for image in concat_images], dim=0)146 # print("input vit image.shape:", concat_images.shape)147 images_features = self.get_model().get_vision_tower()(concat_images) # B_i, N, D148 images_features = torch.split(images_features, image_split_sizes)149 150 if has_video:151 video_split_sizes = [video.shape[0] // mm_local_num_frames for video in concat_videos]152 concat_videos = torch.cat([video.reshape(video.shape[0] // mm_local_num_frames, mm_local_num_frames, video.shape[1], video.shape[2], video.shape[3]) for video in concat_videos], dim=0)153 # print("input vit video.shape:", concat_videos.shape)154 videos_features = self.get_model().get_vision_tower()(concat_videos) # B_v, N, D155 videos_features = [v.reshape(-1, v.shape[-2] // mm_local_num_frames, v.shape[-1]) for v in torch.split(videos_features, video_split_sizes)]156 157 158 all_videos_or_images_features = []159 img_idx = 0160 vid_idx = 0161 162 for idx in range(bs):163 164 if idx in video_idx_in_batch:165 feat = self.get_model().mm_projector(videos_features[vid_idx], compress=True, local_num_frames=getattr(self.config, "mm_local_num_frames", -1))166 167 vid_idx += 1168 else:169 feat = self.get_model().mm_projector(images_features[img_idx], compress=False)170 img_idx += 1171 # print("video_idx_in_batch:", video_idx_in_batch)172 all_videos_or_images_features.append(feat)173 174 if has_video:175 assert vid_idx == len(videos_features), f"vid: {vid_idx} != {len(videos_features)}"176 if has_image:177 assert img_idx == len(images_features), f"img: {img_idx} != {len(images_features)}"178 179 return all_videos_or_images_features180 181 182 183 def prepare_inputs_labels_for_multimodal(self, input_ids, position_ids, attention_mask, past_key_values, labels, images, modalities=["image"], image_sizes=None):184 assert type(modalities) is list, modalities185 186 vision_tower = self.get_vision_tower()187 # rank_print(modalities)188 if vision_tower is None or images is None or input_ids.shape[1] == 1:189 return input_ids, position_ids, attention_mask, past_key_values, None, labels190 191 if type(images) is list or images.ndim == 5:192 if type(images) is list:193 images = [x.unsqueeze(0) if x.ndim == 3 else x for x in images]194 195 video_idx_in_batch = []196 for _ in range(len(modalities)):197 if modalities[_] == "video":198 video_idx_in_batch.append(_)199 200 images_list = []201 for image in images:202 if image.ndim == 4:203 images_list.append(image)204 else:205 images_list.append(image.unsqueeze(0))206 207 208 vision_encode_type = getattr(self.config, "vision_encode_type", "image")209 mm_patch_merge_type = getattr(self.config, "mm_patch_merge_type", "flat")210 image_aspect_ratio = getattr(self.config, "image_aspect_ratio", "square")211 frame_aspect_ratio = getattr(self.config, "frame_aspect_ratio", "square")212 mm_newline_position = getattr(self.config, "mm_newline_position", "nothing")213 214 215 if vision_encode_type == "video_image": # video backbone, process video with compress216 image_features = self.encode_video_image(images_list, video_idx_in_batch=video_idx_in_batch)217 else:218 raise NotImplementedError(vision_encode_type)219 220 221 if mm_patch_merge_type == "flat":222 image_features = [x.flatten(0, 1) for x in image_features]223 elif mm_patch_merge_type.startswith("spatial"):224 new_image_features = []225 for image_idx, image_feature in enumerate(image_features):226 227 if image_idx in video_idx_in_batch: # video operations228 229 if "anyres" in frame_aspect_ratio:230 raise NotImplementedError231 else:232 frame_feature = image_feature233 234 if "pad" in mm_patch_merge_type:235 if mm_newline_position == 'one_token':236 frame_feature = frame_feature.flatten(0, 1)237 if "unpad" in mm_patch_merge_type:238 frame_feature = torch.cat((frame_feature, self.model.image_newline[None].to(frame_feature.device)), dim=0)239 else:240 frame_feature = torch.cat((frame_feature, self.model.frame_newline[None].to(frame_feature.device)), dim=0)241 elif mm_newline_position == 'nothing':242 frame_feature = frame_feature.flatten(0, 1)243 else:244 raise NotImplementedError("add pad please!!")245 else:246 frame_feature = frame_feature.flatten(0, 1)247 248 # print(f"final video frame_feature.shape: {frame_feature.shape}")249 image_feature = frame_feature250 251 elif image_feature.shape[0] > 1: # multi patches and multi images operations252 base_image_feature = image_feature[0]253 image_feature = image_feature[1:]254 origin_size = image_feature.shape255 256 height = width = self.get_model().mm_projector.num_image_patches_per_side 257 assert height * width == base_image_feature.shape[0], f"height:{height}, width: {width}, base_image_feature: {base_image_feature.shape}"258 259 if "anyres_max" in image_aspect_ratio:260 matched_anyres_max_num_patches = re.match(r"anyres_max_(\d+)", image_aspect_ratio)261 if matched_anyres_max_num_patches:262 max_num_patches = int(matched_anyres_max_num_patches.group(1))263 264 if "anyres" in image_aspect_ratio:265 if hasattr(self.get_vision_tower(), "image_size"):266 vision_tower_image_size = self.get_vision_tower().image_size267 else:268 raise ValueError("vision_tower_image_size is not found in the vision tower.")269 try:270 num_patch_width, num_patch_height = get_anyres_image_grid_shape(image_sizes[image_idx], self.config.image_grid_pinpoints, vision_tower_image_size, max_resolutions=None)271 except Exception as e:272 print(f"Error: {e}")273 raise e274 # num_patch_width, num_patch_height = 2, 2275 276 image_feature = image_feature.view(num_patch_height, num_patch_width, height, width, -1)277 else:278 raise NotImplementedError(image_aspect_ratio)279 image_feature = image_feature.view(2, 2, height, width, -1)280 281 if "maxpool2x2" in mm_patch_merge_type:282 raise NotImplementedError283 elif "unpad" in mm_patch_merge_type and "anyres_max" in image_aspect_ratio and matched_anyres_max_num_patches:284 raise NotImplementedError285 elif "unpad" in mm_patch_merge_type:286 raise NotImplementedError287 else:288 image_feature = image_feature.permute(0, 2, 1, 3, 4).contiguous()289 image_feature = image_feature.flatten(0, 3)290 if "nobase" in mm_patch_merge_type:291 pass292 else:293 try:294 image_feature = torch.cat((base_image_feature, image_feature), dim=0)295 except Exception as e:296 raise ValueError(f"{num_patch_width} {num_patch_height} now: base_image_feature: {base_image_feature.shape}, {image_feature.shape}, image_sizes[image_idx]: {image_sizes[image_idx]}, origin_size: {origin_size}, {image_sizes[image_idx]}, {self.config.image_grid_pinpoints}, {vision_tower_image_size}")297 else: # single image operations298 image_feature = image_feature[0]299 if "unpad" in mm_patch_merge_type:300 image_feature = torch.cat((image_feature, self.model.image_newline[None]), dim=0)301 302 # print(f"image/video_feature.shape: {image_feature.shape}")303 new_image_features.append(image_feature)304 image_features = new_image_features305 else:306 raise ValueError(f"Unexpected mm_patch_merge_type: {self.config.mm_patch_merge_type}")307 else:308 # raise NotImplementedError(f"images.shape={images.shape}, modalities={modalities}")309 image_features = self.encode_image(images)310 311 # TODO: image start / end is not implemented here to support pretraining.312 if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(self.config, "mm_use_im_start_end", False):313 raise NotImplementedError314 # print(f"Total images len(image_features: {len(image_features)}")315 316 # Let's just add dummy tensors if they do not exist,317 # it is a headache to deal with None all the time.318 # But it is not ideal, and if you have a better idea,319 # please open an issue / submit a PR, thanks.320 _labels = labels321 _position_ids = position_ids322 _attention_mask = attention_mask323 if attention_mask is None:324 attention_mask = torch.ones_like(input_ids, dtype=torch.bool)325 else:326 attention_mask = attention_mask.bool()327 if position_ids is None:328 position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)329 if labels is None:330 labels = torch.full_like(input_ids, IGNORE_INDEX)331 332 333 input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]334 labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]335 336 new_input_embeds = []337 new_labels = []338 cur_image_idx = 0339 340 mm_llm_compress = getattr(self.config, "mm_llm_compress", False)341 342 if mm_llm_compress:343 self.model.llm_compress_type = getattr(self.config, "llm_compress_type", "attention")344 self.model.llm_compress_layer_list = getattr(self.config, "llm_compress_layer_list", [8, 16, 24])345 self.model.llm_image_token_ratio_list = getattr(self.config, "llm_image_token_ratio_list", [1.0, 0.5, 0.25, 0.125])346 first_image_token_position = []347 text_prompt_lens = []348 else:349 self.model.llm_compress_type = "attention"350 self.model.llm_compress_layer_list = []351 self.model.llm_image_token_ratio_list = []352 first_image_token_position = []353 text_prompt_lens = []354 355 # rank_print("Inserting Images embedding")356 for batch_idx, cur_input_ids in enumerate(input_ids):357 num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()358 359 if mm_llm_compress:360 ####### copy from pdrop, only support single image/video NOTE ##################361 # record image position for further dropping362 image_index = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist()363 assert len(image_index) == 1, f"Only support singe/video: {image_index}"364 if image_index == []:365 first_image_token_position.append(-1)366 else:367 first_image_token_position.append(image_index[0])368 369 370 # record input instruction length in inference mode371 if not self.training: 372 if image_index == []:373 assert num_images == 0, num_images374 else:375 assert num_images == 1, f"num_images={num_images}"376 text_prompt_lens.append(cur_input_ids.shape[0] - num_images) # consider image place holder377 378 ###############################################379 380 # print(f"num_images={num_images}")381 if num_images == 0:382 cur_image_features = image_features[cur_image_idx]383 cur_input_embeds_1 = self.get_model().embed_tokens(cur_input_ids)384 cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)385 new_input_embeds.append(cur_input_embeds)386 new_labels.append(labels[batch_idx])387 cur_image_idx += 1388 continue389 390 image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]391 cur_input_ids_noim = []392 cur_labels = labels[batch_idx]393 cur_labels_noim = []394 for i in range(len(image_token_indices) - 1):395 cur_input_ids_noim.append(cur_input_ids[image_token_indices[i] + 1 : image_token_indices[i + 1]])396 cur_labels_noim.append(cur_labels[image_token_indices[i] + 1 : image_token_indices[i + 1]])397 split_sizes = [x.shape[0] for x in cur_labels_noim]398 cur_input_embeds = self.get_model().embed_tokens(torch.cat(cur_input_ids_noim))399 cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)400 cur_new_input_embeds = []401 cur_new_labels = []402 403 for i in range(num_images + 1):404 cur_new_input_embeds.append(cur_input_embeds_no_im[i])405 cur_new_labels.append(cur_labels_noim[i])406 if i < num_images:407 try:408 cur_image_features = image_features[cur_image_idx]409 except IndexError:410 print(f"cur_image_idx={cur_image_idx} is not ok")411 cur_image_features = image_features[cur_image_idx - 1]412 cur_image_idx += 1413 cur_new_input_embeds.append(cur_image_features)414 cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))415 416 cur_new_input_embeds = [x.to(self.device) for x in cur_new_input_embeds]417 418 # import pdb; pdb.set_trace()419 cur_new_input_embeds = torch.cat(cur_new_input_embeds)420 cur_new_labels = torch.cat(cur_new_labels)421 422 new_input_embeds.append(cur_new_input_embeds)423 new_labels.append(cur_new_labels)424 425 426 if mm_llm_compress:427 self.model.first_image_token_position = first_image_token_position428 self.model.text_prompt_lens = text_prompt_lens429 self.model.num_image_token_lens = [image_feature.shape[0] for image_feature in image_features]430 431 # Truncate sequences to max length as image embeddings can make the sequence longer432 tokenizer_model_max_length = getattr(self.config, "tokenizer_model_max_length", None)433 # rank_print("Finishing Inserting")434 435 new_input_embeds = [x[:tokenizer_model_max_length] for x, modality in zip(new_input_embeds, modalities)]436 new_labels = [x[:tokenizer_model_max_length] for x, modality in zip(new_labels, modalities)]437 438 # Combine them439 max_len = max(x.shape[0] for x in new_input_embeds)440 batch_size = len(new_input_embeds)441 442 new_input_embeds_padded = []443 new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)444 attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)445 position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)446 # print("Prepare pos id")447 448 for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):449 cur_len = cur_new_embed.shape[0]450 if getattr(self.config, "tokenizer_padding_side", "right") == "left":451 new_input_embeds_padded.append(torch.cat((torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device), cur_new_embed), dim=0))452 if cur_len > 0:453 new_labels_padded[i, -cur_len:] = cur_new_labels454 attention_mask[i, -cur_len:] = True455 position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)456 else:457 new_input_embeds_padded.append(torch.cat((cur_new_embed, torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)), dim=0))458 if cur_len > 0:459 new_labels_padded[i, :cur_len] = cur_new_labels460 attention_mask[i, :cur_len] = True461 position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)462 463 new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)464 # print("tokenizer padding")465 466 if _labels is None:467 new_labels = None468 else:469 new_labels = new_labels_padded470 471 if _attention_mask is None:472 attention_mask = None473 else:474 attention_mask = attention_mask.to(dtype=_attention_mask.dtype)475 476 if _position_ids is None:477 position_ids = None478 if getattr(self.config, "use_pos_skipping", False) and self.training:479 position_ids = torch.arange(new_input_embeds.size(1), device=new_input_embeds.device).unsqueeze(0).to(new_input_embeds.device)480 split_position = random.randint(0, new_input_embeds.size(1))481 left_add = random.randint(0, self.config.pos_skipping_range)482 right_add = random.randint(left_add, self.config.pos_skipping_range)483 position_ids[:, :split_position] += left_add484 position_ids[:, split_position:] += right_add485 # import pdb; pdb.set_trace()486 # print("Finish preparing")487 return None, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels488 489 def initialize_vision_tokenizer(self, model_args, tokenizer):490 if model_args.mm_use_im_patch_token:491 tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)492 self.resize_token_embeddings(len(tokenizer))493 494 if model_args.mm_use_im_start_end:495 num_new_tokens = tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)496 self.resize_token_embeddings(len(tokenizer))497 498 if num_new_tokens > 0:499 input_embeddings = self.get_input_embeddings().weight.data500 output_embeddings = self.get_output_embeddings().weight.data501 502 input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)503 output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)504 505 input_embeddings[-num_new_tokens:] = input_embeddings_avg506 output_embeddings[-num_new_tokens:] = output_embeddings_avg507 508 if model_args.tune_mm_mlp_adapter:509 for p in self.get_input_embeddings().parameters():510 p.requires_grad = True511 for p in self.get_output_embeddings().parameters():512 p.requires_grad = False513 514 if model_args.pretrain_mm_mlp_adapter:515 mm_projector_weights = torch.load(model_args.pretrain_mm_mlp_adapter, map_location="cpu")516 embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]517 assert num_new_tokens == 2518 if input_embeddings.shape == embed_tokens_weight.shape:519 input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]520 elif embed_tokens_weight.shape[0] == num_new_tokens:521 input_embeddings[-num_new_tokens:] = embed_tokens_weight522 else:523 raise ValueError(f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}.")524 elif model_args.mm_use_im_patch_token:525 if model_args.tune_mm_mlp_adapter:526 for p in self.get_input_embeddings().parameters():527 p.requires_grad = False528 for p in self.get_output_embeddings().parameters():529 p.requires_grad = False530 531 532 533class VideoChatFlashQwenConfig(Qwen2Config):534 model_type = "videochat_flash_qwen"535 536 537class VideoChatFlashQwenModel(LlavaMetaModel, Qwen2Model_Flash):538 config_class = VideoChatFlashQwenConfig539 540 def __init__(self, config: VideoChatFlashQwenConfig):541 super(VideoChatFlashQwenModel, self).__init__(config)542 543 544class VideoChatFlashQwenForCausalLM(LlavaMetaForCausalLM, Qwen2ForCausalLM_Flash):545 config_class = VideoChatFlashQwenConfig546 547 def __init__(self, config):548 # super(Qwen2ForCausalLM, self).__init__(config)549 Qwen2ForCausalLM_Flash.__init__(self, config)550 config.model_type = "videochat_flash_qwen"551 # config.rope_scaling = None552 553 self.model = VideoChatFlashQwenModel(config)554 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)555 # Initialize weights and apply final processing556 self.post_init()557 558 def get_model(self):559 return self.model560 561 def forward(562 self,563 input_ids: torch.LongTensor = None,564 attention_mask: Optional[torch.Tensor] = None,565 position_ids: Optional[torch.LongTensor] = None,566 past_key_values: Optional[List[torch.FloatTensor]] = None,567 inputs_embeds: Optional[torch.FloatTensor] = None,568 labels: Optional[torch.LongTensor] = None,569 use_cache: Optional[bool] = None,570 output_attentions: Optional[bool] = None,571 output_hidden_states: Optional[bool] = None,572 images: Optional[torch.FloatTensor] = None,573 image_sizes: Optional[List[List[int]]] = None,574 return_dict: Optional[bool] = None,575 modalities: Optional[List[str]] = ["image"],576 dpo_forward: Optional[bool] = False,577 cache_position=None,578 ) -> Union[Tuple, CausalLMOutputWithPast]:579 580 if inputs_embeds is None:581 (input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels) = self.prepare_inputs_labels_for_multimodal(input_ids, position_ids, attention_mask, past_key_values, labels, images, modalities, image_sizes)582 583 # print("inputs_embeds.shape:", inputs_embeds.shape)584 if dpo_forward:585 raise NotImplementedError586 else:587 return super().forward(588 input_ids=input_ids,589 attention_mask=attention_mask,590 position_ids=position_ids,591 past_key_values=past_key_values,592 inputs_embeds=inputs_embeds,593 labels=labels,594 use_cache=use_cache,595 output_attentions=output_attentions,596 output_hidden_states=output_hidden_states,597 return_dict=return_dict,598 )599 600 @torch.no_grad()601 def generate(602 self,603 inputs: Optional[torch.Tensor] = None,604 images: Optional[torch.Tensor] = None,605 image_sizes: Optional[torch.Tensor] = None,606 modalities: Optional[List[str]] = ["image"],607 **kwargs,608 ) -> Union[GenerateOutput, torch.LongTensor]:609 position_ids = kwargs.pop("position_ids", None)610 attention_mask = kwargs.pop("attention_mask", None)611 if "inputs_embeds" in kwargs:612 raise NotImplementedError("`inputs_embeds` is not supported")613 614 if images is not None:615 (inputs, position_ids, attention_mask, _, inputs_embeds, _) = self.prepare_inputs_labels_for_multimodal(inputs, position_ids, attention_mask, None, None, images, modalities, image_sizes=image_sizes)616 else:617 self.model.image_token_posi = [-1] 618 self.model.prompt_len = None619 self.model.image_tokens = [0]620 inputs_embeds = self.get_model().embed_tokens(inputs)621 622 return super().generate(position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs)623 624 @torch.no_grad()625 def chat(self,626 video_path,627 tokenizer,628 user_prompt,629 chat_history=None,630 return_history=True,631 max_num_frames=512,632 media_dict=None,633 generation_config={}):634 635 frames, time_msg = load_video(video_path, max_num_frames=max_num_frames, media_dict=media_dict)636 637 image_sizes = [frames[0].shape[:2]]638 639 frames = [self.get_vision_tower().image_processor.preprocess(frames, return_tensors="pt")["pixel_values"].to(self.model.dtype).cuda()]640 641 conv = conv_templates["qwen_2"].copy()642 643 if chat_history is None or len(chat_history) == 0:644 user_prompt = f'{DEFAULT_IMAGE_TOKEN}\n{time_msg.strip()} {user_prompt}'645 else:646 assert DEFAULT_IMAGE_TOKEN in chat_history[0]['content'], chat_history647 for msg in chat_history:648 conv.append_message(msg['role'], msg['content'])649 650 conv.append_message(conv.roles[0], user_prompt)651 conv.append_message(conv.roles[1], None)652 653 prompt = conv.get_prompt()654 655 input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).cuda()656 657 if tokenizer.pad_token_id is None:658 if "qwen" in tokenizer.name_or_path.lower():659 print("Setting pad token to bos token for qwen model.")660 tokenizer.pad_token_id = 151643661 662 attention_masks = input_ids.ne(tokenizer.pad_token_id).long().cuda()663 664 stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2665 keywords = [stop_str]666 stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)667 668 with torch.inference_mode():669 output_ids = self.generate(670 inputs=input_ids,671 images=frames,672 attention_mask=attention_masks,673 modalities=["video"],674 image_sizes=image_sizes,675 use_cache=True,676 stopping_criteria=[stopping_criteria],677 **generation_config678 )679 680 outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip()681 if outputs.endswith(stop_str):682 outputs = outputs[: -len(stop_str)]683 684 outputs = outputs.strip()685 686 # print(f"\033[91m== Question: \033[0m\n{prompt}\n")687 # print(f"\033[91m== Response: \033[0m\n{outputs}\n")688 689 if chat_history is None:690 chat_history = []691 692 chat_history.append({"role":conv.roles[0], "content":user_prompt})693 chat_history.append({"role":conv.roles[1], "content":outputs})694 if return_history:695 return outputs, chat_history696 else:697 return outputs698 699 700 701 def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):702 images = kwargs.pop("images", None)703 image_sizes = kwargs.pop("image_sizes", None)704 inputs = super().prepare_inputs_for_generation(input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs)705 if images is not None:706 inputs["images"] = images707 if image_sizes is not None:708 inputs["image_sizes"] = image_sizes709 return inputs710 711 712AutoConfig.register("videochat_flash_qwen", VideoChatFlashQwenConfig)713AutoModelForCausalLM.register(VideoChatFlashQwenConfig, VideoChatFlashQwenForCausalLM)