update gradio demo

.streamlit/config.toml

@@ -1,7 +0,0 @@
-[server]
-enableStaticServing = false
-enableXsrfProtection = false
-enableCORS = false
-
-[browser] # This ip and port will show in command prompt
-enableCORS = false

README.md

@@ -3,8 +3,8 @@ title: InternVL
 emoji: ⚡
 colorFrom: yellow
 colorTo: gray
-sdk: streamlit
-sdk_version: 1.28.2
+sdk: gradio
+sdk_version: 4.36.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,60 +1,116 @@
-import streamlit as st
-
-st.set_page_config(layout="wide")
-
-hide_streamlit_style = """
-<style>
-    /* Hide the Streamlit header and menu */
-    header {visibility: hidden;}
-</style>
-"""
-
-st.markdown(hide_streamlit_style, unsafe_allow_html=True)
-
-st.markdown(
-    """
-    <style>
-    html, body, .fullScreenFrame, .fullScreenFrame iframe {
-        margin: 0;
-        padding: 0;
-        height: 100%;
-        width: 100%;
-        border: none;
-        display: block;
-        overflow: hidden;
-    }
-
-    .fullScreenFrame {
-        position: fixed;
-        top: 0;
-        left: 0;
-        right: 0;
-        bottom: 0;
-        z-index: 9999;
-    }
-
-    .main .block-container {
-        padding: 0;
-        margin: 0;
-        height: 100vh;
-    }
-
-    /* Hide Streamlit header and footer */
-    header, footer {
-        display: none;
-    }
-    </style>
-    """,
-    unsafe_allow_html=True,
-)
-
-# Embed the external Streamlit webpage
-st.markdown(
-    """
-    <div class="fullScreenFrame">
-        <iframe src="https://internvl.opengvlab.com/"></iframe>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
+import fire
+import subprocess
+import os
+import time
+import signal
+import subprocess
+import atexit
 
+
+def kill_processes_by_cmd_substring(cmd_substring):
+    # execute `ps -ef` and obtain its output
+    result = subprocess.run(["ps", "-ef"], stdout=subprocess.PIPE, text=True)
+    lines = result.stdout.splitlines()
+
+    # visit each line
+    for line in lines:
+        if cmd_substring in line:
+            # extract PID
+            parts = line.split()
+            pid = int(parts[1])
+            print(f"Killing process with PID: {pid}, CMD: {line}")
+            os.kill(pid, signal.SIGTERM)
+
+
+def main(
+    python_path="python",
+    run_controller=True,
+    run_worker=True,
+    run_gradio=True,
+    controller_port=10086,
+    gradio_port=10087,
+    worker_names=[
+        "OpenGVLab/InternVL2-8B",
+    ],
+    run_sd_worker=False,
+    **kwargs,
+):
+    host = "http://0.0.0.0"
+    controller_process = None
+    if run_controller:
+        # python controller.py --host 0.0.0.0 --port 10086
+        cmd_args = [
+            f"{python_path}",
+            "controller.py",
+            "--host",
+            "0.0.0.0",
+            "--port",
+            f"{controller_port}",
+        ]
+        kill_processes_by_cmd_substring(" ".join(cmd_args))
+        print("Launching controller: ", " ".join(cmd_args))
+        controller_process = subprocess.Popen(cmd_args)
+        atexit.register(controller_process.terminate)
+
+    worker_processes = []
+    if run_worker:
+        worker_port = 10088
+        for worker_name in worker_names:
+            cmd_args = [
+                f"{python_path}",
+                "model_worker.py",
+                "--port",
+                f"{worker_port}",
+                "--controller-url",
+                f"{host}:{controller_port}",
+                "--model-path",
+                f"{worker_name}",
+                "--load-8bit",
+            ]
+            kill_processes_by_cmd_substring(" ".join(cmd_args))
+            print("Launching worker: ", " ".join(cmd_args))
+            worker_process = subprocess.Popen(cmd_args)
+            worker_processes.append(worker_process)
+            atexit.register(worker_process.terminate)
+            worker_port += 1
+
+    time.sleep(10)
+    gradio_process = None
+    if run_gradio:
+        #  python gradio_web_server.py --port 10088 --controller-url http://0.0.0.0:10086
+        cmd_args = [
+            f"{python_path}",
+            "gradio_web_server.py",
+            "--port",
+            f"{gradio_port}",
+            "--controller-url",
+            f"{host}:{controller_port}",
+            "--model-list-mode",
+            "reload",
+        ]
+        kill_processes_by_cmd_substring(" ".join(cmd_args))
+        print("Launching gradio: ", " ".join(cmd_args))
+        gradio_process = subprocess.Popen(cmd_args)
+        atexit.register(gradio_process.terminate)
+
+    sd_worker_process = None
+    if run_sd_worker:
+        # python model_worker.py --port 10088 --controller-address http://
+        cmd_args = [f"{python_path}", "sd_worker.py"]
+        kill_processes_by_cmd_substring(" ".join(cmd_args))
+        print("Launching sd_worker: ", " ".join(cmd_args))
+        sd_worker_process = subprocess.Popen(cmd_args)
+        atexit.register(sd_worker_process.terminate)
+
+    for worker_process in worker_processes:
+        worker_process.wait()
+    if controller_process:
+        controller_process.wait()
+    if gradio_process:
+        gradio_process.wait()
+    if sd_worker_process:
+        sd_worker_process.wait()
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

controller.py

@@ -5,9 +5,9 @@ It sends worker addresses to clients.
 import argparse
 import dataclasses
 import json
+import re
 import threading
 import time
-import re
 from enum import Enum, auto
 from typing import List
 
@@ -113,6 +113,8 @@ class Controller:
             model_names.update(w_info.model_names)
 
         def extract_key(s):
+            if 'Pro' in s:
+                return 999
             match = re.match(r'InternVL2-(\d+)B', s)
             if match:
                 return int(match.group(1))

conversation.py

@@ -0,0 +1,259 @@
+import os
+import dataclasses
+import base64
+import copy
+import hashlib
+import datetime
+from io import BytesIO
+from PIL import Image
+from typing import Any, List, Dict, Union
+from dataclasses import field
+
+from utils import LOGDIR
+
+
+def pil2base64(img: Image.Image) -> str:
+    buffered = BytesIO()
+    img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+
+def resize_img(img: Image.Image, max_len: int, min_len: int) -> Image.Image:
+    max_hw, min_hw = max(img.size), min(img.size)
+    aspect_ratio = max_hw / min_hw
+    # max_len, min_len = 800, 400
+    shortest_edge = int(min(max_len / aspect_ratio, min_len, min_hw))
+    longest_edge = int(shortest_edge * aspect_ratio)
+    W, H = img.size
+    if H > W:
+        H, W = longest_edge, shortest_edge
+    else:
+        H, W = shortest_edge, longest_edge
+    return img.resize((W, H))
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+
+    roles: List[str] = field(
+        default_factory=lambda: [
+            Conversation.SYSTEM,
+            Conversation.USER,
+            Conversation.ASSISTANT,
+        ]
+    )
+    mandatory_system_message = "我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。"
+    system_message: str = "请尽可能详细地回答用户的问题。"
+    messages: List[Dict[str, Any]] = field(default_factory=lambda: [])
+    max_image_limit: int = 4
+    skip_next: bool = False
+    streaming_placeholder: str = "▌"
+
+    def get_system_message(self):
+        return self.mandatory_system_message + "\n\n" + self.system_message
+
+    def set_system_message(self, system_message: str):
+        self.system_message = system_message
+        return self
+
+    def get_prompt(self, inlude_image=False):
+        send_messages = [{"role": "system", "content": self.get_system_message()}]
+        # send_messages = []
+        for message in self.messages:
+            if message["role"] == self.USER:
+                user_message = {
+                    "role": self.USER,
+                    "content": message["content"],
+                }
+                if inlude_image and "image" in message:
+                    user_message["image"] = []
+                    for image in message["image"]:
+                        user_message["image"].append(pil2base64(image))
+                send_messages.append(user_message)
+            elif message["role"] == self.ASSISTANT:
+                send_messages.append(
+                    {"role": self.ASSISTANT, "content": message["content"]}
+                )
+            elif message["role"] == self.SYSTEM:
+                send_messages.append(
+                    {
+                        "role": self.SYSTEM,
+                        "content": message["content"],
+                    }
+                )
+            else:
+                raise ValueError(f"Invalid role: {message['role']}")
+        return send_messages
+
+    def append_message(
+        self,
+        role,
+        content,
+        image_list=None,
+    ):
+        self.messages.append(
+            {
+                "role": role,
+                "content": content,
+                "image": [] if image_list is None else image_list,
+                # "filenames": save_filenames,
+            }
+        )
+
+    def get_images(
+        self,
+        return_copy=False,
+        return_base64=False,
+        source: Union[str, None] = None,
+    ):
+        assert source in [self.USER, self.ASSISTANT, None], f"Invalid source: {soure}"
+        images = []
+        for i, msg in enumerate(self.messages):
+            if source and msg["role"] != source:
+                continue
+
+            for image in msg.get("image", []):
+                # org_image = [i.copy() for i in image]
+                if return_copy:
+                    image = image.copy()
+
+                if return_base64:
+                    image = pil2base64(image)
+
+                images.append(image)
+
+        return images
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, msg in enumerate(self.messages):
+            if msg["role"] == self.SYSTEM:
+                continue
+
+            alt_str = (
+                "user upload image" if msg["role"] == self.USER else "output image"
+            )
+            image = msg.get("image", [])
+            if not isinstance(image, list):
+                images = [image]
+            else:
+                images = image
+
+            img_str_list = []
+            for i in range(len(images)):
+                image = resize_img(
+                    images[i],
+                    400,
+                    800,
+                )
+                img_b64_str = pil2base64(image)
+                W, H = image.size
+                img_str = f'<img src="data:image/png;base64,{img_b64_str}" alt="{alt_str}" style="width: {W}px; max-width:none; max-height:none"></img>'
+                img_str = (
+                    f'<img src="data:image/png;base64,{img_b64_str}" alt="{alt_str}" />'
+                )
+                img_str_list.append(img_str)
+
+            if msg["role"] == self.USER:
+                msg_str = " ".join(img_str_list) + msg["content"]
+                ret.append([msg_str, None])
+            else:
+                msg_str = msg["content"] + " ".join(img_str_list)
+                ret[-1][-1] = msg_str
+        return ret
+
+    def update_message(self, role, content, image=None, idx=-1):
+        assert len(self.messages) > 0, "No message in the conversation."
+
+        idx = (idx + len(self.messages)) % len(self.messages)
+
+        assert (
+            self.messages[idx]["role"] == role
+        ), f"Role mismatch: {role} vs {self.messages[idx]['role']}"
+
+        self.messages[idx]["content"] = content
+        if image is not None:
+            if image not in self.messages[idx]["image"]:
+                self.messages[idx]["image"] = []
+            if not isinstance(image, list):
+                image = [image]
+            self.messages[idx]["image"].extend(image)
+
+    def return_last_message(self):
+        return self.messages[-1]["content"]
+
+    def end_of_current_turn(self):
+        assert len(self.messages) > 0, "No message in the conversation."
+        assert (
+            self.messages[-1]["role"] == self.ASSISTANT
+        ), f"It should end with the message from assistant instead of {self.messages[-1]['role']}."
+
+        if self.messages[-1]["content"][-1] != self.streaming_placeholder:
+            return
+
+        self.update_message(self.ASSISTANT, self.messages[-1]["content"][:-1], None)
+
+    def copy(self):
+        return Conversation(
+            mandatory_system_message=self.mandatory_system_message,
+            system_message=self.system_message,
+            roles=copy.deepcopy(self.roles),
+            messages=copy.deepcopy(self.messages),
+        )
+
+    def dict(self):
+        """
+        all_images = state.get_images()
+        all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
+        t = datetime.datetime.now()
+        for image, hash in zip(all_images, all_image_hash):
+            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)
+        """
+        messages = []
+        for message in self.messages:
+            images = []
+            for image in message.get("image", []):
+                filename = self.save_image(image)
+                images.append(filename)
+
+            messages.append(
+                {
+                    "role": message["role"],
+                    "content": message["content"],
+                    "image": images,
+                }
+            )
+            if len(images) == 0:
+                messages[-1].pop("image")
+
+        return {
+            "mandatory_system_message": self.mandatory_system_message,
+            "system_message": self.system_message,
+            "roles": self.roles,
+            "messages": messages,
+        }
+
+    def save_image(self, image: Image.Image) -> str:
+        t = datetime.datetime.now()
+        image_hash = hashlib.md5(image.tobytes()).hexdigest()
+        filename = os.path.join(
+            LOGDIR,
+            "serve_images",
+            f"{t.year}-{t.month:02d}-{t.day:02d}",
+            f"{image_hash}.jpg",
+        )
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+        return filename

gradio_web_server.py

@@ -0,0 +1,824 @@
+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 == "InternVL-Chat-V1-5":
+            return (1, model_name)  # 1 indicates highest precedence
+        elif model_name.startswith("InternVL-Chat-V1-5-"):
+            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():
+    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_models(models)
+
+    logger.info(f"Models: {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 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 query_image_generation(response, sd_worker_url, timeout=15):
+    if not sd_worker_url:
+        return None
+    sd_worker_url = f"{sd_worker_url}/generate_image/"
+    pattern = r"```drawing-instruction\n(.*?)\n```"
+    match = re.search(pattern, response, re.DOTALL)
+    if match:
+        payload = {"caption": match.group(1)}
+        print("drawing-instruction:", payload)
+        response = requests.post(sd_worker_url, json=payload, timeout=timeout)
+        response.raise_for_status()  # 检查HTTP请求是否成功
+        image = Image.open(BytesIO(response.content))
+        return image
+    else:
+        return None
+
+
+def load_demo(url_params, request: gr.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):
+    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, request: gr.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) + (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()
+    if "<ref>" in ai_response:
+        returned_image = find_bounding_boxes(state, ai_response)
+        returned_image = [returned_image] if returned_image else []
+        state.update_message(Conversation.ASSISTANT, ai_response, returned_image)
+    if "```drawing-instruction" in ai_response:
+        returned_image = query_image_generation(
+            ai_response, sd_worker_url=sd_worker_url
+        )
+        returned_image = [returned_image] if returned_image else []
+        state.update_message(Conversation.ASSISTANT, ai_response, returned_image)
+
+    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);
+    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],
+            [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,
+        )
+        submit_btn.click(
+            add_text,
+            [state, textbox, 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,
+        )
+
+        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=11000)
+    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--concurrency-count", type=int, default=10)
+    parser.add_argument(
+        "--model-list-mode", type=str, default="once", choices=["once", "reload"]
+    )
+    parser.add_argument("--sd-worker-url", type=str, default=None)
+    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}")
+
+    models = get_model_list()
+
+    sd_worker_url = args.sd_worker_url
+    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,
+    )

library.py

@@ -1,95 +0,0 @@
-# --------------------------------------------------------
-# InternVL
-# Copyright (c) 2024 OpenGVLab
-# Licensed under The MIT License [see LICENSE for details]
-# Modified from https://github.com/hreikin/streamlit-uploads-library/blob/main/streamlit_uploads_library/library.py
-# --------------------------------------------------------
-
-import logging
-from math import ceil
-
-import streamlit as st
-
-logger = logging.getLogger(__name__)
-
-
-class Library():
-    """Create a simple library out of streamlit widgets.
-
-    Using the library is simple, import `streamlit_uploads_library` and then instantiate the class with the
-    required `directory` variable. Other options can be configured by passing in different variables
-    when instantiating the class.
-
-    Example Usage:
-        python
-        import streamlit as st
-        from library import Library
-
-        st.set_page_config(page_title="Streamlit Uploads Library", layout="wide")
-        default_library = Library(images=pil_images)
-    """
-
-    def __init__(self, images, image_alignment='end', number_of_columns=5):
-        self.images = images
-        self.image_alignment = image_alignment
-        self.number_of_columns = number_of_columns
-        self.root_container = self.create(images=self.images,
-                                          image_alignment=self.image_alignment,
-                                          number_of_columns=self.number_of_columns)
-
-    def create(_self, images, image_alignment, number_of_columns):
-        """Creates a simple library or gallery with columns.
-
-        Creates a library or gallery using columns out of streamlit widgets.
-        """
-        root_container = st.container()
-        with root_container:
-            # To be able to display the images, details and buttons all in one row and aligned
-            # correctly so that images of different sizes don't affect the alignment of the details
-            # and buttons we need do some minor maths and keep track of multiple index values.
-            # First we instantiate some defaults.
-            col_idx = 0
-            filename_idx = 0
-            max_idx = number_of_columns - 1
-            # Get the file list and filename list, work out the total number of files from the
-            # length of the file list.
-            library_files = images
-            num_of_files = len(library_files)
-            # Work out the number of rows required by dividing the number of files by the number of
-            # columns and rounding up using `math.ceil`.
-            num_of_rows_req = ceil(num_of_files / number_of_columns)
-            # Create the required number of rows (st.container).
-            library_rows = list()
-            library_rows_idx = 0
-            for i in range(num_of_rows_req):
-                library_rows.append(st.container())
-            # For each library row we need to create separate rows (st.container) for images,
-            # and rows (st.expander) for details and buttons to keep them in the correct columns.
-            for idx in range(num_of_rows_req):
-                with library_rows[library_rows_idx]:
-                    imgs_columns = list(st.columns(number_of_columns))
-                # Since we are keeping track of the column and filename indexes we can use
-                # those to slice the `library_files` list at the correct points for each row
-                # and then increase or reset the indexes as required.
-                for img in library_files[filename_idx:(filename_idx + number_of_columns)]:
-                    with imgs_columns[col_idx]:
-                        st.image(img, use_column_width='auto')
-                        st.write(
-                            f"""<style>
-                                [data-testid="stHorizontalBlock"] {{
-                                    align-items: {image_alignment};
-                                }}
-                                </style>
-                                """,
-                            unsafe_allow_html=True
-                        )
-                    # Keeps track of the current column, if we reach the `max_idx` we reset it
-                    # to 0 and increase the row index. This combined with the slicing should
-                    # ensure all images, details and buttons are in the correct columns.
-                    if col_idx < max_idx:
-                        col_idx += 1
-                    else:
-                        col_idx = 0
-                        library_rows_idx += 1
-                    filename_idx += 1
-        return root_container

mm_utils.py

@@ -1,102 +0,0 @@
-import base64
-from io import BytesIO
-
-import torch
-from PIL import Image
-from transformers import StoppingCriteria
-
-from .constants import IMAGE_TOKEN_INDEX
-
-
-def load_image_from_base64(image):
-    return Image.open(BytesIO(base64.b64decode(image)))
-
-
-def expand2square(pil_img, background_color):
-    width, height = pil_img.size
-    if width == height:
-        return pil_img
-    elif width > height:
-        result = Image.new(pil_img.mode, (width, width), background_color)
-        result.paste(pil_img, (0, (width - height) // 2))
-        return result
-    else:
-        result = Image.new(pil_img.mode, (height, height), background_color)
-        result.paste(pil_img, ((height - width) // 2, 0))
-        return result
-
-
-def process_images(images, image_processor, model_cfg):
-    image_aspect_ratio = getattr(model_cfg, 'image_aspect_ratio', None)
-    new_images = []
-    if image_aspect_ratio == 'pad':
-        for image in images:
-            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))
-            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]
-            new_images.append(image)
-    else:
-        return image_processor(images, return_tensors='pt')['pixel_values']
-    if all(x.shape == new_images[0].shape for x in new_images):
-        new_images = torch.stack(new_images, dim=0)
-    return new_images
-
-
-def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX,
-                          num_image_tokens=None, return_tensors=None):
-    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split('<image>')]
-
-    def insert_separator(X, sep):
-        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]
-
-    input_ids = []
-    offset = 0
-    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:
-        offset = 1
-        input_ids.append(prompt_chunks[0][0])
-
-    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + num_image_tokens)):
-        input_ids.extend(x[offset:])
-
-    if return_tensors is not None:
-        if return_tensors == 'pt':
-            return torch.tensor(input_ids, dtype=torch.long)
-        raise ValueError(f'Unsupported tensor type: {return_tensors}')
-    return input_ids
-
-
-def get_model_name_from_path(model_path):
-    model_path = model_path.strip('/')
-    model_paths = model_path.split('/')
-    if model_paths[-1].startswith('checkpoint-'):
-        return model_paths[-2] + '_' + model_paths[-1]
-    else:
-        return model_paths[-1]
-
-
-class KeywordsStoppingCriteria(StoppingCriteria):
-    def __init__(self, keywords, tokenizer, input_ids):
-        self.keywords = keywords
-        self.keyword_ids = []
-        self.max_keyword_len = 0
-        for keyword in keywords:
-            cur_keyword_ids = tokenizer(keyword).input_ids
-            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:
-                cur_keyword_ids = cur_keyword_ids[1:]
-            if len(cur_keyword_ids) > self.max_keyword_len:
-                self.max_keyword_len = len(cur_keyword_ids)
-            self.keyword_ids.append(torch.tensor(cur_keyword_ids))
-        self.tokenizer = tokenizer
-        self.start_len = input_ids.shape[1]
-
-    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        assert output_ids.shape[0] == 1, 'Only support batch size 1 (yet)'  # TODO
-        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)
-        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]
-        for keyword_id in self.keyword_ids:
-            if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all():
-                return True
-        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]
-        for keyword in self.keywords:
-            if keyword in outputs:
-                return True
-        return False

model_worker.py

@@ -9,14 +9,15 @@ A model worker executes the model.
 """
 import argparse
 import asyncio
