|
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 |
|
|
|
|
|
@registry.register_model("minigpt_v2") |
|
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='[INST] {} [/INST]', |
|
max_txt_len=300, |
|
end_sym='\n', |
|
lora_r=64, |
|
lora_target_modules=["q_proj", "v_proj"], |
|
lora_alpha=16, |
|
lora_dropout=0.05, |
|
chat_template=False, |
|
use_grad_checkpoint_llm=False, |
|
max_context_len=3800, |
|
low_resource=False, |
|
device_8bit=0, |
|
): |
|
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, |
|
) |
|
|
|
img_f_dim = self.visual_encoder.num_features * 4 |
|
self.llama_proj = nn.Linear( |
|
img_f_dim, self.llama_model.config.hidden_size |
|
) |
|
self.chat_template = chat_template |
|
|
|
if use_grad_checkpoint_llm: |
|
self.llama_model.gradient_checkpointing_enable() |
|
|
|
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) |
|
return inputs_llama, atts_llama |
|
|
|
@classmethod |
|
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", "") |
|
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 |
|
|