| 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): |
| |
| video_mask = np.zeros(max_frames, dtype=np.int64) |
| max_video_length = 0 |
|
|
| |
| video = np.zeros((max_frames, 3, image_resolution, image_resolution), dtype=np.float64) |
|
|
| 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: |
| |
| 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] |
|
|
| max_video_length = max_video_length if max_video_length > slice_len else slice_len |
| if slice_len < 1: |
| pass |
| else: |
| video[:slice_len, ...] = patch_images |
|
|
| return patch_images, video_mask |
| else: |
| print("video path: {} error.".format(video_path)) |
|
|
| video_mask[:max_video_length] = [1] * max_video_length |
|
|
| return torch.from_numpy(video), video_mask |
|
|
| if __name__ == '__main__': |
| |
| model_path = ${model_path} |
| video_path = ${video_path} |
| max_frames = ${max_frames} |
| qs = "Describe the video." |
|
|
| 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) |
|
|
| |
| if video_path is not None: |
| video_frames, _ = _get_rawvideo_dec(video_path, image_processor, max_frames=max_frames) |
|
|
| cur_prompt = qs |
| if model.config.mm_use_im_start_end: |
| qs = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN * MAX_IMAGE_LENGTH + DEFAULT_IM_END_TOKEN + '\n' + qs |
| else: |
| qs = DEFAULT_IMAGE_TOKEN * MAX_IMAGE_LENGTH + '\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) |
