Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine
When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the native stable-diffusion.cpp engine instead of diffusers, with diffusers as the guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM. - diffusion_engine_router: centralised engine selection (built on the existing select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason. - sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface backed by sd-cli, with lazy binary install, registry-driven asset fetch, step-progress parsing, and cancellation. - diffusion_families: per-family single-file VAE + text-encoder asset mapping. - sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled). - routes/inference + gpu_arbiter: drive the active engine via the router; the API now reports the active engine and any fallback reason. - tests for the backend, router, route selection, and cancellation.
This commit is contained in:
Daniel Han
parent
9c5b35c710
commit
7f3c206fa1
11 changed files with 1241 additions and 27 deletions
139
studio/backend/core/inference/diffusion_engine_router.py
Normal file
139
studio/backend/core/inference/diffusion_engine_router.py
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
# 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, 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
|
||||
with _lock:
|
||||
_active_engine_name = name
|
||||
_fallback_reason = reason if name == ENGINE_DIFFUSERS else None
|
||||
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())
|
||||
|
||||
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())
|
||||
|
|
@ -44,6 +44,20 @@ 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)
|
||||
|
||||
|
||||
# Keyed by architecture, not per model variant: a checkpoint's specific base repo
|
||||
|
|
@ -59,6 +73,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
|
||||
|
|
@ -69,6 +88,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",
|
||||
|
|
@ -77,6 +104,12 @@ _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"),
|
||||
),
|
||||
),
|
||||
DiffusionFamily(
|
||||
name = "z-image",
|
||||
|
|
@ -86,6 +119,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"),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -130,6 +167,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.
|
||||
|
||||
|
|
|
|||
|
|
@ -51,8 +51,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.
|
||||
|
|
|
|||
553
studio/backend/core/inference/sd_cpp_backend.py
Normal file
553
studio/backend/core/inference/sd_cpp_backend.py
Normal file
|
|
@ -0,0 +1,553 @@
|
|||
# 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.sd_cpp_args import SdCppGenParams, SdCppModelFiles
|
||||
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
|
||||
|
||||
|
||||
@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."""
|
||||
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.")
|
||||
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,
|
||||
_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],
|
||||
_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
|
||||
state = _SdState(
|
||||
repo_id = repo_id,
|
||||
base_repo = base,
|
||||
family = fam,
|
||||
device = device,
|
||||
files = files,
|
||||
vae_format = fam.sd_cpp_vae_format,
|
||||
native_speed = "off",
|
||||
offload_flags = (),
|
||||
threads = None,
|
||||
)
|
||||
# Probe the binary version (cheap) so a broken install fails the load.
|
||||
engine.version()
|
||||
with self._lock:
|
||||
if self._load_token != _load_token:
|
||||
return # superseded / cancelled
|
||||
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, _ in assets:
|
||||
if 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)
|
||||
vae_extra = (
|
||||
["--vae-format", state.vae_format] if state.vae_format else None
|
||||
)
|
||||
|
||||
self._gen = _SdGen(total_steps = int(steps))
|
||||
images = []
|
||||
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.
|
||||
seed_i = (seed + index) & ((1 << 53) - 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,
|
||||
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 = vae_extra,
|
||||
on_log = self._on_log,
|
||||
cancel_event = cancel,
|
||||
)
|
||||
with Image.open(out_path) as im:
|
||||
images.append(im.copy())
|
||||
if cancel.is_set():
|
||||
raise RuntimeError("Diffusion generation was cancelled.")
|
||||
return {"images": images, "seed": int(seed), "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",
|
||||
"cpu_offload": False,
|
||||
"offload_policy": "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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
@ -199,6 +224,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 +232,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 +249,10 @@ 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 +264,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 +274,10 @@ 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 +303,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 +326,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 +353,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 +380,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:]))
|
||||
|
|
|
|||
|
|
@ -1904,3 +1904,10 @@ 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)",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -10058,24 +10058,32 @@ async def load_diffusion_model(
|
|||
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
|
||||
):
|
||||
from core.inference.diffusion import get_diffusion_backend
|
||||
from core.inference.diffusion_engine_router import annotate_status, select_and_activate_engine
|
||||
from core.inference.gpu_arbiter import acquire_for, DIFFUSION
|
||||
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.
|
||||
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,
|
||||
family_override = request.family_override,
|
||||
)
|
||||
# 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
|
||||
)
|
||||
# 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)
|
||||
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,
|
||||
|
|
@ -10092,7 +10100,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:
|
||||
|
|
@ -10105,9 +10113,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,
|
||||
|
|
@ -10221,27 +10229,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())
|
||||
|
|
|
|||
136
studio/backend/tests/test_diffusion_engine_router.py
Normal file
136
studio/backend/tests/test_diffusion_engine_router.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
# 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 _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")
|
||||
assert _select() == ENGINE_SD_CPP
|
||||
assert r.active_engine_name() == ENGINE_SD_CPP
|
||||
|
||||
|
||||
@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")
|
||||
# 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")
|
||||
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"]
|
||||
|
|
@ -114,6 +114,24 @@ 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)
|
||||
|
|
@ -423,6 +441,54 @@ 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")
|
||||
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.
|
||||
|
|
|
|||
210
studio/backend/tests/test_sd_cpp_backend.py
Normal file
210
studio/backend/tests/test_sd_cpp_backend.py
Normal file
|
|
@ -0,0 +1,210 @@
|
|||
# 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 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,
|
||||
)
|
||||
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]
|
||||
|
||||
|
||||
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()
|
||||
|
|
@ -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
|
||||
|
|
@ -350,12 +353,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"),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue