Spaces:
Sleeping
Sleeping
| """ | |
| A model worker executes the model. | |
| """ | |
| import argparse | |
| import json | |
| import torch | |
| from vcoder_llava.utils import server_error_msg | |
| from vcoder_llava.model.builder import load_pretrained_model | |
| from vcoder_llava.mm_utils import process_images, load_image_from_base64, tokenizer_seg_token, tokenizer_depth_seg_token, tokenizer_image_token, KeywordsStoppingCriteria | |
| from vcoder_llava.constants import ( | |
| IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, | |
| SEG_TOKEN_INDEX, DEFAULT_SEG_TOKEN, | |
| DEPTH_TOKEN_INDEX, DEFAULT_DEPTH_TOKEN, | |
| ) | |
| from transformers import TextIteratorStreamer | |
| from threading import Thread | |
| class Chat: | |
| def __init__(self, model_path, model_base, model_name, | |
| load_8bit, load_4bit, device, logger): | |
| 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} ...") | |
| self.tokenizer, self.model, self.image_processor, self.seg_image_processor, self.depth_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 = 'llava' in self.model_name.lower() | |
| self.is_seg = "vcoder" in self.model_name.lower() | |
| self.is_depth = "ds" in self.model_name.lower() | |
| def generate_stream(self, params): | |
| tokenizer, model, image_processor, seg_image_processor, depth_image_processor = self.tokenizer, self.model, self.image_processor, self.seg_image_processor, self.depth_image_processor | |
| prompt = params["prompt"] | |
| ori_prompt = prompt | |
| images = params.get("images", None) | |
| segs = params.get("segs", None) | |
| depths = params.get("depths", None) | |
| num_image_tokens = 0 | |
| num_seg_tokens = 0 | |
| num_depth_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_vision_tower().num_patches | |
| if segs is not None and len(segs) > 0 and self.is_seg: | |
| if len(segs) != prompt.count(DEFAULT_SEG_TOKEN): | |
| raise ValueError("Number of segs does not match number of <seg> tokens in prompt") | |
| segs = [load_image_from_base64(seg) for seg in segs] | |
| segs = process_images(segs, seg_image_processor, model.config) | |
| if type(segs) is list: | |
| segs = [seg.to(self.model.device, dtype=torch.float16) for seg in segs] | |
| else: | |
| segs = segs.to(self.model.device, dtype=torch.float16) | |
| replace_seg_token = DEFAULT_SEG_TOKEN | |
| prompt = prompt.replace(DEFAULT_SEG_TOKEN, replace_seg_token) | |
| num_seg_tokens = prompt.count(replace_seg_token) * model.get_vision_tower().num_patches | |
| if depths is not None and len(depths) > 0 and self.is_depth: | |
| if len(depths) != prompt.count(DEFAULT_DEPTH_TOKEN): | |
| raise ValueError("Number of depths does not match number of <depth> tokens in prompt") | |
| depths = [load_image_from_base64(depth) for depth in depths] | |
| depths = process_images(depths, depth_image_processor, model.config) | |
| if type(depths) is list: | |
| depths = [depth.to(self.model.device, dtype=torch.float16) for depth in depths] | |
| else: | |
| depths = depths.to(self.model.device, dtype=torch.float16) | |
| replace_depth_token = DEFAULT_DEPTH_TOKEN | |
| prompt = prompt.replace(DEFAULT_DEPTH_TOKEN, replace_depth_token) | |
| num_depth_tokens = prompt.count(replace_depth_token) * model.get_vision_tower().num_patches | |
| else: | |
| depths = None | |
| else: | |
| segs = None | |
| depths = None | |
| else: | |
| images = None | |
| segs = None | |
| depths = None | |
| image_args = {"images": images, "segs": segs, "depths": depths} | |
| else: | |
| images = None | |
| segs = None | |
| depths = 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', 2048) | |
| 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 | |
| if self.is_seg and segs is not None: | |
| if self.is_depth and depths is not None: | |
| input_ids = tokenizer_depth_seg_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, SEG_TOKEN_INDEX, DEPTH_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device) | |
| else: | |
| input_ids = tokenizer_seg_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, SEG_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device) | |
| else: | |
| 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 - num_seg_tokens - num_depth_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 | |
| generated_text = model.generate( | |
| 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 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() | |
| 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("--multi-modal", action="store_true", help="Multimodal mode is automatically detected with model name, please make sure `llava` is included in the model path.") | |
| parser.add_argument("--limit-model-concurrency", type=int, default=5) | |
| parser.add_argument("--stream-interval", type=int, default=1) | |
| parser.add_argument("--no-register", action="store_true") | |
| parser.add_argument("--load-8bit", action="store_true") | |
| parser.add_argument("--load-4bit", action="store_true") | |
| args = parser.parse_args() |