From d0c5cf6e073848caa646b8a633e04c42c0793fae Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:58:39 -0300 Subject: [PATCH] Fix diffusion flag leak, sd-cli orphan, and native family fallback Six correctness fixes to the diffusion stack, found reviewing the merged phase PRs on this branch: - load_pipeline: restore the try/finally guard around the speed/quant/ placement span. A failure after apply_speed_optims (e.g. OOM in quant or the memory plan) left TF32/cudnn flags flipped process-wide and the half-built pipe resident in VRAM. Now restores the flags and frees VRAM on a failed load. - sd-cli Popen binds to the parent (PR_SET_PDEATHSIG via child_popen_kwargs, matching the llama.cpp sites), so a parent crash mid-generation can't orphan it holding VRAM/RAM. - Native begin_load uses the filename-fallback family detector the route validated with, so a local .gguf whose family keyword lives only in the basename no longer dead-ends 400 on a no-GPU host. - Generate error handler matches exact sentinel messages instead of the "cancelled" substring, fixing a 409 misroute and a raw sd-cli output leak. - find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME/STUDIO_HOME like the installer, so a custom Studio home resolves. - Drop the redundant _tf32_prev bookkeeping; snapshot/restore_backend_flags is now the single owner of the TF32/cudnn restore. The two client-state messages are now shared constants so the 409-vs-500 contract can't drift. Adds a regression test for each behavioral fix. --- studio/backend/core/inference/diffusion.py | 151 +++++++++--------- .../core/inference/diffusion_families.py | 24 +++ .../backend/core/inference/diffusion_speed.py | 44 +---- .../backend/core/inference/sd_cpp_backend.py | 17 +- .../backend/core/inference/sd_cpp_engine.py | 16 +- studio/backend/routes/inference.py | 16 +- .../backend/tests/test_diffusion_backend.py | 32 ++++ studio/backend/tests/test_diffusion_routes.py | 31 ++++ studio/backend/tests/test_sd_cpp_backend.py | 11 ++ 9 files changed, 215 insertions(+), 127 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index cba3ba2af5..6d728df035 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -24,8 +24,10 @@ from loggers import get_logger from utils.hardware import clear_gpu_cache from .diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, DiffusionFamily, - detect_family, + detect_family_for_pick, resolve_base_repo, resolve_local_gguf_child, ) @@ -245,22 +247,6 @@ class DiffusionBackend: base, rfilename, hf_token, cancel_event = self._cancel_event ) - @staticmethod - def _detect_family_for_pick( - repo_id: str, gguf_filename: Optional[str], family_override: Optional[str] - ) -> Optional[DiffusionFamily]: - """Detect the family from the repo id, falling back to the combined - path/filename for a direct local .gguf pick. The frontend splits such a - pick into (parent dir, basename), so the family keyword can live only in - the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent - directory carries none; scan it too when the directory alone is - undetectable. Only used as a fallback, so remote 'org/name' picks and - explicit overrides behave exactly as before.""" - fam = detect_family(repo_id, family_override) - if fam is None and gguf_filename and not family_override: - fam = detect_family(f"{repo_id}/{gguf_filename}", family_override) - return fam - def validate_load_request( self, repo_id: str, @@ -277,7 +263,7 @@ class DiffusionBackend: raise ValueError( "gguf_filename is required: this backend loads single-file GGUF checkpoints only." ) - fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override) + fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError( f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." @@ -368,7 +354,7 @@ class DiffusionBackend: # Resolve the base repo and estimate sizes on this thread (both network # calls) so begin_load returns instantly; the bar shows raw bytes until # the total lands. This is the only writer of _loading's fields here. - fam = self._detect_family_for_pick( + fam = detect_family_for_pick( kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override") ) base = _resolve_base_repo( @@ -641,63 +627,78 @@ class DiffusionBackend: quant_active = transformer_quant_engaged is not None or bool(gguf_filename), logger = logger, ) - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = effective_speed, - cache_active = cache_engaged is not None, - logger = logger, - ) - if transformer_quant_engaged is not None and not speed_applied.get("compiled"): - # Promotion above could not engage compile (e.g. the family is not - # compile-friendly, or compile_repeated_blocks failed): the quantized - # transformer is now running eager, which is far slower than the GGUF - # path it replaced. Surface it loudly rather than hiding the regression. - logger.warning( - "diffusion.transformer_quant: %s engaged but the transformer is NOT " - "compiled; eager torchao quant is ~30x slower than GGUF here", - transformer_quant_engaged, + # apply_speed_optims flips the process-global TF32 / cudnn.benchmark + # flags. If a later step here (text-encoder quant, memory plan) then + # raises -- e.g. OOM -- those flags would leak flipped and a subsequent + # `off` load would no longer be bit-identical. Restore the snapshot + # unless we reach the commit (unload restores on the happy path). + committed = False + try: + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + cache_active = cache_engaged is not None, + logger = logger, + ) + if transformer_quant_engaged is not None and not speed_applied.get("compiled"): + # Promotion above could not engage compile (e.g. the family is not + # compile-friendly, or compile_repeated_blocks failed): the quantized + # transformer is now running eager, which is far slower than the GGUF + # path it replaced. Surface it loudly rather than hiding the regression. + logger.warning( + "diffusion.transformer_quant: %s engaged but the transformer is NOT " + "compiled; eager torchao quant is ~30x slower than GGUF here", + transformer_quant_engaged, + ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, - ) - # Apply the placement planned above (from MEASURED free device memory vs - # the model's estimated resident size). apply_memory_plan returns the - # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module - # offload, and tiling is a no-op on a pipeline with no tiling control), so - # status stays honest. The dense fast path already placed the pipe resident; - # for the `none` policy this is an idempotent re-placement. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Apply the placement planned above (from MEASURED free device memory vs + # the model's estimated resident size). apply_memory_plan returns the + # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module + # offload, and tiling is a no-op on a pipeline with no tiling control), so + # status stays honest. The dense fast path already placed the pipe resident; + # for the `none` policy this is an idempotent re-placement. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - speed_optims = tuple(k for k, v in speed_applied.items() if v), - backend_flags_before = backend_flags_before, - text_encoder_quant = te_quant, - transformer_quant = transformer_quant_engaged, - attention_backend = attention_engaged, - transformer_cache = cache_engaged, - ) + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + transformer_quant = transformer_quant_engaged, + attention_backend = attention_engaged, + transformer_cache = cache_engaged, + ) + committed = True + finally: + if not committed: + # Restore the flags AND free the half-built pipe's VRAM: the + # failed load never commits _state, so nothing else reclaims it + # until the next unload. + restore_backend_flags(backend_flags_before) + clear_gpu_cache() logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -854,7 +855,7 @@ class DiffusionBackend: with self._lock: state = self._state if state is None: - raise RuntimeError("No diffusion model is loaded.") + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. # A cancel that arrived before now either nulled _state (we raised # above) or targets an older generation, so nothing is lost. @@ -923,7 +924,7 @@ class DiffusionBackend: # A cancelled denoise returns early with a partial/garbage image; # don't hand it back to be persisted. if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Return the PIL images (not yet encoded): the route embeds each # image's recipe and persists it via the gallery. return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 3afbe9b495..3ac238e399 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -20,6 +20,14 @@ from pathlib import Path, PurePosixPath from typing import Optional +# Runtime->route contract: the RuntimeError messages a backend raises for +# client-recoverable generate states. The /images/generate route matches these +# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines +# must raise them verbatim -- keep them named here, not as scattered literals. +DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded." +DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled." + + @dataclass(frozen = True) class DiffusionFamily: name: str @@ -173,6 +181,22 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff return None +def detect_family_for_pick( + repo_id: str, gguf_filename: Optional[str] = None, override: Optional[str] = None +) -> Optional[DiffusionFamily]: + """``detect_family``, falling back to the combined path/filename for a direct + local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename), + so the family keyword can live only in the filename (e.g. + ``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none; + scan the combined string too when the directory alone is undetectable. Only a + fallback, so remote ``org/name`` picks and explicit overrides behave exactly as + ``detect_family``. Shared by both engines so validation and load can't diverge.""" + fam = detect_family(repo_id, override) + if fam is None and gguf_filename and not override: + fam = detect_family(f"{repo_id}/{gguf_filename}", override) + return fam + + def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: """The companion diffusers repo: caller-supplied if given, else the family fallback.""" base = (base_repo or "").strip() diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 208cb8b513..5f1ef64445 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -161,12 +161,11 @@ def apply_speed_optims( "compiled": False, } mode = normalize_speed_mode(speed_mode) - # TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this - # load isn't max, so a later default/off diffusion load -- or chat inference in the - # same long-lived process -- doesn't silently inherit a prior max load's TF32 and - # lose the bit-identical default the regression harness checks. - if mode != SPEED_MAX: - _restore_tf32(logger) + # TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max, + # cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this + # call and restores on unload / failed load via snapshot_backend_flags / + # restore_backend_flags, so a later `off` load -- or chat inference in the same + # process -- never inherits them. We keep no separate bookkeeping here. if mode == SPEED_OFF: return applied @@ -251,22 +250,9 @@ def _enable_cudnn_benchmark(logger: Any) -> bool: return False -# The TF32 flag values from before the first max load flipped them, so a later -# non-max load / unload can put the process back exactly as it found it (rather than -# forcing a hardcoded default that might clobber another component's choice). -_tf32_prev: Optional[tuple[bool, bool]] = None - - def _enable_tf32(logger: Any) -> bool: - global _tf32_prev try: import torch - - if _tf32_prev is None: - _tf32_prev = ( - torch.backends.cuda.matmul.allow_tf32, - torch.backends.cudnn.allow_tf32, - ) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True return True @@ -275,26 +261,6 @@ def _enable_tf32(logger: Any) -> bool: return False -def restore_tf32(logger: Any = None) -> None: - """Put the process-global TF32 flags back to their pre-max-load values. No-op if - a max load never set them. Called on a non-max load and on unload.""" - _restore_tf32(logger) - - -def _restore_tf32(logger: Any) -> None: - global _tf32_prev - if _tf32_prev is None: - return - try: - import torch - torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0] - torch.backends.cudnn.allow_tf32 = _tf32_prev[1] - except Exception as exc: # noqa: BLE001 — best-effort restore - _warn(logger, "tf32_restore", exc) - finally: - _tf32_prev = None - - def _fuse_qkv(pipe: Any, logger: Any) -> bool: for owner in (pipe, getattr(pipe, "transformer", None)): fn = getattr(owner, "fuse_qkv_projections", None) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 4d0bbcbfa7..88661f2be1 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -35,8 +35,10 @@ from typing import Any, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, DiffusionFamily, - detect_family, + detect_family_for_pick, family_sd_cpp_supported, resolve_base_repo, resolve_local_gguf_child, @@ -242,7 +244,10 @@ class SdCppDiffusionBackend: raise ValueError( "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." ) - fam = detect_family(repo_id, family_override) + # Use the filename-fallback detector the route validated with, so a local + # .gguf pick whose family keyword lives only in the basename doesn't pass + # validation and then dead-end here on a no-GPU (native-routed) host. + fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError(f"Could not infer a diffusion family for '{repo_id}'.") if not family_sd_cpp_supported(fam): @@ -464,7 +469,7 @@ class SdCppDiffusionBackend: with self._lock: state = self._state if state is None: - raise RuntimeError("No diffusion model is loaded.") + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel engine = self._resolve_engine() try: @@ -485,7 +490,7 @@ class SdCppDiffusionBackend: with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: for index in range(max(1, int(batch_size))): if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Distinct seed per batch image (sd-cli is one image/run here), # so a batch is reproducible image-by-image from the base seed. # Mask to sd-cli's int64 range, NOT 53 bits: the request model and @@ -522,7 +527,7 @@ class SdCppDiffusionBackend: images.append(im.copy()) seeds.append(seed_i) if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # ``seeds`` is the per-image seed (each sd-cli run used seed+index), so # the route can persist the real seed for every image in the batch. return { @@ -532,7 +537,7 @@ class SdCppDiffusionBackend: "repo_id": state.repo_id, } except SdCppCancelled as exc: - raise RuntimeError("Diffusion generation was cancelled.") from exc + raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc finally: self._gen = None with self._lock: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 6714151045..557f6fcb3e 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -34,6 +34,7 @@ import time from pathlib import Path from typing import Callable, Optional +from utils.process_lifetime import child_popen_kwargs from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, @@ -125,7 +126,8 @@ def find_sd_cpp_binary() -> Optional[str]: both engines look): 1. ``SD_CLI_PATH`` env -- a direct path to the binary. 2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir. - 3. ``~/.unsloth/stable-diffusion.cpp`` build layouts (the installer target). + 3. the installer target: ``/../stable-diffusion.cpp`` when + that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``. 4. ``./stable-diffusion.cpp`` in-tree build (developer checkout). 5. ``sd-cli`` (then legacy ``sd``) on PATH. """ @@ -151,8 +153,12 @@ def find_sd_cpp_binary() -> Optional[str]: if hit: return hit - # 3. Default install root (sibling of ~/.unsloth/llama.cpp). - hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp")) + # 3. Default install root: the installer's default_install_dir() -- a sibling of + # the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else + # ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves. + studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") + default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth" + hit = _first_file(_layout_candidates(default_base / "stable-diffusion.cpp")) if hit: return hit @@ -342,6 +348,10 @@ class SdCppEngine: # Own session/process group so cancellation/timeout can kill the whole # tree, not just the parent (POSIX only; harmless flag elsewhere). start_new_session = (os.name == "posix"), + # Bind the child to the parent's lifetime (Linux PR_SET_PDEATHSIG), so a + # hard parent crash mid-generation can't orphan sd-cli holding VRAM/RAM -- + # matching every llama.cpp Popen site. Composes with start_new_session. + **child_popen_kwargs(), ) # Drain stdout on a reader thread so the timeout is enforced even when the # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 400dc239fc..e20a2f1a49 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10370,6 +10370,10 @@ async def generate_diffusion_image( ): from core.inference import image_gallery from core.inference.diffusion_engine_router import get_active_diffusion_engine + from core.inference.diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, + ) backend = get_active_diffusion_engine() try: @@ -10385,11 +10389,15 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - # Only "no model loaded" / cancelled are client-state (409). The native - # sd.cpp engine also raises RuntimeError for execution failures (nonzero - # exit, timeout, missing output), which are server errors (500). + # Only "no model loaded" / user-cancelled are client-state (409); both engines + # raise these two EXACT messages. The native sd.cpp engine also raises + # RuntimeError for execution failures (nonzero exit, timeout, missing output) + # whose text can embed the raw sd-cli tail (local paths / argv) -- those are + # server errors (500) returned as a fixed literal, never echoed. Match the + # sentinels exactly, not as a substring, so an sd-cli failure that merely + # contains "cancelled" can't misroute to 409 and leak that output. msg = str(exc) - if "No diffusion model is loaded" in msg or "cancelled" in msg.lower(): + if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG): raise HTTPException(status_code = 409, detail = msg) logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f8d0172173..5ec8bb3954 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -307,6 +307,38 @@ def test_generate_without_load_raises(fake_runtime): backend.generate(prompt = "x") +def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch): + # A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go + # through the load's try/finally and restore the process-global TF32 / cudnn flags, + # so a later `off` load is still bit-identical, and must not commit a partial state. + # Regression: a refactor dropped this guard, leaking the flags on a failed load. + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + + restored: list = [] + cleared: list = [] + monkeypatch.setattr( + "core.inference.diffusion.restore_backend_flags", lambda snap: restored.append(snap) + ) + monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: cleared.append(True)) + monkeypatch.setattr( + "core.inference.diffusion.apply_memory_plan", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("CUDA out of memory")), + ) + + with pytest.raises(RuntimeError, match = "out of memory"): + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + family_override = "z-image", + base_repo = "base/repo", + speed_mode = "max", + ) + assert restored, "restore_backend_flags was not called on the failed-load path" + assert cleared, "clear_gpu_cache was not called on the failed-load path (VRAM leak)" + assert backend._state is None and backend.is_loaded is False + + def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch): from core.inference import diffusion from core.inference.diffusion_families import detect_family diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d0b6997d07..a3c78a21aa 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -281,6 +281,37 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch): assert "CUDA" not in resp.json()["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). Regression: the handler matched "cancelled" as a substring. + 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 + assert resp.json()["detail"] == "Image generation failed." + assert "cancelled" not in resp.json()["detail"] and "models" not in resp.json()["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.") diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 79464dd6be..84e1b37dd2 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -219,6 +219,17 @@ def test_begin_load_requires_gguf_filename(): b.begin_load("unsloth/Z-Image-Turbo-GGUF") +def test_begin_load_resolves_family_from_filename_only(monkeypatch): + # A local .gguf pick whose family keyword lives only in the basename (parent dir + # carries none) must resolve via the same filename fallback the route validated + # with -- not dead-end with "Could not infer" on a native (no-GPU) host. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + monkeypatch.setattr(b, "_run_load", lambda **kwargs: None) # skip the download thread + b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf") + # Validation passed (no ValueError) and the family was inferred from the filename. + assert b._loading is not None and b._loading.repo_id == "/models/gguf-store" + + def test_ensure_binary_returns_found(monkeypatch): monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli") assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli"