Spaces:
Runtime error
Runtime error
import logging | |
import random | |
import torch | |
from torch.cuda.amp import autocast as autocast | |
import torch.nn as nn | |
from minigpt4.common.registry import registry | |
from minigpt4.models.base_model import disabled_train | |
from minigpt4.models.minigpt_base import MiniGPTBase | |
from minigpt4.models.Qformer import BertConfig, BertLMHeadModel | |
class MiniGPTv2(MiniGPTBase): | |
""" | |
MiniGPT-v2 model | |
""" | |
PRETRAINED_MODEL_CONFIG_DICT = { | |
"pretrain": "configs/models/minigpt_v2.yaml", | |
} | |
def __init__( | |
self, | |
vit_model="eva_clip_g", | |
img_size=448, | |
drop_path_rate=0, | |
use_grad_checkpoint=False, | |
vit_precision="fp16", | |
freeze_vit=True, | |
llama_model="", | |
prompt_template='###Human: {} ###Assistant: ', | |
max_txt_len=300, | |
end_sym='\n', | |
lora_r=64, | |
lora_target_modules=['query_key_value','dense'], | |
lora_alpha=16, | |
lora_dropout=0.05, | |
chat_template=False, | |
use_grad_checkpoint_llm=False, | |
max_context_len=3800, | |
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. | |
): | |
super().__init__( | |
vit_model=vit_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, | |
llama_model=llama_model, | |
max_txt_len=max_txt_len, | |
max_context_len=max_context_len, | |
end_sym=end_sym, | |
prompt_template=prompt_template, | |
low_resource=low_resource, | |
device_8bit=device_8bit, | |
lora_r=lora_r, | |
lora_target_modules=lora_target_modules, | |
lora_alpha=lora_alpha, | |
lora_dropout=lora_dropout, | |
) | |
print('Loading Q-Former') | |
self.Qformer, self.query_tokens = self.init_Qformer( | |
num_query_token = 32, vision_width = self.visual_encoder.num_features, freeze = False | |
) | |
self.load_from_pretrained(url_or_filename="https://storage.googleapis.com/sfr-vision-language-research/LAVIS/models/BLIP2/blip2_pretrained_flant5xxl.pth") # load q-former weights here | |
img_f_dim = self.Qformer.config.hidden_size | |
print('Loading Q-Former Done') | |
# img_f_dim = self.visual_encoder.num_features * 4 | |
self.llama_proj = nn.Linear( | |
self.Qformer.config.hidden_size, 4096 | |
) | |
self.llama_proj2 = nn.Linear( | |
4096, self.llama_model.config.hidden_size | |
) | |
self.chat_template = chat_template | |
if use_grad_checkpoint_llm: | |
self.llama_model.gradient_checkpointing_enable() | |
def init_Qformer(cls, num_query_token, vision_width, freeze): | |
encoder_config = BertConfig.from_pretrained("bert-base-uncased") | |
encoder_config.encoder_width = vision_width | |
# insert cross-attention layer every other block | |
encoder_config.add_cross_attention = True | |
encoder_config.cross_attention_freq = 2 | |
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) | |
Qformer.cls = None | |
Qformer.bert.embeddings.word_embeddings = None | |
Qformer.bert.embeddings.position_embeddings = None | |
for layer in Qformer.bert.encoder.layer: | |
layer.output = None | |
layer.intermediate = None | |
if freeze: | |
for name, param in Qformer.named_parameters(): | |
param.requires_grad = False | |
Qformer = Qformer.eval() | |
Qformer.train = disabled_train | |
query_tokens.requires_grad = False | |
logging.info("freeze Qformer") | |
return Qformer, query_tokens | |
def encode_img(self, image): | |
device = image.device | |
if len(image.shape) > 4: | |
image = image.reshape(-1, *image.shape[-3:]) | |
with self.maybe_autocast(): | |
image_embeds = self.ln_vision(self.visual_encoder(image)).to(device) | |
# image_embeds = image_embeds[:, 1:, :] | |
# bs, pn, hs = image_embeds.shape | |
# image_embeds = image_embeds.view(bs, int(pn / 4), int(hs * 4)) | |
# inputs_llama = self.llama_proj(image_embeds) | |
# atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(image.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, | |
) | |
inputs_llama = self.llama_proj(query_output.last_hidden_state) | |
inputs_llama = self.llama_proj2(inputs_llama) | |
atts_llama = torch.ones(inputs_llama.size()[:-1], dtype=torch.long).to(image.device) | |
return inputs_llama, atts_llama | |
def from_config(cls, cfg): | |
vit_model = cfg.get("vit_model", "eva_clip_g") | |
img_size = cfg.get("image_size") | |
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) | |
low_resource = cfg.get("low_resource", False) | |
prompt_template = cfg.get("prompt_template", '[INST] {} [/INST]') | |
max_txt_len = cfg.get("max_txt_len", 300) | |
end_sym = cfg.get("end_sym", '\n') | |
lora_r = cfg.get("lora_r", 64) | |
lora_alpha = cfg.get("lora_alpha", 16) | |
chat_template = cfg.get("chat_template", False) | |
use_grad_checkpoint_llm = cfg.get("use_grad_checkpoint_llm", False) | |
max_context_len = cfg.get("max_context_len", 3800) | |
model = cls( | |
vit_model=vit_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, | |
llama_model=llama_model, | |
prompt_template=prompt_template, | |
max_txt_len=max_txt_len, | |
low_resource=low_resource, | |
end_sym=end_sym, | |
lora_r=lora_r, | |
lora_alpha=lora_alpha, | |
chat_template=chat_template, | |
use_grad_checkpoint_llm=use_grad_checkpoint_llm, | |
max_context_len=max_context_len, | |
) | |
ckpt_path = cfg.get("ckpt", "") # load weights of MiniGPT-4 | |
if ckpt_path: | |
print("Load Minigpt-4-LLM Checkpoint: {}".format(ckpt_path)) | |
ckpt = torch.load(ckpt_path, map_location="cpu") | |
msg = model.load_state_dict(ckpt['model'], strict=False) | |
return model | |