Delete mplug_docowl/serve

mplug_docowl/serve/cli.py

@@ -1,120 +0,0 @@
-import argparse
-import torch
-
-from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from mplug_owl2.conversation import conv_templates, SeparatorStyle
-from mplug_owl2.model.builder import load_pretrained_model
-from mplug_owl2.mm_utils import process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-
-from PIL import Image
-
-import requests
-from PIL import Image
-from io import BytesIO
-from transformers import TextStreamer
-
-
-def disable_torch_init():
-    """
-    Disable the redundant torch default initialization to accelerate model creation.
-    """
-    import torch
-    setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
-    setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
-
-
-def load_image(image_file):
-    if image_file.startswith('http://') or image_file.startswith('https://'):
-        response = requests.get(image_file)
-        image = Image.open(BytesIO(response.content)).convert('RGB')
-    else:
-        image = Image.open(image_file).convert('RGB')
-    return image
-
-
-def main(args):
-    # Model
-    disable_torch_init()
-
-    model_name = get_model_name_from_path(args.model_path)
-    tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, args.load_8bit, args.load_4bit, device=args.device)
-
-    conv_mode = "mplug_owl2"
-
-    if args.conv_mode is not None and conv_mode != args.conv_mode:
-        print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode))
-    else:
-        args.conv_mode = conv_mode
-
-    conv = conv_templates[args.conv_mode].copy()
-    roles = conv.roles
-
-    image = load_image(args.image_file)
-    # Similar operation in model_worker.py
-    image_tensor = process_images([image], image_processor, args)
-    if type(image_tensor) is list:
-        image_tensor = [image.to(model.device, dtype=torch.float16) for image in image_tensor]
-    else:
-        image_tensor = image_tensor.to(model.device, dtype=torch.float16)
-
-    while True:
-        try:
-            inp = input(f"{roles[0]}: ")
-        except EOFError:
-            inp = ""
-        if not inp:
-            print("exit...")
-            break
-
-        print(f"{roles[1]}: ", end="")
-
-        if image is not None:
-            # first message
-            inp = DEFAULT_IMAGE_TOKEN + inp
-            conv.append_message(conv.roles[0], inp)
-            image = None
-        else:
-            # later messages
-            conv.append_message(conv.roles[0], inp)
-        conv.append_message(conv.roles[1], None)
-        prompt = conv.get_prompt()
-
-        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(model.device)
-        stop_str = conv.sep if conv.sep_style not in [SeparatorStyle.TWO, SeparatorStyle.TWO_NO_SYS] else conv.sep2
-        keywords = [stop_str]
-        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
-        streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-        with torch.inference_mode():
-            output_ids = model.generate(
-                input_ids,
-                images=image_tensor,
-                do_sample=True,
-                temperature=args.temperature,
-                max_new_tokens=args.max_new_tokens,
-                streamer=streamer,
-                use_cache=True,
-                stopping_criteria=[stopping_criteria])
-
-        outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
-        conv.messages[-1][-1] = outputs
-
-        if args.debug:
-            print("\n", {"prompt": prompt, "outputs": outputs}, "\n")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-path", type=str, default="facebook/opt-350m")
-    parser.add_argument("--model-base", type=str, default=None)
-    parser.add_argument("--image-file", type=str, required=True)
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--conv-mode", type=str, default=None)
-    parser.add_argument("--temperature", type=float, default=0.2)
-    parser.add_argument("--max-new-tokens", type=int, default=512)
-    parser.add_argument("--load-8bit", action="store_true")
-    parser.add_argument("--load-4bit", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--image-aspect-ratio", type=str, default='pad')
-    args = parser.parse_args()
-    main(args)

mplug_docowl/serve/controller.py

@@ -1,298 +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 mplug_owl2.constants import CONTROLLER_HEART_BEAT_EXPIRATION
-from mplug_owl2.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(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")

mplug_docowl/serve/gradio_web_server.py

@@ -1,460 +0,0 @@
-import argparse
-import datetime
-import json
-import os
-import time
-
-import gradio as gr
-import requests
-
-from mplug_owl2.conversation import (default_conversation, conv_templates,
-                                   SeparatorStyle)
-from mplug_owl2.constants import LOGDIR
-from mplug_owl2.utils import (build_logger, server_error_msg,
-    violates_moderation, moderation_msg)
-import hashlib
-
-
-logger = build_logger("gradio_web_server", "gradio_web_server.log")
-
-headers = {"User-Agent": "mPLUG-Owl2 Client"}
-
-no_change_btn = gr.Button.update()
-enable_btn = gr.Button.update(interactive=True)
-disable_btn = gr.Button.update(interactive=False)
-
-priority = {
-    "vicuna-13b": "aaaaaaa",
-    "koala-13b": "aaaaaab",
-}
-
-
-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.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 = text + '<|image|>'
-            text = '<|image|>' + text
-        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, temperature, 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, state.to_gradio_chatbot()) + (no_change_btn,) * 5
-        return
-
-    if len(state.messages) == state.offset + 2:
-        # First round of conversation
-        template_name = "mplug_owl2"
-        new_state = conv_templates[template_name].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, 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"{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,
-        "temperature": float(temperature),
-        "top_p": float(top_p),
-        "max_new_tokens": min(int(max_new_tokens), 1536),
-        "stop": state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2,
-        "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 as e:
-        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(start_tstamp, 4),
-            "state": state.dict(),
-            "images": all_image_hash,
-            "ip": request.client.host,
-        }
-        fout.write(json.dumps(data) + "\n")
-
-
-title_markdown = ("""
-<h1 align="center"><a href="https://github.com/X-PLUG/mPLUG-Owl"><img src="https://z1.ax1x.com/2023/11/03/piM1rGQ.md.png", alt="mPLUG-Owl" border="0" style="margin: 0 auto; height: 200px;" /></a> </h1>
-
-<h2 align="center"> mPLUG-Owl2: Revolutionizing Multi-modal Large Language Model with Modality Collaboration</h2>
-
-<h5 align="center"> If you like our project, please give us a star ✨ on Github for latest update.  </h2>
-
-<div align="center">
-    <div style="display:flex; gap: 0.25rem;" align="center">
-        <a href='https://github.com/X-PLUG/mPLUG-Owl'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
-        <a href="https://arxiv.org/abs/2304.14178"><img src="https://img.shields.io/badge/Arxiv-2304.14178-red"></a>
-        <a href='https://github.com/X-PLUG/mPLUG-Owl/stargazers'><img src='https://img.shields.io/github/stars/X-PLUG/mPLUG-Owl.svg?style=social'></a>
-    </div>
-</div>
-
-""")
-
-
-tos_markdown = ("""
-### Terms of use
-By using this service, users are required to agree to the following terms:
-The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. The service may collect user dialogue data for future research.
-Please click the "Flag" button if you get any inappropriate answer! We will collect those to keep improving our moderator.
-For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality.
-""")
-
-
-learn_more_markdown = ("""
-### License
-The service is a research preview intended for non-commercial use only, subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of LLaMA, [Terms of Use](https://openai.com/policies/terms-of-use) of the data generated by OpenAI, and [Privacy Practices](https://chrome.google.com/webstore/detail/sharegpt-share-your-chatg/daiacboceoaocpibfodeljbdfacokfjb) of ShareGPT. Please contact us if you find any potential violation.
-""")
-
-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(title="mPLUG-Owl2", theme=gr.themes.Default(), css=block_css) 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/extreme_ironing.jpg", "What is unusual about this image?"],
-                    [f"{cur_dir}/examples/Rebecca_(1939_poster)_Small.jpeg", "What is the name of the movie in the poster?"],
-                ], inputs=[imagebox, textbox])
-
-                with gr.Accordion("Parameters", open=True) as parameter_row:
-                    temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, interactive=True, label="Temperature",)
-                    top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P",)
-                    max_output_tokens = gr.Slider(minimum=0, maximum=1024, value=512, step=64, interactive=True, label="Max output tokens",)
-
-            with gr.Column(scale=8):
-                chatbot = gr.Chatbot(elem_id="Chatbot", label="mPLUG-Owl2 Chatbot", height=600)
-                with gr.Row():
-                    with gr.Column(scale=8):
-                        textbox.render()
-                    with gr.Column(scale=1, min_width=50):
-                        submit_btn = gr.Button(value="Send", variant="primary")
-                with gr.Row(elem_id="buttons") as button_row:
-                    upvote_btn = gr.Button(value="👍  Upvote", interactive=False)
-                    downvote_btn = gr.Button(value="👎  Downvote", interactive=False)
-                    flag_btn = gr.Button(value="⚠️  Flag", interactive=False)
-                    #stop_btn = gr.Button(value="⏹️  Stop Generation", interactive=False)
-                    regenerate_btn = gr.Button(value="🔄  Regenerate", interactive=False)
-                    clear_btn = gr.Button(value="🗑️  Clear", interactive=False)
-
-        if not embed_mode:
-            gr.Markdown(tos_markdown)
-            gr.Markdown(learn_more_markdown)
-        url_params = gr.JSON(visible=False)
-
-        # Register listeners
-        btn_list = [upvote_btn, downvote_btn, flag_btn, regenerate_btn, clear_btn]
-        upvote_btn.click(
-            upvote_last_response,
-            [state, model_selector],
-            [textbox, upvote_btn, downvote_btn, flag_btn],
-            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, temperature, top_p, 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, temperature, top_p, 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, temperature, top_p, 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=False
-    )

