diff --git a/src/blackbox/vlms.py b/src/blackbox/vlms.py index 6765260..b7d37e2 100644 --- a/src/blackbox/vlms.py +++ b/src/blackbox/vlms.py @@ -12,6 +12,7 @@ import requests import base64 import copy import ast +import json import io from PIL import Image @@ -93,7 +94,6 @@ class VLMS(Blackbox): response: a string history: a list """ - if settings: for k in settings: if k not in self.settings: @@ -260,14 +260,19 @@ class VLMS(Blackbox): ## TODO: add support for multiple images and support image in form-data format json_request = True try: - content_type = request.headers['content-type'] + content_type = request.headers.get('content-type', '') + print(content_type) if content_type == 'application/json': data = await request.json() - else: + elif 'multipart/form-data' in content_type: data = await request.form() - json_request = False + json_request = False + else: + body = await request.body() + data = json.loads(body.decode("utf-8")) + except Exception as e: - return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST) + return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST) model_name = data.get("model_name") prompt = data.get("prompt") @@ -283,8 +288,8 @@ class VLMS(Blackbox): else: return JSONResponse(content={"error": "context format error, should be in format of list or Openai_format"}, status_code=status.HTTP_400_BAD_REQUEST) - if json_request: - img_data = data.get("img_data") + if json_request or 'multipart/form-data' not in content_type: + img_data = data.get("img_data") else: img_data = await data.get("img_data").read() if settings: settings = ast.literal_eval(settings)