|
--- |
|
license: llama2 |
|
--- |
|
# Chat-UniVi: Unified Visual Representation Empowers Large Language Models with Image and Video Understanding |
|
|
|
**Paper or resources for more information:** |
|
[[Paper](https://huggingface.co/papers/2311.08046)] [[Code](https://github.com/PKU-YuanGroup/Chat-UniVi)] |
|
|
|
## License |
|
Llama 2 is licensed under the LLAMA 2 Community License, |
|
Copyright (c) Meta Platforms, Inc. All Rights Reserved. |
|
|
|
## 😮 Highlights |
|
|
|
### 💡 Unified visual representation for image and video |
|
We employ **a set of dynamic visual tokens** to uniformly represent images and videos. |
|
This representation framework empowers the model to efficiently utilize **a limited number of visual tokens** to simultaneously capture **the spatial details necessary for images** and **the comprehensive temporal relationship required for videos**. |
|
|
|
### 🔥 Joint training strategy, making LLMs understand both image and video |
|
Chat-UniVi is trained on a mixed dataset containing both images and videos, allowing direct application to tasks involving both mediums without requiring any modifications. |
|
|
|
### 🤗 High performance, complementary learning with image and video |
|
Extensive experimental results demonstrate that Chat-UniVi, as a unified model, consistently outperforms even existing methods exclusively designed for either images or videos. |
|
|
|
|
|
### Inference for Video Understanding |
|
```python |
|
import torch |
|
import os |
|
from ChatUniVi.constants import * |
|
from ChatUniVi.conversation import conv_templates, SeparatorStyle |
|
from ChatUniVi.model.builder import load_pretrained_model |
|
from ChatUniVi.utils import disable_torch_init |
|
from ChatUniVi.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria |
|
from PIL import Image |
|
from decord import VideoReader, cpu |
|
import numpy as np |
|
|
|
|
|
def _get_rawvideo_dec(video_path, image_processor, max_frames=MAX_IMAGE_LENGTH, image_resolution=224, video_framerate=1, s=None, e=None): |
|
# speed up video decode via decord. |
|
|
|
if s is None: |
|
start_time, end_time = None, None |
|
else: |
|
start_time = int(s) |
|
end_time = int(e) |
|
start_time = start_time if start_time >= 0. else 0. |
|
end_time = end_time if end_time >= 0. else 0. |
|
if start_time > end_time: |
|
start_time, end_time = end_time, start_time |
|
elif start_time == end_time: |
|
end_time = start_time + 1 |
|
|
|
if os.path.exists(video_path): |
|
vreader = VideoReader(video_path, ctx=cpu(0)) |
|
else: |
|
print(video_path) |
|
raise FileNotFoundError |
|
|
|
fps = vreader.get_avg_fps() |
|
f_start = 0 if start_time is None else int(start_time * fps) |
|
f_end = int(min(1000000000 if end_time is None else end_time * fps, len(vreader) - 1)) |
|
num_frames = f_end - f_start + 1 |
|
if num_frames > 0: |
|
# T x 3 x H x W |
|
sample_fps = int(video_framerate) |
|
t_stride = int(round(float(fps) / sample_fps)) |
|
|
|
all_pos = list(range(f_start, f_end + 1, t_stride)) |
|
if len(all_pos) > max_frames: |
|
sample_pos = [all_pos[_] for _ in np.linspace(0, len(all_pos) - 1, num=max_frames, dtype=int)] |
|
else: |
|
sample_pos = all_pos |
|
|
|
patch_images = [Image.fromarray(f) for f in vreader.get_batch(sample_pos).asnumpy()] |
|
|
|
patch_images = torch.stack([image_processor.preprocess(img, return_tensors='pt')['pixel_values'][0] for img in patch_images]) |
|
slice_len = patch_images.shape[0] |
|
|
|
return patch_images, slice_len |
|
else: |
|
print("video path: {} error.".format(video_path)) |
|
|
|
|
|
if __name__ == '__main__': |
|
# Model Parameter |
|
model_path = "Chat-UniVi/Chat-UniVi" # or "Chat-UniVi/Chat-UniVi-13B" |
|
video_path = ${video_path} |
|
|
|
# The number of visual tokens varies with the length of the video. "max_frames" is the maximum number of frames. |
|
# When the video is long, we will uniformly downsample the video to meet the frames when equal to the "max_frames". |
|
max_frames = 100 |
|
|
|
# The number of frames retained per second in the video. |
|
video_framerate = 1 |
|
|
|
# Input Text |
|
qs = "Describe the video." |
|
|
|
# Sampling Parameter |
|
conv_mode = "simple" |
|
temperature = 0.2 |
|
top_p = None |
|
num_beams = 1 |
|
|
|
disable_torch_init() |
|
model_path = os.path.expanduser(model_path) |
|
model_name = "ChatUniVi" |
|
tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name) |
|
|
|
mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False) |
|
mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True) |
|
if mm_use_im_patch_token: |
|
tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) |
|
if mm_use_im_start_end: |
|
tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True) |
|
model.resize_token_embeddings(len(tokenizer)) |
|
|
|
vision_tower = model.get_vision_tower() |
|
if not vision_tower.is_loaded: |
|
vision_tower.load_model() |
|
image_processor = vision_tower.image_processor |
|
|
|
if model.config.config["use_cluster"]: |
|
for n, m in model.named_modules(): |
|
m = m.to(dtype=torch.bfloat16) |
|
|
|
# Check if the video exists |
|
if video_path is not None: |
|
video_frames, slice_len = _get_rawvideo_dec(video_path, image_processor, max_frames=max_frames, video_framerate=video_framerate) |
|
|
|
cur_prompt = qs |
|
if model.config.mm_use_im_start_end: |
|
qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN * slice_len + DEFAULT_IM_END_TOKEN + '\n' + qs |
|
else: |
|
qs = DEFAULT_IMAGE_TOKEN * slice_len + '\n' + qs |
|
|
|
conv = conv_templates[conv_mode].copy() |
|
conv.append_message(conv.roles[0], qs) |
|
conv.append_message(conv.roles[1], None) |
|
prompt = conv.get_prompt() |
|
|
|
input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze( |
|
0).cuda() |
|
|
|
stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2 |
|
keywords = [stop_str] |
|
stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) |
|
|
|
with torch.inference_mode(): |
|
output_ids = model.generate( |
|
input_ids, |
|
images=video_frames.half().cuda(), |
|
do_sample=True, |
|
temperature=temperature, |
|
top_p=top_p, |
|
num_beams=num_beams, |
|
output_scores=True, |
|
return_dict_in_generate=True, |
|
max_new_tokens=1024, |
|
use_cache=True, |
|
stopping_criteria=[stopping_criteria]) |
|
|
|
output_ids = output_ids.sequences |
|
input_token_len = input_ids.shape[1] |
|
n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item() |
|
if n_diff_input_output > 0: |
|
print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids') |
|
outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] |
|
outputs = outputs.strip() |
|
if outputs.endswith(stop_str): |
|
outputs = outputs[:-len(stop_str)] |
|
outputs = outputs.strip() |
|
print(outputs) |
|
``` |
|
|
|
### Inference for Image Understanding |
|
```python |
|
import torch |
|
import os |
|
from ChatUniVi.constants import * |
|
from ChatUniVi.conversation import conv_templates, SeparatorStyle |
|
from ChatUniVi.model.builder import load_pretrained_model |
|
from ChatUniVi.utils import disable_torch_init |
|
from ChatUniVi.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria |
|
from PIL import Image |
|
|
|
|
|
if __name__ == '__main__': |
|
# Model Parameter |
|
model_path = "Chat-UniVi/Chat-UniVi" # or "Chat-UniVi/Chat-UniVi-13B" |
|
image_path = ${image_path} |
|
|
|
# Input Text |
|
qs = "Describe the image." |
|
|
|
# Sampling Parameter |
|
conv_mode = "simple" |
|
temperature = 0.2 |
|
top_p = None |
|
num_beams = 1 |
|
|
|
disable_torch_init() |
|
model_path = os.path.expanduser(model_path) |
|
model_name = "ChatUniVi" |
|
tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name) |
|
|
|
mm_use_im_start_end = getattr(model.config, "mm_use_im_start_end", False) |
|
mm_use_im_patch_token = getattr(model.config, "mm_use_im_patch_token", True) |
|
if mm_use_im_patch_token: |
|
tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True) |
|
if mm_use_im_start_end: |
|
tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True) |
|
model.resize_token_embeddings(len(tokenizer)) |
|
|
|
vision_tower = model.get_vision_tower() |
|
if not vision_tower.is_loaded: |
|
vision_tower.load_model() |
|
image_processor = vision_tower.image_processor |
|
|
|
# Check if the video exists |
|
if image_path is not None: |
|
cur_prompt = qs |
|
if model.config.mm_use_im_start_end: |
|
qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + qs |
|
else: |
|
qs = DEFAULT_IMAGE_TOKEN + '\n' + qs |
|
|
|
conv = conv_templates[conv_mode].copy() |
|
conv.append_message(conv.roles[0], qs) |
|
conv.append_message(conv.roles[1], None) |
|
prompt = conv.get_prompt() |
|
|
|
input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda() |
|
|
|
image = Image.open(image_path) |
|
image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0] |
|
|
|
stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2 |
|
keywords = [stop_str] |
|
stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) |
|
|
|
with torch.inference_mode(): |
|
output_ids = model.generate( |
|
input_ids, |
|
images=image_tensor.unsqueeze(0).half().cuda(), |
|
do_sample=True, |
|
temperature=temperature, |
|
top_p=top_p, |
|
num_beams=num_beams, |
|
max_new_tokens=1024, |
|
use_cache=True, |
|
stopping_criteria=[stopping_criteria]) |
|
|
|
input_token_len = input_ids.shape[1] |
|
n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item() |
|
if n_diff_input_output > 0: |
|
print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids') |
|
outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0] |
|
outputs = outputs.strip() |
|
if outputs.endswith(stop_str): |
|
outputs = outputs[:-len(stop_str)] |
|
outputs = outputs.strip() |
|
print(outputs) |
|
``` |
