ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'diffusion-train-precision' into diffusion-train-tab-2

Browse source
This commit is contained in:
Daniel Han 2026-07-04 01:04:40 +00:00
commit 9b8132ec9e
23 changed files with 3942 additions and 213 deletions
Download patch file Download diff file Expand all files Collapse all files

40
scripts/build_prequant_checkpoint.py
View file

@ -55,6 +55,7 @@ def main(argv = None) -> int:
# Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant).
from core.inference.diffusion_transformer_quant import (
FP8_GRANULARITY,
TQ_FP8,
TQ_SCHEMES,
_make_quant_config,
@ -100,25 +101,30 @@ def main(argv = None) -> int:
k: (v.detach().to("cpu") if hasattr(v, "detach") else v)
for k, v in transformer.state_dict().items()
}
metadata = {
"base_model_id": args.base,
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
# The layers skipped for this scheme (int8's M=1 modulation projections; () for
# the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the
# loader reject a checkpoint that would not match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"fast_accum": fast_accum,
"torch_dtype": args.dtype,
"quant_backend": "torchao",
"transformer_class": fam.transformer_class,
"torch_version": torch.__version__,
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
}
# Record the fp8 granularity so the loader can reject a stale per-tensor checkpoint
# (the runtime now requires per-row; see FP8_GRANULARITY).
if scheme == TQ_FP8:
metadata["fp8_granularity"] = FP8_GRANULARITY
ckpt = {
"format": PREQUANT_FORMAT,
"metadata": {
"base_model_id": args.base,
"family": fam.name,
"scheme": scheme,
"min_features": args.min_features,
# The layers skipped for this scheme (int8's M=1 modulation projections; () for
# the scaled_mm schemes) and, for fp8, the baked accumulate mode. Both let the
# loader reject a checkpoint that would not match the runtime path.
"exclude_name_tokens": list(exclude_name_tokens),
"fast_accum": fast_accum,
"torch_dtype": args.dtype,
"quant_backend": "torchao",
"transformer_class": fam.transformer_class,
"torch_version": torch.__version__,
"torchao_version": getattr(torchao, "__version__", "?"),
"diffusers_version": diffusers.__version__,
},
"metadata": metadata,
"state_dict": state_dict,
}

42
studio/backend/core/inference/diffusion_engine_router.py
View file

@ -29,7 +29,12 @@ 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_backend import (
_install_allowed,
_server_binary_runnable,
ensure_sd_cpp_binary,
ensure_sd_server_binary,
)
from core.inference.sd_cpp_engine import (
ENGINE_DIFFUSERS,
ENGINE_SD_CPP,
@ -148,20 +153,37 @@ def select_and_activate_engine(
fam_ok = family_sd_cpp_supported(fam)
binary = None
server_binary = None
if policy_eligible and fam_ok:
binary = ensure_sd_cpp_binary(
# Probe the resident sd-server FIRST: the backend PREFERS it, and an sd-server-only
# install (no sd-cli) must still route to native rather than silently falling back to
# diffusers. Checking it before the sd-cli install also means a server-only host does
# not pay an avoidable sd-cli download. Install the accelerator-matched build (ROCm /
# Vulkan / CUDA) so a forced-native GPU load gets the GPU server, not the CPU one.
server_binary = ensure_sd_server_binary(
allow_install = _install_allowed(),
accelerator = _install_accelerator_for(backend),
)
# 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 server_binary and not _server_binary_runnable(server_binary):
logger.warning(
"sd-server at %s is present but not runnable; not using it", server_binary
)
server_binary = None
# sd-cli is the one-shot fallback. Always LOCATE an existing binary, but only
# auto-INSTALL it when there is no usable server, so a server-only install is not
# forced to also download a CLI it will never use. Probe runnability before
# committing 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.
binary = ensure_sd_cpp_binary(
allow_install = _install_allowed() and server_binary is None,
accelerator = _install_accelerator_for(backend),
)
if binary and SdCppEngine(binary = binary).version() is None:
logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary)
logger.warning("sd-cli at %s is present but not runnable; not using it", binary)
binary = None
native_available = bool(binary) and policy_eligible and fam_ok
native_available = bool(binary or server_binary) and policy_eligible and fam_ok
choice = select_diffusion_engine(
backend, native_available = native_available, prefer_native = prefer_native
)
@ -173,8 +195,8 @@ def select_and_activate_engine(
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"
elif not (binary or server_binary):
reason = "native sd.cpp binary unavailable"
else:
reason = "diffusers selected"
return _activate(ENGINE_DIFFUSERS, reason)

32
studio/backend/core/inference/diffusion_families.py
View file

@ -421,6 +421,38 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str:
return base or fam.base_repo
# Default (steps, guidance) per model for callers that can't pass them — namely
# the OpenAI /v1/images/generations endpoint, whose spec has no step/guidance
# knobs. Distilled "turbo/schnell" models want few steps and no CFG; the full
# "dev" models want more steps and real CFG. Matched by substring, most specific
# first — the same scheme and values as the UI's MODEL_DEFAULTS table
# (studio/frontend/src/features/images/images-page.tsx); keep the two in sync.
_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = (
("z-image-turbo", 9, 0.0),
("flux.1-schnell", 4, 0.0),
("flux.1", 28, 3.5),
("flux.2-klein", 4, 0.0),
("qwen-image", 20, 4.0),
("z-image", 20, 4.0),
)
# Unrecognised model: distilled few-step / no-CFG shape, matching the UI fallback.
_GENERATION_DEFAULT_FALLBACK = (9, 0.0)
def default_generation_params(*identifiers: Optional[str]) -> tuple[int, float]:
"""Default ``(steps, guidance)`` for a loaded model. The first identifier that
names a known model wins (the repo id, then the resolved base repo), so a
local-path load — whose repo id is just a filesystem path that may not name
the model — still resolves via its base repo. Within an identifier, keys are
matched as substrings, most specific first (the same scheme as the UI)."""
for identifier in identifiers:
needle = (identifier or "").lower()
for key, steps, guidance in _GENERATION_DEFAULTS:
if key in needle:
return steps, guidance
return _GENERATION_DEFAULT_FALLBACK
def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]:
"""The hosted pre-quantized transformer repo for ``scheme`` in this family, or None."""
for entry_scheme, repo_id in fam.prequant_repos:

16
studio/backend/core/inference/diffusion_prequant.py
View file

@ -268,6 +268,22 @@ def _validate_checkpoint(
if meta.get("scheme") != scheme:
_warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
return False
# fp8 REQUIRES per-row granularity (per-tensor collapses outlier-heavy DiTs to noise). A
# checkpoint built before that fix carries the old per-tensor layout and either omits
# ``fp8_granularity`` or records something other than per-row; reject it so the loader
# falls back to rebuilding/re-quantising instead of installing a broken fp8 transformer.
from .diffusion_transformer_quant import FP8_GRANULARITY, TQ_FP8
if scheme == TQ_FP8 and meta.get("fp8_granularity") != FP8_GRANULARITY:
_warn(
logger,
scheme,
ValueError(
f"fp8 checkpoint granularity {meta.get('fp8_granularity')!r} != "
f"{FP8_GRANULARITY!r} (stale per-tensor artifact); rebuild it"
),
)
return False
ckpt_base = meta.get("base_model_id")
if base:
# A checkpoint whose keys happen to match a different base can load strict=True and

34
studio/backend/core/inference/diffusion_transformer_quant.py
View file

@ -37,6 +37,12 @@ TQ_AUTO = "auto"
TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8)
TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES
# The fp8 weight/activation granularity the runtime config uses (see _make_quant_config):
# per-ROW is REQUIRED for correctness on outlier-heavy DiTs. Stamped into a pre-quantized
# fp8 checkpoint's metadata at build time and required by the loader, so a stale checkpoint
# baked with the old per-TENSOR layout is rejected and rebuilt rather than reproducing noise.
FP8_GRANULARITY = "per_row"
# Skip linears whose in/out features are below this. A small share of the FLOPs, so leaving
# them bf16 costs ~nothing and keeps quality a touch higher.
DEFAULT_MIN_LINEAR_FEATURES = 512
@ -298,19 +304,29 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
if scheme == TQ_INT8:
return Int8DynamicActivationInt8WeightConfig()
if scheme == TQ_FP8:
# Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation
# cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate
# than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real
# win there. Data-center HBM parts default to the higher-precision accumulate.
# fast accumulate is a precision (not overflow) tradeoff and stays below the fp8
# quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations).
# Per-ROW granularity (per-token activation scale + per-output-channel weight scale)
# is REQUIRED for correctness, not just quality. torchao's default fp8 granularity is
# per-TENSOR: one scale for the whole activation tensor. A DiT with extreme activation
# outliers breaks under that -- z-image's MLP activations peak near 6.6e4, and a single
# such outlier forces a tensor-wide scale that pushes every normal value (~1-30) below
# the fp8 resolution, so they quantise to ~0 and the denoise collapses to pure noise
# (measured on B200: per-tensor fp8 = noise, per-row fp8 = matches bf16). Per-row
# confines each outlier to its own token/channel. This is also why int8 dynamic was
# always fine here: it is per-token by default. The per-row scaled_mm is probed by
# _smoke_probe, so an arch/build without it falls through the ladder to int8.
#
# fast accumulate (fp8 only) is chosen by GPU class unless forced: consumer / workstation
# cards (GDDR) run the fp8 tensor cores ~2x faster with FP16 (fast) accumulate than FP32
# (e.g. ~838 vs ~419 TFLOPS on RTX 50xx); data-center HBM parts keep precise accumulate.
from torchao.quantization import PerRow
try:
from torchao.float8 import Float8MMConfig
return Float8DynamicActivationFloat8WeightConfig(
mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum))
granularity = PerRow(),
mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)),
)
except Exception: # noqa: BLE001 — older torchao without the explicit knob
return Float8DynamicActivationFloat8WeightConfig()
except Exception: # noqa: BLE001 — older torchao without the explicit mm knob
return Float8DynamicActivationFloat8WeightConfig(granularity = PerRow())
if scheme == TQ_NVFP4:
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig

136
studio/backend/core/inference/sd_cpp_args.py
View file

