|
import numpy as np |
|
|
|
import torch |
|
|
|
from longvu.builder import load_pretrained_model |
|
from longvu.constants import ( |
|
DEFAULT_IMAGE_TOKEN, |
|
IMAGE_TOKEN_INDEX, |
|
) |
|
from longvu.conversation import conv_templates, SeparatorStyle |
|
from longvu.mm_datautils import ( |
|
KeywordsStoppingCriteria, |
|
process_images, |
|
tokenizer_image_token, |
|
) |
|
from decord import cpu, VideoReader |
|
|
|
version = "qwen" |
|
model_name = "cambrian_qwen" |
|
input_model_local_path = "./checkpoints/longvu_qwen" |
|
|
|
device = "cuda:7" |
|
|
|
tokenizer, model, image_processor, context_len = load_pretrained_model( |
|
input_model_local_path, None, model_name, device=device |
|
) |
|
model.get_model().config.tokenizer_model_max_length = 8192 |
|
model.get_model().config.inference_max_length = 128 |
|
model.config.use_cache = True |
|
print(model.device) |
|
|
|
model.eval() |
|
|
|
video_path = "./examples/video1.mp4" |
|
qs = "Describe this video in detail" |
|
|
|
vr = VideoReader(video_path, ctx=cpu(0), num_threads=1) |
|
fps = float(vr.get_avg_fps()) |
|
frame_indices = np.array( |
|
[ |
|
i |
|
for i in range( |
|
0, |
|
len(vr), |
|
round(fps), |
|
) |
|
] |
|
) |
|
video = [] |
|
for frame_index in frame_indices: |
|
img = vr[frame_index].asnumpy() |
|
video.append(img) |
|
video = np.stack(video) |
|
image_sizes = [video[0].shape[:2]] |
|
video = process_images(video, image_processor, model.config) |
|
video = [item.unsqueeze(0) for item in video] |
|
|
|
qs = DEFAULT_IMAGE_TOKEN + "\n" + qs |
|
|
|
conv = conv_templates[version].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) |
|
.to(model.device) |
|
) |
|
|
|
if "llama3" in version: |
|
input_ids = input_ids[0][1:].unsqueeze(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=video, |
|
image_sizes=image_sizes, |
|
do_sample=False, |
|
temperature=0.2, |
|
max_new_tokens=128, |
|
use_cache=True, |
|
stopping_criteria=[stopping_criteria], |
|
) |
|
|
|
if isinstance(output_ids, tuple): |
|
output_ids = output_ids[0] |
|
|
|
pred = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0].strip() |
|
|
|
print("pred: ", pred, flush=True) |
|
|
