diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index 0c874e4b66..2d2cc11343 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -59,8 +59,24 @@ def resolve_diffusion_device_target() -> DiffusionDeviceTarget: (CUDA -> XPU -> MPS -> CPU). On Apple Silicon Studio reports MLX/CPU when its product backend is gated on the ``mlx`` package, but diffusers runs on PyTorch's MPS backend, so those cases still fall through to the MPS probe. + + Torch is optional here: on a CPU-only install without PyTorch the native + stable-diffusion.cpp engine still runs (it shells out to sd-cli), so a missing + torch reports a torch-free CPU target instead of crashing the whole + ``/images/load`` before the engine router can select the native backend. """ - import torch + try: + import torch + except Exception: + return DiffusionDeviceTarget( + device = "cpu", + dtype = None, + backend = "cpu", + vendor = None, + supports_model_cpu_offload = False, + supports_default_torch_compile = False, + supports_pinned_transfer = False, + ) try: from utils.hardware import DeviceType, get_device diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 96407a66b5..c8db7f81d2 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -30,7 +30,12 @@ from typing import Any, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported from core.inference.sd_cpp_backend import _install_allowed, ensure_sd_cpp_binary -from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP, select_diffusion_engine +from core.inference.sd_cpp_engine import ( + ENGINE_DIFFUSERS, + ENGINE_SD_CPP, + SdCppEngine, + select_diffusion_engine, +) from loggers import get_logger logger = get_logger(__name__) @@ -68,17 +73,25 @@ def active_engine_name() -> str: def _activate(name: str, reason: Optional[str]) -> Any: global _active_engine_name, _fallback_reason + # Switching engines: unload the one being deactivated first, or its model + # stays resident but unreachable (the arbiter evictor only targets the active + # engine), leaking 10+ GB and defeating the chat<->diffusion handoff. The unload + # itself (freeing 10+ GB / syncing CUDA) is slow, so resolve the engine under the + # lock but run unload() OUTSIDE it -- holding _lock across a slow unload would + # block every other selection caller. + engine_to_unload = None + old_name = None with _lock: - # Switching engines: unload the one being deactivated first, or its model - # stays resident but unreachable (the arbiter evictor only targets the active - # engine), leaking 10+ GB and defeating the chat<->diffusion handoff. if name != _active_engine_name: - try: - get_active_diffusion_engine().unload() - except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch - logger.warning("failed to unload previous engine %s: %s", _active_engine_name, exc) + engine_to_unload = get_active_diffusion_engine() + old_name = _active_engine_name _active_engine_name = name _fallback_reason = reason if name == ENGINE_DIFFUSERS else None + if engine_to_unload is not None: + try: + engine_to_unload.unload() + except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch + logger.warning("failed to unload previous engine %s: %s", old_name, exc) if name == ENGINE_SD_CPP: logger.info("diffusion engine: sd_cpp") else: @@ -113,6 +126,13 @@ def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] binary = None if policy_eligible and fam_ok: binary = ensure_sd_cpp_binary(allow_install = _install_allowed()) + # Probe runnability here, before committing the route to native: a present but + # non-runnable binary (wrong arch, missing shared libs, no execute bit) would + # otherwise pass as available and only fail inside the background load, instead + # of falling back to diffusers now. + if binary and SdCppEngine(binary = binary).version() is None: + logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary) + binary = None native_available = bool(binary) and policy_eligible and fam_ok choice = select_diffusion_engine( diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index f82335e385..db36627b78 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -234,6 +234,10 @@ class SdCppDiffusionBackend: transformer_cache_threshold: Optional[float] = None, ) -> dict[str, Any]: """Validate, then fetch assets on a daemon thread. Returns at once.""" + # An empty / whitespace token is "no token": passing "" verbatim to HfApi / + # hf_hub_download is treated as an explicit (invalid) credential and breaks the + # anonymous fallback for public repos. + hf_token = hf_token.strip() if hf_token and hf_token.strip() else None if not gguf_filename: raise ValueError( "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." @@ -248,6 +252,11 @@ class SdCppDiffusionBackend: with self._lock: if self._loading is not None and self._loading.error is None: raise RuntimeError("A diffusion load is already in progress.") + # A superseding load must stop any in-flight generation, or the old sd-cli + # keeps running against the previous model and can still return / persist an + # image after the new load has started (matches unload()'s cancel). + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() self._load_token += 1 token = self._load_token self._cancel_event.clear() @@ -465,7 +474,12 @@ class SdCppDiffusionBackend: raise RuntimeError("Diffusion generation was cancelled.") # Distinct seed per batch image (sd-cli is one image/run here), # so a batch is reproducible image-by-image from the base seed. - seed_i = (seed + index) & ((1 << 53) - 1) + # Mask to sd-cli's int64 range, NOT 53 bits: the request model and + # the diffusers backend both accept large explicit seeds, so a tight + # 2**53 mask would silently truncate them (2**53 -> 0) and collide + # distinct requested seeds onto the same image. Randomly-drawn seeds + # above are already 53-bit (JS-safe); explicit seeds pass through. + seed_i = (seed + index) & ((1 << 63) - 1) out_path = str(Path(tmpdir) / f"img_{index}.png") params = SdCppGenParams( prompt = prompt, diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py index 2fe4b197bc..07d495757b 100644 --- a/studio/backend/tests/test_diffusion_engine_router.py +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -47,6 +47,12 @@ def _set_binary(monkeypatch, path): monkeypatch.setattr(r, "ensure_sd_cpp_binary", lambda **_: path) +def _set_runnable(monkeypatch, version = "sd-cli v0"): + """Stub the runnability probe so a stubbed binary path is treated as executable + (the router now probes ``SdCppEngine(...).version()`` before committing to native).""" + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: version)) + + def _select(fam_name = "z-image"): """Activate the engine for a family and return which engine was chosen.""" r.select_and_activate_engine(detect_family(fam_name)) @@ -59,10 +65,21 @@ def _select(fam_name = "z-image"): def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch): _set_device(monkeypatch, "cpu") _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) assert _select() == ENGINE_SD_CPP assert r.active_engine_name() == ENGINE_SD_CPP +def test_present_but_not_runnable_binary_falls_back(monkeypatch): + # A binary that exists but cannot run (version() -> None) must fall back to + # diffusers at selection, not commit native and fail inside the load. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None)) + assert _select() == ENGINE_DIFFUSERS + assert "binary unavailable" in (r.active_status()["fallback_reason"] or "") + + @pytest.mark.parametrize("gpu", ["cuda", "rocm", "xpu"]) def test_gpu_backends_use_diffusers(monkeypatch, gpu): _set_device(monkeypatch, gpu) @@ -90,6 +107,7 @@ def test_sd_cpp_disabled_uses_diffusers(monkeypatch): def test_mps_default_diffusers_but_optin_sd_cpp(monkeypatch): _set_device(monkeypatch, "mps") _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) # Default: MPS is not native-eligible -> diffusers. assert _select() == ENGINE_DIFFUSERS # Opt in: MPS routes to sd.cpp. @@ -115,6 +133,7 @@ def test_missing_binary_falls_back(monkeypatch): def test_force_sd_cpp_on_gpu_when_binary_present(monkeypatch): _set_device(monkeypatch, "cuda") _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp") assert _select() == ENGINE_SD_CPP diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 0569cf1385..d47ab6a876 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -468,6 +468,11 @@ def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path): lambda: SimpleNamespace(backend = "cpu", device = "cpu"), ) monkeypatch.setattr(engine_router, "ensure_sd_cpp_binary", lambda **_: "/x/sd-cli") + # The router now 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.