@ -326,6 +326,142 @@ def build_sd_cpp_upscale_command(
return cmd
def build_sd_cpp_server_command(
binary: str,
files: SdCppModelFiles,
*,
host: str,
port: int,
vae_format: Optional[str] = None,
offload: Optional[list[str]] = None,
native_speed: Optional[str] = None,
threads: Optional[int] = None,
scratch_dir: Optional[str] = None,
verbose: bool = False,
extra_args: Optional[list[str]] = None,
) -> list[str]:
"""Build the ``sd-server`` argv: model + hardware/server flags only.
``sd-server`` (stable-diffusion.cpp ``examples/server``) loads the model once at
spawn from the SAME flags ``sd-cli`` takes (``--diffusion-model`` / ``--vae`` /
the text encoders / ``--vae-format`` / offload + speed), and adds ``--listen-ip``
/ ``--listen-port``. Per-generation parameters (prompt, size, steps, seed, cfg,
sampler, batch) are NOT here -- they go in each ``/sdcpp/v1/img_gen`` request, so
one resident process serves many generations without reloading the weights.
``offload`` / ``native_speed`` map to the exact same sd.cpp flags as the one-shot
engine (``--offload-to-cpu`` / ``--diffusion-fa`` / ...), verified to be accepted
by ``sd-server --help``. ``scratch_dir`` (if given) is pointed at by the LoRA /
hires-upscaler / embeddings directory flags: sd-server's img_gen handler recursively
iterates those dirs, and an unset / missing dir makes it fail the request, so we give
it a real (empty) directory. ``extra_args`` is appended last so a power user can
override anything (sd.cpp's parser is last-wins).
"""
if not files.diffusion_model:
raise ValueError("diffusion_model path is required")
cmd: list[str] = [binary, "--diffusion-model", files.diffusion_model]
for flag, value in (
("--vae", files.vae),
("--clip_l", files.clip_l),
("--clip_g", files.clip_g),
("--t5xxl", files.t5xxl),
("--llm", files.llm),
("--qwen2vl", files.qwen2vl),
):
if value:
cmd += [flag, value]
if vae_format:
cmd += ["--vae-format", vae_format]
cmd += ["--listen-ip", str(host), "--listen-port", str(int(port))]
if scratch_dir:
cmd += [
"--lora-model-dir",
scratch_dir,
"--hires-upscalers-dir",
scratch_dir,
"--embd-dir",
scratch_dir,
]
if threads is not None:
cmd += ["--threads", str(int(threads))]
offload = list(offload or [])
if offload:
cmd += offload
# De-dup speed flags against offload (offload may already include --diffusion-fa).
cmd += [f for f in native_speed_flags(native_speed) if f not in offload]
if verbose:
cmd += ["-v"]
if extra_args:
cmd += list(extra_args)
return cmd
def build_img_gen_request(
*,
prompt: str,
negative_prompt: Optional[str] = None,
width: int = 1024,
height: int = 1024,
steps: Optional[int] = None,
seed: Optional[int] = None,
batch_count: int = 1,
sample_method: Optional[str] = None,
flow_shift: Optional[float] = None,
cfg_scale: Optional[float] = None,
distilled_guidance: Optional[float] = None,
output_format: str = "png",
lora: Optional[list[dict]] = None,
) -> dict:
"""Build the ``POST /sdcpp/v1/img_gen`` JSON body for one text-to-image request.
The native ``sdcpp`` API takes the whole batch in one request (``batch_count``),
so a batch reuses the resident model with no reload. Sampling lives under
``sample_params``; guidance is split exactly like the one-shot engine's
``_map_guidance``: a FLUX distilled value goes to ``guidance.distilled_guidance``,
a real classifier-free scale goes to ``guidance.txt_cfg``. Only set keys are
emitted so the server applies its own defaults for the rest.
"""
if not str(prompt).strip():
raise ValueError("prompt is required")
guidance: dict = {}
if cfg_scale is not None:
guidance["txt_cfg"] = float(cfg_scale)
if distilled_guidance is not None:
guidance["distilled_guidance"] = float(distilled_guidance)
sample_params: dict = {}
if steps is not None:
sample_params["sample_steps"] = int(steps)
if sample_method:
sample_params["sample_method"] = str(sample_method)
if flow_shift is not None:
sample_params["flow_shift"] = float(flow_shift)
if guidance:
sample_params["guidance"] = guidance
req: dict = {
"prompt": prompt,
"negative_prompt": negative_prompt or "",
"width": int(width),
"height": int(height),
"batch_count": max(1, int(batch_count)),
"output_format": output_format,
}
if seed is not None:
req["seed"] = int(seed)
if sample_params:
req["sample_params"] = sample_params
# Structured LoRA list: the sdcpp API resolves each ``path`` against the server's
# scanned ``--lora-model-dir`` (prompt-embedded ``<lora:>`` tags are intentionally
# unsupported server-side), so LoRAs must be staged there and named here.
if lora:
req["lora"] = lora
return req
def _fmt_float(value: float) -> str:
"""Compact float -> str: drop a trailing ``.0`` so ``1.0`` -> ``1`` (sd-cli
accepts both, but the tidy form keeps logged commands readable)."""

628
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -50,13 +50,22 @@ from core.inference.diffusion_memory import (
OFFLOAD_NONE,
OFFLOAD_SEQUENTIAL,
)
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags
from core.inference.sd_cpp_args import (
SdCppGenParams,
SdCppModelFiles,
build_img_gen_request,
offload_flags,
)
from core.inference.sd_cpp_engine import (
SdCppCancelled,
SdCppEngine,
find_sd_cpp_binary,
find_sd_server_binary,
runtime_env,
)
from core.inference.sd_cpp_server import SdCppServer
from loggers import get_logger
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
@ -69,6 +78,45 @@ _STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)")
# download / extract / chmod.
_install_lock = threading.Lock()
# sd-server accepts at most this many images per img_gen job; larger Studio batches
# (the request model allows up to 32) are split into chunks of this size, the way the
# one-shot path did them one image at a time.
_MAX_SERVER_BATCH = 8
# Per-image wall-clock budget for a server job, so a batch gets a timeout proportional to
# its image count (matching the one-shot path, where each image had its own budget) rather
# than one fixed deadline the whole batch has to finish within.
_SERVER_PER_IMAGE_TIMEOUT_S = 1800.0
def _server_binary_runnable(binary: str) -> bool:
"""Best-effort probe that ``binary`` can actually execute (not just exist).
Runs ``<binary> --help`` with the same runtime env the server will use, so a present
but unrunnable build (wrong arch, missing shared libs, no execute bit) is caught before
a multi-GB asset download. Conservative: only a clear "cannot launch" signal (OSError,
or the dynamic-loader exit codes 126/127) returns False; anything else is treated as
runnable so a quirky ``--help`` exit code never blocks a working binary."""
import subprocess
try:
proc = subprocess.run(
[binary, "--help"],
capture_output = True,
timeout = 20,
env = runtime_env(binary),
**windows_hidden_subprocess_kwargs(),
)
except OSError:
return False # cannot exec at all (wrong arch / no execute bit / missing loader)
except Exception: # noqa: BLE001 -- timeout or anything odd: don't block on a flaky probe
return True
# A negative return code is a signal death (e.g. -4 SIGILL from an incompatible
# prebuilt on an older CPU): the binary launches but immediately crashes, so treat it
# as unavailable and let the load fall back to diffusers instead of routing to a
# server that will die on startup.
return proc.returncode >= 0 and proc.returncode not in (126, 127)
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.
@ -106,9 +154,51 @@ def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu"
return None
def ensure_sd_server_binary(
*, allow_install: bool = True, accelerator: str = "cpu"
) -> Optional[str]:
"""Path to a usable ``sd-server`` binary, installing the prebuilt once if needed.
Unlike ``ensure_sd_cpp_binary``, this installs when *sd-server specifically* is
missing -- even if an ``sd-cli`` from an older install is already present -- so an
existing one-shot install is upgraded to the persistent server (the prebuilt archive
ships both). Returns None when it is absent and cannot be installed; the backend then
uses the one-shot fallback. Never raises.
"""
found = find_sd_server_binary()
if found:
return found
if not allow_install:
return None
with _install_lock:
found = find_sd_server_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-server installer import failed: %s", exc)
return None
try:
_install(accelerator = accelerator) # extracts sd-cli AND sd-server
except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back
logger.warning("sd-server auto-install failed: %s", exc)
return None
return find_sd_server_binary()
@dataclass(frozen = True)
class _SdState:
"""The loaded native checkpoint: resolved asset paths + run settings."""
"""The loaded native checkpoint: resolved asset paths + run settings.
``server`` is the resident ``sd-server`` process (the model is loaded once, inside
it) when ``mode == "server"``; in the ``"oneshot"`` fallback it is ``None`` and each
generation re-runs ``sd-cli``."""
repo_id: str
base_repo: str
@ -121,6 +211,8 @@ class _SdState:
threads: Optional[int] = None
sampling_method: Optional[str] = None
flow_shift: Optional[float] = None
server: Optional[SdCppServer] = None
mode: str = "server"
# Token kept so LoRA adapters selected at generate time can be fetched from the Hub.
hf_token: Optional[str] = None
@ -194,11 +286,19 @@ class SdCppDiffusionBackend:
self._lock = threading.Lock()
self._generate_lock = threading.Lock()
self._engine = engine # resolved lazily on first load so import stays cheap
# An engine passed in is an EXPLICIT injection (the test seam / escape hatch) and
# pins one-shot mode; an engine cached later by a runtime fallback must NOT, so a
# now-available server can still be used on the next load.
self._engine_injected = engine is not None
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
# The sd-server being started for an in-flight load, before it is committed to
# _state. Tracked so an unload / superseding load can stop it mid-startup instead
# of leaving it loading (and holding the generate lock) for the whole timeout.
self._pending_server: Optional[SdCppServer] = None
self._gen: Optional[_SdGen] = None
@property
@ -215,6 +315,38 @@ class SdCppDiffusionBackend:
self._engine = SdCppEngine(binary = binary)
return self._engine
def _resolve_backend(self) -> tuple[str, Optional[str], Optional[SdCppEngine]]:
"""Pick the native execution mode: ("server", binary, None) or ("oneshot", None, engine).
The persistent ``sd-server`` is preferred (load once, serve many). The one-shot
``sd-cli`` is the fallback for older / custom builds that lack the server target.
An explicitly injected engine forces one-shot (the unit-test seam and an escape
hatch), so a test never spawns a real server or triggers an install. A lazily
cached fallback engine does NOT force one-shot: once a resident server becomes
available (installed, or a per-model start that previously failed now works), the
next load can use it, instead of being pinned to one-shot for the whole session.
"""
if self._engine_injected and self._engine is not None:
return "oneshot", None, self._resolve_engine()
# Install the sd-server build matching the resolved device backend (ROCm / Vulkan /
# CUDA), not the default CPU build: a forced/enabled native load on a GPU host must
# not silently fetch the plain-CPU server. Lazy import avoids an import cycle with
# the router, which imports this backend during engine selection.
from core.inference.diffusion_engine_router import _install_accelerator_for
accelerator = _install_accelerator_for(
getattr(resolve_diffusion_device_target(), "backend", "cpu")
)
server_binary = ensure_sd_server_binary(
allow_install = _install_allowed(), accelerator = accelerator
)
if server_binary is not None:
return "server", server_binary, None
logger.warning(
"sd-server not found; falling back to one-shot sd-cli (reloads the model per image)."
)
return "oneshot", None, self._resolve_engine()
# ── Background load + progress ─────────────────────────────────────────
def begin_load(
@ -308,10 +440,35 @@ class SdCppDiffusionBackend:
_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()
# Resolve the backend mode (persistent sd-server preferred, one-shot sd-cli
# fallback) and binary up front so an install / missing-binary failure
# surfaces before the multi-GB asset pull.
mode, server_binary, engine = self._resolve_backend()
if mode == "server":
# Probe the server binary before the multi-GB asset pull: a present but
# unrunnable build (wrong arch / missing libs) would otherwise download
# everything and only then fail to start. If it cannot run, fall back to
# the one-shot engine now (when it is usable), else surface the failure.
assert server_binary is not None
if not _server_binary_runnable(server_binary):
logger.warning(
"sd-server at %s is present but not runnable; trying one-shot sd-cli.",
server_binary,
)
try:
usable = self._resolve_engine().version() is not None
except Exception: # noqa: BLE001
usable = False
if not usable:
raise RuntimeError("sd-server binary is present but not runnable.")
mode, server_binary, engine = "oneshot", None, self._resolve_engine()
if mode == "oneshot":
# Probe the binary: version() returns None when the present binary cannot
# run (bad perms / missing libs), so fail now rather than commit a "ready"
# state that crashes on the first generation.
assert engine is not None
if engine.version() is None:
raise RuntimeError("sd-cli binary is present but not runnable.")
assets = self._asset_specs(repo_id, gguf_filename, fam)
self._set_expected_bytes(assets, hf_token)
@ -328,38 +485,21 @@ class SdCppDiffusionBackend:
)
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.
# sd_cpp / MPS), since on CPU the weights are 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,
hf_token = hf_token,
)
# 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.
native_speed = _native_speed_for(speed_mode)
# Tear down any previously-loaded model, then commit the new one. A generation
# that started during the (slow) asset download is still running against the OLD
# model: abort it and WAIT on _generate_lock for it to exit before swapping, or
# a stale run could finish afterward and persist an image from the previous
# model. For server mode we stop the old server and start (load) the new one
# HERE, under _generate_lock, so generation never races a half-loaded server and
# two resident models never coexist. _generate_lock is taken only now, not during
# the download, so the long fetch never serialises against generation.
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled
@ -369,6 +509,79 @@ class SdCppDiffusionBackend:
with self._lock:
if self._load_token != _load_token:
return # superseded / cancelled while waiting
old_state = self._state
self._state = None # the old model is being torn down
if old_state is not None and old_state.server is not None:
old_state.server.stop()
server: Optional[SdCppServer] = None
if mode == "server":
assert server_binary is not None
server = SdCppServer(server_binary)
# Publish the not-yet-committed server so unload() / a superseding load
# can stop it mid-startup (SdCppServer.stop aborts the readiness wait
# without waiting on the lifecycle lock), instead of it loading for the
# full startup timeout while holding the generate lock.
with self._lock:
self._pending_server = server
try:
# Blocks until the server has loaded the model and is answering
# (its readiness check); raises with the log tail on a failed load.
server.start(
files,
vae_format = fam.sd_cpp_vae_format,
offload = list(offload),
native_speed = native_speed,
threads = None,
)
except SdCppCancelled:
# Startup was aborted by an unload / superseding load: stop the
# half-started server and bail (the outer handler returns cleanly).
server.stop()
raise
except Exception as start_exc: # noqa: BLE001
# A present-but-unusable sd-server must be no worse than the
# one-shot engine: fall back to sd-cli when it is usable, else
# surface the server error.
logger.warning(
"sd-server failed to start (%s); falling back to one-shot sd-cli.",
start_exc,
)
server.stop()
server = None
try:
usable = self._resolve_engine().version() is not None
except Exception: # noqa: BLE001
usable = False
if not usable:
raise start_exc
mode = "oneshot"
finally:
with self._lock:
if self._pending_server is server:
self._pending_server = None
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,
offload_flags = offload,
threads = None,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
server = server,
mode = mode,
hf_token = hf_token,
)
with self._lock:
if self._load_token != _load_token:
# Superseded / unloaded while we were loading: discard the server
# we just started so it doesn't leak (and keep _state unloaded).
if server is not None:
server.stop()
return
self._state = state
self._loading = None
except SdCppCancelled:
@ -496,11 +709,13 @@ class SdCppDiffusionBackend:
upscale: Optional[float] = None,
# Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity.
reference_images: Optional[list[str]] = None,
# LoRA adapters as (id, weight) pairs; resolved + materialized into a managed dir
# and injected as <lora:ALIAS:w> prompt tags per sd-cli run. None/empty = no LoRA.
# LoRA adapters as (id, weight) pairs; resolved up front, then applied per engine
# path: <lora:ALIAS:w> prompt tags for one-shot sd-cli, structured `lora` entries
# for the resident sd-server. None/empty = no LoRA.
loras: Optional[list[tuple[str, float]]] = None,
# ControlNet: accepted for interface parity with the diffusers backend. Native sd.cpp
# ControlNet (sd-cli --control-net) is a follow-up; a request is rejected clearly.
# Accepted for the uniform engine interface; the guard below rejects it on the native
# engine (ControlNet is diffusers-only) like img2img/inpaint, so a direct API call with
# ControlNet set fails clearly instead of TypeError'ing on an unexpected kwarg.
controlnet: Optional[tuple[str, str, str, float, float, float]] = None,
) -> dict[str, Any]:
import tempfile
@ -534,8 +749,17 @@ class SdCppDiffusionBackend:
state = self._state
if state is None:
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
# A resident server can exit while idle; if a client generates without first
# polling status, drop the stale loaded state and report not-loaded so it gets
# the recoverable reload path instead of a 500 from img_gen (not running).
if (
state.mode == "server"
and state.server is not None
and not state.server.is_alive()
):
self._state = None
raise RuntimeError(DIFFUSION_NOT_LOADED_MSG)
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)
@ -543,11 +767,10 @@ class SdCppDiffusionBackend:
seed = int(seed)
cfg_scale, flux_guidance = _map_guidance(state.family, guidance)
# Resolve any selected LoRA adapters up front (downloads land in the HF
# cache; a bad id fails here as a clear 400 before we spawn sd-cli).
# Drop weight-0 rows BEFORE the support gate: LoraSpec documents weight 0
# as disabling the adapter (the diffusers path treats it as empty), so a
# request carrying only disabled rows must stay a no-op even on a family
# where native LoRA is unsupported, rather than 400 on a dead selection.
# cache; a bad id fails here as a clear 400 before we generate). Drop
# weight-0 rows BEFORE the support gate: weight 0 disables an adapter, so a
# request carrying only disabled rows stays a no-op even on a family where
# native LoRA is unsupported, rather than 400 on a dead selection.
lora_resolved: list = []
active_loras = [(i, w) for (i, w) in (loras or []) if w != 0]
if active_loras:
@ -564,71 +787,41 @@ class SdCppDiffusionBackend:
lora_resolved = diffusion_lora.resolve_specs(
active_loras, hf_token = state.hf_token, cancel_event = cancel
)
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:
# Materialize selected LoRAs into a managed dir sd-cli can scan, and
# inject matching <lora:ALIAS:w> tags into the prompt (deduped against
# any the user typed). Empty -> prompt/dir unchanged.
eff_prompt = prompt
lora_dir: Optional[str] = None
if lora_resolved:
materialized = diffusion_lora.materialize_native_dir(
lora_resolved, Path(tmpdir) / "loras"
)
eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized)
lora_dir = str(Path(tmpdir) / "loras")
for index in range(max(1, int(batch_size))):
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# 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 = eff_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,
lora_dir = lora_dir,
lora_apply_mode = "auto" if lora_dir else None,
)
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 state.mode == "server" and state.server is not None:
images, seeds = self._generate_server(
state,
prompt = prompt,
negative_prompt = negative_prompt,
width = width,
height = height,
steps = steps,
seed = seed,
batch_size = batch_size,
cfg_scale = cfg_scale,
flux_guidance = flux_guidance,
lora_resolved = lora_resolved,
cancel = cancel,
)
else:
images, seeds = self._generate_oneshot(
state,
prompt = prompt,
negative_prompt = negative_prompt,
width = width,
height = height,
steps = steps,
seed = seed,
batch_size = batch_size,
cfg_scale = cfg_scale,
flux_guidance = flux_guidance,
lora_resolved = lora_resolved,
cancel = cancel,
)
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# ``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.
# ``seeds`` is the per-image seed (image i used seed+i), so the route can
# persist the real seed for every image in the batch.
return {
"images": images,
"seed": int(seed),
@ -643,6 +836,195 @@ class SdCppDiffusionBackend:
if self._active_generate_cancel is cancel:
self._active_generate_cancel = None
def _generate_server(
self,
state: _SdState,
*,
prompt: str,
negative_prompt: Optional[str],
width: int,
height: int,
steps: int,
seed: int,
batch_size: int,
cfg_scale: Optional[float],
flux_guidance: Optional[float],
lora_resolved: list,
cancel: threading.Event,
) -> tuple[list, list[int]]:
"""Generate via the resident sd-server (no model reload).
A batch larger than the server's per-job limit is split into chunks: the server
rejects a batch_count above _MAX_SERVER_BATCH, and the one-shot path served large
batches image-by-image, so preserve that. The base seed is masked to sd.cpp's
signed-int64 range (the request model / diffusers accept larger seeds), and each
chunk is submitted at base+offset so the per-image seeds stay reproducible. Each
chunk gets a timeout proportional to its image count so a slow CPU batch is not
cancelled partway through on one fixed deadline.
LoRA on the server goes through the structured ``lora`` request field, NOT prompt
tags (the sdcpp API intentionally ignores ``<lora:>`` in the prompt). Selected
adapters are staged into the server's ``--lora-model-dir`` scratch dir, which the
server rescans per request, and referenced by their staged filename."""
import io
import os
import shutil
from PIL import Image
from core.inference import diffusion_lora
assert state.server is not None
total = max(1, int(batch_size))
# sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a
# large explicit seed is not rejected / wrapped inconsistently by the server.
base_seed = int(seed) & ((1 << 63) - 1)
images: list = []
seeds: list[int] = []
# Stage selected LoRAs into a per-request subdir of the server's lora-model-dir so a
# previous request's adapters can't leak into this one; reference them by the path
# relative to that dir (what the server's recursive scan resolves against). The
# subdir is removed after the batch. supports_lora already gated the family upstream.
lora_payload: Optional[list[dict]] = None
lora_stage: Optional[Path] = None
if lora_resolved:
server_lora_dir = state.server.lora_dir
if server_lora_dir:
lora_stage = Path(server_lora_dir) / f"gen_{os.urandom(6).hex()}"
materialized = diffusion_lora.materialize_native_dir(lora_resolved, lora_stage)
lora_payload = [
{"path": f"{lora_stage.name}/{Path(m.path).name}", "multiplier": float(m.weight)}
for m in materialized
]
try:
for offset in range(0, total, _MAX_SERVER_BATCH):
if cancel.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
count = min(_MAX_SERVER_BATCH, total - offset)
chunk_seed = (base_seed + offset) & ((1 << 63) - 1)
payload = build_img_gen_request(
prompt = prompt,
negative_prompt = negative_prompt or None,
width = int(width),
height = int(height),
steps = int(steps),
seed = chunk_seed,
batch_count = count,
sample_method = state.sampling_method,
flow_shift = state.flow_shift,
cfg_scale = cfg_scale,
distilled_guidance = flux_guidance,
lora = lora_payload,
)
blobs = state.server.img_gen(
payload,
on_step = self._on_log,
cancel_event = cancel,
total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count,
)
# All-or-nothing per chunk, like the one-shot path: if the server returns fewer
# blobs than requested (e.g. one image in the batch failed to encode), fail
# rather than silently dropping images from the user's requested batch.
if not cancel.is_set() and len(blobs) != count:
raise RuntimeError(
f"sd-server returned {len(blobs)} of {count} requested images in the batch."
)
images.extend(Image.open(io.BytesIO(b)).convert("RGB") for b in blobs)
# sd.cpp advances the seed per image within a job, so report chunk_seed+i.
seeds.extend((chunk_seed + i) & ((1 << 63) - 1) for i in range(len(blobs)))
finally:
if lora_stage is not None:
shutil.rmtree(lora_stage, ignore_errors = True)
return images, seeds
def _generate_oneshot(
self,
state: _SdState,
*,
prompt: str,
negative_prompt: Optional[str],
width: int,
height: int,
steps: int,
seed: int,
batch_size: int,
cfg_scale: Optional[float],
flux_guidance: Optional[float],
lora_resolved: list,
cancel: threading.Event,
) -> tuple[list, list[int]]:
"""Fallback path: re-run one-shot sd-cli per image (reloads the model each time).
LoRA on the one-shot path uses sd-cli's own mechanism: materialize the selected
adapters into a ``--lora-model-dir`` and inject matching ``<lora:ALIAS:w>`` tags
into the prompt (sd-cli parses and strips them). supports_lora already gated the
family upstream, so a non-empty ``lora_resolved`` is safe to apply here."""
import tempfile
from PIL import Image
from core.inference import diffusion_lora
engine = self._resolve_engine()
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))]
images = []
seeds: list[int] = []
with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir:
# Materialize selected LoRAs into a managed dir sd-cli can scan, and inject
# matching <lora:ALIAS:w> tags into the prompt (deduped against any the user
# typed). Empty -> prompt/dir unchanged.
eff_prompt = prompt
lora_dir: Optional[str] = None
if lora_resolved:
materialized = diffusion_lora.materialize_native_dir(
lora_resolved, Path(tmpdir) / "loras"
)
eff_prompt = diffusion_lora.inject_prompt_tags(prompt, materialized)
lora_dir = str(Path(tmpdir) / "loras")
for index in range(max(1, int(batch_size))):
if cancel.is_set():
raise RuntimeError(DIFFUSION_CANCELLED_MSG)
# Distinct seed per batch image, reproducible image-by-image from the base
# seed. Mask to int64, NOT 53 bits: the request model and the diffusers
# backend both accept large explicit seeds, so a tight 2**53 mask would
# truncate them and collide distinct requested seeds onto the same image.
seed_i = (seed + index) & ((1 << 63) - 1)
out_path = str(Path(tmpdir) / f"img_{index}.png")
params = SdCppGenParams(
prompt = eff_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,
lora_dir = lora_dir,
lora_apply_mode = "auto" if lora_dir else None,
)
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)
return images, seeds
def _on_log(self, line: str) -> None:
gen = self._gen
if gen is None or gen.total_steps <= 0:
@ -680,13 +1062,40 @@ class SdCppDiffusionBackend:
with self._lock:
if self._active_generate_cancel is not None:
self._active_generate_cancel.set()
state = self._state
self._state = None
self._load_token += 1
self._loading = None
# A load may be mid server.start() with the server not yet committed to _state;
# grab it too so we can stop it (its startup is abortable) instead of leaving it
# loading for the full startup timeout.
pending = self._pending_server
self._pending_server = None
# Stop the resident server outside the lock (terminate can take a few seconds). A
# mid-flight generation had its cancel event set above, so its poll loop unwinds
# as the process goes away.
if state is not None and state.server is not None:
state.server.stop()
if pending is not None and pending is not (state.server if state else None):
pending.stop()
return self.status()
def status(self) -> dict[str, Any]:
state = self._state
# A resident sd-server can exit after load (OOM-killed / crashed while idle). If so,
# drop the stale loaded state so status reports not-loaded and clients reload,
# instead of every generation failing with a 500 against a dead process.
if (
state is not None
and state.mode == "server"
and state.server is not None
and not state.server.is_alive()
):
logger.warning("sd-server exited after load; clearing loaded state")
with self._lock:
if self._state is state:
self._state = None
state = None
if state is None:
return {
"loaded": False,
@ -706,6 +1115,7 @@ class SdCppDiffusionBackend:
"attention_backend": None,
"transformer_cache": None,
"engine": "sd_cpp",
"native_mode": None,
"supports_lora": False,
"supports_controlnet": False,
"workflows": [],
@ -740,8 +1150,10 @@ class SdCppDiffusionBackend:
model_kind = "gguf",
transformer_quant = None,
),
# Native ControlNet (sd-cli --control-net) is a follow-up; off for now.
# ControlNet is diffusers-only; the native engine's generate() rejects it.
"supports_controlnet": False,
# "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli.
"native_mode": state.mode,
# The native engine supports plain text-to-image only (generate() rejects
# img2img / inpaint / reference / upscale), so advertise just txt2img. Without
# this the status omits workflows, the UI reads [], and it disables the Create

81
studio/backend/core/inference/sd_cpp_engine.py
View file

