Spaces:
Running
on
A10G
Running
on
A10G
| import logging | |
| import random | |
| import torch | |
| from torch.cuda.amp import autocast as autocast | |
| import torch.nn as nn | |
| from video_llama.common.registry import registry | |
| from video_llama.models.blip2 import Blip2Base, disabled_train | |
| from video_llama.models.modeling_llama import LlamaForCausalLM | |
| # from video_llama.models.Qformer import BertEncoder | |
| from transformers import LlamaTokenizer,BertConfig | |
| # from transformers.models.bert.modeling_bert import BertEncoder | |
| import einops | |
| import copy | |
| import os | |
| from video_llama.models.Qformer import BertConfig, BertLMHeadModel | |
| # from flamingo_pytorch import PerceiverResampler | |
| class VideoLLAMA(Blip2Base): | |
| """ | |
| BLIP2 GPT-LLAMA model. | |
| """ | |
| PRETRAINED_MODEL_CONFIG_DICT = { | |
| "pretrain_vicuna": "configs/models/video_llama.yaml", | |
| } | |
| def init_video_Qformer(cls, num_query_token, vision_width,num_hidden_layers =2): | |
| encoder_config = BertConfig.from_pretrained("bert-base-uncased") | |
| encoder_config.num_hidden_layers = num_hidden_layers | |
| encoder_config.encoder_width = vision_width | |
| # insert cross-attention layer every other block | |
| encoder_config.add_cross_attention = True | |
| encoder_config.cross_attention_freq = 1 | |
| encoder_config.query_length = num_query_token | |
| Qformer = BertLMHeadModel(config=encoder_config) | |
| query_tokens = nn.Parameter( | |
| torch.zeros(1, num_query_token, encoder_config.hidden_size) | |
| ) | |
| query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range) | |
| return Qformer, query_tokens | |
| def __init__( | |
| self, | |
| vit_model="eva_clip_g", | |
| q_former_model="https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth", | |
| img_size=224, | |
| drop_path_rate=0, | |
| use_grad_checkpoint=False, | |
| vit_precision="fp16", | |
| freeze_vit=True, | |
| freeze_qformer=True, | |
| num_query_token=32, | |
| llama_model="", | |
| prompt_path="", | |
| prompt_template="", | |
| max_txt_len=32, | |
| end_sym='\n', | |
| low_resource=False, # use 8 bit and put vit in cpu | |
| device_8bit=0, # the device of 8bit model should be set when loading and cannot be changed anymore. | |
| frozen_llama_proj=True, | |
| llama_proj_model='', | |
| fusion_header_type= "seqTransf", | |
| max_frame_pos= 32, | |
| fusion_head_layers = 2, | |
| num_video_query_token = 32, | |
| ): | |
| super().__init__() | |
| self.tokenizer = self.init_tokenizer() | |
| self.low_resource = low_resource | |
| print('Loading VIT') | |
| self.visual_encoder, self.ln_vision = self.init_vision_encoder( | |
| vit_model, img_size, drop_path_rate, use_grad_checkpoint, vit_precision | |
| ) | |
| if freeze_vit: | |
| for name, param in self.visual_encoder.named_parameters(): | |
| param.requires_grad = False | |
| self.visual_encoder = self.visual_encoder.eval() | |
| self.visual_encoder.train = disabled_train | |
| for name, param in self.ln_vision.named_parameters(): | |
| param.requires_grad = False | |
| self.ln_vision = self.ln_vision.eval() | |
| self.ln_vision.train = disabled_train | |
| logging.info("freeze vision encoder") | |
| print('Loading VIT Done') | |
| print('Loading Q-Former') | |
| self.Qformer, self.query_tokens = self.init_Qformer( | |
| num_query_token, self.visual_encoder.num_features | |
| ) | |
| self.Qformer.cls = None | |
| self.Qformer.bert.embeddings.word_embeddings = None | |
| self.Qformer.bert.embeddings.position_embeddings = None | |
| for layer in self.Qformer.bert.encoder.layer: | |
| layer.output = None | |
| layer.intermediate = None | |
| self.load_from_pretrained(url_or_filename=q_former_model) | |
| if freeze_qformer: | |
| for name, param in self.Qformer.named_parameters(): | |
| param.requires_grad = False | |
| self.Qformer = self.Qformer.eval() | |
| self.Qformer.train = disabled_train | |
| self.query_tokens.requires_grad = False | |
| logging.info("freeze Qformer") | |
| logging.info('Loading Q-Former Done') | |
| logging.info('Loading LLAMA Tokenizer') | |
| self.llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model, use_fast=False, use_auth_token=os.environ["API_TOKEN"]) | |
| if self.llama_tokenizer.pad_token is None: | |
| self.llama_tokenizer.pad_token = self.llama_tokenizer.eos_token | |
| DEFAULT_IMAGE_PATCH_TOKEN = '<ImageHere>' | |
| self.llama_tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) | |
| self.IMAGE_PATCH_TOKEN_ID = self.llama_tokenizer.get_vocab()[DEFAULT_IMAGE_PATCH_TOKEN] | |
| logging.info('Loading LLAMA Model') | |
| if self.low_resource: | |
| self.llama_model = LlamaForCausalLM.from_pretrained( | |
| llama_model, | |
| torch_dtype=torch.float16, | |
| load_in_8bit=True, | |
| device_map={'': device_8bit}, | |
| use_auth_token=os.environ["API_TOKEN"] | |
| ) | |
| else: | |
| self.llama_model = LlamaForCausalLM.from_pretrained( | |
| llama_model, | |
| torch_dtype=torch.float16,use_auth_token=os.environ["API_TOKEN"] | |
| ) | |
| for name, param in self.llama_model.named_parameters(): | |
| param.requires_grad = False | |
| logging.info('Loading LLAMA Done') | |
| logging.info('Loading LLAMA proj') | |
| self.llama_proj = nn.Linear( | |
| self.Qformer.config.hidden_size, self.llama_model.config.hidden_size | |
| ) | |
| if llama_proj_model: | |
| print("load llama proj weight: {}".format(llama_proj_model)) | |
| llama_proj_weight = torch.load(llama_proj_model, map_location="cpu") | |
| msg = model.load_state_dict(llama_proj_weight['model'], strict=False) | |
| if frozen_llama_proj: | |
| # todo frozen llama_proj | |
| for name, param in self.llama_proj.named_parameters(): | |
| param.requires_grad = False | |
| logging.info('LLAMA proj is frozen') | |
| else: | |
| for name, param in self.llama_proj.named_parameters(): | |
| param.requires_grad = True | |
| logging.info('LLAMA proj is not frozen') | |
| logging.info('Loading llama_proj Done') | |
| self.max_txt_len = max_txt_len | |
| self.end_sym = end_sym | |
| if prompt_path: | |
| with open(prompt_path, 'r') as f: | |
| raw_prompts = f.read().splitlines() | |
| filted_prompts = [raw_prompt for raw_prompt in raw_prompts if "<ImageHere>" in raw_prompt] | |
| self.prompt_list = [prompt_template.format(p) for p in filted_prompts] | |
| print('Load {} training prompts'.format(len(self.prompt_list))) | |
| print('Prompt Example \n{}'.format(random.choice(self.prompt_list))) | |
| else: | |
| self.prompt_list = [] | |
| self.video_frame_position_embedding = nn.Embedding(max_frame_pos, self.Qformer.config.hidden_size) | |
| self.num_video_query_token = num_video_query_token | |
| self.video_Qformer,self.video_query_tokens = self.init_video_Qformer(num_query_token = num_video_query_token,\ | |
| vision_width=self.Qformer.config.hidden_size, num_hidden_layers =2) | |
| self.video_Qformer.cls = None | |
| self.video_Qformer.bert.embeddings.word_embeddings = None | |
| self.video_Qformer.bert.embeddings.position_embeddings = None | |
| for layer in self.video_Qformer.bert.encoder.layer: | |
| layer.output = None | |
| layer.intermediate = None | |
| def vit_to_cpu(self): | |
| self.ln_vision.to("cpu") | |
| self.ln_vision.float() | |
| self.visual_encoder.to("cpu") | |
| self.visual_encoder.float() | |
| def encode_img(self, image): | |
| device = image.device | |
| # if self.low_resource: | |
| # self.vit_to_cpu() | |
| # image = image.to("cpu") | |
| # input shape b,c,t,h,w | |
| batch_size,_,time_length,_,_ = image.size() | |
| image = einops.rearrange(image, 'b c t h w -> (b t) c h w') | |
| with self.maybe_autocast(): | |
| # embed image features with blip2, out: (b t) q h | |
| image_embeds = self.ln_vision(self.visual_encoder(image)).to(device) | |
| image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device) | |
| query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1) | |
| query_output = self.Qformer.bert( | |
| query_embeds=query_tokens, | |
| encoder_hidden_states=image_embeds, | |
| encoder_attention_mask=image_atts, | |
| return_dict=True, | |
| ) | |
| # add frame_pos embedding | |
| position_ids = torch.arange(time_length, dtype=torch.long, device=query_tokens.device) | |
| position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) | |
| frame_position_embeddings = self.video_frame_position_embedding(position_ids) | |
| q_hidden_state = query_output.last_hidden_state | |
| frame_position_embeddings = frame_position_embeddings.unsqueeze(-2) | |
| frame_hidden_state = einops.rearrange(q_hidden_state, '(b t) q h -> b t q h',b=batch_size,t=time_length) | |
| frame_hidden_state = frame_position_embeddings + frame_hidden_state | |
| # frame attention | |
| frame_hidden_state = einops.rearrange(frame_hidden_state, 'b t q h -> b (t q) h',b=batch_size,t=time_length) | |
| frame_atts = torch.ones(frame_hidden_state.size()[:-1], dtype=torch.long).to(device) | |
| video_query_tokens = self.video_query_tokens.expand(frame_hidden_state.shape[0], -1, -1) | |
| # print('attention') | |
| # print(video_query_tokens.size()) | |
| # print(frame_hidden_state.size()) | |
| video_query_output = self.video_Qformer.bert( | |
| query_embeds=video_query_tokens, | |
| encoder_hidden_states=frame_hidden_state, | |
| encoder_attention_mask=frame_atts, | |
| return_dict=True, | |
| ) | |
| video_hidden = video_query_output.last_hidden_state | |
| inputs_llama = self.llama_proj(video_hidden) | |
| atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(image_embeds.device) | |
| return inputs_llama, atts_llama | |
| def prompt_wrap(self, img_embeds, atts_img, prompt): | |
| if prompt: | |
| batch_size = img_embeds.shape[0] | |
| # print(prompt) | |
| p_before, p_after = prompt.split('<ImageHere>') | |
| p_before_tokens = self.llama_tokenizer( | |
| p_before, return_tensors="pt", add_special_tokens=False).to(img_embeds.device) | |
| p_after_tokens = self.llama_tokenizer( | |
| p_after, return_tensors="pt", add_special_tokens=False).to(img_embeds.device) | |
| p_before_embeds = self.llama_model.model.embed_tokens(p_before_tokens.input_ids).expand(batch_size, -1, -1) | |
| p_after_embeds = self.llama_model.model.embed_tokens(p_after_tokens.input_ids).expand(batch_size, -1, -1) | |
| wrapped_img_embeds = torch.cat([p_before_embeds, img_embeds, p_after_embeds], dim=1) | |
| wrapped_atts_img = atts_img[:, :1].expand(-1, wrapped_img_embeds.shape[1]) | |
| return wrapped_img_embeds, wrapped_atts_img | |
| else: | |
| return img_embeds, atts_img | |
| def forward(self, samples): | |
| if 'conv_type' in samples.keys() and samples['conv_type']=='multi': | |
| num_patch_tokens = self.num_video_query_token | |
| im_patch_token_id = self.IMAGE_PATCH_TOKEN_ID | |
| image = samples["images"] | |
| input_ids = samples['input_ids'] | |
| if len(image.size())==4: | |
| time = 1 | |
| image = einops.repeat(image, 'b c h w -> b c t h w',t = time) | |
| img_embeds, atts_img = self.encode_img(image) | |
| temp_input_ids = copy.deepcopy(input_ids) | |
| temp_input_ids[temp_input_ids == im_patch_token_id] = 0 | |
| temp_input_embedding = self.llama_model.model.embed_tokens(temp_input_ids) | |
| new_input_embeds=[] | |
| cur_image_idx = 0 | |
| for cur_input_ids, cur_input_embeds in zip(input_ids, temp_input_embedding): | |
| cur_image_features = img_embeds[cur_image_idx] | |
| if (cur_input_ids == im_patch_token_id).sum() != num_patch_tokens: | |
| raise ValueError("The number of image patch tokens should be the same as the number of image patches.") | |
| masked_indices = torch.where(cur_input_ids == im_patch_token_id)[0] | |
| mask_index_start = masked_indices[0] | |
| if (masked_indices != torch.arange(mask_index_start, mask_index_start+num_patch_tokens, device=masked_indices.device, dtype=masked_indices.dtype)).any(): | |
| raise ValueError("The image patch tokens should be consecutive.") | |
| cur_new_input_embeds = torch.cat((cur_input_embeds[:mask_index_start], cur_image_features, cur_input_embeds[mask_index_start+num_patch_tokens:]), dim=0) | |
| new_input_embeds.append(cur_new_input_embeds) | |
| cur_image_idx+=1 | |
| inputs_embeds = torch.stack(new_input_embeds, dim=0) | |
| targets = samples['labels'] | |
| attention_mask = samples['attention_mask'] | |
| with self.maybe_autocast(): | |
| outputs = self.llama_model( | |
| inputs_embeds=inputs_embeds, | |
| attention_mask=attention_mask, | |
| return_dict=True, | |
| labels=targets, | |
| ) | |
| loss = outputs.loss | |
| return {"loss": loss} | |
| else: | |
| image = samples["image"] | |
| if len(image.size()) != 5: | |
| time = 1 | |
| image = einops.repeat(image, 'b c h w -> b c t h w',t = time) | |
| img_embeds, atts_img = self.encode_img(image) | |
| if self.prompt_list: | |
| prompt = random.choice(self.prompt_list) | |
| img_embeds, atts_img = self.prompt_wrap(img_embeds, atts_img, prompt) | |
| self.llama_tokenizer.padding_side = "right" | |
| text = [t + self.end_sym for t in samples["text_input"]] | |
| to_regress_tokens = self.llama_tokenizer( | |
| text, | |
| return_tensors="pt", | |
| padding="longest", | |
| truncation=True, | |
| max_length=self.max_txt_len, | |
| add_special_tokens=False | |
| ).to(image.device) | |
| targets = to_regress_tokens.input_ids.masked_fill( | |
| to_regress_tokens.input_ids == self.llama_tokenizer.pad_token_id, -100 | |
| ) | |
| empty_targets = ( | |
| torch.ones([atts_img.shape[0], atts_img.shape[1]+1], | |
| dtype=torch.long).to(image.device).fill_(-100) # plus one for bos | |
| ) | |
| targets = torch.cat([empty_targets, targets], dim=1) | |
| batch_size = img_embeds.shape[0] | |
| bos = torch.ones([batch_size, 1], | |
| dtype=to_regress_tokens.input_ids.dtype, | |
| device=to_regress_tokens.input_ids.device) * self.llama_tokenizer.bos_token_id | |
| bos_embeds = self.llama_model.model.embed_tokens(bos) | |
| atts_bos = atts_img[:, :1] | |
| to_regress_embeds = self.llama_model.model.embed_tokens(to_regress_tokens.input_ids) | |
| inputs_embeds = torch.cat([bos_embeds, img_embeds, to_regress_embeds], dim=1) | |
| attention_mask = torch.cat([atts_bos, atts_img, to_regress_tokens.attention_mask], dim=1) | |
| with self.maybe_autocast(): | |
| outputs = self.llama_model( | |
| inputs_embeds=inputs_embeds, | |
| attention_mask=attention_mask, | |
| return_dict=True, | |
| labels=targets, | |
| ) | |
| loss = outputs.loss | |
| return {"loss": loss} | |
| def from_config(cls, cfg): | |
| vit_model = cfg.get("vit_model", "eva_clip_g") | |
| q_former_model = cfg.get("q_former_model", "https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth") | |
| img_size = cfg.get("image_size") | |
| num_query_token = cfg.get("num_query_token") | |
| llama_model = cfg.get("llama_model") | |
| drop_path_rate = cfg.get("drop_path_rate", 0) | |
| use_grad_checkpoint = cfg.get("use_grad_checkpoint", False) | |
| vit_precision = cfg.get("vit_precision", "fp16") | |
| freeze_vit = cfg.get("freeze_vit", True) | |
| freeze_qformer = cfg.get("freeze_qformer", True) | |
| low_resource = cfg.get("low_resource", False) | |
| device_8bit = cfg.get("device_8bit", 0) | |
| prompt_path = cfg.get("prompt_path", "") | |
| prompt_template = cfg.get("prompt_template", "") | |
| max_txt_len = cfg.get("max_txt_len", 32) | |
| end_sym = cfg.get("end_sym", '\n') | |
| frozen_llama_proj = cfg.get("frozen_llama_proj", True) | |
| llama_proj_model = cfg.get("llama_proj_model", '') | |
| fusion_header_type = cfg.get("fusion_header_type", 'seqTransf') | |
| max_frame_pos = cfg.get("max_frame_pos", 32) | |
| fusion_head_layers = cfg.get("fusion_head_layers", 2) | |
| num_video_query_token = cfg.get("num_video_query_token", 32) | |
| model = cls( | |
| vit_model=vit_model, | |
| q_former_model=q_former_model, | |
| img_size=img_size, | |
| drop_path_rate=drop_path_rate, | |
| use_grad_checkpoint=use_grad_checkpoint, | |
| vit_precision=vit_precision, | |
| freeze_vit=freeze_vit, | |
| freeze_qformer=freeze_qformer, | |
| num_query_token=num_query_token, | |
| llama_model=llama_model, | |
| prompt_path=prompt_path, | |
| prompt_template=prompt_template, | |
| max_txt_len=max_txt_len, | |
| end_sym=end_sym, | |
| low_resource=low_resource, | |
| device_8bit=device_8bit, | |
| fusion_header_type=fusion_header_type, | |
| max_frame_pos=max_frame_pos, | |
| fusion_head_layers=fusion_head_layers, | |
| frozen_llama_proj=frozen_llama_proj, | |
| num_video_query_token=num_video_query_token | |
| ) | |
| ckpt_path = cfg.get("ckpt", "") # load weights of MiniGPT-4 | |
| if ckpt_path: | |
| print("Load BLIP2-LLM Checkpoint: {}".format(ckpt_path)) | |
| ckpt = torch.load(ckpt_path, map_location="cpu") | |
| msg = model.load_state_dict(ckpt['model'], strict=False) | |
| return model | |