optimize

controller.log

@@ -29,3 +29,11 @@
 2025-03-23 15:04:32 | ERROR | stderr | INFO:     Waiting for application startup.
 2025-03-23 15:04:32 | ERROR | stderr | INFO:     Application startup complete.
 2025-03-23 15:04:32 | ERROR | stderr | INFO:     Uvicorn running on http://0.0.0.0:10000 (Press CTRL+C to quit)
+2025-03-24 14:06:11 | INFO | controller | args: Namespace(host='0.0.0.0', port=10000, dispatch_method='shortest_queue')
+2025-03-24 14:06:11 | INFO | controller | Init controller
+2025-03-24 14:06:11 | ERROR | stderr | INFO:     Started server process [95537]
+2025-03-24 14:06:11 | ERROR | stderr | INFO:     Waiting for application startup.
+2025-03-24 14:06:11 | ERROR | stderr | INFO:     Application startup complete.
+2025-03-24 14:06:11 | ERROR | stderr | ERROR:    [Errno 48] error while attempting to bind on address ('0.0.0.0', 10000): address already in use
+2025-03-24 14:06:11 | ERROR | stderr | INFO:     Waiting for application shutdown.
+2025-03-24 14:06:11 | ERROR | stderr | INFO:     Application shutdown complete.

model_worker_ad9563.log

@@ -1,17 +0,0 @@
-2025-03-23 15:01:04 | INFO | model_worker | args: Namespace(host='0.0.0.0', port=40000, worker_address='http://localhost:40000', controller_address='http://localhost:10000', model_name='/home/agent_h/data/starvector-1b-im2svg', multi_modal=False, limit_model_concurrency=5, stream_interval=1, no_register=False, openai_api_key='EMPTY', vllm_base_url='http://localhost:8000')
-2025-03-23 15:01:04 | INFO | model_worker | Loading the model /home/agent_h/data/starvector-1b-im2svg on worker ad9563 ...
-2025-03-23 15:01:04 | INFO | model_worker | Register to controller
-2025-03-23 15:01:04 | ERROR | stderr | INFO:     Started server process [48407]
-2025-03-23 15:01:04 | ERROR | stderr | INFO:     Waiting for application startup.
-2025-03-23 15:01:04 | ERROR | stderr | INFO:     Application startup complete.
-2025-03-23 15:01:04 | ERROR | stderr | INFO:     Uvicorn running on http://0.0.0.0:40000 (Press CTRL+C to quit)
-2025-03-23 15:01:19 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:01:34 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:01:49 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:02:04 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:02:19 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:02:34 | INFO | model_worker | Send heart beat. Models: ['/home/agent_h/data/starvector-1b-im2svg']. Semaphore: None. global_counter: 0
-2025-03-23 15:02:45 | ERROR | stderr | INFO:     Shutting down
-2025-03-23 15:02:45 | ERROR | stderr | INFO:     Waiting for application shutdown.
-2025-03-23 15:02:45 | ERROR | stderr | INFO:     Application shutdown complete.
-2025-03-23 15:02:45 | ERROR | stderr | INFO:     Finished server process [48407]

star-vector-dev/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

star-vector-dev/.gitignore

@@ -0,0 +1,181 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Other
+*vscode*
+*egg*
+*nfs*
+*conv.json*
+*rebuttal*
+*.log*
+*remove_files*
+*wandb*
+*tmp*
+*vscode*
+*.csv
+*avoid_samples*
+*logs*
+*results*
+*.pickle
+*.pkl
+*internal*
+*test.png*
+assets/reward_assets

start.sh

@@ -2,7 +2,6 @@
 
 bash -c "$SSH_TUNNEL_CMD_1" &
 
-echo "SSH tunnel started, PID: $SSH_PID"
 python -m starvector.serve.vllm_api_gradio.controller --host 0.0.0.0 --port 10000 &
 python -m starvector.serve.vllm_api_gradio.model_worker --host 0.0.0.0 --controller http://localhost:10000 --port 40000 --worker http://localhost:40000 --model-name /home/agent_h/data/starvector-1b-im2svg --vllm-base-url http://localhost:8000 &
 python -m starvector.serve.vllm_api_gradio.gradio_web_server --controller http://localhost:10000 --model-list-mode reload --port 7860

starvector/serve/controller.py

@@ -0,0 +1,293 @@
+"""
+A controller manages distributed workers.
+It sends worker addresses to clients.
+"""
+import argparse
+import asyncio
+import dataclasses
+from enum import Enum, auto
+import json
+import logging
+import time
+from typing import List, Union
+import threading
+
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+import numpy as np
+import requests
+import uvicorn
+
+from starvector.serve.constants import CONTROLLER_HEART_BEAT_EXPIRATION
+from starvector.serve.util import build_logger, server_error_msg
+
+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)
+
+        return list(model_names)
+
+    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
+
+            # Check status before returning
+            while True:
+                pt = np.random.choice(np.arange(len(worker_names)),
+                    p=worker_speeds)
+                worker_name = worker_names[pt]
+
+                if self.get_worker_status(worker_name):
+                    break
+                else:
+                    self.remove_worker(worker_name)
+                    worker_speeds[pt] = 0
+                    norm = np.sum(worker_speeds)
+                    if norm < 1e-4:
+                        return ""
+                    worker_speeds = worker_speeds / norm
+                    continue
+            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="localhost")
+    parser.add_argument("--port", type=int, default=21001)
+    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")