@ -51,6 +51,9 @@ logger = logging.getLogger(__name__)
# target is ``sd-cli``; older builds shipped ``sd`` -- both are probed on PATH.
_BINARY_STEM = "sd-cli"
_LEGACY_STEM = "sd"
# The persistent HTTP server target (stable-diffusion.cpp ``examples/server``). It
# ships next to ``sd-cli`` in both the prebuilt archives and the cmake build tree.
_SERVER_STEM = "sd-server"
class SdCppCancelled(RuntimeError):
@ -119,11 +122,11 @@ def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[
return env
def _layout_candidates(root: Path) -> list[Path]:
"""sd-cli locations under a stable-diffusion.cpp checkout/install ``root``,
def _layout_candidates(root: Path, stem: str = _BINARY_STEM) -> list[Path]:
"""``stem`` locations under a stable-diffusion.cpp checkout/install ``root``,
highest priority first: the cmake ``build/bin`` tree, then a Windows Release
subdir, then the root itself."""
name = _binary_name(_BINARY_STEM)
name = _binary_name(stem)
cands = [
root / "build" / "bin" / name,
root / "build" / "bin" / "Release" / name,
@ -133,37 +136,39 @@ def _layout_candidates(root: Path) -> list[Path]:
return cands
def find_sd_cpp_binary() -> Optional[str]:
"""Locate the ``sd-cli`` binary, or None.
def _first_file(paths: list[Path]) -> Optional[str]:
for p in paths:
try:
if p.is_file():
return str(p)
except OSError:
continue
return None
Search order (mirrors the llama.cpp finder so a Studio install lands where
both engines look):
1. ``SD_CLI_PATH`` env -- a direct path to the binary.
def _find_binary(
*, direct_env: str, path_stems: tuple[str, ...], layout_stem: str
) -> Optional[str]:
"""Shared finder for the stable-diffusion.cpp binaries.
Search order (mirrors the llama.cpp finder so a Studio install lands where every
binary is looked for):
1. ``direct_env`` -- a direct path to the binary.
2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir.
3. The default install root build layouts (the installer target); honors
``UNSLOTH_STUDIO_HOME`` / ``STUDIO_HOME``, else ``~/.unsloth/stable-diffusion.cpp``.
4. ``./stable-diffusion.cpp`` in-tree build (developer checkout).
5. ``sd-cli`` (then legacy ``sd``) on PATH.
5. ``path_stems`` on PATH (in order).
"""
def _first_file(paths: list[Path]) -> Optional[str]:
for p in paths:
try:
if p.is_file():
return str(p)
except OSError:
continue
return None
# 1. Direct binary path.
env_bin = os.environ.get("SD_CLI_PATH")
env_bin = os.environ.get(direct_env)
if env_bin and Path(env_bin).is_file():
return env_bin
# 2. Custom install dir.
custom = os.environ.get("UNSLOTH_SD_CPP_PATH")
if custom:
hit = _first_file(_layout_candidates(Path(custom)))
hit = _first_file(_layout_candidates(Path(custom), layout_stem))
if hit:
return hit
@ -177,27 +182,55 @@ def find_sd_cpp_binary() -> Optional[str]:
if studio_home
else Path.home() / ".unsloth" / "stable-diffusion.cpp"
)
hit = _first_file(_layout_candidates(default_root))
hit = _first_file(_layout_candidates(default_root, layout_stem))
if hit:
return hit
# 4. In-tree developer build: <repo_root>/stable-diffusion.cpp.
try:
project_root = Path(__file__).resolve().parents[4]
hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp"))
hit = _first_file(_layout_candidates(project_root / "stable-diffusion.cpp", layout_stem))
if hit:
return hit
except (OSError, IndexError):
pass
# 5. PATH.
for stem in (_BINARY_STEM, _LEGACY_STEM):
for stem in path_stems:
on_path = shutil.which(stem)
if on_path:
return on_path
return None
def find_sd_cpp_binary() -> Optional[str]:
"""Locate the one-shot ``sd-cli`` binary (env ``SD_CLI_PATH``), or None.
Probes ``sd-cli`` then legacy ``sd`` on PATH. This is the fallback engine once the
persistent ``sd-server`` exists; it also still backs the ESRGAN upscale mode.
"""
return _find_binary(
direct_env = "SD_CLI_PATH",
path_stems = (_BINARY_STEM, _LEGACY_STEM),
layout_stem = _BINARY_STEM,
)
def find_sd_server_binary() -> Optional[str]:
"""Locate the persistent ``sd-server`` binary (env ``SD_SERVER_PATH``), or None.
Same precedence as ``find_sd_cpp_binary`` but keyed to the ``sd-server`` stem, so
a Studio install (prebuilt archive or cmake build, both of which ship ``sd-server``
next to ``sd-cli``) is found in the same places. Preferred over the one-shot CLI:
it loads the model once and serves many generations without reloading from disk.
"""
return _find_binary(
direct_env = "SD_SERVER_PATH",
path_stems = (_SERVER_STEM,),
layout_stem = _SERVER_STEM,
)
class SdCppEngine:
"""A thin handle over a located ``sd-cli`` binary.

510
studio/backend/core/inference/sd_cpp_server.py Normal file
View file

@ -0,0 +1,510 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persistent ``sd-server`` (stable-diffusion.cpp) process manager.
The native diffusion tier used to shell out to one-shot ``sd-cli`` per image, which
reloaded the multi-GB GGUF from disk every generation. ``sd-server`` (the upstream
``examples/server`` target) loads the model once at spawn and serves many generations
over HTTP, exactly like the chat backend's persistent ``llama-server``. This manager
owns ONLY the process + HTTP lifecycle; the backend (``sd_cpp_backend.py``) still owns
asset resolution, request validation, and the public Studio surface.
Shape mirrors ``core/rag/embed_llama_server.py``:
* ``start`` -- pick a free loopback port, spawn the server (model loads here),
drain stdout on a daemon thread, poll until ready.
* ``img_gen`` -- POST ``/sdcpp/v1/img_gen`` (the whole batch in one request),
poll the async job to a terminal state, return image bytes.
* ``stop`` -- SIGTERM -> wait -> SIGKILL, join the drain thread (idempotent).
Readiness is real: upstream ``main.cpp`` loads the model BEFORE it binds the port and
prints ``listening on:``, so a 200 from ``GET /v1/models`` (a trivial handler) means the
model is loaded; a load failure exits the process before listening and is surfaced with
the captured log tail. (The richer ``/sdcpp/v1/capabilities`` handler can block in some
builds, so it is not used for readiness.) The job JSON has no per-step field, so step progress is recovered by
parsing the server's stdout (the same ``N/M`` lines ``sd-cli`` emits), routed to the
active generation's callback.
Import-light on purpose (no torch / diffusers / PIL), so selecting the native tier on
a CPU box never drags the GPU stack into the process.
"""
from __future__ import annotations
import atexit
import base64
import logging
import shutil
import socket
import subprocess
import tempfile
import threading
import time
from collections import deque
from typing import Any, Callable, Optional
import httpx
from core.inference.sd_cpp_args import SdCppModelFiles, build_sd_cpp_server_command
from core.inference.sd_cpp_engine import SdCppCancelled, runtime_env
from utils.native_path_leases import child_env_without_native_path_secret
from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = logging.getLogger(__name__)
# httpx transport errors meaning "the server is gone / connection refused" -- treated
# as "not ready yet" while polling readiness, and as a fatal "server died" mid-request.
_TRANSPORT_ERRORS = (
httpx.ConnectError,
httpx.ReadError,
httpx.RemoteProtocolError,
httpx.WriteError,
)
# Readiness probe. Upstream binds the port only AFTER the model is loaded, so any 200
# means ready. We use /v1/models (a trivial, always-fast handler) rather than
# /sdcpp/v1/capabilities: the capabilities handler can block in some builds (it enumerates
# model metadata), which would stall readiness even though the server is up.
_READY_PATH = "/v1/models"
# Native async sdcpp API.
_IMG_GEN_PATH = "/sdcpp/v1/img_gen"
_JOBS_PATH = "/sdcpp/v1/jobs"
_TERMINAL_OK = "completed"
_TERMINAL_FAIL = "failed"
_TERMINAL_CANCELLED = "cancelled"
# After a cancel is requested, how long to let the server reflect it in job status before
# abandoning the poll. The native cancel is best-effort, so without this cap a server that
# ignores/loses the cancel would keep this call (and the backend's generate lock) alive
# until the job finishes naturally, blocking a superseding load from swapping the model.
_CANCEL_GRACE_S = 5.0
class SdCppServer:
"""A resident ``sd-server`` subprocess plus the HTTP client that drives it."""
def __init__(
self,
binary: str,
*,
host: str = "127.0.0.1",
) -> None:
self.binary = binary
self.host = host
self.port: Optional[int] = None
self._process: Optional[subprocess.Popen] = None
# Fixed-size, thread-safe tail buffer: the drain thread appends while lifecycle /
# request threads read it for diagnostics, so a deque(maxlen) is safer and cheaper
# than a list with manual slicing.
self._tail: deque[str] = deque(maxlen = 200)
self._stdout_thread: Optional[threading.Thread] = None
self._lifecycle_lock = threading.Lock()
# Set (lock-free) by stop() so a blocking start()/readiness wait can be aborted
# promptly without waiting on the lifecycle lock start() holds.
self._abort = threading.Event()
# Set for the duration of a generation so the continuous stdout drain can feed
# the active request's step-progress callback; cleared in img_gen's finally.
self._step_listener: Optional[Callable[[str], None]] = None
# trust_env=False: this client only ever talks to the loopback sd-server, so it must
# not route through HTTP_PROXY/HTTPS_PROXY (a proxy without 127.0.0.1 in NO_PROXY
# would break readiness/generation). Matches the local llama-server clients.
self._client = httpx.Client(timeout = 30.0, trust_env = False)
self._scratch_dir: Optional[str] = None
self._stopped = False
atexit.register(self.stop)
# ── lifecycle ────────────────────────────────────────────────────────────
@property
def base_url(self) -> str:
return f"http://{self.host}:{self.port}"
@property
def lora_dir(self) -> Optional[str]:
"""The server's ``--lora-model-dir`` scratch dir. Adapters staged here are picked
up per request (the sdcpp ``img_gen`` route refreshes its LoRA scan each call), so
server-mode LoRA works by materializing here and referencing the files via the
structured ``lora`` request field (the server ignores ``<lora:>`` prompt tags)."""
return self._scratch_dir
def is_alive(self) -> bool:
return self._process is not None and self._process.poll() is None
@staticmethod
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return int(s.getsockname()[1])
def start(
self,
files: SdCppModelFiles,
*,
vae_format: Optional[str] = None,
offload: Optional[list[str]] = None,
native_speed: Optional[str] = None,
threads: Optional[int] = None,
env: Optional[dict[str, str]] = None,
startup_timeout: float = 600.0,
) -> None:
"""Spawn the server (which loads the model) and block until it is ready.
Raises ``RuntimeError`` (with the captured log tail) if the process exits during
startup or never answers within ``startup_timeout``. Holds the lifecycle lock so
a concurrent start/stop can't interleave.
"""
with self._lifecycle_lock:
# A stop()/unload that raced in AFTER the backend published this server as
# _pending_server but BEFORE start() took the lock has already set _abort and
# closed the httpx client. Honor that delivered stop instead of clearing the
# abort and spawning a model process the cancelled load would then leak.
if self._stopped or self._abort.is_set():
raise SdCppCancelled("sd-server start was cancelled before launch.")
self._abort.clear()
port = self._find_free_port()
# An empty scratch dir for sd-server's LoRA / upscaler / embeddings scans
# (it recursively iterates them per request and errors on a missing dir).
self._scratch_dir = tempfile.mkdtemp(prefix = "sdcpp_dirs_")
cmd = build_sd_cpp_server_command(
self.binary,
files,
host = self.host,
port = port,
vae_format = vae_format,
offload = list(offload or []),
native_speed = native_speed,
threads = threads,
scratch_dir = self._scratch_dir,
verbose = True, # sd-server prints the per-step sampling lines we parse
)
run_env = runtime_env(self.binary, child_env_without_native_path_secret())
if env:
run_env.update(env)
logger.info("starting sd-server: %s", " ".join(cmd))
# Clear in place: reassigning to [] drops the deque(maxlen=200) bound, so the
# continuous stdout drain would then grow the tail without limit for the whole
# resident-server lifetime.
self._tail.clear()
self._spawn_error: Optional[Exception] = None
spawned = threading.Event()
# Spawn INSIDE the drain thread, which then reads stdout for the process's whole
# lifetime. child_popen_kwargs() sets PR_SET_PDEATHSIG, which on Linux is bound to
# the CREATING THREAD -- so the child must be created by a thread that outlives it,
# or a transient spawner thread ending would kill the server. The drain thread is
# exactly that long-lived owner; it dies only when the process exits or the
# interpreter goes away (the case we DO want to reap the GPU-resident server).
def _own_process() -> None:
try:
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
errors = "replace",
env = run_env,
**windows_hidden_subprocess_kwargs(),
**child_popen_kwargs(),
)
except Exception as exc: # noqa: BLE001 -- surface the spawn failure to start()
self._spawn_error = exc
spawned.set()
return
self._process = proc
self.port = port
adopt_pid(proc.pid) # so a global shutdown sweep also reaps it
spawned.set()
self._drain_stdout(proc)
# stdout closed == the process exited; reap it so it is not left a zombie
# until the next stop()/reload.
try:
proc.wait(timeout = 5)
except Exception: # noqa: BLE001
pass
self._stdout_thread = threading.Thread(
target = _own_process, daemon = True, name = "sd-server-owner"
)
self._stdout_thread.start()
spawned.wait()
if self._spawn_error is not None:
self._dispose()
raise RuntimeError(f"failed to spawn sd-server: {self._spawn_error}")
if not self._wait_ready(startup_timeout):
tail = "\n".join(list(self._tail)[-30:])
aborted = self._abort.is_set()
self._kill_locked()
self._dispose()
if aborted:
raise SdCppCancelled("sd-server startup was cancelled.")
raise RuntimeError("sd-server failed to become ready. Last output:\n" + tail[:2000])
def _wait_ready(
self,
timeout: float,
interval: float = 0.5,
) -> bool:
"""Poll ``/v1/models`` until 200; bail early if the process exits.
Upstream binds the port only AFTER the model is loaded, so a 200 here is a true
ready signal (no half-loaded race)."""
deadline = time.monotonic() + timeout
url = f"{self.base_url}{_READY_PATH}"
while time.monotonic() < deadline:
# A concurrent stop() (unload / superseding load) sets _abort so this wait can
# bail without holding the model-load hostage for the full startup_timeout.
if self._abort.is_set():
logger.info("sd-server startup aborted before ready")
return False
if not self.is_alive():
code = None if self._process is None else self._process.returncode
logger.error("sd-server exited early during load (code %s)", code)
return False
try:
if self._client.get(url, timeout = 2.0).status_code == 200:
return True
except (*_TRANSPORT_ERRORS, httpx.TimeoutException):
pass
time.sleep(interval)
logger.error("sd-server readiness timed out after %ss", timeout)
return False
def _drain_stdout(self, proc: subprocess.Popen) -> None:
"""Drain stdout so the pipe never deadlocks; keep a tail for diagnostics and
feed each line to the active generation's step callback."""
try:
assert proc.stdout is not None
for raw in proc.stdout:
line = raw.rstrip()
if not line:
continue
self._tail.append(line) # deque(maxlen) discards the oldest automatically
logger.debug("[sd-server] %s", line)
cb = self._step_listener
if cb is not None:
try:
cb(line)
except Exception: # noqa: BLE001 -- a progress callback must never break drain
pass
except Exception: # noqa: BLE001 -- drain thread must never raise (pipe closed at teardown)
pass
def stop(self) -> None:
"""Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP
client + atexit handler. Idempotent."""
# Signal abort BEFORE contending for the lifecycle lock: a concurrent start() holds
# that lock for the whole (up to startup_timeout) readiness wait, so setting the
# event lets that wait bail immediately instead of stop() blocking behind it.
self._abort.set()
self._stopped = True
with self._lifecycle_lock:
self._kill_locked()
self._dispose()
def _dispose(self) -> None:
"""Release per-instance resources (on stop / failed start). The backend never
reuses a disposed server, so this closes the pooled httpx client and drops the
atexit handler that would otherwise pin every reloaded instance for the session."""
try:
atexit.unregister(self.stop)
except Exception: # noqa: BLE001
pass
try:
self._client.close()
except Exception: # noqa: BLE001
pass
if self._scratch_dir:
shutil.rmtree(self._scratch_dir, ignore_errors = True)
self._scratch_dir = None
def _kill_locked(self) -> None:
proc = self._process
if proc is None:
return
pid = proc.pid
try:
proc.terminate()
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
logger.warning("sd-server did not exit on SIGTERM; killing")
try:
proc.kill()
proc.wait(timeout = 5)
except Exception: # noqa: BLE001 -- best-effort teardown
pass
except Exception as exc: # noqa: BLE001
logger.warning("error terminating sd-server: %s", exc)
finally:
forget_pid(pid)
self._process = None
self.port = None
if self._stdout_thread is not None:
self._stdout_thread.join(timeout = 2)
self._stdout_thread = None
# ── generation ───────────────────────────────────────────────────────────
def img_gen(
self,
payload: dict[str, Any],
*,
on_step: Optional[Callable[[str], None]] = None,
cancel_event: Optional[threading.Event] = None,
poll_interval: float = 0.4,
submit_timeout: float = 60.0,
total_timeout: float = 1800.0,
) -> list[bytes]:
"""Submit one async ``img_gen`` job, poll it to completion, return image bytes.
``on_step`` receives each server stdout line (for the step bar). ``cancel_event``,
when set, cancels the job via the native endpoint and raises ``SdCppCancelled``.
Raises ``RuntimeError`` on submit/poll failures (including the server dying), with
the log tail attached.
"""
# If the server was already stopped for a cancel/unload/superseding load that set
# the cancel event before this submit began, report it as a cancellation (which the
# route maps to a client-state 409) rather than a generic "server died" 500.
if self._stopped or not self.is_alive():
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
raise RuntimeError("sd-server is not running.")
self._step_listener = on_step
job_id: Optional[str] = None
try:
# Submit -> 202 Accepted + job id.
try:
resp = self._client.post(
f"{self.base_url}{_IMG_GEN_PATH}", json = payload, timeout = submit_timeout
)
except (*_TRANSPORT_ERRORS, httpx.TimeoutException) as exc:
raise RuntimeError(self._died_message("img_gen submit", exc)) from exc
if resp.status_code == 429:
raise RuntimeError("sd-server job queue is full (HTTP 429).")
if resp.status_code not in (200, 202):
raise RuntimeError(
f"sd-server img_gen submit -> {resp.status_code}: {resp.text[:500]}"
)
try:
job = resp.json()
except ValueError as exc:
raise RuntimeError(
f"sd-server img_gen returned a non-JSON submit response: {exc}"
) from exc
if not isinstance(job, dict):
raise RuntimeError(
f"sd-server img_gen returned an unexpected submit response type: {type(job)}"
)
job_id = job.get("id")
if not job_id:
raise RuntimeError(f"sd-server img_gen returned no job id: {job}")
# Poll the job to a terminal state.
deadline = time.monotonic() + total_timeout
cancel_sent_at: Optional[float] = None
while True:
if cancel_event is not None and cancel_event.is_set():
if cancel_sent_at is None:
self.cancel(job_id)
cancel_sent_at = time.monotonic()
elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S:
# The best-effort cancel was not reflected in job status within the
# grace window; abandon the poll so the caller can stop the server
# instead of holding the generate lock until the job finishes.
raise SdCppCancelled("sd-server generation was cancelled.")
if not self.is_alive():
# If we're unwinding a cancel (e.g. unload killed the server), surface a
# clean cancellation rather than a generic "server died" error.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.")
raise RuntimeError(self._died_message("img_gen poll", None))
if time.monotonic() > deadline:
# Best-effort cancel, then tear the server down: current sd-server does
# not interrupt an already-generating job (cancel_generating=false / 409),
# so leaving it up would keep denoising the abandoned job and block later
# generations/reloads behind it. Stopping frees the slot; the backend sees
# the dead server on the next generate and takes the recoverable reload path.
self.cancel(job_id)
self.stop()
raise RuntimeError(f"sd-server generation timed out after {total_timeout}s")
try:
jr = self._client.get(f"{self.base_url}{_JOBS_PATH}/{job_id}", timeout = 10.0)
except (*_TRANSPORT_ERRORS, httpx.TimeoutException):
time.sleep(poll_interval)
continue
except RuntimeError as exc:
# A concurrent stop()/unload closes the shared httpx client; httpx then
# raises a plain RuntimeError ("client has been closed") that is NOT a
# transport error. When we are being cancelled, report it as a clean
# cancellation (route -> 409) instead of a generic 500 generation failure.
if cancel_event is not None and cancel_event.is_set():
raise SdCppCancelled("sd-server generation was cancelled.") from exc
raise
if jr.status_code in (404, 410):
raise RuntimeError(f"sd-server job {job_id} is gone (HTTP {jr.status_code}).")
if jr.status_code != 200:
time.sleep(poll_interval)
continue
try:
jd = jr.json()
except ValueError as exc:
raise RuntimeError(f"sd-server job status was not JSON: {exc}") from exc
if not isinstance(jd, dict):
raise RuntimeError(
f"sd-server job status returned an unexpected response type: {type(jd)}"
)
status = jd.get("status")
if status == _TERMINAL_OK:
return self._decode_images(jd)
if status == _TERMINAL_FAIL:
err = jd.get("error") or {}
raise RuntimeError(
"sd-server generation failed: "
f"{err.get('code', 'error')}: {err.get('message', '')}".strip()
)
if status == _TERMINAL_CANCELLED:
raise SdCppCancelled("sd-server generation was cancelled.")
time.sleep(poll_interval)
finally:
self._step_listener = None
def cancel(self, job_id: str) -> None:
"""Best-effort native cancel of an in-flight job."""
try:
self._client.post(f"{self.base_url}{_JOBS_PATH}/{job_id}/cancel", timeout = 5.0)
except Exception: # noqa: BLE001 -- cancel is best-effort
pass
@staticmethod
def _decode_images(job: dict[str, Any]) -> list[bytes]:
# Defensive against an unexpected response shape (a misbehaving/older server):
# verify each level is the type we index before calling dict/list methods.
result = job.get("result") if isinstance(job, dict) else None
images = result.get("images") if isinstance(result, dict) else None
items = [it for it in images if isinstance(it, dict)] if isinstance(images, list) else []
out: list[bytes] = []
for item in sorted(items, key = lambda d: d.get("index", 0)):
b64 = item.get("b64_json")
if not b64:
continue
try:
out.append(base64.b64decode(b64))
except Exception as exc: # noqa: BLE001
raise RuntimeError(f"sd-server returned an undecodable image: {exc}") from exc
if not out:
raise RuntimeError("sd-server completed the job but returned no images.")
return out
def _died_message(self, where: str, exc: Optional[Exception]) -> str:
tail = "\n".join(list(self._tail)[-20:])
base = f"sd-server connection lost during {where}"
if not self.is_alive():
code = None if self._process is None else self._process.returncode
base += f" (process exited, code {code})"
if exc is not None:
base += f": {exc}"
if tail:
base += "\nLast output:\n" + tail[:1500]
return base

