support openai interface (#1037)

Co-authored-by: yihuiwen <yihuiwen@sensetime.com>

Author: black-eleven

lightx2v/server/api/openai_images.py

@@ -0,0 +1,255 @@
+import asyncio
+import base64
+import re
+import time
+import uuid
+from pathlib import Path
+from typing import Literal, Optional
+
+from fastapi import APIRouter, File, Form, HTTPException, Request, UploadFile
+from loguru import logger
+from pydantic import BaseModel, Field
+
+from ..schema import ImageTaskRequest
+from ..task_manager import TaskStatus, task_manager
+from .deps import get_services
+
+router = APIRouter()
+
+_SIZE_PATTERN = re.compile(r"^\s*(\d+)\s*x\s*(\d+)\s*$", re.IGNORECASE)
+
+
+class OpenAIImageGenerationRequest(BaseModel):
+    prompt: str = Field(..., description="Text prompt")
+    model: Optional[str] = Field(default=None, description="Ignored for compatibility")
+    n: int = Field(default=1, description="Number of images, currently only supports 1")
+    size: Optional[str] = Field(default=None, description="Image size, e.g. 1024x1024")
+    response_format: Literal["b64_json"] = Field(default="b64_json")
+    user: Optional[str] = Field(default=None, description="Ignored for compatibility")
+    seed: Optional[int] = Field(default=None, description="Optional random seed")
+
+
+class OpenAIImageResponse(BaseModel):
+    created: int
+    data: list[dict[str, str]]
+
+
+def _write_file_sync(file_path: Path, content: bytes) -> None:
+    with open(file_path, "wb") as buffer:
+        buffer.write(content)
+
+
+def _shape_from_size(size: str) -> tuple[int, int]:
+    match = _SIZE_PATTERN.match(size)
+    if not match:
+        raise ValueError("size must be in WxH format, e.g. 1024x1024")
+    width = int(match.group(1))
+    height = int(match.group(2))
+    if width <= 0 or height <= 0:
+        raise ValueError("size width and height must be positive")
+    return width, height
+
+
+async def _wait_task_result_png(task_id: str, timeout_seconds: int, poll_interval_seconds: float) -> bytes:
+    start_time = time.monotonic()
+    while True:
+        task_status = task_manager.get_task_status(task_id)
+        if not task_status:
+            raise HTTPException(status_code=500, detail=f"Task status not found: {task_id}")
+
+        status = task_status.get("status")
+        if status == TaskStatus.COMPLETED.value:
+            result_png = task_manager.get_task_result_png(task_id)
+            if result_png:
+                return result_png
+            raise HTTPException(status_code=500, detail=f"Task completed but no in-memory image found: {task_id}")
+
+        if status == TaskStatus.FAILED.value:
+            raise HTTPException(status_code=500, detail=task_status.get("error", "Task failed"))
+
+        if status == TaskStatus.CANCELLED.value:
+            raise HTTPException(status_code=409, detail=task_status.get("error", "Task cancelled"))
+
+        if (time.monotonic() - start_time) > timeout_seconds:
+            task_manager.cancel_task(task_id)
+            raise HTTPException(status_code=504, detail=f"Task {task_id} timed out after {timeout_seconds} seconds")
+
+        await asyncio.sleep(poll_interval_seconds)
+
+
+async def _watch_client_disconnect(request: Request, task_id: str, poll_interval_seconds: float = 0.2) -> bool:
+    while True:
+        if await request.is_disconnected():
+            task_manager.cancel_task(task_id)
+            logger.info(f"Client disconnected, task {task_id} cancelled")
+            return True
+        await asyncio.sleep(poll_interval_seconds)
+
+
+async def _run_sync_image_task(request: Request, message: ImageTaskRequest) -> bytes:
+    task_id = None
+    timeout_seconds = 600
+    poll_interval_seconds = 0.5
+
+    try:
+        message.prefer_memory_result = True
+        task_id = task_manager.create_task(message)
+        message.task_id = task_id
+
+        wait_task = asyncio.create_task(_wait_task_result_png(task_id, timeout_seconds, poll_interval_seconds))
+        disconnect_task = asyncio.create_task(_watch_client_disconnect(request, task_id))
+
+        done, pending = await asyncio.wait({wait_task, disconnect_task}, return_when=asyncio.FIRST_COMPLETED)
+        for pending_task in pending:
+            pending_task.cancel()
+        await asyncio.gather(*pending, return_exceptions=True)
+
+        if disconnect_task in done and disconnect_task.result():
+            if not wait_task.done():
+                wait_task.cancel()
+                await asyncio.gather(wait_task, return_exceptions=True)
+            raise HTTPException(status_code=499, detail=f"Client disconnected, task {task_id} cancelled")
+
+        return wait_task.result()
+    except RuntimeError as e:
+        raise HTTPException(status_code=503, detail=str(e))
+    except HTTPException:
+        raise
+    except asyncio.CancelledError:
+        if task_id:
+            task_manager.cancel_task(task_id)
+        raise
+    except Exception as e:
+        logger.error(f"Failed to run OpenAI-compatible image task: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+def _build_url_response(request: Request, task_id: str, image_bytes: bytes) -> str:
+    services = get_services()
+    assert services.file_service is not None, "File service is not initialized"
+
+    file_name = f"{task_id}.png"
+    output_path = services.file_service.output_video_dir / file_name
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    _write_file_sync(output_path, image_bytes)
+
+    base = str(request.base_url).rstrip("/")
+    return f"{base}/v1/files/download/{file_name}"
+
+
+def _build_openai_response(request: Request, task_id: str, image_bytes: bytes, response_format: Literal["url", "b64_json"]):
+    if response_format == "b64_json":
+        return OpenAIImageResponse(created=int(time.time()), data=[{"b64_json": base64.b64encode(image_bytes).decode("utf-8")}])
+
+    return OpenAIImageResponse(created=int(time.time()), data=[{"url": _build_url_response(request, task_id, image_bytes)}])
+
+
+def _build_image_task_request(
+    prompt: str,
+    *,
+    negative_prompt: str = "",
+    seed: Optional[int] = None,
+    target_shape: Optional[list[int]] = None,
+    image_path: str = "",
+    image_mask_path: str = "",
+) -> ImageTaskRequest:
+    payload = {
+        "prompt": prompt,
+        "negative_prompt": negative_prompt,
+        "image_path": image_path,
+        "image_mask_path": image_mask_path,
+    }
+    if target_shape:
+        payload["target_shape"] = target_shape
+    if seed is not None:
+        payload["seed"] = seed
+    return ImageTaskRequest(**payload)
+
+
+@router.post("/generations", response_model=OpenAIImageResponse)
+async def create_openai_image_generation(request: Request, body: OpenAIImageGenerationRequest):
+    if body.n != 1:
+        raise HTTPException(status_code=400, detail="Only n=1 is currently supported")
+    if not body.prompt.strip():
+        raise HTTPException(status_code=400, detail="prompt is required")
+
+    target_shape = None
+    if body.size:
+        try:
+            width, height = _shape_from_size(body.size)
+            target_shape = [height, width]
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    message = _build_image_task_request(
+        prompt=body.prompt,
+        seed=body.seed,
+        target_shape=target_shape,
+    )
+
+    result_png = await _run_sync_image_task(request, message)
+    return _build_openai_response(request, message.task_id, result_png, body.response_format)
+
+
+async def _save_upload_file(file: UploadFile, target_dir: Path) -> str:
+    if not file.filename:
+        raise HTTPException(status_code=400, detail="Uploaded file has no filename")
+
+    file_extension = Path(file.filename).suffix or ".png"
+    unique_filename = f"{uuid.uuid4()}{file_extension}"
+    file_path = target_dir / unique_filename
+
+    content = await file.read()
+    if not content:
+        raise HTTPException(status_code=400, detail=f"Uploaded file is empty: {file.filename}")
+    await asyncio.to_thread(_write_file_sync, file_path, content)
+    return str(file_path)
+
+
+@router.post("/edits", response_model=OpenAIImageResponse)
+async def create_openai_image_edit(
+    request: Request,
+    image: UploadFile = File(...),
+    prompt: str = Form(...),
+    mask: UploadFile | None = File(default=None),
+    model: str | None = Form(default=None),
+    n: int = Form(default=1),
+    size: str | None = Form(default=None),
+    response_format: Literal["url", "b64_json"] = Form(default="url"),
+    user: str | None = Form(default=None),
+    negative_prompt: str = Form(default=""),
+    seed: int | None = Form(default=None),
+):
+    _ = model, user
+    if n != 1:
+        raise HTTPException(status_code=400, detail="Only n=1 is currently supported")
+    if not prompt.strip():
+        raise HTTPException(status_code=400, detail="prompt is required")
+
+    services = get_services()
+    assert services.file_service is not None, "File service is not initialized"
+
+    target_shape = None
+    if size:
+        try:
+            width, height = _shape_from_size(size)
+            target_shape = [height, width]
+        except ValueError as e:
+            raise HTTPException(status_code=400, detail=str(e))
+
+    image_path = await _save_upload_file(image, services.file_service.input_image_dir)
+    image_mask_path = ""
+    if mask is not None:
+        image_mask_path = await _save_upload_file(mask, services.file_service.input_image_dir)
+
+    message = _build_image_task_request(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        seed=seed,
+        target_shape=target_shape,
+        image_path=image_path,
+        image_mask_path=image_mask_path,
+    )
+
+    result_png = await _run_sync_image_task(request, message)
+    return _build_openai_response(request, message.task_id, result_png, response_format)

lightx2v/server/api/router.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter
 
 from .files import router as files_router
+from .openai_images import router as openai_images_router
 from .service_routes import router as service_router
 from .tasks import common_router, image_router, video_router
 
@@ -19,6 +20,7 @@ def create_api_router() -> APIRouter:
     tasks_router.post("/", response_model_exclude_unset=True, deprecated=True)(create_video_task)
 
     api_router.include_router(tasks_router)
+    api_router.include_router(openai_images_router, prefix="/v1/images", tags=["openai-images"])
     api_router.include_router(files_router, prefix="/v1/files", tags=["files"])
     api_router.include_router(service_router, prefix="/v1/service", tags=["service"])
 

scripts/server/test_openai_images_client.py

@@ -0,0 +1,115 @@
+import argparse
+import base64
+from pathlib import Path
+from typing import Any
+
+import requests
+
+try:
+    from openai import OpenAI  # pyright: ignore[reportMissingImports]
+except ImportError:
+    OpenAI = None  # type: ignore[assignment]
+
+
+def _extract_data_item(response: Any) -> dict[str, Any]:
+    if not hasattr(response, "data") or not response.data:
+        raise RuntimeError(f"Invalid OpenAI images response: {response}")
+    item = response.data[0]
+    if hasattr(item, "model_dump"):
+        return item.model_dump()  # openai pydantic object
+    if isinstance(item, dict):
+        return item
+    raise RuntimeError(f"Unsupported data item type: {type(item)!r}")
+
+
+def _save_image_from_item(item: dict[str, Any], output_path: Path) -> Path:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    if "b64_json" in item and item["b64_json"]:
+        image_bytes = base64.b64decode(item["b64_json"])
+        output_path.write_bytes(image_bytes)
+        return output_path
+
+    if "url" in item and item["url"]:
+        resp = requests.get(item["url"], timeout=120)
+        resp.raise_for_status()
+        output_path.write_bytes(resp.content)
+        return output_path
+
+    raise RuntimeError(f"Response item has neither b64_json nor url: {item}")
+
+
+def run_generate(client: Any, args: argparse.Namespace) -> Path:
+    response = client.images.generate(
+        model=args.model,
+        prompt=args.prompt,
+        size=args.size,
+        response_format=args.response_format,
+    )
+    item = _extract_data_item(response)
+    print(f"[generate] response item: {item}")
+    return _save_image_from_item(item, Path(args.output_dir) / "generate.png")
+
+
+def run_edit(client: Any, args: argparse.Namespace) -> Path:
+    if not args.image:
+        raise ValueError("--image is required for edit mode")
+
+    image_path = Path(args.image)
+    if not image_path.exists():
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+
+    with image_path.open("rb") as image_file:
+        kwargs = {
+            "model": args.model,
+            "image": image_file,
+            "prompt": args.edit_prompt or args.prompt,
+            "size": args.size,
+            "response_format": args.response_format,
+        }
+        if args.mask:
+            mask_path = Path(args.mask)
+            if not mask_path.exists():
+                raise FileNotFoundError(f"Mask file not found: {mask_path}")
+            with mask_path.open("rb") as mask_file:
+                response = client.images.edit(mask=mask_file, **kwargs)
+        else:
+            response = client.images.edit(**kwargs)
+
+    item = _extract_data_item(response)
+    print(f"[edit] response item: {item}")
+    return _save_image_from_item(item, Path(args.output_dir) / "edit.png")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Test OpenAI-compatible image APIs on LightX2V server.")
+    parser.add_argument("--base_url", type=str, default="http://127.0.0.1:8000/v1", help="OpenAI-compatible base URL")
+    parser.add_argument("--api_key", type=str, default="dummy-key", help="OpenAI API key placeholder")
+    parser.add_argument("--model", type=str, default="gpt-image-1", help="Model name (for compatibility only)")
+    parser.add_argument("--mode", choices=["generate", "edit", "all"], default="all", help="Test mode")
+    parser.add_argument("--prompt", type=str, default="a futuristic city at sunset", help="Prompt for generation")
+    parser.add_argument("--edit_prompt", type=str, default="", help="Prompt for edit (defaults to --prompt)")
+    parser.add_argument("--size", type=str, default="1024x1024", help="Image size, e.g. 1024x1024")
+    parser.add_argument("--response_format", choices=["url", "b64_json"], default="url", help="OpenAI response format")
+    parser.add_argument("--image", type=str, default="", help="Input image path for edit mode")
+    parser.add_argument("--mask", type=str, default="", help="Optional mask image path for edit mode")
+    parser.add_argument("--output_dir", type=str, default="outputs/openai_images_test", help="Directory to save outputs")
+    args = parser.parse_args()
+
+    if OpenAI is None:
+        raise RuntimeError("Missing dependency: openai. Please install it with `pip install openai`.")
+
+    client = OpenAI(api_key=args.api_key, base_url=args.base_url)
+
+    output_paths: list[Path] = []
+    if args.mode in ("generate", "all"):
+        output_paths.append(run_generate(client, args))
+    if args.mode in ("edit", "all"):
+        output_paths.append(run_edit(client, args))
+
+    for path in output_paths:
+        print(f"[saved] {path}")
+
+
+if __name__ == "__main__":
+    main()