first commit

.gitignore

@@ -0,0 +1,150 @@
+# 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/
+pip-wheel-metadata/
+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/
+
+# 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
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.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
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Logs
+serve_images/
+
+# 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/
+
+# Ruff
+.ruff_cache/
+
+# IDE caches
+.idea/
+.vscode/
+
+# Mac OS
+.DS_Store
+
+# Tokens
+.hf_token
+
+# Scratch & Caches
+__scratch/
+scratch/
+cache/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024-present, Toyota Research Institute.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Makefile

@@ -0,0 +1,17 @@
+.PHONY: help check autoformat
+.DEFAULT: help
+
+# Generates a useful overview/help message for various make features - add to this as necessary!
+help:
+	@echo "make check"
+	@echo "    Run code style and linting (black, ruff) *without* changing files!"
+	@echo "make autoformat"
+	@echo "    Run code styling (black, ruff) and update in place - committing with pre-commit also does this."
+
+check:
+	black --check .
+	ruff check --show-source .
+
+autoformat:
+	black .
+	ruff check --fix --show-fixes .

README.md

@@ -0,0 +1,48 @@
+# VLM Demo
+
+> *VLM Demo*: Lightweight repo for chatting with models loaded into *VLM Bench*.
+
+---
+
+## Installation
+
+This repository 
+
+```bash
+git clone [email protected]:TRI-ML/vlm-demo.git
+cd vlm-demo
+pip install -e .
+```
+
+This repository also requires that the `vlm-bench` package (`vlbench`) and 
+`prismatic-vlms` package (`prisma`) are installed in the current environment.
+These can both be installed from source from the following git repos:
+
+`vlm-bench`: `https://github.com/TRI-ML/vlm-bench`
+`prismatic-vlms`: `https://github.com/TRI-ML/prismatic-vlms`
+
+## Usage
+
+Start Gradio Controller: `serve/gradio_controller.py`
+Start Gradio Web Server: `serve/gradio_web_server.py`
+Run interactive demo: `interactive_demo.py`
+
+To run the demo, run the following commands:
+
+Start Gradio Controller: `python -m serve.controller --host 0.0.0.0 --port 10000`
+Start Gradio Web Server: `python -m serve.gradio_web_server --controller http://localhost:10000 --model-list-mode reload --share`
+Run interactive demo: `CUDA_VISIBLE_DEVICES=0 python -m interactive_demo  --port 40000  --model_dir <PATH TO MODEL CKPT>`
+
+## Contributing
+
+Before committing to the repository, *make sure to set up your dev environment!*
+
+Here are the basic development environment setup guidelines:
+
++ Fork/clone the repository, performing an editable installation. Make sure to install with the development dependencies
+  (e.g., `pip install -e ".[dev]"`); this will install `black`, `ruff`, and `pre-commit`.
+
++ Install `pre-commit` hooks (`pre-commit install`).
+
++ Branch for the specific feature/issue, issuing PR against the upstream repository for review.
+

interactive_demo.py

