from typing import List, Optional, Tuple
import torch
from PIL import Image
from torch import nn
from transformers import (
CLIPVisionModel,
GenerationMixin,
PreTrainedModel,
PreTrainedTokenizer,
)
from .catty import split_image_with_catty
from .configuration_points_chat import POINTSChatConfig
from .dynamic_high_resolution import split_image
from .modeling_llama import CustomLlamaForCausalLM
class POINTSChatModel(PreTrainedModel, GenerationMixin):
config_class = POINTSChatConfig
_no_split_modules = ["CLIPVisionModel", "LLamaDecoderLayer"]
"""Chat model for POINTS.
Official implementation of the paper "POINTS: Improving Your Vision-language Model with Affordable Strategies" # noqa: E501
paper: https://huggingface.co/papers/2409.04828
Args:
config (PretrainedConfig): The model config.
"""
def __init__(self, config: POINTSChatConfig) -> None:
super().__init__(config)
self.general_vit = CLIPVisionModel(config.vision_config)
self.ocr_vit = CLIPVisionModel(config.vision_config)
self.llm = CustomLlamaForCausalLM(config.llm_config)
self.vision_projector = nn.Sequential(
nn.Linear(config.vision_config.hidden_size *
4, config.llm_config.hidden_size),
nn.GELU(),
nn.Linear(config.llm_config.hidden_size,
config.llm_config.hidden_size)
)
def apply_chat_template(self, prompt: str, image_num: int) -> str:
"""Apply the Yi-1.5-Chat template to the prompt.
Args:
prompt (str): The prompt to apply the template to.
image_num (int): The number of the image in the prompt.
Returns:
str: The prompt with the template applied.
"""
image_tokens = ('<|endoftext|>' * 144) * image_num
prompt = f'<|im_start|>user\n{image_tokens}{prompt}<|im_end|>\n<|im_start|>assistant\n' # noqa: E501
return prompt
def pixel_shuffle(self, feature_map: torch.Tensor,
scale_factor: float = 0.5) -> torch.Tensor:
"""Implementation of pixel shuffle.
Merge several patches into a single patch by concatenating
them across the channel dimension. Therefore, we can reduce
the image sequence length. In POINTS, we merge 2x2 adjacent
patches into a single patch.
Args:
feature_map (torch.Tensor): The feature map to be pixel
shuffled.
scale_factor (float, optional): The scale factor for the
"""
# taken from https://huggingface.co/OpenGVLab/InternVL-Chat-V1-5/blob/main/modeling_internvl_chat.py#L187 # noqa
n, w, h, c = feature_map.size()
# N, W, H, C --> N, W, H * scale, C // scale
feature_map = feature_map.view(
n, w, int(h * scale_factor), int(c / scale_factor))
# N, W, H * scale, C // scale --> N, H * scale, W, C // scale
feature_map = feature_map.permute(0, 2, 1, 3).contiguous()
# N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
feature_map = feature_map.view(
n,
int(h * scale_factor),
int(w * scale_factor),
int(c / (scale_factor * scale_factor)),
)
feature_map = feature_map.permute(0, 2, 1, 3).contiguous()
return feature_map
def extract_image_features(self, images: torch.Tensor,
vision_encoder: str = 'general_vit') -> torch.Tensor: # noqa: E501
"""Extract the image features from the vision encoder.
Args:
images (torch.Tensor): The images to extract the features from.
vision_encoder (str, optional): The vision encoder to use.
Defaults to 'general_vit'.
Returns:
torch.Tensor: The extracted image features.
"""
if vision_encoder == 'general_vit':
image_features = self.general_vit(
images, output_hidden_states=True
)
else:
image_features = self.ocr_vit(
images, output_hidden_states=True
)
image_features = image_features.hidden_states[-2]
image_features = image_features[:, 1:]
image_features = image_features.reshape(-1, 24, 24, 1024)
image_features = self.pixel_shuffle(image_features, 0.5)
image_features = image_features.view(-1, 144, 4096)
image_features = self.vision_projector(image_features)
return image_features
def get_pos_mapping(self, pos: List[list]) -> Tuple[dict, int]:
"""Get the position mapping for the images.
Args:
pos (List[list]): The position of the images in the prompt.
Returns:
Tuple[dict, int]: The position mapping and the
total number of images.
"""
mapping = {}
total_images = 0
for i, (start, end) in enumerate(pos):
num_image = int((end - start) / 144)
mapping[i] = num_image
total_images += num_image
return mapping, total_images
@torch.no_grad()
def chat(self, pixel_values: Image, prompt: str,
tokenizer: PreTrainedTokenizer,
image_processor, catty: bool = True,
generation_config: dict = None,
max_splits: int = 8) -> str:
"""Generate a response to the input prompt.
Args:
pixel_values (Image): The input image.
prompt (str): The input prompt.
tokenizer (PreTrainedTokenizer): The tokenizer to use.
image_processor: The image processor to use.
catty (bool, optional): Whether to use catty. Defaults to True.
generation_config (dict, optional): The generation config.
Defaults to None.
max_splits (int, optional): The maximum number of splits.
Defaults to 8.
Returns:
str: The generated response.
"""
if catty:
cropped_images = split_image_with_catty(pixel_values,
do_resize=True,
max_crop_slices=max_splits)
else:
cropped_images = split_image(pixel_values, max_splits=max_splits)
prompt = self.apply_chat_template(prompt, len(cropped_images))
cropped_images = image_processor.preprocess(
cropped_images, return_tensors='pt')['pixel_values']
cropped_images = cropped_images.to(self.device)
# extract features with general_vit
general_vit_features = self.extract_image_features(
cropped_images, vision_encoder='general_vit')
# extract features with ocr_vit
ocr_vit_features = self.extract_image_features(
cropped_images, vision_encoder='ocr_vit')
image_features = 0.5 * general_vit_features + 0.5 * ocr_vit_features
model_inputs = tokenizer(prompt, return_tensors='pt')
input_ids = model_inputs['input_ids'].to(self.device)
attention_mask = model_inputs['attention_mask'].to(self.device)
# stop token
eos_token_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
# image token
image_token_id = tokenizer.convert_tokens_to_ids("<|endoftext|>")
generation_config.update(
{
'eos_token_id': eos_token_id,
}
)
outputs = self.generate(
input_ids=input_ids,
attention_mask=attention_mask,
image_features=[image_features],
image_token_id=image_token_id,
**generation_config
)
response = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
return response
def generate(self,
input_ids: torch.LongTensor,
attention_mask: torch.LongTensor,
image_features: List[torch.Tensor],
image_token_id: int,
generation_config: Optional[dict] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**generate_kwargs) -> torch.LongTensor:
input_embeddings = self.llm.lm.embed_in(input_ids)
batch_size = input_ids.shape[0]
assert len(image_features) == batch_size
for i in range(batch_size):
special_pos = input_ids[i] == image_token_id
pos = (special_pos[:-1] != special_pos[1:]).nonzero() + 1
if pos.shape[0] % 2 != 0:
# when the sequence is <image><caption>
# we need to add a dummy token
pos = torch.cat([torch.tensor([[0]]).to(pos.device), pos])
pos = pos.reshape(-1, 2).tolist()
pos_mapping, total_images = self.get_pos_mapping(pos)
assert total_images == len(image_features[i])
img_offset = 0
for j, (start, end) in enumerate(pos):
num_images = pos_mapping[j]
input_embeddings[i, start:end] = torch.cat(
[image_features[i][img_offset+k]
for k in range(num_images)],
dim=0
)
img_offset += num_images
outputs = self.llm.generate(
inputs_embeds=input_embeddings,
attention_mask=attention_mask,
generation_config=generation_config,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
use_cache=True,
**generate_kwargs
)
return outputs