Studio: add OpenAI-compatible /v1/images/generations endpoint (#6686)
* Studio: add OpenAI-compatible /v1/images/generations endpoint * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to the loaded base repo for image-generation defaults * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restrict the image-generation 503 to genuine unload races * Route /v1/images/generations through the active diffusion engine --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
This commit is contained in:
oobabooga
GitHub
parent
ca782d102e
commit
8c5a00e0ca
4 changed files with 650 additions and 0 deletions
|
|
@ -209,6 +209,38 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
|
|||
return base or fam.base_repo
|
||||
|
||||
|
||||
# Default (steps, guidance) per model for callers that can't pass them — namely
|
||||
# the OpenAI /v1/images/generations endpoint, whose spec has no step/guidance
|
||||
# knobs. Distilled "turbo/schnell" models want few steps and no CFG; the full
|
||||
# "dev" models want more steps and real CFG. Matched by substring, most specific
|
||||
# first — the same scheme and values as the UI's MODEL_DEFAULTS table
|
||||
# (studio/frontend/src/features/images/images-page.tsx); keep the two in sync.
|
||||
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
|
||||
("z-image-turbo", 9, 0.0),
|
||||
("flux.1-schnell", 4, 0.0),
|
||||
("flux.1", 28, 3.5),
|
||||
("flux.2-klein", 4, 0.0),
|
||||
("qwen-image", 20, 4.0),
|
||||
("z-image", 20, 4.0),
|
||||
)
|
||||
# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
|
||||
_GENERATION_DEFAULT_FALLBACK = (9, 0.0)
|
||||
|
||||
|
||||
def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]:
|
||||
"""Default ``(steps, guidance)`` for a loaded model. The first identifier that
|
||||
names a known model wins (the repo id, then the resolved base repo), so a
|
||||
local-path load — whose repo id is just a filesystem path that may not name
|
||||
the model — still resolves via its base repo. Within an identifier, keys are
|
||||
matched as substrings, most specific first (the same scheme as the UI)."""
|
||||
for identifier in identifiers:
|
||||
needle = (identifier or "").lower()
|
||||
for key, steps, guidance in _GENERATION_DEFAULTS:
|
||||
if key in needle:
|
||||
return steps, guidance
|
||||
return _GENERATION_DEFAULT_FALLBACK
|
||||
|
||||
|
||||
def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]:
|
||||
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None."""
|
||||
for entry_scheme, repo_id in fam.prequant_repos:
|
||||
|
|
|
|||
|
|
@ -1928,3 +1928,68 @@ class DiffusionStatusResponse(BaseModel):
|
|||
None,
|
||||
description = "Why diffusers was chosen over the native sd.cpp engine (null when none)",
|
||||
)
|
||||
|
||||
|
||||
# ── OpenAI-compatible images API (POST /v1/images/generations) ──
|
||||
#
|
||||
# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf
|
||||
# OpenAI clients work unchanged. The loaded image GGUF stands in for the model;
|
||||
# GPT-image-only knobs (quality, style, background, output_format, ...) are
|
||||
# accepted and ignored, exactly as dall-e-2 ignores them. The size string is
|
||||
# parsed and `stream` is rejected in the route, where the diffusion backend is
|
||||
# in reach; everything Pydantic can check declaratively lives here.
|
||||
|
||||
|
||||
class ImageGenerationRequest(BaseModel):
|
||||
"""OpenAI ``CreateImageRequest`` for ``POST /v1/images/generations``.
|
||||
|
||||
``prompt`` is the only required field, per the spec. Unlisted OpenAI fields
|
||||
are ignored (Pydantic's default), matching dall-e-2's treatment of the
|
||||
GPT-image-only parameters."""
|
||||
|
||||
prompt: str = Field(..., min_length = 1, description = "Text description of the image(s).")
|
||||
model: Optional[str] = Field(
|
||||
None, description = "Model id (informational; the loaded image model is used)."
|
||||
)
|
||||
n: int = Field(1, ge = 1, le = 10, description = "Number of images to generate (1-10).")
|
||||
size: str = Field(
|
||||
"auto", description = "'auto' or '<width>x<height>' (256-2048, each a multiple of 16)."
|
||||
)
|
||||
response_format: Literal["url", "b64_json"] = Field(
|
||||
"url", description = "Return each image as a URL or a base64-encoded PNG."
|
||||
)
|
||||
user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).")
|
||||
# gpt-image-only; declared so we can reject it with a clear error instead of
|
||||
# silently returning JSON to a client that asked for an SSE stream.
|
||||
stream: Optional[bool] = Field(
|
||||
None, description = "Streaming image generation is not supported; omit or set false."
|
||||
)
|
||||
|
||||
@field_validator("n", "size", "response_format", mode = "before")
|
||||
@classmethod
|
||||
def _null_means_default(cls, value, info):
|
||||
# OpenAI marks these nullable WITH a default, so an explicit null means
|
||||
# "use the default" — coalesce it instead of 400-ing a spec-valid body.
|
||||
if value is None:
|
||||
return cls.model_fields[info.field_name].default
|
||||
return value
|
||||
|
||||
|
||||
class ImageGenerationData(BaseModel):
|
||||
"""One image in an ``ImagesResponse`` (OpenAI ``Image``). Exactly one of
|
||||
``url`` / ``b64_json`` is set, per the request's ``response_format``; the
|
||||
route serializes with ``exclude_none`` so the unused key is omitted."""
|
||||
|
||||
b64_json: Optional[str] = Field(
|
||||
None, description = "Base64-encoded PNG (response_format=b64_json)."
|
||||
)
|
||||
url: Optional[str] = Field(None, description = "URL to the PNG bytes (response_format=url).")
|
||||
|
||||
|
||||
class ImageGenerationResponse(BaseModel):
|
||||
"""OpenAI ``ImagesResponse``. dall-e-shaped: the GPT-image-only top-level
|
||||
fields (background/output_format/size/quality/usage) are omitted, since our
|
||||
sizes wouldn't satisfy their fixed enums and we report no token usage."""
|
||||
|
||||
created: int = Field(..., description = "Unix timestamp (seconds) the images were created.")
|
||||
data: list[ImageGenerationData] = Field(..., description = "The generated images.")
|
||||
|
|
|
|||
|
|
@ -1062,6 +1062,9 @@ from models.inference import (
|
|||
DiffusionLoadProgressResponse,
|
||||
GalleryImage,
|
||||
GalleryListResponse,
|
||||
ImageGenerationRequest,
|
||||
ImageGenerationData,
|
||||
ImageGenerationResponse,
|
||||
LoadResponse,
|
||||
LoadProgressResponse,
|
||||
UnloadResponse,
|
||||
|
|
@ -11301,3 +11304,167 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub
|
|||
async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)):
|
||||
from core.inference.diffusion_engine_router import get_active_diffusion_engine
|
||||
return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress())
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
# OpenAI-compatible images API (POST /v1/images/generations)
|
||||
#
|
||||
# The inference router is mounted at both /api/inference and /v1, so this also
|
||||
# answers /v1/images/generations for off-the-shelf OpenAI clients. It maps
|
||||
# OpenAI's CreateImageRequest onto the in-process diffusion backend and returns
|
||||
# an ImagesResponse. Studio's own Image tab uses the richer /images/generate
|
||||
# route above; this is the spec-shaped surface, and the single error boundary
|
||||
# mapping backend exceptions to OpenAI error envelopes (the global /v1 handler
|
||||
# wraps HTTPException detail into the envelope).
|
||||
# ──────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample
|
||||
# x 2x patch); the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all
|
||||
# satisfy this. Mirrors DiffusionGenerateRequest's width/height bounds so both
|
||||
# generate paths accept the same geometry.
|
||||
_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$")
|
||||
_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048
|
||||
# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both
|
||||
# "no image model" responses stay identical.
|
||||
_NO_IMAGE_MODEL_MSG = "No image model loaded. Load an image model first."
|
||||
|
||||
|
||||
def _parse_openai_image_size(size: str) -> tuple[int, int]:
|
||||
"""OpenAI ``size`` -> (width, height). ``auto``/empty -> 1024x1024 (~1MP, what
|
||||
these models target). Raises ValueError with a client-facing message."""
|
||||
text = (size or "").strip().lower()
|
||||
if text in ("", "auto"):
|
||||
return 1024, 1024
|
||||
match = _IMAGE_SIZE_RE.match(text)
|
||||
if not match:
|
||||
raise ValueError("size must be 'auto' or '<width>x<height>', e.g. '1024x1024'.")
|
||||
width, height = int(match.group(1)), int(match.group(2))
|
||||
for label, value in (("width", width), ("height", height)):
|
||||
if not _IMAGE_DIM_MIN <= value <= _IMAGE_DIM_MAX:
|
||||
raise ValueError(f"size {label} must be between {_IMAGE_DIM_MIN} and {_IMAGE_DIM_MAX}.")
|
||||
if value % 16 != 0:
|
||||
raise ValueError(f"size {label} must be a multiple of 16.")
|
||||
return width, height
|
||||
|
||||
|
||||
def _absolute_image_url(request: Request, relative: str) -> str:
|
||||
"""Join a relative gallery path onto the request's own scheme+host, for the
|
||||
response_format=url links. Like every Studio route, the target needs the
|
||||
bearer token; b64_json avoids that for clients that can't carry it."""
|
||||
return str(request.base_url).rstrip("/") + relative
|
||||
|
||||
|
||||
@router.post(
|
||||
"/images/generations",
|
||||
response_model = ImageGenerationResponse,
|
||||
response_model_exclude_none = True,
|
||||
)
|
||||
async def openai_image_generations(
|
||||
body: ImageGenerationRequest,
|
||||
request: Request,
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""OpenAI-compatible text-to-image (POST /v1/images/generations).
|
||||
|
||||
Generates ``n`` images from ``prompt`` on the loaded diffusion model and
|
||||
returns them as URLs (default) or base64 PNGs per ``response_format``. Steps
|
||||
and guidance have no OpenAI knob, so they default per loaded model."""
|
||||
from core.inference import image_gallery
|
||||
from core.inference.diffusion_engine_router import get_active_diffusion_engine
|
||||
from core.inference.diffusion_families import default_generation_params
|
||||
|
||||
if body.stream:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = openai_error_body(
|
||||
"Streaming image generation is not supported.", status = 400, param = "stream"
|
||||
),
|
||||
)
|
||||
try:
|
||||
width, height = _parse_openai_image_size(body.size)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(
|
||||
status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size")
|
||||
)
|
||||
|
||||
# Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same
|
||||
# accessor /images/generate uses, so a model loaded on the native engine isn't
|
||||
# wrongly reported unloaded here.
|
||||
backend = get_active_diffusion_engine()
|
||||
status = backend.status()
|
||||
if not status.get("loaded"):
|
||||
# Mirror /v1/completions and /v1/embeddings, which 503 when their backend
|
||||
# isn't loaded; the global handler turns this into the OpenAI envelope.
|
||||
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
|
||||
|
||||
# Fall back to the resolved base repo so a local-path load (whose repo_id is a
|
||||
# filesystem path) still gets the right per-model steps/guidance.
|
||||
steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo"))
|
||||
try:
|
||||
result = await asyncio.to_thread(
|
||||
backend.generate,
|
||||
prompt = body.prompt,
|
||||
width = width,
|
||||
height = height,
|
||||
steps = steps,
|
||||
guidance = guidance,
|
||||
batch_size = body.n,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope)
|
||||
# A RuntimeError with the model now unloaded means it was evicted/unloaded
|
||||
# between the readiness check above and the call (a transient race): 503.
|
||||
# Every other failure (CUDA OOM, a diffusers shape/device error, both also
|
||||
# RuntimeError) is a real 500, and its raw message must not reach the client.
|
||||
if isinstance(exc, RuntimeError) and not backend.is_loaded:
|
||||
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
|
||||
logger.error("openai_images.generate_failed: %s", exc)
|
||||
raise HTTPException(status_code = 500, detail = "Image generation failed.")
|
||||
|
||||
created = int(time.time())
|
||||
want_b64 = body.response_format == "b64_json"
|
||||
# Persist each image with its full recipe embedded, like /images/generate, so
|
||||
# the response_format=url links resolve and the images show up in the gallery.
|
||||
recipe = {
|
||||
"prompt": body.prompt,
|
||||
"negative_prompt": None,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"steps": steps,
|
||||
"guidance": guidance,
|
||||
# The batch shares one base seed, so restoring a batch_index>0 sibling needs the
|
||||
# original batch_size to replay it (same as /images/generate); persist it.
|
||||
"batch_size": body.n,
|
||||
"model": result.get("repo_id"),
|
||||
"created_at": float(created),
|
||||
}
|
||||
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed
|
||||
# per image (returned in ``seeds``), so record each image's own seed, like
|
||||
# /images/generate, or a native batch_index>0 image shows the wrong seed.
|
||||
per_image_seeds = result.get("seeds")
|
||||
|
||||
def _persist() -> list[ImageGenerationData]:
|
||||
items: list[ImageGenerationData] = []
|
||||
for index, image in enumerate(result["images"]):
|
||||
seed = (
|
||||
per_image_seeds[index]
|
||||
if per_image_seeds and index < len(per_image_seeds)
|
||||
else result["seed"]
|
||||
)
|
||||
record = image_gallery.save(image, {**recipe, "batch_index": index, "seed": seed})
|
||||
if want_b64:
|
||||
encoded = image_gallery.image_b64(record["id"])
|
||||
if encoded is None: # vanished between write and read — fail the call
|
||||
raise RuntimeError("generated image could not be read back for encoding")
|
||||
items.append(ImageGenerationData(b64_json = encoded))
|
||||
else:
|
||||
items.append(ImageGenerationData(url = _absolute_image_url(request, record["url"])))
|
||||
return items
|
||||
|
||||
try:
|
||||
data = await asyncio.to_thread(_persist)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger.error("openai_images.persist_failed: %s", exc)
|
||||
raise HTTPException(status_code = 500, detail = "Failed to save the generated image.")
|
||||
|
||||
return ImageGenerationResponse(created = created, data = data)
|
||||
|
|
|
|||
386
studio/backend/tests/test_openai_images_generations_route.py
Normal file
386
studio/backend/tests/test_openai_images_generations_route.py
Normal file
|
|
@ -0,0 +1,386 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""FastAPI round-trip tests for the OpenAI-compatible POST /v1/images/generations.
|
||||
|
||||
The diffusion backend and image gallery are replaced with light fakes, so these
|
||||
exercise the route wiring, OpenAI param mapping, validation, error envelopes, and
|
||||
response shape without torch, diffusers, weights, or a GPU. The pure helpers
|
||||
(`_parse_openai_image_size`, `default_generation_params`) are unit-tested directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import core.inference.diffusion as diffusion_module
|
||||
import core.inference.diffusion_engine_router as engine_router
|
||||
import core.inference.image_gallery as gallery_module
|
||||
from auth.authentication import get_current_subject
|
||||
from core.inference.diffusion_families import default_generation_params
|
||||
from routes.inference import router, _parse_openai_image_size
|
||||
from utils.api_errors import install_api_error_handlers
|
||||
|
||||
|
||||
# ── pure helpers ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repo_id, expected",
|
||||
[
|
||||
("unsloth/Z-Image-Turbo-GGUF", (9, 0.0)), # turbo entry, before the z-image fallback
|
||||
("unsloth/Z-Image-GGUF", (20, 4.0)),
|
||||
("unsloth/FLUX.1-schnell-GGUF", (4, 0.0)), # schnell entry, before the flux.1 entry
|
||||
("black-forest-labs/FLUX.1-dev", (28, 3.5)),
|
||||
("unsloth/FLUX.2-klein-4B-GGUF", (4, 0.0)),
|
||||
("unsloth/Qwen-Image-2512-GGUF", (20, 4.0)),
|
||||
("some/unknown-model", (9, 0.0)), # fallback
|
||||
("", (9, 0.0)),
|
||||
],
|
||||
)
|
||||
def test_default_generation_params(repo_id, expected):
|
||||
assert default_generation_params(repo_id) == expected
|
||||
|
||||
|
||||
def test_default_generation_params_specificity_ordering():
|
||||
# The "-turbo" / "-schnell" entries must win over their broader siblings; a
|
||||
# reorder that broke this would silently mis-default.
|
||||
assert default_generation_params("x/Z-Image-Turbo") != default_generation_params("x/Z-Image")
|
||||
assert default_generation_params("x/FLUX.1-schnell") != default_generation_params(
|
||||
"x/FLUX.1-dev"
|
||||
)
|
||||
|
||||
|
||||
def test_default_generation_params_falls_back_to_base_repo():
|
||||
# A local-path load: repo_id is a filesystem path that names no model, so the
|
||||
# resolved base repo is what identifies it (and distinguishes dev from schnell).
|
||||
assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-dev") == (28, 3.5)
|
||||
assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-schnell") == (
|
||||
4,
|
||||
0.0,
|
||||
)
|
||||
assert default_generation_params("/models/my-ckpt", "Qwen/Qwen-Image") == (20, 4.0)
|
||||
# repo_id wins when it already names the model; base repo is only a fallback.
|
||||
assert default_generation_params("unsloth/Z-Image-Turbo-GGUF", "Tongyi-MAI/Z-Image") == (9, 0.0)
|
||||
# Nothing identifiable -> fallback; None identifiers are skipped.
|
||||
assert default_generation_params(None, None) == (9, 0.0)
|
||||
assert default_generation_params("/models/x", None) == (9, 0.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"size, expected",
|
||||
[
|
||||
("auto", (1024, 1024)),
|
||||
("", (1024, 1024)),
|
||||
("AUTO", (1024, 1024)),
|
||||
("512x512", (512, 512)),
|
||||
("512x256", (512, 256)),
|
||||
("1792x1024", (1792, 1024)), # dall-e-3 named size: must pass the bounds
|
||||
("1024x1792", (1024, 1792)),
|
||||
(" 256 x 256 ", (256, 256)),
|
||||
],
|
||||
)
|
||||
def test_parse_image_size_ok(size, expected):
|
||||
assert _parse_openai_image_size(size) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size", ["abc", "100x100", "4096x4096", "300x300", "512", "x512", "0x0"])
|
||||
def test_parse_image_size_rejects(size):
|
||||
with pytest.raises(ValueError):
|
||||
_parse_openai_image_size(size)
|
||||
|
||||
|
||||
# ── route round-trip ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class _FakeBackend:
|
||||
def __init__(
|
||||
self,
|
||||
loaded = True,
|
||||
repo_id = "unsloth/Z-Image-Turbo-GGUF",
|
||||
base_repo = None,
|
||||
generate_error = None,
|
||||
unload_on_generate = False,
|
||||
native_seeds = False,
|
||||
) -> None:
|
||||
self._loaded = loaded
|
||||
self._repo_id = repo_id
|
||||
self._base_repo = base_repo
|
||||
# Model the native sd.cpp engine, which returns a distinct seed per image.
|
||||
self._native_seeds = native_seeds
|
||||
# When set, generate() raises this; unload_on_generate flips is_loaded off
|
||||
# first, to model the eviction/unload race vs an in-pipeline failure (OOM).
|
||||
self._generate_error = generate_error
|
||||
self._unload_on_generate = unload_on_generate
|
||||
self.calls = []
|
||||
|
||||
@property
|
||||
def is_loaded(self):
|
||||
return self._loaded
|
||||
|
||||
def status(self):
|
||||
return {
|
||||
"loaded": self._loaded,
|
||||
"repo_id": self._repo_id if self._loaded else None,
|
||||
"family": "z-image" if self._loaded else None,
|
||||
"base_repo": self._base_repo if self._loaded else None,
|
||||
"device": "cpu",
|
||||
"dtype": "float32",
|
||||
"cpu_offload": False,
|
||||
}
|
||||
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
prompt,
|
||||
width,
|
||||
height,
|
||||
steps,
|
||||
guidance,
|
||||
batch_size = 1,
|
||||
):
|
||||
if not self._loaded:
|
||||
raise RuntimeError("No diffusion model is loaded.")
|
||||
if self._generate_error is not None:
|
||||
if self._unload_on_generate:
|
||||
self._loaded = False
|
||||
raise self._generate_error
|
||||
self.calls.append(
|
||||
dict(
|
||||
prompt = prompt,
|
||||
width = width,
|
||||
height = height,
|
||||
steps = steps,
|
||||
guidance = guidance,
|
||||
batch_size = batch_size,
|
||||
)
|
||||
)
|
||||
out = {
|
||||
"images": [object() for _ in range(batch_size)],
|
||||
"seed": 4242,
|
||||
"repo_id": self._repo_id,
|
||||
}
|
||||
if self._native_seeds:
|
||||
out["seeds"] = [4242 + i for i in range(batch_size)]
|
||||
return out
|
||||
|
||||
|
||||
def _make_client(backend):
|
||||
store = {}
|
||||
|
||||
def _save(image, meta):
|
||||
image_id = f"img{len(store)}"
|
||||
record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"}
|
||||
store[image_id] = record
|
||||
return record
|
||||
|
||||
app = FastAPI()
|
||||
install_api_error_handlers(app)
|
||||
app.include_router(router, prefix = "/v1")
|
||||
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
||||
return TestClient(app), store, _save
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(monkeypatch):
|
||||
backend = _FakeBackend()
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None)
|
||||
cli.backend = backend # type: ignore[attr-defined]
|
||||
return cli
|
||||
|
||||
|
||||
def _post(client, body):
|
||||
return client.post("/v1/images/generations", json = body)
|
||||
|
||||
|
||||
def test_url_response_shape(client):
|
||||
resp = _post(client, {"prompt": "a sloth", "size": "256x256"})
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
assert set(body.keys()) == {"created", "data"}
|
||||
assert isinstance(body["created"], int) and body["created"] > 0
|
||||
assert len(body["data"]) == 1
|
||||
item = body["data"][0]
|
||||
assert "url" in item and "b64_json" not in item # exclude_none drops the unused key
|
||||
assert item["url"].endswith("/file")
|
||||
# Z-Image-Turbo defaults (9 steps, 0 guidance) flow into the backend call.
|
||||
assert client.backend.calls[0] == dict(
|
||||
prompt = "a sloth", width = 256, height = 256, steps = 9, guidance = 0.0, batch_size = 1
|
||||
)
|
||||
|
||||
|
||||
def test_b64_response_shape(client):
|
||||
resp = _post(client, {"prompt": "a sloth", "size": "256x256", "response_format": "b64_json"})
|
||||
assert resp.status_code == 200
|
||||
item = resp.json()["data"][0]
|
||||
assert "b64_json" in item and "url" not in item
|
||||
assert item["b64_json"] == "QUJD"
|
||||
|
||||
|
||||
def test_local_load_uses_base_repo_for_defaults(monkeypatch):
|
||||
# repo_id is a local path that names no model; base_repo identifies FLUX.1-dev,
|
||||
# so the route must pick 28 steps / 3.5 guidance, not the 9/0 fallback.
|
||||
backend = _FakeBackend(repo_id = "/models/my-flux", base_repo = "black-forest-labs/FLUX.1-dev")
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
|
||||
assert resp.status_code == 200
|
||||
assert backend.calls[0]["steps"] == 28 and backend.calls[0]["guidance"] == 3.5
|
||||
|
||||
|
||||
def test_pipeline_runtime_error_is_sanitized_500(monkeypatch):
|
||||
# A RuntimeError raised inside the pipeline while the model stays loaded (e.g.
|
||||
# CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that
|
||||
# echoes the raw exception text.
|
||||
oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)")
|
||||
backend = _FakeBackend(generate_error = oom)
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
|
||||
assert resp.status_code == 500
|
||||
assert resp.json()["error"]["message"] == "Image generation failed."
|
||||
assert "CUDA" not in resp.text # raw exception text must not leak
|
||||
|
||||
|
||||
def test_unload_race_returns_503(monkeypatch):
|
||||
# The model is evicted/unloaded between the readiness check and the call: the
|
||||
# RuntimeError with is_loaded now False is the one case that maps to 503.
|
||||
backend = _FakeBackend(
|
||||
generate_error = RuntimeError("No diffusion model is loaded."),
|
||||
unload_on_generate = True,
|
||||
)
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
|
||||
assert resp.status_code == 503
|
||||
err = resp.json()["error"]
|
||||
assert err["type"] == "api_error"
|
||||
# The 503 carries the fixed sanitized message, not the raw exception text.
|
||||
assert err["message"] == "No image model loaded. Load an image model first."
|
||||
|
||||
|
||||
def test_non_runtime_pipeline_error_is_500(monkeypatch):
|
||||
# A non-RuntimeError from the pipeline (the model stays loaded) must not route
|
||||
# to the 503 branch (which is gated on isinstance RuntimeError) -> sanitized 500.
|
||||
backend = _FakeBackend(generate_error = ValueError("bad tensor shape"))
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
|
||||
assert resp.status_code == 500
|
||||
assert "shape" not in resp.text
|
||||
|
||||
|
||||
def test_n_maps_to_batch(client):
|
||||
resp = _post(client, {"prompt": "p", "size": "256x256", "n": 3})
|
||||
assert resp.status_code == 200
|
||||
assert len(resp.json()["data"]) == 3
|
||||
assert client.backend.calls[0]["batch_size"] == 3
|
||||
|
||||
|
||||
def test_batch_persists_batch_size(monkeypatch):
|
||||
# n>1 must persist batch_size in each gallery record so the Studio restore path
|
||||
# can replay a batch_index>0 sibling (which shares the batch's single seed).
|
||||
backend = _FakeBackend()
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3})
|
||||
assert resp.status_code == 200
|
||||
records = sorted(store.values(), key = lambda r: r["batch_index"])
|
||||
assert [r["batch_index"] for r in records] == [0, 1, 2]
|
||||
assert all(r["batch_size"] == 3 for r in records)
|
||||
|
||||
|
||||
def test_uses_active_engine_not_diffusers_singleton(monkeypatch):
|
||||
# On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the
|
||||
# diffusers singleton. The route must query get_active_diffusion_engine (like
|
||||
# /images/generate) or it 503s a model that is loaded and usable.
|
||||
active = _FakeBackend(loaded = True) # the active (e.g. sd_cpp) engine, loaded
|
||||
idle_diffusers = _FakeBackend(loaded = False) # diffusers singleton, empty
|
||||
monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: active)
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: idle_diffusers)
|
||||
cli, store, _save = _make_client(active)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
|
||||
assert resp.status_code == 200
|
||||
assert len(active.calls) == 1 # the active engine did the work, not the idle singleton
|
||||
|
||||
|
||||
def test_native_batch_persists_per_image_seed(monkeypatch):
|
||||
# The native sd.cpp engine returns a distinct seed per image (base+index) in
|
||||
# "seeds"; each gallery record must store its own seed (like /images/generate),
|
||||
# not the shared base, or a restored batch_index>0 image shows the wrong seed.
|
||||
backend = _FakeBackend(native_seeds = True)
|
||||
monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3})
|
||||
assert resp.status_code == 200
|
||||
records = sorted(store.values(), key = lambda r: r["batch_index"])
|
||||
assert [r["seed"] for r in records] == [4242, 4243, 4244]
|
||||
|
||||
|
||||
def test_null_fields_coalesce_to_defaults(client):
|
||||
# OpenAI marks n/size/response_format nullable-with-default: null -> default.
|
||||
resp = _post(client, {"prompt": "p", "n": None, "size": None, "response_format": None})
|
||||
assert resp.status_code == 200
|
||||
assert len(resp.json()["data"]) == 1
|
||||
assert "url" in resp.json()["data"][0]
|
||||
assert client.backend.calls[0]["width"] == 1024 # size null -> auto -> 1024
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body, param",
|
||||
[
|
||||
({"size": "256x256"}, "prompt"), # missing prompt
|
||||
({"prompt": "", "size": "256x256"}, "prompt"), # empty prompt
|
||||
({"prompt": "p", "size": "300x300"}, "size"), # not multiple of 16
|
||||
({"prompt": "p", "size": "abc"}, "size"), # unparseable
|
||||
({"prompt": "p", "stream": True}, "stream"), # streaming unsupported
|
||||
],
|
||||
)
|
||||
def test_validation_400_with_param(client, body, param):
|
||||
resp = _post(client, body)
|
||||
assert resp.status_code == 400
|
||||
err = resp.json()["error"]
|
||||
assert err["type"] == "invalid_request_error"
|
||||
assert err["param"] == param
|
||||
for k in ("message", "code"):
|
||||
assert k in err
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n", [0, 11, -1])
|
||||
def test_n_out_of_range_400(client, n):
|
||||
resp = _post(client, {"prompt": "p", "size": "256x256", "n": n})
|
||||
assert resp.status_code == 400
|
||||
assert resp.json()["error"]["type"] == "invalid_request_error"
|
||||
|
||||
|
||||
def test_no_model_loaded_503(monkeypatch):
|
||||
backend = _FakeBackend(loaded = False)
|
||||
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
||||
cli, store, _save = _make_client(backend)
|
||||
monkeypatch.setattr(gallery_module, "save", _save)
|
||||
resp = cli.post("/v1/images/generations", json = {"prompt": "p"})
|
||||
assert resp.status_code == 503
|
||||
# 503 still wears the OpenAI envelope (api_error) on the /v1 surface.
|
||||
assert resp.json()["error"]["type"] == "api_error"
|
||||
|
||||
|
||||
def test_auth_required():
|
||||
backend = _FakeBackend()
|
||||
app = FastAPI()
|
||||
install_api_error_handlers(app)
|
||||
app.include_router(router, prefix = "/v1")
|
||||
# No dependency override: the real auth dependency runs and rejects.
|
||||
resp = TestClient(app).post("/v1/images/generations", json = {"prompt": "p"})
|
||||
assert resp.status_code in (401, 403)
|
||||
Loading…
Add table
Add a link
Reference in a new issue