LanguageBind/Video-LLaVA
234
1import shutil2import subprocess3 4import torch5import gradio as gr6from fastapi import FastAPI7import os8from PIL import Image9import tempfile10from decord import VideoReader, cpu11from transformers import TextStreamer12 13from llava.constants import DEFAULT_X_TOKEN, X_TOKEN_INDEX14from llava.conversation import conv_templates, SeparatorStyle, Conversation15from llava.serve.gradio_utils import Chat, tos_markdown, learn_more_markdown, title_markdown, block_css16 17 18def save_image_to_local(image):19 filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.jpg')20 image = Image.open(image)21 image.save(filename)22 # print(filename)23 return filename24 25 26def save_video_to_local(video_path):27 filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.mp4')28 shutil.copyfile(video_path, filename)29 return filename30 31 32def generate(image1, video, textbox_in, first_run, state, state_, images_tensor):33 flag = 134 if not textbox_in:35 if len(state_.messages) > 0:36 textbox_in = state_.messages[-1][1]37 state_.messages.pop(-1)38 flag = 039 else:40 return "Please enter instruction"41 42 image1 = image1 if image1 else "none"43 video = video if video else "none"44 # assert not (os.path.exists(image1) and os.path.exists(video))45 46 if type(state) is not Conversation:47 state = conv_templates[conv_mode].copy()48 state_ = conv_templates[conv_mode].copy()49 images_tensor = [[], []]50 51 first_run = False if len(state.messages) > 0 else True52 53 text_en_in = textbox_in.replace("picture", "image")54 55 # images_tensor = [[], []]56 image_processor = handler.image_processor57 if os.path.exists(image1) and not os.path.exists(video):58 tensor = image_processor.preprocess(image1, return_tensors='pt')['pixel_values'][0]59 # print(tensor.shape)60 tensor = tensor.to(handler.model.device, dtype=dtype)61 images_tensor[0] = images_tensor[0] + [tensor]62 images_tensor[1] = images_tensor[1] + ['image']63 print(torch.cuda.memory_allocated())64 print(torch.cuda.max_memory_allocated())65 video_processor = handler.video_processor66 if not os.path.exists(image1) and os.path.exists(video):67 tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]68 # print(tensor.shape)69 tensor = tensor.to(handler.model.device, dtype=dtype)70 images_tensor[0] = images_tensor[0] + [tensor]71 images_tensor[1] = images_tensor[1] + ['video']72 print(torch.cuda.memory_allocated())73 print(torch.cuda.max_memory_allocated())74 if os.path.exists(image1) and os.path.exists(video):75 tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]76 # print(tensor.shape)77 tensor = tensor.to(handler.model.device, dtype=dtype)78 images_tensor[0] = images_tensor[0] + [tensor]79 images_tensor[1] = images_tensor[1] + ['video']80 81 82 tensor = image_processor.preprocess(image1, return_tensors='pt')['pixel_values'][0]83 # print(tensor.shape)84 tensor = tensor.to(handler.model.device, dtype=dtype)85 images_tensor[0] = images_tensor[0] + [tensor]86 images_tensor[1] = images_tensor[1] + ['image']87 print(torch.cuda.memory_allocated())88 print(torch.cuda.max_memory_allocated())89 90 91 92 if os.path.exists(image1) and not os.path.exists(video):93 text_en_in = DEFAULT_X_TOKEN['IMAGE'] + '\n' + text_en_in94 if not os.path.exists(image1) and os.path.exists(video):95 text_en_in = DEFAULT_X_TOKEN['VIDEO'] + '\n' + text_en_in96 if os.path.exists(image1) and os.path.exists(video):97 text_en_in = DEFAULT_X_TOKEN['VIDEO'] + '\n' + text_en_in + '\n' + DEFAULT_X_TOKEN['IMAGE']98 99 text_en_out, state_ = handler.generate(images_tensor, text_en_in, first_run=first_run, state=state_)100 state_.messages[-1] = (state_.roles[1], text_en_out)101 102 text_en_out = text_en_out.split('#')[0]103 textbox_out = text_en_out104 105 show_images = ""106 if os.path.exists(image1):107 filename = save_image_to_local(image1)108 show_images += f'<img src="./file={filename}" style="display: inline-block;width: 250px;max-height: 400px;">'109 if os.path.exists(video):110 filename = save_video_to_local(video)111 show_images += f'<video controls playsinline width="500" style="display: inline-block;" src="./file={filename}"></video>'112 113 if flag:114 state.append_message(state.roles[0], textbox_in + "\n" + show_images)115 state.append_message(state.roles[1], textbox_out)116 torch.cuda.empty_cache()117 return (state, state_, state.to_gradio_chatbot(), False, gr.update(value=None, interactive=True), images_tensor, gr.update(value=image1 if os.path.exists(image1) else None, interactive=True), gr.update(value=video if os.path.exists(video) else None, interactive=True))118 119def regenerate(state, state_):120 state.messages.pop(-1)121 state_.messages.pop(-1)122 if len(state.messages) > 0:123 return state, state_, state.to_gradio_chatbot(), False124 return (state, state_, state.to_gradio_chatbot(), True)125 126 127def clear_history(state, state_):128 state = conv_templates[conv_mode].copy()129 state_ = conv_templates[conv_mode].copy()130 return (gr.update(value=None, interactive=True),131 gr.update(value=None, interactive=True),\132 gr.update(value=None, interactive=True),\133 True, state, state_, state.to_gradio_chatbot(), [[], []])134 135 136 137conv_mode = "llava_v1"138model_path = 'LanguageBind/Video-LLaVA-7B'139device = 'cuda'140load_8bit = False141load_4bit = True142dtype = torch.float16143handler = Chat(model_path, conv_mode=conv_mode, load_8bit=load_8bit, load_4bit=load_8bit, device=device)144# handler.model.to(dtype=dtype)145if not os.path.exists("temp"):146 os.makedirs("temp")147 148print(torch.cuda.memory_allocated())149print(torch.cuda.max_memory_allocated())150 151app = FastAPI()152 153textbox = gr.Textbox(154 show_label=False, placeholder="Enter text and press ENTER", container=False155 )156with gr.Blocks(title='Video-LLaVA🚀', theme=gr.themes.Default(), css=block_css) as demo:157 gr.Markdown(title_markdown)158 state = gr.State()159 state_ = gr.State()160 first_run = gr.State()161 images_tensor = gr.State()162 163 with gr.Row():164 with gr.Column(scale=3):165 image1 = gr.Image(label="Input Image", type="filepath")166 video = gr.Video(label="Input Video")167 168 cur_dir = os.path.dirname(os.path.abspath(__file__))169 gr.Examples(170 examples=[171 [172 f"{cur_dir}/examples/extreme_ironing.jpg",173 "What is unusual about this image?",174 ],175 [176 f"{cur_dir}/examples/waterview.jpg",177 "What are the things I should be cautious about when I visit here?",178 ],179 [180 f"{cur_dir}/examples/desert.jpg",181 "If there are factual errors in the questions, point it out; if not, proceed answering the question. What’s happening in the desert?",182 ],183 ],184 inputs=[image1, textbox],185 )186 187 with gr.Column(scale=7):188 chatbot = gr.Chatbot(label="Video-LLaVA", bubble_full_width=True).style(height=750)189 with gr.Row():190 with gr.Column(scale=8):191 textbox.render()192 with gr.Column(scale=1, min_width=50):193 submit_btn = gr.Button(194 value="Send", variant="primary", interactive=True195 )196 with gr.Row(elem_id="buttons") as button_row:197 upvote_btn = gr.Button(value="👍 Upvote", interactive=True)198 downvote_btn = gr.Button(value="👎 Downvote", interactive=True)199 flag_btn = gr.Button(value="⚠️ Flag", interactive=True)200 # stop_btn = gr.Button(value="⏹️ Stop Generation", interactive=False)201 regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=True)202 clear_btn = gr.Button(value="🗑️ Clear history", interactive=True)203 204 with gr.Row():205 gr.Examples(206 examples=[207 [208 f"{cur_dir}/examples/sample_img_8.png",209 f"{cur_dir}/examples/sample_demo_8.mp4",210 "Are the image and the video depicting the same place?",211 ],212 [213 f"{cur_dir}/examples/sample_img_22.png",214 f"{cur_dir}/examples/sample_demo_22.mp4",215 "Are the instruments in the pictures used in the video?",216 ],217 [218 f"{cur_dir}/examples/sample_img_13.png",219 f"{cur_dir}/examples/sample_demo_13.mp4",220 "Does the flag in the image appear in the video?",221 ],222 ],223 inputs=[image1, video, textbox],224 )225 gr.Examples(226 examples=[227 [228 f"{cur_dir}/examples/sample_demo_1.mp4",229 "Why is this video funny?",230 ],231 [232 f"{cur_dir}/examples/sample_demo_7.mp4",233 "Create a short fairy tale with a moral lesson inspired by the video.",234 ],235 [236 f"{cur_dir}/examples/sample_demo_8.mp4",237 "Where is this video taken from? What place/landmark is shown in the video?",238 ],239 [240 f"{cur_dir}/examples/sample_demo_12.mp4",241 "What does the woman use to split the logs and how does she do it?",242 ],243 [244 f"{cur_dir}/examples/sample_demo_18.mp4",245 "Describe the video in detail.",246 ],247 [248 f"{cur_dir}/examples/sample_demo_22.mp4",249 "Describe the activity in the video.",250 ],251 ],252 inputs=[video, textbox],253 )254 gr.Markdown(tos_markdown)255 gr.Markdown(learn_more_markdown)256 257 submit_btn.click(generate, [image1, video, textbox, first_run, state, state_, images_tensor],258 [state, state_, chatbot, first_run, textbox, images_tensor, image1, video])259 260 regenerate_btn.click(regenerate, [state, state_], [state, state_, chatbot, first_run]).then(261 generate, [image1, video, textbox, first_run, state, state_, images_tensor], [state, state_, chatbot, first_run, textbox, images_tensor, image1, video])262 263 clear_btn.click(clear_history, [state, state_],264 [image1, video, textbox, first_run, state, state_, chatbot, images_tensor])265 266# app = gr.mount_gradio_app(app, demo, path="/")267demo.launch()268 269 270# uvicorn llava.serve.gradio_web_server:app271