prathammk/Conversational-Image-Video-Recognition-ChatBot
0
1# from .demo_modelpart import InferenceDemo2import gradio as gr3import os4 5# import time6import cv27 8 9# import copy10import torch11 12import spaces13import numpy as np14 15from llava import conversation as conversation_lib16from llava.constants import DEFAULT_IMAGE_TOKEN17 18 19from llava.constants import (20 IMAGE_TOKEN_INDEX,21 DEFAULT_IMAGE_TOKEN,22 DEFAULT_IM_START_TOKEN,23 DEFAULT_IM_END_TOKEN,24)25from llava.conversation import conv_templates, SeparatorStyle26from llava.model.builder import load_pretrained_model27from llava.utils import disable_torch_init28from llava.mm_utils import (29 tokenizer_image_token,30 get_model_name_from_path,31 KeywordsStoppingCriteria,32)33 34from PIL import Image35 36import requests37from PIL import Image38from io import BytesIO39from transformers import TextStreamer40 41import gradio as gr42import gradio_client43import subprocess44import sys45 46def install_gradio_4_35_0():47 current_version = gr.__version__48 if current_version != "4.35.0":49 print(f"Current Gradio version: {current_version}")50 print("Installing Gradio 4.35.0...")51 subprocess.check_call([sys.executable, "-m", "pip", "install", "gradio==4.35.0", "--force-reinstall"])52 print("Gradio 4.35.0 installed successfully.")53 else:54 print("Gradio 4.35.0 is already installed.")55 56# Call the function to install Gradio 4.35.0 if needed57install_gradio_4_35_0()58 59import gradio as gr60import gradio_client61print(f"Gradio version: {gr.__version__}")62print(f"Gradio-client version: {gradio_client.__version__}")63 64class InferenceDemo(object):65 def __init__(66 self, args, model_path, tokenizer, model, image_processor, context_len67 ) -> None:68 disable_torch_init()69 70 self.tokenizer, self.model, self.image_processor, self.context_len = (71 tokenizer,72 model,73 image_processor,74 context_len,75 )76 77 if "llama-2" in model_name.lower():78 conv_mode = "llava_llama_2"79 elif "v1" in model_name.lower():80 conv_mode = "llava_v1"81 elif "mpt" in model_name.lower():82 conv_mode = "mpt"83 elif "qwen" in model_name.lower():84 conv_mode = "qwen_1_5"85 else:86 conv_mode = "llava_v0"87 88 if args.conv_mode is not None and conv_mode != args.conv_mode:89 print(90 "[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}".format(91 conv_mode, args.conv_mode, args.conv_mode92 )93 )94 else:95 args.conv_mode = conv_mode96 self.conv_mode = conv_mode97 self.conversation = conv_templates[args.conv_mode].copy()98 self.num_frames = args.num_frames99 100 101def is_valid_video_filename(name):102 video_extensions = ["avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg"]103 104 ext = name.split(".")[-1].lower()105 106 if ext in video_extensions:107 return True108 else:109 return False110 111 112def sample_frames(video_file, num_frames):113 video = cv2.VideoCapture(video_file)114 total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))115 interval = total_frames // num_frames116 frames = []117 for i in range(total_frames):118 ret, frame = video.read()119 pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))120 if not ret:121 continue122 if i % interval == 0:123 frames.append(pil_img)124 video.release()125 return frames126 127 128def load_image(image_file):129 if image_file.startswith("http") or image_file.startswith("https"):130 response = requests.get(image_file)131 if response.status_code == 200:132 image = Image.open(BytesIO(response.content)).convert("RGB")133 else:134 print("failed to load the image")135 else:136 print("Load image from local file")137 print(image_file)138 image = Image.open(image_file).convert("RGB")139 140 return image141 142 143def clear_history(history):144 145 our_chatbot.conversation = conv_templates[our_chatbot.conv_mode].copy()146 147 return None148 149 150def clear_response(history):151 for index_conv in range(1, len(history)):152 # loop until get a text response from our model.153 conv = history[-index_conv]154 if not (conv[0] is None):155 break156 question = history[-index_conv][0]157 history = history[:-index_conv]158 return history, question159 160 161# def print_like_dislike(x: gr.LikeData):162# print(x.index, x.value, x.liked)163 164 165def add_message(history, message):166 # history=[]167 global our_chatbot168 if len(history) == 0:169 our_chatbot = InferenceDemo(170 args, model_path, tokenizer, model, image_processor, context_len171 )172 173 for x in message["files"]:174 history.append(((x,), None))175 if message["text"] is not None:176 history.append((message["text"], None))177 return history, gr.MultimodalTextbox(value=None, interactive=False)178 179 180@spaces.GPU181def bot(history):182 text = history[-1][0]183 images_this_term = []184 text_this_term = ""185 # import pdb;pdb.set_trace()186 num_new_images = 0187 for i, message in enumerate(history[:-1]):188 if type(message[0]) is tuple:189 images_this_term.append(message[0][0])190 if is_valid_video_filename(message[0][0]):191 num_new_images += our_chatbot.num_frames192 else:193 num_new_images += 1194 else:195 num_new_images = 0196 197 # for message in history[-i-1:]:198 # images_this_term.append(message[0][0])199 200 assert len(images_this_term) > 0, "must have an image"201 # image_files = (args.image_file).split(',')202 # image = [load_image(f) for f in images_this_term if f]203 image_list = []204 for f in images_this_term:205 if is_valid_video_filename(f):206 image_list += sample_frames(f, our_chatbot.num_frames)207 else:208 image_list.append(load_image(f))209 image_tensor = [210 our_chatbot.image_processor.preprocess(f, return_tensors="pt")["pixel_values"][211 0212 ]213 .half()214 .to(our_chatbot.model.device)215 for f in image_list216 ]217 218 image_tensor = torch.stack(image_tensor)219 image_token = DEFAULT_IMAGE_TOKEN * num_new_images220 # if our_chatbot.model.config.mm_use_im_start_end:221 # inp = DEFAULT_IM_START_TOKEN + image_token + DEFAULT_IM_END_TOKEN + "\n" + inp222 # else:223 inp = text224 inp = image_token + "\n" + inp225 our_chatbot.conversation.append_message(our_chatbot.conversation.roles[0], inp)226 # image = None227 our_chatbot.conversation.append_message(our_chatbot.conversation.roles[1], None)228 prompt = our_chatbot.conversation.get_prompt()229 230 input_ids = (231 tokenizer_image_token(232 prompt, our_chatbot.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"233 )234 .unsqueeze(0)235 .to(our_chatbot.model.device)236 )237 stop_str = (238 our_chatbot.conversation.sep239 if our_chatbot.conversation.sep_style != SeparatorStyle.TWO240 else our_chatbot.conversation.sep2241 )242 keywords = [stop_str]243 stopping_criteria = KeywordsStoppingCriteria(244 keywords, our_chatbot.tokenizer, input_ids245 )246 streamer = TextStreamer(247 our_chatbot.tokenizer, skip_prompt=True, skip_special_tokens=True248 )249 print(our_chatbot.model.device)250 print(input_ids.device)251 print(image_tensor.device)252 # import pdb;pdb.set_trace()253 with torch.inference_mode():254 output_ids = our_chatbot.model.generate(255 input_ids,256 images=image_tensor,257 do_sample=True,258 temperature=0.2,259 max_new_tokens=1024,260 streamer=streamer,261 use_cache=False,262 stopping_criteria=[stopping_criteria],263 )264 265 outputs = our_chatbot.tokenizer.decode(output_ids[0]).strip()266 if outputs.endswith(stop_str):267 outputs = outputs[: -len(stop_str)]268 our_chatbot.conversation.messages[-1][-1] = outputs269 270 history[-1] = [text, outputs]271 272 return history273 274 275txt = gr.Textbox(276 scale=4,277 show_label=False,278 placeholder="Enter text and press enter.",279 container=False,280)281 282with gr.Blocks(283 css=".message-wrap.svelte-1lcyrx4>div.svelte-1lcyrx4 img {min-width: 40px}",284) as demo:285 286 # Informations287 title_markdown = """288 # LLaVA-NeXT Interleave289 [[Blog]](https://llava-vl.github.io/blog/2024-06-16-llava-next-interleave/) [[Code]](https://github.com/LLaVA-VL/LLaVA-NeXT) [[Model]](https://huggingface.co/lmms-lab/llava-next-interleave-7b)290 Note: The internleave checkpoint is updated (Date: Jul. 24, 2024), the wrong checkpiont is used before.291 """292 tos_markdown = """293 ### TODO!. Terms of use294 By using this service, users are required to agree to the following terms:295 The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.296 Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.297 For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.298 """299 learn_more_markdown = """300 ### TODO!. License301 The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.302 """303 models = [304 "LLaVA-Interleave-7B",305 ]306 cur_dir = os.path.dirname(os.path.abspath(__file__))307 gr.Markdown(title_markdown)308 with gr.Column():309 with gr.Row():310 chatbot = gr.Chatbot([], elem_id="chatbot", bubble_full_width=False)311 312 with gr.Row():313 upvote_btn = gr.Button(value="👍 Upvote", interactive=True)314 downvote_btn = gr.Button(value="👎 Downvote", interactive=True)315 flag_btn = gr.Button(value="⚠️ Flag", interactive=True)316 # stop_btn = gr.Button(value="⏹️ Stop Generation", interactive=True)317 regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=True)318 clear_btn = gr.Button(value="🗑️ Clear history", interactive=True)319 320 chat_input = gr.MultimodalTextbox(321 interactive=True,322 file_types=["image", "video"],323 placeholder="Enter message or upload file...",324 show_label=False,325 )326 327 print(cur_dir)328 gr.Examples(329 examples=[330 # [331 # {332 # "text": "<image> <image> <image> Which image shows a different mood of character from the others?",333 # "files": [f"{cur_dir}/examples/examples_image12.jpg", f"{cur_dir}/examples/examples_image13.jpg", f"{cur_dir}/examples/examples_image14.jpg"]334 # },335 # {336 # "text": "Please pay attention to the movement of the object from the first image to the second image, then write a HTML code to show this movement.",337 # "files": [338 # f"{cur_dir}/examples/code1.jpeg",339 # f"{cur_dir}/examples/code2.jpeg",340 # ],341 # }342 # ],343 [344 {345 "files": [346 f"{cur_dir}/examples/shub.jpg",347 f"{cur_dir}/examples/shuc.jpg",348 f"{cur_dir}/examples/shud.jpg",349 ],350 "text": "what is fun about the images?",351 }352 ],353 [354 {355 "files": [356 f"{cur_dir}/examples/iphone-15-price-1024x576.jpg",357 f"{cur_dir}/examples/dynamic-island-1024x576.jpg",358 f"{cur_dir}/examples/iphone-15-colors-1024x576.jpg",359 f"{cur_dir}/examples/Iphone-15-Usb-c-charger-1024x576.jpg",360 f"{cur_dir}/examples/A-17-processors-1024x576.jpg",361 ],362 "text": "The images are the PPT of iPhone 15 review. can you summarize the main information?",363 }364 ],365 [366 {367 "files": [368 f"{cur_dir}/examples/fangao3.jpeg",369 f"{cur_dir}/examples/fangao2.jpeg",370 f"{cur_dir}/examples/fangao1.jpeg",371 ],372 "text": "Do you kown who draw these paintings?",373 }374 ],375 [376 {377 "files": [378 f"{cur_dir}/examples/oprah-winfrey-resume.png",379 f"{cur_dir}/examples/steve-jobs-resume.jpg",380 ],381 "text": "Hi, there are two candidates, can you provide a brief description for each of them for me?",382 }383 ],384 [385 {386 "files": [387 f"{cur_dir}/examples/original_bench.jpeg",388 f"{cur_dir}/examples/changed_bench.jpeg",389 ],390 "text": "How to edit image1 to make it look like image2?",391 }392 ],393 [394 {395 "files": [396 f"{cur_dir}/examples/twitter2.jpeg",397 f"{cur_dir}/examples/twitter3.jpeg",398 f"{cur_dir}/examples/twitter4.jpeg",399 ],400 "text": "Please write a twitter blog post with the images.",401 }402 ]403 404 ],405 inputs=[chat_input],406 label="Compare images: "407 )408 409 chat_msg = chat_input.submit(410 add_message, [chatbot, chat_input], [chatbot, chat_input]411 )412 bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")413 bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])414 415 # chatbot.like(print_like_dislike, None, None)416 clear_btn.click(417 fn=clear_history, inputs=[chatbot], outputs=[chatbot], api_name="clear_all"418 )419 420 421demo.queue()422 423if __name__ == "__main__":424 import argparse425 426 argparser = argparse.ArgumentParser()427 argparser.add_argument("--server_name", default="0.0.0.0", type=str)428 argparser.add_argument("--port", default="6123", type=str)429 argparser.add_argument(430 "--model_path", default="lmms-lab/llava-next-interleave-qwen-7b", type=str431 )432 # argparser.add_argument("--model-path", type=str, default="facebook/opt-350m")433 argparser.add_argument("--model-base", type=str, default=None)434 argparser.add_argument("--num-gpus", type=int, default=1)435 argparser.add_argument("--conv-mode", type=str, default=None)436 argparser.add_argument("--temperature", type=float, default=0.2)437 argparser.add_argument("--max-new-tokens", type=int, default=512)438 argparser.add_argument("--num_frames", type=int, default=16)439 argparser.add_argument("--load-8bit", action="store_true")440 argparser.add_argument("--load-4bit", action="store_true")441 argparser.add_argument("--debug", action="store_true")442 443 args = argparser.parse_args()444 445 model_path = args.model_path446 filt_invalid = "cut"447 model_name = get_model_name_from_path(args.model_path)448 tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit)449 model=model.to(torch.device('cuda'))450 our_chatbot = None451 demo.launch()452