-import base64
+
 import json
-import os
+import math
 import threading
 import time
 import uuid
+import traceback
 from functools import partial
-from io import BytesIO
+
 from threading import Thread
 
 import requests
@@ -28,33 +29,36 @@ from fastapi import BackgroundTasks, FastAPI, Request
 from fastapi.responses import StreamingResponse
 from PIL import Image
 from torchvision.transforms.functional import InterpolationMode
-from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          TextIteratorStreamer)
-from utils import build_logger, pretty_print_semaphore, server_error_msg
+from transformers import AutoModel, AutoTokenizer, TextIteratorStreamer
+from utils import (
+    build_logger,
+    pretty_print_semaphore,
+    server_error_msg,
+    load_image_from_base64,
+)
+import spaces
 
 worker_id = str(uuid.uuid4())[:6]
-logger = build_logger('model_worker', f'model_worker_{worker_id}.log')
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
 global_counter = 0
 model_semaphore = None
 
 
-def load_image_from_base64(image):
-    return Image.open(BytesIO(base64.b64decode(image)))
-
-
 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)
-    ])
+    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_diff = float("inf")
     best_ratio = (1, 1)
     area = width * height
     for ratio in target_ratios:
@@ -69,19 +73,26 @@ def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_
     return best_ratio
 
 
