Update chat.py

chat.py

@@ -1,395 +1,205 @@
+"""
+A model worker executes the model.
+"""
 import argparse
-import datetime
 import json
-import os
-import time
-
-import gradio as gr
-import hashlib
-
-from vcoder_llava.vcoder_conversation import (default_conversation, conv_templates,
-                                   SeparatorStyle)
-from vcoder_llava.constants import LOGDIR
-from vcoder_llava.utils import (build_logger, server_error_msg,
-                          violates_moderation, moderation_msg)
-from chat import Chat
-
-
-logger = build_logger("gradio_app", "gradio_web_server.log")
-
-headers = {"User-Agent": "VCoder Client"}
-
-no_change_btn = gr.Button.update()
-enable_btn = gr.Button.update(interactive=True)
-disable_btn = gr.Button.update(interactive=False)
-
-priority = {
-    "vicuna-13b": "aaaaaaa",
-    "koala-13b": "aaaaaab",
-}
-
-
-def get_conv_log_filename():
-    t = datetime.datetime.now()
-    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
-    return name
-
-
-get_window_url_params = """
-function() {
-    const params = new URLSearchParams(window.location.search);
-    url_params = Object.fromEntries(params);
-    console.log(url_params);
-    return url_params;
-    }
-"""
-
-
-def load_demo_refresh_model_list(request: gr.Request):
-    logger.info(f"load_demo. ip: {request.client.host}")
-    state = default_conversation.copy()
-    dropdown_update = gr.Dropdown.update(
-        choices=models,
-        value=models[0] if len(models) > 0 else ""
-    )
-    return state, dropdown_update
-
-
-def vote_last_response(state, vote_type, model_selector, request: gr.Request):
-    with open(get_conv_log_filename(), "a") as fout:
-        data = {
-            "tstamp": round(time.time(), 4),
-            "type": vote_type,
-            "model": model_selector,
-            "state": state.dict(),
-        }
-        fout.write(json.dumps(data) + "\n")
-
-
-def upvote_last_response(state, model_selector, request: gr.Request):
-    vote_last_response(state, "upvote", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-
-def downvote_last_response(state, model_selector, request: gr.Request):
-    vote_last_response(state, "downvote", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-
-def flag_last_response(state, model_selector, request: gr.Request):
-    vote_last_response(state, "flag", model_selector, request)
-    return ("",) + (disable_btn,) * 3
-
-def regenerate(state, image_process_mode, seg_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, prev_human_msg[1][3], seg_process_mode, None, None)
-    state.skip_next = False
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
-
-
-def clear_history(request: gr.Request):
-    state = default_conversation.copy()
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
-
-
-def add_text(state, text, image, image_process_mode, seg, seg_process_mode, depth, depth_process_mode, request: gr.Request):
-    logger.info(f"add_text. len: {len(text)}")
-    if len(text) <= 0 and image is None:
-        state.skip_next = True
-        return (state, state.to_gradio_chatbot(), "", None, None) + (no_change_btn,) * 5
-    if args.moderate:
-        flagged = violates_moderation(text)
-        if flagged:
-            state.skip_next = True
-            return (state, state.to_gradio_chatbot(), moderation_msg, None, None) + (
-                no_change_btn,) * 5
-
-    text = text[:1576]  # Hard cut-off
-    if image is not None:
-        text = text[:1200]  # Hard cut-off for images
-        if '<image>' not in text:
-            text = '<image>\n' + text
-        if seg is not None:
-            if '<seg>' not in text:
-                text = '<seg>\n' + text
-    
-        text = (text, image, image_process_mode, seg, seg_process_mode, None, None)
-        if len(state.get_images(return_pil=True)) > 0:
-            state = default_conversation.copy()
-    state.append_message(state.roles[0], text)
-    state.append_message(state.roles[1], None)
-    state.skip_next = False
-    return (state, state.to_gradio_chatbot(), "", None, None) + (disable_btn,) * 5
-
-
-def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request: gr.Request):
-    start_tstamp = time.time()
-    model_name = model_selector
-
-    if state.skip_next:
-        # This generate call is skipped due to invalid inputs
-        yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
-        return
-
-    if len(state.messages) == state.offset + 2:
-        # First round of conversation
-        if "llava" in model_name.lower():
-            template_name = "llava_v1"
-        new_state = conv_templates[template_name].copy()
-        new_state.append_message(new_state.roles[0], state.messages[-2][1])
-        new_state.append_message(new_state.roles[1], None)
-        state = new_state
-    
-    # Construct prompt
-    prompt = state.get_prompt()
-
-    all_images = state.get_images(return_pil=True)
-    all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
-    for image, hash in zip(all_images, all_image_hash):
-        t = datetime.datetime.now()
-        filename = os.path.join(LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg")
-        if not os.path.isfile(filename):
-            os.makedirs(os.path.dirname(filename), exist_ok=True)
-            image.save(filename)
-    
-    all_segs = state.get_segs(return_pil=True)
-    all_seg_hash = [hashlib.md5(seg.tobytes()).hexdigest() for seg in all_segs]
-    for seg, hash in zip(all_segs, all_seg_hash):
-        t = datetime.datetime.now()
-        filename = os.path.join(LOGDIR, "serve_segs", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{hash}.jpg")
-        if not os.path.isfile(filename):
-            os.makedirs(os.path.dirname(filename), exist_ok=True)
-            seg.save(filename)
-
-    # Make requests
-    pload = {
-        "model": model_name,
-        "prompt": prompt,
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "max_new_tokens": min(int(max_new_tokens), 1536),
-        "stop": state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2,
-        "images": f'List of {len(state.get_images())} images: {all_image_hash}',
-        "segs": f'List of {len(state.get_segs())} segs: {all_seg_hash}',
-    }
-    logger.info(f"==== request ====\n{pload}")
-
-    pload['images'] = state.get_images()
-    pload['segs'] = state.get_segs()
-
-    state.messages[-1][-1] = "▌"
-    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
-
-
-    try:
-        # Stream output
-        response = chat.generate_stream_gate(pload)
-        for chunk in response:
-            if chunk:
-                data = json.loads(chunk.decode())
-                if data["error_code"] == 0:
-                    output = data["text"][len(prompt):].strip()
-                    state.messages[-1][-1] = output + "▌"
-                    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+import torch
+
+from vcoder_llava.utils import server_error_msg
+from vcoder_llava.model.builder import load_pretrained_model
+from vcoder_llava.mm_utils import process_images, load_image_from_base64, tokenizer_seg_token, tokenizer_depth_seg_token, tokenizer_image_token, KeywordsStoppingCriteria
+from vcoder_llava.constants import (
+    IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN,
+    SEG_TOKEN_INDEX, DEFAULT_SEG_TOKEN,
+    DEPTH_TOKEN_INDEX, DEFAULT_DEPTH_TOKEN
+)
+from transformers import TextIteratorStreamer
+
+class Chat:
+    def __init__(self, model_path, model_base, model_name,
+                 load_8bit, load_4bit, device, logger):
+        if model_path.endswith("/"):
+            model_path = model_path[:-1]
+        if model_name is None:
+            model_paths = model_path.split("/")
+            if model_paths[-1].startswith('checkpoint-'):
+                self.model_name = model_paths[-2] + "_" + model_paths[-1]
+            else:
+                self.model_name = model_paths[-1]
+        else:
+            self.model_name = model_name
+
+        self.device = device
+        logger.info(f"Loading the model {self.model_name} ...")
+        self.tokenizer, self.model, self.image_processor, self.seg_image_processor, self.depth_image_processor, self.context_len = load_pretrained_model(
+            model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)
+        self.is_multimodal = 'llava' in self.model_name.lower()
+        self.is_seg = "seg_llava" in self.model_name.lower()
+        self.is_depth = False
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor, seg_image_processor, depth_image_processor = self.tokenizer, self.model, self.image_processor, self.seg_image_processor, self.depth_image_processor
+
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+        segs = params.get("segs", None)
+        depths = params.get("depths", None)
+        num_image_tokens = 0
+        num_seg_tokens = 0
+        num_depth_tokens = 0
+        if images is not None and len(images) > 0 and self.is_multimodal:
+            if len(images) > 0:
+                if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
+                    raise ValueError("Number of images does not match number of <image> tokens in prompt")
+
+                images = [load_image_from_base64(image) for image in images]
+                images = process_images(images, image_processor, model.config)
+
+                if type(images) is list:
+                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
                 else:
-                    output = data["text"] + f" (error_code: {data['error_code']})"
-                    state.messages[-1][-1] = output
-                    yield (state, state.to_gradio_chatbot()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-                    return
-                time.sleep(0.03)
-    except Exception:
-        gr.Warning(server_error_msg)
-        state.messages[-1][-1] = server_error_msg
-        yield (state, state.to_gradio_chatbot()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-        return
-
-    state.messages[-1][-1] = state.messages[-1][-1][:-1]
-    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
-
-    finish_tstamp = time.time()
-    logger.info(f"{output}")
-
-    with open(get_conv_log_filename(), "a") as fout:
-        data = {
-            "tstamp": round(finish_tstamp, 4),
-            "type": "chat",
-            "model": model_name,
-            "start": round(start_tstamp, 4),
-            "finish": round(start_tstamp, 4),
-            "state": state.dict(),
-            "images": all_image_hash,
-            "segs": all_seg_hash,
-            "ip": request.client.host,
-        }
-        fout.write(json.dumps(data) + "\n")
-
-
-title = "<h1 style='margin-bottom: -10px; text-align: center'>VCoder: Versatile Vision Encoders for Multimodal Large Language Models</h1>"
-# style='
-description = "<p style='font-size: 16px; margin: 5px; font-weight: w300; text-align: center'> <a href='https://praeclarumjj3.github.io/' style='text-decoration:none' target='_blank'>Jitesh Jain, </a> <a href='https://jwyang.github.io/' style='text-decoration:none' target='_blank'>Jianwei Yang, <a href='https://www.humphreyshi.com/home' style='text-decoration:none' target='_blank'>Humphrey Shi</a></p>" \
-            + "<p style='font-size: 16px; margin: 5px; font-weight: w600; text-align: center'> <a href='https://praeclarumjj3.github.io/vcoder/' target='_blank'>Project Page</a> | <a href='https://praeclarumjj3.github.io/vcoder/' target='_blank'>Video</a> | <a href='https://arxiv.org/abs/2211.06220' target='_blank'>ArXiv Paper</a> | <a href='https://github.com/SHI-Labs/VCoder' target='_blank'>Github Repo</a></p>" \
-            + "<p style='text-align: center; font-size: 16px; margin: 5px; font-weight: w300;'> [Note: Please click on Regenerate button if you are unsatisfied with the generated response. You may find screenshots of our demo trials <a href='https://github.com/SHI-Labs/VCoder/blob/main/images/' style='text-decoration:none' target='_blank'>here</a>.]</p>" \
-            + "<p style='text-align: center; font-size: 16px; margin: 5px; font-weight: w300;'> [Note: You can obtain segmentation maps for your image using the <a href='https://huggingface.co/spaces/shi-labs/OneFormer' style='text-decoration:none' target='_blank'>OneFormer Demo</a>. Please click on Regenerate button if you are unsatisfied with the generated response. You may find screenshots of our demo trials <a href='https://github.com/SHI-Labs/VCoder/blob/main/images/' style='text-decoration:none' target='_blank'>here</a>.]</p>"
-
-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.
-""")
-
-
-learn_more_markdown = ("""
-### License
-The service is a research preview intended for non-commercial use only, subject to the [License](https://huggingface.co/lmsys/vicuna-7b-v1.5) of Vicuna-v1.5, [License](https://github.com/haotian-liu/LLaVA/blob/main/LICENSE) of LLaVA, [Terms of Use](https://cocodataset.org/#termsofuse) of the COCO dataset, [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.
-""")
-
-block_css = """
-
-#buttons button {
-    min-width: min(120px,100%);
-}
-
-"""
-
-def build_demo(embed_mode):
-    
-    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()
-
-        if not embed_mode:
-            gr.Markdown(title)
-            gr.Markdown(description)
-
-        with gr.Row():
-            with gr.Column(scale=3):
-                with gr.Row(elem_id="model_selector_row"):
-                    model_selector = gr.Dropdown(
-                        choices=models,
-                        value=models[0] if len(models) > 0 else "",
-                        interactive=True,
-                        show_label=False,
-                        container=False)
-
-                # with gr.Row():
-                imagebox = gr.Image(type="pil", label="Image Input")
-                image_process_mode = gr.Radio(
-                    ["Crop", "Resize", "Pad", "Default"],
-                    value="Default",
-                    label="Preprocess for non-square image", visible=False)
-
-                segbox = gr.Image(type="pil", label="Seg Map")
-                seg_process_mode = gr.Radio(
-                    ["Crop", "Resize", "Pad", "Default"],
-                    value="Default",
-                    label="Preprocess for non-square Seg Map", visible=False)
-
-                with gr.Accordion("Parameters", open=False) as parameter_row:
-                    temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.8, step=0.1, interactive=True, label="Temperature",)
-                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, 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="VCoder Chatbot", height=550)
-                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)
-
-        cur_dir = os.path.dirname(os.path.abspath(__file__))
-        gr.Examples(examples=[
-            [f"{cur_dir}/examples/people.jpg", f"{cur_dir}/examples/people_pan.png", "What objects can be seen in the image?", "0.9", "1.0"],
-            [f"{cur_dir}/examples/corgi.jpg", f"{cur_dir}/examples/corgi_pan.png", "What objects can be seen in the image?", "0.6", "0.7"], 
-            [f"{cur_dir}/examples/friends.jpg", f"{cur_dir}/examples/friends_pan.png", "Can you count the number of people in the image?", "0.8", "0.9"], 
-            [f"{cur_dir}/examples/friends.jpg", f"{cur_dir}/examples/friends_pan.png", "What is happening in the image?", "0.8", "0.9"],
-            [f"{cur_dir}/examples/suits.jpg", f"{cur_dir}/examples/suits_pan.png", "What objects can be seen in the image?", "0.5", "0.5"], 
-            [f"{cur_dir}/examples/suits.jpg", f"{cur_dir}/examples/suits_ins.png", "What objects can be seen in the image?", "0.5", "0.5"],
-        ], inputs=[imagebox, segbox, textbox, temperature, top_p])
-        
-        if not embed_mode:
-            gr.Markdown(tos_markdown)
-            gr.Markdown(learn_more_markdown)
-
-        # Register listeners
-        btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
-        upvote_btn.click(upvote_last_response,
-            [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn])
-        downvote_btn.click(downvote_last_response,
-            [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn])
-        flag_btn.click(flag_last_response,
-            [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn])
-        regenerate_btn.click(regenerate, [state, image_process_mode, seg_process_mode],
-            [state, chatbot, textbox, imagebox, segbox] + btn_list).then(
-            http_bot, [state, model_selector, temperature, top_p, max_output_tokens],
-            [state, chatbot] + btn_list)
-        clear_btn.click(clear_history, None, [state, chatbot, textbox, imagebox, segbox] + btn_list)
-
-        textbox.submit(add_text, [state, textbox, imagebox, image_process_mode, segbox, seg_process_mode], [state, chatbot, textbox, imagebox, segbox] + btn_list
-            ).then(http_bot, [state, model_selector, temperature, top_p, max_output_tokens],
-                   [state, chatbot] + btn_list)
-        submit_btn.click(add_text, [state, textbox, imagebox, image_process_mode, segbox, seg_process_mode], [state, chatbot, textbox, imagebox, segbox] + btn_list
-            ).then(http_bot, [state, model_selector, temperature, top_p, max_output_tokens],
-                   [state, chatbot] + btn_list)
-
-        demo.load(load_demo_refresh_model_list, None, [state, model_selector])
-
-    return demo
+                    images = images.to(self.model.device, dtype=torch.float16)
+                
+                replace_token = DEFAULT_IMAGE_TOKEN
+                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
+                num_image_tokens = prompt.count(replace_token) * model.get_vision_tower().num_patches
+
+                if segs is not None and len(segs) > 0 and self.is_seg:
+                    if len(segs) != prompt.count(DEFAULT_SEG_TOKEN):
+                        raise ValueError("Number of segs does not match number of <seg> tokens in prompt")
+                    
+                    segs = [load_image_from_base64(seg) for seg in segs]
+                    segs = process_images(segs, seg_image_processor, model.config)
+                
+                    if type(segs) is list:
+                        segs = [seg.to(self.model.device, dtype=torch.float16) for seg in segs]
+                    else:
+                        segs = segs.to(self.model.device, dtype=torch.float16)
+                    
+                    replace_seg_token = DEFAULT_SEG_TOKEN
+                    prompt = prompt.replace(DEFAULT_SEG_TOKEN, replace_seg_token)
+                    num_seg_tokens = prompt.count(replace_seg_token) * model.get_vision_tower().num_patches
+
+                    if depths is not None and len(depths) > 0 and self.is_depth:
+                        if len(depths) != prompt.count(DEFAULT_DEPTH_TOKEN):
+                            raise ValueError("Number of depths does not match number of <depth> tokens in prompt")
+                        
+                        depths = [load_image_from_base64(depth) for depth in depths]
+                        depths = process_images(depths, depth_image_processor, model.config)
+                    
+                        if type(depths) is list:
+                            depths = [depth.to(self.model.device, dtype=torch.float16) for depth in depths]
+                        else:
+                            depths = depths.to(self.model.device, dtype=torch.float16)
+                        
+                        replace_depth_token = DEFAULT_DEPTH_TOKEN
+                        prompt = prompt.replace(DEFAULT_DEPTH_TOKEN, replace_depth_token)
+                        num_depth_tokens = prompt.count(replace_depth_token) * model.get_vision_tower().num_patches
+                    else:
+                        depths = None
+                else:
+                    segs = None
+                    depths = None
+            else:
+                images = None
+                segs = None
+                depths = None
+            image_args = {"images": images, "segs": segs, "depths": depths}
+        else:
+            images = None
+            segs = None
+            depths = None
+            image_args = {}
+
+        temperature = float(params.get("temperature", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        max_context_length = getattr(model.config, 'max_position_embeddings', 2048)
+        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)
+        stop_str = params.get("stop", None)
+        do_sample = True if temperature > 0.001 else False
+
+        if self.is_seg:
+            if self.is_depth:
+                input_ids = tokenizer_depth_seg_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, SEG_TOKEN_INDEX, DEPTH_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
+            else:
+                input_ids = tokenizer_seg_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, SEG_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
+        else:
+            input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
+        keywords = [stop_str]
+        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
+        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 - num_seg_tokens - num_depth_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
+
+        generated_text = model.generate(
+            inputs=input_ids,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_p=top_p,
+            max_new_tokens=max_new_tokens,
+            streamer=streamer,
+            stopping_criteria=[stopping_criteria],
+            use_cache=True,
+            **image_args
+        )
+        # thread.start()
+
+        generated_text = ori_prompt
+        for new_text in streamer:
+            generated_text += new_text
+            if generated_text.endswith(stop_str):
+                generated_text = generated_text[:-len(stop_str)]
+            yield json.dumps({"text": generated_text, "error_code": 0}).encode()
+
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode()
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode()
+        except Exception as e:
+            print("Caught Unknown Error", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode()
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="shi-labs/vcoder_ds_llava-v1.5-13b")
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str,
+        default="http://localhost:21002")
+    parser.add_argument("--controller-address", type=str,
+        default="http://localhost:21001")
+    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
     parser.add_argument("--model-base", type=str, default=None)
     parser.add_argument("--model-name", type=str)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--multi-modal", action="store_true", help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.")
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
     parser.add_argument("--load-8bit", action="store_true")
     parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--share", action="store_true")
-    parser.add_argument("--moderate", action="store_true")
-    parser.add_argument("--embed", action="store_true")
-    parser.add_argument("--concurrency-count", type=int, default=10)
-    parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--port", type=int)
     args = parser.parse_args()
-    logger.info(f"args: {args}")
-
-    if args.model_name is None:
-        model_paths = args.model_path.split("/")
-        if model_paths[-1].startswith('checkpoint-'):
-            model_name = model_paths[-2] + "_" + model_paths[-1]
-        else:
-            model_name = model_paths[-1]
-    else:
-        model_name = args.model_name
-
-    models = [model_name]
-    chat = Chat(
-        args.model_path,
-        args.model_base,
-        args.model_name,
-        args.load_8bit,
-        args.load_4bit,
-        args.device,
-        logger
-    )
-
-    logger.info(args)
-    demo = build_demo(args.embed)
-    demo.queue(
-        concurrency_count=args.concurrency_count,
-        api_open=False
-    ).launch(
-        server_name=args.host,
-        server_port=args.port,
-        share=args.share
-    )