diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 5beb8e36b6..7f72743d35 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -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: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 11d3e5025c..04b70e7e11 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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 'x' (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.") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 67e541a8d8..7400a9f657 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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 'x', 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) diff --git a/studio/backend/tests/test_openai_images_generations_route.py b/studio/backend/tests/test_openai_images_generations_route.py new file mode 100644 index 0000000000..f82586a173 --- /dev/null +++ b/studio/backend/tests/test_openai_images_generations_route.py @@ -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)