Fix ZeroGPT issue.

app.py

@@ -3,8 +3,6 @@ import datetime
 import hashlib
 import json
 import os
-import subprocess
-import sys
 import time
 
 import gradio as gr
@@ -13,22 +11,17 @@ import requests
 from constants import LOGDIR
 from conversation import (default_conversation, conv_templates,
                                    SeparatorStyle)
-from utils import (build_logger, server_error_msg,
-    violates_moderation, moderation_msg)
+from utils import (build_logger, server_error_msg)
 
 
 logger = build_logger("gradio_web_server", "gradio_web_server.log")
 
-headers = {"User-Agent": "LLaVA Client"}
+from model_worker import ModelWorker
 
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
 disable_btn = gr.Button(interactive=False)
 
-priority = {
-    "vicuna-13b": "aaaaaaa",
-    "koala-13b": "aaaaaab",
-}
 
 
 def get_conv_log_filename():
@@ -36,17 +29,6 @@ def get_conv_log_filename():
     name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
     return name
 
-
-def get_model_list():
-    ret = requests.post(args.controller_url + "/refresh_all_workers")
-    assert ret.status_code == 200
-    ret = requests.post(args.controller_url + "/list_models")
-    models = ret.json()["models"]
-    models.sort(key=lambda x: priority.get(x, x))
-    logger.info(f"Models: {models}")
-    return models
-
-
 get_window_url_params = """
 function() {
     const params = new URLSearchParams(window.location.search);
@@ -60,24 +42,11 @@ function() {
 def load_demo(url_params, request: gr.Request):
     logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
 
+    global worker
     dropdown_update = gr.Dropdown(visible=True)
-    if "model" in url_params:
-        model = url_params["model"]
-        if model in models:
-            dropdown_update = gr.Dropdown(value=model, visible=True)
-
-    state = default_conversation.copy()
-    return state, dropdown_update
-
+    worker = ModelWorker(model_path, None, model_name, True, lora_path)
 
-def load_demo_refresh_model_list(request: gr.Request):
-    logger.info(f"load_demo. ip: {request.client.host}")
-    models = get_model_list()
     state = default_conversation.copy()
-    dropdown_update = gr.Dropdown(
-        choices=models,
-        value=models[0] if len(models) > 0 else ""
-    )
     return state, dropdown_update
 
 
@@ -132,12 +101,6 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
     if len(text) <= 0 and image is None:
         state.skip_next = True
         return (state, state.to_gradio_chatbot(), "", 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) + (
-                no_change_btn,) * 5
 
     text = text[:1536]  # Hard cut-off
     if image is not None:
@@ -152,7 +115,6 @@ def add_text(state, text, image, image_process_mode, request: gr.Request):
     state.skip_next = False
     return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
 
-
 def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request: gr.Request):
     logger.info(f"http_bot. ip: {request.client.host}")
     start_tstamp = time.time()
@@ -204,18 +166,6 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
         new_state.append_message(new_state.roles[1], None)
         state = new_state
 
-    # Query worker address
-    controller_url = args.controller_url
-    ret = requests.post(controller_url + "/get_worker_address",
-            json={"model": model_name})
-    worker_addr = ret.json()["address"]
-    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
-
-    # No available worker
-    if worker_addr == "":
-        state.messages[-1][-1] = server_error_msg
-        yield (state, state.to_gradio_chatbot(), disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-        return
 
     # Construct prompt
     prompt = state.get_prompt()
@@ -248,9 +198,7 @@ def http_bot(state, model_selector, temperature, top_p, max_new_tokens, request:
 
     try:
         # Stream output
-        response = requests.post(worker_addr + "/worker_generate_stream",
-            headers=headers, json=pload, stream=True, timeout=10)
-        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+        for chunk in worker.generate_stream_gate(pload):
             if chunk:
                 data = json.loads(chunk.decode())
                 if data["error_code"] == 0:
@@ -312,14 +260,12 @@ block_css = """
 
 """
 
-def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
+def build_demo(cur_dir=None, concurrency_count=10):
     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_markdown)
-
+        gr.Markdown(title_markdown)
         with gr.Row():
             with gr.Column(scale=2):
                 # add a description
@@ -381,9 +327,8 @@ This is the demo for Dr-LLaVA. So far it could only be used for H&E stained Bone
                     regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
                     clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
 
-        if not embed_mode:
-            gr.Markdown(tos_markdown)
-            gr.Markdown(learn_more_markdown)
+        gr.Markdown(tos_markdown)
+        gr.Markdown(learn_more_markdown)
         url_params = gr.JSON(visible=False)
 
         # Register listeners
@@ -444,84 +389,29 @@ This is the demo for Dr-LLaVA. So far it could only be used for H&E stained Bone
             [state, chatbot] + btn_list,
             concurrency_limit=concurrency_count
         )
-
-        if args.model_list_mode == "once":
-            demo.load(
-                load_demo,
-                [url_params],
-                [state, model_selector],
-                js=get_window_url_params
-            )
-        elif args.model_list_mode == "reload":
-            demo.load(
-                load_demo_refresh_model_list,
-                None,
-                [state, model_selector],
-                queue=False
-            )
-        else:
-            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
-
+    
+        demo.load(
+            load_demo,
+            [url_params],
+            [state, model_selector],
+            js=get_window_url_params
+        )
     return demo
 
-def start_controller():
-    logger.info("Starting the controller")
-    controller_command = [
-        "python",
-        "-m",
-        "controller",
-        "--host",
-        "0.0.0.0",
-        "--port",
-        "10000",
-    ]
-    return subprocess.Popen(controller_command)
-
-def start_worker():
-    logger.info(f"Starting the model worker")
-    worker_command = [
-        "python",
-        "-m",
-        "model_worker",
-        "--host",
-        "0.0.0.0",
-        "--controller",
-        "http://localhost:10000",
-        "--load-bf16",
-        "--model-name",
-        "llava-rlhf-13b-v1.5-336",
-        "--model-path",
-        "daviddaytw/Dr-LLaVA-sft",
-        "--lora-path",
-        "daviddaytw/Dr-LLaVA-lora-adapter",
-    ]
-    return subprocess.Popen(worker_command)
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int)
-    parser.add_argument("--controller-url", type=str, default="http://localhost:10000")
     parser.add_argument("--concurrency-count", type=int, default=16)
-    parser.add_argument(
-        "--model-list-mode", type=str, default="reload", choices=["once", "reload"]
-    )
     parser.add_argument("--share", action="store_true")
-    parser.add_argument("--moderate", action="store_true")
-    parser.add_argument("--embed", action="store_true")
     args = parser.parse_args()
     logger.info(f"args: {args}")
 
-    controller_proc = start_controller()
-    worker_proc = start_worker()
-
-    # Wait for worker and controller to start
-    time.sleep(60)
-
-    models = get_model_list()
-
-    logger.info(args)
-    demo = build_demo(args.embed, concurrency_count=args.concurrency_count)
+    models = ['llava-rlhf-13b-v1.5-336']
+    model_path = 'daviddaytw/Dr-LLaVA-sft'
+    model_name = 'llava-rlhf-13b-v1.5-336'
+    lora_path = 'daviddaytw/Dr-LLaVA-lora-adapter'
+    demo = build_demo(concurrency_count=args.concurrency_count)
     demo.queue(
         api_open=False
     ).launch(

model_worker.py

@@ -1,26 +1,14 @@
 """
 A model worker executes the model.
 """
-import argparse
-import asyncio
 import json
-import time
-import threading
 import uuid
-
-from fastapi import FastAPI, Request, BackgroundTasks
-from fastapi.responses import StreamingResponse
-import requests
 import torch
-import uvicorn
-from functools import partial
 import spaces
 
 from peft import PeftModel
 
-from llava.constants import WORKER_HEART_BEAT_INTERVAL
-from llava.utils import (build_logger, server_error_msg,
-    pretty_print_semaphore)
+from llava.utils import (build_logger, server_error_msg)
 from model_builder import load_pretrained_model
 from llava.mm_utils import process_images, load_image_from_base64, tokenizer_image_token, KeywordsStoppingCriteria
 from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
@@ -37,20 +25,8 @@ global_counter = 0
 model_semaphore = None
 
 
-def heart_beat_worker(controller):
-
-    while True:
-        time.sleep(WORKER_HEART_BEAT_INTERVAL)
-        controller.send_heart_beat()
-
-
 class ModelWorker:
-    def __init__(self, controller_addr, worker_addr,
-                 worker_id, no_register,
-                 model_path, model_base, model_name,
-                 load_8bit, load_4bit, load_bf16, lora_path):
-        self.controller_addr = controller_addr
-        self.worker_addr = worker_addr
+    def __init__(self, model_path, model_base, model_name, load_bf16, lora_path):
         self.worker_id = worker_id
         if model_path.endswith("/"):
             model_path = model_path[:-1]
@@ -65,7 +41,7 @@ class ModelWorker:
 
         logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
         self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
-            model_path, model_base, self.model_name, load_8bit, load_4bit, load_bf16=load_bf16)
+            model_path, model_base, self.model_name, False, False, load_bf16=load_bf16)
         self.is_multimodal = 'llava' in self.model_name.lower()
         self.load_bf16 = load_bf16
 
@@ -76,59 +52,7 @@ class ModelWorker:
                 torch_device='cpu',
                 device_map="cpu",
             )
-
-        if not no_register:
-            self.register_to_controller()
-            self.heart_beat_thread = threading.Thread(
-                target=heart_beat_worker, args=(self,))
-            self.heart_beat_thread.start()
-
-    def register_to_controller(self):
-        logger.info("Register to controller")
-
-        url = self.controller_addr + "/register_worker"
-        data = {
-            "worker_name": self.worker_addr,
-            "check_heart_beat": True,
-            "worker_status": self.get_status()
-        }
-        r = requests.post(url, json=data)
-        assert r.status_code == 200
-
-    def send_heart_beat(self):
-        logger.info(f"Send heart beat. Models: {[self.model_name]}. "
-                    f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
-                    f"global_counter: {global_counter}")
-
-        url = self.controller_addr + "/receive_heart_beat"
-
-        while True:
-            try:
-                ret = requests.post(url, json={
-                    "worker_name": self.worker_addr,
-                    "queue_length": self.get_queue_length()}, timeout=5)
-                exist = ret.json()["exist"]
-                break
-            except requests.exceptions.RequestException as e:
-                logger.error(f"heart beat error: {e}")
-            time.sleep(5)
-
-        if not exist:
-            self.register_to_controller()
-
-    def get_queue_length(self):
-        if model_semaphore is None:
-            return 0
-        else:
-            return args.limit_model_concurrency - model_semaphore._value + (len(
-                model_semaphore._waiters) if model_semaphore._waiters is not None else 0)
-
-    def get_status(self):
-        return {
-            "model_names": [self.model_name],
-            "speed": 1,
-            "queue_length": self.get_queue_length(),
-        }
+        self.model.to("cuda:0")
 
     @spaces.GPU
     def generate_stream(self, params):
@@ -232,71 +156,7 @@ class ModelWorker:
             }
             yield json.dumps(ret).encode() + b"\0"
 
-
-app = FastAPI()
-
-
 def release_model_semaphore(fn=None):
     model_semaphore.release()
     if fn is not None:
         fn()
-
-
-@app.post("/worker_generate_stream")
-async def generate_stream(request: Request):
-    global model_semaphore, global_counter
-    global_counter += 1
-    params = await request.json()
-
-    if model_semaphore is None:
-        model_semaphore = asyncio.Semaphore(args.limit_model_concurrency)
-    await model_semaphore.acquire()
-    worker.send_heart_beat()
-    generator = worker.generate_stream_gate(params)
-    background_tasks = BackgroundTasks()
-    background_tasks.add_task(partial(release_model_semaphore, fn=worker.send_heart_beat))
-    return StreamingResponse(generator, background=background_tasks)
-
-
-@app.post("/worker_get_status")
-async def get_status(request: Request):
-    return worker.get_status()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    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("--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("--load-bf16", action="store_true")
-    parser.add_argument("--lora-path", type=str, default=None)
-    args = parser.parse_args()
-    logger.info(f"args: {args}")
-
-    if args.multi_modal:
-        logger.warning("Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.")
-
-    worker = ModelWorker(args.controller_address,
-                         args.worker_address,
-                         worker_id,
-                         args.no_register,
-                         args.model_path,
-                         args.model_base,
-                         args.model_name,
-                         args.load_8bit,
-                         args.load_4bit,
-                         args.load_bf16,
-                         args.lora_path)
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info")