Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race

Codex review on the native-engine routing:

- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
  unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
  never touches the GPU, so that needlessly tore down the resident chat model. The
  handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
  CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.

- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
  _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
  the setting was unverifiable. status now derives them from state.offload_flags
  (still 'none' on CPU, where the flags are empty).

- _run_load committed the new state without cancelling/waiting on a generation that
  started during the (slow) asset download, so a stale sd-cli run against the OLD
  model could finish afterward and persist an image from the previous model once the
  new load reported ready. The commit now signals the in-flight cancel and waits on
  _generate_lock before swapping _state (taken only at commit, so the download never
  serialises against generation), mirroring the diffusers load path.

Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.

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,

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,

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]

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"