add diffusers 版本升级为0.5.0

Author: shell-nlp

gpt_server/model_worker/base/base_model_worker.py

@@ -189,3 +189,6 @@ def transcription(self, params):
 
     def generate_voice_stream(self, params):
         raise NotImplementedError
+
+    def get_image_output(self, params):
+        raise NotImplementedError

gpt_server/model_worker/base/model_worker_base.py

@@ -55,7 +55,9 @@ def __init__(
         multimodal: bool = False,
     ):
         is_vision = False
-        if model_type != "asr" and model_type != "tts":
+        if model_type in ["image"]:
+            pass
+        elif model_type not in ["asr", "tts"]:
             try:
                 self.model_config = AutoConfig.from_pretrained(
                     model_path, trust_remote_code=True
@@ -406,6 +408,16 @@ async def api_get_embeddings(request: Request):
     return JSONResponse(content=embedding)
 
 
+@app.post("/worker_get_image_output")
+async def api_get_embeddings(request: Request):
+    params = await request.json()
+    await acquire_worker_semaphore()
+    logger.debug(f"params {params}")
+    embedding = await worker.get_image_output(params)
+    release_worker_semaphore()
+    return JSONResponse(content=embedding)
+
+
 @app.post("/worker_get_classify")
 async def api_get_classify(request: Request):
     params = await request.json()

gpt_server/model_worker/flux.py

@@ -0,0 +1,71 @@
+import asyncio
+
+import io
+import os
+from typing import List
+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
+
+root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+class FluxWorker(ModelWorkerBase):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        conv_template: str = None,  # type: ignore
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+            model_type="image",
+        )
+        backend = os.environ["backend"]
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipe = FluxPipeline.from_pretrained(
+            model_path, torch_dtype=torch.bfloat16
+        ).to(self.device)
+
+        logger.warning(f"模型：{model_names[0]}")
+
+    async def get_image_output(self, params):
+        prompt = params["prompt"]
+        image = self.pipe(
+            prompt,
+            height=1024,
+            width=1024,
+            guidance_scale=3.5,
+            num_inference_steps=50,
+            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},
+            },
+        }
+        return result
+
+
+if __name__ == "__main__":
+    FluxWorker.run()

gpt_server/model_worker/utils.py

@@ -4,6 +4,15 @@
 import base64
 import io
 
+from PIL.Image import Image
+
+
+# 转换为Base64
+def pil_to_base64(pil_img: Image, format: str = "PNG"):
+    buffered = io.BytesIO()
+    pil_img.save(buffered, format=format)  # 明确指定PNG格式
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
 
 def extract_base64(data_url: str):
     """从Data URL中提取纯Base64数据"""

gpt_server/openai_api_protocol/custom_api_protocol.py

@@ -14,6 +14,12 @@
 from pydantic import Field, BaseModel
 
 
+class ImagesGenRequest(BaseModel):
+    prompt: str
+    model: str
+    output_format: str  # png, jpeg, or webp
+
+
 # copy from https://github.com/remsky/Kokoro-FastAPI/blob/master/api/src/routers/openai_compatible.py
 class OpenAISpeechRequest(BaseModel):
     model: str = Field(

gpt_server/serving/openai_api_server.py

@@ -713,7 +713,35 @@ async def generate_completion(payload: Dict[str, Any], worker_addr: str):
     ModerationsRequest,
     SpeechRequest,
     OpenAISpeechRequest,
+    ImagesGenRequest,
 )
+
+
+async def get_images_gen(payload: Dict[str, Any]):
+    model_name = payload["model"]
+    worker_addr = get_worker_address(model_name)
+
+    transcription = await fetch_remote(
+        worker_addr + "/worker_get_image_output", payload
+    )
+    return json.loads(transcription)
+
+
+@app.post("/v1/images/generations", dependencies=[Depends(check_api_key)])
+async def speech(request: ImagesGenRequest):
+    error_check_ret = check_model(request)
+    if error_check_ret is not None:
+        return error_check_ret
+
+    payload = {
+        "model": request.model,
+        "prompt": request.prompt,
+        "output_format": request.output_format,
+    }
+    result = await get_images_gen(payload=payload)
+    return result
+
+
 import edge_tts
 import uuid
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "gpt_server"
-version = "0.4.7"
+version = "0.5.0"
 description = "gpt_server是一个用于生产级部署LLMs或Embedding的开源框架。"
 readme = "README.md"
 license = { text = "Apache 2.0" }
@@ -28,6 +28,7 @@ dependencies = [
     "sglang[all]>=0.4.6.post5",
     "flashinfer-python",
     "flashtts>=0.1.7",
+    "diffusers>=0.33.1",
 ]
 
 [tool.uv]