-def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
+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)
+        (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)
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size
+    )
 
     # calculate the target width and height
     target_width = image_size * target_aspect_ratio[0]
@@ -96,7 +107,7 @@ def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnai
             (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
+            ((i // (target_width // image_size)) + 1) * image_size,
         )
         # split the image
         split_img = resized_img.crop(box)
@@ -114,78 +125,163 @@ def heart_beat_worker(controller):
         controller.send_heart_beat()
 
 
+def split_model(model_name):
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    num_layers = {
+        "InternVL2-8B": 32,
+        "InternVL2-26B": 48,
+        "InternVL2-40B": 60,
+        "InternVL2-Llama3-76B": 80,
+        "InternVL2-78B": 80,
+        "InternVL2-Pro": 80,
+    }[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
+
+
 class ModelWorker:
-    def __init__(self, controller_addr, worker_addr, worker_id, model_path, model_name,
-                 load_8bit, device, context_len=8192):
+    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('/'):
+        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]
+            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
 
-        logger.info(f'Loading the model {self.model_name} on worker {worker_id} ...')
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
-        if device == 'auto':
-            os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
-            # This can make distributed deployment work properly
-            self.model = AutoModelForCausalLM.from_pretrained(
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_path, trust_remote_code=True, use_fast=False
+        )
+        tokens_to_keep = ["<box>", "</box>", "<ref>", "</ref>"]
+        tokenizer.additional_special_tokens = [
+            item
+            for item in tokenizer.additional_special_tokens
+            if item not in tokens_to_keep
+        ]
+        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.float16,
-                device_map='auto',
-                trust_remote_code=True).eval()
+                torch_dtype=torch.bfloat16,
+                device_map=device_map,
+                trust_remote_code=True,
+            ).eval()
         else:
