9a933d375b
Three fixes from review. The 0.39-only pipeline classes (Flux2Klein, Z-Image, Krea 2, LTX-2, HunyuanImage) were resolved by getattr deep in the load, so on the older diffusers that packaging still allows on Python 3.9 -- diffusers dropped 3.9 in 0.38 and this project still supports it, so the 0.39 floor has to be conditional or the extra becomes unresolvable -- an advertised model failed with a bare AttributeError after its checkpoint had already been downloaded. Krea 2 already guarded itself this way; assert_pipeline_class_available now runs the same check for every image and video family from validation, before any fetch, and names the version and the fix. WebM export accumulated the whole VP9 output in a BytesIO and returned it as one bytes object that the response held again. The request caps allow 2048x2048 for 1024 frames, so an export runs to hundreds of MB and concurrent clicks could exhaust the process, while the MP4 route beside it already streamed from disk. transcode_to_file encodes to a temp file and the route returns a FileResponse with a background unlink, so nothing large is resident. A conditioned generation's recipe carried only the txt2img fields, so the gallery presented an inpaint or upscale result as a complete Create recipe and restoring it replayed an unrelated text-to-image request. The images themselves are still not persisted (user uploads with their own lifetime), but the workflow and its scalars are, restore reapplies them, and the toast now names the inputs that have to be supplied again instead of silently landing on Create. Reported by Codex.
1139 lines
47 KiB
Python
1139 lines
47 KiB
Python
# 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 diffusion image routes.
|
|
|
|
The diffusion backend is replaced with a lightweight fake, so these exercise the
|
|
route wiring, validation (422), error mapping, and response shapes without torch,
|
|
diffusers, weights, or a GPU.
|
|
"""
|
|
|
|
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.gpu_arbiter as gpu_arbiter
|
|
import core.inference.image_gallery as gallery_module
|
|
from auth.authentication import get_current_subject
|
|
from routes.inference import studio_router
|
|
|
|
|
|
class _FakeBackend:
|
|
def __init__(self) -> None:
|
|
self.loaded = False
|
|
# Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload route reads this
|
|
# to keep DIFFUSION ownership while a concurrent load is still loading.
|
|
self.loading: tuple = ()
|
|
|
|
@property
|
|
def is_loaded(self) -> bool:
|
|
return self.loaded
|
|
|
|
def loading_repo_ids(self) -> tuple:
|
|
return tuple(self.loading)
|
|
|
|
def validate_load_request(
|
|
self,
|
|
model_path,
|
|
*,
|
|
gguf_filename = None,
|
|
family_override = None,
|
|
model_kind = None,
|
|
base_repo = None,
|
|
):
|
|
# Mirror the real backend's cheap validation so the route's validate-before-evict ordering is
|
|
# exercised.
|
|
from core.inference.diffusion import resolve_model_kind
|
|
from core.inference.diffusion_families import detect_family
|
|
|
|
kind = resolve_model_kind(gguf_filename, model_kind)
|
|
if kind in ("gguf", "single_file") and not gguf_filename:
|
|
raise ValueError("a single-file checkpoint name is required.")
|
|
# Non-GGUF loads are gated to unsloth/* (or a local path), like the real backend.
|
|
if kind != "gguf" and not model_path.lower().startswith("unsloth/"):
|
|
raise ValueError(
|
|
f"Non-GGUF diffusion loads are restricted to unsloth/* repos; got '{model_path}'."
|
|
)
|
|
# A client-supplied base_repo clears the same trust bar as the real backend, so the route's
|
|
# validate-before-evict rejects an untrusted companion base.
|
|
if base_repo and base_repo.strip() and not base_repo.lower().startswith("unsloth/"):
|
|
raise ValueError(
|
|
f"base_repo is restricted to unsloth/* repos (or a local path); got '{base_repo}'."
|
|
)
|
|
fam = detect_family(model_path, family_override)
|
|
if fam is None:
|
|
raise ValueError(f"Could not infer a diffusion family for '{model_path}'.")
|
|
return fam
|
|
|
|
def begin_load(self, model_path, **kwargs):
|
|
# The real backend loads on a thread; the fake completes instantly.
|
|
self.loaded = True
|
|
self.last_load_kwargs = dict(kwargs)
|
|
return {
|
|
"loaded": True,
|
|
"repo_id": model_path,
|
|
"family": "z-image",
|
|
"base_repo": kwargs.get("base_repo") or "base/repo",
|
|
"device": "cpu",
|
|
"dtype": "float32",
|
|
"cpu_offload": False,
|
|
"offload_policy": "none",
|
|
"vae_tiling": False,
|
|
"memory_mode": kwargs.get("memory_mode") or "auto",
|
|
}
|
|
|
|
def load_progress(self):
|
|
return {
|
|
"phase": "ready" if self.loaded else None,
|
|
"bytes_downloaded": 0,
|
|
"bytes_total": 0,
|
|
"fraction": 1.0 if self.loaded else 0.0,
|
|
"error": None,
|
|
}
|
|
|
|
def generate(
|
|
self,
|
|
*,
|
|
seed = None,
|
|
batch_size = 1,
|
|
prompts = None,
|
|
seeds = None,
|
|
**kwargs,
|
|
):
|
|
if not self.loaded:
|
|
raise RuntimeError("No diffusion model is loaded.")
|
|
if prompts is not None or seeds is not None:
|
|
# List-driven batch: the LIST sets the image count and each image's own seed (batch_size is only a
|
|
# per-forward cap), exactly as the real engine reports it.
|
|
base = seeds[0] if seeds else (seed if seed is not None else 4242)
|
|
count = len(prompts) if prompts is not None else len(seeds)
|
|
per_image = seeds if seeds is not None else [base + i for i in range(count)]
|
|
return {
|
|
"images": [object() for _ in range(count)],
|
|
"seed": base,
|
|
"seeds": list(per_image),
|
|
"repo_id": "x/z-image",
|
|
}
|
|
# The real backend returns the PIL images and the route persists them; the fake returns sentinels
|
|
# since image_gallery is stubbed in the fixture.
|
|
return {
|
|
"images": [object() for _ in range(batch_size)],
|
|
"seed": seed if seed is not None else 4242,
|
|
"repo_id": "x/z-image",
|
|
# The real backend reports the workflow it resolved; the recipe records it.
|
|
"workflow": (
|
|
"inpaint"
|
|
if kwargs.get("mask_image")
|
|
else ("img2img" if kwargs.get("init_image") else "txt2img")
|
|
),
|
|
}
|
|
|
|
def generate_progress(self):
|
|
# Idle by default; the persist-window override lives in the route, not here.
|
|
return {"active": False, "step": 0, "total_steps": 0, "fraction": 0.0, "eta_seconds": None}
|
|
|
|
def unload(self):
|
|
self.loaded = False
|
|
return _unloaded_status()
|
|
|
|
def status(self):
|
|
return {**_unloaded_status(), "loaded": self.loaded}
|
|
|
|
|
|
def _unloaded_status():
|
|
return {
|
|
"loaded": False,
|
|
"repo_id": None,
|
|
"family": None,
|
|
"base_repo": None,
|
|
"device": None,
|
|
"dtype": None,
|
|
"cpu_offload": False,
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def client(monkeypatch, tmp_path):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
# Neutralise the engine router so the routes deterministically drive this fake (diffusers) backend
|
|
# regardless of the host's real device, and never attempt a native sd.cpp install. The selection
|
|
# logic is covered in test_diffusion_engine_router.py.
|
|
import core.inference.diffusion_engine_router as engine_router
|
|
|
|
# Delegate to whatever get_diffusion_backend currently returns, so per-test re-patches of the
|
|
# backend still flow through the routes.
|
|
monkeypatch.setattr(
|
|
engine_router,
|
|
"select_and_activate_engine",
|
|
lambda fam, **kw: diffusion_module.get_diffusion_backend(),
|
|
)
|
|
monkeypatch.setattr(
|
|
engine_router,
|
|
"get_active_diffusion_engine",
|
|
lambda: diffusion_module.get_diffusion_backend(),
|
|
)
|
|
monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers")
|
|
monkeypatch.setattr(engine_router, "_fallback_reason", None)
|
|
# Isolate from the real GPU arbiter: reset ownership and stub the evictors so the load route's
|
|
# acquire_for() never touches live backend singletons.
|
|
monkeypatch.setattr(gpu_arbiter, "_owner", None)
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None)
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None)
|
|
|
|
# In-memory gallery backed by tmp files, so routes exercise persistence wiring without PIL/real
|
|
# disk under studio_root.
|
|
store: dict[str, dict] = {}
|
|
|
|
def _save(image, meta):
|
|
image_id = f"img{len(store)}"
|
|
(tmp_path / f"{image_id}.png").write_bytes(b"PNG")
|
|
record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"}
|
|
store[image_id] = record
|
|
return record
|
|
|
|
def _clear():
|
|
n = len(store)
|
|
store.clear()
|
|
return n
|
|
|
|
monkeypatch.setattr(gallery_module, "save", _save)
|
|
monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None)
|
|
|
|
def _list_images(
|
|
limit = None,
|
|
offset = 0,
|
|
*,
|
|
valid = None,
|
|
):
|
|
ordered = sorted(store.values(), key = lambda r: r.get("created_at", 0.0), reverse = True)
|
|
if valid is not None:
|
|
ordered = [r for r in ordered if valid(r)]
|
|
return ordered[offset:] if limit is None else ordered[offset : offset + limit]
|
|
|
|
monkeypatch.setattr(gallery_module, "list_images", _list_images)
|
|
monkeypatch.setattr(
|
|
gallery_module,
|
|
"image_path",
|
|
lambda i: (tmp_path / f"{i}.png") if i in store else None,
|
|
)
|
|
# The serve route resolves through owned_image_path; the fake store only holds owned records, so a
|
|
# stem not in it is treated as foreign and refused, like the real guard.
|
|
monkeypatch.setattr(
|
|
gallery_module,
|
|
"owned_image_path",
|
|
lambda i: (tmp_path / f"{i}.png") if i in store else None,
|
|
)
|
|
monkeypatch.setattr(gallery_module, "delete", lambda i: store.pop(i, None) is not None)
|
|
monkeypatch.setattr(gallery_module, "clear", _clear)
|
|
|
|
app = FastAPI()
|
|
app.include_router(studio_router, prefix = "/api/inference")
|
|
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
|
return TestClient(app)
|
|
|
|
|
|
def test_load_generate_status_unload_roundtrip(client):
|
|
loaded = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "unsloth/Z-Image-Turbo-GGUF",
|
|
"gguf_filename": "z-image-turbo-Q4_K_S.gguf",
|
|
"base_repo": "unsloth/Z-Image-base",
|
|
},
|
|
)
|
|
assert loaded.status_code == 200
|
|
body = loaded.json()
|
|
assert body["loaded"] is True and body["family"] == "z-image"
|
|
|
|
assert client.get("/api/inference/images/status").json()["loaded"] is True
|
|
|
|
gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7})
|
|
assert gen.status_code == 200
|
|
# One persisted record carrying the full recipe back.
|
|
images = gen.json()["images"]
|
|
assert len(images) == 1
|
|
img = images[0]
|
|
assert img["seed"] == 7 and img["prompt"] == "a sloth" and img["id"]
|
|
|
|
# The image is now listable, fetchable, and deletable.
|
|
listed = client.get("/api/inference/images/gallery").json()["images"]
|
|
assert [i["id"] for i in listed] == [img["id"]]
|
|
assert client.get(img["url"]).status_code == 200
|
|
assert client.delete(img["url"].removesuffix("/file")).status_code == 200
|
|
assert client.get("/api/inference/images/gallery").json()["images"] == []
|
|
|
|
unloaded = client.post("/api/inference/images/unload")
|
|
assert unloaded.status_code == 200 and unloaded.json()["loaded"] is False
|
|
assert client.get("/api/inference/images/status").json()["loaded"] is False
|
|
|
|
|
|
def test_gallery_serve_refuses_unowned_id(client):
|
|
# The serve route resolves through the ownership guard, so a guessed stem for a PNG the gallery
|
|
# does not own is a 404, not a stream of foreign bytes.
|
|
assert client.get("/api/inference/images/gallery/family-photo/file").status_code == 404
|
|
|
|
|
|
def test_generate_holds_progress_active_during_persist(client, monkeypatch):
|
|
# generate-progress must stay active while a finished generation is still writing its gallery
|
|
# record, so a concurrent reload's mount probe keeps polling instead of refreshing the gallery
|
|
# early. Probe the persist counter from inside the save call, and confirm it clears afterwards.
|
|
import core.inference.image_gallery as gallery_module
|
|
import routes.inference as inf
|
|
|
|
client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "unsloth/Z-Image-Turbo-GGUF",
|
|
"gguf_filename": "z-image-turbo-Q4_K_S.gguf",
|
|
"base_repo": "unsloth/Z-Image-base",
|
|
},
|
|
)
|
|
|
|
# Idle before any generation.
|
|
assert client.get("/api/inference/images/generate-progress").json()["active"] is False
|
|
|
|
seen = {}
|
|
real_save = gallery_module.save
|
|
|
|
def _probe_save(image, meta):
|
|
seen["during"] = inf._diffusion_persist_active
|
|
return real_save(image, meta)
|
|
|
|
monkeypatch.setattr(gallery_module, "save", _probe_save)
|
|
|
|
gen = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7})
|
|
assert gen.status_code == 200
|
|
# Active while the record was being persisted, and back to idle once the route returned.
|
|
assert seen["during"] >= 1
|
|
assert inf._diffusion_persist_active == 0
|
|
assert client.get("/api/inference/images/generate-progress").json()["active"] is False
|
|
|
|
|
|
def test_load_rejects_untrusted_base_repo(client):
|
|
# A trusted GGUF model_path paired with an untrusted remote base_repo is rejected at the route
|
|
# (validate runs before the GPU handoff), so an authenticated client cannot make the server fetch
|
|
# and deserialize an arbitrary companion repo.
|
|
r = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "unsloth/Z-Image-Turbo-GGUF",
|
|
"gguf_filename": "z-image-turbo-Q4_K_S.gguf",
|
|
"base_repo": "evil/companions",
|
|
},
|
|
)
|
|
assert r.status_code == 400
|
|
assert "base_repo" in r.json()["detail"]
|
|
assert client.get("/api/inference/images/status").json()["loaded"] is False
|
|
|
|
|
|
def test_unload_keeps_ownership_when_a_model_is_still_resident(client, monkeypatch):
|
|
# The unload route must drop DIFFUSION ownership only when nothing is resident. If a concurrent
|
|
# load re-established residency while the slow unload ran, releasing would clear the newer claim
|
|
# and a later chat load would skip eviction and OOM.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
gpu_arbiter._owner = gpu_arbiter.DIFFUSION
|
|
|
|
# Simulate a concurrent load having re-loaded: unload leaves the engine resident.
|
|
backend.loaded = True
|
|
monkeypatch.setattr(backend, "unload", lambda: {**_unloaded_status(), "loaded": True})
|
|
r = client.post("/api/inference/images/unload")
|
|
assert r.status_code == 200
|
|
assert gpu_arbiter.current_owner() == gpu_arbiter.DIFFUSION # ownership retained
|
|
|
|
# The normal case (nothing resident after unload) still releases ownership.
|
|
monkeypatch.setattr(backend, "unload", lambda: {**_unloaded_status(), "loaded": False})
|
|
backend.loaded = False
|
|
r = client.post("/api/inference/images/unload")
|
|
assert r.status_code == 200
|
|
assert gpu_arbiter.current_owner() is None
|
|
|
|
|
|
def test_unload_keeps_ownership_when_a_load_is_in_flight(client, monkeypatch):
|
|
# A concurrent /images/load re-acquires DIFFUSION and starts a background load, so the engine is
|
|
# NOT is_loaded yet but a load IS in flight. The unload route must keep ownership on the in-flight
|
|
# state alone, since is_loaded stays False for the whole download/finalize window.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
gpu_arbiter._owner = gpu_arbiter.DIFFUSION
|
|
|
|
backend.loaded = False
|
|
backend.loading = ("unsloth/z-image-turbo",)
|
|
monkeypatch.setattr(backend, "unload", lambda: {**_unloaded_status(), "loaded": False})
|
|
r = client.post("/api/inference/images/unload")
|
|
assert r.status_code == 200
|
|
assert gpu_arbiter.current_owner() == gpu_arbiter.DIFFUSION # ownership retained for the load
|
|
|
|
backend.loading = ()
|
|
|
|
|
|
def test_generate_batch_size_persists_each_image(client):
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
resp = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "p", "batch_size": 3, "seed": 5},
|
|
)
|
|
assert resp.status_code == 200
|
|
images = resp.json()["images"]
|
|
assert len(images) == 3
|
|
assert all(i["seed"] == 5 for i in images) # the batch shares one seed
|
|
assert len({i["id"] for i in images}) == 3 # but each is a distinct record
|
|
assert len(client.get("/api/inference/images/gallery").json()["images"]) == 3
|
|
|
|
|
|
def test_generate_seed_list_records_replay_from_each_own_seed(client):
|
|
# A seeds LIST sets each image's own seed, so the recipe must NOT claim the base seed + the
|
|
# request's batch_size: restore prefers batch_seed (`batch_seed ?? seed`), which would regenerate
|
|
# seed 5 for the seed-99 image.
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
resp = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "p", "seeds": [5, 99]},
|
|
)
|
|
assert resp.status_code == 200
|
|
images = resp.json()["images"]
|
|
assert [i["seed"] for i in images] == [5, 99]
|
|
assert [i["batch_seed"] for i in images] == [5, 99] # replays THIS image, not the base
|
|
assert [i["batch_size"] for i in images] == [1, 1] # as a single image, not a batch
|
|
|
|
|
|
def test_generate_prompt_list_records_each_prompt_and_seed(client):
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
resp = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "unused", "prompts": ["a cat", "a dog"], "seed": 10},
|
|
)
|
|
assert resp.status_code == 200
|
|
images = resp.json()["images"]
|
|
assert [i["prompt"] for i in images] == ["a cat", "a dog"]
|
|
assert [i["seed"] for i in images] == [10, 11]
|
|
assert [i["batch_seed"] for i in images] == [10, 11]
|
|
assert [i["batch_size"] for i in images] == [1, 1]
|
|
|
|
|
|
def test_generate_legacy_batch_still_records_the_base_seed_and_size(client):
|
|
# The batch_size path is unchanged: those images DO share one base seed, so restore must replay
|
|
# the whole batch, not the derived per-image seed.
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
resp = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "p", "batch_size": 3, "seed": 5},
|
|
)
|
|
images = resp.json()["images"]
|
|
assert all(i["batch_seed"] == 5 for i in images)
|
|
assert all(i["batch_size"] == 3 for i in images)
|
|
assert [i["batch_index"] for i in images] == [0, 1, 2]
|
|
|
|
|
|
def test_generate_request_rejects_zero_denoise_strength():
|
|
# strength 0 does NOT keep the source: every diffusers img2img/inpaint pipeline derives t_start =
|
|
# steps - int(steps * strength), so 0 leaves zero denoising steps (FLUX/Qwen/Z-Image raise, SDXL
|
|
# img2img crashes on empty latents). Reject it as a 422 up front.
|
|
import pydantic
|
|
|
|
from models.inference import DiffusionGenerateRequest
|
|
|
|
with pytest.raises(pydantic.ValidationError):
|
|
DiffusionGenerateRequest(prompt = "x", strength = 0.0)
|
|
assert DiffusionGenerateRequest(prompt = "x", strength = 0.1).strength == 0.1
|
|
assert DiffusionGenerateRequest(prompt = "x", strength = 1.0).strength == 1.0
|
|
assert DiffusionGenerateRequest(prompt = "x").strength is None # unset stays the pipe default
|
|
|
|
|
|
def test_gallery_pagination(client):
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
client.post("/api/inference/images/generate", json = {"prompt": "p", "batch_size": 5, "seed": 1})
|
|
page1 = client.get("/api/inference/images/gallery?limit=2&offset=0").json()
|
|
assert len(page1["images"]) == 2 and page1["has_more"] is True
|
|
last = client.get("/api/inference/images/gallery?limit=2&offset=4").json()
|
|
assert len(last["images"]) == 1 and last["has_more"] is False
|
|
|
|
|
|
def test_generate_rejects_non_multiple_of_16(client):
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
# Odd, and a multiple of 8 that isn't a multiple of 16: both rejected, since Z-Image requires
|
|
# dimensions divisible by 16.
|
|
for bad in (1001, 1000):
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": bad})
|
|
assert resp.status_code == 422, bad
|
|
# A multiple of 16 is accepted.
|
|
ok = client.post("/api/inference/images/generate", json = {"prompt": "p", "width": 1024})
|
|
assert ok.status_code == 200
|
|
|
|
|
|
def test_generate_rejects_batch_seed_past_json_safe_range(client):
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
# A seed at the cap with a batch derives per-image seeds past the JSON-safe range, so the request
|
|
# is rejected.
|
|
over = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "p", "seed": 2**53 - 1, "batch_size": 2},
|
|
)
|
|
assert over.status_code == 422
|
|
# The top-of-batch seed lands exactly on the cap: still JSON-safe, so accepted.
|
|
ok = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {"prompt": "p", "seed": 2**53 - 2, "batch_size": 2},
|
|
)
|
|
assert ok.status_code == 200
|
|
|
|
|
|
def test_non_gguf_load_restricted_to_unsloth(client):
|
|
# gguf_filename is optional now; with none the load is a full-pipeline kind, gated to unsloth/*
|
|
# repos, so a non-unsloth repo is rejected with a 400.
|
|
resp = client.post("/api/inference/images/load", json = {"model_path": "x/z-image"})
|
|
assert resp.status_code == 400
|
|
assert "unsloth" in resp.json()["detail"].lower()
|
|
|
|
|
|
def test_pipeline_load_allowed_for_unsloth_repo(client):
|
|
# An unsloth/* repo with no filename loads as a full diffusers pipeline, so the route forwards
|
|
# model_kind="pipeline" to begin_load.
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "unsloth/Z-Image-Turbo-unsloth-bnb-4bit"}
|
|
)
|
|
assert resp.status_code == 200
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
assert backend.last_load_kwargs["model_kind"] == "pipeline"
|
|
assert backend.last_load_kwargs.get("gguf_filename") is None
|
|
|
|
|
|
def test_generate_without_load_returns_409(client):
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
|
assert resp.status_code == 409
|
|
|
|
|
|
def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
|
|
# A loaded model that fails mid-pipeline (CUDA OOM, a RuntimeError) is a server failure: 500 with
|
|
# a message built from FIXED text, not a 409 and not the raw exception. The class of failure is
|
|
# named so the page can suggest something; none of the engine's own text is echoed, since it can
|
|
# carry local paths and argv.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
backend.loaded = True
|
|
|
|
def _oom(**kwargs):
|
|
raise RuntimeError(
|
|
"CUDA out of memory. Tried to allocate 20.00 GiB at /home/u/models/x.safetensors"
|
|
)
|
|
|
|
monkeypatch.setattr(backend, "generate", _oom)
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
|
assert resp.status_code == 500
|
|
detail = resp.json()["detail"]
|
|
assert detail.startswith("Image generation failed.")
|
|
assert "ran out of memory" in detail
|
|
for leak in ("CUDA", "20.00 GiB", "/home/u", "safetensors"):
|
|
assert leak not in detail
|
|
|
|
|
|
def test_generate_native_process_death_names_the_engine_not_its_output(client, monkeypatch):
|
|
# What a Metal host hits: the native renderer aborts inside its own text encoder. The page used
|
|
# to show only "Image generation failed."; it now says which component died, while the abort
|
|
# backtrace (full of local paths) stays in the server log.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
backend.loaded = True
|
|
|
|
def _abort(**kwargs):
|
|
raise RuntimeError(
|
|
"sd-server connection lost during img_gen poll (process exited, code -6)\n"
|
|
"Last output:\n0 sd-server ggml_abort + 156 at /Users/me/.cache/sd-cli"
|
|
)
|
|
|
|
monkeypatch.setattr(backend, "generate", _abort)
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
|
assert resp.status_code == 500
|
|
detail = resp.json()["detail"]
|
|
assert "native image renderer stopped" in detail
|
|
for leak in ("ggml_abort", "/Users/me", "img_gen", "code -6"):
|
|
assert leak not in detail
|
|
|
|
|
|
def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch):
|
|
# A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled" must stay a
|
|
# sanitized 500, not misroute to 409 and echo that output (path/arg leak).
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
backend.loaded = True
|
|
|
|
def _fail(**kwargs):
|
|
raise RuntimeError("sd-cli exited 1. Last output:\nop cancelled at /home/u/models/x.gguf")
|
|
|
|
monkeypatch.setattr(backend, "generate", _fail)
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
|
assert resp.status_code == 500
|
|
detail = resp.json()["detail"]
|
|
assert detail.startswith("Image generation failed.")
|
|
assert "cancelled" not in detail and "models" not in detail and "/home/u" not in detail
|
|
|
|
|
|
def test_generate_user_cancellation_returns_409(client, monkeypatch):
|
|
# The exact cancellation sentinel both engines raise is client-state (409).
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
backend.loaded = True
|
|
|
|
def _cancel(**kwargs):
|
|
raise RuntimeError("Diffusion generation was cancelled.")
|
|
|
|
monkeypatch.setattr(backend, "generate", _cancel)
|
|
resp = client.post("/api/inference/images/generate", json = {"prompt": "p"})
|
|
assert resp.status_code == 409
|
|
assert resp.json()["detail"] == "Diffusion generation was cancelled."
|
|
|
|
|
|
def test_load_unknown_family_returns_400(client, monkeypatch):
|
|
def _raise(*a, **k):
|
|
raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.")
|
|
|
|
backend = _FakeBackend()
|
|
# Validation runs in the pre-flight (before the GPU is taken), so that is where an unsupported
|
|
# model is rejected now.
|
|
backend.validate_load_request = _raise
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 400
|
|
assert "isn't a supported image-generation model" in resp.json()["detail"]
|
|
|
|
|
|
def test_load_validation_failure_does_not_evict_chat(client, monkeypatch):
|
|
# A rejected image-model pick must not tear down the user's loaded chat model: validation runs
|
|
# before acquire_for, so chat keeps the GPU on a 400.
|
|
monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT)
|
|
evicted = []
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True))
|
|
|
|
backend = _FakeBackend()
|
|
|
|
def _raise(*a, **k):
|
|
raise ValueError("'x/y' isn't a supported image-generation model.")
|
|
|
|
backend.validate_load_request = _raise
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 400
|
|
assert evicted == [] # chat backend was never evicted
|
|
assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT
|
|
|
|
|
|
def test_load_refused_during_training_does_not_evict_chat(client, monkeypatch):
|
|
# An image load while training is active is refused (409) before the GPU is taken, so the training
|
|
# run and the loaded chat model are both untouched.
|
|
import core.training as core_training
|
|
|
|
monkeypatch.setattr(gpu_arbiter, "_owner", gpu_arbiter.CHAT)
|
|
evicted = []
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: evicted.append(True))
|
|
|
|
class _Training:
|
|
def is_training_active(self):
|
|
return True
|
|
|
|
monkeypatch.setattr(core_training, "get_training_backend", lambda: _Training())
|
|
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"},
|
|
)
|
|
assert resp.status_code == 409
|
|
assert "training" in resp.json()["detail"].lower()
|
|
assert evicted == [] # chat backend was never evicted
|
|
assert gpu_arbiter.current_owner() == gpu_arbiter.CHAT
|
|
|
|
|
|
def test_load_progress_route(client):
|
|
# Before load: idle.
|
|
idle = client.get("/api/inference/images/load-progress")
|
|
assert idle.status_code == 200 and idle.json()["phase"] is None
|
|
# After load: the fake reports ready.
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
ready = client.get("/api/inference/images/load-progress")
|
|
assert ready.json()["phase"] == "ready"
|
|
|
|
|
|
def test_routes_require_auth():
|
|
# No dependency override: the auth dependency must reject the request.
|
|
app = FastAPI()
|
|
app.include_router(studio_router, prefix = "/api/inference")
|
|
unauth = TestClient(app)
|
|
assert unauth.get("/api/inference/images/status").status_code in (401, 403)
|
|
|
|
|
|
def test_invalid_family_returns_400_without_evicting_chat(client):
|
|
# An undetectable family fails validation BEFORE the GPU handoff, so the arbiter is never acquired
|
|
# and a loaded chat model would not be evicted.
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/y", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 400
|
|
assert "family" in resp.json()["detail"]
|
|
assert gpu_arbiter._owner is None
|
|
|
|
|
|
def test_validate_filenotfound_maps_to_400_without_eviction(client, monkeypatch):
|
|
def _raise_fnf(*a, **k):
|
|
raise FileNotFoundError("'q.gguf' not found under /models/x.")
|
|
|
|
backend = _FakeBackend()
|
|
backend.validate_load_request = _raise_fnf
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "/models/x", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 400
|
|
assert gpu_arbiter._owner is None
|
|
|
|
|
|
def test_memory_mode_threads_through_to_backend(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "memory_mode": "low_vram"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["memory_mode"] == "low_vram"
|
|
assert backend.last_load_kwargs.get("memory_mode") == "low_vram"
|
|
|
|
|
|
def test_transformer_quant_threads_through_to_backend(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "auto"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert backend.last_load_kwargs.get("transformer_quant") == "auto"
|
|
|
|
|
|
def test_transformer_quant_fast_accum_threads_through(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"transformer_quant": "fp8",
|
|
"transformer_quant_fast_accum": False,
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False
|
|
|
|
|
|
def test_transformer_prequant_path_threads_through(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"transformer_quant": "fp8",
|
|
"transformer_prequant_path": "/data/zimage_fp8.pt",
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert backend.last_load_kwargs.get("transformer_prequant_path") == "/data/zimage_fp8.pt"
|
|
|
|
|
|
def test_attention_backend_threads_through(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"attention_backend": "cudnn",
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert backend.last_load_kwargs.get("attention_backend") == "cudnn"
|
|
|
|
|
|
def test_invalid_attention_backend_returns_422(client):
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "attention_backend": "bogus"},
|
|
)
|
|
assert resp.status_code == 422
|
|
|
|
|
|
def test_prequant_path_doc_describes_allowlist_not_toggle():
|
|
# The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a directory allowlist,
|
|
# not a =1 toggle (bare on/off tokens are dropped), so operators following the doc don't get every
|
|
# request silently refused.
|
|
from models.inference import DiffusionLoadRequest
|
|
|
|
desc = DiffusionLoadRequest.model_fields["transformer_prequant_path"].description
|
|
assert "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH" in desc
|
|
assert "=1" not in desc
|
|
assert "allowlist" in desc.lower() or "director" in desc.lower()
|
|
|
|
|
|
def test_transformer_cache_threads_through(client, monkeypatch):
|
|
backend = _FakeBackend()
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"transformer_cache": "fbcache",
|
|
"transformer_cache_threshold": 0.1,
|
|
},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert backend.last_load_kwargs.get("transformer_cache") == "fbcache"
|
|
assert backend.last_load_kwargs.get("transformer_cache_threshold") == 0.1
|
|
|
|
|
|
def test_invalid_transformer_cache_returns_422(client):
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"transformer_cache": "deepcache",
|
|
},
|
|
)
|
|
assert resp.status_code == 422
|
|
|
|
|
|
def test_out_of_range_cache_threshold_returns_422(client):
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {
|
|
"model_path": "x/z-image",
|
|
"gguf_filename": "q.gguf",
|
|
"transformer_cache_threshold": 1.5,
|
|
},
|
|
)
|
|
assert resp.status_code == 422
|
|
|
|
|
|
def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path):
|
|
"""End-to-end through the REAL router: a CPU host with an available binary routes
|
|
the load to the native sd.cpp engine and the response reports engine=sd_cpp."""
|
|
from types import SimpleNamespace
|
|
|
|
import core.inference.diffusion_engine_router as engine_router
|
|
import core.inference.sd_cpp_backend as sd_backend
|
|
|
|
for e in (
|
|
"UNSLOTH_DIFFUSION_ENGINE",
|
|
"UNSLOTH_DIFFUSION_SD_CPP",
|
|
"UNSLOTH_DIFFUSION_SD_CPP_MPS",
|
|
"UNSLOTH_DIFFUSION_SD_CPP_INSTALL",
|
|
):
|
|
monkeypatch.delenv(e, raising = False)
|
|
|
|
validator = _FakeBackend() # supplies validate_load_request (and is the diffusers fallback)
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: validator)
|
|
# Force the router's decision inputs: CPU device + an available binary.
|
|
monkeypatch.setattr(
|
|
engine_router,
|
|
"resolve_diffusion_device_target",
|
|
lambda: SimpleNamespace(backend = "cpu", device = "cpu"),
|
|
)
|
|
monkeypatch.setattr(engine_router, "ensure_sd_cpp_binary", lambda **_: "/x/sd-cli")
|
|
# The router probes runnability before committing to native; treat the stub binary as executable.
|
|
monkeypatch.setattr(
|
|
engine_router, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: "sd-cli v0")
|
|
)
|
|
monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers")
|
|
monkeypatch.setattr(engine_router, "_fallback_reason", None)
|
|
# The native backend the router will activate.
|
|
sd_fake = _FakeBackend()
|
|
monkeypatch.setattr(sd_backend, "get_sd_cpp_backend", lambda: sd_fake)
|
|
|
|
monkeypatch.setattr(gpu_arbiter, "_owner", None)
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None)
|
|
monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None)
|
|
|
|
app = FastAPI()
|
|
app.include_router(studio_router, prefix = "/api/inference")
|
|
app.dependency_overrides[get_current_subject] = lambda: "test-user"
|
|
client = TestClient(app)
|
|
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "z.gguf"},
|
|
)
|
|
assert resp.status_code == 200
|
|
body = resp.json()
|
|
assert body["engine"] == "sd_cpp"
|
|
assert body["fallback_reason"] is None
|
|
assert sd_fake.loaded is True # the native engine actually received the load
|
|
|
|
|
|
def test_invalid_transformer_quant_returns_422_without_eviction(client):
|
|
# An unsupported transformer_quant is rejected by the request schema (Literal), so the GPU is
|
|
# never acquired and no chat model is evicted.
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "int2"},
|
|
)
|
|
assert resp.status_code == 422
|
|
assert gpu_arbiter._owner is None
|
|
|
|
|
|
def test_invalid_memory_mode_returns_422_without_eviction(client):
|
|
# An unsupported memory_mode is rejected by the request schema (Literal), so the GPU is never
|
|
# acquired and no chat model is evicted.
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "memory_mode": "ultra"},
|
|
)
|
|
assert resp.status_code == 422
|
|
assert gpu_arbiter._owner is None
|
|
|
|
|
|
def test_in_progress_returns_409_after_validation_passes(client, monkeypatch):
|
|
def _busy(*a, **k):
|
|
raise RuntimeError("A diffusion load is already in progress.")
|
|
|
|
backend = _FakeBackend()
|
|
backend.begin_load = _busy
|
|
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
|
|
# Pin the resolved device to cuda: the route only takes the arbiter for non-CPU loads, so on a
|
|
# CPU-only host the ownership assert below would never hold.
|
|
import types as _types
|
|
|
|
import core.inference.diffusion_device as devmod
|
|
|
|
monkeypatch.setattr(
|
|
devmod,
|
|
"resolve_diffusion_device_target",
|
|
lambda: _types.SimpleNamespace(device = "cuda"),
|
|
)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "q.gguf"},
|
|
)
|
|
assert resp.status_code == 409
|
|
# Validation passed first, so the GPU WAS acquired before begin_load reported busy.
|
|
assert gpu_arbiter._owner == gpu_arbiter.DIFFUSION
|
|
|
|
|
|
def _force_engine(monkeypatch, backend, *, engine_name, device):
|
|
"""Pin engine selection + device so the load route's arbiter gating is deterministic."""
|
|
import types as _types
|
|
|
|
import core.inference.diffusion_device as devmod
|
|
import core.inference.diffusion_engine_router as router
|
|
|
|
monkeypatch.setattr(router, "select_and_activate_engine", lambda fam, **kw: backend)
|
|
monkeypatch.setattr(router, "active_engine_name", lambda: engine_name)
|
|
monkeypatch.setattr(
|
|
devmod, "resolve_diffusion_device_target", lambda: _types.SimpleNamespace(device = device)
|
|
)
|
|
acquired: list = []
|
|
|
|
def _fake_acquire(role, register = None):
|
|
# Mirror the real arbiter: record the handoff and run the (registered) load under it.
|
|
acquired.append(role)
|
|
return register() if register is not None else None
|
|
|
|
monkeypatch.setattr(gpu_arbiter, "acquire_for", _fake_acquire)
|
|
return acquired
|
|
|
|
|
|
def test_cpu_native_load_skips_gpu_arbiter(client, monkeypatch):
|
|
# A native sd.cpp load on a pure-CPU host never touches the GPU, so the route must NOT evict the
|
|
# resident chat model: the arbiter handoff is skipped.
|
|
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
|
|
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
acquired = _force_engine(monkeypatch, backend, engine_name = ENGINE_SD_CPP, device = "cpu")
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 200
|
|
assert acquired == [] # no arbiter handoff for a CPU native load
|
|
|
|
|
|
def test_gpu_native_load_takes_arbiter(client, monkeypatch):
|
|
# A force-native sd.cpp load on a GPU box DOES use the GPU, so the arbiter is acquired, like the
|
|
# always-GPU diffusers path.
|
|
from core.inference.sd_cpp_engine import ENGINE_SD_CPP
|
|
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
acquired = _force_engine(monkeypatch, backend, engine_name = ENGINE_SD_CPP, device = "cuda")
|
|
resp = client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
assert resp.status_code == 200
|
|
assert acquired == [gpu_arbiter.DIFFUSION]
|
|
|
|
|
|
def test_images_info_lists_every_family(client):
|
|
# The pure info endpoint is hardware-independent (no load required): one entry per auto-policy
|
|
# family, each with the quant estimates the UI shows.
|
|
from core.inference.diffusion_auto_policy import _FAMILY_BF16_GB
|
|
|
|
resp = client.get("/api/inference/images/info")
|
|
assert resp.status_code == 200
|
|
families = resp.json()["families"]
|
|
assert {f["family"] for f in families} == set(_FAMILY_BF16_GB)
|
|
sample = families[0]
|
|
est = sample["estimated_resident_gb"]
|
|
# Quantised estimates undercut bf16, and nvfp4 undercuts int8 (matching the pure helper).
|
|
assert est["int8"] < est["bf16"]
|
|
assert est["nvfp4"] < est["int8"]
|
|
|
|
|
|
def test_status_passes_through_resolved(client, monkeypatch):
|
|
# The additive `resolved` provenance record round-trips through the status route so the frontend
|
|
# can render the "Auto: X" badges.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
resolved = {
|
|
"speed_mode": {"value": "eager", "source": "auto", "reason": "per-kind default"},
|
|
"transformer_quant": {"value": "int8", "source": "explicit", "reason": "requested"},
|
|
"cpu_offload": {"value": False, "source": "auto", "reason": "from the memory plan"},
|
|
"transformer_cache": {"value": None, "source": "auto", "reason": "few-step model"},
|
|
}
|
|
monkeypatch.setattr(
|
|
backend, "status", lambda: {**_unloaded_status(), "loaded": True, "resolved": resolved}
|
|
)
|
|
body = client.get("/api/inference/images/status").json()
|
|
assert body["resolved"] == resolved
|
|
assert body["resolved"]["speed_mode"]["source"] == "auto"
|
|
# The cpu_offload value stays a real boolean (not coerced to a string).
|
|
assert body["resolved"]["cpu_offload"]["value"] is False
|
|
|
|
|
|
def test_status_resolved_defaults_to_null(client):
|
|
# A backend status without a `resolved` key leaves the additive field null (older backends and the
|
|
# unloaded state).
|
|
body = client.get("/api/inference/images/status").json()
|
|
assert body["resolved"] is None
|
|
|
|
|
|
def test_download_plan_forwards_the_load_time_controls(client, monkeypatch):
|
|
# The plan drives the staged download, so it must be computed from the SAME configuration the load
|
|
# will run with. The dense-quant prefetch decision reads the memory policy, the prequant path and
|
|
# the adapter selection as well as speed/quant: dropping them stages the base transformer/ shards
|
|
# for a low-VRAM load that never opens them, or omits them for a baked-LoRA load that does.
|
|
backend = diffusion_module.get_diffusion_backend()
|
|
seen: dict = {}
|
|
|
|
def _plan(model_path, **kwargs):
|
|
seen["model_path"] = model_path
|
|
seen.update(kwargs)
|
|
return {"entries": [], "total_bytes": 0}
|
|
|
|
monkeypatch.setattr(backend, "download_plan", _plan, raising = False)
|
|
|
|
resp = client.post(
|
|
"/api/inference/images/download-plan",
|
|
json = {
|
|
"model_path": "unsloth/FLUX.1-dev-GGUF",
|
|
"gguf_filename": "flux1-dev-Q4_K_M.gguf",
|
|
"model_kind": "gguf",
|
|
"hf_token": "hf_secret",
|
|
"speed_mode": "off",
|
|
"transformer_quant": "int8",
|
|
"memory_mode": "low_vram",
|
|
"cpu_offload": True,
|
|
"loras": [{"id": "unsloth/some-lora", "weight": 0.8}],
|
|
},
|
|
)
|
|
|
|
assert resp.status_code == 200
|
|
assert seen["hf_token"] == "hf_secret"
|
|
assert seen["speed_mode"] == "off"
|
|
assert seen["transformer_quant"] == "int8"
|
|
assert seen["memory_mode"] == "low_vram"
|
|
assert seen["cpu_offload"] is True
|
|
assert len(seen["loras"] or []) == 1
|
|
|
|
|
|
def test_load_refused_when_only_the_diffusion_probe_can_be_read(client, monkeypatch):
|
|
# The two training probes are independent. An LLM backend that raises used to short-circuit the
|
|
# guard entirely, so an image load sailed past a KNOWN-active diffusion trainer and contended
|
|
# with it for VRAM. An unreadable LLM state must not disable the diffusion interlock.
|
|
import core.training as core_training
|
|
import routes.inference as inference_routes
|
|
|
|
class _Broken:
|
|
def is_training_active(self):
|
|
raise RuntimeError("training backend unavailable")
|
|
|
|
monkeypatch.setattr(core_training, "get_training_backend", lambda: _Broken())
|
|
monkeypatch.setattr(inference_routes, "_diffusion_training_active", lambda: True)
|
|
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"},
|
|
)
|
|
assert resp.status_code == 409
|
|
assert "training" in resp.json()["detail"].lower()
|
|
|
|
# With neither trainer active the unreadable LLM probe still must not block the load.
|
|
monkeypatch.setattr(inference_routes, "_diffusion_training_active", lambda: False)
|
|
resp = client.post(
|
|
"/api/inference/images/load",
|
|
json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"},
|
|
)
|
|
assert resp.status_code == 200
|
|
|
|
|
|
def test_recipe_records_the_conditioned_workflow_settings(client, monkeypatch):
|
|
# A conditioned generation's recipe used to carry only the txt2img fields, so the gallery
|
|
# presented an inpaint result as a complete Create recipe and restoring it replayed an
|
|
# unrelated text-to-image request. The images are still not persisted (user uploads with their
|
|
# own lifetime), but what ran IS, so the client can say which inputs to supply again.
|
|
import base64
|
|
import io
|
|
|
|
from PIL import Image
|
|
|
|
client.post(
|
|
"/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"}
|
|
)
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (8, 8), (120, 30, 90)).save(buf, format = "PNG")
|
|
px = base64.b64encode(buf.getvalue()).decode()
|
|
gen = client.post(
|
|
"/api/inference/images/generate",
|
|
json = {
|
|
"prompt": "a sloth",
|
|
"seed": 7,
|
|
"init_image": px,
|
|
"mask_image": px,
|
|
"strength": 0.42,
|
|
},
|
|
)
|
|
assert gen.status_code == 200
|
|
img = gen.json()["images"][0]
|
|
assert img["workflow"] == "inpaint"
|
|
assert img["strength"] == 0.42
|
|
# A plain txt2img still records its own workflow and leaves the conditioning fields empty.
|
|
plain = client.post("/api/inference/images/generate", json = {"prompt": "a sloth", "seed": 7})
|
|
assert plain.status_code == 200
|
|
plain_img = plain.json()["images"][0]
|
|
assert plain_img["workflow"] == "txt2img"
|
|
assert plain_img["strength"] is None and plain_img["upscale"] is None
|