Views
No views yet
1import io
2
3import argparse
4import numpy as np
5import torch
6from decord import cpu, VideoReader, bridge
7from transformers import AutoModelForCausalLM, AutoTokenizer
8
9MODEL_PATH = "THUDM/cogvlm2-llama3-caption"
10
11DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
12TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[
13 0] >= 8 else torch.float16
14
15parser = argparse.ArgumentParser(description="CogVLM2-Video CLI Demo")
16parser.add_argument('--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=0)
17args = parser.parse_args([])
18
19
20def load_video(video_data, strategy='chat'):
21 bridge.set_bridge('torch')
22 mp4_stream = video_data
23 num_frames = 24
24 decord_vr = VideoReader(io.BytesIO(mp4_stream), ctx=cpu(0))
25
26 frame_id_list = None
27 total_frames = len(decord_vr)
28 if strategy == 'base':
29 clip_end_sec = 60
30 clip_start_sec = 0
31 start_frame = int(clip_start_sec * decord_vr.get_avg_fps())
32 end_frame = min(total_frames,
33 int(clip_end_sec * decord_vr.get_avg_fps())) if clip_end_sec is not None else total_frames
34 frame_id_list = np.linspace(start_frame, end_frame - 1, num_frames, dtype=int)
35 elif strategy == 'chat':
36 timestamps = decord_vr.get_frame_timestamp(np.arange(total_frames))
37 timestamps = [i[0] for i in timestamps]
38 max_second = round(max(timestamps)) + 1
39 frame_id_list = []
40 for second in range(max_second):
41 closest_num = min(timestamps, key=lambda x: abs(x - second))
42 index = timestamps.index(closest_num)
43 frame_id_list.append(index)
44 if len(frame_id_list) >= num_frames:
45 break
46
47 video_data = decord_vr.get_batch(frame_id_list)
48 video_data = video_data.permute(3, 0, 1, 2)
49 return video_data
50
51
52tokenizer = AutoTokenizer.from_pretrained(
53 MODEL_PATH,
54 trust_remote_code=True,
55)
56
57model = AutoModelForCausalLM.from_pretrained(
58 MODEL_PATH,
59 torch_dtype=TORCH_TYPE,
60 trust_remote_code=True
61).eval().to(DEVICE)
62
63
64def predict(prompt, video_data, temperature):
65 strategy = 'chat'
66
67 video = load_video(video_data, strategy=strategy)
68
69 history = []
70 query = prompt
71 inputs = model.build_conversation_input_ids(
72 tokenizer=tokenizer,
73 query=query,
74 images=[video],
75 history=history,
76 template_version=strategy
77 )
78 inputs = {
79 'input_ids': inputs['input_ids'].unsqueeze(0).to('cuda'),
80 'token_type_ids': inputs['token_type_ids'].unsqueeze(0).to('cuda'),
81 'attention_mask': inputs['attention_mask'].unsqueeze(0).to('cuda'),
82 'images': [[inputs['images'][0].to('cuda').to(TORCH_TYPE)]],
83 }
84 gen_kwargs = {
85 "max_new_tokens": 2048,
86 "pad_token_id": 128002,
87 "top_k": 1,
88 "do_sample": False,
89 "top_p": 0.1,
90 "temperature": temperature,
91 }
92 with torch.no_grad():
93 outputs = model.generate(**inputs, **gen_kwargs)
94 outputs = outputs[:, inputs['input_ids'].shape[1]:]
95 response = tokenizer.decode(outputs[0], skip_special_tokens=True)
96 return response
97
98
99def test():
100 prompt = "Please describe this video in detail."
101 temperature = 0.1
102 video_data = open('test.mp4', 'rb').read()
103 response = predict(prompt, video_data, temperature)
104 print(response)
105
106
107if __name__ == '__main__':
108 test()@article{yang2024cogvideox,
title={CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer},
author={Yang, Zhuoyi and Teng, Jiayan and Zheng, Wendi and Ding, Ming and Huang, Shiyu and Xu, Jiazheng and Yang, Yuanming and Hong, Wenyi and Zhang, Xiaohan and Feng, Guanyu and others},
journal={arXiv preprint arXiv:2408.06072},
year={2024}
}