增强图片生成，支持 base64 和 url

Author: shell-nlp

gpt_server/model_worker/flux.py

@@ -3,12 +3,14 @@
 import io
 import os
 from typing import List
+import uuid
 from loguru import logger
 import shortuuid
 from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
 from gpt_server.model_worker.utils import pil_to_base64
 import torch
 from diffusers import FluxPipeline
+from gpt_server.utils import STATIC_DIR
 
 root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
@@ -44,6 +46,7 @@ def __init__(
 
     async def get_image_output(self, params):
         prompt = params["prompt"]
+        response_format = params.get("response_format", "b64_json")
         image = self.pipe(
             prompt,
             height=1024,
@@ -53,17 +56,39 @@ async def get_image_output(self, params):
             max_sequence_length=512,
             generator=torch.Generator(self.device).manual_seed(0),
         ).images[0]
-        base64 = pil_to_base64(pil_img=image)
-        result = {
-            "created": shortuuid.random(),
-            "data": [{"b64_json": base64}],
-            "usage": {
-                "total_tokens": 0,
-                "input_tokens": 0,
-                "output_tokens": 0,
-                "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
-            },
-        }
+        result = {}
+        if response_format == "b64_json":
+            # Convert PIL image to base64
+            base64 = pil_to_base64(pil_img=image)
+            result = {
+                "created": shortuuid.random(),
+                "data": [{"b64_json": base64}],
+                "usage": {
+                    "total_tokens": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                },
+            }
+            return result
+        elif response_format == "url":
+            # 生成唯一文件名（避免冲突）
+            file_name = str(uuid.uuid4()) + ".png"
+            save_path = STATIC_DIR / file_name
+            image.save(save_path, format="PNG")
+            WORKER_PORT = os.environ["WORKER_PORT"]
+            WORKER_HOST = os.environ["WORKER_HOST"]
+            url = f"http://{WORKER_HOST}:{WORKER_PORT}/static/{file_name}"
+            result = {
+                "created": shortuuid.random(),
+                "data": [{"url": url}],
+                "usage": {
+                    "total_tokens": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                },
+            }
         return result
 
 

gpt_server/openai_api_protocol/custom_api_protocol.py

@@ -21,9 +21,13 @@ class ImagesGenRequest(BaseModel):
         default="png",
         description="png, jpeg, or webp",
     )
-    model_type: Literal["t2v", "t2i"] = Field(
-        default="t2i",
-        description="t2v: 文生视频 t2i: 文生图",
+    # model_type: Literal["t2v", "t2i"] = Field(
+    #     default="t2i",
+    #     description="t2v: 文生视频 t2i: 文生图",
+    # )
+    response_format: Literal["url", "b64_json"] = Field(
+        default="url",
+        description="生成图像时返回的格式。必须为“ur”或“b64_json”之一。URL仅在图像生成后60分钟内有效。",
     )
 
 

gpt_server/serving/openai_api_server.py

@@ -736,7 +736,7 @@ async def speech(request: ImagesGenRequest):
         "model": request.model,
         "prompt": request.prompt,
         "output_format": request.output_format,
-        "model_type": request.model_type,
+        "response_format": request.response_format,
     }
     result = await get_images_gen(payload=payload)
     return result

tests/test_image_gen.py

@@ -2,9 +2,14 @@
 from openai import OpenAI
 
 client = OpenAI(api_key="EMPTY", base_url="http://localhost:8082/v1")
-
-img = client.images.generate(model="flux", prompt="A red pig")
-
+# 两种响应方式
+## response_format = "url"    默认为 url
+img = client.images.generate(model="flux", prompt="A red pig", response_format="url")
+print(img.data[0])
+## response_format = "b64_json"
+img = client.images.generate(
+    model="flux", prompt="A red pig", response_format="b64_json"
+)
 image_bytes = base64.b64decode(img.data[0].b64_json)
 with open("output.png", "wb") as f:
     f.write(image_bytes)