@@ -0,0 +1,289 @@
+"""
+interactive_demo.py
+
+Entry point for all VLM-Bench interactive demos; specify model and get a gradio UI where you can chat with it!
+
+This file is heavily adapted from the script used to serve models in the LLaVa repo:
+https://github.com/haotian-liu/LLaVA/blob/main/llava/serve/model_worker.py. It is
+modified to ensure compatibility with our Prismatic models.
+"""
+import asyncio
+import json
+import os
+import threading
+import time
+import uuid
+from dataclasses import dataclass
+from functools import partial
+from pathlib import Path
+from typing import Union
+
+import draccus
+import requests
+import torch
+import uvicorn
+from accelerate.utils import set_seed
+from fastapi import BackgroundTasks, FastAPI, Request
+from fastapi.responses import StreamingResponse
+from llava.constants import WORKER_HEART_BEAT_INTERVAL
+from llava.mm_utils import load_image_from_base64
+from llava.utils import server_error_msg
+from torchvision.transforms import Compose
+
+from vlbench.models import load_vlm
+from vlbench.overwatch import initialize_overwatch
+from serve import INTERACTION_MODES_MAP, MODEL_ID_TO_NAME
+
+GB = 1 << 30
+worker_id = str(uuid.uuid4())[:6]
+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, vlm, model_base, model_name):
+        self.controller_addr = controller_addr
+        self.worker_addr = worker_addr
+        self.worker_id = worker_id
+        self.model_name = model_name
+
+        # logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
+        self.vlm = vlm
+        self.tokenizer, self.model, self.image_processor, self.context_len = (
+            vlm.tokenizer,
+            vlm.model,
+            vlm.image_processor,
+            vlm.max_length,
+        )
+
+        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:
+                pass
+                # 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 (
+                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):
+        prompt = params["prompt"]
+        ori_prompt = prompt
+        images = params.get("images", None)
+
+        temperature = params.get("temperature", 0.2)
+        max_new_tokens = params.get("max_new_tokens", 2048)
+        interaction_mode = INTERACTION_MODES_MAP[params.get("interaction_mode", "Chat")]
+
+        if temperature != 0:
+            self.vlm.set_generate_kwargs(
+                {"do_sample": True, "max_new_tokens": max_new_tokens, "temperature": temperature}
+            )
+        else:
+            self.vlm.set_generate_kwargs({"do_sample": False, "max_new_tokens": max_new_tokens})
+
+        if images is not None and len(images) == 1:
+            images = [load_image_from_base64(image) for image in images]
+        else:
+            raise NotImplementedError("Only supports queries with one image for now")
+
+        if interaction_mode == "chat":
+            question_prompt = [prompt]
+        else:
+            prompt_fn = self.vlm.get_prompt_fn(interaction_mode)
+            if interaction_mode != "captioning":
+                question_prompt = [prompt_fn(prompt)]
+            else:
+                question_prompt = [prompt_fn()]
+
+        if isinstance(self.image_processor, Compose) or hasattr(self.image_processor, "is_prismatic"):
+            # This is a standard `torchvision.transforms` object or custom PrismaticVLM wrapper
+            pixel_values = self.image_processor(images[0].convert("RGB"))
+        else:
+            # Assume `image_transform` is a HF ImageProcessor...
+            pixel_values = self.image_processor(images[0].convert("RGB"), return_tensors="pt")["pixel_values"][0]
+
+        generated_text = self.vlm.generate_answer(torch.unsqueeze(pixel_values.cuda(), 0), question_prompt)[0]
+        generated_text = generated_text.split("USER")[0].split("ASSISTANT")[0]
+        yield json.dumps({"text": ori_prompt + 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(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()
+
+
+# Initialize Overwatch =>> Wraps `logging.Logger` and `accelerate.PartialState`
+overwatch = initialize_overwatch(__name__)
+
+
+@dataclass
+class DemoConfig:
+    # fmt: off
+
+    # === Model Parameters =>> Quartz ===
+    model_family: str = "quartz"                    # Model family to load from in < `quartz` | `llava-v15` | ... >
+    model_id: str = "llava-v1.5-7b"     # Model ID to load and run (instance of `model_family`)
+    model_dir: Path = (                             # Path to model checkpoint to load --> should be self-contained
+        "resize-naive-siglip-vit-l-16-384px-no-align-2-epochs+13b+stage-finetune+x7"
+    )
+
+    # === Model Parameters =>> Official LLaVa ===
+    # model_family: str = "llava-v15"
+    # model_id: str = "llava-v1.5-13b"
+    # model_dir: Path = "liuhaotian/llava-v1.5-13b"
+
+    # Model Worker Parameters
+    host: str = "0.0.0.0"
+    port: int = 40000
+    controller_address: str = "http://localhost:10000"
+    model_base: str = "llava-v15"
+    limit_model_concurrency: int = 5
+    stream_interval: int = 1
+    no_register: bool = False
+
+    # Inference Parameters
+    device_batch_size: int = 1                      # Device Batch Size set to 1 until LLaVa/HF LLaMa fixes bugs!
+    num_workers: int = 2                            # Number of Dataloader Workers (on each process)
+
+    # HF Hub Credentials (for LLaMa-2)
+    hf_token: Union[str, Path] = Path(".hf_token")  # Environment variable or Path to HF Token
+
+    # Randomness
+    seed: int = 21                                  # Random Seed (for reproducibility)
+
+    def __post_init__(self) -> None:
+        if self.model_family == "quartz":
+            self.model_name = MODEL_ID_TO_NAME[str(self.model_dir)]
+            self.run_dir = Path("/mnt/fsx/x-onyx-vlms/runs") / self.model_dir
+        elif self.model_family in {"instruct-blip", "llava", "llava-v15"}:
+            self.model_name = MODEL_ID_TO_NAME[self.model_id]
+            self.run_dir = self.model_dir
+        else:
+            raise ValueError(f"Run Directory for `{self.model_family = }` does not exist!")
+        self.worker_address = f"http://localhost:{self.port}"
+
+    # fmt: on
+
+
+@draccus.wrap()
+def interactive_demo(cfg: DemoConfig):
+    # overwatch.info(f"Starting Evaluation for Dataset `{cfg.dataset.dataset_id}` w/ Model `{cfg.model_id}`")
+    set_seed(cfg.seed)
+
+    # Build the VLM --> Download/Load Pretrained Model from Checkpoint
+    overwatch.info("Initializing VLM =>> Bundling Models, Image Processors, and Tokenizer")
+    hf_token = cfg.hf_token.read_text().strip() if isinstance(cfg.hf_token, Path) else os.environ[cfg.hf_token]
+    vlm = load_vlm(cfg.model_family, cfg.model_id, cfg.run_dir, hf_token=hf_token)
+
+    global worker
+    global limit_model_concurrency
+    limit_model_concurrency = cfg.limit_model_concurrency
+    worker = ModelWorker(
+        cfg.controller_address, cfg.worker_address, worker_id, cfg.no_register, vlm, cfg.model_base, cfg.model_name
+    )
+    uvicorn.run(app, host=cfg.host, port=cfg.port, log_level="info")
+
+
+if __name__ == "__main__":
+    interactive_demo()

pyproject.toml

@@ -0,0 +1,64 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "vldemo"
+authors = [
+    {name = "Siddharth Karamcheti", email="[email protected]"}
+]
+description = "VLM Demo: Interactive Demo for VLMs"
+version = "0.0.1"
+readme = "README.md"
+requires-python = ">=3.8"
+keywords = ["machine learning"]
+license = {file = "LICENSE"}
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+
+]
+
+[project.optional-dependencies]
+dev = [
+    "black",
+    "gpustat",
+    "ipython",
+    "pre-commit",
+    "ruff",
+]
+
+[project.urls]
+homepage = "https://github.com/TRI-ML/vlm-demo"
+repository = "https://github.com/TRI-ML/vlm-demo"
+documentation = "https://github.com/TRI-ML/vlm-demo"
+
+[tool.setuptools.packages.find]
+where = ["."]
+exclude = ["cache"]
+
+[tool.black]
+line-length = 121
+target-version = ["py38", "py39", "py310"]
+preview = true
+
+[tool.ruff]
+line-length = 121
+target-version = "py38"
+select = ["A", "B", "C90", "E", "F", "I", "RUF", "W"]
+ignore = ["B008", "F722"]
+
+[tool.ruff.per-file-ignores]
+"__init__.py" = ["E402", "F401"]

serve/__init__.py

@@ -0,0 +1,44 @@
+from collections import OrderedDict
+
+
+# Arrange keys in display priority order (high --> low)
+MODEL_ID_TO_NAME = OrderedDict(
+    [
+        (
+            "llava-lvis4v-lrv+lvis4v-lrv-resize-naive-clip-vit-l-14-336px-no-align-2-epochs-llama2pure+13b+stage-finetune+x7",
+            "Prism-CLIP 13B",
+        ),
+        (
+            "llava-lvis4v-lrv+lvis4v-lrv-resize-naive-clip-vit-l-14-336px-no-align-2-epochs-llama2pure+7b+stage-finetune+x7",
+            "Prism-CLIP 7B",
+        ),
+        (
+            "resize-naive-clip-vit-l-14-336px-no-align-llama2pure+13b+stage-finetune+x7",
+            "Prism-CLIP 13B (Controlled)",
+        ),
+        (
+            "resize-naive-clip-vit-l-14-336px-no-align-llama2pure+7b+stage-finetune+x7",
+            "Prism-CLIP 7B (Controlled)",
+        ),
+        (
+            "resize-naive-clip-vit-l-14-336px-no-align+13b+stage-finetune+x7",
+            "Prism-CLIP 13B (Controlled) - Chat",
+        ),
+        (
+            "resize-naive-clip-vit-l-14-336px-no-align+7b+stage-finetune+x7",
+            "Prism-CLIP 7B (Controlled) - Chat",
+        ),
+        ("llava-v1.5-7b", "LLaVA 1.5: 7B"),
+        ("llava-v1.5-13b", "LLaVA 1.5: 13B"),
+    ]
+)
+
+INTERACTION_MODES_MAP = OrderedDict(
+    [
+        ("Chat", "chat"),
+        ("Captioning", "captioning"),
+        ("Bounding Box Prediction", "bbox_pred"),
+        ("Visual Question Answering", "vqa"),
+        ("True/False Visual Question Answering", "true_false"),
+    ]
+)

serve/controller.py

@@ -0,0 +1,298 @@
+"""
+controller.py
+A controller manages distributed workers.
+It sends worker addresses to clients.
+
+This file is exactly copied from
+https://github.com/haotian-liu/LLaVA/blob/main/llava/serve/controller.py.
+"""
+import argparse
+import dataclasses
+import json
+import threading
+import time
+from enum import Enum, auto
+from typing import List
+
+import numpy as np
+import requests
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import StreamingResponse
+from llava.constants import CONTROLLER_HEART_BEAT_EXPIRATION
+from llava.utils 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("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_lottery(self, model_name: str):
+        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
+
+    def get_worker_address_shortest_queue(self, model_name: str):
+        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
+
+    def get_worker_address(self, model_name: str):
+        if self.dispatch_method == DispatchMethod.LOTTERY:
+            return self.get_worker_address_lottery(model_name)
+        elif self.dispatch_method == DispatchMethod.SHORTEST_QUEUE:
+            return self.get_worker_address_shortest_queue(model_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:
+            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():
+    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")

serve/gradio_web_server.py

@@ -0,0 +1,462 @@
+"""
+gradio_web_server.py
+
+Entry point for all VLM-Bench interactive demos; specify model and get a gradio UI where you can chat with it!
+
+This file is copied from the script used to define the gradio web server in the LLaVa codebase:
+https://github.com/haotian-liu/LLaVA/blob/main/llava/serve/gradio_web_server.py with only very minor
+modifications.
+"""
+
+import argparse
+import datetime
+import hashlib
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+from llava.constants import LOGDIR
+from llava.conversation import conv_templates, default_conversation
+from llava.utils import build_logger, moderation_msg, server_error_msg, violates_moderation
+
+from serve import INTERACTION_MODES_MAP, MODEL_ID_TO_NAME
+
+logger = build_logger("gradio_web_server", "gradio_web_server.log")
+
+headers = {"User-Agent": "PrismaticVLMs Client"}
+
+no_change_btn = gr.Button.update()
+enable_btn = gr.Button.update(interactive=True)
+disable_btn = gr.Button.update(interactive=False)
+
+
+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 = sorted(
+        models, key=lambda x: list(MODEL_ID_TO_NAME.values()).index(x) if x in MODEL_ID_TO_NAME.values() else len(models)
+    )
+    logger.info(f"Models: {models}")
+    return models
+
+
+get_window_url_params = """
+function() {
+    const params = new URLSearchParams(window.location.search);
+    url_params = Object.fromEntries(params);
+    console.log(url_params);
+    return url_params;
+    }
+"""
+
+
+def 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
+
+
+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.update(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, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def clear_history(request: gr.Request):
+    logger.info(f"clear_history. ip: {request.client.host}")
+    state = default_conversation.copy()
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def add_text(state, text, image, image_process_mode, request: gr.Request):
+    logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
+    if len(text) <= 0 and image is None:
+        state.skip_next = True
+        return (state, state.to_gradio_chatbot(), "", None) + (no_change_btn,) * 5
+    if args.moderate:
+        flagged = violates_moderation(text)
+        if flagged:
+            state.skip_next = True
+            return (state, state.to_gradio_chatbot(), moderation_msg, None) + (no_change_btn,) * 5
+
+    text = text[:1536]  # Hard cut-off
+    if image is not None:
+        text = text[:1200]  # Hard cut-off for images
+        if "<image>" not in text:
+            # text = '<Image><image></Image>' + text
+            text = text + "\n<image>"
+        text = (text, image, image_process_mode)
+        if len(state.get_images(return_pil=True)) > 0:
+            state = default_conversation.copy()
+    state.append_message(state.roles[0], text)
+    state.append_message(state.roles[1], None)
+    state.skip_next = False
+    return (state, state.to_gradio_chatbot(), "", None) + (disable_btn,) * 5
+
+
+def http_bot(state, model_selector, interaction_mode, temperature, 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, state.to_gradio_chatbot()) + (no_change_btn,) * 5
+        return
+
+    if len(state.messages) == state.offset + 2:
+        # First round of conversation
+        # (Note): Hardcoding llava_v1 conv template for now
+        new_state = conv_templates["llava_v1"].copy()
+        new_state.append_message(new_state.roles[0], state.messages[-2][1])
+        new_state.append_message(new_state.roles[1], None)
+        state = new_state
+
+    # Query worker address
+    controller_url = args.controller_url
+    ret = requests.post(controller_url + "/get_worker_address", json={"model": model_name})
+    worker_addr = ret.json()["address"]
+    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+
+    # No available worker
+    if worker_addr == "":
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot(), disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+        return
+
+    # Construct prompt
+    prompt = state.get_prompt()
+
+    all_images = state.get_images(return_pil=True)
+    all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images]
+    for image, im_hash in zip(all_images, all_image_hash):
+        t = datetime.datetime.now()
+        filename = os.path.join(LOGDIR, "serve_images", f"{t.year}-{t.month:02d}-{t.day:02d}", f"{im_hash}.jpg")
+        if not os.path.isfile(filename):
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            image.save(filename)
+
+    # Make requests
+    pload = {
+        "model": model_name,
+        "prompt": prompt,
+        "interaction_mode": interaction_mode,
+        "temperature": float(temperature),
+        "max_new_tokens": int(max_new_tokens),
+        "images": f"List of {len(state.get_images())} images: {all_image_hash}",
+    }
+    logger.info(f"==== request ====\n{pload}")
+
+    pload["images"] = state.get_images()
+
+    state.messages[-1][-1] = "▌"
+    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+
+    try:
+        # Stream output
+        response = requests.post(
+            worker_addr + "/worker_generate_stream", headers=headers, json=pload, stream=True, timeout=10
+        )
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                if data["error_code"] == 0:
+                    output = data["text"][len(prompt) :].strip()
+                    state.messages[-1][-1] = output + "▌"
+                    yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
+                else:
+                    output = data["text"] + f" (error_code: {data['error_code']})"
+                    state.messages[-1][-1] = output
+                    yield (state, state.to_gradio_chatbot()) + (
+                        disable_btn,
+                        disable_btn,
+                        disable_btn,
+                        enable_btn,
+                        enable_btn,
+                    )
+                    return
+                time.sleep(0.03)
+    except requests.exceptions.RequestException:
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot()) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
+        return
+
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
+
+    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),
+            "state": state.dict(),
+            "images": all_image_hash,
+            "ip": request.client.host,
+        }
+        fout.write(json.dumps(data) + "\n")
+
+
+title_markdown = """
+# Prismatic VLMs: Investigating the Design Space of Visually-Conditioned Language Models
+[[Project Page](TODO)] [[Code](TODO)]
+[[Models](TODO)]
+| 📚 [[Paper](TODO)]
+"""
+
+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. This website
+is heavily inspired by the website released by [LLaVA](https://github.com/haotian-liu/LLaVA).
+"""
+
+
+learn_more_markdown = """
+### License
+The service is a research preview intended for non-commercial use only, subject to the model
+[License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA,
+[Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI,
+and [Privacy Practices]
+(https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb)
+of ShareGPT. Please contact us if you find any potential violation.
+"""
+
+block_css = """
+
+#buttons button {
+    min-width: min(120px,100%);
+}
+
+"""
+
+
+def build_demo(embed_mode):
+    textbox = gr.Textbox(show_label=False, placeholder="Enter text and press ENTER", container=False)
+
+    with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="stone")) as demo:
+        state = gr.State()
+
+        if not embed_mode:
+            gr.Markdown(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(
+                    ["Crop", "Resize", "Pad", "Default"],
+                    value="Default",
+                    label="Preprocess for non-square image",
+                    visible=False,
+                )
+
+                cur_dir = os.path.dirname(os.path.abspath(__file__))
+                gr.Examples(
+                    examples=[
+                        [f"{cur_dir}/examples/cows_in_pasture.png", "How many cows are in this image?"],
+                        [
+                            f"{cur_dir}/examples/monkey_knives.png",
+                            "What is happening in this image?",
+                        ],
+                    ],
+                    inputs=[imagebox, textbox],
+                )
+
+                with gr.Accordion("Parameters", open=False):
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.2,
+                        step=0.1,
+                        interactive=True,
+                        label="Temperature",
+                    )
+                    max_output_tokens = gr.Slider(
+                        minimum=0,
+                        maximum=4096,
+                        value=2048,
+                        step=64,
+                        interactive=True,
+                        label="Max output tokens",
+                    )
+
+                with gr.Accordion("Interaction Mode", open=False):
+                    interaction_modes = list(INTERACTION_MODES_MAP.keys())
+                    interaction_mode = gr.Dropdown(
+                        choices=interaction_modes,
+                        value=interaction_modes[0] if len(interaction_modes) > 0 else "Chat",
+                        interactive=True,
+                        show_label=False,
+                        container=False,
+                    )
+
+            with gr.Column(scale=8):
+                chatbot = gr.Chatbot(elem_id="chatbot", label="PrismaticVLMs Chatbot", height=550)
+                with gr.Row():
+                    with gr.Column(scale=8):
+                        textbox.render()
+                    with gr.Column(scale=1, min_width=50):
+                        submit_btn = gr.Button(value="Generate", variant="primary")
+                with gr.Row(elem_id="buttons"):
+                    # upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
+                    # downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
+                    # flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
+                    # stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
+                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
+                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
+
+        if not embed_mode:
+            gr.Markdown(tos_markdown)
+            gr.Markdown(learn_more_markdown)
+        url_params = gr.JSON(visible=False)
+
+        # Register listeners
+        btn_list = [regenerate_btn, clear_btn]
+        # upvote_btn.click(
+        #     upvote_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
+        # )
+        # downvote_btn.click(
+        #     downvote_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
+        # )
+        # flag_btn.click(
+        #     flag_last_response, [state, model_selector], [textbox, upvote_btn, downvote_btn, flag_btn], queue=False
+        # )
+
+        regenerate_btn.click(
+            regenerate, [state, image_process_mode], [state, chatbot, textbox, imagebox, *btn_list], queue=False
+        ).then(
+            http_bot,
+            [state, model_selector, interaction_mode, temperature, max_output_tokens],
+            [state, chatbot, *btn_list],
+        )
+
+        clear_btn.click(clear_history, None, [state, chatbot, textbox, imagebox, *btn_list], queue=False)
+
+        textbox.submit(
+            add_text,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox, *btn_list],
+            queue=False,
+        ).then(
+            http_bot,
+            [state, model_selector, interaction_mode, temperature, max_output_tokens],
+            [state, chatbot, *btn_list],
+        )
+
+        submit_btn.click(
+            add_text,
+            [state, textbox, imagebox, image_process_mode],
+            [state, chatbot, textbox, imagebox, *btn_list],
+            queue=False,
+        ).then(
+            http_bot,
+            [state, model_selector, interaction_mode, temperature, max_output_tokens],
+            [state, chatbot, *btn_list],
+        )
+
+        if args.model_list_mode == "once":
+            demo.load(load_demo, [url_params], [state, model_selector], _js=get_window_url_params, queue=False)
+        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=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
+    )