diff --git a/src/blackbox/vlms.py b/src/blackbox/vlms.py index 9ac90d5..a5b3370 100644 --- a/src/blackbox/vlms.py +++ b/src/blackbox/vlms.py @@ -12,13 +12,13 @@ import requests import base64 import copy import ast +import random +from time import time import io from PIL import Image from lmdeploy.serve.openai.api_client import APIClient -import io -from PIL import Image -from lmdeploy.serve.openai.api_client import APIClient + def is_base64(value) -> bool: try: @@ -50,8 +50,8 @@ class VLMS(Blackbox): - ignore_eos (bool): indicator for ignoring eos - skip_special_tokens (bool): Whether or not to remove special tokens in the decoding. Default to be True.""" - self.url = vlm_config.url - + self.model_dict = vlm_config.urls + self.model_url = None self.temperature: float = 0.7 self.top_p:float = 1 self.max_tokens: (int |None) = 512 @@ -81,7 +81,7 @@ class VLMS(Blackbox): data = args[0] return isinstance(data, list) - def processing(self, prompt:str, images:str | bytes, settings: dict, model_name: Optional[str] = None, user_context: List[dict] = None) -> str: + def processing(self, prompt:str | None, images:str | bytes | None, settings: dict, model_name: Optional[str] = None, user_context: List[dict] = None) -> str: """ Args: prompt: a string query to the model. @@ -105,6 +105,9 @@ class VLMS(Blackbox): else: settings = {} + if not prompt: + prompt = '你是一个辅助机器人,请就此图做一个简短的概括性描述,包括图中的主体物品及状态,不超过50字。' if images else '你好' + # Transform the images into base64 format where openai format need. if images: if is_base64(images): # image as base64 str @@ -148,7 +151,11 @@ class VLMS(Blackbox): # 'content': '图片中主要展示了一只老虎,它正在绿色的草地上休息。草地上有很多可以让人坐下的地方,而且看起来相当茂盛。背景比较模糊,可能是因为老虎的影响,让整个图片的其他部分都变得不太清晰了。' # } # ] - api_client = APIClient(self.url) + + user_context = self.keep_last_k_images(user_context,k = 1) + if self.model_url is None: self.model_url = self._get_model_url(model_name) + + api_client = APIClient(self.model_url) # api_client = APIClient("http://10.6.80.91:23333") model_name = api_client.available_models[0] # Reformat input into openai format to request. @@ -198,10 +205,28 @@ class VLMS(Blackbox): print(item["choices"][0]["message"]['content']) responses += item["choices"][0]["message"]['content'] # total_token_usage += item['usage']['total_tokens'] # 'usage': {'prompt_tokens': *, 'total_tokens': *, 'completion_tokens': *} - user_context = messages + [{'role': 'assistant', 'content': responses}] + self.custom_print(user_context) return responses, user_context + def _get_model_url(self,model_name:str | None): + available_models = {} + for model, url in self.model_dict.items(): + try: + response = requests.get(url,timeout=3) + if response.status_code == 200: + available_models[model] = url + except Exception as e: + # print(e) + pass + if not available_models: print("There are no available running models and please check your endpoint urls.") + if model_name and model_name in available_models: + return available_models[model_name] + else: + model = random.choice(list(available_models.keys())) + print(f"No such model {model_name}, using {model} instead.") if model_name else print(f"Using random model {model}.") + return available_models[model] + def _into_openai_format(self, context:List[list]) -> List[dict]: """ Convert the data into openai format. @@ -255,7 +280,35 @@ class VLMS(Blackbox): return user_context + def keep_last_k_images(self, user_context: list, k:int=2): + count = 0 + result =[] + for item in user_context[::-1]: + if item['role'] == 'user' and len(item['content']) > 1: + for idx, info in enumerate(item['content']): + if info['type'] in ('image_url','image') and count >= k: + item['content'].pop(idx) + # item['content'].insert(idx, {'type': 'text', 'text': ''}) + elif info['type'] in ('image_url','image') and count < k: + count += 1 + else: + continue + result.append(item) + return result[::-1] + + def custom_print(self, user_context: list): + result = [] + for item in user_context: + if item['role'] == 'user': + for idx, info in enumerate(item['content']): + if info['type'] in ('image_url','image'): + item['content'].pop(idx) + item['content'].insert(idx, {'type': 'image', 'image': '####'}) + else: + continue + result.append(item) + print(result) async def fast_api_handler(self, request: Request) -> Response: ## TODO: add support for multiple images and support image in form-data format json_request = True @@ -291,8 +344,8 @@ class VLMS(Blackbox): if prompt is None: return JSONResponse(content={'error': "Question is required"}, status_code=status.HTTP_400_BAD_REQUEST) - if model_name is None or model_name.isspace(): - model_name = "Qwen-VL-Chat" + # if model_name is None or model_name.isspace(): + # model_name = "Qwen-VL-Chat" response, history = self.processing(prompt, img_data,settings, model_name,user_context=user_context) # jsonresp = str(JSONResponse(content={"response": self.processing(prompt, img_data, model_name)}).body, "utf-8") diff --git a/src/configuration.py b/src/configuration.py index 651cb95..0c6e0cb 100644 --- a/src/configuration.py +++ b/src/configuration.py @@ -153,4 +153,4 @@ class VLMConf(): @inject def __init__(self, config: Configuration) -> None: - self.url = config.get("vlms.url") + self.urls = config.get("vlms.urls")