optimize

starvector/serve/controller.py

@@ -1,293 +0,0 @@
-"""
-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

@@ -1,432 +0,0 @@
-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

@@ -1,562 +0,0 @@
-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=8192, value=8192, 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

@@ -1,269 +0,0 @@
-"""
-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

@@ -231,7 +231,6 @@ def http_bot(state, task_selector, text_caption, model_selector, num_beams, temp
 
                     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.01)
     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)