mplug_docowl/serve/model_worker.py

@@ -1,342 +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 mplug_docowl.utils import (build_logger, server_error_msg,
-    pretty_print_semaphore)
-
-from mplug_docowl.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN,WORKER_HEART_BEAT_INTERVAL
-from mplug_docowl.conversation import conv_templates, SeparatorStyle
-from mplug_docowl.model.builder import load_pretrained_model
-from mplug_docowl.mm_utils import load_image_from_base64, process_images, tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria
-from mplug_docowl.processor import DocProcessor
-
-
-from transformers import TextIteratorStreamer
-from threading import Thread
-
-
-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 DocOwlInfer():
-    def __init__(self, ckpt_path, anchors='grid_9', add_global_img=True, load_8bit=False, load_4bit=False):
-        model_name = get_model_name_from_path(ckpt_path)
-        ic(model_name)
-        self.tokenizer, self.model, _, _ = load_pretrained_model(ckpt_path, None, model_name, load_8bit=load_8bit, load_4bit=load_4bit, device="cuda")
-        self.doc_image_processor = DocProcessor(image_size=448, anchors=anchors, add_global_img=add_global_img, add_textual_crop_indicator=True)
-        self.streamer = TextStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    def inference(self, image, query):
-        image_tensor, patch_positions, text = self.doc_image_processor(images=image, query='<|image|>'+query)
-        image_tensor = image_tensor.to(self.model.device, dtype=torch.float16)
-        patch_positions = patch_positions.to(self.model.device)
-
-        # ic(image_tensor.shape, patch_positions.shape, text)
-
-        conv = conv_templates["mplug_owl2"].copy()
-        roles = conv.roles # ("USER", "ASSISTANT")
-
-        conv.append_message(conv.roles[0], text)
-        conv.append_message(conv.roles[1], None)
-        prompt = conv.get_prompt()
-
-        # ic(prompt)
-
-        input_ids = tokenizer_image_token(prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.model.device)
-        
-        # ic(input_ids)
-
-        stop_str = conv.sep2
-        keywords = [stop_str]
-        stopping_criteria = KeywordsStoppingCriteria(keywords, self.tokenizer, input_ids)
-
-        with torch.inference_mode():
-            output_ids = self.model.generate(
-                input_ids,
-                images=image_tensor,
-                patch_positions=patch_positions,
-                do_sample=False,
-                temperature=1.0,
-                max_new_tokens=512,
-                streamer=self.streamer,
-                use_cache=True,
-                stopping_criteria=[stopping_criteria])
-
-        outputs = self.tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
-
-        return outputs.replace('</s>', '')
-
-# TODO: adapt for docowl infer
-class ModelWorker:
-    def __init__(self, controller_addr, worker_addr,
-                 worker_id, no_register,
-                 model_path, model_base, model_name,
-                 resolution, anchors, add_global_img,
-                 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]
-        
-        self.model_name = get_model_name_from_path(ckpt_path)
-
-        self.device = device
-        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")
-
-        self.tokenizer, self.model, _, self.context_len = load_pretrained_model(
-            model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)
-        
-        self.resolution=resolution
-        self.token_num_each_img = (self.resolution/14)*(self.resolution/14)/self.model.get_model().vison2text.conv_patch
-        self.doc_image_processor = DocProcessor(image_size=resolution, anchors=anchors, add_global_img=add_global_img, add_textual_crop_indicator=True)
-       
-
-        self.is_multimodal = True
-
-        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 = self.tokenizer, self.model, self.image_processor
-
-        prompt = params["prompt"]
-        ori_prompt = prompt
-        images = params.get("images", None)
-        num_image_tokens = 0
-        if images is not None and len(images) > 0 and self.is_multimodal:
-            if len(images) > 0:
-                
-                """if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
-                    raise ValueError("Number of images does not match number of <|image|> tokens in prompt")
-
-                images = [load_image_from_base64(image) for image in images]
-                images = process_images(images, image_processor, model.config)
-
-                if type(images) is list:
-                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
-                else:
-                    images = images.to(self.model.device, dtype=torch.float16)"""
-                
-                # docowl only support 1 image, so only keep the last image
-                image = images[-1]
-                assert prompt.count(DEFAULT_IMAGE_TOKEN) == 1
-
-                image_tensor, patch_positions, prompt = self.doc_image_processor(images=image, query=prompt)
-                image_tensor = image_tensor.to(self.model.device, dtype=torch.float16)
-                patch_positions = patch_positions.to(self.model.device)
-
-                replace_token = DEFAULT_IMAGE_TOKEN
-                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
-                num_image_tokens = prompt.count(replace_token) * (self.token_num_each_img+1)
-            else:
-                images = None
-                patch_positions = None
-            image_args = {"images": images, "patch_positions":patch_positions}
-        else:
-            images = None
-            image_args = {}
-
-        temperature = float(params.get("temperature", 1.0))
-        top_p = float(params.get("top_p", 1.0))
-        max_context_length = getattr(model.config, 'max_position_embeddings', 4096)
-        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)
-        stop_str = params.get("stop", None)
-        do_sample = True if temperature > 0.001 else False
-
-        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
-        keywords = [stop_str]
-        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
-
-        max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
-
-        if max_new_tokens < 1:
-            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
-            return
-
-        thread = Thread(target=model.generate, kwargs=dict(
-            inputs=input_ids,
-            do_sample=do_sample,
-            temperature=temperature,
-            top_p=top_p,
-            max_new_tokens=max_new_tokens,
-            streamer=streamer,
-            stopping_criteria=[stopping_criteria],
-            use_cache=True,
-            **image_args
-        ))
-        thread.start()
-
-        generated_text = ori_prompt
-        for new_text in streamer:
-            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="facebook/opt-350m")
-    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("--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}")
-
-
-    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")

