From 9a2cc341c5d3da1c4e6a310652ad8c3bbae27242 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 28 Jun 2026 05:43:45 +0000 Subject: [PATCH] Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler Address review feedback on #6724: - engine router: unload the engine being deactivated on a switch, so the old model is not left resident-but-unreachable (the evictor only targets the active engine). - generate route: sd.cpp execution errors (nonzero exit / timeout / missing output) now map to 500, not 409 (which only means not-loaded / cancelled). - native batch: return per-image seeds and persist the actual seed for each image so every batch image is reproducible. - Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the stable-diffusion.cpp docs; other families keep sd-cli defaults. - honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload offload flags on the native load instead of hardcoding them off. - fail the load when the sd-cli binary is present but not runnable (version() now returns None on exec error / nonzero exit). - size estimate: only treat the transformer asset as a possible local path. --- .../core/inference/diffusion_engine_router.py | 8 ++ .../core/inference/diffusion_families.py | 9 +++ .../backend/core/inference/sd_cpp_backend.py | 74 ++++++++++++++++--- .../backend/core/inference/sd_cpp_engine.py | 13 +++- studio/backend/routes/inference.py | 24 ++++-- studio/backend/tests/test_sd_cpp_backend.py | 13 ++++ studio/backend/tests/test_sd_cpp_engine.py | 4 +- 7 files changed, 122 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py index 60aa14bc9e..96407a66b5 100644 --- a/studio/backend/core/inference/diffusion_engine_router.py +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -69,6 +69,14 @@ def active_engine_name() -> str: def _activate(name: str, reason: Optional[str]) -> Any: global _active_engine_name, _fallback_reason 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) _active_engine_name = name _fallback_reason = reason if name == ENGINE_DIFFUSERS else None if name == ENGINE_SD_CPP: diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index b1f3be98a4..df3246e2b0 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -58,6 +58,12 @@ class DiffusionFamily: # autoencoder, None (auto) otherwise. sd_cpp_vae_format: Optional[str] = None sd_cpp_text_encoders: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) + # Family-specific sd-cli sampler settings, applied on the native path so the + # output matches the model's supported invocation (e.g. Qwen-Image needs + # --sampling-method euler --flow-shift 3 per stable-diffusion.cpp's docs). None + # leaves sd-cli's defaults (correct for the distilled flux/z-image families). + sd_cpp_sampling_method: Optional[str] = None + sd_cpp_flow_shift: Optional[float] = None # Keyed by architecture, not per model variant: a checkpoint's specific base repo @@ -114,6 +120,9 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( "qwen2vl", ), ), + # Qwen-Image's supported sd.cpp invocation (docs/qwen_image.md). + sd_cpp_sampling_method = "euler", + sd_cpp_flow_shift = 3.0, ), DiffusionFamily( name = "z-image", diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 5838566c3d..3fbad30571 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -41,7 +41,13 @@ from core.inference.diffusion_families import ( resolve_base_repo, resolve_local_gguf_child, ) -from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles +from core.inference.diffusion_memory import ( + OFFLOAD_GROUP, + OFFLOAD_MODEL, + OFFLOAD_NONE, + OFFLOAD_SEQUENTIAL, +) +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags from core.inference.sd_cpp_engine import ( SdCppCancelled, SdCppEngine, @@ -110,6 +116,26 @@ class _SdState: native_speed: str = "off" offload_flags: tuple[str, ...] = () threads: Optional[int] = None + sampling_method: Optional[str] = None + flow_shift: Optional[float] = None + + +def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: + """Map the diffusers memory knobs onto an sd-cli offload policy. Only meaningful + off-CPU (forced sd_cpp / MPS); on CPU everything is resident in RAM anyway.""" + mode = (memory_mode or "").strip().lower() + if mode == "low_vram": + return OFFLOAD_SEQUENTIAL + if mode == "balanced": + return OFFLOAD_GROUP + if cpu_offload and mode in ("", "auto"): + return OFFLOAD_MODEL + return OFFLOAD_NONE + + +def _native_speed_for(speed_mode: Optional[str]) -> str: + mode = (speed_mode or "off").strip().lower() + return mode if mode in ("default", "max") else "off" @dataclass @@ -235,6 +261,9 @@ class SdCppDiffusionBackend: base = base, fam = fam, hf_token = hf_token, + cpu_offload = cpu_offload, + memory_mode = memory_mode, + speed_mode = speed_mode, _load_token = token, ), daemon = True, @@ -249,6 +278,9 @@ class SdCppDiffusionBackend: base: str, fam: DiffusionFamily, hf_token: Optional[str], + cpu_offload: bool = False, + memory_mode: Optional[str] = None, + speed_mode: Optional[str] = None, _load_token: int, ) -> None: try: @@ -271,6 +303,12 @@ class SdCppDiffusionBackend: qwen2vl = paths.get("qwen2vl"), ) device = resolve_diffusion_device_target().device + # Honor the requested speed everywhere; offload only off-CPU (forced + # sd_cpp / MPS), since on CPU sd-cli is resident in RAM and the offload + # flags are no-ops. + offload: tuple[str, ...] = () + if device != "cpu": + offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload))) state = _SdState( repo_id = repo_id, base_repo = base, @@ -278,12 +316,17 @@ class SdCppDiffusionBackend: device = device, files = files, vae_format = fam.sd_cpp_vae_format, - native_speed = "off", - offload_flags = (), + native_speed = _native_speed_for(speed_mode), + offload_flags = offload, threads = None, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, ) - # Probe the binary version (cheap) so a broken install fails the load. - engine.version() + # Probe the binary: version() returns None when the present binary cannot + # run (bad permissions / missing shared libs), so fail the load now rather + # 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.") with self._lock: if self._load_token != _load_token: return # superseded / cancelled @@ -319,8 +362,10 @@ class SdCppDiffusionBackend: try: from huggingface_hub import HfApi api = HfApi(token = hf_token) - for repo, fn, _ in assets: - if Path(repo).expanduser().exists(): + for repo, fn, kind in assets: + # Only the transformer can be a local path; for the others ``repo`` is + # an HF id (a same-named local dir must not skip the size estimate). + if kind == "diffusion_model" and Path(repo).expanduser().exists(): continue try: info = api.get_paths_info(repo, paths = [fn], expand = False) @@ -405,10 +450,15 @@ class SdCppDiffusionBackend: else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) - vae_extra = ["--vae-format", state.vae_format] if state.vae_format else None + extra_args: list[str] = [] + if state.vae_format: + extra_args += ["--vae-format", state.vae_format] + if state.flow_shift is not None: + extra_args += ["--flow-shift", repr(float(state.flow_shift))] self._gen = _SdGen(total_steps = int(steps)) images = [] + seeds: list[int] = [] with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: for index in range(max(1, int(batch_size))): if cancel.is_set(): @@ -426,6 +476,7 @@ class SdCppDiffusionBackend: cfg_scale = cfg_scale, guidance = flux_guidance, seed = seed_i, + sampling_method = state.sampling_method, batch_count = 1, ) engine.generate( @@ -435,15 +486,18 @@ class SdCppDiffusionBackend: offload = list(state.offload_flags) or None, native_speed = state.native_speed, threads = state.threads, - extra_args = vae_extra, + extra_args = extra_args or None, on_log = self._on_log, cancel_event = cancel, ) with Image.open(out_path) as im: images.append(im.copy()) + seeds.append(seed_i) if cancel.is_set(): raise RuntimeError("Diffusion generation was cancelled.") - return {"images": images, "seed": int(seed), "repo_id": state.repo_id} + # ``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 {"images": images, "seed": int(seed), "seeds": seeds, "repo_id": state.repo_id} except SdCppCancelled as exc: raise RuntimeError("Diffusion generation was cancelled.") from exc finally: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 399b1ffbdc..6714151045 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -189,7 +189,10 @@ class SdCppEngine: return bool(self.binary) and Path(self.binary).is_file() def version(self, *, timeout: float = 10.0) -> Optional[str]: - """First line of ``sd-cli --version``, cached. None if it can't run.""" + """First line of ``sd-cli --version``, cached on success. ``None`` when the + binary is absent OR present-but-unrunnable (exec error / nonzero exit, e.g. + missing shared libraries / bad permissions), so callers can fail a load early + instead of committing a "ready" state that crashes on first generation.""" if not self.is_available(): return None if self._version is not None: @@ -204,10 +207,12 @@ class SdCppEngine: check = False, env = runtime_env(self.binary), ) - text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip() - self._version = text.splitlines()[0] if text else "" except (OSError, subprocess.SubprocessError): - self._version = "" + return None + if res.returncode != 0: + return None + text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip() + self._version = text.splitlines()[0] if text else "" return self._version def generate( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1bad4c48be..1363dca458 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10127,20 +10127,28 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - # No model loaded (or unloaded mid-flight) — a client-state problem. - raise HTTPException(status_code = 409, detail = str(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). + msg = str(exc) + if "No diffusion model is loaded" in msg or "cancelled" in msg.lower(): + raise HTTPException(status_code = 409, detail = msg) + logger.error("diffusion.generate_failed: %s", exc) + raise HTTPException(status_code = 500, detail = "Image generation failed.") except Exception as exc: logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") - # Persist each image with its full recipe embedded. The whole batch shares - # one seed (drawn sequentially from one generator) and one timestamp (the - # images are generated together), so their gallery order is stable on reload. + # Persist each image with its full recipe embedded. The diffusers batch shares + # one seed (drawn sequentially from one generator); the native sd.cpp batch uses a + # distinct seed per image and returns them in ``seeds`` so each is reproducible. created_at = time.time() + per_image_seeds = result.get("seeds") def _persist() -> list[dict]: records = [] for index, image in enumerate(result["images"]): + seed = per_image_seeds[index] if per_image_seeds and index < len(per_image_seeds) else result["seed"] records.append( image_gallery.save( image, @@ -10151,9 +10159,9 @@ async def generate_diffusion_image( "height": request.height, "steps": request.steps, "guidance": request.guidance, - "seed": result["seed"], - # Position within the batch: images here share a seed + timestamp, - # so the export filename needs this to stay unique. + "seed": seed, + # Position within the batch: shared timestamp, so the export + # filename needs this to stay unique. "batch_index": index, "model": result.get("repo_id"), "created_at": created_at, diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 3d0f3f4691..c6fbf7243e 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -74,6 +74,8 @@ def _loaded_backend(fam_name = "z-image", engine = None): diffusion_model = "/m/z.gguf", vae = "/m/vae.safetensors", llm = "/m/llm.safetensors" ), vae_format = fam.sd_cpp_vae_format, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, ) return b @@ -153,6 +155,17 @@ def test_generate_returns_images_and_seed(): assert len(eng.calls) == 2 seeds = [params.seed for _, params, _, _ in eng.calls] assert seeds == [123, 124] + # The per-image seeds are returned so the route can persist each one. + assert out["seeds"] == [123, 124] + + +def test_generate_qwen_passes_sampling_args(): + eng = _FakeEngine() + b = _loaded_backend(fam_name = "qwen-image", engine = eng) + b.generate(prompt = "x", steps = 20, guidance = 4.0, seed = 1) + _, params, _, kw = eng.calls[0] + assert params.sampling_method == "euler" # Qwen's supported sd.cpp sampler + assert "--flow-shift" in (kw.get("extra_args") or []) def test_generate_raises_when_not_loaded(): diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index af3f015830..054574d905 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -95,7 +95,9 @@ def test_engine_version_parsed_and_cached(tmp_path, monkeypatch): def _fake_run(*_a, **_k): calls["n"] += 1 - return types.SimpleNamespace(stdout = "stable-diffusion.cpp version master-721\n", stderr = "") + return types.SimpleNamespace( + stdout = "stable-diffusion.cpp version master-721\n", stderr = "", returncode = 0 + ) monkeypatch.setattr(eng.subprocess, "run", _fake_run) assert e.version() == "stable-diffusion.cpp version master-721"