Spaces:
Running
Running
| import torch, torchvision, transformers, collections | |
| from dataclasses import asdict | |
| from torchvision.io import read_video | |
| from models import build_model_and_tokenizer, parse_args, fast_greedy_generate | |
| logger = transformers.logging.get_logger('liveinfer') | |
| # python -m demo.cli --resume_from_checkpoint ... | |
| class LiveInfer: | |
| def __init__(self, ) -> None: | |
| args = parse_args() | |
| args.resume_from_checkpoint = 'live1+_aug_2e/' | |
| args.attn_implementation = 'sdpa' | |
| self.model, self.tokenizer = build_model_and_tokenizer(is_training=False, set_vision_inside=True, **asdict(args)) | |
| self.model.to('cuda') | |
| # visual | |
| self.hidden_size = self.model.config.hidden_size | |
| self.frame_fps = args.frame_fps | |
| self.frame_interval = 1 / self.frame_fps | |
| self.frame_resolution = self.model.config.frame_resolution | |
| self.frame_num_tokens = self.model.config.frame_num_tokens | |
| self.frame_v_placeholder = self.model.config.v_placeholder * self.frame_num_tokens | |
| self.frame_token_interval_id = self.model.config.frame_token_interval_id | |
| self.frame_placeholder_ids = torch.tensor(self.model.config.v_placeholder_id).repeat(self.model.config.frame_num_tokens).reshape(1,-1) | |
| # generation | |
| self.system_prompt = args.system_prompt | |
| self.inplace_output_ids = torch.zeros(1, 100, device='cuda', dtype=torch.long) | |
| self.frame_token_interval_threshold = 0.725 | |
| self.eos_token_id = self.model.config.eos_token_id | |
| self._start_ids = self.tokenizer.apply_chat_template([{'role': 'system', 'content': self.system_prompt}], add_stream_prompt=True, return_tensors='pt').to('cuda') | |
| self._added_stream_prompt_ids = self.tokenizer.apply_chat_template([{}], add_stream_prompt=True, return_tensors='pt').to('cuda') | |
| self._added_stream_generation_ids = self.tokenizer.apply_chat_template([{}], add_stream_generation_prompt=True, return_tensors='pt').to('cuda') | |
| # app | |
| self.reset() | |
| def _call_for_response(self, video_time, query): | |
| if query is not None: | |
| self.last_ids = self.tokenizer.apply_chat_template([{'role': 'user', 'content': query}], add_stream_query_prompt=True, add_generation_prompt=True, return_tensors='pt').to('cuda') | |
| else: | |
| assert self.last_ids == 933, f'{self.last_ids} != 933' # HACK, 933 = ]\n | |
| self.last_ids = self._added_stream_generation_ids | |
| inputs_embeds = self.model.get_input_embeddings()(self.last_ids) | |
| output_ids, self.past_key_values = fast_greedy_generate(model=self.model, inputs_embeds=inputs_embeds, past_key_values=self.past_key_values, eos_token_id=self.eos_token_id, inplace_output_ids=self.inplace_output_ids) | |
| self.last_ids = output_ids[:, -1:] | |
| if query: | |
| query = f'(Video Time = {video_time}s) User: {query}' | |
| response = f'(Video Time = {video_time}s) Assistant:{self.tokenizer.decode(output_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)}' | |
| return query, response | |
| def _call_for_streaming(self, ): | |
| while self.frame_embeds_queue: | |
| # 1. if query is before next frame, response | |
| if self.query_queue and self.frame_embeds_queue[0][0] > self.query_queue[0][0]: | |
| video_time, query = self.query_queue.popleft() | |
| return video_time, query | |
| video_time, frame_embeds = self.frame_embeds_queue.popleft() | |
| if not self.past_key_values: | |
| self.last_ids = self._start_ids | |
| elif self.last_ids == self.eos_token_id: | |
| self.last_ids = torch.cat([self.last_ids, self._added_stream_prompt_ids], dim=1) | |
| inputs_embeds = torch.cat([ | |
| self.model.get_input_embeddings()(self.last_ids).view(1, -1, self.hidden_size), | |
| frame_embeds.view(1, -1, self.hidden_size), | |
| ], dim=1) | |
| outputs = self.model(inputs_embeds=inputs_embeds, use_cache=True, past_key_values=self.past_key_values) | |
| self.past_key_values = outputs.past_key_values | |
| # 2. if the same time, response after frame at that time | |
| if self.query_queue and video_time >= self.query_queue[0][0]: | |
| video_time, query = self.query_queue.popleft() | |
| return video_time, query | |
| # 3. if the next is frame but next is not interval, then response | |
| next_score = outputs.logits[:,-1:].softmax(dim=-1) | |
| if next_score[:,:,self.frame_token_interval_id] < self.frame_token_interval_threshold: | |
| next_score[:,:,self.frame_token_interval_id].zero_() | |
| self.last_ids = next_score.argmax(dim=-1) | |
| if self.last_ids != self.frame_token_interval_id: | |
| return video_time, None | |
| return None, None | |
| def reset(self, ): | |
| self.query_queue = collections.deque() | |
| self.frame_embeds_queue = collections.deque() | |
| self.video_time = 0 | |
| self.last_frame_idx = -1 | |
| self.video_tensor = None | |
| self.last_ids = torch.tensor([[]], device='cuda', dtype=torch.long) | |
| self.past_key_values = None | |
| def input_query_stream(self, query, history=None, video_time=None): | |
| if video_time is None: | |
| self.query_queue.append((self.video_time, query)) | |
| else: | |
| self.query_queue.append((video_time, query)) | |
| if not self.past_key_values: | |
| return f'(NOTE: No video stream here. Please select or upload a video. Then the assistant will answer "{query} (at {self.video_time}s)" in the video stream)' | |
| return f'(NOTE: Received "{query}" (at {self.video_time}s). Please wait until previous frames have been processed)' | |
| def input_video_stream(self, video_time): | |
| frame_idx = int(video_time * self.frame_fps) | |
| if frame_idx > self.last_frame_idx: | |
| ranger = range(self.last_frame_idx + 1, frame_idx + 1) | |
| frames_embeds = self.model.visual_embed(self.video_tensor[ranger]).split(self.frame_num_tokens) | |
| self.frame_embeds_queue.extend([(r / self.frame_fps, frame_embeds) for r, frame_embeds in zip(ranger, frames_embeds)]) | |
| self.last_frame_idx = frame_idx | |
| self.video_time = video_time | |
| def load_video(self, video_path): | |
| self.video_tensor = read_video(video_path, pts_unit='sec', output_format='TCHW')[0].to('cuda') | |
| self.num_video_frames = self.video_tensor.size(0) | |
| self.video_duration = self.video_tensor.size(0) / self.frame_fps | |
| logger.warning(f'{video_path} -> {self.video_tensor.shape}, {self.frame_fps} FPS') | |
| def __call__(self, ): | |
| while not self.frame_embeds_queue: | |
| continue | |
| video_time, query = self._call_for_streaming() | |
| response = None | |
| if video_time is not None: | |
| query, response = self._call_for_response(video_time, query) | |
| return query, response |