-            self.model = AutoModelForCausalLM.from_pretrained(
+            self.model = AutoModel.from_pretrained(
                 model_path,
                 load_in_8bit=load_8bit,
-                torch_dtype=torch.float16,
-                trust_remote_code=True).eval()
-        if not load_8bit and not device == 'auto':
+                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,))
+            target=heart_beat_worker, args=(self,)
+        )
         self.heart_beat_thread.start()
 
+    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')
+        logger.info("Register to controller")
 
-        url = self.controller_addr + '/register_worker'
+        url = self.controller_addr + "/register_worker"
         data = {
-            'worker_name': self.worker_addr,
-            'check_heart_beat': True,
-            'worker_status': self.get_status()
+            "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}')
+        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'
+        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']
+                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}')
+                logger.error(f"heart beat error: {e}")
             time.sleep(5)
 
         if not exist:
@@ -195,80 +291,115 @@ class ModelWorker:
         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)
+            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(),
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
         }
 
+    @spaces.GPU
     @torch.inference_mode()
     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']
+        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 = 1
+        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:
+            if message["role"] == "user":
+                prefix = ""
+                if "image" in message:
                     max_input_tile_temp = []
-                    for image_str in message['image']:
+                    for image_str in message["image"]:
                         pil_images.append(load_image_from_base64(image_str))