starvector/serve/gradio_demo_with_updated_gradio.py

@@ -0,0 +1,432 @@
+import argparse
+import datetime
+import json
+import os
+import time
+import gradio as gr
+import requests
+from starvector.serve.conversation import default_conversation
+from starvector.serve.constants import LOGDIR, CLIP_QUERY_LENGTH
+from starvector.serve.util import (build_logger, server_error_msg)
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+headers = {"User-Agent": "StarVector Client"}
+
+no_change_btn = gr.Button()
+enable_btn = gr.Button(interactive=True)
+disable_btn = gr.Button(interactive=False)
+
+priority = {
+    "starvector-1.4b": "aaaaaaa",
+}
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+def get_model_list():
+    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    assert ret.status_code == 200
+    ret = requests.post(args.controller_url + "/list_models")
+    models = ret.json()["models"]
+    models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"Models: {models}")
+    return models
+
+get_window_url_params = """
+function() {
+    const params = new URLSearchParams(window.location.search);
+    url_params = Object.fromEntries(params);
+    console.log(url_params);
+    return url_params;
+    }
+"""
+
+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 = default_conversation.copy()
+    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 = default_conversation.copy()
+    dropdown_update = gr.Dropdown(
+        choices=models,
+        value=models[0] if len(models) > 0 else ""
+    )
+    return state, dropdown_update
+
+def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "model": model_selector,
+            "state": state.dict(),
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+def upvote_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"upvote. ip: {request.client.host}")
+    vote_last_response(state, "upvote", model_selector, request)
+    return ("",) + (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, "downvote", model_selector, request)
+    return ("",) + (disable_btn,) * 3
+
+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)
+    return ("",) + (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
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, None, None, None) + (disable_btn,) * 6
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = default_conversation.copy()
+    return (state, None, None) + (disable_btn,) * 6
+
+def send_image(state, image, image_process_mode, request: gr.Request):
+    logger.info(f"send_image. ip: {request.client.host}.")
+    state.stop_sampling = False
+    if image is None:
+        state.skip_next = True
+        return (state, None, None, image) + (no_change_btn,) * 6
+
+    if image is not None:
+        text = (image, image_process_mode)
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], "▌")
+    state.skip_next = False
+    msg = state.to_gradio_svg_code()[0][1]
+    return (state, msg, state.to_gradio_svg_render(), image) + (no_change_btn,) * 6
+
+def stop_sampling(state, image, request: gr.Request):
+    logger.info(f"stop_sampling. ip: {request.client.host}")
+    state.stop_sampling = True
+    return (state, None, None, image) + (disable_btn,) * 6
+
+def http_bot(state, model_selector, num_beams, temperature, len_penalty, top_p, max_new_tokens, request: gr.Request):
+    logger.info(f"http_bot. ip: {request.client.host}")
+    start_tstamp = time.time()
+    model_name = model_selector
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (state, None, None) + (no_change_btn,) * 6
+        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
+        yield (state, None, None, disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn)
+        return
+
+    # Construct prompt
+    prompt = state.get_prompt()
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "num_beams": int(num_beams),
+        "temperature": float(temperature),
+        "len_penalty": float(len_penalty),
+        "top_p": float(top_p),
+        "max_new_tokens": min(int(max_new_tokens), 8192-CLIP_QUERY_LENGTH),
+    }
+    logger.info(f"==== request ====\n{pload}")
+
+    pload['images'] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, disable_btn, disable_btn, enable_btn)
+
+    try:
+        # Stream output
+        if state.stop_sampling:
+            state.messages[1][-1] = "▌"
+            yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn)
+            return
+        
+        response = requests.post(worker_addr + "/worker_generate_stream",
+            headers=headers, json=pload, stream=True, timeout=100)
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    # output = data["text"].strip().replace('<', '&lt;').replace('>', '&gt;') # trick to avoid the SVG getting rendered
+                    output = data["text"].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    st = state.to_gradio_svg_code()
+                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, enable_btn)
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    
+                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn)
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException as e:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, None, None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn)
+        return
+
+    yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (enable_btn,) * 6
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "svg": state.messages[-1][-1],
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+title_markdown = ("""
+# 💫 StarVector: Generating Scalable Vector Graphics Code from Images and Text
+[[Project Page](https://starvector.github.io)] [[Code](https://github.com/joanrod/star-vector)] [[Model](https://huggingface.co/joanrodai/starvector-1.4b)] | 📚 [[StarVector](https://arxiv.org/abs/2312.11556)] 
+""")
+
+sub_title_markdown = (""" Throw an image and vectorize it! The model expects vector-like images to generate the corresponding svg code.""")
+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. Please contact us if you find any potential violation.
+""")
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+.gradio-container{
+    max-width: 1200px!important
+}
+
+#svg_render{
+    padding: 20px !important;
+}
+
+#svg_code{
+    height: 200px !important; 
+    overflow: scroll !important; 
+    white-space: unset !important; 
+    flex-shrink: unset !important;
+}
+
+
+h1{display: flex;align-items: center;justify-content: center;gap: .25em}
+*{transition: width 0.5s ease, flex-grow 0.5s ease}
+"""
+
+def build_demo(embed_mode, concurrency_count=10):
+    with gr.Blocks(title="StarVector", theme=gr.themes.Default(), css=block_css) as demo:
+        state = gr.State()
+        if not embed_mode:
+            gr.Markdown(title_markdown)
+            gr.Markdown(sub_title_markdown)
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Row(elem_id="model_selector_row"):
+                    model_selector = gr.Dropdown(
+                        choices=models,
+                        value=models[0] if len(models) > 0 else "",
+                        interactive=True,
+                        show_label=False,
+                        container=False)
+                imagebox = gr.Image(type="pil")
+                image_process_mode = gr.Radio(
+                    ["Resize", "Pad", "Default"],
+                    value="Pad",
+                    label="Preprocess for non-square image", visible=False)
+
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                gr.Examples(examples=[
+                    [f"{cur_dir}/examples/sample-4.png"],
+                    [f"{cur_dir}/examples/sample-7.png"],
+                    [f"{cur_dir}/examples/sample-16.png"],
+                    [f"{cur_dir}/examples/sample-17.png"],
+                    [f"{cur_dir}/examples/sample-18.png"],
+                    [f"{cur_dir}/examples/sample-0.png"],
+                    [f"{cur_dir}/examples/sample-1.png"],
+                    [f"{cur_dir}/examples/sample-6.png"],
+                ], inputs=[imagebox])
+                
+                with gr.Column(scale=1, min_width=50):
+                    submit_btn = gr.Button(value="Send", variant="primary")
+
+                with gr.Accordion("Parameters", open=True) as parameter_row:
+                    num_beams = gr.Slider(minimum=1, maximum=10, value=1, step=1, interactive=True, label="Num Beams", visible=False,)
+                    temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.8, step=0.05, interactive=True, label="Temperature",)
+                    len_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=0.6, step=0.05, interactive=True, label="Length Penalty",)
+                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, interactive=True, label="Top P",)
+                    max_output_tokens = gr.Slider(minimum=0, maximum=8192, value=2000, step=64, interactive=True, label="Max output tokens",)
+
+            with gr.Column(scale=8):
+                with gr.Row():
+                    svg_code = gr.Code(label="SVG Code", elem_id='svg_code', min_width=200, interactive=False, lines=5)
+                with gr.Row():
+                     gr.Image(width=50, height=256, label="Rendered SVG", elem_id='svg_render')
+                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, visible=False)
+                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False, visible=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, stop_btn]
+        upvote_btn.click(
+            upvote_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+        downvote_btn.click(
+            downvote_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+        flag_btn.click(
+            flag_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+
+        regenerate_btn.click(
+            regenerate,
+            [state, image_process_mode],
+            [state, svg_code, svg_render, imagebox] + btn_list,
+            queue=False
+        ).then(
+            http_bot,
+            [state, model_selector, num_beams, temperature, len_penalty,  top_p, max_output_tokens],
+            [state, svg_code, svg_render] + btn_list,
+            concurrency_limit=concurrency_count
+        )
+
+        submit_btn.click(
+            send_image,
+            [state, imagebox, image_process_mode],
+            [state, svg_code, svg_render, imagebox] + btn_list,
+            queue=False
+        ).then(
+            http_bot,
+            [state, model_selector, num_beams, temperature, len_penalty, top_p, max_output_tokens],
+            [state, svg_code, svg_render] + btn_list,
+            concurrency_limit=concurrency_count
+        )
+
+        clear_btn.click(
+            clear_history,
+            None,
+            [state, svg_code, svg_render] + btn_list,
+            queue=False
+        )
+
+        stop_btn.click(
+            stop_sampling,
+            [state, imagebox],
+            [state, imagebox] + btn_list,
+            queue=False
+        ).then(
+            clear_history,
+            None,
+            [state, svg_code, svg_render] + btn_list,
+            queue=False
+        )
+
+        if args.model_list_mode == "once":
+            demo.load(
+                load_demo,
+                [url_params],
+                [state, model_selector],
+                _js=get_window_url_params,
+            )
+        elif args.model_list_mode == "reload":
+            demo.load(
+                load_demo_refresh_model_list,
+                None,
+                [state, model_selector],
+                queue=False
+            )
+        else:
+            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
+
+    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)
+    parser.add_argument("--controller-url", type=str, default="http://localhost:21001")
+    parser.add_argument("--concurrency-count", type=int, default=15)
+    parser.add_argument("--model-list-mode", type=str, default="once", 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}")
+
+    models = get_model_list()
+
+    logger.info(args)
+    demo = build_demo(args.embed, concurrency_count=args.concurrency_count)
+    demo.queue(
+        api_open=False
+    ).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share
+    )

starvector/serve/gradio_web_server.py

@@ -0,0 +1,562 @@
+import argparse
+import datetime
+import json
+import os
+import time
+import gradio as gr
+import requests
+from starvector.serve.conversation import default_conversation
+from starvector.serve.constants import LOGDIR, CLIP_QUERY_LENGTH
+from starvector.serve.util import (build_logger, server_error_msg)
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+headers = {"User-Agent": "StarVector Client"}
+
+no_change_btn = gr.Button.update()
+enable_btn = gr.Button.update(interactive=True)
+disable_btn = gr.Button.update(interactive=False)
+
+priority = {
+    "starvector-1b-im2svg": "aaaaaaa",
+}
+
+def get_conv_log_filename():
+    t = datetime.datetime.now()
+    name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
+    return name
+
+def get_model_list():
+    ret = requests.post(args.controller_url + "/refresh_all_workers")
+    assert ret.status_code == 200
+    ret = requests.post(args.controller_url + "/list_models")
+    models = ret.json()["models"]
+    models.sort(key=lambda x: priority.get(x, x))
+    logger.info(f"Models: {models}")
+    return models
+
+def load_demo(url_params, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
+
+    dropdown_update = gr.Dropdown.update(visible=True)
+    if "model" in url_params:
+        model = url_params["model"]
+        if model in models:
+            dropdown_update = gr.Dropdown.update(
+                value=model, visible=True)
+
+    state = default_conversation.copy()
+    return state, dropdown_update
+
+mapping_model_task = {
+    'Image2SVG': 'im2svg',
+    'Text2SVG': 'text2svg'
+}
+
+def get_models_dropdown_from_task(task):
+    models = get_model_list()
+    models = [model for model in models if mapping_model_task[task] in model]
+    dropdown_update = gr.Dropdown.update(
+        choices=models,
+        value=models[0] if len(models) > 0 else ""
+    )
+    return dropdown_update
+    
+
+def load_demo_refresh_model_list(task, request: gr.Request):
+    logger.info(f"load_demo. ip: {request.client.host}")
+    dropdown_update = get_models_dropdown_from_task(task)
+    state = default_conversation.copy()
+    return state, dropdown_update
+
+def vote_last_response(state, vote_type, model_selector, request: gr.Request):
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(time.time(), 4),
+            "type": vote_type,
+            "model": model_selector,
+            "state": state.dict(),
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+def upvote_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"upvote. ip: {request.client.host}")
+    vote_last_response(state, "upvote", model_selector, request)
+    return ("",) + (disable_btn,) * 7
+
+def downvote_last_response(state, model_selector, request: gr.Request):
+    logger.info(f"downvote. ip: {request.client.host}")
+    vote_last_response(state, "downvote", model_selector, request)
+    return ("",) + (disable_btn,) * 7
+
+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)
+    return ("",) + (disable_btn,) * 7
+
+def regenerate(state, image_process_mode, request: gr.Request):
+    logger.info(f"regenerate. ip: {request.client.host}")
+    state.messages[-1][-1] = None
+    prev_human_msg = state.messages[-2]
+    if type(prev_human_msg[1]) in (tuple, list):
+        prev_human_msg[1] = (prev_human_msg[1][:2], image_process_mode)
+    state.skip_next = False
+    return (state, None, None, None) + (disable_btn,) * 7
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = default_conversation.copy()
+    return (state, None, None) + (disable_btn,) * 7
+
+def send_data(state, image, image_process_mode, text_caption, task, request: gr.Request):
+    logger.info(f"send_data. ip: {request.client.host}.")
+    if task == 'Image2SVG':
+        if image is None:
+            state.skip_next = True
+            return (state, None, None, image) + (no_change_btn,) * 7
+
+        if image is not None:
+            image_message = (image, image_process_mode)
+        state.append_message(state.roles[0], image_message)
+        state.append_message(state.roles[1], "▌")
+        state.skip_next = False
+        msg = state.to_gradio_svg_code()[0][1]
+        return (state, msg, state.to_gradio_svg_render(), image) + (no_change_btn,) * 7
+    else:
+        if text_caption is None:
+            state.skip_next = True
+            return (state, None, None, image) + (no_change_btn,) * 7
+
+        state.append_message(state.roles[0], text_caption)
+        state.append_message(state.roles[1], "▌")
+        state.skip_next = False
+        msg = state.to_gradio_svg_code()[0][1]
+        return (state, msg, state.to_gradio_svg_render(), image) + (no_change_btn,) * 7
+
+def download_files(state, request: gr.Request):
+    logger.info(f"download_files. ip: {request.client.host}")
+    svg_str, image = state.download_files()
+    
+    # TODO: Figure out how to download the SVG in the users browser, idk how to do it now 
+            
+def update_task(task):
+    dropdown_update = get_models_dropdown_from_task(task)
+
+    if task == "Text2SVG":
+        return 1.0, 0.9, 0.95, dropdown_update
+    else:
+        return 0.6, 0.9, 0.95, dropdown_update
+    
+    
+def stop_sampling(state, image, request: gr.Request):
+    logger.info(f"stop_sampling. ip: {request.client.host}")
+    state.stop_sampling = True
+    return (state, None, None, image) + (disable_btn,) * 7
+
+def http_bot(state, task_selector, text_caption, model_selector, num_beams, temperature, len_penalty, top_p, max_new_tokens, request: gr.Request):
+    logger.info(f"http_bot. ip: {request.client.host}")
+    start_tstamp = time.time()
+    model_name = model_selector
+
+    if state.skip_next:
+        # This generate call is skipped due to invalid inputs
+        yield (state, None, None) + (no_change_btn,) * 7
+        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
+        yield (state, None, None, disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, disable_btn)
+        return
+
+    # Construct prompt
+    if task_selector == "Image2SVG":
+        prompt = state.get_image_prompt()
+    else:
+        prompt = text_caption
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "num_beams": int(num_beams),
+        "temperature": float(temperature),
+        "len_penalty": float(len_penalty),
+        "top_p": float(top_p),
+        "max_new_tokens": min(int(max_new_tokens), 8192-CLIP_QUERY_LENGTH),
+    }
+    logger.info(f"==== request ====\n{pload}")
+
+    pload['images'] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+
+    try:
+        # Stream output
+        if state.stop_sampling:
+            state.messages[1][-1] = "▌"
+            yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, enable_btn)
+            return
+        
+        response = requests.post(worker_addr + "/worker_generate_stream",
+            headers=headers, json=pload, stream=True, timeout=10)
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    # output = data["text"].strip().replace('<', '&lt;').replace('>', '&gt;') # trick to avoid the SVG getting rendered
+                    output = data["text"].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    st = state.to_gradio_svg_code()
+                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, enable_btn, enable_btn)
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    
+                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, disable_btn)
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException as e:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, None, None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, disable_btn)
+        return
+
+    yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (enable_btn,) * 7
+
+    finish_tstamp = time.time()
+    logger.info(f"{output}")
+
+    with open(get_conv_log_filename(), "a") as fout:
+        data = {
+            "tstamp": round(finish_tstamp, 4),
+            "type": "chat",
+            "model": model_name,
+            "start": round(start_tstamp, 4),
+            "finish": round(finish_tstamp, 4),
+            "svg": state.messages[-1][-1],
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+title_markdown = ("""
+# 💫 StarVector: Generating Scalable Vector Graphics Code from Images and Text
+
+[[Project Page](https://starvector.github.io)] [[Code](https://github.com/joanrod/star-vector)] [[Model](https://huggingface.co/joanrodai/starvector-1.4b)] | 📚 [[StarVector](https://arxiv.org/abs/2312.11556)]""")
+
+sub_title_markdown = ("""**How does it work?** Select the task you want to perform, and the model will be automatically set. For **Text2SVG**, introduce a prompt in Text Caption. For **Image2SVG**, select an image and vectorize it. \ 
+**Note**: The current model works on vector-like images like icons and or vector-like designs.""")
+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. Please contact us if you find any potential violation.
+""")
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+.gradio-container{
+    max-width: 1200px!important
+}
+
+.ͼ1 .cm-content {
+    white-space: unset !important; 
+    flex-shrink: unset !important;
+}
+
+.ͼ2p .cm-scroller {
+    max-height: 200px;
+    overflow: scroll;
+}
+
+#svg_render{
+    padding: 20px !important;
+}
+
+#submit_btn{
+    max-height: 40px;
+}
+
+.selector{
+    max-height: 100px;
+}
+h1{display: flex;align-items: center;justify-content: center;gap: .25em}
+*{transition: width 0.5s ease, flex-grow 0.5s ease}
+"""
+def build_demo(embed_mode):
+    svg_render = gr.Image(label="Rendered SVG", elem_id='svg_render', height=300)
+    svg_code = gr.Code(label="SVG Code", elem_id='svg_code', interactive=True, lines=5)
+
+    with gr.Blocks(title="StarVector", theme=gr.themes.Default(), css=block_css) as demo:
+        state = gr.State()
+        if not embed_mode:
+            gr.Markdown(title_markdown)
+            gr.Markdown(sub_title_markdown)
+        with gr.Row():
+            with gr.Column(scale=4):
+                task_selector = gr.Dropdown(
+                    choices=["Image2SVG", "Text2SVG"],
+                    value="Image2SVG",
+                    label="Task",
+                    interactive=True,
+                    show_label=True,
+                    container=True,
+                    elem_id="task_selector",
+                    elem_classes=["selector"],
+                )
+                model_selector = gr.Dropdown(
+                    choices=models,
+                    value=models[0] if len(models) > 0 else "",
+                    label="Model",
+                    interactive=True,
+                    show_label=True,
+                    container=True,
+                    elem_classes=["selector"],
+                )
+                    
+                imagebox = gr.Image(type="pil", visible=True, elem_id="imagebox")
+                image_process_mode = gr.Radio(
+                    ["Resize", "Pad", "Default"],
+                    value="Pad",
+                    label="Preprocess for non-square image", visible=False)
+                
+                # Text input
+                text_caption = gr.Textbox(label="Text Caption", visible=True, value="The icon of a yellow star", elem_id="text_caption")
+                
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                gr.Examples(examples=[
+                    [f"{cur_dir}/examples/sample-4.png"],
+                    [f"{cur_dir}/examples/sample-7.png"],
+                    [f"{cur_dir}/examples/sample-16.png"],
+                    [f"{cur_dir}/examples/sample-17.png"],
+                    [f"{cur_dir}/examples/sample-18.png"],
+                    [f"{cur_dir}/examples/sample-0.png"],
+                    [f"{cur_dir}/examples/sample-1.png"],
+                    [f"{cur_dir}/examples/sample-6.png"],
+                ], inputs=[imagebox], elem_id="examples")
+                
+                submit_btn = gr.Button(value="Send", variant="primary", elem_id="submit_btn", interactive=True)
+
+                with gr.Accordion("Parameters", open=False):
+                    num_beams = gr.Slider(minimum=1, maximum=10, value=1, step=1, interactive=True, label="Num Beams", visible=False,)
+                    temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.9, step=0.05, interactive=True, label="Temperature",)
+                    len_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=0.6, step=0.05, interactive=True, label="Length Penalty",)
+                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, interactive=True, label="Top P",)
+                    max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=1024, step=64, interactive=True, label="Max output tokens",)
+
+            with gr.Column(scale=9):
+                with gr.Row():
+                    svg_code.render()
+                with gr.Row():
+                    svg_render.render()
+
+                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, visible=False)
+                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False, visible=False)
+                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)  
+                    download_btn = gr.Button(value="Download SVG", interactive=False, visible=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, stop_btn, download_btn]
+        upvote_btn.click(
+            upvote_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+        downvote_btn.click(
+            downvote_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+        flag_btn.click(
+            flag_last_response,
+            [state, model_selector],
+            [upvote_btn, downvote_btn, flag_btn],
+            queue=False
+        )
+
+        regenerate_btn.click(
+            regenerate,
+            [state, image_process_mode],
+            [state, svg_code, svg_render, imagebox] + btn_list,
+            queue=False
+        ).then(
+            http_bot,
+            [state, task_selector, text_caption, model_selector, num_beams, temperature, len_penalty,  top_p, max_output_tokens],
+            [state, svg_code, svg_render] + btn_list)
+
+        submit_btn.click(
+            send_data,
+            [state, imagebox, image_process_mode, text_caption, task_selector],
+            [state, svg_code, svg_render, imagebox] + btn_list,
+            queue=False
+        ).then(
+            http_bot,
+            [state, task_selector, text_caption, model_selector, num_beams, temperature, len_penalty, top_p, max_output_tokens],
+            [state, svg_code, svg_render] + btn_list
+        )
+
+        clear_btn.click(
+            clear_history,
+            None,
+            [state, svg_code, svg_render] + btn_list,
+            queue=False
+        )
+
+        stop_btn.click(
+            stop_sampling,
+            [state, imagebox],
+            [state, imagebox] + btn_list,
+            queue=False
+        ).then(
+            clear_history,
+            None,
+            [state, svg_code, svg_render] + btn_list,
+            queue=False
+        )
+        
+        download_btn.click(
+            download_files,
+            [state],
+            None,
+            queue=False
+        )
+        task_selector.change(
+            update_task,            
+            inputs=[task_selector],
+            outputs=[len_penalty, temperature, top_p, model_selector],
+            queue=False,
+            _js="""
+                    function(task) {
+                        var imageBoxElement = document.getElementById("imagebox");
+                        var textCaptionElement = document.getElementById("text_caption");
+                        var examplesElement = document.getElementById("examples");
+                        if (task === "Text2SVG") {
+                            imageBoxElement.style.display = "none";
+                            textCaptionElement.style.display = "block";
+                            examplesElement.style.display = "none";
+                        } else if (task === "Image2SVG") {
+                            imageBoxElement.style.display = "block";
+                            textCaptionElement.style.display = "none";
+                            examplesElement.style.display = "block";
+                        }
+                        return task;
+                    }
+                """
+        )
+        
+        if args.model_list_mode == "once":
+            demo.load(
+                load_demo,
+                [url_params, task_selector],
+                [state, model_selector],
+                _js="""
+                    function() {
+                        const params = new URLSearchParams(window.location.search);
+                        url_params = Object.fromEntries(params);
+                        console.log(url_params);
+                        return url_params;
+                        
+                        }
+                    """,
+                queue=False
+            )
+        elif args.model_list_mode == "reload":
+            demo.load(
+                load_demo_refresh_model_list,
+                [task_selector],
+                [state, model_selector],
+                _js="""
+                    function(task) {
+                        var textCaptionElement = document.getElementById("text_caption");
+                        var autoScrollBottom = true;
+                        textCaptionElement.style.display = "none";
+                        function updateScroll(){
+                            if (autoScrollBottom) {
+                                var element = document.getElementsByClassName("cm-scroller")[0];
+                                element.scrollTop = element.scrollHeight;
+                            }
+                        }
+                        function handleScroll() {
+                            var element = document.getElementsByClassName("cm-scroller")[0];
+                            //if (element.scrollHeight - element.scrollTop === element.clientHeight) {
+                            if (element.scrollHeight - (element.scrollTop + element.clientHeight) < 0.2*(element.scrollTop)) {
+                                // User has scrolled to the bottom, enable auto-scrolling
+                                autoScrollBottom = true;
+                                console.log("bottom");
+                            } else {
+                                console.log("not bottom");
+                                // User has scrolled away from the bottom, disable auto-scrolling
+                                autoScrollBottom = false;
+                            }
+                        }
+                        setInterval(updateScroll,500);
+                        var element = document.getElementsByClassName("cm-scroller")[0];
+                        element.addEventListener("scroll", handleScroll);
+    
+                        return task;
+                    }
+
+                    """,
+                queue=False,
+            )
+                
+        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)
+    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("--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()
+
+    logger.info(args)
+    demo = build_demo(args.embed)
+    demo.queue(
+        concurrency_count=args.concurrency_count,
+        api_open=False
+    ).launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share
+    )

starvector/serve/model_worker.py

@@ -0,0 +1,269 @@
+"""
+A model worker executes the model.
+"""
+import argparse
+import asyncio
+import json
+import time
+import threading
+import uuid
+from fastapi import FastAPI, Request, BackgroundTasks
+from fastapi.responses import StreamingResponse
+import requests
+import torch
+import uvicorn
+from functools import partial
+from starvector.serve.constants import WORKER_HEART_BEAT_INTERVAL, CLIP_QUERY_LENGTH
+from starvector.serve.util import (build_logger, server_error_msg,
+    pretty_print_semaphore)
+from starvector.model.builder import load_pretrained_model
+from starvector.serve.util import process_images, load_image_from_base64
+from threading import Thread
+from transformers import TextIteratorStreamer
+
+GB = 1 << 30
+
+worker_id = str(uuid.uuid4())[:6]
+logger = build_logger("model_worker", f"model_worker_{worker_id}.log")
+global_counter = 0
+model_semaphore = None
+
+def heart_beat_worker(controller):
+    while True:
+        time.sleep(WORKER_HEART_BEAT_INTERVAL)
+        controller.send_heart_beat()
+
+class ModelWorker:
+    def __init__(self, controller_addr, worker_addr,
+                 worker_id, no_register,
+                 model_path, model_base, model_name,
+                 load_8bit, load_4bit, device):
+        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
+
+        if "text2svg" in self.model_name.lower():
+            self.task = "Text2SVG"
+        elif "im2svg" in self.model_name.lower():
+            self.task = "Image2SVG"
+            
+        self.device = device
+        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
+            model_path, device=self.device, load_in_8bit=load_8bit, load_in_4bit=load_4bit)
+        self.model.to(torch.bfloat16)
+        self.is_multimodal = 'starvector' in self.model_name.lower()
+
+        if not no_register:
+            self.register_to_controller()
+            self.heart_beat_thread = threading.Thread(
+                target=heart_beat_worker, args=(self,))
+            self.heart_beat_thread.start()
+
+    def register_to_controller(self):
+        logger.info("Register to controller")
+
+        url = self.controller_addr + "/register_worker"
+        data = {
+            "worker_name": self.worker_addr,
+            "check_heart_beat": True,
+            "worker_status": self.get_status()
+        }
+        r = requests.post(url, json=data)
+        assert r.status_code == 200
+
+    def send_heart_beat(self):
+        logger.info(f"Send heart beat. Models: {[self.model_name]}. "
+                    f"Semaphore: {pretty_print_semaphore(model_semaphore)}. "
+                    f"global_counter: {global_counter}")
+
+        url = self.controller_addr + "/receive_heart_beat"
+
+        while True:
+            try:
+                ret = requests.post(url, json={
+                    "worker_name": self.worker_addr,
+                    "queue_length": self.get_queue_length()}, timeout=5)
+                exist = ret.json()["exist"]
+                break
+            except requests.exceptions.RequestException as e:
+                logger.error(f"heart beat error: {e}")
+            time.sleep(5)
+
+        if not exist:
+            self.register_to_controller()
+
+    def get_queue_length(self):
+        if model_semaphore is None:
+            return 0
+        else:
+            return args.limit_model_concurrency - model_semaphore._value + (len(
+                model_semaphore._waiters) if model_semaphore._waiters is not None else 0)
+
+    def get_status(self):
+        return {
+            "model_names": [self.model_name],
+            "speed": 1,
+            "queue_length": self.get_queue_length(),
+        }
+
+    @torch.inference_mode()
+    def generate_stream(self, params):
+        tokenizer, model, image_processor, task = self.tokenizer, self.model, self.image_processor, self.task
+                 
+        num_beams = int(params.get("num_beams", 1))
+        temperature = float(params.get("temperature", 1.0))
+        len_penalty = float(params.get("len_penalty", 1.0))
+        top_p = float(params.get("top_p", 1.0))
+        max_context_length = getattr(model.config, 'max_position_embeddings', 8192)
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=False, skip_special_tokens=True, timeout=15)
+        prompt = params["prompt"]
+             
+        if task == "Image2SVG":
+            images = params.get("images", None)
+            for b64_image in images: 
+                if b64_image is not None and self.is_multimodal:
+                    image = load_image_from_base64(b64_image)            
+                    image = process_images(image, image_processor)
+                    image = image.to(self.model.device, dtype=torch.float16)
+                else:
+                    image = None
+           
+            max_new_tokens = min(int(params.get("max_new_tokens", 256)), 8192)
+            max_new_tokens = min(max_new_tokens, max_context_length - CLIP_QUERY_LENGTH)
+            pre_pend = prompt 
+            batch = {}
+            batch["image"] = image
+            generate_method = model.model.generate_im2svg
+        else:
+            max_new_tokens = min(int(params.get("max_new_tokens", 128)), 8192)
+            pre_pend = ""
+            batch = {}
+            batch['caption'] = [prompt]
+            # White PIL image
+            batch['image'] = torch.zeros((3, 256, 256), dtype=torch.float16).to(self.model.device)
+            generate_method = model.model.generate_text2svg
+            
+        if max_new_tokens < 1:
+            yield json.dumps({"text": prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
+            return
+
+        thread = Thread(target=generate_method, kwargs=dict(
+            batch=batch,
+            prompt=prompt,
+            use_nucleus_sampling=True,
+            num_beams=num_beams,
+            temperature=temperature,
+            length_penalty=len_penalty,
+            top_p=top_p,
+            max_length=max_new_tokens,
+            streamer=streamer,
+        ))
+        thread.start()
+
+        generated_text = pre_pend
+        for new_text in streamer:
+            if new_text == " ":
+                continue
+            generated_text += new_text
+            # if generated_text.endswith(stop_str):
+            #     generated_text = generated_text[:-len(stop_str)]
+            yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0"
+    
+    def generate_stream_gate(self, params):
+        try:
+            for x in self.generate_stream(params):
+                yield x
+        except ValueError as e:
+            print("Caught ValueError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except torch.cuda.CudaError as e:
+            print("Caught torch.cuda.CudaError:", e)
+            ret = {
+                "text": server_error_msg,
+                "error_code": 1,
+            }
+            yield json.dumps(ret).encode() + b"\0"
+        except Exception as e:
+            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="localhost")
+    parser.add_argument("--port", type=int, default=21002)
+    parser.add_argument("--worker-address", type=str,
+        default="http://localhost:21002")
+    parser.add_argument("--controller-address", type=str,
+        default="http://localhost:21001")
+    parser.add_argument("--model-path", type=str, default="joanrodai/starvector-1.4b")
+    parser.add_argument("--model-base", type=str, default=None)
+    parser.add_argument("--model-name", type=str)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--multi-modal", action="store_true", help="Multimodal mode is automatically detected with model name, please make sure `starvector` is included in the model path.")
+    parser.add_argument("--limit-model-concurrency", type=int, default=5)
+    parser.add_argument("--stream-interval", type=int, default=1)
+    parser.add_argument("--no-register", action="store_true")
+    parser.add_argument("--load-8bit", action="store_true")
+    parser.add_argument("--load-4bit", action="store_true")
+    args = parser.parse_args()
+    logger.info(f"args: {args}")
+
+    if args.multi_modal:
+        logger.warning("Multimodal mode is automatically detected with model name, please make sure `starvector` is included in the model path.")
+
+    worker = ModelWorker(args.controller_address,
+                         args.worker_address,
+                         worker_id,
+                         args.no_register,
+                         args.model_path,
+                         args.model_base,
+                         args.model_name,
+                         args.load_8bit,
+                         args.load_4bit,
+                         args.device)
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")

starvector/serve/vllm_api_gradio/gradio_web_server.py

@@ -204,7 +204,6 @@ def http_bot(state, task_selector, text_caption, model_selector, num_beams, temp
 
     state.messages[-1][-1] = "▌"
     yield (state, state.messages[-1][-1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
-
     try:
         # Stream output
         if state.stop_sampling:
@@ -214,23 +213,33 @@ def http_bot(state, task_selector, text_caption, model_selector, num_beams, temp
         
         response = requests.post(worker_addr + "/worker_generate_stream",
             headers=headers, json=pload, stream=True, timeout=10)
+
+        update_interval = 2  # seconds
+        last_update_time = time.time()
+
         for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
             if chunk:
                 data = json.loads(chunk.decode())
                 if data["error_code"] == 0:
-                    # output = data["text"].strip().replace('<', '&lt;').replace('>', '&gt;') # trick to avoid the SVG getting rendered
                     output = data["text"].strip()
                     state.messages[-1][-1] = output + "▌"
-                    st = state.to_gradio_svg_code()
-                    # Explicitly set the string value without HTML escaping
-                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, enable_btn, enable_btn)
+
+                    # Only update if sufficient time has passed
+                    current_time = time.time()
+                    if current_time - last_update_time >= update_interval:
+                        st = state.to_gradio_svg_code()
+                        yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn,) * 7
+                        last_update_time = current_time
                 else:
+                    # handle errors and yield immediately if needed
                     output = data["text"] + f" (error_code: {data['error_code']})"
                     state.messages[-1][-1] = output
                     st = state.to_gradio_svg_code()
-
-                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, disable_btn)
+                    yield (state, st[-1][1], state.to_gradio_svg_render()) + (disable_btn,) * 7
                     return
+        # Final yield to ensure the last state is rendered
+        st = state.to_gradio_svg_code()
+        yield (state, st[-1][1], state.to_gradio_svg_render()) + (enable_btn,) * 7
     except requests.exceptions.RequestException as e:
         state.messages[-1][-1] = server_error_msg
         yield (state, None, None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn, disable_btn, disable_btn)
@@ -576,7 +585,7 @@ def build_demo(embed_mode):
                     temperature = gr.Slider(minimum=0.0, maximum=2.0, value=0.2, step=0.05, interactive=True, label="Temperature",)
                     len_penalty = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.05, interactive=True, label="Length Penalty",)
                     top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, interactive=True, label="Top P",)
-                    max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=1024, step=64, interactive=True, label="Max output tokens",)
+                    max_output_tokens = gr.Slider(minimum=0, maximum=8192, value=2048, step=64, interactive=True, label="Max output tokens",)
 
             with gr.Column(scale=9):
                 with gr.Row():

starvector/serve/vllm_api_gradio/model_worker.py

@@ -117,7 +117,7 @@ class ModelWorker:
         temperature = float(params.get("temperature", 1.0))
         len_penalty = float(params.get("len_penalty", 1.0))
         top_p = float(params.get("top_p", 1.0))
-        max_context_length = 1000
+        max_context_length = 8192
 
         # prompt = params["prompt"]
         prompt = "<svg "
@@ -132,6 +132,8 @@ class ModelWorker:
 
             max_new_tokens = min(int(params.get("max_new_tokens", 256)), 8192)
             max_new_tokens = min(max_new_tokens, max_context_length - CLIP_QUERY_LENGTH)
+            # log max new token
+            logger.info(f"max_new_tokens: {max_new_tokens}")
 
             # Use the chat completions endpoint
             vllm_endpoint = f"{self.vllm_base_url}/v1/chat/completions"