import gradio as gr import os import torch import spaces from llava import conversation as conversation_lib from llava.constants import IMAGE_TOKEN_IDX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN 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 tokenizer_image_token, get_model_name_from_path, process_images from PIL import Image import argparse from transformers import TextIteratorStreamer from threading import Thread import subprocess # Install flash attention, skipping CUDA build if necessary subprocess.run( "pip install flash-attn --no-build-isolation", env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"}, shell=True, ) # os.environ['GRADIO_TEMP_DIR'] = './gradio_tmp' no_change_btn = gr.Button() enable_btn = gr.Button(interactive=True) disable_btn = gr.Button(interactive=False) argparser = argparse.ArgumentParser() argparser.add_argument("--server_name", default="0.0.0.0", type=str) argparser.add_argument("--port", default="6324", type=str) argparser.add_argument("--model-path", default="umd-vt-nyu/clip-evaclip-und-gen-pretrain", type=str) argparser.add_argument("--model-base", type=str, default=None) argparser.add_argument("--num-gpus", type=int, default=1) argparser.add_argument("--conv-mode", type=str, default="llama3") argparser.add_argument("--temperature", type=float, default=0.2) argparser.add_argument("--max-new-tokens", type=int, default=512) argparser.add_argument("--num_frames", type=int, default=16) argparser.add_argument("--load-8bit", action="store_true") argparser.add_argument("--load-4bit", action="store_true") argparser.add_argument("--debug", action="store_true") args = argparser.parse_args() model_path = args.model_path conv_mode = args.conv_mode filt_invalid="cut" model_name = get_model_name_from_path(args.model_path) model_kwargs = { "use_cache": False, "trust_remote_code": True, "torch_dtype": torch.bfloat16, "attn_implementation": "sdpa" } tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, device_map="cuda:0", **model_kwargs) our_chatbot = None def upvote_last_response(state): return ("",) + (disable_btn,) * 3 def downvote_last_response(state): return ("",) + (disable_btn,) * 3 def flag_last_response(state): return ("",) + (disable_btn,) * 3 def clear_history(): state =conv_templates[conv_mode].copy() return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5 def add_text(state, imagebox, textbox, image_process_mode): if state is None: state = conv_templates[conv_mode].copy() if imagebox is not None: textbox = DEFAULT_IMAGE_TOKEN + '\n' + textbox image = Image.open(imagebox).convert('RGB') if imagebox is not None: textbox = (textbox, image, image_process_mode) state.append_message(state.roles[0], textbox) state.append_message(state.roles[1], None) yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) def delete_text(state, image_process_mode): state.messages[-1][-1] = None prev_human_msg = state.messages[-2] if type(prev_human_msg[1]) in (tuple, list): prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode) yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) def regenerate(state, image_process_mode): state.messages[-1][-1] = None prev_human_msg = state.messages[-2] if type(prev_human_msg[1]) in (tuple, list): prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode) state.skip_next = False return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5 @spaces.GPU def generate(state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens): prompt = state.get_prompt() images = state.get_images(return_pil=True) #prompt, image_args = process_image(prompt, images) ori_prompt = prompt num_image_tokens = 0 if images is not None and len(images) > 0: if len(images) > 0: if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN): raise ValueError("Number of images does not match number of tokens in prompt") #images = [load_image_from_base64(image) for image in images] image_sizes = [image.size for image in images] images = process_images(images, image_processor, model.config) if type(images) is list: images = [image.to(model.device, dtype=torch.float16) for image in images] else: images = images.to(model.device, dtype=torch.float16) else: images = None image_sizes = None image_args = {"images": images, "image_sizes": image_sizes} else: images = None image_args = {} max_context_length = getattr(model.config, 'max_position_embeddings', 2048) max_new_tokens = 512 do_sample = True if temperature > 0.001 else False stop_str = state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2 input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_IDX, return_tensors='pt').unsqueeze(0).to(model.device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15) max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens) if max_new_tokens < 1: # yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0" return thread = Thread(target=model.generate, kwargs=dict( inputs=input_ids, do_sample=do_sample, temperature=temperature, top_p=top_p, max_new_tokens=max_new_tokens, streamer=streamer, use_cache=True, pad_token_id=tokenizer.eos_token_id, **image_args )) thread.start() generated_text = '' for new_text in streamer: generated_text += new_text if generated_text.endswith(stop_str): generated_text = generated_text[:-len(stop_str)] state.messages[-1][-1] = generated_text yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn) yield (state, state.to_gradio_chatbot(), "", None) + (enable_btn,) * 5 torch.cuda.empty_cache() txt = gr.Textbox( scale=4, show_label=False, placeholder="Enter text and press enter.", container=False, ) title_markdown = (""" # Florence-llama [[Code](TBD)] [[Model](TBD)] | 📚 [[Arxiv](TBD)]] """) # title_markdown = (""" # # Florence-llama # """) tos_markdown = (""" ### Terms of use By using this service, users are required to agree to the following terms: 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. Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality. """) learn_more_markdown = (""" ### License The service is a research preview intended for non-commercial use only, subject to the. Please contact us if you find any potential violation. """) block_css = """ #buttons button { min-width: min(120px,100%); } """ textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False) with gr.Blocks(title="llava", theme=gr.themes.Default(), css=block_css) as demo: state = gr.State() gr.Markdown(title_markdown) with gr.Row(): with gr.Column(scale=3): imagebox = gr.Image(label="Input Image", type="filepath") image_process_mode = gr.Radio( ["Crop", "Resize", "Pad", "Default"], value="Default", label="Preprocess for non-square image", visible=False) cur_dir = os.path.dirname(os.path.abspath(__file__)) # gr.Examples(examples=[ # [f"{cur_dir}/assets/health-insurance.png", "Under which circumstances do I need to be enrolled in mandatory health insurance if I am an international student?"], # [f"{cur_dir}/assets/leasing-apartment.png", "I don't have any 3rd party renter's insurance now. Do I need to get one for myself?"], # [f"{cur_dir}/assets/nvidia.jpeg", "Who is the person in the middle?"], # [f"{cur_dir}/assets/animal-compare.png", "Are these two pictures showing the same kind of animal?"], # [f"{cur_dir}/assets/georgia-tech.jpeg", "Where is this photo taken?"] # ], inputs=[imagebox, textbox], cache_examples=False) gr.Examples(examples=[ [f"{cur_dir}/assets/animal-compare.png", "Are these two pictures showing the same kind of animal?"] ], inputs=[imagebox, textbox], cache_examples=False) with gr.Accordion("Parameters", open=False) as parameter_row: temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",) top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",) max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",) with gr.Column(scale=8): chatbot = gr.Chatbot( elem_id="chatbot", label="llava Chatbot", height=650, layout="panel", ) with gr.Row(): with gr.Column(scale=8): textbox.render() with gr.Column(scale=1, min_width=50): submit_btn = gr.Button(value="Send", variant="primary") with gr.Row(elem_id="buttons") as button_row: upvote_btn = gr.Button(value="👍 Upvote", interactive=False) downvote_btn = gr.Button(value="👎 Downvote", interactive=False) flag_btn = gr.Button(value="⚠ī¸ Flag", interactive=False) #stop_btn = gr.Button(value="⏚ī¸ Stop Generation", interactive=False) regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=False) clear_btn = gr.Button(value="🗑ī¸ Clear", interactive=False) gr.Markdown(tos_markdown) gr.Markdown(learn_more_markdown) url_params = gr.JSON(visible=False) # Register listeners btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn] upvote_btn.click( upvote_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) downvote_btn.click( downvote_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) flag_btn.click( flag_last_response, [state], [textbox, upvote_btn, downvote_btn, flag_btn] ) clear_btn.click( clear_history, None, [state, chatbot, textbox, imagebox] + btn_list, queue=False ) regenerate_btn.click( delete_text, [state, image_process_mode], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) textbox.submit( add_text, [state, imagebox, textbox, image_process_mode], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) submit_btn.click( add_text, [state, imagebox, textbox, image_process_mode], [state, chatbot, textbox, imagebox] + btn_list, ).then( generate, [state, imagebox, textbox, image_process_mode, temperature, top_p, max_output_tokens], [state, chatbot, textbox, imagebox] + btn_list, ) demo.queue( status_update_rate=10, api_open=False ).launch()