70
studio/backend/models/inference.py
View file

@ -2089,6 +2089,11 @@ class DiffusionStatusResponse(BaseModel):
"txt2img, img2img, inpaint. Empty when nothing is loaded or on the native engine.",
)
engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp")
native_mode: Optional[str] = Field(
None,
description = "Native sd.cpp execution mode: server (resident sd-server) | oneshot "
"(per-image sd-cli) | null (diffusers engine)",
)
fallback_reason: Optional[str] = Field(
None,
description = "Why diffusers was chosen over the native sd.cpp engine (null when none)",
@ -2105,3 +2110,68 @@ class DiffusionStatusResponse(BaseModel):
"picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False "
"for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.",
)
# ── OpenAI-compatible images API (POST /v1/images/generations) ──
#
# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf
# OpenAI clients work unchanged. The loaded image GGUF stands in for the model;
# GPT-image-only knobs (quality, style, background, output_format, ...) are
# accepted and ignored, exactly as dall-e-2 ignores them. The size string is
# parsed and `stream` is rejected in the route, where the diffusion backend is
# in reach; everything Pydantic can check declaratively lives here.
class ImageGenerationRequest(BaseModel):
"""OpenAI ``CreateImageRequest`` for ``POST /v1/images/generations``.
``prompt`` is the only required field, per the spec. Unlisted OpenAI fields
are ignored (Pydantic's default), matching dall-e-2's treatment of the
GPT-image-only parameters."""
prompt: str = Field(..., min_length = 1, description = "Text description of the image(s).")
model: Optional[str] = Field(
None, description = "Model id (informational; the loaded image model is used)."
)
n: int = Field(1, ge = 1, le = 10, description = "Number of images to generate (1-10).")
size: str = Field(
"auto", description = "'auto' or '<width>x<height>' (256-2048, each a multiple of 16)."
)
response_format: Literal["url", "b64_json"] = Field(
"url", description = "Return each image as a URL or a base64-encoded PNG."
)
user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).")
# gpt-image-only; declared so we can reject it with a clear error instead of
# silently returning JSON to a client that asked for an SSE stream.
stream: Optional[bool] = Field(
None, description = "Streaming image generation is not supported; omit or set false."
)
@field_validator("n", "size", "response_format", mode = "before")
@classmethod
def _null_means_default(cls, value, info):
# OpenAI marks these nullable WITH a default, so an explicit null means
# "use the default" — coalesce it instead of 400-ing a spec-valid body.
if value is None:
return cls.model_fields[info.field_name].default
return value
class ImageGenerationData(BaseModel):
"""One image in an ``ImagesResponse`` (OpenAI ``Image``). Exactly one of
``url`` / ``b64_json`` is set, per the request's ``response_format``; the
route serializes with ``exclude_none`` so the unused key is omitted."""
b64_json: Optional[str] = Field(
None, description = "Base64-encoded PNG (response_format=b64_json)."
)
url: Optional[str] = Field(None, description = "URL to the PNG bytes (response_format=url).")
class ImageGenerationResponse(BaseModel):
"""OpenAI ``ImagesResponse``. dall-e-shaped: the GPT-image-only top-level
fields (background/output_format/size/quality/usage) are omitted, since our
sizes wouldn't satisfy their fixed enums and we report no token usage."""
created: int = Field(..., description = "Unix timestamp (seconds) the images were created.")
data: list[ImageGenerationData] = Field(..., description = "The generated images.")

167
studio/backend/routes/inference.py
View file

@ -1062,6 +1062,9 @@ from models.inference import (
DiffusionLoadProgressResponse,
GalleryImage,
GalleryListResponse,
ImageGenerationRequest,
ImageGenerationData,
ImageGenerationResponse,
LoadResponse,
LoadProgressResponse,
UnloadResponse,
@ -11375,3 +11378,167 @@ async def diffusion_load_progress(current_subject: str = Depends(get_current_sub
async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.diffusion_engine_router import get_active_diffusion_engine
return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress())
# ──────────────────────────────────────────────────────────────────────────
# OpenAI-compatible images API (POST /v1/images/generations)
#
# The inference router is mounted at both /api/inference and /v1, so this also
# answers /v1/images/generations for off-the-shelf OpenAI clients. It maps
# OpenAI's CreateImageRequest onto the in-process diffusion backend and returns
# an ImagesResponse. Studio's own Image tab uses the richer /images/generate
# route above; this is the spec-shaped surface, and the single error boundary
# mapping backend exceptions to OpenAI error envelopes (the global /v1 handler
# wraps HTTPException detail into the envelope).
# ──────────────────────────────────────────────────────────────────────────
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample
# x 2x patch); the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all
# satisfy this. Mirrors DiffusionGenerateRequest's width/height bounds so both
# generate paths accept the same geometry.
_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$")
_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048
# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both
# "no image model" responses stay identical.
_NO_IMAGE_MODEL_MSG = "No image model loaded. Load an image model first."
def _parse_openai_image_size(size: str) -> tuple[int, int]:
"""OpenAI ``size`` -> (width, height). ``auto``/empty -> 1024x1024 (~1MP, what
these models target). Raises ValueError with a client-facing message."""
text = (size or "").strip().lower()
if text in ("", "auto"):
return 1024, 1024
match = _IMAGE_SIZE_RE.match(text)
if not match:
raise ValueError("size must be 'auto' or '<width>x<height>', e.g. '1024x1024'.")
width, height = int(match.group(1)), int(match.group(2))
for label, value in (("width", width), ("height", height)):
if not _IMAGE_DIM_MIN <= value <= _IMAGE_DIM_MAX:
raise ValueError(f"size {label} must be between {_IMAGE_DIM_MIN} and {_IMAGE_DIM_MAX}.")
if value % 16 != 0:
raise ValueError(f"size {label} must be a multiple of 16.")
return width, height
def _absolute_image_url(request: Request, relative: str) -> str:
"""Join a relative gallery path onto the request's own scheme+host, for the
response_format=url links. Like every Studio route, the target needs the
bearer token; b64_json avoids that for clients that can't carry it."""
return str(request.base_url).rstrip("/") + relative
@router.post(
"/images/generations",
response_model = ImageGenerationResponse,
response_model_exclude_none = True,
)
async def openai_image_generations(
body: ImageGenerationRequest,
request: Request,
current_subject: str = Depends(get_current_subject),
):
"""OpenAI-compatible text-to-image (POST /v1/images/generations).
Generates ``n`` images from ``prompt`` on the loaded diffusion model and
returns them as URLs (default) or base64 PNGs per ``response_format``. Steps
and guidance have no OpenAI knob, so they default per loaded model."""
from core.inference import image_gallery
from core.inference.diffusion_engine_router import get_active_diffusion_engine
from core.inference.diffusion_families import default_generation_params
if body.stream:
raise HTTPException(
status_code = 400,
detail = openai_error_body(
"Streaming image generation is not supported.", status = 400, param = "stream"
),
)
try:
width, height = _parse_openai_image_size(body.size)
except ValueError as exc:
raise HTTPException(
status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size")
)
# Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same
# accessor /images/generate uses, so a model loaded on the native engine isn't
# wrongly reported unloaded here.
backend = get_active_diffusion_engine()
status = backend.status()
if not status.get("loaded"):
# Mirror /v1/completions and /v1/embeddings, which 503 when their backend
# isn't loaded; the global handler turns this into the OpenAI envelope.
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
# Fall back to the resolved base repo so a local-path load (whose repo_id is a
# filesystem path) still gets the right per-model steps/guidance.
steps, guidance = default_generation_params(status.get("repo_id"), status.get("base_repo"))
try:
result = await asyncio.to_thread(
backend.generate,
prompt = body.prompt,
width = width,
height = height,
steps = steps,
guidance = guidance,
batch_size = body.n,
)
except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope)
# A RuntimeError with the model now unloaded means it was evicted/unloaded
# between the readiness check above and the call (a transient race): 503.
# Every other failure (CUDA OOM, a diffusers shape/device error, both also
# RuntimeError) is a real 500, and its raw message must not reach the client.
if isinstance(exc, RuntimeError) and not backend.is_loaded:
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
logger.error("openai_images.generate_failed: %s", exc)
raise HTTPException(status_code = 500, detail = "Image generation failed.")
created = int(time.time())
want_b64 = body.response_format == "b64_json"
# Persist each image with its full recipe embedded, like /images/generate, so
# the response_format=url links resolve and the images show up in the gallery.
recipe = {
"prompt": body.prompt,
"negative_prompt": None,
"width": width,
"height": height,
"steps": steps,
"guidance": guidance,
# The batch shares one base seed, so restoring a batch_index>0 sibling needs the
# original batch_size to replay it (same as /images/generate); persist it.
"batch_size": body.n,
"model": result.get("repo_id"),
"created_at": float(created),
}
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed
# per image (returned in ``seeds``), so record each image's own seed, like
# /images/generate, or a native batch_index>0 image shows the wrong seed.
per_image_seeds = result.get("seeds")
def _persist() -> list[ImageGenerationData]:
items: list[ImageGenerationData] = []
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"]
)
record = image_gallery.save(image, {**recipe, "batch_index": index, "seed": seed})
if want_b64:
encoded = image_gallery.image_b64(record["id"])
if encoded is None: # vanished between write and read — fail the call
raise RuntimeError("generated image could not be read back for encoding")
items.append(ImageGenerationData(b64_json = encoded))
else:
items.append(ImageGenerationData(url = _absolute_image_url(request, record["url"])))
return items
try:
data = await asyncio.to_thread(_persist)
except Exception as exc: # noqa: BLE001
logger.error("openai_images.persist_failed: %s", exc)
raise HTTPException(status_code = 500, detail = "Failed to save the generated image.")
return ImageGenerationResponse(created = created, data = data)

14
studio/backend/tests/test_diffusion_engine_router.py
View file

@ -32,6 +32,10 @@ def _clean_env_and_state(monkeypatch):
"get_active_diffusion_engine",
lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}),
)
# Default: no resident sd-server (so existing tests exercise the sd-cli path only) and
# a stubbed runnability probe, so neither reaches the real install/exec path.
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: None)
monkeypatch.setattr(r, "_server_binary_runnable", lambda *_a, **_k: True)
yield
@ -70,6 +74,16 @@ def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch):
assert r.active_engine_name() == ENGINE_SD_CPP
def test_cpu_with_only_sd_server_picks_sd_cpp(monkeypatch):
# An sd-server-only install (no runnable sd-cli) must still route to native: the
# backend prefers the resident server, so a runnable sd-server is native availability.
_set_device(monkeypatch, "cpu")
_set_binary(monkeypatch, None) # no sd-cli
monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None))
monkeypatch.setattr(r, "ensure_sd_server_binary", lambda **_: "/usr/bin/sd-server")
assert _select() == 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.

23
studio/backend/tests/test_diffusion_prequant.py
View file

@ -133,9 +133,13 @@ def _stub_torch_accelerate(
def _good_ckpt(scheme = "fp8", base = "Tongyi-MAI/Z-Image-Turbo"):
meta = {"scheme": scheme, "base_model_id": base}
# fp8 checkpoints must record per-row granularity or the loader rejects them as stale.
if scheme == "fp8":
meta["fp8_granularity"] = "per_row"
return {
"format": PREQUANT_FORMAT,
"metadata": {"scheme": scheme, "base_model_id": base},
"metadata": meta,
"state_dict": {"weight": object()},
}
@ -220,6 +224,23 @@ def test_load_base_mismatch_is_none(monkeypatch, tmp_path):
assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None
def test_load_fp8_stale_per_tensor_is_rejected(monkeypatch, tmp_path):
# A pre-fix fp8 checkpoint has no fp8_granularity (old per-tensor layout); it must be
# rejected so the loader rebuilds instead of reproducing the noise failure.
stale = _good_ckpt(scheme = "fp8")
del stale["metadata"]["fp8_granularity"]
assert _load(monkeypatch, tmp_path, stale, scheme = "fp8") is None
# An explicit per-tensor granularity is likewise rejected.
per_tensor = _good_ckpt(scheme = "fp8")
per_tensor["metadata"]["fp8_granularity"] = "per_tensor"
assert _load(monkeypatch, tmp_path, per_tensor, scheme = "fp8") is None
def test_load_int8_ignores_fp8_granularity(monkeypatch, tmp_path):
# The granularity gate is fp8-only: an int8 checkpoint without it still loads.
assert _load(monkeypatch, tmp_path, _good_ckpt(scheme = "int8"), scheme = "int8") is not None
def test_load_missing_base_metadata_is_none(monkeypatch, tmp_path):
# A checkpoint whose keys happen to match a different base can load strict=True and then
# render from the wrong weights, so a base was requested but none recorded must be refused.

17
studio/backend/tests/test_diffusion_transformer_quant.py
View file

@ -415,6 +415,23 @@ def test_resolve_fast_accum(monkeypatch):
assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card)
def test_fp8_config_uses_per_row_granularity():
"""FP8 must use PerRow (per-token activation + per-channel weight) scaling. torchao's
default is per-TENSOR: on a DiT with extreme activation outliers (z-image's ~6.6e4) one
outlier forces a tensor-wide scale that pushes normal values below fp8 resolution and the
denoise collapses to noise. This is the regression guard for that fix (validated on B200:
per-tensor fp8 = noise, per-row fp8 = matches bf16)."""
torchao_quant = pytest.importorskip("torchao.quantization")
per_row = getattr(torchao_quant, "PerRow", None)
if per_row is None:
pytest.skip("torchao build without PerRow granularity")
cfg = tq._make_quant_config(TQ_FP8)
gran = getattr(cfg, "granularity", None)
assert gran is not None, "fp8 config must set an explicit granularity, not torchao's default"
grans = gran if isinstance(gran, (list, tuple)) else [gran]
assert grans and all(isinstance(g, per_row) for g in grans), f"expected all PerRow, got {gran}"
def test_quantize_transformer_applies_and_marks(monkeypatch):
monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8)
seen: dict = {}

386
studio/backend/tests/test_openai_images_generations_route.py Normal file
View file

