add ERNIE-Image OpenAI image generation API endpoint

vllm_omni/entrypoints/openai/ernie_image_api_server.py

+# SPDX-License-Identifier: Apache-2.0
+"""Standalone OpenAI-compatible ERNIE-Image API server."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import base64
+import io
+import os
+import random
+import time
+from enum import Enum
+from http import HTTPStatus
+
+import numpy as np
+import torch
+import uvicorn
+from diffusers import ErnieImagePipeline
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field, field_validator
+
+DEFAULT_MODEL_PATH = "/root/private_data/chenyh/models/ERNIE-Image"
+os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
+
+
+class ResponseFormat(str, Enum):
+    B64_JSON = "b64_json"
+
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str = Field(..., description="Text description of image")
+    model: str | None = Field(default=None, description="Optional model id")
+    n: int = Field(default=1, ge=1, le=10)
+    size: str | None = Field(default=None, description="Image dimensions in WIDTHxHEIGHT format")
+    response_format: ResponseFormat = Field(default=ResponseFormat.B64_JSON)
+    user: str | None = Field(default=None)
+    negative_prompt: str | None = Field(default=None)
+    num_inference_steps: int | None = Field(default=None, ge=1, le=200)
+    guidance_scale: float | None = Field(default=None, ge=0.0, le=20.0)
+    seed: int | None = Field(default=None)
+
+    @field_validator("size")
+    @classmethod
+    def validate_size(cls, value: str | None) -> str | None:
+        if value is None:
+            return value
+        parse_size(value)
+        return value
+
+
+class ImageData(BaseModel):
+    b64_json: str | None = Field(default=None)
+    url: str | None = Field(default=None)
+    revised_prompt: str | None = Field(default=None)
+
+
+class ImageGenerationResponse(BaseModel):
+    created: int
+    data: list[ImageData]
+    output_format: str | None = None
+    size: str | None = None
+
+
+def parse_size(size_str: str) -> tuple[int, int]:
+    parts = size_str.split("x")
+    if len(parts) != 2:
+        raise ValueError("size must be in WIDTHxHEIGHT format, for example 1024x1024")
+    try:
+        width = int(parts[0])
+        height = int(parts[1])
+    except ValueError as exc:
+        raise ValueError("size width and height must be integers") from exc
+    if width <= 0 or height <= 0:
+        raise ValueError("size width and height must be positive")
+    return width, height
+
+
+def encode_image_base64(image) -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    buffer.seek(0)
+    return base64.b64encode(buffer.read()).decode("utf-8")
+
+
+class ErnieImageService:
+    def __init__(self, model_path: str, torch_dtype: str = "bfloat16") -> None:
+        dtype = getattr(torch, torch_dtype)
+        self.model_path = model_path
+        self.pipe = ErnieImagePipeline.from_pretrained(model_path, torch_dtype=dtype)
+        self.pipe = self.pipe.to("cuda")
+        self.pipe.transformer.eval()
+        self.pipe.vae.eval()
+        self.pipe.text_encoder.eval()
+        self.pipe.pe.eval()
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def _set_seed(seed: int) -> torch.Generator:
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.use_deterministic_algorithms(True)
+        torch.backends.cudnn.benchmark = False
+        return torch.Generator(device="cuda").manual_seed(seed)
+
+    async def generate(self, request: ImageGenerationRequest) -> ImageGenerationResponse:
+        width, height = (parse_size(request.size) if request.size else (1024, 1024))
+        steps = request.num_inference_steps if request.num_inference_steps is not None else 50
+        guidance_scale = request.guidance_scale if request.guidance_scale is not None else 5.0
+
+        image_data: list[ImageData] = []
+        async with self._lock:
+            for idx in range(request.n):
+                seed = random.randint(0, 2**31 - 1) if request.seed is None else request.seed + idx
+                generator = self._set_seed(seed)
+                with torch.inference_mode():
+                    output = self.pipe(
+                        prompt=request.prompt,
+                        negative_prompt=request.negative_prompt or None,
+                        height=height,
+                        width=width,
+                        num_inference_steps=steps,
+                        guidance_scale=guidance_scale,
+                        generator=generator,
+                    )
+
+                revised_prompts = getattr(output, "revised_prompts", None)
+                image_data.append(
+                    ImageData(
+                        b64_json=encode_image_base64(output.images[0]),
+                        revised_prompt=(revised_prompts[0] if revised_prompts else None),
+                    )
+                )
+
+        return ImageGenerationResponse(
+            created=int(time.time()),
+            data=image_data,
+            output_format="png",
+            size=f"{width}x{height}",
+        )
+
+
+def build_app(service: ErnieImageService) -> FastAPI:
+    app = FastAPI(title="ERNIE-Image OpenAI API", version="0.1.0")
+
+    @app.get("/healthz")
+    async def healthz() -> dict[str, str]:
+        return {"status": "healthy"}
+
+    @app.post("/v1/images/generations", response_model=ImageGenerationResponse)
+    async def generate_images(request: ImageGenerationRequest) -> ImageGenerationResponse:
+        if not request.prompt or not request.prompt.strip():
+            raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="prompt cannot be empty")
+        if request.response_format != ResponseFormat.B64_JSON:
+            raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail="Only b64_json is supported")
+        try:
+            return await service.generate(request)
+        except HTTPException:
+            raise
+        except ValueError as exc:
+            raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail=str(exc)) from exc
+        except Exception as exc:
+            raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR, detail=f"Image generation failed: {exc}") from exc
+
+    return app
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run ERNIE-Image OpenAI-compatible API server")
+    parser.add_argument("--host", default="0.0.0.0", help="Server host")
+    parser.add_argument("--port", type=int, default=8091, help="Server port")
+    parser.add_argument("--model-path", default=DEFAULT_MODEL_PATH, help="Path to ERNIE-Image model")
+    parser.add_argument("--torch-dtype", default="bfloat16", choices=["float16", "bfloat16", "float32"])
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    service = ErnieImageService(model_path=args.model_path, torch_dtype=args.torch_dtype)
+    app = build_app(service)
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()