""" A model worker executes the model. """ import argparse import asyncio import json import time import threading import uuid import requests import torch 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") class ModelWorker: def __init__(self, model_path, model_base, model_name, load_8bit, load_4bit, device): 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 @torch.inference_mode() def predict_stream(self, params): tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor prompt = params["prompt"] + "The quality of the image is" 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 = {} input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device) logits = model.forward( input_ids=input_ids, use_cache=True, **image_args).logits[0,-1] print(logits.shape) softmax_logits = torch.softmax(logits[[1781,6588,6460]], 0) print(tokenizer(["good", "average", "poor"])) fake_streamer = [] for id_, word in enumerate(["good", "average", "poor"]): stream_ = f"Probability of {word} quality: {softmax_logits[id_].item():.4f};\n" fake_streamer.append(stream_) quality_score = 0.5 * softmax_logits[1] + softmax_logits[0] stream_ = f"Quality score: {quality_score:.4f} (range [0,1])." fake_streamer.append(stream_) generated_text = ori_prompt.replace("The quality of the image is", "") for new_text in fake_streamer: generated_text += new_text yield json.dumps({"text": generated_text, "error_code": 0}).encode() @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() def predict_stream_gate(self, params): try: for x in self.predict_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() 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() except Exception as e: print("Caught Unknown Error", e) ret = { "text": server_error_msg, "error_code": 1, } yield json.dumps(ret).encode() 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() 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() except Exception as e: print("Caught Unknown Error", e) ret = { "text": server_error_msg, "error_code": 1, } yield json.dumps(ret).encode()