diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index db36627b78..4d0bbcbfa7 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -336,11 +336,25 @@ class SdCppDiffusionBackend: # than commit a "ready" state that crashes on the first generation. if engine.version() is None: raise RuntimeError("sd-cli binary is present but not runnable.") + # A generation that started during the (slow) asset download is still running + # against the OLD model. Abort it, then WAIT on _generate_lock for it to exit + # before publishing the new state -- otherwise that stale sd-cli run can finish + # afterward and persist an image from the previous model once this load reports + # ready (mirrors the diffusers load_pipeline commit). _generate_lock is taken + # only here, not during the download, so the long fetch never serialises against + # generation; the inner token re-check guards an unload/newer load arriving while + # we waited. with self._lock: if self._load_token != _load_token: return # superseded / cancelled - self._state = state - self._loading = None + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + with self._generate_lock: + with self._lock: + if self._load_token != _load_token: + return # superseded / cancelled while waiting + self._state = state + self._loading = None except SdCppCancelled: return except Exception as exc: # noqa: BLE001 -- surfaced via load_progress @@ -596,8 +610,12 @@ class SdCppDiffusionBackend: "base_repo": state.base_repo, "device": state.device, "dtype": "gguf", - "cpu_offload": False, - "offload_policy": "none", + # Reflect the offload flags actually passed to sd-cli, so a balanced/low_vram + # (or cpu_offload) load is verifiable from status instead of always reading + # "none". On CPU _run_load leaves offload_flags empty (the flags are no-ops), + # so this correctly stays "none" there. + "cpu_offload": bool(state.offload_flags), + "offload_policy": "active" if state.offload_flags else "none", "vae_tiling": False, "memory_mode": None, "speed_mode": state.native_speed, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 15c46d4132..496a406574 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10058,8 +10058,14 @@ async def load_diffusion_model( request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) ): from core.inference.diffusion import get_diffusion_backend - from core.inference.diffusion_engine_router import annotate_status, select_and_activate_engine + from core.inference.diffusion_device import resolve_diffusion_device_target + from core.inference.diffusion_engine_router import ( + active_engine_name, + annotate_status, + select_and_activate_engine, + ) from core.inference.gpu_arbiter import acquire_for, DIFFUSION + from core.inference.sd_cpp_engine import ENGINE_SD_CPP from utils.native_path_leases import redact_native_paths backend = get_diffusion_backend() @@ -10077,9 +10083,16 @@ async def load_diffusion_model( # installing the sd-cli binary if needed -- all BEFORE evicting chat, so a # native fallback never strands a half-loaded state. engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token) - # Now take the GPU from the chat backend, then kick the (slow) load onto a - # background thread and return at once — the client polls images/load-progress. - await asyncio.to_thread(acquire_for, DIFFUSION) + # Take the GPU from the chat backend only when this load will actually use it. + # diffusers always does; a *force-native* sd.cpp load on a CUDA/XPU/MPS box does + # too. But a native sd.cpp load on a pure-CPU host never touches the GPU, so + # acquiring would evict the resident chat model for nothing -- skip the handoff. + device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) + needs_gpu = active_engine_name() != ENGINE_SD_CPP or device != "cpu" + if needs_gpu: + # Then kick the (slow) load onto a background thread and return at once -- + # the client polls images/load-progress. + await asyncio.to_thread(acquire_for, DIFFUSION) status_dict = await asyncio.to_thread( engine.begin_load, request.model_path, diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d47ab6a876..4fc9d66c34 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -535,3 +535,48 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch): 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 = [] + monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + 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 + # (same as 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] diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index c6fbf7243e..79464dd6be 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -6,6 +6,7 @@ from __future__ import annotations import threading +import types import pytest from PIL import Image @@ -236,3 +237,74 @@ def test_unload_clears_state_and_signals_cancel(): assert st["loaded"] is False assert cancel.is_set() assert b._cancel_event.is_set() + + +def test_status_reports_offload_when_flags_active(): + # status must reflect the offload flags actually passed to sd-cli, not always "none", + # so a balanced/low_vram (or cpu_offload) load is verifiable. + b = _loaded_backend() + # No flags (CPU default) -> none. + assert b.status()["offload_policy"] == "none" and b.status()["cpu_offload"] is False + # Flags present (off-CPU offload) -> reported active. + s = b._state + b._state = bk._SdState( + repo_id = s.repo_id, + base_repo = s.base_repo, + family = s.family, + device = "cuda", + files = s.files, + offload_flags = ("--vae-on-cpu", "--clip-on-cpu"), + ) + st = b.status() + assert st["cpu_offload"] is True and st["offload_policy"] == "active" + + +def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): + # A generation that started during the asset download is still running against the OLD + # model. _run_load must cancel it AND wait on _generate_lock before committing the new + # state, or a stale sd-cli run finishes afterward and persists an image from the previous + # model once the new load reports ready. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family("z-image") + monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) + monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None) + monkeypatch.setattr( + b, + "_fetch_assets", + lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, + ) + # Avoid importing torch from the worker thread (its first import deadlocks off the main + # thread -- a test artifact, not a production path); the device only needs to be CPU here. + monkeypatch.setattr( + bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") + ) + + b._load_token = 5 + cancel = threading.Event() + b._active_generate_cancel = cancel # a generation is "in flight" + + committed = threading.Event() + + def _load(): + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 5, + ) + committed.set() + + b._generate_lock.acquire() # simulate the live denoise holding _generate_lock + try: + threading.Thread(target = _load, daemon = True).start() + # The commit must block behind the live generation and not publish the new state, + # but must already have signalled the in-flight cancel. + assert not committed.wait(0.5) + assert b._state is None + assert cancel.is_set() + finally: + b._generate_lock.release() + assert committed.wait(5) # only now does the commit run + assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF"