MAGAer13/mPLUG-Owl2

1"""2A model worker executes the model.3"""4import argparse5import asyncio6import json7import time8import threading9import uuid10 11import requests12import torch13from functools import partial14 15from mplug_owl2.constants import WORKER_HEART_BEAT_INTERVAL16from mplug_owl2.utils import (build_logger, server_error_msg,17    pretty_print_semaphore)18from mplug_owl2.model.builder import load_pretrained_model19from mplug_owl2.mm_utils import process_images, load_image_from_base64, tokenizer_image_token, KeywordsStoppingCriteria20from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN21from transformers import TextIteratorStreamer22from threading import Thread23 24GB = 1 << 3025 26worker_id = str(uuid.uuid4())[:6]27logger = build_logger("model_worker", f"model_worker_{worker_id}.log")28 29class ModelWorker:30    def __init__(self, model_path, model_base, model_name, load_8bit, load_4bit, device):31        self.worker_id = worker_id32        if model_path.endswith("/"):33            model_path = model_path[:-1]34        if model_name is None:35            model_paths = model_path.split("/")36            if model_paths[-1].startswith('checkpoint-'):37                self.model_name = model_paths[-2] + "_" + model_paths[-1]38            else:39                self.model_name = model_paths[-1]40        else:41            self.model_name = model_name42 43        self.device = device44        logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...")45        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(46            model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device)47        self.is_multimodal = True48 49    @torch.inference_mode()50    def generate_stream(self, params):51        tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor52 53        prompt = params["prompt"]54        ori_prompt = prompt55        images = params.get("images", None)56        num_image_tokens = 057        if images is not None and len(images) > 0 and self.is_multimodal:58            if len(images) > 0:59                if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN):60                    raise ValueError("Number of images does not match number of <|image|> tokens in prompt")61 62                images = [load_image_from_base64(image) for image in images]63                images = process_images(images, image_processor, model.config)64 65                if type(images) is list:66                    images = [image.to(self.model.device, dtype=torch.float16) for image in images]67                else:68                    images = images.to(self.model.device, dtype=torch.float16)69 70                replace_token = DEFAULT_IMAGE_TOKEN71                prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token)72 73                num_image_tokens = prompt.count(replace_token) * (model.get_model().visual_abstractor.config.num_learnable_queries + 1)74            else:75                images = None76            image_args = {"images": images}77        else:78            images = None79            image_args = {}80 81        temperature = float(params.get("temperature", 1.0))82        top_p = float(params.get("top_p", 1.0))83        max_context_length = getattr(model.config, 'max_position_embeddings', 4096)84        max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024)85        stop_str = params.get("stop", None)86        do_sample = True if temperature > 0.001 else False87 88        input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device)89        keywords = [stop_str]90        stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)91        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)92 93        max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens)94 95        if max_new_tokens < 1:96            yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0"97            return98 99        thread = Thread(target=model.generate, kwargs=dict(100            inputs=input_ids,101            do_sample=do_sample,102            temperature=temperature,103            top_p=top_p,104            max_new_tokens=max_new_tokens,105            streamer=streamer,106            stopping_criteria=[stopping_criteria],107            use_cache=True,108            **image_args109        ))110        thread.start()111 112        generated_text = ori_prompt113        for new_text in streamer:114            generated_text += new_text115            if generated_text.endswith(stop_str):116                generated_text = generated_text[:-len(stop_str)]117            yield json.dumps({"text": generated_text, "error_code": 0}).encode()118 119    def generate_stream_gate(self, params):120        try:121            for x in self.generate_stream(params):122                yield x123        except ValueError as e:124            print("Caught ValueError:", e)125            ret = {126                "text": server_error_msg,127                "error_code": 1,128            }129            yield json.dumps(ret).encode() 130        except torch.cuda.CudaError as e:131            print("Caught torch.cuda.CudaError:", e)132            ret = {133                "text": server_error_msg,134                "error_code": 1,135            }136            yield json.dumps(ret).encode()137        except Exception as e:138            print("Caught Unknown Error", e)139            ret = {140                "text": server_error_msg,141                "error_code": 1,142            }143            yield json.dumps(ret).encode()