3D-GRAND / llava /model /llava_arch.py
jedyang97's picture
initial demo
947767a
raw
history blame
20.1 kB
# Copyright 2023 Haotian Liu
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import ABC, abstractmethod
import torch
import torch.nn as nn
import transformers
from llava.model.utils import get_w
from .multimodal_encoder.builder import build_vision_tower
from llava.constants import (
GROUND_TOKEN,
IGNORE_INDEX,
IMAGE_TOKEN_INDEX,
DEFAULT_IMAGE_PATCH_TOKEN,
DEFAULT_IM_START_TOKEN,
DEFAULT_IM_END_TOKEN,
PROFILE_RUNTIME,
)
import time
from transformers.utils import logging
logger = logging.get_logger("transformers")
class LlavaMetaModel:
def __init__(self, config):
super(LlavaMetaModel, self).__init__(config)
if hasattr(config, "mm_vision_tower"):
self.vision_tower = build_vision_tower(config, delay_load=True)
if self.vision_tower is not None:
self.mm_projector = nn.Linear(
self.vision_tower.hidden_size, config.hidden_size
) # placeholder, this will be re-initialized later in initialize_vision_modules()
def get_vision_tower(self):
vision_tower = getattr(self, "vision_tower", None)
if type(vision_tower) is list:
vision_tower = vision_tower[0]
return vision_tower
def initialize_vision_modules(self, model_args, fsdp=None):
vision_tower = model_args.vision_tower
mm_vision_select_layer = model_args.mm_vision_select_layer
mm_vision_select_feature = model_args.mm_vision_select_feature
pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
pretrain_vision_tower = model_args.pretrain_vision_tower
self.config.mm_vision_tower = vision_tower
if hasattr(self, "vision_tower"):
del self.vision_tower
torch.cuda.empty_cache()
vision_tower = build_vision_tower(model_args)
if vision_tower is None:
return
if fsdp is not None and len(fsdp) > 0:
self.vision_tower = [vision_tower]
else:
self.vision_tower = vision_tower
# add these model args to HF config so that they can be saved (used for loading checkpoint)
self.config.use_mm_proj = True
self.config.mm_hidden_size = vision_tower.hidden_size
self.config.mm_vision_select_layer = mm_vision_select_layer
self.config.mm_vision_select_feature = mm_vision_select_feature
self.config.num_points = model_args.num_points
self.config.feature_dim = model_args.feature_dim
self.config.num_latents = model_args.num_latents
self.config.d_latents = model_args.d_latents
self.config.num_cross_attention_heads = model_args.num_cross_attention_heads
self.config.position_encoding_type = model_args.position_encoding_type
if hasattr(self, "mm_projector"):
del self.mm_projector
torch.cuda.empty_cache()
self.mm_projector = nn.Linear(self.config.mm_hidden_size, self.config.hidden_size)
if pretrain_mm_mlp_adapter is not None:
mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
torch.cuda.empty_cache()
class LlavaMetaForCausalLM(ABC):
@abstractmethod
def get_model(self):
pass
def get_vision_tower(self):
return self.get_model().get_vision_tower()
def encode_images(self, images):
vision_features_before_mm_projection = self.get_model().get_vision_tower()(
images
) # for minkowski, the output of this step will be float32
vision_features_before_mm_projection = vision_features_before_mm_projection.to(
dtype=self.dtype
) # convert back to the dtype of the LLM (bfloat16 in most cases), no-op if the dtype is already the same
vision_features = self.get_model().mm_projector(
vision_features_before_mm_projection
) # vision_features and mm_projector are both float32
return vision_features, vision_features_before_mm_projection
def prepare_inputs_labels_for_multimodal(
self, input_ids, attention_mask, past_key_values, labels, images
):
vision_tower = self.get_vision_tower()
if vision_tower is None or images is None or input_ids.shape[1] == 1:
if (
past_key_values is not None
and vision_tower is not None
and images is not None
and input_ids.shape[1] == 1
):
attention_mask = torch.ones(
(attention_mask.shape[0], past_key_values[-1][-1].shape[-2] + 1),
dtype=attention_mask.dtype,
device=attention_mask.device,
)
vision_features_before_mm_projection = images
return (
input_ids,
attention_mask,
past_key_values,
None,
labels,
vision_features_before_mm_projection,
)
start_time_encode_images = time.time()
if not isinstance(images, SparseTensor) and (type(images) is list or images.ndim == 5):
concat_images = torch.cat([image for image in images], dim=0)
vision_features, vision_features_before_mm_projection = self.encode_images(
concat_images
)
split_sizes = [image.shape[0] for image in images]
vision_features = torch.split(vision_features, split_sizes, dim=0)
vision_features = [x.flatten(0, 1) for x in vision_features]
else:
vision_features, vision_features_before_mm_projection = self.encode_images(images)
if PROFILE_RUNTIME:
logger.info(f"Time to encode images: {time.time() - start_time_encode_images}")
start_time_for_loop = time.time()
new_input_embeds = []
new_labels = [] if labels is not None else None
cur_image_idx = 0
for batch_idx, cur_input_ids in enumerate(input_ids):
if (cur_input_ids == IMAGE_TOKEN_INDEX).sum() == 0:
# multimodal LLM, but the current sample is not multimodal
cur_input_embeds = self.get_model().embed_tokens(cur_input_ids)
cur_input_embeds = (
cur_input_embeds
+ (0.0 * self.get_model().mm_projector(vision_tower.dummy_feature)).sum()
)
new_input_embeds.append(cur_input_embeds)
if labels is not None:
new_labels.append(labels[batch_idx])
cur_image_idx += 1
continue
image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
cur_new_input_embeds = []
if labels is not None:
cur_labels = labels[batch_idx]
cur_new_labels = []
assert cur_labels.shape == cur_input_ids.shape
# The following while loop looks for all image tokens in the current sentence
# and replace them with the corresponding image features.
while image_token_indices.numel() > 0:
cur_vision_features = vision_features[cur_image_idx]
image_token_start = image_token_indices[0]
if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
self.config, "mm_use_im_start_end", False
):
cur_new_input_embeds.append(
self.get_model()
.embed_tokens(cur_input_ids[: image_token_start - 1])
.detach()
)
cur_new_input_embeds.append(
self.get_model().embed_tokens(
cur_input_ids[image_token_start - 1 : image_token_start]
)
)
cur_new_input_embeds.append(cur_vision_features)
cur_new_input_embeds.append(
self.get_model().embed_tokens(
cur_input_ids[image_token_start + 1 : image_token_start + 2]
)
)
if labels is not None:
cur_new_labels.append(cur_labels[:image_token_start])
cur_new_labels.append(
torch.full(
(cur_vision_features.shape[0],),
IGNORE_INDEX,
device=labels.device,
dtype=labels.dtype,
)
)
cur_new_labels.append(cur_labels[image_token_start : image_token_start + 1])
cur_labels = cur_labels[image_token_start + 2 :]
else:
cur_new_input_embeds.append(
self.get_model().embed_tokens(cur_input_ids[:image_token_start])
)
cur_new_input_embeds.append(cur_vision_features)
if labels is not None:
cur_new_labels.append(cur_labels[:image_token_start])
cur_new_labels.append(
torch.full(
(cur_vision_features.shape[0],),
IGNORE_INDEX,
device=labels.device,
dtype=labels.dtype,
)
)
cur_labels = cur_labels[image_token_start + 1 :]
cur_image_idx += 1
if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
self.config, "mm_use_im_start_end", False
):
cur_input_ids = cur_input_ids[image_token_start + 2 :]
else:
cur_input_ids = cur_input_ids[image_token_start + 1 :]
image_token_indices = torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0]
if cur_input_ids.numel() > 0:
if getattr(self.config, "tune_mm_mlp_adapter", False) and getattr(
self.config, "mm_use_im_start_end", False
):
cur_new_input_embeds.append(
self.get_model().embed_tokens(cur_input_ids).detach()
)
else:
cur_new_input_embeds.append(self.get_model().embed_tokens(cur_input_ids))
if labels is not None:
cur_new_labels.append(cur_labels)
cur_new_input_embeds = [x.to(device=self.device) for x in cur_new_input_embeds]
cur_new_input_embeds = torch.cat(cur_new_input_embeds, dim=0)
new_input_embeds.append(cur_new_input_embeds)
if labels is not None:
cur_new_labels = torch.cat(cur_new_labels, dim=0)
new_labels.append(cur_new_labels)
if PROFILE_RUNTIME:
logger.info(f"Time for loop: {time.time() - start_time_for_loop}")
start_time_paddding = time.time()
# pad all sentences in batch to the same length
if any(x.shape != new_input_embeds[0].shape for x in new_input_embeds):
max_len = max(x.shape[0] for x in new_input_embeds)
new_input_embeds_align = []
for cur_new_embed in new_input_embeds:
cur_new_embed = torch.cat(
(
cur_new_embed,
torch.zeros(
(max_len - cur_new_embed.shape[0], cur_new_embed.shape[1]),
dtype=cur_new_embed.dtype,
device=cur_new_embed.device,
),
),
dim=0,
)
new_input_embeds_align.append(cur_new_embed)
new_input_embeds = torch.stack(new_input_embeds_align, dim=0)
if labels is not None:
new_labels_align = []
_new_labels = new_labels
for cur_new_label in new_labels:
cur_new_label = torch.cat(
(
cur_new_label,
torch.full(
(max_len - cur_new_label.shape[0],),
IGNORE_INDEX,
dtype=cur_new_label.dtype,
device=cur_new_label.device,
),
),
dim=0,
)
new_labels_align.append(cur_new_label)
new_labels = torch.stack(new_labels_align, dim=0)
if attention_mask is not None:
new_attention_mask = []
for cur_attention_mask, cur_new_labels, cur_new_labels_align in zip(
attention_mask, _new_labels, new_labels
):
new_attn_mask_pad_left = torch.full(
(cur_new_labels.shape[0] - labels.shape[1],),
True,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
new_attn_mask_pad_right = torch.full(
(cur_new_labels_align.shape[0] - cur_new_labels.shape[0],),
False,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
cur_new_attention_mask = torch.cat(
(
new_attn_mask_pad_left,
cur_attention_mask,
new_attn_mask_pad_right,
),
dim=0,
)
new_attention_mask.append(cur_new_attention_mask)
attention_mask = torch.stack(new_attention_mask, dim=0)
assert attention_mask.shape == new_labels.shape
else:
new_input_embeds = torch.stack(new_input_embeds, dim=0)
if labels is not None:
new_labels = torch.stack(new_labels, dim=0)
if attention_mask is not None:
new_attn_mask_pad_left = torch.full(
(
attention_mask.shape[0],
new_input_embeds.shape[1] - input_ids.shape[1],
),
True,
dtype=attention_mask.dtype,
device=attention_mask.device,
)
attention_mask = torch.cat((new_attn_mask_pad_left, attention_mask), dim=1)
assert attention_mask.shape == new_input_embeds.shape[:2]
if PROFILE_RUNTIME:
logger.info(f"Time padding: {time.time() - start_time_paddding}")
return (
None,
attention_mask,
past_key_values,
new_input_embeds,
new_labels,
vision_features_before_mm_projection,
)
def initialize_vision_tokenizer(self, model_args, tokenizer):
if model_args.mm_use_im_patch_token:
tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
self.resize_token_embeddings(len(tokenizer))
if model_args.mm_use_im_start_end:
(
num_new_tokens,
input_embeddings,
output_embeddings,
) = self.add_special_tokens_and_resize_embeddings(
special_tokens_list=[DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN],
tokenizer=tokenizer,
)
if model_args.tune_mm_mlp_adapter:
for p in self.get_input_embeddings().parameters():
p.requires_grad = True
for p in self.get_output_embeddings().parameters():
p.requires_grad = False
if model_args.pretrain_mm_mlp_adapter:
mm_projector_weights = torch.load(
model_args.pretrain_mm_mlp_adapter, map_location="cpu"
)
embed_tokens_weight = mm_projector_weights["model.embed_tokens.weight"]
assert num_new_tokens == 2
if input_embeddings.shape == embed_tokens_weight.shape:
input_embeddings[-num_new_tokens:] = embed_tokens_weight[-num_new_tokens:]
elif embed_tokens_weight.shape[0] == num_new_tokens:
input_embeddings[-num_new_tokens:] = embed_tokens_weight
else:
raise ValueError(
f"Unexpected embed_tokens_weight shape. Pretrained: {embed_tokens_weight.shape}. Current: {input_embeddings.shape}. Numer of new tokens: {num_new_tokens}."
)
elif model_args.mm_use_im_patch_token:
if model_args.tune_mm_mlp_adapter:
for p in self.get_input_embeddings().parameters():
p.requires_grad = False
for p in self.get_output_embeddings().parameters():
p.requires_grad = False
# add special tokens if bbox_tokenization_type is location_tokens
if model_args.bbox_tokenization_type == "location_tokens":
num_special_tokens = model_args.num_voxels_per_axis_for_location_tokens**3
self.add_special_tokens_and_resize_embeddings(
special_tokens_list=[f"<loc_{i}>" for i in range(num_special_tokens)],
tokenizer=tokenizer,
)
for p in self.get_input_embeddings().parameters():
p.requires_grad = True
for p in self.get_output_embeddings().parameters():
p.requires_grad = True
elif model_args.bbox_tokenization_type == "ground_token":
self.add_special_tokens_and_resize_embeddings(
special_tokens_list=[GROUND_TOKEN], tokenizer=tokenizer
)
for p in self.get_input_embeddings().parameters():
p.requires_grad = True
for p in self.get_output_embeddings().parameters():
p.requires_grad = True
# add special token to input id mapping
self.config.added_special_token_to_input_id = tokenizer.get_added_vocab()
def add_special_tokens_and_resize_embeddings(
self,
special_tokens_list: list[str],
tokenizer: transformers.PreTrainedTokenizer,
):
num_new_tokens = tokenizer.add_tokens(special_tokens_list, 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
return num_new_tokens, input_embeddings, output_embeddings