mplug_docowl/serve/model_worker_bak.py

@@ -1,278 +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 mplug_owl2.constants import WORKER_HEART_BEAT_INTERVAL
-from mplug_owl2.utils import (build_logger, server_error_msg,
-    pretty_print_semaphore)
-from mplug_owl2.model.builder import load_pretrained_model
-from mplug_owl2.mm_utils import process_images, load_image_from_base64, tokenizer_image_token, KeywordsStoppingCriteria
-from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
-from transformers import TextIteratorStreamer
-from threading import Thread
-
-
-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
-
-        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, model_base, self.model_name, load_8bit, load_4bit, device=self.device)
-        self.is_multimodal = True
-
-        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 = self.tokenizer, self.model, self.image_processor
-
-        prompt = params["prompt"]
-        ori_prompt = prompt
-        images = params.get("images", None)
-        num_image_tokens = 0
-        if images is not None and len(images) > 0 and self.is_multimodal:
-            if len(images) > 0:
-                if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):
-                    raise ValueError("Number of images does not match number of <|image|> tokens in prompt")
-
-                images = [load_image_from_base64(image) for image in images]
-                images = process_images(images, image_processor, model.config)
-
-                if type(images) is list:
-                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]
-                else:
-                    images = images.to(self.model.device, dtype=torch.float16)
-
-                replace_token = DEFAULT_IMAGE_TOKEN
-                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)
-
-                num_image_tokens = prompt.count(replace_token) * (model.get_model().visual_abstractor.config.num_learnable_queries + 1)
-            else:
-                images = None
-            image_args = {"images": images}
-        else:
-            images = None
-            image_args = {}
-
-        temperature = float(params.get("temperature", 1.0))
-        top_p = float(params.get("top_p", 1.0))
-        max_context_length = getattr(model.config, 'max_position_embeddings', 4096)
-        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)
-        stop_str = params.get("stop", None)
-        do_sample = True if temperature > 0.001 else False
-
-        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)
-        keywords = [stop_str]
-        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
-
-        max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)
-
-        if max_new_tokens < 1:
-            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"
-            return
-
-        thread = Thread(target=model.generate, kwargs=dict(
-            inputs=input_ids,
-            do_sample=do_sample,
-            temperature=temperature,
-            top_p=top_p,
-            max_new_tokens=max_new_tokens,
-            streamer=streamer,
-            stopping_criteria=[stopping_criteria],
-            use_cache=True,
-            **image_args
-        ))
-        thread.start()
-
-        generated_text = ori_prompt
-        for new_text in streamer:
-            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="facebook/opt-350m")
-    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("--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}")
-
-
-    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")

mplug_docowl/serve/register_workers.py

@@ -1,26 +0,0 @@
-"""
-Manually register workers.
-
-Usage:
-python3 -m fastchat.serve.register_worker --controller http://localhost:21001 --worker-name http://localhost:21002
-"""
-
-import argparse
-
-import requests
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--controller-address", type=str)
-    parser.add_argument("--worker-name", type=str)
-    parser.add_argument("--check-heart-beat", action="store_true")
-    args = parser.parse_args()
-
-    url = args.controller_address + "/register_worker"
-    data = {
-        "worker_name": args.worker_name,
-        "check_heart_beat": args.check_heart_beat,
-        "worker_status": None,
-    }
-    r = requests.post(url, json=data)
-    assert r.status_code == 200