-                        prefix += f'Image-{global_image_cnt}: <image>\n\n'
+                        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'])))
+                        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, ])
+                content = prefix + message["content"]
+                history.append(
+                    [
+                        content,
+                    ]
+                )
             else:
-                history[-1].append(message['content'])
+                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
+            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.'
-        logger.info(f'max_input_tile_list: {max_input_tile_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):
+            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)
+                        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.float16)
-            logger.info(f'Split images to {pixel_values.shape}')
+            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=False, timeout=10)
+        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,
@@ -279,53 +410,61 @@ class ModelWorker:
             top_p=top_p,
             streamer=streamer,
         )
-        logger.info(history)
-        logger.info(f'Generation config: {generation_config}')
-        try:
-            thread = Thread(target=self.model.chat, kwargs=dict(
+        logger.info(f"Generation config: {generation_config}")
+
+        thread = Thread(
+            target=self.model.chat,
+            kwargs=dict(
                 tokenizer=self.tokenizer,
                 pixel_values=pixel_values,
                 question=question,
                 history=history,
                 return_history=False,
                 generation_config=generation_config,
-            ))
-            thread.start()
-    
-            generated_text = ''
-            for new_text in streamer:
-                generated_text += new_text
-                yield json.dumps({'text': generated_text.replace(self.model.conv_template.sep, ''),
-                                  'error_code': 0}).encode() + b'\0'
-            self.model.system_message = old_system_message
-        except:
-            torch.cuda.empty_cache()
+            ),
+        )
+        thread.start()
+
+        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)
+            print("Caught ValueError:", e)
+            traceback.print_exc()
             ret = {
-                'text': server_error_msg,
-                'error_code': 1,
+                "text": server_error_msg,
+                "error_code": 1,
             }
