import argparse import os import random import numpy as np import torch import torch.backends.cudnn as cudnn import gradio as gr import argparse import torch from llava.constants import ( IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, IMAGE_PLACEHOLDER, ) from llava.conversation import conv_templates, SeparatorStyle from llava.model.builder import load_pretrained_model from llava.utils import disable_torch_init from llava.mm_utils import ( process_images, tokenizer_image_token, get_model_name_from_path, ) from PIL import Image import requests from PIL import Image from io import BytesIO import re from llava.chat import Chat, conv_llava_v1 # imports modules for registration def parse_args(): parser = argparse.ArgumentParser(description="Demo") parser.add_argument("--model-path", type=str, default="mqt_llava_weight") parser.add_argument("--model-base", type=str, default=None) # parser.add_argument("--image-file", type=str, required=True) # parser.add_argument("--query", type=str, required=True) parser.add_argument("--conv-mode", type=str, default='llava_v1') parser.add_argument("--sep", type=str, default=",") parser.add_argument("--temperature", type=float, default=0) parser.add_argument("--top_p", type=float, default=None) parser.add_argument("--num_beams", type=int, default=1) parser.add_argument("--max_new_tokens", type=int, default=512) parser.add_argument("--num-visual-tokens", type=int, default=256) parser.add_argument("--gpu-id", type=int, default=0) args = parser.parse_args() return args # ======================================== # Model Initialization # ======================================== print('Initializing Chat') args = parse_args() if torch.cuda.is_available(): device='cuda:{}'.format(args.gpu_id) else: device=torch.device('cpu') disable_torch_init() model_name = get_model_name_from_path(args.model_path) tokenizer, model, image_processor, context_len = load_pretrained_model( args.model_path, args.model_base, model_name, device_map=device, device=device ) # vis_processor_cfg = cfg.datasets_cfg.cc_sbu_align.vis_processor.train # vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg) chat = Chat(model, tokenizer, image_processor, args, device=device) print('Initialization Finished') # ======================================== # Gradio Setting # ======================================== def gradio_reset(chat_state, img_list): if chat_state is not None: chat_state.messages = [] if img_list is not None: img_list = [] return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False),gr.update(value="Upload & Start Chat", interactive=True), chat_state, img_list def upload_img(gr_img, text_input, chat_state): if gr_img is None: return None, None, gr.update(interactive=True), chat_state, None chat_state = conv_llava_v1.copy() #CONV_VISION.copy() img_list = [] llm_message = chat.upload_img(gr_img, chat_state, img_list) return gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Chatting", interactive=False), chat_state, img_list def gradio_ask(user_message, chatbot, chat_state): if len(user_message) == 0: return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state chat.ask(user_message, chat_state) chatbot = chatbot + [[user_message, None]] return '', chatbot, chat_state def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature, num_visual_tokens): llm_message = chat.answer(conv=chat_state, img_list=img_list, num_beams=num_beams, temperature=temperature, num_visual_tokens=num_visual_tokens, ) #[0] chatbot[-1][1] = llm_message[0] return chatbot, chat_state, img_list title = """

Demo of MQT-LLaVA

""" description = """

This is the demo of MQT-LLaVA. Upload your images and start chatting!
To use example questions, click example image, hit upload & start chat, and press enter on your keyboard in the chatbox.
Due to limited memory constraint, we only support single turn conversation. To ask multiple questions, hit Restart and upload your image!

""" article = """

""" #TODO show examples below with gr.Blocks() as demo: gr.Markdown(title) gr.Markdown(description) gr.Markdown(article) with gr.Row(): with gr.Column(scale=0.5): image = gr.Image(type="pil") upload_button = gr.Button(value="Upload & Start Chat", interactive=True, variant="primary") clear = gr.Button("Restart 🔄") num_visual_tokens = gr.Slider( minimum=1, maximum=256, value=256, step=1, interactive=True, label="Number of visual tokens", ) temperature = gr.Slider( minimum=0.1, maximum=2.0, value=1.0, step=0.1, interactive=True, label="Temperature", ) num_beams = gr.Slider( minimum=1, maximum=10, value=5, step=1, interactive=True, label="beam search numbers", ) with gr.Column(): chat_state = gr.State() img_list = gr.State() chatbot = gr.Chatbot(label='MQT-LLaVA') text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False) gr.Examples(examples=[ [f"images/extreme_ironing.jpg", "What is unusual about this image?"], [f"images/waterview.jpg", "What are the things I should be cautious about when I visit here?"], ], inputs=[image, text_input]) upload_button.click(upload_img, [image, text_input, chat_state], [image, text_input, upload_button, chat_state, img_list]) text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then( gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature, num_visual_tokens], [chatbot, chat_state, img_list] ) clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list], queue=False) demo.launch(enable_queue=True)