@ -0,0 +1,386 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""FastAPI round-trip tests for the OpenAI-compatible POST /v1/images/generations.
The diffusion backend and image gallery are replaced with light fakes, so these
exercise the route wiring, OpenAI param mapping, validation, error envelopes, and
response shape without torch, diffusers, weights, or a GPU. The pure helpers
(`_parse_openai_image_size`, `default_generation_params`) are unit-tested directly.
"""
from __future__ import annotations
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
import core.inference.diffusion as diffusion_module
import core.inference.diffusion_engine_router as engine_router
import core.inference.image_gallery as gallery_module
from auth.authentication import get_current_subject
from core.inference.diffusion_families import default_generation_params
from routes.inference import router, _parse_openai_image_size
from utils.api_errors import install_api_error_handlers
# ── pure helpers ────────────────────────────────────────────────────────
@pytest.mark.parametrize(
"repo_id, expected",
[
("unsloth/Z-Image-Turbo-GGUF", (9, 0.0)), # turbo entry, before the z-image fallback
("unsloth/Z-Image-GGUF", (20, 4.0)),
("unsloth/FLUX.1-schnell-GGUF", (4, 0.0)), # schnell entry, before the flux.1 entry
("black-forest-labs/FLUX.1-dev", (28, 3.5)),
("unsloth/FLUX.2-klein-4B-GGUF", (4, 0.0)),
("unsloth/Qwen-Image-2512-GGUF", (20, 4.0)),
("some/unknown-model", (9, 0.0)), # fallback
("", (9, 0.0)),
],
)
def test_default_generation_params(repo_id, expected):
assert default_generation_params(repo_id) == expected
def test_default_generation_params_specificity_ordering():
# The "-turbo" / "-schnell" entries must win over their broader siblings; a
# reorder that broke this would silently mis-default.
assert default_generation_params("x/Z-Image-Turbo") != default_generation_params("x/Z-Image")
assert default_generation_params("x/FLUX.1-schnell") != default_generation_params(
"x/FLUX.1-dev"
)
def test_default_generation_params_falls_back_to_base_repo():
# A local-path load: repo_id is a filesystem path that names no model, so the
# resolved base repo is what identifies it (and distinguishes dev from schnell).
assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-dev") == (28, 3.5)
assert default_generation_params("/models/my-ckpt", "black-forest-labs/FLUX.1-schnell") == (
4,
0.0,
)
assert default_generation_params("/models/my-ckpt", "Qwen/Qwen-Image") == (20, 4.0)
# repo_id wins when it already names the model; base repo is only a fallback.
assert default_generation_params("unsloth/Z-Image-Turbo-GGUF", "Tongyi-MAI/Z-Image") == (9, 0.0)
# Nothing identifiable -> fallback; None identifiers are skipped.
assert default_generation_params(None, None) == (9, 0.0)
assert default_generation_params("/models/x", None) == (9, 0.0)
@pytest.mark.parametrize(
"size, expected",
[
("auto", (1024, 1024)),
("", (1024, 1024)),
("AUTO", (1024, 1024)),
("512x512", (512, 512)),
("512x256", (512, 256)),
("1792x1024", (1792, 1024)), # dall-e-3 named size: must pass the bounds
("1024x1792", (1024, 1792)),
(" 256 x 256 ", (256, 256)),
],
)
def test_parse_image_size_ok(size, expected):
assert _parse_openai_image_size(size) == expected
@pytest.mark.parametrize("size", ["abc", "100x100", "4096x4096", "300x300", "512", "x512", "0x0"])
def test_parse_image_size_rejects(size):
with pytest.raises(ValueError):
_parse_openai_image_size(size)
# ── route round-trip ────────────────────────────────────────────────────
class _FakeBackend:
def __init__(
self,
loaded = True,
repo_id = "unsloth/Z-Image-Turbo-GGUF",
base_repo = None,
generate_error = None,
unload_on_generate = False,
native_seeds = False,
) -> None:
self._loaded = loaded
self._repo_id = repo_id
self._base_repo = base_repo
# Model the native sd.cpp engine, which returns a distinct seed per image.
self._native_seeds = native_seeds
# When set, generate() raises this; unload_on_generate flips is_loaded off
# first, to model the eviction/unload race vs an in-pipeline failure (OOM).
self._generate_error = generate_error
self._unload_on_generate = unload_on_generate
self.calls = []
@property
def is_loaded(self):
return self._loaded
def status(self):
return {
"loaded": self._loaded,
"repo_id": self._repo_id if self._loaded else None,
"family": "z-image" if self._loaded else None,
"base_repo": self._base_repo if self._loaded else None,
"device": "cpu",
"dtype": "float32",
"cpu_offload": False,
}
def generate(
self,
*,
prompt,
width,
height,
steps,
guidance,
batch_size = 1,
):
if not self._loaded:
raise RuntimeError("No diffusion model is loaded.")
if self._generate_error is not None:
if self._unload_on_generate:
self._loaded = False
raise self._generate_error
self.calls.append(
dict(
prompt = prompt,
width = width,
height = height,
steps = steps,
guidance = guidance,
batch_size = batch_size,
)
)
out = {
"images": [object() for _ in range(batch_size)],
"seed": 4242,
"repo_id": self._repo_id,
}
if self._native_seeds:
out["seeds"] = [4242 + i for i in range(batch_size)]
return out
def _make_client(backend):
store = {}
def _save(image, meta):
image_id = f"img{len(store)}"
record = {**meta, "id": image_id, "url": f"/api/inference/images/gallery/{image_id}/file"}
store[image_id] = record
return record
app = FastAPI()
install_api_error_handlers(app)
app.include_router(router, prefix = "/v1")
app.dependency_overrides[get_current_subject] = lambda: "test-user"
return TestClient(app), store, _save
@pytest.fixture
def client(monkeypatch):
backend = _FakeBackend()
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
monkeypatch.setattr(gallery_module, "image_b64", lambda i: "QUJD" if i in store else None)
cli.backend = backend # type: ignore[attr-defined]
return cli
def _post(client, body):
return client.post("/v1/images/generations", json = body)
def test_url_response_shape(client):
resp = _post(client, {"prompt": "a sloth", "size": "256x256"})
assert resp.status_code == 200
body = resp.json()
assert set(body.keys()) == {"created", "data"}
assert isinstance(body["created"], int) and body["created"] > 0
assert len(body["data"]) == 1
item = body["data"][0]
assert "url" in item and "b64_json" not in item # exclude_none drops the unused key
assert item["url"].endswith("/file")
# Z-Image-Turbo defaults (9 steps, 0 guidance) flow into the backend call.
assert client.backend.calls[0] == dict(
prompt = "a sloth", width = 256, height = 256, steps = 9, guidance = 0.0, batch_size = 1
)
def test_b64_response_shape(client):
resp = _post(client, {"prompt": "a sloth", "size": "256x256", "response_format": "b64_json"})
assert resp.status_code == 200
item = resp.json()["data"][0]
assert "b64_json" in item and "url" not in item
assert item["b64_json"] == "QUJD"
def test_local_load_uses_base_repo_for_defaults(monkeypatch):
# repo_id is a local path that names no model; base_repo identifies FLUX.1-dev,
# so the route must pick 28 steps / 3.5 guidance, not the 9/0 fallback.
backend = _FakeBackend(repo_id = "/models/my-flux", base_repo = "black-forest-labs/FLUX.1-dev")
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 200
assert backend.calls[0]["steps"] == 28 and backend.calls[0]["guidance"] == 3.5
def test_pipeline_runtime_error_is_sanitized_500(monkeypatch):
# A RuntimeError raised inside the pipeline while the model stays loaded (e.g.
# CUDA OOM, a RuntimeError subclass) must be a sanitized 500, not a 503 that
# echoes the raw exception text.
oom = RuntimeError("CUDA out of memory. Tried to allocate 20.00 GiB (GPU 0; 47.5 GiB total)")
backend = _FakeBackend(generate_error = oom)
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 500
assert resp.json()["error"]["message"] == "Image generation failed."
assert "CUDA" not in resp.text # raw exception text must not leak
def test_unload_race_returns_503(monkeypatch):
# The model is evicted/unloaded between the readiness check and the call: the
# RuntimeError with is_loaded now False is the one case that maps to 503.
backend = _FakeBackend(
generate_error = RuntimeError("No diffusion model is loaded."),
unload_on_generate = True,
)
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 503
err = resp.json()["error"]
assert err["type"] == "api_error"
# The 503 carries the fixed sanitized message, not the raw exception text.
assert err["message"] == "No image model loaded. Load an image model first."
def test_non_runtime_pipeline_error_is_500(monkeypatch):
# A non-RuntimeError from the pipeline (the model stays loaded) must not route
# to the 503 branch (which is gated on isinstance RuntimeError) -> sanitized 500.
backend = _FakeBackend(generate_error = ValueError("bad tensor shape"))
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 500
assert "shape" not in resp.text
def test_n_maps_to_batch(client):
resp = _post(client, {"prompt": "p", "size": "256x256", "n": 3})
assert resp.status_code == 200
assert len(resp.json()["data"]) == 3
assert client.backend.calls[0]["batch_size"] == 3
def test_batch_persists_batch_size(monkeypatch):
# n>1 must persist batch_size in each gallery record so the Studio restore path
# can replay a batch_index>0 sibling (which shares the batch's single seed).
backend = _FakeBackend()
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3})
assert resp.status_code == 200
records = sorted(store.values(), key = lambda r: r["batch_index"])
assert [r["batch_index"] for r in records] == [0, 1, 2]
assert all(r["batch_size"] == 3 for r in records)
def test_uses_active_engine_not_diffusers_singleton(monkeypatch):
# On a no-GPU host the loaded model lives behind the native sd_cpp engine, not the
# diffusers singleton. The route must query get_active_diffusion_engine (like
# /images/generate) or it 503s a model that is loaded and usable.
active = _FakeBackend(loaded = True) # the active (e.g. sd_cpp) engine, loaded
idle_diffusers = _FakeBackend(loaded = False) # diffusers singleton, empty
monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: active)
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: idle_diffusers)
cli, store, _save = _make_client(active)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256"})
assert resp.status_code == 200
assert len(active.calls) == 1 # the active engine did the work, not the idle singleton
def test_native_batch_persists_per_image_seed(monkeypatch):
# The native sd.cpp engine returns a distinct seed per image (base+index) in
# "seeds"; each gallery record must store its own seed (like /images/generate),
# not the shared base, or a restored batch_index>0 image shows the wrong seed.
backend = _FakeBackend(native_seeds = True)
monkeypatch.setattr(engine_router, "get_active_diffusion_engine", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p", "size": "256x256", "n": 3})
assert resp.status_code == 200
records = sorted(store.values(), key = lambda r: r["batch_index"])
assert [r["seed"] for r in records] == [4242, 4243, 4244]
def test_null_fields_coalesce_to_defaults(client):
# OpenAI marks n/size/response_format nullable-with-default: null -> default.
resp = _post(client, {"prompt": "p", "n": None, "size": None, "response_format": None})
assert resp.status_code == 200
assert len(resp.json()["data"]) == 1
assert "url" in resp.json()["data"][0]
assert client.backend.calls[0]["width"] == 1024 # size null -> auto -> 1024
@pytest.mark.parametrize(
"body, param",
[
({"size": "256x256"}, "prompt"), # missing prompt
({"prompt": "", "size": "256x256"}, "prompt"), # empty prompt
({"prompt": "p", "size": "300x300"}, "size"), # not multiple of 16
({"prompt": "p", "size": "abc"}, "size"), # unparseable
({"prompt": "p", "stream": True}, "stream"), # streaming unsupported
],
)
def test_validation_400_with_param(client, body, param):
resp = _post(client, body)
assert resp.status_code == 400
err = resp.json()["error"]
assert err["type"] == "invalid_request_error"
assert err["param"] == param
for k in ("message", "code"):
assert k in err
@pytest.mark.parametrize("n", [0, 11, -1])
def test_n_out_of_range_400(client, n):
resp = _post(client, {"prompt": "p", "size": "256x256", "n": n})
assert resp.status_code == 400
assert resp.json()["error"]["type"] == "invalid_request_error"
def test_no_model_loaded_503(monkeypatch):
backend = _FakeBackend(loaded = False)
monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend)
cli, store, _save = _make_client(backend)
monkeypatch.setattr(gallery_module, "save", _save)
resp = cli.post("/v1/images/generations", json = {"prompt": "p"})
assert resp.status_code == 503
# 503 still wears the OpenAI envelope (api_error) on the /v1 surface.
assert resp.json()["error"]["type"] == "api_error"
def test_auth_required():
backend = _FakeBackend()
app = FastAPI()
install_api_error_handlers(app)
app.include_router(router, prefix = "/v1")
# No dependency override: the real auth dependency runs and rejects.
resp = TestClient(app).post("/v1/images/generations", json = {"prompt": "p"})
assert resp.status_code in (401, 403)

103
studio/backend/tests/test_sd_cpp_args.py
View file

@ -21,7 +21,9 @@ from core.inference.sd_cpp_args import (
SdCppGenParams,
SdCppModelFiles,
SdCppUpscaleParams,
build_img_gen_request,
build_sd_cpp_command,
build_sd_cpp_server_command,
build_sd_cpp_upscale_command,
native_speed_flags,
offload_flags,
@ -364,3 +366,104 @@ def test_build_upscale_requires_input_and_model():
SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""),
output_path = "/o.png",
)
# ── sd-server spawn command ──────────────────────────────────────────────────
def test_server_command_has_model_and_listen_but_no_request_params():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server", files, host = "127.0.0.1", port = 5678, vae_format = "flux2"
)
assert _pair(cmd, "--diffusion-model") == "/m/z.gguf"
assert _pair(cmd, "--vae") == "/m/ae.sft"
assert _pair(cmd, "--llm") == "/m/q.gguf"
assert _pair(cmd, "--vae-format") == "flux2"
assert _pair(cmd, "--listen-ip") == "127.0.0.1"
assert _pair(cmd, "--listen-port") == "5678"
# Per-request parameters must NOT be baked into the spawn command.
for flag in (
"--prompt",
"--seed",
"--steps",
"--cfg-scale",
"--guidance",
"--width",
"--height",
"--batch-count",
):
assert flag not in cmd
def test_server_command_maps_offload_and_speed_and_dedupes():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server",
files,
host = "127.0.0.1",
port = 1,
offload = ["--offload-to-cpu", "--diffusion-fa"],
native_speed = "default", # would add --diffusion-fa again
threads = 8,
)
assert _pair(cmd, "--threads") == "8"
assert cmd.count("--diffusion-fa") == 1 # de-duped against offload
assert "--offload-to-cpu" in cmd
def test_server_command_scratch_dir_expands_to_lora_upscaler_embd():
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
cmd = build_sd_cpp_server_command(
"/bin/sd-server", files, host = "127.0.0.1", port = 1, scratch_dir = "/tmp/scratch"
)
assert _pair(cmd, "--lora-model-dir") == "/tmp/scratch"
assert _pair(cmd, "--hires-upscalers-dir") == "/tmp/scratch"
assert _pair(cmd, "--embd-dir") == "/tmp/scratch"
# Absent when not requested.
bare = build_sd_cpp_server_command("/bin/sd-server", files, host = "127.0.0.1", port = 1)
assert "--lora-model-dir" not in bare and "--hires-upscalers-dir" not in bare
def test_server_command_requires_diffusion_model():
with pytest.raises(ValueError):
build_sd_cpp_server_command(
"/bin/sd-server", SdCppModelFiles(diffusion_model = ""), host = "127.0.0.1", port = 1
)
# ── img_gen request body ─────────────────────────────────────────────────────
def test_img_gen_request_maps_core_fields():
req = build_img_gen_request(
prompt = "a fox",
negative_prompt = "blurry",
width = 512,
height = 768,
steps = 8,
seed = 42,
batch_count = 3,
sample_method = "euler",
cfg_scale = 4.0,
)
assert req["prompt"] == "a fox" and req["negative_prompt"] == "blurry"
assert req["width"] == 512 and req["height"] == 768
assert req["seed"] == 42 and req["batch_count"] == 3
assert req["sample_params"]["sample_steps"] == 8
assert req["sample_params"]["sample_method"] == "euler"
assert req["sample_params"]["guidance"]["txt_cfg"] == 4.0
assert req["output_format"] == "png"
def test_img_gen_request_flux_uses_distilled_guidance():
req = build_img_gen_request(prompt = "x", steps = 4, distilled_guidance = 3.5, flow_shift = 3.0)
g = req["sample_params"]["guidance"]
assert g["distilled_guidance"] == 3.5
assert "txt_cfg" not in g
assert req["sample_params"]["flow_shift"] == 3.0
def test_img_gen_request_requires_prompt():
with pytest.raises(ValueError):
build_img_gen_request(prompt = " ", steps = 4)

419
studio/backend/tests/test_sd_cpp_backend.py
View file

@ -77,10 +77,70 @@ def _loaded_backend(fam_name = "z-image", engine = None):
vae_format = fam.sd_cpp_vae_format,
sampling_method = fam.sd_cpp_sampling_method,
flow_shift = fam.sd_cpp_flow_shift,
mode = "oneshot", # this fixture injects an engine, so it exercises the one-shot path
)
return b
class _FakeServer:
"""Stands in for SdCppServer: records the spawn + one img_gen per whole batch."""
def __init__(self, binary):
self.binary = binary
self.started = None
self.stopped = False
self.payloads = []
self.timeouts = []
self.alive = True
self.lora_dir = None # set by a test to the server's --lora-model-dir scratch dir
def is_alive(self):
return self.alive and not self.stopped
def start(
self,
files,
*,
vae_format = None,
offload = None,
native_speed = None,
threads = None,
):
self.started = dict(
files = files,
vae_format = vae_format,
offload = offload,
native_speed = native_speed,
threads = threads,
)
def img_gen(
self,
payload,
*,
on_step = None,
cancel_event = None,
total_timeout = None,
):
import io as _io
self.payloads.append(payload)
self.timeouts.append(total_timeout)
if on_step is not None:
steps = payload.get("sample_params", {}).get("sample_steps", 0)
on_step(f" {steps}/{steps}")
n = int(payload.get("batch_count", 1))
blobs = []
for i in range(n):
buf = _io.BytesIO()
Image.new("RGB", (1, 1), (i, i, i)).save(buf, format = "PNG")
blobs.append(buf.getvalue())
return blobs
def stop(self):
self.stopped = True
# ── asset resolution ──────────────────────────────────────────────────────────
@ -321,6 +381,260 @@ def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch):
assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF"
# ── persistent sd-server mode ──────────────────────────────────────────────────
def test_resolve_backend_prefers_server(monkeypatch):
b = SdCppDiffusionBackend() # no injected engine
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
mode, binary, engine = b._resolve_backend()
assert mode == "server" and binary == "/x/sd-server" and engine is None
def test_resolve_backend_injected_engine_forces_oneshot():
b = SdCppDiffusionBackend(engine = _FakeEngine())
mode, binary, engine = b._resolve_backend()
assert mode == "oneshot" and binary is None and engine is not None
def test_resolve_backend_falls_back_to_oneshot_without_server(monkeypatch):
b = SdCppDiffusionBackend()
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: None)
monkeypatch.setattr(bk, "_install_allowed", lambda: False) # don't attempt a real install
monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli")
mode, binary, engine = b._resolve_backend()
assert mode == "oneshot" and engine is not None
def test_resolve_backend_cached_fallback_engine_does_not_pin_oneshot(monkeypatch):
# A lazily cached fallback engine (NOT an explicit injection) must not force one-shot:
# once a server is available again, the next load can use it.
b = SdCppDiffusionBackend() # no injected engine
b._engine = _FakeEngine() # simulate a prior lazy one-shot fallback caching the engine
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
mode, binary, engine = b._resolve_backend()
assert mode == "server" and binary == "/x/sd-server" and engine is None
def _run_server_load(
monkeypatch,
b,
servers,
fam_name = "z-image",
):
fam = detect_family(fam_name)
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
# The fake binary path is not a real executable; skip the up-front runnability probe.
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
def _factory(binary):
s = _FakeServer(binary)
servers.append(s)
return s
monkeypatch.setattr(bk, "SdCppServer", _factory)
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"},
)
monkeypatch.setattr(
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
)
b._load_token = 1
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 = 1,
)
def test_server_load_spawns_once_and_status_reports_mode(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
assert len(servers) == 1
assert servers[0].started is not None # the model is loaded once, at spawn
assert b._state is not None and b._state.mode == "server" and b._state.server is servers[0]
assert b.status()["native_mode"] == "server"
def test_server_generate_uses_one_request_for_whole_batch(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 7, batch_size = 3)
assert len(out["images"]) == 3
assert all(isinstance(im, Image.Image) for im in out["images"])
# ONE job for the whole batch (no per-image model reload), unlike the one-shot path.
assert len(servers[0].payloads) == 1
assert servers[0].payloads[0]["batch_count"] == 3
assert out["seed"] == 7 and out["seeds"] == [7, 8, 9]
# step progress was driven from the server's stdout line.
assert b._gen is None # cleared after generate
def test_server_generate_splits_batches_above_server_limit(monkeypatch):
# A batch above the server's per-job limit is chunked (the one-shot path did these
# image-by-image); each chunk gets a timeout proportional to its image count.
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 100, batch_size = 10)
assert len(out["images"]) == 10
counts = [p["batch_count"] for p in servers[0].payloads]
assert counts == [bk._MAX_SERVER_BATCH, 10 - bk._MAX_SERVER_BATCH] # [8, 2]
# Each chunk's timeout scales with its image count, not one fixed batch deadline.
assert servers[0].timeouts == [
bk._SERVER_PER_IMAGE_TIMEOUT_S * 8,
bk._SERVER_PER_IMAGE_TIMEOUT_S * 2,
]
# Seeds run contiguously across chunks (chunk 2 submitted at base + 8).
assert out["seeds"] == list(range(100, 110))
assert servers[0].payloads[1]["seed"] == 108
def test_server_generate_masks_large_seed(monkeypatch):
# sd.cpp's image seed is signed int64; a larger explicit seed must be masked before it
# reaches the server (the request model / diffusers accept up to 2**64 - 1).
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
out = b.generate(prompt = "x", width = 64, height = 64, steps = 4, seed = 2**64 - 1, batch_size = 1)
assert servers[0].payloads[0]["seed"] <= (1 << 63) - 1
assert all(s <= (1 << 63) - 1 for s in out["seeds"])
def test_status_clears_when_server_died(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
assert b.status()["loaded"] is True
servers[0].alive = False # the resident server crashed / was OOM-killed
st = b.status()
assert st["loaded"] is False
assert b._state is None # stale state was dropped so clients reload
def test_server_generate_progress_from_stdout(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
seen = {}
class _WatchServer(_FakeServer):
def img_gen(
self,
payload,
*,
on_step = None,
cancel_event = None,
total_timeout = None,
):
on_step(" 4/8")
seen["mid"] = b.generate_progress()
return super().img_gen(
payload, on_step = on_step, cancel_event = cancel_event, total_timeout = total_timeout
)
b._state = bk._SdState(
repo_id = b._state.repo_id,
base_repo = b._state.base_repo,
family = b._state.family,
device = b._state.device,
files = b._state.files,
vae_format = b._state.vae_format,
sampling_method = b._state.sampling_method,
flow_shift = b._state.flow_shift,
server = _WatchServer("/x/sd-server"),
mode = "server",
)
b.generate(prompt = "x", steps = 8, seed = 1)
assert seen["mid"]["step"] == 4 and seen["mid"]["total_steps"] == 8
def test_server_unload_stops_server(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
st = b.unload()
assert st["loaded"] is False
assert servers[0].stopped is True
assert b._state is None
def test_server_reload_stops_old_server_before_new(monkeypatch):
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
# A second load must tear down the first server and start a fresh one.
b._load_token = 2
fam = detect_family("z-image")
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 = 2,
)
assert len(servers) == 2
assert servers[0].stopped is True # old server stopped
assert b._state.server is servers[1] and servers[1].stopped is False
def test_server_start_failure_falls_back_to_oneshot(monkeypatch):
# A present-but-broken sd-server must not fail the load when sd-cli works.
b = SdCppDiffusionBackend()
monkeypatch.setattr(bk, "find_sd_server_binary", lambda: "/x/sd-server")
# Probe passes; the failure we exercise here is in start(), not the up-front probe.
monkeypatch.setattr(bk, "_server_binary_runnable", lambda *_a, **_k: True)
class _BadServer:
def __init__(self, binary):
self.stopped = False
def start(self, *a, **k):
raise RuntimeError("sd-server broken")
def stop(self):
self.stopped = True
monkeypatch.setattr(bk, "SdCppServer", _BadServer)
fake = _FakeEngine()
monkeypatch.setattr(b, "_resolve_engine", lambda: fake)
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"},
)
monkeypatch.setattr(
bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu")
)
fam = detect_family("z-image")
b._load_token = 1
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 = 1,
)
assert b._state is not None and b._state.mode == "oneshot" and b._state.server is None
# and it can still generate via the one-shot engine
out = b.generate(prompt = "x", steps = 4, seed = 1)
assert len(out["images"]) == 1 and len(fake.calls) == 1
def test_run_load_redacts_paths_in_progress_error(monkeypatch):
# A load failure surfaced via load_progress() must run through redact_native_paths, the
# same scrub the diffusers load path applies, so a registered native path can't leak.
@ -355,3 +669,108 @@ def test_run_load_redacts_paths_in_progress_error(monkeypatch):
with npl._REDACTION_LOCK:
if secret_root in npl._NATIVE_PATH_REDACTIONS:
npl._NATIVE_PATH_REDACTIONS.remove(secret_root)
# ── LoRA (native engine) ────────────────────────────────────────────────────────
def _fake_materialize(resolved, dest):
"""Stand-in for diffusion_lora.materialize_native_dir: write a stub file per adapter
into ``dest`` and return the resolved list pointing at the written paths (mirroring the
real helper's contract without touching the Hub / real weights)."""
from pathlib import Path as _P
from core.inference import diffusion_lora as dl
dest.mkdir(parents = True, exist_ok = True)
out = []
for r in resolved:
p = _P(dest) / f"{r.alias}.safetensors"
p.write_bytes(b"stub")
out.append(dl.ResolvedLora(r.id, r.alias, str(p), r.fmt, r.weight))
return out
def _patch_lora(monkeypatch, resolved, supported = True):
from core.inference import diffusion_lora as dl
monkeypatch.setattr(dl, "supports_lora", lambda **k: supported)
monkeypatch.setattr(dl, "resolve_specs", lambda specs, **k: list(resolved))
monkeypatch.setattr(dl, "materialize_native_dir", _fake_materialize)
def test_generate_oneshot_applies_loras_via_prompt_tags(monkeypatch):
# One-shot sd-cli LoRA: adapters materialized into a --lora-model-dir and selected with
# <lora:ALIAS:w> tags injected into the prompt (the real inject_prompt_tags runs here).
from core.inference import diffusion_lora as dl
eng = _FakeEngine()
b = _loaded_backend(engine = eng) # mode = "oneshot"
_patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.8)])
b.generate(prompt = "a fox", steps = 4, seed = 1, loras = [("id1", 0.8)])
_, params, _, _ = eng.calls[0]
assert params.lora_dir is not None and params.lora_apply_mode == "auto"
assert "<lora:myalias:0.8>" in params.prompt
def test_generate_server_stages_loras_and_sends_structured_field(monkeypatch, tmp_path):
# Server-mode LoRA rides the structured `lora` request field (the sdcpp API ignores
# <lora:> prompt tags): adapters staged into the server's --lora-model-dir, referenced
# by their path relative to it + the validated multiplier.
from pathlib import Path as _P
from core.inference import diffusion_lora as dl
b = SdCppDiffusionBackend()
servers: list = []
_run_server_load(monkeypatch, b, servers)
servers[0].lora_dir = str(tmp_path)
_patch_lora(monkeypatch, [dl.ResolvedLora("id1", "myalias", "/x/a.safetensors", "safetensors", 0.7)])
b.generate(prompt = "x", steps = 4, seed = 1, batch_size = 1, loras = [("id1", 0.7)])
payload = servers[0].payloads[0]
assert "lora" in payload and len(payload["lora"]) == 1
assert payload["lora"][0]["multiplier"] == 0.7
assert payload["lora"][0]["path"].endswith("myalias.safetensors")
assert "<lora:" not in payload["prompt"] # no prompt-tag mechanism on the server
# The per-request stage subdir under the server's lora dir is removed after the batch.
assert not list(_P(tmp_path).glob("gen_*"))
def test_generate_rejects_loras_on_unsupported_family(monkeypatch):
b = _loaded_backend(engine = _FakeEngine())
_patch_lora(monkeypatch, [], supported = False)
with pytest.raises(ValueError, match = "LoRA is not supported"):
b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 1.0)])
def test_generate_zero_weight_loras_are_noop(monkeypatch):
# weight-0 rows are dropped BEFORE the support gate, so a request carrying only disabled
# adapters stays a no-op even on a family where native LoRA is unsupported.
eng = _FakeEngine()
b = _loaded_backend(engine = eng)
_patch_lora(monkeypatch, [], supported = False) # would raise if the gate were reached
b.generate(prompt = "x", steps = 4, seed = 1, loras = [("id1", 0.0)])
_, params, _, _ = eng.calls[0]
assert params.lora_dir is None # nothing applied
def test_generate_rejects_controlnet_on_native_engine():
# ControlNet is diffusers-only. The route passes `controlnet` to whichever engine is
# active, so the native backend must reject it with a clean ValueError (-> 400) rather
# than TypeError on an unexpected kwarg (-> opaque 500).
b = _loaded_backend(engine = _FakeEngine())
with pytest.raises(ValueError, match = "ControlNet is not yet supported on the native"):
b.generate(prompt = "x", steps = 4, seed = 1, controlnet = ("id", "img", "canny", 1.0, 0.0, 1.0))
def test_generate_rejects_image_conditioned_on_native_engine():
# img2img / inpaint / reference / upscale are likewise diffusers-only; a direct API call
# with an init image on the native engine gets a clean ValueError, not a silent txt2img.
b = _loaded_backend(engine = _FakeEngine())
with pytest.raises(ValueError, match = "not yet supported on the native"):
b.generate(prompt = "x", steps = 4, seed = 1, init_image = "data:image/png;base64,AAAA")
def test_status_native_reports_supports_controlnet_false():
b = _loaded_backend()
assert b.status()["supports_controlnet"] is False

53
studio/backend/tests/test_sd_cpp_engine.py
View file

@ -24,6 +24,7 @@ from core.inference.sd_cpp_engine import (
ENGINE_SD_CPP,
SdCppEngine,
find_sd_cpp_binary,
find_sd_server_binary,
runtime_env,
select_diffusion_engine,
)
@ -75,6 +76,58 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch):
assert find_sd_cpp_binary() is None
# ── sd-server discovery ──────────────────────────────────────────────────────
def _clear_server_env(monkeypatch):
monkeypatch.delenv("SD_SERVER_PATH", raising = False)
monkeypatch.delenv("SD_CLI_PATH", raising = False)
monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False)
def test_find_server_prefers_sd_server_path_env(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
binary = tmp_path / "sd-server"
binary.write_text("x")
monkeypatch.setenv("SD_SERVER_PATH", str(binary))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() == str(binary)
def test_find_server_build_layout(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
root = tmp_path / "sdcpp"
built = root / "build" / "bin" / "sd-server"
built.parent.mkdir(parents = True)
built.write_text("x")
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() == str(built)
def test_find_server_path_fallback(tmp_path, monkeypatch):
_clear_server_env(monkeypatch)
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
monkeypatch.setattr(
eng.shutil, "which", lambda stem: "/usr/bin/sd-server" if stem == "sd-server" else None
)
assert find_sd_server_binary() == "/usr/bin/sd-server"
def test_find_server_not_confused_with_sd_cli(tmp_path, monkeypatch):
# A tree that has only sd-cli must NOT be reported as an sd-server (and vice versa),
# so the backend correctly falls back to one-shot when only the CLI is present.
_clear_server_env(monkeypatch)
root = tmp_path / "sdcpp"
(root / "build" / "bin").mkdir(parents = True)
(root / "build" / "bin" / "sd-cli").write_text("x")
monkeypatch.setenv("UNSLOTH_SD_CPP_PATH", str(root))
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
assert find_sd_server_binary() is None
assert find_sd_cpp_binary() == str(root / "build" / "bin" / "sd-cli")
# ── availability / version ──────────────────────────────────────────────────

263
studio/backend/tests/test_sd_cpp_install.py
View file

@ -148,11 +148,12 @@ def test_pinned_tag_default_and_override(monkeypatch):
assert _pinned_tag() is None
def test_repo_default_and_mirror_override(monkeypatch):
def test_repo_default_and_override(monkeypatch):
monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False)
assert _repo() == DEFAULT_REPO == "leejet/stable-diffusion.cpp"
monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "unslothai/stable-diffusion.cpp")
assert _repo() == "unslothai/stable-diffusion.cpp"
# Default is the Unsloth mirror; the env override can point back to leejet upstream.
assert _repo() == DEFAULT_REPO == "unslothai/stable-diffusion.cpp"
monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", "leejet/stable-diffusion.cpp")
assert _repo() == "leejet/stable-diffusion.cpp"
# ── sha256 integrity check ──────────────────────────────────────────────────
@ -306,3 +307,257 @@ def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch):
binary.parent.mkdir(parents = True)
binary.write_bytes(b"x")
assert eng.find_sd_cpp_binary() == str(binary)
# ── Unsloth mirror: default source + the CPU/Apple asset set it publishes ─────
_TAG = "master-741-484baa4"
# Exactly what unslothai/stable-diffusion.cpp's CI publishes (CPU + Apple only; GPU
# hosts run diffusers and never reach the native installer).
_MIRROR_ASSETS = [
f"sd-{_TAG}-bin-Darwin-macOS-arm64.zip",
f"sd-{_TAG}-bin-Darwin-macOS-x86_64.zip",
f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip",
f"sd-{_TAG}-bin-Linux-Ubuntu-24.04-aarch64.zip",
f"sd-{_TAG}-bin-win-cpu-x64.zip",
]
def _mresolve(
system,
machine,
accelerator = "auto",
):
return resolve_release_asset(
_MIRROR_ASSETS, system = system, machine = machine, accelerator = accelerator
)
def test_default_repo_is_the_unsloth_mirror():
assert sdmod.DEFAULT_REPO == "unslothai/stable-diffusion.cpp"
assert sdmod.UPSTREAM_FALLBACK_REPO == "leejet/stable-diffusion.cpp"
def test_mirror_matrix_resolves_every_cpu_apple_host():
assert _mresolve("Darwin", "arm64") == f"sd-{_TAG}-bin-Darwin-macOS-arm64.zip"
assert _mresolve("Darwin", "x86_64") == f"sd-{_TAG}-bin-Darwin-macOS-x86_64.zip"
# WSL reports as Linux x86_64, so this also covers WSL.
assert _mresolve("Linux", "x86_64") == f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip"
assert _mresolve("Linux", "aarch64") == f"sd-{_TAG}-bin-Linux-Ubuntu-24.04-aarch64.zip"
assert _mresolve("Windows", "AMD64") == f"sd-{_TAG}-bin-win-cpu-x64.zip"
assert _mresolve("Windows", "AMD64", "cpu") == f"sd-{_TAG}-bin-win-cpu-x64.zip"
# ── mirror -> upstream fallback in install() ─────────────────────────────────
def _stub_two_repos(monkeypatch, *, mirror_serves, upstream_serves, zip_bytes, digest):
"""Stub _fetch_release to serve (or 404) per repo, so install()'s mirror->upstream
fallback can be exercised without network. Host pinned to Linux x86_64."""
monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False)
monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False)
mirror_name = f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip"
upstream_name = "sd-master-741-484baa4-bin-Linux-Ubuntu-24.04-x86_64.zip"
def _rel(name):
return {
"tag_name": _TAG,
"assets": [
{
"name": name,
"browser_download_url": f"https://example.invalid/{name}",
"digest": digest,
}
],
}
def fake_fetch(
tag = None,
*,
repo = None,
token = None,
timeout = 30.0,
allow_latest = True,
):
r = repo or sdmod.DEFAULT_REPO
if r == sdmod.DEFAULT_REPO:
if mirror_serves:
return _rel(mirror_name)
raise urllib.error.HTTPError(f"https://api/{r}", 404, "not found", None, None)
if upstream_serves:
return _rel(upstream_name)
raise urllib.error.HTTPError(f"https://api/{r}", 404, "not found", None, None)
monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch)
monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zip_bytes))
monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux")
monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64")
return mirror_name, upstream_name
def test_install_uses_mirror_when_available(tmp_path, monkeypatch, capsys):
zb = _zip_with_sd_cli()
_stub_two_repos(
monkeypatch,
mirror_serves = True,
upstream_serves = True,
zip_bytes = zb,
digest = "sha256:" + hashlib.sha256(zb).hexdigest(),
)
sd_cli = install(install_dir = tmp_path)
assert sd_cli.name == "sd-cli" and sd_cli.is_file()
assert "unslothai/stable-diffusion.cpp" in capsys.readouterr().out
def test_install_falls_back_to_upstream_when_mirror_missing(tmp_path, monkeypatch, capsys):
zb = _zip_with_sd_cli()
_stub_two_repos(
monkeypatch,
mirror_serves = False,
upstream_serves = True,
zip_bytes = zb,
digest = "sha256:" + hashlib.sha256(zb).hexdigest(),
)
sd_cli = install(install_dir = tmp_path)
assert sd_cli.name == "sd-cli" and sd_cli.is_file()
captured = capsys.readouterr()
# The repo-fallback diagnostic goes to stderr so --print-asset's stdout stays a single
# asset line; the "source ..." install note still prints to stdout.
assert "falling back to leejet/stable-diffusion.cpp" in captured.err
assert "source leejet/stable-diffusion.cpp" in captured.out
def test_install_errors_when_neither_source_serves(tmp_path, monkeypatch):
_stub_two_repos(
monkeypatch, mirror_serves = False, upstream_serves = False, zip_bytes = b"", digest = ""
)
with pytest.raises(RuntimeError, match = "No prebuilt sd-cli"):
install(install_dir = tmp_path)
# ── explicit GPU accelerator on a CPU-only mirror -> no CPU substitution ──────
def test_mirror_windows_gpu_accel_is_no_match_not_cpu():
# The mirror ships only a CPU win zip. An explicit --accelerator cuda/vulkan/rocm must
# NOT silently resolve to that CPU build; it returns None so install() falls back to
# upstream, which does build the accelerated asset.
for accel in ("cuda", "vulkan", "rocm"):
assert _mresolve("Windows", "AMD64", accel) is None
# auto / cpu still resolve to the CPU build.
assert _mresolve("Windows", "AMD64", "cpu") == f"sd-{_TAG}-bin-win-cpu-x64.zip"
def test_mirror_linux_gpu_accel_is_no_match_not_cpu():
for accel in ("cuda", "vulkan", "rocm"):
assert _mresolve("Linux", "x86_64", accel) is None
assert _mresolve("Linux", "x86_64", "cpu") == f"sd-{_TAG}-bin-Linux-Ubuntu-22.04-x86_64.zip"
def test_upstream_full_matrix_still_resolves_gpu_accel():
# Regression guard: the "explicit GPU accel -> None on miss" change must not break
# the upstream matrix, which does publish the accelerated builds.
assert _resolve("Windows", "AMD64", "cuda") == "sd-master-8caa3f9-bin-win-cuda12-x64.zip"
assert _resolve("Linux", "x86_64", "vulkan").endswith("x86_64-vulkan.zip")
assert "rocm" in _resolve("Linux", "x86_64", "rocm")
# ── explicit repo override suppresses the upstream fallback ───────────────────
def test_explicit_repo_override_equal_to_default_suppresses_fallback(tmp_path, monkeypatch):
# A user who pins UNSLOTH_SD_CPP_REPO (even to the default value) must get exactly that
# repo -- no surprise leejet substitution -- so a missing release errors instead.
_stub_two_repos(
monkeypatch, mirror_serves = False, upstream_serves = True, zip_bytes = b"", digest = ""
)
monkeypatch.setenv("UNSLOTH_SD_CPP_REPO", sdmod.DEFAULT_REPO)
with pytest.raises(RuntimeError, match = "No prebuilt sd-cli"):
install(install_dir = tmp_path)
# ── pinned tag missing on mirror -> pinned upstream before mirror latest ──────
def test_pinned_tag_prefers_upstream_pin_over_mirror_latest(tmp_path, monkeypatch, capsys):
# The mirror lacks the pinned tag but has a newer latest; upstream has the pin. The
# pinned upstream build must win over the unpinned mirror-latest (reproducibility).
monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False)
monkeypatch.setenv("UNSLOTH_SD_CPP_TAG", "master-999-pinned")
zb = _zip_with_sd_cli()
digest = "sha256:" + hashlib.sha256(zb).hexdigest()
def _rel(name, tag):
return {
"tag_name": tag,
"assets": [
{
"name": name,
"browser_download_url": f"https://example.invalid/{name}",
"digest": digest,
}
],
}
mirror_latest = "sd-master-000-latest-bin-Linux-Ubuntu-22.04-x86_64.zip"
upstream_pinned = "sd-master-999-pinned-bin-Linux-Ubuntu-24.04-x86_64.zip"
def fake_fetch(
tag = None,
*,
repo = None,
token = None,
timeout = 30.0,
allow_latest = True,
):
r = repo or sdmod.DEFAULT_REPO
if r == sdmod.DEFAULT_REPO:
if tag == "master-999-pinned": # mirror lacks the pin
if not allow_latest:
return None
return _rel(mirror_latest, "master-000-latest")
return _rel(mirror_latest, "master-000-latest")
# upstream HAS the pin
if tag == "master-999-pinned":
return _rel(upstream_pinned, "master-999-pinned")
return _rel(upstream_pinned, "master-999-pinned")
monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch)
monkeypatch.setattr(sdmod, "_download", lambda url, dest, **k: dest.write_bytes(zb))
monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux")
monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64")
install(install_dir = tmp_path)
out = capsys.readouterr().out
assert "source leejet/stable-diffusion.cpp release master-999-pinned" in out
# ── --print-asset routes through the primary/upstream fallback ────────────────
def test_print_asset_uses_upstream_fallback(monkeypatch, capsys):
# A host the mirror does not build (Linux Vulkan) must print the upstream asset that a
# real install would fetch, not "no matching prebuilt".
monkeypatch.delenv("UNSLOTH_SD_CPP_REPO", raising = False)
monkeypatch.delenv("UNSLOTH_SD_CPP_TAG", raising = False)
def fake_fetch(
tag = None,
*,
repo = None,
token = None,
timeout = 30.0,
allow_latest = True,
):
r = repo or sdmod.DEFAULT_REPO
if r == sdmod.DEFAULT_REPO:
return {"tag_name": _TAG, "assets": [{"name": n} for n in _MIRROR_ASSETS]}
return {"tag_name": "master-8caa3f9", "assets": [{"name": n} for n in _ASSETS]}
monkeypatch.setattr(sdmod, "_fetch_release", fake_fetch)
monkeypatch.setattr(sdmod.platform, "system", lambda: "Linux")
monkeypatch.setattr(sdmod.platform, "machine", lambda: "x86_64")
rc = sdmod.main(["--print-asset", "--accelerator", "vulkan"])
out = capsys.readouterr().out
assert rc == 0
assert "vulkan" in out and "no matching prebuilt" not in out