-            yield json.dumps(ret).encode() + b'\0'
+            yield json.dumps(ret).encode() + b"\0"
         except torch.cuda.CudaError as e:
-            print('Caught torch.cuda.CudaError:', e)
+            traceback.print_exc()
+            print("Caught torch.cuda.CudaError:", e)
             ret = {
-                'text': server_error_msg,
-                'error_code': 1,
+                "text": server_error_msg,
+                "error_code": 1,
             }
-            yield json.dumps(ret).encode() + b'\0'
+            yield json.dumps(ret).encode() + b"\0"
         except Exception as e:
-            print('Caught Unknown Error', e)
+            traceback.print_exc()
+            print("Caught Unknown Error", e)
             ret = {
-                'text': server_error_msg,
-                'error_code': 1,
+                "text": server_error_msg,
+                "error_code": 1,
             }
-            yield json.dumps(ret).encode() + b'\0'
+            yield json.dumps(ret).encode() + b"\0"
 
 
 app = FastAPI()
@@ -337,7 +476,7 @@ def release_model_semaphore(fn=None):
         fn()
 
 
-@app.post('/worker_generate_stream')
+@app.post("/worker_generate_stream")
 async def generate_stream(request: Request):
     global model_semaphore, global_counter
     global_counter += 1
@@ -349,35 +488,39 @@ async def generate_stream(request: Request):
     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))
