# 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()