Spaces:
Paused
Paused
import gc | |
import math | |
import timm | |
import torch | |
from torch import Tensor | |
import torch.nn as nn | |
from torch.nn import CrossEntropyLoss | |
from typing import List, Optional, Tuple, Union | |
from transformers import AutoConfig, AutoModelForCausalLM | |
from transformers import MistralForCausalLM, MistralModel, MistralConfig | |
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast | |
from omnilmm.model.utils import build_transform | |
from omnilmm.model.resampler import Resampler | |
DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>" | |
DEFAULT_IM_START_TOKEN = "<im_start>" | |
DEFAULT_IM_END_TOKEN = "<im_end>" | |
class OmniLMMConfig(MistralConfig): | |
model_type = "omnilmm" | |
class Identity(torch.nn.Identity): | |
def forward(self, input: Tensor, **kwargs) -> Tensor: | |
return super().forward(input) | |
def create_vision_module(config): | |
vision_tower = timm.create_model('eva02_enormous_patch14_clip_224.laion2b_plus', | |
pretrained=False, | |
num_classes=0, | |
dynamic_img_size=True, | |
dynamic_img_pad=True) | |
if isinstance(vision_tower, timm.models.VisionTransformer): | |
if vision_tower.attn_pool is not None: | |
vision_tower.attn_pool = Identity() | |
# use 2nd last layer's output | |
vision_tower.blocks[-1] = Identity() | |
embed_dim = config.hidden_size | |
resampler = Resampler( | |
grid_size=int(math.sqrt(config.num_query)), | |
embed_dim=embed_dim, | |
num_heads=embed_dim // 128, | |
kv_dim=vision_tower.embed_dim, | |
) | |
return vision_tower, resampler | |
class OmniLMMModel(MistralModel): | |
config_class = OmniLMMConfig | |
def __init__(self, config: OmniLMMConfig, mm_vision_tower=None, mm_hidden_size=None, tune_clip=True): | |
super(OmniLMMModel, self).__init__(config) | |
if hasattr(config, "mm_vision_tower"): | |
vision_tower, resampler = create_vision_module(config) | |
# print(__file__, 'skip loading vision tower weights') | |
# HACK: for FSDP | |
self.vision_tower = [vision_tower] | |
self.resampler = resampler | |
if tune_clip: | |
self.vision_tower = self.vision_tower[0] | |
self.vision_config = lambda x: None | |
def initialize_vision_modules(self, vision_tower, no_randaug, num_query, image_size, tune_clip=False): | |
self.config.mm_vision_tower = vision_tower | |
self.config.use_mm_proj = True | |
self.config.num_query = num_query | |
self.config.image_size = image_size | |
if not hasattr(self, 'vision_tower'): | |
vision_tower, resampler = create_vision_module(self.config) | |
state_dict = torch.load( | |
'/tt/data/public/multimodal/multimodal_model_ckpts/timm/eva02_enormous_patch14_clip_224.laion2b_plus.pt') | |
vision_tower.load_state_dict(state_dict, strict=False) | |
del state_dict | |
gc.collect() | |
else: | |
if isinstance(self.vision_tower, list): | |
vision_tower = self.vision_tower[0] | |
else: | |
vision_tower = self.vision_tower | |
resampler = self.resampler | |
self.vision_tower = vision_tower if tune_clip else [vision_tower] | |
self.resampler = resampler | |
train_img_transform = build_transform( | |
is_train=True, randaug=not no_randaug, input_size=self.config.image_size, std_mode='OPENAI_CLIP') | |
eval_img_transform = build_transform( | |
is_train=False, input_size=self.config.image_size, std_mode='OPENAI_CLIP') | |
return dict( | |
image_processor=(train_img_transform, eval_img_transform), | |
image_token_len=num_query, | |
vision_config=self.vision_config | |
) | |
def get_vision_embedding(self, pixel_values): | |
if isinstance(self.vision_tower, list): | |
vision_tower = self.vision_tower[0] # HACK: for FSDP | |
else: | |
vision_tower = self.vision_tower | |
dtype = vision_tower.pos_embed.data.dtype | |
vision_embedding = vision_tower.forward_features( | |
pixel_values.type(dtype)) | |
if hasattr(vision_tower, 'num_prefix_tokens') and vision_tower.num_prefix_tokens > 0: | |
vision_embedding = vision_embedding[:, | |
vision_tower.num_prefix_tokens:] | |
res = self.resampler(vision_embedding) | |
return res | |
def get_vllm_embedding(self, data): | |
if 'vision_hidden_states' not in data: | |
pixel_values_list = data['pixel_values'] | |
vision_hidden_states = [] | |
for pixel_values in pixel_values_list: | |
if len(pixel_values) > 0: | |
vision_hidden_states.append(self.get_vision_embedding(pixel_values.unsqueeze(0))[0]) | |
else: | |
vision_hidden_states.append([]) | |
else: | |
vision_hidden_states = data['vision_hidden_states'] | |
#vllm_embedding = self.llm.model.embed_tokens(data['input_ids']) * self.llm.config.scale_emb | |
inputs_embeds = self.embed_tokens(data['input_ids']) | |
vision_hidden_states = [i.type(inputs_embeds.dtype) | |
if isinstance(i, torch.Tensor) else i for i in vision_hidden_states | |
] | |
# HACK: replace back original embeddings for LLaVA pretraining | |
orig_embeds_params = getattr(self, 'orig_embeds_params', None) | |
new_input_embeds = [] | |
cur_image_idx = 0 | |
for cur_input_ids, cur_input_embeds in zip(data['input_ids'], inputs_embeds): | |
if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0: | |
# multimodal LLM, but the current sample is not multimodal | |
cur_input_embeds = cur_input_embeds + (0. * dummy_image_features).sum() | |
new_input_embeds.append(cur_input_embeds) | |
continue | |
if self.vision_config.use_im_start_end: | |
cur_image_features = vision_hidden_states[cur_image_idx] | |
num_patches = cur_image_features.shape[0] | |
if (cur_input_ids == self.vision_config.im_start_token).sum() != (cur_input_ids == self.vision_config.im_end_token).sum(): | |
raise ValueError( | |
"The number of image start tokens and image end tokens should be the same.") | |
image_start_tokens = torch.where( | |
cur_input_ids == self.vision_config.im_start_token)[0] | |
for image_start_token_pos in image_start_tokens: | |
cur_image_features = vision_hidden_states[cur_image_idx].to( | |
device=cur_input_embeds.device) | |
num_patches = cur_image_features.shape[0] | |
if cur_input_ids[image_start_token_pos + num_patches + 1] != self.vision_config.im_end_token: | |
raise ValueError( | |
"The image end token should follow the image start token.") | |
if orig_embeds_params is not None: | |
cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos].detach(), cur_input_embeds[image_start_token_pos:image_start_token_pos+1], cur_image_features, | |
cur_input_embeds[image_start_token_pos + num_patches + 1:image_start_token_pos + num_patches + 2], cur_input_embeds[image_start_token_pos + num_patches + 2:].detach()), dim=0) | |
else: | |
cur_new_input_embeds = torch.cat( | |
(cur_input_embeds[:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:]), dim=0) | |
cur_image_idx += 1 | |
new_input_embeds.append(cur_new_input_embeds) | |
else: | |
raise NotImplementedError | |
inputs_embeds = torch.stack(new_input_embeds, dim=0) | |
return inputs_embeds, vision_hidden_states | |
def forward( | |
self, | |
input_ids: torch.LongTensor = None, | |
attention_mask: Optional[torch.Tensor] = None, | |
past_key_values: Optional[List[torch.FloatTensor]] = None, | |
inputs_embeds: Optional[torch.FloatTensor] = None, | |
use_cache: Optional[bool] = None, | |
output_attentions: Optional[bool] = None, | |
output_hidden_states: Optional[bool] = None, | |
images: Optional[torch.FloatTensor] = None, | |
return_dict: Optional[bool] = None, | |
**kwargs | |
) -> Union[Tuple, BaseModelOutputWithPast]: | |
# HACK: replace back original embeddings for LLaVA pretraining | |
orig_embeds_params = getattr(self, 'orig_embeds_params', None) | |
if inputs_embeds is None and past_key_values is None: | |
inputs_embeds = self.embed_tokens(input_ids) | |
vision_tower = getattr(self, 'vision_tower', None) | |
if vision_tower is not None and (input_ids.shape[1] != 1 or self.training) and images is not None: | |
if type(images) is list: | |
image_features = [] | |
for image in images: | |
image_forward_out = self.get_vision_embedding(image.unsqueeze(0))[ | |
0] | |
image_features.append(image_forward_out) | |
else: | |
image_features = self.get_vision_embedding(images) | |
dummy_image_features = torch.zeros( | |
self.config.num_query, | |
self.config.hidden_size, | |
device=inputs_embeds.device, | |
dtype=inputs_embeds.dtype) | |
new_input_embeds = [] | |
cur_image_idx = 0 | |
for cur_input_ids, cur_input_embeds in zip(input_ids, inputs_embeds): | |
if (cur_input_ids == self.vision_config.im_patch_token).sum() == 0: | |
# multimodal LLM, but the current sample is not multimodal | |
cur_input_embeds = cur_input_embeds + \ | |
(0. * dummy_image_features).sum() | |
new_input_embeds.append(cur_input_embeds) | |
continue | |
if self.vision_config.use_im_start_end: | |
cur_image_features = image_features[cur_image_idx] | |
num_patches = cur_image_features.shape[0] | |
if (cur_input_ids == self.vision_config.im_start_token).sum() != (cur_input_ids == self.vision_config.im_end_token).sum(): | |
raise ValueError( | |
"The number of image start tokens and image end tokens should be the same.") | |
image_start_tokens = torch.where( | |
cur_input_ids == self.vision_config.im_start_token)[0] | |
for image_start_token_pos in image_start_tokens: | |
cur_image_features = image_features[cur_image_idx].to( | |
device=cur_input_embeds.device) | |
num_patches = cur_image_features.shape[0] | |
if cur_input_ids[image_start_token_pos + num_patches + 1] != self.vision_config.im_end_token: | |
raise ValueError( | |
"The image end token should follow the image start token.") | |
if orig_embeds_params is not None: | |
cur_new_input_embeds = torch.cat((cur_input_embeds[:image_start_token_pos].detach(), cur_input_embeds[image_start_token_pos:image_start_token_pos+1], cur_image_features, | |
cur_input_embeds[image_start_token_pos + num_patches + 1:image_start_token_pos + num_patches + 2], cur_input_embeds[image_start_token_pos + num_patches + 2:].detach()), dim=0) | |
else: | |
cur_new_input_embeds = torch.cat( | |
(cur_input_embeds[:image_start_token_pos+1], cur_image_features, cur_input_embeds[image_start_token_pos + num_patches + 1:]), dim=0) | |
cur_image_idx += 1 | |
new_input_embeds.append(cur_new_input_embeds) | |
else: | |
raise NotImplementedError | |
inputs_embeds = torch.stack(new_input_embeds, dim=0) | |
input_ids = None | |
return super(OmniLMMModel, self).forward( | |
input_ids=input_ids, attention_mask=attention_mask, past_key_values=past_key_values, | |
inputs_embeds=inputs_embeds, use_cache=use_cache, | |
output_attentions=output_attentions, output_hidden_states=output_hidden_states, | |
return_dict=return_dict, | |
**kwargs | |
) | |
class OmniLMMForCausalLM(MistralForCausalLM): | |
config_class = OmniLMMConfig | |
def __init__(self, config, mm_vision_tower=None, tune_clip=True): | |
super(MistralForCausalLM, self).__init__(config) | |
self.model = OmniLMMModel( | |
config, mm_vision_tower=mm_vision_tower, tune_clip=tune_clip) | |
self.lm_head = nn.Linear( | |
config.hidden_size, config.vocab_size, bias=False) | |
# Initialize weights and apply final processing | |
self.post_init() | |
def forward( | |
self, | |
input_ids: torch.LongTensor = None, | |
attention_mask: Optional[torch.Tensor] = None, | |
past_key_values: Optional[List[torch.FloatTensor]] = None, | |
inputs_embeds: Optional[torch.FloatTensor] = None, | |
labels: Optional[torch.LongTensor] = None, | |
use_cache: Optional[bool] = None, | |
output_attentions: Optional[bool] = None, | |
output_hidden_states: Optional[bool] = None, | |
images: Optional[torch.FloatTensor] = None, | |
return_dict: Optional[bool] = None, | |
**kwargs | |
) -> Union[Tuple, CausalLMOutputWithPast]: | |
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions | |
output_hidden_states = ( | |
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states | |
) | |
return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
# print(f'@@@ At forward, labels: {labels.shape}-{labels}', flush=True) | |
# print(f'@@@ At forward, input_ids: {input_ids.shape}-{input_ids}', flush=True) | |
# print(f'@@@ At forward, input_ids: {attention_mask.shape}-{attention_mask}', flush=True) | |
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) | |
outputs = self.model( | |
input_ids=input_ids, | |
attention_mask=attention_mask, | |
past_key_values=past_key_values, | |
inputs_embeds=inputs_embeds, | |
use_cache=use_cache, | |
output_attentions=output_attentions, | |
output_hidden_states=output_hidden_states, | |
return_dict=return_dict, | |
images=images, | |
**kwargs | |
) | |
hidden_states = outputs[0] | |
logits = self.lm_head(hidden_states) | |
loss = None | |
if labels is not None: | |
# Shift so that tokens < n predict n | |
shift_logits = logits[..., :-1, :].contiguous() | |
shift_labels = labels[..., 1:].contiguous() | |
# Flatten the tokens | |
loss_fct = CrossEntropyLoss() | |
shift_logits = shift_logits.view(-1, self.config.vocab_size) | |
shift_labels = shift_labels.view(-1) | |
# Enable model/pipeline parallelism | |
shift_labels = shift_labels.to(shift_logits.device) | |
loss = loss_fct(shift_logits, shift_labels) | |
if not return_dict: | |
output = (logits,) + outputs[1:] | |
return (loss,) + output if loss is not None else output | |
return CausalLMOutputWithPast( | |
loss=loss, | |
logits=logits, | |
past_key_values=outputs.past_key_values, | |
hidden_states=outputs.hidden_states, | |
attentions=outputs.attentions, | |
) | |
# TODO could be removed for generate_vllm() | |
def prepare_inputs_for_generation( | |
self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs | |
): | |
if past_key_values: | |
input_ids = input_ids[:, -1:] | |
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step | |
if inputs_embeds is not None and past_key_values is None: | |
model_inputs = {"inputs_embeds": inputs_embeds} | |
else: | |
model_inputs = {"input_ids": input_ids} | |
model_inputs.update( | |
{ | |
"past_key_values": past_key_values, | |
"use_cache": kwargs.get("use_cache"), | |
"attention_mask": attention_mask, | |
"images": kwargs.get("images", None), | |
} | |
) | |
return model_inputs | |
def generate_vllm( | |
self, | |
input_ids: torch.LongTensor = None, | |
images: Optional[torch.FloatTensor] = None, | |
vision_hidden_states=None, | |
return_vision_hidden_states=False, | |
**kwargs | |
): | |
model_inputs = {'input_ids': input_ids} | |
if vision_hidden_states is None: | |
model_inputs['pixel_values'] = images | |
else: | |
model_inputs['vision_hidden_states'] = vision_hidden_states | |
with torch.inference_mode(): | |
inputs_embeds, vision_hidden_states = self.model.get_vllm_embedding(model_inputs) | |
result = self.generate( | |
inputs_embeds=inputs_embeds, | |
**kwargs | |
) | |
if return_vision_hidden_states: | |
return result, vision_hidden_states | |
return result | |
def initialize_vision_tokenizer(self, mm_use_im_start_end, tokenizer, device, | |
tune_mm_mlp_adapter=False): | |
self.model.vision_config.use_im_start_end = mm_use_im_start_end | |
tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) | |
self.resize_token_embeddings(len(tokenizer)) | |
if mm_use_im_start_end: | |
num_new_tokens = tokenizer.add_tokens( | |
[DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True) | |
self.resize_token_embeddings(len(tokenizer)) | |
self.model.vision_config.im_start_token, self.model.vision_config.im_end_token = tokenizer.convert_tokens_to_ids( | |
[DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN]) | |
if num_new_tokens > 0: | |
input_embeddings = self.get_input_embeddings().weight.data | |
output_embeddings = self.get_output_embeddings().weight.data | |
input_embeddings_avg = input_embeddings[:-num_new_tokens].mean( | |
dim=0, keepdim=True) | |
output_embeddings_avg = output_embeddings[:-num_new_tokens].mean( | |
dim=0, keepdim=True) | |
input_embeddings[-num_new_tokens:] = input_embeddings_avg | |
output_embeddings[-num_new_tokens:] = output_embeddings_avg | |
# for new sft data | |
num_new_tokens = tokenizer.add_tokens( | |
['<box>', '</box>', '<ref>', '</ref>', '<quad>', '</quad>'], special_tokens=True) | |
self.resize_token_embeddings(len(tokenizer)) | |
if num_new_tokens > 0: | |
input_embeddings = self.get_input_embeddings().weight.data | |
output_embeddings = self.get_output_embeddings().weight.data | |
input_embeddings_avg = input_embeddings[:-num_new_tokens].mean( | |
dim=0, keepdim=True) | |
output_embeddings_avg = output_embeddings[:-num_new_tokens].mean( | |
dim=0, keepdim=True) | |
input_embeddings[-num_new_tokens:] = input_embeddings_avg | |
output_embeddings[-num_new_tokens:] = output_embeddings_avg | |
if tune_mm_mlp_adapter: | |
self.model.orig_embeds_params = [ | |
self.get_input_embeddings().weight.data.clone().to(device=device)] | |
for p in self.get_input_embeddings().parameters(): | |
p.requires_grad = True | |
for p in self.get_output_embeddings().parameters(): | |
p.requires_grad = False | |
self.model.vision_config.im_patch_token = tokenizer.convert_tokens_to_ids( | |
[DEFAULT_IMAGE_PATCH_TOKEN])[0] | |
print(f'Tokenizer: {tokenizer}\n patch_token_id: {self.model.vision_config.im_patch_token}, visoin_config: {self.model.vision_config}', flush=True) | |
# exit() | |
AutoConfig.register("omnilmm", OmniLMMConfig) | |
AutoModelForCausalLM.register(OmniLMMConfig, OmniLMMForCausalLM) | |