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 new file mode 100644 index 0000000000..c8db7f81d2 --- /dev/null +++ b/studio/backend/core/inference/diffusion_engine_router.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Selects the diffusion engine (diffusers vs native sd.cpp) for the live route. + +The image routes drive one engine at a time. On a CUDA/ROCm/XPU GPU that engine is +the diffusers ``DiffusionBackend`` (the default, and the only path with the torchao +fast-quant / compile stack). With no usable GPU (CPU, or Apple MPS when explicitly +enabled) it is the native ``SdCppDiffusionBackend``, which is faster and far lighter +on RAM there. The choice is made once at load time and remembered, so ``generate`` / +``unload`` / ``status`` / progress all act on the same engine the load committed to. + +Selection is centralised here and built on the existing pure ``select_diffusion_engine`` +decision; this module adds the policy around it (env opt-out, MPS gating, per-family +native-asset support, lazy binary availability) and records why a fallback happened. + +Env knobs (one canonical interpretation each): + UNSLOTH_DIFFUSION_ENGINE=auto|diffusers|sd_cpp force an engine (auto = decide) + UNSLOTH_DIFFUSION_SD_CPP=auto|0|1 enable/disable the native route + UNSLOTH_DIFFUSION_SD_CPP_MPS=0|1 allow native on Apple MPS (default off) + UNSLOTH_DIFFUSION_SD_CPP_INSTALL=auto|0|1 allow lazy binary install (in sd_cpp_backend) +""" + +from __future__ import annotations + +import os +import threading +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, + SdCppEngine, + select_diffusion_engine, +) +from loggers import get_logger + +logger = get_logger(__name__) + +_DISABLE_TOKENS = frozenset({"0", "off", "false", "no"}) +_ENABLE_TOKENS = frozenset({"1", "on", "true", "yes"}) + +# The engine the current (or most recent) load committed to, and why a non-native +# choice was made. Mutated only under _lock during selection. +_lock = threading.Lock() +_active_engine_name: str = ENGINE_DIFFUSERS +_fallback_reason: Optional[str] = None + + +def _engine_config() -> tuple[str, str, bool]: + forced = os.environ.get("UNSLOTH_DIFFUSION_ENGINE", "auto").strip().lower() + sd_cpp = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP", "auto").strip().lower() + mps = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP_MPS", "0").strip().lower() in _ENABLE_TOKENS + return forced, sd_cpp, mps + + +def get_active_diffusion_engine() -> Any: + """The engine object the active selection points at (defaults to diffusers).""" + if _active_engine_name == ENGINE_SD_CPP: + from core.inference.sd_cpp_backend import get_sd_cpp_backend + return get_sd_cpp_backend() + from core.inference.diffusion import get_diffusion_backend + + return get_diffusion_backend() + + +def active_engine_name() -> str: + return _active_engine_name + + +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: + if name != _active_engine_name: + 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: + logger.info("diffusion engine: diffusers (%s)", reason or "selected") + return get_active_diffusion_engine() + + +def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] = None) -> Any: + """Pick + activate the engine for loading ``fam`` on this host; return the engine. + + Falls back to diffusers (recording a reason) whenever the native route is + disabled, the device has a usable GPU, MPS is not enabled, the family has no + native asset mapping, or the sd-cli binary is unavailable -- always BEFORE the + slow load begins, so a fallback never strands a half-native load. + """ + forced, sd_cpp_pref, mps_enabled = _engine_config() + + if forced == ENGINE_DIFFUSERS: + return _activate(ENGINE_DIFFUSERS, "forced (UNSLOTH_DIFFUSION_ENGINE=diffusers)") + + prefer_native = forced == ENGINE_SD_CPP + if sd_cpp_pref in _DISABLE_TOKENS and not prefer_native: + return _activate(ENGINE_DIFFUSERS, "native engine disabled (UNSLOTH_DIFFUSION_SD_CPP=0)") + + target = resolve_diffusion_device_target() + backend = target.backend + # Policy: CPU is always native-eligible; MPS only when explicitly enabled; a GPU + # backend (cuda/rocm/xpu) never is, unless the user force-selects sd_cpp. + policy_eligible = backend == "cpu" or (backend == "mps" and mps_enabled) or prefer_native + fam_ok = family_sd_cpp_supported(fam) + + 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( + backend, native_available = native_available, prefer_native = prefer_native + ) + if choice == ENGINE_SD_CPP: + return _activate(ENGINE_SD_CPP, None) + + # Explain the diffusers choice for status/telemetry. + if not policy_eligible: + reason = f"GPU backend '{backend}' uses diffusers" + elif not fam_ok: + reason = f"family '{fam.name}' has no native sd.cpp asset mapping" + elif not binary: + reason = "sd-cli binary unavailable" + else: + reason = "diffusers selected" + return _activate(ENGINE_DIFFUSERS, reason) + + +def annotate_status(status: dict[str, Any]) -> dict[str, Any]: + """Tag a backend status dict with the active engine + any fallback reason.""" + out = dict(status) + out["engine"] = _active_engine_name + out["fallback_reason"] = _fallback_reason + return out + + +def active_status() -> dict[str, Any]: + """The active engine's status, annotated with which engine + any fallback reason.""" + return annotate_status(get_active_diffusion_engine().status()) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index d5068cbf49..3afbe9b495 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -45,6 +45,26 @@ class DiffusionFamily: # materialising the dense bf16 transformer on the GPU (much lower load VRAM + a # smaller download). Empty until checkpoints are hosted -> behaviour is unchanged. prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple) + # Native (stable-diffusion.cpp) single-file assets, used only when the no-GPU + # sd.cpp engine is selected (CPU / Apple). The transformer GGUF is shared with + # the diffusers path; sd-cli additionally needs a single-file VAE and text + # encoder(s), because the diffusers base repo ships those sharded and sd-cli + # cannot read that layout. Each asset is a hashable (repo_id, filename) the + # backend fetches with hf_hub_download. ``sd_cpp_text_encoders`` carries a + # trailing SdCppModelFiles field name (clip_l / t5xxl / llm / qwen2vl / clip_g) + # so the backend maps each file onto the right sd-cli flag. Empty -> the family + # has no native mapping and the sd.cpp route falls back to diffusers. + sd_cpp_vae: Optional[tuple[str, str]] = None + # VAE latent-format override for sd-cli (--vae-format): "flux2" for the FLUX.2 + # 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 @@ -60,6 +80,11 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", aliases = ("flux1", "flux-1"), + sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"), + sd_cpp_text_encoders = ( + ("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"), + ("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "t5xxl"), + ), ), # FLUX.2-klein is a distinct pipeline (Flux2KleinPipeline) with a Qwen3 text # encoder, not the Mistral-based Flux2Pipeline; it must precede a generic @@ -70,6 +95,14 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-klein-4B", aliases = ("flux2-klein",), + # FLUX.2 uses a distinct 32-channel autoencoder; sd-cli needs the latent + # format override. The single-file VAE ships in Comfy-Org/flux2-dev (the + # klein-4B repo only has a sharded diffusers VAE). Shares Qwen3-4B with z-image. + sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"), + sd_cpp_vae_format = "flux2", + sd_cpp_text_encoders = ( + ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"), + ), ), DiffusionFamily( name = "qwen-image", @@ -78,6 +111,19 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( base_repo = "Qwen/Qwen-Image", cfg_kwarg = "true_cfg_scale", aliases = ("qwen_image", "qwenimage"), + sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"), + # The Qwen2.5-VL text encoder as a Q4_K_M GGUF keeps the CPU RAM win (the + # bf16 safetensors encoder is ~15 GB). sd-cli's --qwen2vl is an alias of --llm. + sd_cpp_text_encoders = ( + ( + "unsloth/Qwen2.5-VL-7B-Instruct-GGUF", + "Qwen2.5-VL-7B-Instruct-Q4_K_M.gguf", + "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", @@ -87,6 +133,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("zimage", "z_image"), # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, + sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"), + sd_cpp_text_encoders = ( + ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"), + ), ), ) @@ -137,6 +187,13 @@ def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: return None +def family_sd_cpp_supported(fam: DiffusionFamily) -> bool: + """True when the family has the single-file VAE + text-encoder mapping the + native sd.cpp engine needs. A family without it can only run on diffusers, so + the no-GPU route falls back rather than routing to sd-cli.""" + return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders) + + def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: """Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes. diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index daa0ef4e94..a6c79bf1e8 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -55,8 +55,10 @@ def _evict_chat() -> None: def _evict_diffusion() -> None: - from core.inference.diffusion import get_diffusion_backend - get_diffusion_backend().unload() + # Unload whichever engine the router has active (diffusers or native sd.cpp), so a + # chat acquire frees the right one. + from core.inference.diffusion_engine_router import get_active_diffusion_engine + get_active_diffusion_engine().unload() # Patchable in tests via monkeypatch.setitem. diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py new file mode 100644 index 0000000000..4d0bbcbfa7 --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -0,0 +1,661 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Native stable-diffusion.cpp diffusion backend (the no-GPU tier). + +``SdCppDiffusionBackend`` presents the SAME public surface the image routes use on +the diffusers ``DiffusionBackend`` (``begin_load`` / ``load_progress`` / ``generate`` +/ ``generate_progress`` / ``unload`` / ``status``), but is backed by the ``sd-cli`` +subprocess (``SdCppEngine``) instead of an in-process diffusers pipeline. The engine +router (``diffusion_engine_router.py``) selects this backend only when no usable +CUDA/ROCm/XPU GPU is present, where it is measurably faster and far lighter on RAM +than diffusers (see outputs/sdcpp_cpu). + +It reuses the transformer GGUF the diffusers path already downloads and additionally +fetches the per-family single-file VAE + text encoders declared in +``diffusion_families`` (sd-cli cannot read the sharded diffusers components). The +binary is installed lazily on first use; if it is unavailable or the family has no +native mapping, the router falls back to diffusers, so this backend is only ever +asked to run requests it can serve. + +Import-light on purpose: no torch / diffusers here, so selecting it on a CPU box +does not drag the heavy GPU stack into the process. +""" + +from __future__ import annotations + +import logging +import os +import re +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional + +from core.inference.diffusion_device import resolve_diffusion_device_target +from core.inference.diffusion_families import ( + DiffusionFamily, + detect_family, + family_sd_cpp_supported, + resolve_base_repo, + resolve_local_gguf_child, +) +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, + find_sd_cpp_binary, +) +from loggers import get_logger + +logger = get_logger(__name__) + +# A sampling-progress line like " 4/4" / "[ 12/ 28]" / "sampling: 50%|...| 14/28". +# We only trust a match whose denominator equals the requested step count, so an +# unrelated "1/100" elsewhere in the log can't move the bar. +_STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)") + +# Serialises the one-time binary install so concurrent first-loads don't race on the +# download / extract / chmod. +_install_lock = threading.Lock() + + +def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu") -> Optional[str]: + """Path to a usable ``sd-cli`` binary, installing the prebuilt once if needed. + + Returns the binary path, or None when it is absent and cannot be installed + (install disabled, no network, unsupported platform). Never raises -- a None + return is the router's signal to fall back to diffusers. + """ + found = find_sd_cpp_binary() + if found: + return found + if not allow_install: + return None + with _install_lock: + # Re-check inside the lock: a concurrent first-load may have installed it. + found = find_sd_cpp_binary() + if found: + return found + try: + import sys + + studio_dir = Path(__file__).resolve().parents[3] # .../studio + if str(studio_dir) not in sys.path: + sys.path.insert(0, str(studio_dir)) + from install_sd_cpp_prebuilt import install as _install + except Exception as exc: # noqa: BLE001 -- import path / module issues are non-fatal + logger.warning("sd-cli installer import failed: %s", exc) + return None + try: + path = _install(accelerator = accelerator) + logger.info("sd-cli installed at %s", path) + return str(path) + except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back + logger.warning("sd-cli auto-install failed: %s", exc) + return None + + +@dataclass(frozen = True) +class _SdState: + """The loaded native checkpoint: resolved asset paths + run settings.""" + + repo_id: str + base_repo: str + family: DiffusionFamily + device: str + files: SdCppModelFiles + vae_format: Optional[str] = None + 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 +class _SdLoading: + """An in-flight asset download, polled for progress.""" + + repo_id: str + base_repo: str + expected_bytes: int = 0 + downloaded_bytes: int = 0 + error: Optional[str] = None + + +@dataclass +class _SdGen: + """An in-flight generation, updated from parsed sd-cli progress lines.""" + + total_steps: int + step: int = 0 + first_step_at: float = 0.0 + eta_seconds: Optional[float] = None + + +def _estimate_eta(total_steps: int, step: int, first_step_at: float, now: float) -> Optional[float]: + steps_since_first = step - 1 + if not first_step_at or steps_since_first <= 0: + return None + per_step = (now - first_step_at) / steps_since_first + return max(0.0, (total_steps - step) * per_step) + + +def _map_guidance( + fam: DiffusionFamily, guidance: Optional[float] +) -> tuple[Optional[float], Optional[float]]: + """(cfg_scale, guidance) for sd-cli from the single diffusers ``guidance`` value. + + FLUX families take a distilled embedded ``--guidance``; everyone else uses real + classifier-free ``--cfg-scale``. A distilled 0/1 means CFG off (sd-cli's 1.0); a + value > 1 is real CFG. Mirrors the engine mapping validated in the CPU benchmark. + """ + if fam.name in ("flux.1", "flux.2-klein"): + return None, (float(guidance) if guidance is not None else None) + cfg = float(guidance) if (guidance is not None and guidance > 1.0) else 1.0 + return cfg, None + + +class SdCppDiffusionBackend: + """Native sd.cpp backend with the diffusers ``DiffusionBackend`` method surface.""" + + def __init__(self, engine: Optional[SdCppEngine] = None) -> None: + self._lock = threading.Lock() + self._generate_lock = threading.Lock() + self._engine = engine # resolved lazily on first load so import stays cheap + self._state: Optional[_SdState] = None + self._loading: Optional[_SdLoading] = None + self._load_token = 0 + self._cancel_event = threading.Event() + self._active_generate_cancel: Optional[threading.Event] = None + self._gen: Optional[_SdGen] = None + + @property + def is_loaded(self) -> bool: + return self._state is not None + + def _resolve_engine(self) -> SdCppEngine: + """The SdCppEngine, installing the binary on first use. Raises if unusable.""" + if self._engine is not None and self._engine.is_available(): + return self._engine + binary = ensure_sd_cpp_binary(allow_install = _install_allowed()) + if not binary: + raise RuntimeError("sd-cli (stable-diffusion.cpp) binary is unavailable.") + self._engine = SdCppEngine(binary = binary) + return self._engine + + # ── Background load + progress ───────────────────────────────────────── + + def begin_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + cpu_offload: bool = False, + memory_mode: Optional[str] = None, + speed_mode: Optional[str] = None, + # diffusers-only knobs accepted (so the route calls both engines uniformly) + # and ignored -- sd.cpp has no torchao quant / SDPA dispatcher / fbcache. + text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, + transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + 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." + ) + fam = detect_family(repo_id, 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): + raise ValueError(f"Family '{fam.name}' has no native sd.cpp asset mapping.") + + base = resolve_base_repo(fam, base_repo) + 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() + self._loading = _SdLoading(repo_id = repo_id, base_repo = base) + + threading.Thread( + target = self._run_load, + kwargs = dict( + repo_id = repo_id, + gguf_filename = gguf_filename, + 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, + ).start() + return self.status() + + def _run_load( + self, + *, + repo_id: str, + gguf_filename: str, + 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: + # Ensure the binary up front so an install failure surfaces before the + # multi-GB asset pull (the router also pre-checks, but a forced reload here + # must not silently download then fail at generate). + engine = self._resolve_engine() + + assets = self._asset_specs(repo_id, gguf_filename, fam) + self._set_expected_bytes(assets, hf_token) + paths = self._fetch_assets(assets, hf_token) + + files = SdCppModelFiles( + diffusion_model = paths["diffusion_model"], + vae = paths.get("vae"), + clip_l = paths.get("clip_l"), + clip_g = paths.get("clip_g"), + t5xxl = paths.get("t5xxl"), + llm = paths.get("llm"), + 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, + family = fam, + device = device, + files = files, + vae_format = fam.sd_cpp_vae_format, + 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() 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.") + # 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 + 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 + if self._load_token != _load_token: + return + logger.error("sd_cpp.load_failed: %s", exc) + with self._lock: + if self._load_token == _load_token and self._loading is not None: + self._loading.error = str(exc) + + def _asset_specs( + self, repo_id: str, gguf_filename: str, fam: DiffusionFamily + ) -> list[tuple[str, str, str]]: + """(repo, filename, kind) for every file sd-cli needs. ``kind`` is the + SdCppModelFiles field; the transformer reuses the diffusers GGUF.""" + specs: list[tuple[str, str, str]] = [(repo_id, gguf_filename, "diffusion_model")] + if fam.sd_cpp_vae: + specs.append((fam.sd_cpp_vae[0], fam.sd_cpp_vae[1], "vae")) + for terepo, tefile, kind in fam.sd_cpp_text_encoders: + specs.append((terepo, tefile, kind)) + return specs + + def _set_expected_bytes( + self, assets: list[tuple[str, str, str]], hf_token: Optional[str] + ) -> None: + """Best-effort total download size for the progress bar (0 if unknown).""" + total = 0 + try: + from huggingface_hub import HfApi + api = HfApi(token = hf_token) + 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) + for it in info: + total += int(getattr(it, "size", 0) or 0) + except Exception: # noqa: BLE001 -- one missing size is non-fatal + continue + except Exception: # noqa: BLE001 -- estimate is best-effort + total = 0 + loading = self._loading + if loading is not None: + loading.expected_bytes = total + + def _fetch_assets( + self, assets: list[tuple[str, str, str]], hf_token: Optional[str] + ) -> dict[str, str]: + """Download every asset (cancellable), returning kind -> local path.""" + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + paths: dict[str, str] = {} + for repo, fn, kind in assets: + if self._cancel_event.is_set(): + raise SdCppCancelled("load cancelled") + local_root = Path(repo).expanduser() + if kind == "diffusion_model" and local_root.exists(): + path = str(resolve_local_gguf_child(local_root, fn)) + else: + path = hf_hub_download_with_xet_fallback( + repo, fn, hf_token, cancel_event = self._cancel_event + ) + paths[kind] = path + with self._lock: + if self._loading is not None: + try: + self._loading.downloaded_bytes += os.path.getsize(path) + except OSError: + pass + return paths + + def load_progress(self) -> dict[str, Any]: + loading = self._loading + if loading is not None and loading.error: + return _progress("error", error = loading.error) + if loading is None: + return _progress("ready" if self._state is not None else None) + downloaded = loading.downloaded_bytes + expected = loading.expected_bytes + if expected > 0 and downloaded >= expected * 0.999: + return _progress("finalizing", min(downloaded, expected), expected, 1.0) + fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 + return _progress("downloading", downloaded, expected, fraction) + + # ── Generate ─────────────────────────────────────────────────────────── + + def generate( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + steps: int = 9, + guidance: float = 0.0, + seed: Optional[int] = None, + batch_size: int = 1, + ) -> dict[str, Any]: + import tempfile + + from PIL import Image + + cancel = threading.Event() + with self._generate_lock: + with self._lock: + state = self._state + if state is None: + raise RuntimeError("No diffusion model is loaded.") + self._active_generate_cancel = cancel + engine = self._resolve_engine() + try: + if seed is None: + seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) + else: + seed = int(seed) + cfg_scale, flux_guidance = _map_guidance(state.family, guidance) + 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(): + 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. + # 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, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + cfg_scale = cfg_scale, + guidance = flux_guidance, + seed = seed_i, + sampling_method = state.sampling_method, + batch_count = 1, + ) + engine.generate( + state.files, + params, + output_path = out_path, + offload = list(state.offload_flags) or None, + native_speed = state.native_speed, + threads = state.threads, + 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.") + # ``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: + self._gen = None + with self._lock: + if self._active_generate_cancel is cancel: + self._active_generate_cancel = None + + def _on_log(self, line: str) -> None: + gen = self._gen + if gen is None or gen.total_steps <= 0: + return + for a, b in _STEP_RE.findall(line): + if int(b) == gen.total_steps: + now = time.time() + gen.step = min(int(a), gen.total_steps) + if gen.first_step_at == 0.0: + gen.first_step_at = now + gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now) + + def generate_progress(self) -> dict[str, Any]: + gen = self._gen + if gen is None or gen.total_steps <= 0: + return { + "active": False, + "step": 0, + "total_steps": 0, + "fraction": 0.0, + "eta_seconds": None, + } + return { + "active": True, + "step": gen.step, + "total_steps": gen.total_steps, + "fraction": min(gen.step / gen.total_steps, 1.0), + "eta_seconds": gen.eta_seconds, + } + + # ── Unload / status ────────────────────────────────────────────────────── + + def unload(self) -> dict[str, Any]: + self._cancel_event.set() + with self._lock: + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + self._state = None + self._load_token += 1 + self._loading = None + return self.status() + + def status(self) -> dict[str, Any]: + state = self._state + if state is None: + return { + "loaded": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + "offload_policy": None, + "vae_tiling": False, + "memory_mode": None, + "speed_mode": None, + "speed_optims": [], + "text_encoder_quant": None, + "transformer_quant": None, + "attention_backend": None, + "transformer_cache": None, + "engine": "sd_cpp", + } + return { + "loaded": True, + "repo_id": state.repo_id, + "family": state.family.name, + "base_repo": state.base_repo, + "device": state.device, + "dtype": "gguf", + # 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, + "speed_optims": [], + "text_encoder_quant": None, + "transformer_quant": None, + "attention_backend": None, + "transformer_cache": None, + "engine": "sd_cpp", + } + + +def _install_allowed() -> bool: + """Whether lazy binary install is permitted (UNSLOTH_DIFFUSION_SD_CPP_INSTALL).""" + val = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP_INSTALL", "auto").strip().lower() + return val not in ("0", "off", "false", "no") + + +def _progress( + phase: Optional[str], + bytes_downloaded: int = 0, + bytes_total: int = 0, + fraction: float = 0.0, + *, + error: Optional[str] = None, +) -> dict[str, Any]: + return { + "phase": phase, + "bytes_downloaded": bytes_downloaded, + "bytes_total": bytes_total, + "fraction": fraction, + "error": error, + } + + +_sd_cpp_backend: Optional[SdCppDiffusionBackend] = None + + +def get_sd_cpp_backend() -> SdCppDiffusionBackend: + global _sd_cpp_backend + if _sd_cpp_backend is None: + _sd_cpp_backend = SdCppDiffusionBackend() + return _sd_cpp_backend diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 186ea51631..6714151045 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -26,6 +26,7 @@ import logging import os import queue import shutil +import signal import subprocess import sys import threading @@ -50,6 +51,30 @@ _BINARY_STEM = "sd-cli" _LEGACY_STEM = "sd" +class SdCppCancelled(RuntimeError): + """A generation was cancelled via its ``cancel_event`` (unload / superseding load + / arbiter eviction). Distinct from a generation *failure* so the caller can keep + cancellation semantics (no diffusers fallback, no error surfaced as a crash).""" + + +def _terminate(proc: "subprocess.Popen") -> None: + """Hard-stop an sd-cli process (and any children). On POSIX the process is its + own session leader (``start_new_session``), so kill the whole group; otherwise + fall back to killing just the process.""" + if proc.poll() is not None: + return + try: + if os.name == "posix": + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + else: + proc.kill() + except Exception: # noqa: BLE001 -- killpg can miss (no pgid / already gone); fall back + try: + proc.kill() + except Exception: # noqa: BLE001 -- best-effort teardown + pass + + def _binary_name(stem: str) -> str: return f"{stem}.exe" if sys.platform == "win32" else stem @@ -164,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: @@ -179,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( @@ -199,6 +229,7 @@ class SdCppEngine: timeout: Optional[float] = 1800.0, env: Optional[dict[str, str]] = None, on_log: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, ) -> Path: """Run one ``sd-cli`` generation; return the written image path. @@ -206,7 +237,9 @@ class SdCppEngine: (``--diffusion-fa`` etc.), de-duplicated against the offload flags that may already include them. Raises ``RuntimeError`` if the binary is missing, the process exits nonzero, or no output file is produced. ``on_log`` (if given) - receives each line of sd-cli's progress output as it arrives. + receives each line of sd-cli's progress output as it arrives. ``cancel_event`` + (if given) is polled while the child runs; when set, the process tree is + killed and ``SdCppCancelled`` is raised. """ offload = list(offload or []) speed = [f for f in native_speed_flags(native_speed) if f not in offload] @@ -221,7 +254,14 @@ class SdCppEngine: verbose = verbose, extra_args = merged_extra, ) - return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + return self._run( + cmd, + output_path, + timeout = timeout, + env = env, + on_log = on_log, + cancel_event = cancel_event, + ) def upscale( self, @@ -233,6 +273,7 @@ class SdCppEngine: timeout: Optional[float] = 1800.0, env: Optional[dict[str, str]] = None, on_log: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, ) -> Path: """Upscale an image with an ESRGAN model; return the written path.""" cmd = build_sd_cpp_upscale_command( @@ -242,7 +283,14 @@ class SdCppEngine: verbose = verbose, extra_args = extra_args, ) - return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + return self._run( + cmd, + output_path, + timeout = timeout, + env = env, + on_log = on_log, + cancel_event = cancel_event, + ) # ── internals ───────────────────────────────────────────────────────────── @@ -268,11 +316,13 @@ class SdCppEngine: timeout: Optional[float], env: Optional[dict[str, str]], on_log: Optional[Callable[[str], None]], + cancel_event: Optional[threading.Event] = None, ) -> Path: """Run an sd-cli argv, stream output, and return the produced image path. - Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output. - Shared by ``generate`` and ``upscale``. + Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output, and + ``SdCppCancelled`` when ``cancel_event`` fires. Shared by ``generate`` and + ``upscale``. """ out = Path(output_path) base = dict(os.environ) @@ -289,6 +339,9 @@ class SdCppEngine: text = True, errors = "replace", env = run_env, + # 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"), ) # 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 @@ -313,8 +366,13 @@ class SdCppEngine: stdout_done = False try: while True: + # Cancellation (unload / superseding load / arbiter eviction): kill the + # process tree and signal the caller it was cancelled, not a failure. + if cancel_event is not None and cancel_event.is_set() and proc.poll() is None: + _terminate(proc) + raise SdCppCancelled("sd-cli generation was cancelled.") if deadline is not None and time.monotonic() >= deadline and proc.poll() is None: - proc.kill() + _terminate(proc) raise RuntimeError(f"sd-cli timed out after {timeout}s") try: line = line_q.get(timeout = 0.1) @@ -335,7 +393,7 @@ class SdCppEngine: ret = proc.wait(timeout = 5.0) finally: if proc.poll() is None: - proc.kill() + _terminate(proc) if ret != 0: raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 21e9760aed..2c1701f6ea 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1912,3 +1912,8 @@ class DiffusionStatusResponse(BaseModel): "_native_cudnn), or null for the default SDPA", ) transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") + engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp") + fallback_reason: Optional[str] = Field( + None, + description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 38dd2d1f20..400dc239fc 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10299,15 +10299,22 @@ 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_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() try: # Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, - # missing local GGUF) must not evict a working chat model and then 400. - # validate_load_request does local-path I/O, so run it off the event loop. - await asyncio.to_thread( + # missing local GGUF) must not evict a working chat model and then 400. The + # validated family also drives engine selection below. + fam = await asyncio.to_thread( backend.validate_load_request, request.model_path, gguf_filename = request.gguf_filename, @@ -10317,11 +10324,22 @@ async def load_diffusion_model( # compete with the training subprocess for VRAM. The chat path does the # same via _guard_chat_load_against_training; this is its image sibling. _guard_diffusion_load_against_training() - # 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) + # Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU), + # 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) + # 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( - backend.begin_load, + engine.begin_load, request.model_path, gguf_filename = request.gguf_filename, base_repo = request.base_repo, @@ -10338,7 +10356,7 @@ async def load_diffusion_model( transformer_cache = request.transformer_cache, transformer_cache_threshold = request.transformer_cache_threshold, ) - return DiffusionStatusResponse(**status_dict) + return DiffusionStatusResponse(**annotate_status(status_dict)) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) except RuntimeError as exc: @@ -10351,9 +10369,9 @@ async def generate_diffusion_image( request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject) ): from core.inference import image_gallery - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import get_active_diffusion_engine - backend = get_diffusion_backend() + backend = get_active_diffusion_engine() try: result = await asyncio.to_thread( backend.generate, @@ -10367,26 +10385,32 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - if not backend.is_loaded: - # The only genuine client-state 409: nothing is loaded to generate with. - raise HTTPException(status_code = 409, detail = "No diffusion model is loaded.") - # A pipeline RuntimeError (CUDA OOM, shape/device) is a server failure; fall - # through to the sanitized 500 instead of echoing raw exception text (which - # would 409 an OOM as retryable and leak VRAM totals / tensor shapes). + # 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, @@ -10397,9 +10421,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, # The batch shares one seed, so reproducing image batch_index>0 # needs the original batch_size: persist it so restore can replay. @@ -10476,27 +10500,27 @@ async def clear_gallery_images(current_subject: str = Depends(get_current_subjec @studio_router.post("/images/unload", response_model = DiffusionStatusResponse) async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import annotate_status, get_active_diffusion_engine from core.inference.gpu_arbiter import release, DIFFUSION - status_dict = await asyncio.to_thread(get_diffusion_backend().unload) + status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload) release(DIFFUSION) - return DiffusionStatusResponse(**status_dict) + return DiffusionStatusResponse(**annotate_status(status_dict)) @studio_router.get("/images/status", response_model = DiffusionStatusResponse) async def diffusion_status(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionStatusResponse(**get_diffusion_backend().status()) + from core.inference.diffusion_engine_router import active_status + return DiffusionStatusResponse(**active_status()) @studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress()) + from core.inference.diffusion_engine_router import get_active_diffusion_engine + return DiffusionLoadProgressResponse(**get_active_diffusion_engine().load_progress()) @studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse) async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress()) + from core.inference.diffusion_engine_router import get_active_diffusion_engine + return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress()) diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py new file mode 100644 index 0000000000..07d495757b --- /dev/null +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the diffusion engine router (diffusers vs native sd.cpp selection).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from core.inference import diffusion_engine_router as r +from core.inference.diffusion_families import detect_family +from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP + +_ENVS = ( + "UNSLOTH_DIFFUSION_ENGINE", + "UNSLOTH_DIFFUSION_SD_CPP", + "UNSLOTH_DIFFUSION_SD_CPP_MPS", + "UNSLOTH_DIFFUSION_SD_CPP_INSTALL", +) + + +@pytest.fixture(autouse = True) +def _clean_env_and_state(monkeypatch): + for e in _ENVS: + monkeypatch.delenv(e, raising = False) + # A light status-capable stub so neither selection nor active_status() imports the + # heavy diffusers/sd.cpp backends; the active engine NAME comes from module state. + monkeypatch.setattr( + r, + "get_active_diffusion_engine", + lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}), + ) + yield + + +def _set_device(monkeypatch, backend): + monkeypatch.setattr( + r, + "resolve_diffusion_device_target", + lambda: SimpleNamespace(backend = backend, device = backend), + ) + + +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)) + return r.active_engine_name() + + +# ── core selection matrix ───────────────────────────────────────────────────── + + +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) + _set_binary(monkeypatch, "/usr/bin/sd-cli") # even with a binary, GPU stays diffusers + assert _select() == ENGINE_DIFFUSERS + assert "uses diffusers" in (r.active_status()["fallback_reason"] or "") + + +def test_forced_diffusers_overrides_cpu(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "diffusers") + assert _select() == ENGINE_DIFFUSERS + assert "forced" in (r.active_status()["fallback_reason"] or "") + + +def test_sd_cpp_disabled_uses_diffusers(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP", "0") + assert _select() == ENGINE_DIFFUSERS + assert "disabled" in (r.active_status()["fallback_reason"] or "") + + +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. + monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_MPS", "1") + assert _select() == ENGINE_SD_CPP + + +def test_unsupported_family_falls_back(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setattr(r, "family_sd_cpp_supported", lambda fam: False) + assert _select() == ENGINE_DIFFUSERS + assert "no native sd.cpp asset mapping" in (r.active_status()["fallback_reason"] or "") + + +def test_missing_binary_falls_back(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) # install unavailable + assert _select() == ENGINE_DIFFUSERS + assert "binary unavailable" in (r.active_status()["fallback_reason"] or "") + + +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 + + +def test_force_sd_cpp_without_binary_falls_back(monkeypatch): + _set_device(monkeypatch, "cuda") + _set_binary(monkeypatch, None) + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp") + assert _select() == ENGINE_DIFFUSERS + + +# ── active_status annotation ────────────────────────────────────────────────── + + +def test_active_status_injects_engine_and_reason(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) + _select() # -> diffusers fallback (no binary) + st = r.active_status() + assert st["engine"] == ENGINE_DIFFUSERS + assert st["fallback_reason"] and "binary unavailable" in st["fallback_reason"] diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d03c48b477..d0b6997d07 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -114,6 +114,26 @@ def _unloaded_status(): def client(monkeypatch, tmp_path): backend = _FakeBackend() monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Neutralise the engine router so the routes deterministically drive this fake + # (diffusers) backend regardless of the host's real device, and never attempt a + # native sd.cpp install/download. The router's selection logic is covered in + # test_diffusion_engine_router.py; one route-level sd_cpp test lives below. + import core.inference.diffusion_engine_router as engine_router + + # Delegate to whatever get_diffusion_backend currently returns, so per-test + # re-patches of the backend still flow through the routes. + monkeypatch.setattr( + engine_router, + "select_and_activate_engine", + lambda fam, **kw: diffusion_module.get_diffusion_backend(), + ) + monkeypatch.setattr( + engine_router, + "get_active_diffusion_engine", + lambda: diffusion_module.get_diffusion_backend(), + ) + monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers") + monkeypatch.setattr(engine_router, "_fallback_reason", None) # Isolate from the real GPU arbiter: reset ownership and stub the evictors so # the load route's acquire_for() never touches live backend singletons. monkeypatch.setattr(gpu_arbiter, "_owner", None) @@ -501,6 +521,62 @@ def test_out_of_range_cache_threshold_returns_422(client): assert resp.status_code == 422 +def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path): + """End-to-end through the REAL router: a CPU host with an available binary routes + the load to the native sd.cpp engine and the response reports engine=sd_cpp.""" + from types import SimpleNamespace + + import core.inference.diffusion_engine_router as engine_router + import core.inference.sd_cpp_backend as sd_backend + + for e in ( + "UNSLOTH_DIFFUSION_ENGINE", + "UNSLOTH_DIFFUSION_SD_CPP", + "UNSLOTH_DIFFUSION_SD_CPP_MPS", + "UNSLOTH_DIFFUSION_SD_CPP_INSTALL", + ): + monkeypatch.delenv(e, raising = False) + + validator = _FakeBackend() # supplies validate_load_request (and is the diffusers fallback) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: validator) + # Force the router's decision inputs: CPU device + an available binary. + monkeypatch.setattr( + engine_router, + "resolve_diffusion_device_target", + 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. + sd_fake = _FakeBackend() + monkeypatch.setattr(sd_backend, "get_sd_cpp_backend", lambda: sd_fake) + + monkeypatch.setattr(gpu_arbiter, "_owner", None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None) + + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + client = TestClient(app) + + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "z.gguf"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["engine"] == "sd_cpp" + assert body["fallback_reason"] is None + assert sd_fake.loaded is True # the native engine actually received the load + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted. @@ -537,3 +613,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 new file mode 100644 index 0000000000..79464dd6be --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the native sd.cpp diffusion backend (the no-GPU engine).""" + +from __future__ import annotations + +import threading +import types + +import pytest +from PIL import Image + +from core.inference import sd_cpp_backend as bk +from core.inference.diffusion_families import detect_family +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles +from core.inference.sd_cpp_backend import ( + SdCppDiffusionBackend, + _map_guidance, + ensure_sd_cpp_binary, +) +from core.inference.sd_cpp_engine import SdCppCancelled + + +class _FakeEngine: + """Stands in for SdCppEngine: writes a 1x1 PNG and records the args.""" + + def __init__( + self, + *, + fail = None, + cancel_on_call = False, + ): + self.calls = [] + self.fail = fail + self.cancel_on_call = cancel_on_call + + def is_available(self): + return True + + def version(self, **_): + return "fake sd-cli" + + def generate( + self, + files, + params, + *, + output_path, + cancel_event = None, + **kw, + ): + self.calls.append((files, params, output_path, kw)) + if self.cancel_on_call and cancel_event is not None: + cancel_event.set() + if self.fail is not None: + raise self.fail + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("cancelled") + Image.new("RGB", (1, 1), (10, 20, 30)).save(output_path) + from pathlib import Path + + return Path(output_path) + + +def _loaded_backend(fam_name = "z-image", engine = None): + b = SdCppDiffusionBackend(engine = engine or _FakeEngine()) + fam = detect_family(fam_name) + b._state = bk._SdState( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + base_repo = fam.base_repo, + family = fam, + device = "cpu", + files = SdCppModelFiles( + 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 + + +# ── asset resolution ────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "fam_name,expect_kinds", + [ + ("flux.1", {"diffusion_model", "vae", "clip_l", "t5xxl"}), + ("z-image", {"diffusion_model", "vae", "llm"}), + ("qwen-image", {"diffusion_model", "vae", "qwen2vl"}), + ("flux.2-klein", {"diffusion_model", "vae", "llm"}), + ], +) +def test_asset_specs_cover_required_files(fam_name, expect_kinds): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family(fam_name) + specs = b._asset_specs("unsloth/x-GGUF", "x-Q4_K_M.gguf", fam) + kinds = {kind for _, _, kind in specs} + assert kinds == expect_kinds + # Every spec has a non-empty repo + filename. + assert all(repo and fn for repo, fn, _ in specs) + # The transformer reuses the requested GGUF, not a registry file. + tr = [s for s in specs if s[2] == "diffusion_model"][0] + assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf" + + +# ── guidance mapping ────────────────────────────────────────────────────────── + + +def test_map_guidance_flux_uses_distilled_guidance(): + cfg, g = _map_guidance(detect_family("flux.1"), 3.5) + assert cfg is None and g == 3.5 + + +def test_map_guidance_cfg_family_off_when_distilled(): + # qwen-image uses real CFG; a distilled 0 -> CFG off (1.0), a >1 value passes through. + assert _map_guidance(detect_family("qwen-image"), 0.0) == (1.0, None) + assert _map_guidance(detect_family("qwen-image"), 4.0) == (4.0, None) + + +# ── status ──────────────────────────────────────────────────────────────────── + + +def test_status_unloaded_reports_sd_cpp_engine(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + st = b.status() + assert st["loaded"] is False and st["engine"] == "sd_cpp" + + +def test_status_loaded_shape(): + b = _loaded_backend() + st = b.status() + assert st["loaded"] is True + assert st["engine"] == "sd_cpp" + assert st["family"] == "z-image" + assert st["device"] == "cpu" + # diffusers-only fields are present (route response parity) but null. + for k in ("transformer_quant", "attention_backend", "transformer_cache", "text_encoder_quant"): + assert st[k] is None + + +# ── generate ────────────────────────────────────────────────────────────────── + + +def test_generate_returns_images_and_seed(): + eng = _FakeEngine() + b = _loaded_backend(engine = eng) + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 123, batch_size = 2) + assert out["seed"] == 123 + assert out["repo_id"] == "unsloth/Z-Image-Turbo-GGUF" + assert len(out["images"]) == 2 + assert all(isinstance(im, Image.Image) for im in out["images"]) + # One sd-cli run per batch image, each a distinct seed from the base. + 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(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + with pytest.raises(RuntimeError, match = "No diffusion model is loaded"): + b.generate(prompt = "x") + + +def test_generate_passes_vae_format_for_flux2(): + eng = _FakeEngine() + b = _loaded_backend(fam_name = "flux.2-klein", engine = eng) + b.generate(prompt = "x", steps = 4, seed = 1) + _, _, _, kw = eng.calls[0] + assert kw.get("extra_args") == ["--vae-format", "flux2"] + + +def test_generate_cancellation_raises_cancelled_not_failure(): + # The engine cancels mid-run; the backend surfaces a cancellation, not a crash. + eng = _FakeEngine(cancel_on_call = True) + b = _loaded_backend(engine = eng) + with pytest.raises(RuntimeError, match = "cancelled"): + b.generate(prompt = "x", steps = 8, seed = 5) + + +def test_generate_progress_tracks_parsed_steps(): + b = _loaded_backend() + b._gen = bk._SdGen(total_steps = 8) + b._on_log(" sampling 4/8 done") + p = b.generate_progress() + assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8 + # A fraction with a different denominator must not move the bar. + b._on_log("loaded 1/3 tensors") + assert b.generate_progress()["step"] == 4 + + +# ── load validation + binary install ────────────────────────────────────────── + + +def test_begin_load_rejects_unsupported_family(monkeypatch): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + # A family with no native asset mapping must be rejected (router falls back). + monkeypatch.setattr(bk, "family_sd_cpp_supported", lambda fam: False) + with pytest.raises(ValueError, match = "no native sd.cpp asset mapping"): + b.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z.gguf") + + +def test_begin_load_requires_gguf_filename(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "gguf_filename is required"): + b.begin_load("unsloth/Z-Image-Turbo-GGUF") + + +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" + + +def test_ensure_binary_install_disabled_returns_none(monkeypatch): + monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: None) + assert ensure_sd_cpp_binary(allow_install = False) is None + + +def test_unload_clears_state_and_signals_cancel(): + cancel = threading.Event() + b = _loaded_backend() + b._active_generate_cancel = cancel + st = b.unload() + 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" diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index a879252761..054574d905 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -78,7 +78,10 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): # ── availability / version ────────────────────────────────────────────────── -def test_engine_unavailable_when_no_binary(): +def test_engine_unavailable_when_no_binary(monkeypatch): + # Force the "no binary anywhere" condition so the test is hermetic even on a host + # that happens to have sd-cli installed. + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) assert e.is_available() is False assert e.version() is None @@ -92,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" @@ -350,12 +355,13 @@ def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): assert "--upscale-model" in _FakePopen.captured_cmd -def test_upscale_raises_when_binary_missing(): +def test_upscale_raises_when_binary_missing(monkeypatch, tmp_path): + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) with pytest.raises(RuntimeError, match = "not found"): e.upscale( SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), - output_path = "/tmp/x.png", + output_path = str(tmp_path / "x.png"), )