Remove deprecated API functions and update model initialization in app.py

api.py

@@ -1,33 +0,0 @@
-# --------------------------------------------------------
-# InternVL
-# Copyright (c) 2024 OpenGVLab
-# Licensed under The MIT License [see LICENSE for details]
-# --------------------------------------------------------
-
-import base64
-import json
-from io import BytesIO
-
-import requests
-from PIL import Image
-
-
-def get_model_list(controller_url):
-    ret = requests.post(controller_url + '/refresh_all_workers')
-    assert ret.status_code == 200
-    ret = requests.post(controller_url + '/list_models')
-    models = ret.json()['models']
-    return models
-
-
-def get_selected_worker_ip(controller_url, selected_model):
-    ret = requests.post(controller_url + '/get_worker_address',
-            json={'model': selected_model})
-    worker_addr = ret.json()['address']
-    return worker_addr
-
-
-def pil_image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format='PNG')
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')

app.py

@@ -24,11 +24,14 @@ from utils import (
     load_image_from_base64,
     get_log_filename,
 )
+from threading import Thread
+import torch
 from conversation import Conversation
+from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
 
 logger = build_logger("gradio_web_server", "gradio_web_server.log")
 
-headers = {"User-Agent": "InternVL-Chat Client"}
+headers = {"User-Agent": "Vintern-Chat Client"}
 
 no_change_btn = gr.Button()
 enable_btn = gr.Button(interactive=True)
@@ -62,64 +65,6 @@ def init_state(state=None):
         del state
     return Conversation()
 
-
-def find_bounding_boxes(state, response):
-    pattern = re.compile(r"<ref>\s*(.*?)\s*</ref>\s*<box>\s*(\[\[.*?\]\])\s*</box>")
-    matches = pattern.findall(response)
-    results = []
-    for match in matches:
-        results.append((match[0], eval(match[1])))
-    returned_image = None
-    latest_image = state.get_images(source=state.USER)[-1]
-    returned_image = latest_image.copy()
-    width, height = returned_image.size
-    draw = ImageDraw.Draw(returned_image)
-    for result in results:
-        line_width = max(1, int(min(width, height) / 200))
-        random_color = (
-            random.randint(0, 128),
-            random.randint(0, 128),
-            random.randint(0, 128),
-        )
-        category_name, coordinates = result
-        coordinates = [
-            (
-                float(x[0]) / 1000,
-                float(x[1]) / 1000,
-                float(x[2]) / 1000,
-                float(x[3]) / 1000,
-            )
-            for x in coordinates
-        ]
-        coordinates = [
-            (
-                int(x[0] * width),
-                int(x[1] * height),
-                int(x[2] * width),
-                int(x[3] * height),
-            )
-            for x in coordinates
-        ]
-        for box in coordinates:
-            draw.rectangle(box, outline=random_color, width=line_width)
-            font = ImageFont.truetype("assets/SimHei.ttf", int(20 * line_width / 2))
-            text_size = font.getbbox(category_name)
-            text_width, text_height = (
-                text_size[2] - text_size[0],
-                text_size[3] - text_size[1],
-            )
-            text_position = (box[0], max(0, box[1] - text_height))
-            draw.rectangle(
-                [
-                    text_position,
-                    (text_position[0] + text_width, text_position[1] + text_height),
-                ],
-                fill=random_color,
-            )
-            draw.text(text_position, category_name, fill="white", font=font)
-    return returned_image if len(matches) > 0 else None
-
-
 def vote_last_response(state, liked, request: gr.Request):
     conv_data = {
         "tstamp": round(time.time(), 4),
@@ -220,6 +165,15 @@ def add_text(state, message, system_prompt, request: gr.Request):
         disable_btn,
     ) * 5
 