+    background_tasks.add_task(
+        partial(release_model_semaphore, fn=worker.send_heart_beat)
+    )
     return StreamingResponse(generator, background=background_tasks)
 
 
-@app.post('/worker_get_status')
+@app.post("/worker_get_status")
 async def get_status(request: Request):
     return worker.get_status()
 
 
-if __name__ == '__main__':
+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-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-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')
+    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_address,
-                         args.worker_address,
-                         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')
+    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")

requirements.txt

@@ -1,4 +1,14 @@
-opencv-python
-streamlit_image_select
-streamlit==1.36.0
-flask
+diffusers==0.29.2
+fastapi==0.111.1
+filelock==3.15.4
+fire==0.6.0
+gradio==4.38.1
+numpy==2.0.1
+Pillow==10.4.0
+pydantic==2.8.2
+Requests==2.32.3
+spaces==0.28.3
+torch==2.0.1
+torchvision==0.15.2
+transformers==4.37.2
+uvicorn==0.30.3

utils.py

@@ -1,13 +1,22 @@
+from ast import Dict
 import logging
 import logging.handlers
 import os
 import sys
-
+import base64
+from PIL import Image
+from io import BytesIO
+import json
 import requests
 from constants import LOGDIR
+import datetime
 
-server_error_msg = '**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**'
-moderation_msg = 'YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN.'
+server_error_msg = (
+    "**NETWORK ERROR DUE TO HIGH TRAFFIC. PLEASE REGENERATE OR REFRESH THIS PAGE.**"
+)
+moderation_msg = (
+    "YOUR INPUT VIOLATES OUR CONTENT MODERATION GUIDELINES. PLEASE TRY AGAIN."
+)
 
 handler = None
 
@@ -16,8 +25,8 @@ def build_logger(logger_name, logger_filename):
     global handler
 
     formatter = logging.Formatter(
-        fmt='%(asctime)s | %(levelname)s | %(name)s | %(message)s',
-        datefmt='%Y-%m-%d %H:%M:%S',
+        fmt="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
     )
 
     # Set the format of root handlers
@@ -26,12 +35,12 @@ def build_logger(logger_name, logger_filename):
     logging.getLogger().handlers[0].setFormatter(formatter)
 
     # Redirect stdout and stderr to loggers