417
studio/backend/tests/test_sd_cpp_server.py Normal file
View file

@ -0,0 +1,417 @@
# 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 persistent sd-server process manager (SdCppServer).
Hermetic: subprocess.Popen and the httpx client are faked, so nothing spawns a real
binary or opens a socket beyond the free-port probe."""
from __future__ import annotations
import base64
import io
import threading
import pytest
from PIL import Image
from core.inference import sd_cpp_server as srv
from core.inference.sd_cpp_args import SdCppModelFiles
from core.inference.sd_cpp_engine import SdCppCancelled
from core.inference.sd_cpp_server import SdCppServer
_FILES = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/vae.sft", llm = "/m/llm.sft")
def _png_b64(shade: int) -> str:
buf = io.BytesIO()
Image.new("RGB", (1, 1), (shade, shade, shade)).save(buf, format = "PNG")
return base64.b64encode(buf.getvalue()).decode()
class _FakePopen:
"""Minimal Popen stand-in. stdout yields the scripted lines then BLOCKS until the
process is terminated/killed/exited -- mirroring a real child that holds its pipe
open for its lifetime (so the owner/drain thread stays alive, as in production)."""
def __init__(
self,
lines = (),
exit_code = None,
):
self.pid = 4242
self._lines = list(lines)
self._exit = exit_code # None == alive
self.returncode = exit_code
self.terminated = False
self.killed = False
self._done = threading.Event()
if exit_code is not None:
self._done.set()
@property
def stdout(self):
def _gen():
for ln in self._lines:
yield ln
self._done.wait() # hold the pipe open until the process ends
return _gen()
def poll(self):
return self._exit
def terminate(self):
self.terminated = True
self._exit = 0
self.returncode = 0
self._done.set()
def wait(self, timeout = None):
self._done.wait(timeout)
if self._exit is None:
self._exit = 0
self.returncode = 0
return self.returncode
def kill(self):
self.killed = True
self._exit = -9
self.returncode = -9
self._done.set()
class _Resp:
def __init__(
self,
status_code,
payload = None,
text = "",
bad_json = False,
):
self.status_code = status_code
self._payload = payload if payload is not None else {}
self.text = text
self._bad_json = bad_json
def json(self):
if self._bad_json:
raise ValueError("not json")
return self._payload
class _FakeClient:
def __init__(
self,
*,
get = None,
post = None,
):
self._get = get or (lambda url: _Resp(200, {}))
self._post = post or (lambda url, json: _Resp(202, {"id": "job1"}))
self.get_urls = []
self.post_calls = []
self.closed = False
def get(
self,
url,
timeout = None,
):
self.get_urls.append(url)
return self._get(url)
def post(
self,
url,
json = None,
timeout = None,
):
self.post_calls.append((url, json))
return self._post(url, json)
def close(self):
self.closed = True
@pytest.fixture
def patched(monkeypatch):
"""Neutralise process-lifetime side effects for the manager under test."""
monkeypatch.setattr(srv, "adopt_pid", lambda pid: None)
monkeypatch.setattr(srv, "forget_pid", lambda pid: None)
monkeypatch.setattr(srv, "child_popen_kwargs", lambda: {})
monkeypatch.setattr(srv, "windows_hidden_subprocess_kwargs", lambda: {})
return monkeypatch
def _server_with(popen, client):
s = SdCppServer("/x/sd-server")
s._client = client
# Attach the fake process + port so generation tests can run without start().
s._process = popen
s.port = 1234
return s
# ── start / readiness ──────────────────────────────────────────────────────────
def test_start_becomes_ready_when_capabilities_200(patched):
popen = _FakePopen(lines = ["loading model", "listening on: http://127.0.0.1:1"])
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(
popen, _FakeClient(get = lambda url: _Resp(200, {"model": {"path": "/m/z.gguf"}}))
)
s.start(_FILES, startup_timeout = 5.0)
assert s.is_alive() is True
assert s.port is not None
def test_start_fails_fast_when_process_exits(patched):
# Model load failed -> process exits before listening; start must raise with the tail.
popen = _FakePopen(lines = ["error: bad model"], exit_code = 1)
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
# Capabilities never answers (connection refused) -> readiness relies on exit detection.
s = _server_with(
popen, _FakeClient(get = lambda url: (_ for _ in ()).throw(srv.httpx.ConnectError("refused")))
)
with pytest.raises(RuntimeError, match = "failed to become ready"):
s.start(_FILES, startup_timeout = 2.0)
# ── generation ───────────────────────────────────────────────────────────────
def _completed_job(images_b64):
return _Resp(
200,
{
"status": "completed",
"result": {"images": [{"index": i, "b64_json": b} for i, b in enumerate(images_b64)]},
},
)
def test_img_gen_returns_image_bytes_in_index_order(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobA"}),
# result images deliberately out of order -> manager must sort by index.
get = lambda url: _Resp(
200,
{
"status": "completed",
"result": {
"images": [
{"index": 1, "b64_json": _png_b64(200)},
{"index": 0, "b64_json": _png_b64(50)},
]
},
},
),
),
)
blobs = s.img_gen({"prompt": "x", "batch_count": 2, "sample_params": {"sample_steps": 4}})
assert len(blobs) == 2
first = Image.open(io.BytesIO(blobs[0])).convert("RGB").getpixel((0, 0))
assert first == (50, 50, 50) # index 0 first
def test_img_gen_failed_job_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobF"}),
get = lambda url: _Resp(
200, {"status": "failed", "error": {"code": "x", "message": "boom"}}
),
),
)
with pytest.raises(RuntimeError, match = "generation failed.*boom"):
s.img_gen({"prompt": "x"})
def test_img_gen_queue_full_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(429, text = "busy")))
with pytest.raises(RuntimeError, match = "queue is full"):
s.img_gen({"prompt": "x"})
def test_img_gen_cancel_posts_cancel_and_raises(patched):
popen = _FakePopen()
cancel = threading.Event()
cancel.set() # already cancelled before the first poll
client = _FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobC"}),
get = lambda url: _Resp(
200, {"status": "cancelled", "error": {"code": "cancelled", "message": "c"}}
),
)
s = _server_with(popen, client)
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel)
assert any(url.endswith("/cancel") for url, _ in client.post_calls)
def test_img_gen_detects_server_death(patched):
popen = _FakePopen()
def _die_get(url):
popen._exit = 137 # the process died between submit and poll
return _Resp(200, {"status": "generating"})
s = _server_with(
popen, _FakeClient(post = lambda url, json: _Resp(202, {"id": "jobD"}), get = _die_get)
)
with pytest.raises(RuntimeError, match = "connection lost|process exited"):
s.img_gen({"prompt": "x"})
# ── stdout routing + stop ──────────────────────────────────────────────────────
def test_drain_routes_lines_to_step_listener_and_tail(patched):
s = SdCppServer("/x/sd-server")
seen = []
s._step_listener = seen.append
# exit_code set so stdout ends after the scripted lines (a live fake would block).
s._drain_stdout(_FakePopen(lines = ["sampling 1/8", "", "sampling 8/8", "done"], exit_code = 0))
assert "sampling 1/8" in seen and "sampling 8/8" in seen
assert "" not in seen # blank lines skipped
assert s._tail[-1] == "done"
def test_stop_is_idempotent_and_terminates(patched):
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
client = _FakeClient(get = lambda url: _Resp(200, {}))
s = _server_with(popen, client)
s.start(_FILES, startup_timeout = 5.0)
s.stop()
assert popen.terminated is True
assert s.is_alive() is False
assert client.closed is True # stop() releases the pooled HTTP client
s.stop() # second call must not raise
def test_img_gen_submit_error_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(400, text = "bad params")))
with pytest.raises(RuntimeError, match = "submit -> 400"):
s.img_gen({"prompt": "x"})
def test_img_gen_malformed_submit_json_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, bad_json = True)))
with pytest.raises(RuntimeError, match = "non-JSON submit"):
s.img_gen({"prompt": "x"})
def test_img_gen_empty_result_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobE"}),
get = lambda url: _Resp(200, {"status": "completed", "result": {"images": []}}),
),
)
with pytest.raises(RuntimeError, match = "no images"):
s.img_gen({"prompt": "x"})
def test_img_gen_rejected_after_stop(patched):
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {})))
s.start(_FILES, startup_timeout = 5.0)
s.stop()
with pytest.raises(RuntimeError, match = "not running"):
s.img_gen({"prompt": "x"})
# ── cancellation + defensive parsing (review follow-ups) ───────────────────────
def test_img_gen_cancelled_before_submit_reports_cancellation(patched):
# The server was stopped for a cancel/unload before submit; with the cancel event set
# this must surface as a cancellation (route -> 409), not a generic "not running" 500.
popen = _FakePopen()
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
s = _server_with(popen, _FakeClient(get = lambda url: _Resp(200, {})))
s.start(_FILES, startup_timeout = 5.0)
s.stop()
cancel = threading.Event()
cancel.set()
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel)
def test_img_gen_abandons_when_cancel_not_honored(patched):
# A best-effort cancel the server ignores must not pin this call (and the generate
# lock) until natural completion: after the grace window it raises cancellation.
patched.setattr(srv, "_CANCEL_GRACE_S", 0.0)
popen = _FakePopen()
cancel = threading.Event()
cancel.set()
client = _FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobG"}),
get = lambda url: _Resp(200, {"status": "generating"}), # never terminal
)
s = _server_with(popen, client)
with pytest.raises(SdCppCancelled):
s.img_gen({"prompt": "x"}, cancel_event = cancel, poll_interval = 0.01)
def test_img_gen_non_dict_submit_json_raises(patched):
popen = _FakePopen()
s = _server_with(popen, _FakeClient(post = lambda url, json: _Resp(202, ["not", "a", "dict"])))
with pytest.raises(RuntimeError, match = "unexpected submit response"):
s.img_gen({"prompt": "x"})
def test_img_gen_non_dict_status_json_raises(patched):
popen = _FakePopen()
s = _server_with(
popen,
_FakeClient(
post = lambda url, json: _Resp(202, {"id": "jobH"}),
get = lambda url: _Resp(200, ["unexpected"]),
),
)
with pytest.raises(RuntimeError, match = "unexpected response type"):
s.img_gen({"prompt": "x"}, poll_interval = 0.01)
def test_decode_images_tolerates_unexpected_shapes():
# A misbehaving/older server can return non-dict result/images/items; _decode_images
# must raise a clean "no images" rather than an AttributeError on .get().
for job in ({"result": ["x"]}, {"result": {"images": "nope"}}, {"result": {"images": [1, 2]}}):
with pytest.raises(RuntimeError, match = "no images"):
SdCppServer._decode_images(job)
def test_start_aborted_by_concurrent_stop(patched):
# A stop() during the readiness wait must abort start() promptly (without waiting out
# the startup timeout) and surface as a cancellation.
popen = _FakePopen(lines = ["loading model"])
patched.setattr(srv.subprocess, "Popen", lambda *a, **k: popen)
def _never_ready(url):
raise srv.httpx.ConnectError("refused")
s = _server_with(popen, _FakeClient(get = _never_ready))
def _stop_soon():
import time as _t
_t.sleep(0.2)
s.stop()
threading.Thread(target = _stop_soon, daemon = True).start()
with pytest.raises(SdCppCancelled):
s.start(_FILES, startup_timeout = 30.0)

510
studio/frontend/src/features/images/diffusion-train-dialog.tsx Normal file
View file

@ -0,0 +1,510 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useRef, useState } from "react";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store";
import { toast } from "@/lib/toast";
import {
type DiffusionTrainingInfo,
type DiffusionTrainingStatus,
getDiffusionTrainingInfo,
getDiffusionTrainingStatus,
startDiffusionTraining,
stopDiffusionTraining,
uploadDiffusionDataset,
} from "./api";
// The two official SDXL bases the backend allowlists for non-GGUF loads. Everything the
// dropdown offers is trainable; "custom" is the escape hatch for local SDXL checkpoints.
const SDXL_BASES: Array<{ id: string; label: string }> = [
{ id: "stabilityai/stable-diffusion-xl-base-1.0", label: "SDXL Base 1.0 (best quality)" },
{ id: "stabilityai/sdxl-turbo", label: "SDXL Turbo (fast, good for quick tests)" },
];
const CUSTOM_BASE = "__custom__";
const UPLOAD_DATASET = "__upload__";
const DATASET_FILE_ACCEPT = ".png,.jpg,.jpeg,.webp,.bmp,.txt,.caption,.jsonl";
const selectClass =
"h-8 w-full rounded-md border border-input bg-background px-2 text-xs";
// A self-contained "Train an SDXL LoRA" dialog. It posts to /api/train/diffusion/start
// and polls /status while open, so it never blocks the page and works whether or not a
// model is loaded for generation. Only SDXL is trainable today; the backend refuses
// known non-SDXL picks instantly, and the base-model dropdown keeps users on safe picks.
export function DiffusionTrainDialog({
open,
onOpenChange,
defaultBaseModel,
onTrainingComplete,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
defaultBaseModel?: string;
// Called once when a run finishes so the page can rescan the LoRA picker.
onTrainingComplete?: () => void;
}) {
const [baseChoice, setBaseChoice] = useState(SDXL_BASES[0].id);
const [customBase, setCustomBase] = useState("");
const [info, setInfo] = useState<DiffusionTrainingInfo | null>(null);
const [dataset, setDataset] = useState<string>(UPLOAD_DATASET);
const [uploadName, setUploadName] = useState("my-images");
const [uploading, setUploading] = useState(false);
const fileInputRef = useRef<HTMLInputElement | null>(null);
const [outputDir, setOutputDir] = useState("");
const [instancePrompt, setInstancePrompt] = useState("");
const [showAdvanced, setShowAdvanced] = useState(false);
const [steps, setSteps] = useState(500);
const [learningRate, setLearningRate] = useState(0.0001);
const [rank, setRank] = useState(16);
const [resolution, setResolution] = useState(1024);
const [batchSize, setBatchSize] = useState(1);
const [precision, setPrecision] = useState<"bf16" | "fp16" | "no">("bf16");
const [starting, setStarting] = useState(false);
const [status, setStatus] = useState<DiffusionTrainingStatus | null>(null);
// The dialog stays mounted (ImagesPage is keep-alive), so seed per-open state here:
// the base-model choice from the currently loaded SDXL pipeline (when there is one),
// and the dataset list from the backend.
const refreshInfo = useCallback(async (): Promise<DiffusionTrainingInfo | null> => {
try {
const i = await getDiffusionTrainingInfo();
setInfo(i);
return i;
} catch {
return null; // older backends: keep the upload-only flow usable
}
}, []);
useEffect(() => {
if (!open) return;
if (defaultBaseModel) {
setBaseChoice(defaultBaseModel);
}
void refreshInfo().then((i) => {
// Preselect the only dataset, or the freshest-looking state: with no datasets
// yet, the picker sits on "Upload new images".
setDataset((cur) => {
if (cur !== UPLOAD_DATASET && i?.datasets.some((d) => d.name === cur)) return cur;
return i && i.datasets.length > 0 ? i.datasets[0].name : UPLOAD_DATASET;
});
});
}, [open, defaultBaseModel, refreshInfo]);
const poll = useCallback(async () => {
try {
setStatus(await getDiffusionTrainingStatus());
} catch {
// Best-effort; a failed poll should not surface an error while the dialog is open.
}
}, []);
// Poll status only while the dialog is open.
useEffect(() => {
if (!open) return;
void poll();
const id = window.setInterval(() => void poll(), 1500);
return () => window.clearInterval(id);
}, [open, poll]);
const active = Boolean(status?.active) || status?.status === "running";
const completed = status?.status === "completed";
const pct =
status && status.total_steps > 0
? Math.min(100, Math.round((status.step / status.total_steps) * 100))
: 0;
// Notify the parent exactly once when a run reaches "completed", so it can rescan the
// LoRA picker (a LoRA trained while a model is loaded is otherwise invisible until a
// model swap re-runs the discovery effect).
const [notifiedComplete, setNotifiedComplete] = useState(false);
useEffect(() => {
if (status?.status === "completed" && !notifiedComplete) {
setNotifiedComplete(true);
onTrainingComplete?.();
} else if (status?.status === "running" && notifiedComplete) {
setNotifiedComplete(false); // arm again for the next run
}
}, [status?.status, notifiedComplete, onTrainingComplete]);
const selectedDataset =
dataset !== UPLOAD_DATASET ? info?.datasets.find((d) => d.name === dataset) : undefined;
const onUpload = useCallback(async () => {
const files = Array.from(fileInputRef.current?.files ?? []);
if (files.length === 0) {
toast.error("Choose the images to upload first.");
return;
}
const name = uploadName.trim();
if (!name) {
toast.error("Give the dataset a folder name, e.g. my-style-photos.");
return;
}
setUploading(true);
try {
const res = await uploadDiffusionDataset(name, files);
toast.success(
`Uploaded ${res.uploaded} file${res.uploaded === 1 ? "" : "s"} - ` +
`"${res.name}" now has ${res.image_count} images`,
);
if (fileInputRef.current) fileInputRef.current.value = "";
await refreshInfo();
setDataset(res.name);
} catch (e) {
toast.error(e instanceof Error ? e.message : "Upload failed");
} finally {
setUploading(false);
}
}, [uploadName, refreshInfo]);
const onStart = useCallback(async () => {
const baseModel = (baseChoice === CUSTOM_BASE ? customBase : baseChoice).trim();
if (!baseModel) {
toast.error("Pick a base model (or fill in the custom repo/path).");
return;
}
if (dataset === UPLOAD_DATASET) {
toast.error("Upload your training images first (or pick an existing dataset).");
return;
}
if (!outputDir.trim()) {
toast.error("Name the adapter (this becomes its folder under Studio outputs).");
return;
}
if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) {
toast.error(
"These images have no captions - add a trigger prompt so the trainer knows " +
"what to learn (it becomes the caption for every image).",
);
return;
}
// Mirror the backend's numeric validation so obvious mistakes are caught before the
// request (the backend returns 400 for these; catching here gives a clearer message).
if (steps < 1) return toast.error("Steps must be at least 1.");
if (rank < 1) return toast.error("LoRA rank must be at least 1.");
if (resolution < 64 || resolution % 8 !== 0) {
return toast.error("Resolution must be a multiple of 8 and at least 64.");
}
if (batchSize < 1) return toast.error("Batch size must be at least 1.");
if (learningRate <= 0) return toast.error("Learning rate must be greater than 0.");
setStarting(true);
try {
await startDiffusionTraining({
base_model: baseModel,
data_dir: dataset,
output_dir: outputDir.trim(),
instance_prompt: instancePrompt.trim() || undefined,
resolution,
train_steps: steps,
learning_rate: learningRate,
train_batch_size: batchSize,
lora_rank: rank,
mixed_precision: precision,
// Forward the saved Hub token so a gated/private SDXL base can be trained (the
// image load flow already sends it, so a model you can load, you can also train).
hf_token: hfApiToken(getHfToken()) || undefined,
});
toast.success("Training started");
void poll();
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to start training");
} finally {
setStarting(false);
}
}, [
baseChoice,
customBase,
dataset,
selectedDataset,
outputDir,
instancePrompt,
resolution,
steps,
learningRate,
batchSize,
rank,
precision,
poll,
]);
const onStop = useCallback(async () => {
try {
await stopDiffusionTraining();
toast.success("Stop requested; finishing the current step.");
void poll();
} catch (e) {
toast.error(e instanceof Error ? e.message : "Failed to stop training");
}
}, [poll]);
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="flex max-h-[85vh] max-w-lg flex-col">
<DialogHeader>
<DialogTitle>Train an SDXL LoRA</DialogTitle>
<DialogDescription>
Teach SDXL a style, character, or subject from your own images. The finished
adapter shows up in this page&apos;s LoRA picker. Only SDXL can be trained for
now - FLUX, Qwen-Image and Z-Image load LoRAs but can&apos;t train them yet.
</DialogDescription>
</DialogHeader>
<div className="grid gap-3 overflow-y-auto py-2 pr-1">
{/* 1. Base model: a constrained dropdown instead of free text, so the "SDXL
only" rule is embodied by the control. Custom stays available for local
SDXL checkpoints; a known non-SDXL pick is refused instantly by the API. */}
<div className="grid gap-1.5">
<Label className="text-xs">Base model to train on</Label>
<select
value={baseChoice}
onChange={(e) => setBaseChoice(e.target.value)}
className={selectClass}
>
{SDXL_BASES.map((b) => (
<option key={b.id} value={b.id}>
{b.label}
</option>
))}
{defaultBaseModel && !SDXL_BASES.some((b) => b.id === defaultBaseModel) && (
<option value={defaultBaseModel}>Loaded model: {defaultBaseModel}</option>
)}
<option value={CUSTOM_BASE}>Custom SDXL repo or local path...</option>
</select>
{baseChoice === CUSTOM_BASE && (
<Input
value={customBase}
placeholder="my-org/my-sdxl-finetune or /path/to/sdxl-pipeline"
spellCheck={false}
onChange={(e) => setCustomBase(e.target.value)}
className="h-8 text-xs"
/>
)}
</div>
{/* 2. Training images: pick an existing dataset folder or upload straight from
the browser - no shell access or Studio-home knowledge needed. */}
<div className="grid gap-1.5">
<Label className="text-xs">Training images</Label>
<select
value={dataset}
onChange={(e) => setDataset(e.target.value)}
className={selectClass}
>
{(info?.datasets ?? []).map((d) => (
<option key={d.name} value={d.name}>
{d.name} ({d.image_count} image{d.image_count === 1 ? "" : "s"}
{d.caption_count > 0 ? `, ${d.caption_count} captions` : ""})
</option>
))}
<option value={UPLOAD_DATASET}>Upload new images...</option>
</select>
{dataset === UPLOAD_DATASET && (
<div className="grid gap-1.5 rounded-md border border-dashed border-border p-2">
<Input
value={uploadName}
placeholder="my-style-photos"
spellCheck={false}
onChange={(e) => setUploadName(e.target.value)}
className="h-8 text-xs"
aria-label="New dataset name"
/>
<div className="flex items-center gap-2">
<input
ref={fileInputRef}
type="file"
multiple
accept={DATASET_FILE_ACCEPT}
className="min-w-0 flex-1 text-xs file:mr-2 file:rounded-md file:border-0 file:bg-muted file:px-2 file:py-1 file:text-xs"
aria-label="Training images"
/>
<Button
type="button"
size="sm"
variant="secondary"
className="h-8 shrink-0"
onClick={onUpload}
disabled={uploading}
>
{uploading ? "Uploading..." : "Upload"}
</Button>
</div>
<p className="text-[11px] text-muted-foreground">
10-50 images work well. Optional captions: a .txt per image (same
filename) or a metadata.jsonl; without them the trigger prompt below
captions every image. You can upload more into the same name later.
</p>
</div>
)}
{selectedDataset && selectedDataset.caption_count === 0 && (
<p className="text-[11px] text-muted-foreground">
No caption files in this dataset - the trigger prompt below will be used
as the caption for every image.
</p>
)}
</div>
{/* 3. What to call the result + how to trigger it. */}
<div className="grid gap-1.5">
<Label className="text-xs">Adapter name</Label>
<Input
value={outputDir}
placeholder="my-style-lora"
spellCheck={false}
onChange={(e) => setOutputDir(e.target.value)}
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Trigger prompt (how you&apos;ll invoke the style later)</Label>
<Input
value={instancePrompt}
placeholder="a photo in SKS style"
onChange={(e) => setInstancePrompt(e.target.value)}
className="h-8 text-xs"
/>
</div>
{/* 4. Hyperparameters, collapsed: the defaults suit a first run, and hiding
them keeps the primary flow at three decisions. */}
<button
type="button"
className="w-fit text-xs text-muted-foreground underline-offset-2 hover:underline"
onClick={() => setShowAdvanced((s) => !s)}
aria-expanded={showAdvanced}
>
{showAdvanced ? "Hide training settings" : "Training settings (defaults suit a first run)"}
</button>
{showAdvanced && (
<>
<div className="grid grid-cols-2 gap-3 sm:grid-cols-4">
<div className="grid gap-1.5">
<Label className="text-xs">Steps</Label>
<Input
type="number"
min={1}
value={steps}
onChange={(e) => setSteps(Number(e.target.value) || 1)}
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">LoRA rank</Label>
<Input
type="number"
min={1}
value={rank}
onChange={(e) => setRank(Number(e.target.value) || 1)}
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Resolution</Label>
<Input
type="number"
min={64}
step={64}
value={resolution}
onChange={(e) => setResolution(Number(e.target.value) || 1024)}
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Batch</Label>
<Input
type="number"
min={1}
value={batchSize}
onChange={(e) => setBatchSize(Number(e.target.value) || 1)}
className="h-8 text-xs"
/>
</div>
</div>
<div className="grid grid-cols-2 gap-3">
<div className="grid gap-1.5">
<Label className="text-xs">Learning rate</Label>
<Input
type="number"
step={0.00001}
min={0}
value={learningRate}
onChange={(e) => setLearningRate(Number(e.target.value) || 0.0001)}
className="h-8 text-xs"
/>
</div>
<div className="grid gap-1.5">
<Label className="text-xs">Precision</Label>
<select
value={precision}
onChange={(e) => setPrecision(e.target.value as "bf16" | "fp16" | "no")}
className={selectClass}
>
<option value="bf16">bf16 (default)</option>
<option value="fp16">fp16 (older GPUs)</option>
<option value="no">fp32 (no mixed)</option>
</select>
</div>
</div>
</>
)}
{status && status.status !== "idle" && (
<div className="rounded-lg border border-border bg-muted/30 p-3 text-xs">
<div className="mb-1 flex items-center justify-between">
<span className="font-medium capitalize">{status.status}</span>
<span className="text-muted-foreground">
{status.total_steps > 0 ? `${status.step}/${status.total_steps}` : ""}
</span>
</div>
<div className="mb-2 h-1.5 w-full overflow-hidden rounded-full bg-border">
<div className="h-full bg-primary transition-all" style={{ width: `${pct}%` }} />
</div>
<div className="text-muted-foreground">
{completed
? "Adapter ready - find it in the LoRA picker on this page."
: status.message}
{status.loss != null && !completed && <> · loss {status.loss.toFixed(4)}</>}
{status.lora_path && (
<div className="mt-1 break-all text-[11px]">Saved: {status.lora_path}</div>
)}
</div>
</div>
)}
</div>
<DialogFooter>
{active ? (
<Button type="button" variant="destructive" onClick={onStop}>
Stop
</Button>
) : completed ? (
<>
<Button type="button" variant="secondary" onClick={() => onOpenChange(false)}>
Done - open the LoRA picker
</Button>
<Button type="button" onClick={onStart} disabled={starting || uploading}>
{starting ? "Starting..." : "Train another"}
</Button>
</>
) : (
<Button type="button" onClick={onStart} disabled={starting || uploading}>
{starting ? "Starting..." : "Start training"}
</Button>
)}
</DialogFooter>
</DialogContent>
</Dialog>
);
}

1
studio/frontend/src/features/images/images-page.tsx
View file

@ -2006,7 +2006,6 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
)}
</div>
</div>
{/* Train mode: the full-page training workspace. Kept unmounted in Create mode so its
polling stops; Create's own state (gallery, model, workflow) is untouched. */}
{pageMode === "train" ? (

193
studio/install_sd_cpp_prebuilt.py
View file

@ -38,14 +38,18 @@ import zipfile
from pathlib import Path
from typing import Optional, Sequence
# Default upstream source. Overridable with UNSLOTH_SD_CPP_REPO so a pinned unslothai
# mirror (built the same way as unslothai/llama.cpp's prebuilts) can be used without a
# code change once it exists; otherwise this falls back to leejet upstream.
DEFAULT_REPO = "leejet/stable-diffusion.cpp"
# Default source: the Unsloth-built mirror, whose CPU/Apple prebuilts are compiled and
# published by unslothai/stable-diffusion.cpp (the same way unslothai/llama.cpp ships its
# prebuilts). Override with UNSLOTH_SD_CPP_REPO to point elsewhere (e.g. back to leejet).
# GPU hosts never reach here -- they run diffusers -- so only CPU/Apple assets are needed.
DEFAULT_REPO = "unslothai/stable-diffusion.cpp"
# Upstream we fall back to if the mirror can't serve this host (mirror release missing, or
# a host we don't yet build): resolve against leejet so native install still works.
UPSTREAM_FALLBACK_REPO = "leejet/stable-diffusion.cpp"
# Pinned release tag for REPRODUCIBILITY: "releases/latest" silently swaps the binary
# under users on every upstream push. Override with UNSLOTH_SD_CPP_TAG; set it empty to
# track latest. If the pinned tag is gone upstream, install falls back to latest.
DEFAULT_TAG = "master-737-3b6c9ca"
# under users on every push. Override with UNSLOTH_SD_CPP_TAG; set it empty to track
# latest. If the pinned tag is gone, install falls back to that repo's latest.
DEFAULT_TAG = "master-741-484baa4"
# Back-compat alias (some callers/tests import REPO).
REPO = DEFAULT_REPO
@ -118,14 +122,24 @@ def resolve_release_asset(
pool = [a for a in zips if "bin-win" in a.lower()]
token = _WINDOWS_ACCEL_TOKEN.get(accel, accel)
sel = [a for a in pool if token in a.lower()]
if not sel: # fall back to a plain avx2 CPU build
sel = [a for a in pool if "avx2" in a.lower()]
return sel[0] if sel else (pool[0] if pool else None)
if sel:
return sel[0]
# An EXPLICIT GPU accelerator with no matching asset is a real miss, not a CPU
# request: return None so the caller can fall back to a repo that builds it,
# rather than silently installing a CPU build for a --accelerator cuda request.
if accel in ("cuda", "vulkan", "rocm"):
return None
# auto / cpu -> a plain avx2 CPU build, else any windows build.
cpu = [a for a in pool if "avx2" in a.lower()]
return cpu[0] if cpu else (pool[0] if pool else None)
# linux (and anything else unix-like)
pool = [a for a in zips if "linux" in a.lower() and any(t in a.lower() for t in arch)]
if accel in _LINUX_ACCEL_TOKEN:
sel = [a for a in pool if _LINUX_ACCEL_TOKEN[accel] in a.lower()]
if accel in ("cuda", "vulkan", "rocm"):
# Explicit GPU accelerator: require its marker, else no match (let the caller
# fall back) -- never hand back a plain CPU build for a GPU request.
marker = _LINUX_ACCEL_TOKEN.get(accel, accel)
sel = [a for a in pool if marker in a.lower()]
else: # auto / cpu -> the plain build with no accelerator marker
sel = [a for a in pool if not any(m in a.lower() for m in _LINUX_ACCEL_MARKERS)]
return sel[0] if sel else None
@ -137,10 +151,14 @@ def _fetch_release(
repo: Optional[str] = None,
token: Optional[str] = None,
timeout: float = 30.0,
) -> dict:
"""GET a release JSON from GitHub. With ``tag`` set, fetch that exact release (and fall
back to latest if the tag is gone upstream); otherwise fetch latest. ``token`` is
optional and lifts the API rate limit."""
allow_latest: bool = True,
) -> Optional[dict]:
"""GET a release JSON from GitHub. With ``tag`` set, fetch that exact release; otherwise
fetch latest. ``token`` is optional and lifts the API rate limit.
When the pinned ``tag`` is missing (404): if ``allow_latest`` fall back to that repo's
latest, else return ``None`` so the caller can try the SAME pin on another repo before
settling for any repo's unpinned latest."""
repo = repo or _repo()
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
@ -155,9 +173,11 @@ def _fetch_release(
if tag:
try:
return _get(f"{base}/tags/{tag}")
except urllib.error.HTTPError as exc: # pinned tag removed upstream -> latest
except urllib.error.HTTPError as exc: # pinned tag removed -> maybe latest
if exc.code != 404:
raise
if not allow_latest:
return None
print(
f"sd-cli: pinned tag {tag} not found on {repo}; falling back to latest", flush = True
)
@ -213,6 +233,17 @@ def _locate_sd_cli(root: Path) -> Optional[Path]:
return None
def _locate_sd_server(root: Path) -> Optional[Path]:
"""The persistent ``sd-server`` binary in the extracted tree, if the archive ships
one (modern stable-diffusion.cpp releases do). Best-effort: the native backend
falls back to one-shot ``sd-cli`` when it is absent."""
name = "sd-server.exe" if sys.platform == "win32" else "sd-server"
for p in root.rglob(name):
if p.is_file():
return p
return None
def _download(
url: str,
dest: Path,
@ -272,6 +303,90 @@ def _maybe_fetch_windows_cudart(release: dict, chosen: str, target: Path) -> Non
dest.unlink(missing_ok = True)
def _resolve_repo_asset(
repo: str,
tag: Optional[str],
accelerator: str,
token: Optional[str],
*,
allow_latest: bool = True,
) -> tuple[Optional[dict], Optional[str]]:
"""Fetch ``repo``'s release and pick the asset for this host. Returns
``(release, asset_name)`` or ``(None, None)`` when the repo has no usable release
(fetch failed, or the pinned tag is missing and ``allow_latest`` is False) or no
asset for this host, so the caller can fall back."""
try:
release = _fetch_release(tag, repo = repo, token = token, allow_latest = allow_latest)
except Exception as exc: # noqa: BLE001 - network / rate limit -> fall back
print(f"sd-cli: {repo} release fetch failed ({exc})", flush = True)
return None, None
if release is None: # pinned tag missing and the latest fallback was withheld
return None, None
names = [a["name"] for a in (release.get("assets") or [])]
chosen = resolve_release_asset(
names,
system = platform.system(),
machine = platform.machine(),
accelerator = accelerator,
)
return release, chosen
def _resolve_with_fallback(
accelerator: str, token: Optional[str]
) -> tuple[str, Optional[dict], Optional[str]]:
"""Resolve ``(used_repo, release, asset_name)`` for this host across the primary repo
and -- only when the built-in default is in use and the user did not pin a repo -- the
upstream fallback.
Ordering guarantees reproducibility: a pinned tag is tried EXACTLY on every candidate
repo before any repo's unpinned latest, so a mirror that is missing the pinned release
prefers the pinned upstream build over an unpinned mirror-latest. Returns
``(primary, None, None)`` when nothing serves this host. Shared by ``install`` and
``--print-asset`` so both honour the same fallback."""
tag = _pinned_tag()
primary = _repo()
# Fall back to upstream ONLY when the user did not pin a repo (env unset) and the
# built-in default is in use: an explicit UNSLOTH_SD_CPP_REPO (even one equal to the
# default) gets exactly that repo, with no surprise upstream substitution.
repo_pinned = bool((os.environ.get("UNSLOTH_SD_CPP_REPO") or "").strip())
allow_upstream = (
not repo_pinned and primary == DEFAULT_REPO and DEFAULT_REPO != UPSTREAM_FALLBACK_REPO
)
# (repo, tag_to_fetch, allow_latest). With a pin set, try the exact pin on every repo
# first (allow_latest = False), then each repo's latest (tag = None).
attempts: list[tuple[str, Optional[str], bool]] = []
if tag:
attempts.append((primary, tag, False))
if allow_upstream:
attempts.append((UPSTREAM_FALLBACK_REPO, tag, False))
attempts.append((primary, None, True))
if allow_upstream:
attempts.append((UPSTREAM_FALLBACK_REPO, None, True))
else:
attempts.append((primary, None, True))
if allow_upstream:
attempts.append((UPSTREAM_FALLBACK_REPO, None, True))
for repo, want_tag, allow_latest in attempts:
release, chosen = _resolve_repo_asset(
repo, want_tag, accelerator, token, allow_latest = allow_latest
)
if release is not None and chosen:
if repo != primary:
# Diagnostic goes to stderr, not stdout: --print-asset documents its
# stdout as the asset name only, so a caller parsing it as a single line
# must not see this fallback log mixed in.
print(
f"falling back to {repo} for {platform.system()}/{platform.machine()}",
file = sys.stderr,
flush = True,
)
return repo, release, chosen
return primary, None, None
def install(
*,
install_dir: Optional[Path] = None,
@ -280,25 +395,22 @@ def install(
) -> Path:
"""Download + extract the prebuilt for this host. Returns the sd-cli path.
Raises ``RuntimeError`` if no asset matches the host (the caller should then
build from source) or the archive has no ``sd-cli``.
Resolves against the Unsloth mirror (``DEFAULT_REPO``) first; if the mirror can't
serve this host (release missing, or a host we don't build) AND the default repo is
in use, falls back to leejet upstream so native install still works. Raises
``RuntimeError`` only when neither source has an asset for the host, or the archive
has no ``sd-cli``.
"""
target = install_dir or default_install_dir()
release = _fetch_release(_pinned_tag(), token = token)
print(f"sd-cli: source {_repo()} release {release.get('tag_name', '?')}", flush = True)
names = [a["name"] for a in release.get("assets", [])]
chosen = resolve_release_asset(
names,
system = platform.system(),
machine = platform.machine(),
accelerator = accelerator,
)
if not chosen:
used_repo, release, chosen = _resolve_with_fallback(accelerator, token)
if release is None or not chosen:
raise RuntimeError(
f"No prebuilt sd-cli for {platform.system()}/{platform.machine()} "
f"(accelerator={accelerator}). Build from source: "
f"https://github.com/{_repo()}"
f"(accelerator={accelerator}) from {used_repo}. Build from source: "
f"https://github.com/{used_repo}"
)
print(f"sd-cli: source {used_repo} release {release.get('tag_name', '?')}", flush = True)
asset = next(a for a in release["assets"] if a["name"] == chosen)
url = asset["browser_download_url"]
target.mkdir(parents = True, exist_ok = True)
@ -323,6 +435,13 @@ def install(
if sys.platform != "win32":
_make_executable(sd_cli)
print(f"installed sd-cli -> {sd_cli}", flush = True)
# The same archive ships the persistent sd-server; make it runnable too so the
# native backend can prefer it (load once, serve many) over one-shot sd-cli.
sd_server = _locate_sd_server(target)
if sd_server is not None and sys.platform != "win32":
_make_executable(sd_server)
if sd_server is not None:
print(f"installed sd-server -> {sd_server}", flush = True)
return sd_cli
@ -338,14 +457,10 @@ def main(argv: Optional[list[str]] = None) -> int:
args = p.parse_args(argv)
if args.print_asset:
release = _fetch_release(_pinned_tag())
names = [a["name"] for a in release.get("assets", [])]
chosen = resolve_release_asset(
names,
system = platform.system(),
machine = platform.machine(),
accelerator = args.accelerator,
)
# Route through the same primary/fallback resolution as install(), so a host the
# mirror does not build (e.g. a Linux Vulkan request) reports the upstream asset
# it would actually download instead of a false "no matching prebuilt".
_used, _release, chosen = _resolve_with_fallback(args.accelerator, None)
print(chosen or "(no matching prebuilt; build from source)")
return 0 if chosen else 2