IVGSZ commited on
Commit
6274907
·
verified ·
1 Parent(s): 07d4ad9

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +1 -1
  2. demo.py +141 -0
app.py CHANGED
@@ -93,7 +93,7 @@ def clear_history(state, state_):
93
  True, state, state_, state.to_gradio_chatbot(), [])
94
 
95
 
96
- conv_mode = "simple"
97
  handler = Chat(model_path, conv_mode=conv_mode, load_4bit=load_4bit, load_8bit=load_8bit)
98
  if not os.path.exists("temp"):
99
  os.makedirs("temp")
 
93
  True, state, state_, state.to_gradio_chatbot(), [])
94
 
95
 
96
+ conv_mode = "vicuna_v1"
97
  handler = Chat(model_path, conv_mode=conv_mode, load_4bit=load_4bit, load_8bit=load_8bit)
98
  if not os.path.exists("temp"):
99
  os.makedirs("temp")
demo.py ADDED
@@ -0,0 +1,141 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from ..constants import *
3
+ from ..conversation import conv_templates, SeparatorStyle
4
+ from ..model.builder import load_pretrained_model
5
+ from ..utils import disable_torch_init
6
+ from ..mm_utils import tokenizer_image_token, KeywordsStoppingCriteria, get_model_name_from_path
7
+ from PIL import Image
8
+ import os
9
+ from decord import VideoReader, cpu
10
+ import numpy as np
11
+
12
+
13
+ class Chat:
14
+ def __init__(self, model_path, conv_mode="simple", load_8bit=False, load_4bit=False):
15
+ disable_torch_init()
16
+ model_name = get_model_name_from_path(model_path)
17
+ self.tokenizer, self.model, self.image_processor, context_len = load_pretrained_model(model_path, None, model_name, load_8bit=load_8bit, load_4bit=load_4bit)
18
+
19
+ mm_use_im_start_end = getattr(self.model.config, "mm_use_im_start_end", False)
20
+ mm_use_im_patch_token = getattr(self.model.config, "mm_use_im_patch_token", True)
21
+ if mm_use_im_patch_token:
22
+ self.tokenizer.add_tokens([DEFAULT_IMAGE_PATCH_TOKEN], special_tokens=True)
23
+ if mm_use_im_start_end:
24
+ self.tokenizer.add_tokens([DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True)
25
+ self.model.resize_token_embeddings(len(self.tokenizer))
26
+
27
+ vision_tower = self.model.get_vision_tower()
28
+ if not vision_tower.is_loaded:
29
+ vision_tower.load_model()
30
+
31
+ self.image_processor = vision_tower.image_processor
32
+ self.conv_mode = conv_mode
33
+ print(self.model)
34
+
35
+ def get_prompt(self, qs, state):
36
+ state.append_message(state.roles[0], qs)
37
+ state.append_message(state.roles[1], None)
38
+ return state
39
+
40
+ def _get_rawvideo_dec(self, video_path, image_processor, max_frames=MAX_IMAGE_LENGTH, image_resolution=224,
41
+ video_framerate=1, s=None, e=None):
42
+ if s is None:
43
+ start_time, end_time = None, None
44
+ else:
45
+ start_time = int(s)
46
+ end_time = int(e)
47
+ start_time = start_time if start_time >= 0. else 0.
48
+ end_time = end_time if end_time >= 0. else 0.
49
+ if start_time > end_time:
50
+ start_time, end_time = end_time, start_time
51
+ elif start_time == end_time:
52
+ end_time = start_time + 1
53
+
54
+ if os.path.exists(video_path):
55
+ vreader = VideoReader(video_path, ctx=cpu(0))
56
+ else:
57
+ print(video_path)
58
+ raise FileNotFoundError
59
+
60
+ fps = vreader.get_avg_fps()
61
+ f_start = 0 if start_time is None else int(start_time * fps)
62
+ f_end = int(min(1000000000 if end_time is None else end_time * fps, len(vreader) - 1))
63
+ num_frames = f_end - f_start + 1
64
+ if num_frames > 0:
65
+ sample_fps = int(video_framerate)
66
+ t_stride = int(round(float(fps) / sample_fps))
67
+
68
+ all_pos = list(range(f_start, f_end + 1, t_stride))
69
+ if len(all_pos) > max_frames:
70
+ sample_pos = [all_pos[_] for _ in np.linspace(0, len(all_pos) - 1, num=max_frames, dtype=int)]
71
+ else:
72
+ sample_pos = all_pos
73
+
74
+ patch_images = [Image.fromarray(f) for f in vreader.get_batch(sample_pos).asnumpy()]
75
+ return patch_images
76
+
77
+ @torch.inference_mode()
78
+ def generate(self, images_tensor: list, prompt: str, first_run: bool, state):
79
+ tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor
80
+
81
+ state = self.get_prompt(prompt, state)
82
+ prompt = state.get_prompt()
83
+ print(prompt)
84
+
85
+ images_tensor = torch.stack(images_tensor, dim=0)
86
+ input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
87
+
88
+ temperature = 0.2
89
+ max_new_tokens = 1024
90
+
91
+ stop_str = conv_templates[self.conv_mode].copy().sep if conv_templates[self.conv_mode].copy().sep_style != SeparatorStyle.TWO else \
92
+ conv_templates[self.conv_mode].copy().sep2
93
+ keywords = [stop_str]
94
+ stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
95
+
96
+ with torch.inference_mode():
97
+ output_ids = model.generate(
98
+ input_ids,
99
+ images=images_tensor,
100
+ do_sample=True,
101
+ temperature=temperature,
102
+ num_beams=1,
103
+ max_new_tokens=max_new_tokens,
104
+ use_cache=True,
105
+ stopping_criteria=[stopping_criteria])
106
+
107
+ input_token_len = input_ids.shape[1]
108
+ n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
109
+ if n_diff_input_output > 0:
110
+ print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
111
+ outputs = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=True)[0]
112
+ outputs = outputs.strip()
113
+ if outputs.endswith(stop_str):
114
+ outputs = outputs[:-len(stop_str)]
115
+ outputs = outputs.strip()
116
+
117
+ print('response', outputs)
118
+ return outputs, state
119
+
120
+
121
+
122
+ title_markdown = ("""
123
+ <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
124
+ <div>
125
+ <h1 >Flash-VStream: Memory-Based Real-Time Understanding for Long Video Streams</h1>
126
+ </div>
127
+ </div>
128
+ <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
129
+ <div style="display:flex; gap: 0.25rem;" align="center">
130
+ <a href="https://invinciblewyq.github.io/vstream-page/"><img src='https://img.shields.io/badge/Project-Page-Green'></a>
131
+ <a href="https://arxiv.org/abs/2406.08085v1"><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
132
+ <a href='https://github.com/IVGSZ/Flash-VStream'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
133
+ </div>
134
+ </div>
135
+ """)
136
+
137
+ block_css = """
138
+ #buttons button {
139
+ min-width: min(120px,100%);
140
+ }
141
+ """