-    stdout_logger = logging.getLogger('stdout')
+    stdout_logger = logging.getLogger("stdout")
     stdout_logger.setLevel(logging.INFO)
     sl = StreamToLogger(stdout_logger, logging.INFO)
     sys.stdout = sl
 
-    stderr_logger = logging.getLogger('stderr')
+    stderr_logger = logging.getLogger("stderr")
     stderr_logger.setLevel(logging.ERROR)
     sl = StreamToLogger(stderr_logger, logging.ERROR)
     sys.stderr = sl
@@ -45,7 +54,8 @@ def build_logger(logger_name, logger_filename):
         os.makedirs(LOGDIR, exist_ok=True)
         filename = os.path.join(LOGDIR, logger_filename)
         handler = logging.handlers.TimedRotatingFileHandler(
-            filename, when='D', utc=True)
+            filename, when="D", utc=True
+        )
         handler.setFormatter(formatter)
 
         for name, item in logging.root.manager.loggerDict.items():
@@ -59,33 +69,34 @@ class StreamToLogger(object):
     """
     Fake file-like stream object that redirects writes to a logger instance.
     """
+
     def __init__(self, logger, log_level=logging.INFO):
         self.terminal = sys.stdout
         self.logger = logger
         self.log_level = log_level
-        self.linebuf = ''
+        self.linebuf = ""
 
     def __getattr__(self, attr):
         return getattr(self.terminal, attr)
 
     def write(self, buf):
         temp_linebuf = self.linebuf + buf
-        self.linebuf = ''
+        self.linebuf = ""
         for line in temp_linebuf.splitlines(True):
             # From the io.TextIOWrapper docs:
             #   On output, if newline is None, any '\n' characters written
             #   are translated to the system default line separator.
             # By default sys.stdout.write() expects '\n' newlines and then
             # translates them so this is still cross platform.
-            if line[-1] == '\n':
+            if line[-1] == "\n":
                 self.logger.log(self.log_level, line.rstrip())
             else:
                 self.linebuf += line
 
     def flush(self):
-        if self.linebuf != '':
+        if self.linebuf != "":
             self.logger.log(self.log_level, self.linebuf.rstrip())
-        self.linebuf = ''
+        self.linebuf = ""
 
 
 def disable_torch_init():
@@ -93,23 +104,26 @@ def disable_torch_init():
     Disable the redundant torch default initialization to accelerate model creation.
     """
     import torch
-    setattr(torch.nn.Linear, 'reset_parameters', lambda self: None)
-    setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None)
+
+    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
+    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
 
 
 def violates_moderation(text):
     """
     Check whether the text violates OpenAI moderation API.
     """
-    url = 'https://api.openai.com/v1/moderations'
-    headers = {'Content-Type': 'application/json',
-               'Authorization': 'Bearer ' + os.environ['OPENAI_API_KEY']}
-    text = text.replace('\n', '')
-    data = '{' + '"input": ' + f'"{text}"' + '}'
-    data = data.encode('utf-8')
+    url = "https://api.openai.com/v1/moderations"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer " + os.environ["OPENAI_API_KEY"],
+    }
+    text = text.replace("\n", "")
+    data = "{" + '"input": ' + f'"{text}"' + "}"
+    data = data.encode("utf-8")
     try:
         ret = requests.post(url, headers=headers, data=data, timeout=5)
-        flagged = ret.json()['results'][0]['flagged']
+        flagged = ret.json()["results"][0]["flagged"]
     except requests.exceptions.RequestException as e:
         flagged = False
     except KeyError as e:
@@ -120,5 +134,30 @@ def violates_moderation(text):
 
 def pretty_print_semaphore(semaphore):
     if semaphore is None:
-        return 'None'
-    return f'Semaphore(value={semaphore._value}, locked={semaphore.locked()})'
+        return "None"
+    return f"Semaphore(value={semaphore._value}, locked={semaphore.locked()})"
+
+
+def load_image_from_base64(image):
+    return Image.open(BytesIO(base64.b64decode(image)))
+
+
+def get_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+
+def data_wrapper(data):
+    if isinstance(data, bytes):
+        return data
+    elif isinstance(data, Image.Image):
+        buffered = BytesIO()
+        data.save(buffered, format="PNG")
+        return buffered.getvalue()
+    elif isinstance(data, str):
+        return data.encode()
+    elif isinstance(data, Dict):
+        return json.dumps(data).encode()
+    else:
+        raise ValueError(f"Unsupported data type: {type(data)}")