add openai-api-image

Author: sangchengmeng

lightllm/server/api_http.py

@@ -19,15 +19,20 @@
 import asyncio
 import collections
 import time
+import json
 import uvloop
+import requests
+import base64
 import os
+from io import BytesIO
 import pickle
 from .build_prompt import build_prompt, init_tokenizer
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 import ujson as json
 from http import HTTPStatus
 import uuid
+from PIL import Image
 import multiprocessing as mp
 from typing import AsyncGenerator, Union
 from typing import Callable
@@ -40,6 +45,7 @@
 from .httpserver_for_pd_master.manager import HttpServerManagerForPDMaster
 from .api_lightllm import lightllm_get_score, lightllm_pd_generate_stream
 from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.server.embed_cache.utils import image2base64
 
 from .api_models import (
     ChatCompletionRequest,
@@ -230,6 +236,38 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
         return create_error_response(HTTPStatus.BAD_REQUEST, "The function call feature is not supported")
 
     created_time = int(time.time())
+    
+    multimodal_params_dict = {"images": []}
+    for message in request.messages:
+        if isinstance(message.content, list):
+            texts = []
+            for content in message.content:
+                if content.type == 'text' and content.text:
+                    texts.append(content.text)
+                elif content.type == 'image_url' and content.image_url is not None:
+                    img = content.image_url.url
+                    if img.startswith("http://") or img.startswith("https://"):
+                        response = requests.get(img, stream=True, timeout=2)
+                        data = image2base64(response.raw)
+                    elif img.startswith("file://"):
+                        data = image2base64(img[7:])
+                    elif img.startswith("data:image"):
+                        # "data:image/jpeg;base64,{base64_image}"
+                        data_str = img.split(";", 1)[1]
+                        if data_str.startswith("base64,"):
+                            data = data_str[7:]
+                        else :
+                            raise ValueError("Unrecognized image input.")
+                    else:
+                        raise ValueError("Unrecognized image input. Supports local path, http url, base64, and PIL.Image.")
+
+                    multimodal_params_dict["images"].append({
+                        "type": "base64",
+                        "data": data
+                    })
+                    
+            message.content = "\n".join(texts)
+    
     prompt = await build_prompt(request)
     sampling_params_dict = {
         "do_sample": request.do_sample,
@@ -249,7 +287,7 @@ async def chat_completions(request: ChatCompletionRequest, raw_request: Request)
     sampling_params.init(tokenizer=g_objs.httpserver_manager.tokenizer, **sampling_params_dict)
 
     sampling_params.verify()
-    multimodal_params = MultimodalParams(images=[])
+    multimodal_params = MultimodalParams(**multimodal_params_dict)
 
     results_generator = g_objs.httpserver_manager.generate(
         prompt, sampling_params, multimodal_params, request=raw_request

lightllm/server/api_models.py

@@ -5,10 +5,21 @@
 import uuid
 
 
+class ImageURL(BaseModel):
+    url: str
+
+class MessageContent(BaseModel):
+    type: str
+    text: Optional[str] = None
+    image_url: Optional[ImageURL] = None
+
+class Message(BaseModel):
+    role: str
+    content: Union[str, List[MessageContent]]
+
 class ChatCompletionRequest(BaseModel):
-    # The openai api native parameters
     model: str
-    messages: List[Dict[str, str]]
+    messages: List[Message]
     function_call: Optional[str] = "none"
     temperature: Optional[float] = 1
     top_p: Optional[float] = 1.0

lightllm/server/embed_cache/utils.py

@@ -1,8 +1,9 @@
+import base64
 import torch
 import numpy as np
 from io import BytesIO
 import multiprocessing.shared_memory as shm
-
+from PIL import Image
 
 def tensor2bytes(t):
     # t = t.cpu().numpy().tobytes()
@@ -12,6 +13,13 @@ def tensor2bytes(t):
     buf.seek(0)
     return buf.read()
 
+def image2base64(img_str: str):
+    image_obj = Image.open(img_str)
+    if image_obj.format is None:
+        raise ValueError("No image format found.")
+    buffer = BytesIO()
+    image_obj.save(buffer, format=image_obj.format)
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
 
 def bytes2tensor(b):
     # return torch.from_numpy(np.frombuffer(b, dtype=np.float16)).cuda()