+model_name = "5CD-AI/Vintern-1B-v3_5"
+model = AutoModel.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True,
+).eval().cuda()
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)
+
 
 def http_bot(
     state,
@@ -230,7 +184,7 @@ def http_bot(
     max_input_tiles,
     request: gr.Request,
 ):
-    model_name = 'Vintern-1B-v3'
+    
     logger.info(f"http_bot. ip: {request.client.host}")
     start_tstamp = time.time()
     if hasattr(state, "skip_next") and state.skip_next:
@@ -242,12 +196,8 @@ def http_bot(
         ) + (no_change_btn,) * 5
         return
 
-    worker_addr = os.environ.get("WORKER_ADDR", "")
-    api_token = os.environ.get("API_TOKEN", "")
-    headers = {"Authorization": f"{api_token}", "Content-Type": "application/json"}
-
     # No available worker
-    if worker_addr == "":
+    if model is None:
         # state.messages[-1][-1] = server_error_msg
         state.update_message(Conversation.ASSISTANT, server_error_msg)
         yield (
@@ -265,17 +215,6 @@ def http_bot(
     all_images = state.get_images(source=state.USER)
     all_image_paths = [state.save_image(image) for image in all_images]
 
-    # Make requests
-    pload = {
-        "model": model_name,
-        "messages": state.get_prompt_v2(inlude_image=True, max_dynamic_patch=max_input_tiles),
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "max_tokens": max_new_tokens,
-        "repetition_penalty": repetition_penalty,
-        "stream": True
-    }
-    logger.info(f"==== request ====\n{pload}")
     state.append_message(Conversation.ASSISTANT, state.streaming_placeholder)
     yield (
         state,
@@ -285,26 +224,29 @@ def http_bot(
 
     try:
         # Stream output
-        response = requests.post(worker_addr, json=pload, headers=headers, stream=True, timeout=300)
-        finnal_output = ''
-        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\n"):
-            if chunk:
-                chunk = chunk.decode()
-                if chunk == 'data: [DONE]':
-                    break
-                if chunk.startswith("data:"):
-                    chunk = chunk[5:]
-                    chunk = json.loads(chunk)
-                    output = chunk['choices'][0]['delta']['content']
-                    finnal_output += output
+        # response = requests.post(worker_addr, json=pload, headers=headers, stream=True, timeout=300)
+        streamer = TextIteratorStreamer(
+            tokenizer, skip_prompt=True, skip_special_tokens=True
+        )
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            # Remove <|im_end|> or similar tokens from the output
+            buffer = buffer.replace("<|im_end|>", "")
                 
-                state.update_message(Conversation.ASSISTANT, finnal_output + state.streaming_placeholder, None)
-                yield (
-                    state,
-                    state.to_gradio_chatbot(),
-                    gr.MultimodalTextbox(interactive=False),
-                ) + (disable_btn,) * 5
-    except requests.exceptions.RequestException as e:
+            state.update_message(Conversation.ASSISTANT, buffer + state.streaming_placeholder, None)
+            yield (
+                state,
+                state.to_gradio_chatbot(),
+                gr.MultimodalTextbox(interactive=False),
+            ) + (disable_btn,) * 5
+
+    except Exception as e:
         state.update_message(Conversation.ASSISTANT, server_error_msg, None)
         yield (
             state,
@@ -332,7 +274,7 @@ def http_bot(
     ) + (enable_btn,) * 5
 
     finish_tstamp = time.time()
-    logger.info(f"{finnal_output}")
+    logger.info(f"{buffer}")
     data = {
         "tstamp": round(finish_tstamp, 4),
         "like": None,

controller.py

@@ -1,291 +0,0 @@
-"""
-A controller manages distributed workers.
-It sends worker addresses to clients.
-"""
-import argparse
-import dataclasses
-import json
-import re
-import threading
-import time
-from enum import Enum, auto
-from typing import List
-
-import numpy as np
-import requests
-import uvicorn
-from fastapi import FastAPI, Request
-from starlette.responses import StreamingResponse
-from utils import build_logger, server_error_msg
-
-CONTROLLER_HEART_BEAT_EXPIRATION = 30
-logger = build_logger('controller', 'controller.log')
-
-
-class DispatchMethod(Enum):
-    LOTTERY = auto()
-    SHORTEST_QUEUE = auto()
-
-    @classmethod
-    def from_str(cls, name):
-        if name == 'lottery':
-            return cls.LOTTERY
-        elif name == 'shortest_queue':
-            return cls.SHORTEST_QUEUE
-        else:
-            raise ValueError(f'Invalid dispatch method')
-
-
-@dataclasses.dataclass
-class WorkerInfo:
-    model_names: List[str]
-    speed: int
-    queue_length: int
-    check_heart_beat: bool
-    last_heart_beat: str
-
-
-def heart_beat_controller(controller):
-    while True:
-        time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION)
-        controller.remove_stable_workers_by_expiration()
-
-
-class Controller:
-    def __init__(self, dispatch_method: str):
-        # Dict[str -> WorkerInfo]
-        self.worker_info = {}
-        self.dispatch_method = DispatchMethod.from_str(dispatch_method)
-
-        self.heart_beat_thread = threading.Thread(
-            target=heart_beat_controller, args=(self,))
-        self.heart_beat_thread.start()
-
-        logger.info('Init controller')
-
-    def register_worker(self, worker_name: str, check_heart_beat: bool,
-                        worker_status: dict):
-        if worker_name not in self.worker_info:
-            logger.info(f'Register a new worker: {worker_name}')
-        else:
-            logger.info(f'Register an existing worker: {worker_name}')
-
-        if not worker_status:
-            worker_status = self.get_worker_status(worker_name)
-        if not worker_status:
-            return False
-
-        self.worker_info[worker_name] = WorkerInfo(
-            worker_status['model_names'], worker_status['speed'], worker_status['queue_length'],
-            check_heart_beat, time.time())
-
-        logger.info(f'Register done: {worker_name}, {worker_status}')
-        return True
-
-    def get_worker_status(self, worker_name: str):
-        try:
-            r = requests.post(worker_name + '/worker_get_status', timeout=5)
-        except requests.exceptions.RequestException as e:
-            logger.error(f'Get status fails: {worker_name}, {e}')
-            return None
-
-        if r.status_code != 200:
-            logger.error(f'Get status fails: {worker_name}, {r}')
-            return None
-
-        return r.json()
-
-    def remove_worker(self, worker_name: str):
-        del self.worker_info[worker_name]
-
-    def refresh_all_workers(self):
-        old_info = dict(self.worker_info)
-        self.worker_info = {}
-
-        for w_name, w_info in old_info.items():
-            if not self.register_worker(w_name, w_info.check_heart_beat, None):
-                logger.info(f'Remove stale worker: {w_name}')
-
-    def list_models(self):
-        model_names = set()
-
-        for w_name, w_info in self.worker_info.items():
-            model_names.update(w_info.model_names)
-
-        def extract_key(s):
-            if 'Pro' in s:
-                return 999
-            match = re.match(r'Vintern-(\d+)B', s)
-            if match:
-                return int(match.group(1))
-            return -1
-
-        def custom_sort_key(s):
-            key = extract_key(s)
-            # Return a tuple where -1 will ensure that non-matching items come last
-            return (0 if key != -1 else 1, -key if key != -1 else s)
-
-        sorted_list = sorted(list(model_names), key=custom_sort_key)
-        return sorted_list
-
-    def get_worker_address(self, model_name: str):
-        if self.dispatch_method == DispatchMethod.LOTTERY:
-            worker_names = []
-            worker_speeds = []
-            for w_name, w_info in self.worker_info.items():
-                if model_name in w_info.model_names:
-                    worker_names.append(w_name)
-                    worker_speeds.append(w_info.speed)
-            worker_speeds = np.array(worker_speeds, dtype=np.float32)
-            norm = np.sum(worker_speeds)
-            if norm < 1e-4:
-                return ''
-            worker_speeds = worker_speeds / norm
-            if True:  # Directly return address
-                pt = np.random.choice(np.arange(len(worker_names)),
-                    p=worker_speeds)
-                worker_name = worker_names[pt]
-                return worker_name
-
-        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
-            worker_names = []
-            worker_qlen = []
-            for w_name, w_info in self.worker_info.items():
-                if model_name in w_info.model_names:
-                    worker_names.append(w_name)
-                    worker_qlen.append(w_info.queue_length / w_info.speed)
-            if len(worker_names) == 0:
-                return ''
-            min_index = np.argmin(worker_qlen)
-            w_name = worker_names[min_index]
-            self.worker_info[w_name].queue_length += 1
-            logger.info(f'names: {worker_names}, queue_lens: {worker_qlen}, ret: {w_name}')
-            return w_name
-        else:
-            raise ValueError(f'Invalid dispatch method: {self.dispatch_method}')
-
-    def receive_heart_beat(self, worker_name: str, queue_length: int):
-        if worker_name not in self.worker_info:
-            logger.info(f'Receive unknown heart beat. {worker_name}')
-            return False
-
-        self.worker_info[worker_name].queue_length = queue_length
-        self.worker_info[worker_name].last_heart_beat = time.time()
-        logger.info(f'Receive heart beat. {worker_name}')
-        return True
-
-    def remove_stable_workers_by_expiration(self):
-        expire = time.time() - CONTROLLER_HEART_BEAT_EXPIRATION
-        to_delete = []
-        for worker_name, w_info in self.worker_info.items():
-            if w_info.check_heart_beat and w_info.last_heart_beat < expire:
-                to_delete.append(worker_name)
-
-        for worker_name in to_delete:
-            self.remove_worker(worker_name)
-
-    def worker_api_generate_stream(self, params):
-        worker_addr = self.get_worker_address(params['model'])
-        if not worker_addr:
-            logger.info(f"no worker: {params['model']}")
-            ret = {
-                'text': server_error_msg,
-                'error_code': 2,
-            }
-            yield json.dumps(ret).encode() + b'\0'
-
-        try:
-            response = requests.post(worker_addr + '/worker_generate_stream',
-                json=params, stream=True, timeout=5)
-            for chunk in response.iter_lines(decode_unicode=False, delimiter=b'\0'):
-                if chunk:
-                    yield chunk + b'\0'
-        except requests.exceptions.RequestException as e:
-            logger.info(f'worker timeout: {worker_addr}')
-            ret = {
-                'text': server_error_msg,
-                'error_code': 3,
-            }
-            yield json.dumps(ret).encode() + b'\0'
-
-    # Let the controller act as a worker to achieve hierarchical
-    # management. This can be used to connect isolated sub networks.
-    def worker_api_get_status(self):
-        model_names = set()
-        speed = 0
-        queue_length = 0
-
-        for w_name in self.worker_info:
-            worker_status = self.get_worker_status(w_name)
-            if worker_status is not None:
-                model_names.update(worker_status['model_names'])
-                speed += worker_status['speed']
-                queue_length += worker_status['queue_length']
-
-        return {
-            'model_names': list(model_names),
-            'speed': speed,
-            'queue_length': queue_length,
-        }
-
-
-app = FastAPI()
-
-
-@app.post('/register_worker')
-async def register_worker(request: Request):
-    data = await request.json()
-    controller.register_worker(
-        data['worker_name'], data['check_heart_beat'],
-        data.get('worker_status', None))
-
-
-@app.post('/refresh_all_workers')
-async def refresh_all_workers():
-    models = controller.refresh_all_workers()
-
-
-@app.post('/list_models')
-async def list_models():
-    models = controller.list_models()
-    return {'models': models}
-
-
-@app.post('/get_worker_address')
-async def get_worker_address(request: Request):
-    data = await request.json()
-    addr = controller.get_worker_address(data['model'])
-    return {'address': addr}
-
-
-@app.post('/receive_heart_beat')
-async def receive_heart_beat(request: Request):
-    data = await request.json()
-    exist = controller.receive_heart_beat(
-        data['worker_name'], data['queue_length'])
-    return {'exist': exist}
-
-
-@app.post('/worker_generate_stream')
-async def worker_api_generate_stream(request: Request):
-    params = await request.json()
-    generator = controller.worker_api_generate_stream(params)
-    return StreamingResponse(generator)
-
-
-@app.post('/worker_get_status')
-async def worker_api_get_status(request: Request):
-    return controller.worker_api_get_status()
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--host', type=str, default='0.0.0.0')
-    parser.add_argument('--port', type=int, default=10075)
-    parser.add_argument('--dispatch-method', type=str, choices=[
-        'lottery', 'shortest_queue'], default='shortest_queue')
-    args = parser.parse_args()
-    logger.info(f'args: {args}')
-
-    controller = Controller(args.dispatch_method)
-    uvicorn.run(app, host=args.host, port=args.port, log_level='info')

gradio_web_server.py

@@ -1,761 +0,0 @@
-import argparse
-from ast import parse
-import datetime
-import json
-import os
-import time
-import hashlib
-import re
-
-import gradio as gr
-import requests
-import random
-from filelock import FileLock
-from io import BytesIO
-from PIL import Image, ImageDraw, ImageFont
-
-from constants import LOGDIR
-from utils import (
-    build_logger,
-    server_error_msg,
-    violates_moderation,
-    moderation_msg,
-    load_image_from_base64,
-    get_log_filename,
-)
-from conversation import Conversation
-
-logger = build_logger("gradio_web_server", "gradio_web_server.log")
-
-headers = {"User-Agent": "InternVL-Chat Client"}
-
-no_change_btn = gr.Button()
-enable_btn = gr.Button(interactive=True)
-disable_btn = gr.Button(interactive=False)
-
-
-def write2file(path, content):
-    lock = FileLock(f"{path}.lock")
-    with lock:
-        with open(path, "a") as fout:
-            fout.write(content)
-
-
-def sort_models(models):
-    def custom_sort_key(model_name):
-        # InternVL-Chat-V1-5 should be the first item
-        if model_name == "Vintern-1B-v3":
-            return (1, model_name)  # 1 indicates highest precedence
-        elif model_name.startswith("Vintern-1B-v3"):
-            return (1, model_name)  # 1 indicates highest precedence
-        else:
-            return (0, model_name)  # 0 indicates normal order
-
-    models.sort(key=custom_sort_key, reverse=True)
-    try:  # We have five InternVL-Chat-V1-5 models, randomly choose one to be the first
-        first_three = models[:4]
-        random.shuffle(first_three)
-        models[:4] = first_three
-    except:
-        pass
-    return models
-
-
-def get_model_list():
-    logger.info(f"Call `get_model_list`")
-    ret = requests.post(args.controller_url + "/refresh_all_workers")
-    logger.info(f"status_code from `get_model_list`: {ret.status_code}")
-    assert ret.status_code == 200
-    ret = requests.post(args.controller_url + "/list_models")
-    logger.info(f"status_code from `list_models`: {ret.status_code}")
-    models = ret.json()["models"]
-    models = sort_models(models)
-
-    logger.info(f"Models (from {args.controller_url}): {models}")
-    return models
-
-
-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 init_state(state=None):
-    if state is not None:
-        del state
-    return Conversation()
-
-def load_demo(url_params, request: gr.Request = None):
-    if not request:
-        logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
-
-    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 = init_state()
-    return state, dropdown_update
-
-
-def load_demo_refresh_model_list(request: gr.Request = None):
-    if not request:
-        logger.info(f"load_demo. ip: {request.client.host}")
-    models = get_model_list()
-    state = init_state()
-    dropdown_update = gr.Dropdown(
-        choices=models, value=models[0] if len(models) > 0 else ""
-    )
-    return state, dropdown_update
-
-
-def vote_last_response(state, liked, model_selector, request: gr.Request):
-    conv_data = {
-        "tstamp": round(time.time(), 4),
-        "like": liked,
-        "model": model_selector,
-        "state": state.dict(),
-        "ip": request.client.host,
-    }
-    write2file(get_log_filename(), json.dumps(conv_data) + "\n")
-
-
-def upvote_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"upvote. ip: {request.client.host}")
-    vote_last_response(state, True, model_selector, request)
-    textbox = gr.MultimodalTextbox(value=None, interactive=True)
-    return (textbox,) + (disable_btn,) * 3
-
-
-def downvote_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"downvote. ip: {request.client.host}")
-    vote_last_response(state, False, model_selector, request)
-    textbox = gr.MultimodalTextbox(value=None, interactive=True)
-    return (textbox,) + (disable_btn,) * 3
-
-
-def vote_selected_response(
-    state, model_selector, request: gr.Request, data: gr.LikeData
-):
-    logger.info(
-        f"Vote: {data.liked}, index: {data.index}, value: {data.value} , ip: {request.client.host}"
-    )
-    conv_data = {
-        "tstamp": round(time.time(), 4),
-        "like": data.liked,
-        "index": data.index,
-        "model": model_selector,
-        "state": state.dict(),
-        "ip": request.client.host,
-    }
-    write2file(get_log_filename(), json.dumps(conv_data) + "\n")
-    return
-
-
-def flag_last_response(state, model_selector, request: gr.Request):
-    logger.info(f"flag. ip: {request.client.host}")
-    vote_last_response(state, "flag", model_selector, request)
-    textbox = gr.MultimodalTextbox(value=None, interactive=True)
-    return (textbox,) + (disable_btn,) * 3
-
-
-def regenerate(state, image_process_mode, request: gr.Request):
-    logger.info(f"regenerate. ip: {request.client.host}")
-    # state.messages[-1][-1] = None
-    state.update_message(Conversation.ASSISTANT, None, -1)
-    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
-    textbox = gr.MultimodalTextbox(value=None, interactive=True)
-    return (state, state.to_gradio_chatbot(), textbox) + (disable_btn,) * 5
-
-
-def clear_history(request: gr.Request):
-    logger.info(f"clear_history. ip: {request.client.host}")
-    state = init_state()
-    textbox = gr.MultimodalTextbox(value=None, interactive=True)
-    return (state, state.to_gradio_chatbot(), textbox) + (disable_btn,) * 5
-
-
-def change_system_prompt(state, system_prompt, request: gr.Request):
-    logger.info(f"Change system prompt. ip: {request.client.host}")
-    state.set_system_message(system_prompt)
-    return state
-
-
-def add_text(state, message, system_prompt, model_selector, request: gr.Request):
-    print(f"state: {state}")
-    if not state:
-        state, model_selector = load_demo_refresh_model_list(request)
-    images = message.get("files", [])
-    text = message.get("text", "").strip()
-    logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
-    # import pdb; pdb.set_trace()
-    textbox = gr.MultimodalTextbox(value=None, interactive=False)
-    if len(text) <= 0 and len(images) == 0:
-        state.skip_next = True
-        return (state, state.to_gradio_chatbot(), textbox) + (no_change_btn,) * 5
-    if args.moderate:
-        flagged = violates_moderation(text)
-        if flagged:
-            state.skip_next = True
-            textbox = gr.MultimodalTextbox(
-                value={"text": moderation_msg}, interactive=True
-            )
-            return (state, state.to_gradio_chatbot(), textbox) + (no_change_btn,) * 5
-    images = [Image.open(path).convert("RGB") for path in images]
-
-    if len(images) > 0 and len(state.get_images(source=state.USER)) > 0:
-        state = init_state(state)
-    state.set_system_message(system_prompt)
-    state.append_message(Conversation.USER, text, images)
-    state.skip_next = False
-    return (state, state.to_gradio_chatbot(), textbox, model_selector) + (
-        disable_btn,
-    ) * 5
-
-
-def http_bot(
-    state,
-    model_selector,
-    temperature,
-    top_p,
-    repetition_penalty,
-    max_new_tokens,
-    max_input_tiles,
-    # bbox_threshold,
-    # mask_threshold,
-    request: gr.Request,
-):
-    logger.info(f"http_bot. ip: {request.client.host}")
-    start_tstamp = time.time()
-    model_name = model_selector
-    if hasattr(state, "skip_next") and state.skip_next:
-        # This generate call is skipped due to invalid inputs
-        yield (
-            state,
-            state.to_gradio_chatbot(),
-            gr.MultimodalTextbox(interactive=False),
-        ) + (no_change_btn,) * 5
-        return
-
-    # 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
-        state.update_message(Conversation.ASSISTANT, server_error_msg)
-        yield (
-            state,
-            state.to_gradio_chatbot(),
-            gr.MultimodalTextbox(interactive=False),
-            disable_btn,
-            disable_btn,
-            disable_btn,
-            enable_btn,
-            enable_btn,
-        )
-        return
-
-    all_images = state.get_images(source=state.USER)
-    all_image_paths = [state.save_image(image) for image in all_images]
-
-    # Make requests
-    pload = {
-        "model": model_name,
-        "prompt": state.get_prompt(),
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "max_new_tokens": max_new_tokens,
-        "max_input_tiles": max_input_tiles,
-        # "bbox_threshold": bbox_threshold,
-        # "mask_threshold": mask_threshold,
-        "repetition_penalty": repetition_penalty,
-        "images": f"List of {len(all_images)} images: {all_image_paths}",
-    }
-    logger.info(f"==== request ====\n{pload}")
-    pload.pop("images")
-    pload["prompt"] = state.get_prompt(inlude_image=True)
-    state.append_message(Conversation.ASSISTANT, state.streaming_placeholder)
-    yield (
-        state,
-        state.to_gradio_chatbot(),
-        gr.MultimodalTextbox(interactive=False),
-    ) + (disable_btn,) * 5
-
-    try:
-        # Stream output
-        response = requests.post(
-            worker_addr + "/worker_generate_stream",
-            headers=headers,
-            json=pload,
-            stream=True,
-            timeout=20,
-        )
-        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
-            if chunk:
-                data = json.loads(chunk.decode())
-                if data["error_code"] == 0:
-                    if "text" in data:
-                        output = data["text"].strip()
-                        output += state.streaming_placeholder
-
-                    image = None
-                    if "image" in data:
-                        image = load_image_from_base64(data["image"])
-                        _ = state.save_image(image)
-
-                    state.update_message(Conversation.ASSISTANT, output, image)
-                    yield (
-                        state,
-                        state.to_gradio_chatbot(),
-                        gr.MultimodalTextbox(interactive=False),
-                    ) + (disable_btn,) * 5
-                else:
-                    output = (
-                        f"**{data['text']}**" + f" (error_code: {data['error_code']})"
-                    )
-
-                    state.update_message(Conversation.ASSISTANT, output, None)
-                    yield (
-                        state,
-                        state.to_gradio_chatbot(),
-                        gr.MultimodalTextbox(interactive=True),
-                    ) + (
-                        disable_btn,
-                        disable_btn,
-                        disable_btn,
-                        enable_btn,
-                        enable_btn,
-                    )
-                    return
-    except requests.exceptions.RequestException as e:
-        state.update_message(Conversation.ASSISTANT, server_error_msg, None)
-        yield (
-            state,
-            state.to_gradio_chatbot(),
-            gr.MultimodalTextbox(interactive=True),
-        ) + (
-            disable_btn,
-            disable_btn,
-            disable_btn,
-            enable_btn,
-            enable_btn,
-        )
-        return
-
-    ai_response = state.return_last_message()
-
-    state.end_of_current_turn()
-
-    yield (
-        state,
-        state.to_gradio_chatbot(),
-        gr.MultimodalTextbox(interactive=True),
-    ) + (enable_btn,) * 5
-
-    finish_tstamp = time.time()
-    logger.info(f"{output}")
-    data = {
-        "tstamp": round(finish_tstamp, 4),
-        "like": None,
-        "model": model_name,
-        "start": round(start_tstamp, 4),
-        "finish": round(start_tstamp, 4),
-        "state": state.dict(),
-        "images": all_image_paths,
-        "ip": request.client.host,
-    }
-    write2file(get_log_filename(), json.dumps(data) + "\n")
-
-
-title_html = """
-<h2> <span class="gradient-text" id="text">InternVL2</span><span class="plain-text">: Better than the Best—Expanding Performance Boundaries of Open-Source Multimodal Models with the Progressive Scaling Strategy</span></h2>
-<a href="https://internvl.github.io/blog/2024-07-02-InternVL-2.0/">[📜 InternVL2 Blog]</a> 
-<a href="https://huggingface.co/spaces/OpenGVLab/InternVL">[🤗 HF Demo]</a> 
-<a href="https://github.com/OpenGVLab/InternVL?tab=readme-ov-file#quick-start-with-huggingface">[🚀 Quick Start]</a> 
-<a href="https://github.com/OpenGVLab/InternVL/blob/main/document/How_to_use_InternVL_API.md">[🌐 API]</a> 
-"""
-
-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 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.
-
-### Acknowledgement
-This demo is modified from LLaVA's demo. Thanks for their awesome work!
-"""
-# .gradio-container {margin: 5px 10px 0 10px !important};
-block_css = """
-.gradio-container {margin: 0.1% 1% 0 1% !important; max-width: 98% !important;};
-#buttons button {
-    min-width: min(120px,100%);
-}
-
-.gradient-text {
-    font-size: 28px;
-    width: auto;
-    font-weight: bold;
-    background: linear-gradient(45deg, red, orange, yellow, green, blue, indigo, violet);
-    background-clip: text;
-    -webkit-background-clip: text;
-    color: transparent;
-}
-
-.plain-text {
-    font-size: 22px;
-    width: auto;
-    font-weight: bold;
-}
-"""
-
-js = """
-function createWaveAnimation() {
-    const text = document.getElementById('text');
-    var i = 0;
-    setInterval(function() {
-        const colors = [
-            'red, orange, yellow, green, blue, indigo, violet, purple',
-            'orange, yellow, green, blue, indigo, violet, purple, red',
-            'yellow, green, blue, indigo, violet, purple, red, orange',
-            'green, blue, indigo, violet, purple, red, orange, yellow',
-            'blue, indigo, violet, purple, red, orange, yellow, green',
-            'indigo, violet, purple, red, orange, yellow, green, blue',
-            'violet, purple, red, orange, yellow, green, blue, indigo',
-            'purple, red, orange, yellow, green, blue, indigo, violet',
-        ];
-        const angle = 45;
-        const colorIndex = i % colors.length;
-        text.style.background = `linear-gradient(${angle}deg, ${colors[colorIndex]})`;
-        text.style.webkitBackgroundClip = 'text';
-        text.style.backgroundClip = 'text';
-        text.style.color = 'transparent';
-        text.style.fontSize = '28px';
-        text.style.width = 'auto';
-        text.textContent = 'InternVL2';
-        text.style.fontWeight = 'bold';
-        i += 1;
-    }, 200);
-    const params = new URLSearchParams(window.location.search);
-    url_params = Object.fromEntries(params);
-    // console.log(url_params);
-    // console.log('hello world...');
-    // console.log(window.location.search);
-    // console.log('hello world...');
-    // alert(window.location.search)
-    // alert(url_params);
-    return url_params;
-}
-
-"""
-
-
-def build_demo(embed_mode):
-    textbox = gr.MultimodalTextbox(
-        interactive=True,
-        file_types=["image", "video"],
-        placeholder="Enter message or upload file...",
-        show_label=False,
-    )
-
-    with gr.Blocks(
-        title="InternVL-Chat",
-        theme=gr.themes.Default(),
-        css=block_css,
-    ) as demo:
-        state = gr.State()
-
-        if not embed_mode:
-            # gr.Markdown(title_markdown)
-            gr.HTML(title_html)
-
-        with gr.Row():
-            with gr.Column(scale=2):
-
-                with gr.Row(elem_id="model_selector_row"):
-                    model_selector = gr.Dropdown(
-                        choices=models,
-                        value=models[0] if len(models) > 0 else "",
-                        # value="InternVL-Chat-V1-5",
-                        interactive=True,
-                        show_label=False,
-                        container=False,
-                    )
-
-                with gr.Accordion("System Prompt", open=False) as system_prompt_row:
-                    system_prompt = gr.Textbox(
-                        value="请尽可能详细地回答用户的问题。",
-                        label="System Prompt",
-                        interactive=True,
-                    )
-                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",
-                    )
-                    repetition_penalty = gr.Slider(
-                        minimum=1.0,
-                        maximum=1.5,
-                        value=1.1,
-                        step=0.02,
-                        interactive=True,
-                        label="Repetition penalty",
-                    )
-                    max_output_tokens = gr.Slider(
-                        minimum=0,
-                        maximum=4096,
-                        value=1024,
-                        step=64,
-                        interactive=True,
-                        label="Max output tokens",
-                    )
-                    max_input_tiles = gr.Slider(
-                        minimum=1,
-                        maximum=32,
-                        value=12,
-                        step=1,
-                        interactive=True,
-                        label="Max input tiles (control the image size)",
-                    )
-                examples = gr.Examples(
-                    examples=[
-                        [
-                            {
-                                "files": [
-                                    "gallery/prod_9.jpg",
-                                ],
-                                "text": "What's at the far end of the image?",
-                            }
-                        ],
-                        [
-                            {
-                                "files": [
-                                    "gallery/astro_on_unicorn.png",
-                                ],
-                                "text": "What does this image mean?",
-                            }
-                        ],
-                        [
-                            {
-                                "files": [
-                                    "gallery/prod_12.png",
-                                ],
-                                "text": "What are the consequences of the easy decisions shown in this image?",
-                            }
-                        ],
-                        [
-                            {
-                                "files": [
-                                    "gallery/child_1.jpg",
-                                    "gallery/child_2.jpg",
-                                    f"gallery/child_3.jpg",
-                                ],
-                                "text": "这三帧图片讲述了一件什么事情？",
-                            }
-                        ],
-                    ],
-                    inputs=[textbox],
-                )
-
-            with gr.Column(scale=8):
-                chatbot = gr.Chatbot(
-                    elem_id="chatbot",
-                    label="InternVL2",
-                    height=580,
-                    show_copy_button=True,
-                    show_share_button=True,
-                    avatar_images=[
-                        "assets/human.png",
-                        "assets/assistant.png",
-                    ],
-                    bubble_full_width=False,
-                )
-                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)
-
-        if not embed_mode:
-            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, 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],
-        )
-        chatbot.like(
-            vote_selected_response,
-            [state, model_selector],
-            [],
-        )
-        flag_btn.click(
-            flag_last_response,
-            [state, model_selector],
-            [textbox, upvote_btn, downvote_btn, flag_btn],
-        )
-        regenerate_btn.click(
-            regenerate,
-            [state, system_prompt],
-            [state, chatbot, textbox] + btn_list,
-        ).then(
-            http_bot,
-            [
-                state,
-                model_selector,
-                temperature,
-                top_p,
-                repetition_penalty,
-                max_output_tokens,
-                max_input_tiles,
-                # bbox_threshold,
-                # mask_threshold,
-            ],
-            [state, chatbot, textbox] + btn_list,
-        )
-        clear_btn.click(clear_history, None, [state, chatbot, textbox] + btn_list)
-
-        textbox.submit(
-            add_text,
-            [state, textbox, system_prompt, model_selector],
-            [state, chatbot, textbox, model_selector] + btn_list,
-        ).then(
-            http_bot,
-            [
-                state,
-                model_selector,
-                temperature,
-                top_p,
-                repetition_penalty,
-                max_output_tokens,
-                max_input_tiles,
-                # bbox_threshold,
-                # mask_threshold,
-            ],
-            [state, chatbot, textbox] + btn_list,
-        )
-        submit_btn.click(
-            add_text,
-            [state, textbox, system_prompt, model_selector],
-            [state, chatbot, textbox, model_selector] + btn_list,
-        ).then(
-            http_bot,
-            [
-                state,
-                model_selector,
-                temperature,
-                top_p,
-                repetition_penalty,
-                max_output_tokens,
-                max_input_tiles,
-                # bbox_threshold,
-                # mask_threshold,
-            ],
-            [state, chatbot, textbox] + btn_list,
-        )
-
-        # NOTE: The following code will be not triggered when deployed on HF space.
-        # It's very strange. I don't know why.
-        """
-        if args.model_list_mode == "once":
-            demo.load(
-                load_demo,
-                [url_params],
-                [state, model_selector],
-                js=js,
-            )
-        elif args.model_list_mode == "reload":
-            demo.load(
-                load_demo_refresh_model_list,
-                None,
-                [state, model_selector],
-                js=js,
-            )
-        else:
-            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
-        """
-
-    return demo
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--host", type=str, default="0.0.0.0")
-    parser.add_argument("--port", type=int, default=7860)
-    parser.add_argument("--controller-url", type=str, default=None)
-    parser.add_argument("--concurrency-count", type=int, default=10)
-    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}")
-    if not args.controller_url:
-        args.controller_url = os.environ.get("CONTROLLER_URL", None)
-
-    if not args.controller_url:
-        raise ValueError("controller-url is required.")
-
-    models = get_model_list()
-
-    logger.info(args)
-    demo = build_demo(args.embed)
-    demo.queue(api_open=False).launch(
-        server_name=args.host,
-        server_port=args.port,
-        share=args.share,
-        max_threads=args.concurrency_count,
-    )

model_worker.py

@@ -1,541 +0,0 @@
-# --------------------------------------------------------
-# InternVL
-# Copyright (c) 2024 OpenGVLab
-# Licensed under The MIT License [see LICENSE for details]
-# --------------------------------------------------------
-
-"""
-A model worker executes the model.
-"""
-import spaces
-import os
-import argparse
-import asyncio
-
-import json
-import math
-import threading
-import time
-import uuid
-import traceback
-from functools import partial
-from threading import Thread
-
-import requests
-import torch
-import torchvision.transforms as T
-import uvicorn
-from constants import IMAGENET_MEAN, IMAGENET_STD, WORKER_HEART_BEAT_INTERVAL
-from fastapi import BackgroundTasks, FastAPI, Request
-from fastapi.responses import StreamingResponse
-from PIL import Image
-from torchvision.transforms.functional import InterpolationMode
-from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
-from utils import (
-    build_logger,
-    pretty_print_semaphore,
-    server_error_msg,
-    load_image_from_base64,
-)
-
-
-worker_id = str(uuid.uuid4())[:6]
-logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
-global_counter = 0
-model_semaphore = None
-
-
-def build_transform(input_size):
-    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
-    transform = T.Compose(
-        [
-            T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
-            T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
-            T.ToTensor(),
-            T.Normalize(mean=MEAN, std=STD),
-        ]
-    )
-    return transform
-
-
-def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
-    best_ratio_diff = float("inf")
-    best_ratio = (1, 1)
-    area = width * height
-    for ratio in target_ratios:
-        target_aspect_ratio = ratio[0] / ratio[1]
-        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
-        if ratio_diff < best_ratio_diff:
-            best_ratio_diff = ratio_diff
-            best_ratio = ratio
-        elif ratio_diff == best_ratio_diff:
-            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
-                best_ratio = ratio
-    return best_ratio
-
-
-def dynamic_preprocess(
-    image, min_num=1, max_num=6, image_size=448, use_thumbnail=False
-):
-    orig_width, orig_height = image.size
-    aspect_ratio = orig_width / orig_height
-
-    # calculate the existing image aspect ratio
-    target_ratios = set(
-        (i, j)
-        for n in range(min_num, max_num + 1)
-        for i in range(1, n + 1)
-        for j in range(1, n + 1)
-        if i * j <= max_num and i * j >= min_num
-    )
-    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
-
-    # find the closest aspect ratio to the target
-    target_aspect_ratio = find_closest_aspect_ratio(
-        aspect_ratio, target_ratios, orig_width, orig_height, image_size
-    )
-
-    # calculate the target width and height
-    target_width = image_size * target_aspect_ratio[0]
-    target_height = image_size * target_aspect_ratio[1]
-    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
-
-    # resize the image
-    resized_img = image.resize((target_width, target_height))
-    processed_images = []
-    for i in range(blocks):
-        box = (
-            (i % (target_width // image_size)) * image_size,
-            (i // (target_width // image_size)) * image_size,
-            ((i % (target_width // image_size)) + 1) * image_size,
-            ((i // (target_width // image_size)) + 1) * image_size,
-        )
-        # split the image
-        split_img = resized_img.crop(box)
-        processed_images.append(split_img)
-    assert len(processed_images) == blocks
-    if use_thumbnail and len(processed_images) != 1:
-        thumbnail_img = image.resize((image_size, image_size))
-        processed_images.append(thumbnail_img)
-    return processed_images
-
-
-def heart_beat_worker(controller):
-    while True:
-        time.sleep(WORKER_HEART_BEAT_INTERVAL)
-        controller.send_heart_beat()
-
-
-def split_model(model_name):
-    device_map = {}
-    world_size = torch.cuda.device_count()
-    num_layers = {
-        "Vintern-1B-v3": 24,
-    }[model_name]
-    # Since the first GPU will be used for ViT, treat it as half a GPU.
-    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
-    num_layers_per_gpu = [num_layers_per_gpu] * world_size
-    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
-    layer_cnt = 0
-    for i, num_layer in enumerate(num_layers_per_gpu):
-        for j in range(num_layer):
-            device_map[f"language_model.model.layers.{layer_cnt}"] = i
-            layer_cnt += 1
-    device_map["vision_model"] = 0
-    device_map["mlp1"] = 0
-    device_map["language_model.model.tok_embeddings"] = 0
-    device_map["language_model.model.embed_tokens"] = 0
-    device_map["language_model.output"] = 0
-    device_map["language_model.model.norm"] = 0
-    device_map["language_model.lm_head"] = 0
-    device_map[f"language_model.model.layers.{num_layers - 1}"] = 0
-
-    return device_map
-
-
-def multi_thread_infer(
-    model, tokenizer, pixel_values, question, history, generation_config
-):
-    with torch.no_grad():
-        thread = Thread(
-            target=model.chat,
-            kwargs=dict(
-                tokenizer=tokenizer,
-                pixel_values=pixel_values,
-                question=question,
-                history=history,
-                return_history=False,
-                generation_config=generation_config,
-            ),
-        )
-        thread.start()
-
-
-class ModelWorker:
-    def __init__(
-        self,
-        controller_addr,
-        worker_addr,
-        worker_id,
-        model_path,
-        model_name,
-        load_8bit,
-        device,
-        context_len=8192,
-    ):
-        self.controller_addr = controller_addr
-        self.worker_addr = worker_addr
-        self.worker_id = worker_id
-        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.import_flash_attn()
-        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_path, trust_remote_code=True, use_fast=False
-        )
-        self.tokenizer = tokenizer
-
-        if device == "auto":
-            device_map = split_model(self.model_name)
-            self.model = AutoModel.from_pretrained(
-                model_path,
-                load_in_8bit=load_8bit,
-                torch_dtype=torch.bfloat16,
-                device_map=device_map,
-                trust_remote_code=True,
-            ).eval()
-        else:
-            self.model = AutoModel.from_pretrained(
-                model_path,
-                load_in_8bit=load_8bit,
-                torch_dtype=torch.bfloat16,
-                trust_remote_code=True,
-            ).eval()
-        if not load_8bit and not device == "auto":
-            self.model = self.model.cuda()
-        self.load_8bit = load_8bit
-        self.device = device
-        self.model_path = model_path
-        self.image_size = self.model.config.force_image_size
-        self.context_len = context_len
-        self.register_to_controller()
-        self.heart_beat_thread = threading.Thread(
-            target=heart_beat_worker, args=(self,)
-        )
-        self.heart_beat_thread.start()
-
-    @spaces.GPU(duration=120)
-    def import_flash_attn(self):
-        try:
-            import flash_attn
-        except ImportError:
-
-            def install_flash_attn():
-                os.system(
-                    "FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn==2.5.9.post1 --no-build-isolation"
-                )
-
-            install_flash_attn()
-            # import flash_attn
-
-    def reload_model(self):
-        del self.model
-        torch.cuda.empty_cache()
-        if self.device == "auto":
-            device_map = split_model(self.model_name)
-            self.model = AutoModel.from_pretrained(
-                self.model_path,
-                load_in_8bit=self.load_8bit,
-                torch_dtype=torch.bfloat16,
-                device_map=device_map,
-                trust_remote_code=True,
-            ).eval()
-        else:
-            self.model = AutoModel.from_pretrained(
-                self.model_path,
-                load_in_8bit=self.load_8bit,
-                torch_dtype=torch.bfloat16,
-                trust_remote_code=True,
-            ).eval()
-        if not self.load_8bit and not self.device == "auto":
-            self.model = self.model.cuda()
-
-    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(),
-        }
-
-    def generate_stream(self, params):
-        system_message = params["prompt"][0]["content"]
-        send_messages = params["prompt"][1:]
-        max_input_tiles = params["max_input_tiles"]
-        temperature = params["temperature"]
-        top_p = params["top_p"]
-        max_new_tokens = params["max_new_tokens"]
-        repetition_penalty = params["repetition_penalty"]
-        do_sample = True if temperature > 0.0 else False
-
-        global_image_cnt = 0
-        history, pil_images, max_input_tile_list = [], [], []
-        for message in send_messages:
-            if message["role"] == "user":
-                prefix = ""
-                if "image" in message:
-                    max_input_tile_temp = []
-                    for image_str in message["image"]:
-                        pil_images.append(load_image_from_base64(image_str))
-                        prefix += f"Image-{global_image_cnt + 1}: <image>\n\n"
-                        global_image_cnt += 1
-                        max_input_tile_temp.append(
-                            max(1, max_input_tiles // len(message["image"]))
-                        )
-                    if len(max_input_tile_temp) > 0:
-                        max_input_tile_list.append(max_input_tile_temp)
-                content = prefix + message["content"]
-                history.append(
-                    [
-                        content,
-                    ]
-                )
-            else:
-                history[-1].append(message["content"])
-        question, history = history[-1][0], history[:-1]
-
-        if global_image_cnt == 1:
-            question = question.replace("Image-1: <image>\n\n", "<image>\n")
-            history = [
-                [item[0].replace("Image-1: <image>\n\n", "<image>\n"), item[1]]
-                for item in history
-            ]
-
-        # Create a new list to store processed sublists
-        flattened_list = []
-        # Iterate through all but the last sublist in max_input_tile_list and process them
-        for sublist in max_input_tile_list[:-1]:
-            processed_sublist = [1] * len(
-                sublist
-            )  # Change each element in the sublist to 1
-            flattened_list.extend(
-                processed_sublist
-            )  # Flatten the processed sublist and add to the new list
-        # If max_input_tile_list is not empty, add the last sublist to the new list
-        if max_input_tile_list:
-            flattened_list.extend(max_input_tile_list[-1])
-        max_input_tile_list = flattened_list
-        assert len(max_input_tile_list) == len(
-            pil_images
-        ), "The number of max_input_tile_list and pil_images should be the same."
-
-        old_system_message = self.model.system_message
-        self.model.system_message = system_message
-        image_tiles = []
-        transform = build_transform(input_size=self.image_size)
-        if len(pil_images) > 0:
-            for current_max_input_tiles, pil_image in zip(
-                max_input_tile_list, pil_images
-            ):
-                if self.model.config.dynamic_image_size:
-                    tiles = dynamic_preprocess(
-                        pil_image,
-                        image_size=self.image_size,
-                        max_num=current_max_input_tiles,
-                        use_thumbnail=self.model.config.use_thumbnail,
-                    )
-                else:
-                    tiles = [pil_image]
-                image_tiles += tiles
-            pixel_values = [transform(item) for item in image_tiles]
-            pixel_values = torch.stack(pixel_values).to(
-                self.model.device, dtype=torch.bfloat16
-            )
-            logger.info(f"Split images to {pixel_values.shape}")
-        else:
-            pixel_values = None
-
-        streamer = TextIteratorStreamer(
-            self.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=10
-        )
-        generation_config = dict(
-            num_beams=1,
-            max_new_tokens=max_new_tokens,
-            do_sample=do_sample,
-            temperature=temperature,
-            repetition_penalty=repetition_penalty,
-            max_length=self.context_len,
-            top_p=top_p,
-            streamer=streamer,
-        )
-        logger.info(f"Generation config: {generation_config}")
-        multi_thread_infer(
-            self.model,
-            self.tokenizer,
-            pixel_values,
-            question,
-            history,
-            generation_config,
-        )
-
-        generated_text = ""
-        for new_text in streamer:
-            generated_text += new_text
-            if generated_text.endswith(self.model.conv_template.sep):
-                generated_text = generated_text[: -len(self.model.conv_template.sep)]
-            yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
-        logger.info(
-            f"max_input_tile_list: {max_input_tile_list}, history: {history}, "
-            f"question: {question}, answer: {generated_text}"
-        )
-        self.model.system_message = old_system_message
-
-    def generate_stream_gate(self, params):
-        try:
-            for x in self.generate_stream(params):
-                yield x
-        except ValueError as e:
-            print("Caught ValueError:", e)
-            traceback.print_exc()
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            yield json.dumps(ret).encode() + b"\0"
-        except torch.cuda.CudaError as e:
-            traceback.print_exc()
-            print("Caught torch.cuda.CudaError:", e)
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            yield json.dumps(ret).encode() + b"\0"
-        except Exception as e:
-            traceback.print_exc()
-            print("Caught Unknown Error", e)
-            ret = {
-                "text": server_error_msg,
-                "error_code": 1,
-            }
-            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="0.0.0.0")
-    parser.add_argument("--port", type=int, default=21002)
-    parser.add_argument("--worker-url", type=str, default="http://localhost")
-    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
-    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
-    parser.add_argument("--model-name", type=str)
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--limit-model-concurrency", type=int, default=5)
-    parser.add_argument("--stream-interval", type=int, default=1)
-    parser.add_argument("--load-8bit", action="store_true")
-    args = parser.parse_args()
-    logger.info(f"args: {args}")
-
-    worker = ModelWorker(
-        args.controller_url,
-        args.worker_url + f":{args.port}",
-        worker_id,
-        args.model_path,
-        args.model_name,
-        args.load_8bit,
-        args.device,
-    )
-    uvicorn.run(app, host=args.host, port=args.port, log_level="info")