Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac
Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images.
This commit is contained in:
parent
7ca5573498
commit
a7b8f825da
7 changed files with 1298 additions and 0 deletions
114
scripts/sd_cpp_smoke.py
Normal file
114
scripts/sd_cpp_smoke.py
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""End-to-end smoke for the native stable-diffusion.cpp engine.
|
||||
|
||||
Drives the real ``SdCppEngine`` over a built ``sd-cli`` and a set of split GGUF
|
||||
assets (the same the diffusers path consumes), running one txt2img generation
|
||||
and reporting wall time. This is the GPU/native analogue of
|
||||
``scripts/diffusion_bench.py``: it proves the engine wiring (finder -> command
|
||||
builder -> subprocess -> output PNG) works against real weights.
|
||||
|
||||
Example (Z-Image-Turbo on one GPU):
|
||||
|
||||
SD_CLI_PATH=.../sd-cli CUDA_VISIBLE_DEVICES=6 python scripts/sd_cpp_smoke.py \\
|
||||
--family z-image \\
|
||||
--diffusion-model .../z-image-turbo-Q4_K_M.gguf \\
|
||||
--vae .../ae.safetensors \\
|
||||
--llm .../Qwen3-4B-Instruct-2507-Q4_K_M.gguf \\
|
||||
--memory-mode balanced --steps 8 --cfg-scale 1.0 --width 512 --height 512
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend"
|
||||
if str(_BACKEND_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND_ROOT))
|
||||
|
||||
from core.inference.diffusion_memory import ( # noqa: E402
|
||||
MEMORY_MODE_BALANCED,
|
||||
MEMORY_MODE_FAST,
|
||||
MEMORY_MODE_LOW_VRAM,
|
||||
OFFLOAD_GROUP,
|
||||
OFFLOAD_MODEL,
|
||||
OFFLOAD_NONE,
|
||||
)
|
||||
from core.inference.sd_cpp_args import ( # noqa: E402
|
||||
SdCppGenParams,
|
||||
SdCppModelFiles,
|
||||
offload_flags,
|
||||
)
|
||||
from core.inference.sd_cpp_engine import SdCppEngine, find_sd_cpp_binary # noqa: E402
|
||||
|
||||
# memory-mode (user knob) -> sd.cpp offload policy, matching the diffusers planner.
|
||||
_MODE_TO_POLICY = {
|
||||
MEMORY_MODE_FAST: OFFLOAD_NONE,
|
||||
MEMORY_MODE_BALANCED: OFFLOAD_GROUP,
|
||||
MEMORY_MODE_LOW_VRAM: OFFLOAD_MODEL,
|
||||
}
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description = "Native sd-cli engine smoke test.")
|
||||
p.add_argument("--binary", default = None, help = "sd-cli path (else env / finder)")
|
||||
p.add_argument("--family", default = "z-image")
|
||||
p.add_argument("--diffusion-model", required = True)
|
||||
p.add_argument("--vae", default = None)
|
||||
p.add_argument("--clip_l", default = None)
|
||||
p.add_argument("--t5xxl", default = None)
|
||||
p.add_argument("--llm", default = None)
|
||||
p.add_argument("--qwen2vl", default = None)
|
||||
p.add_argument("--prompt", default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed")
|
||||
p.add_argument("--negative-prompt", default = None)
|
||||
p.add_argument("--width", type = int, default = 512)
|
||||
p.add_argument("--height", type = int, default = 512)
|
||||
p.add_argument("--steps", type = int, default = 8)
|
||||
p.add_argument("--cfg-scale", type = float, default = 1.0)
|
||||
p.add_argument("--seed", type = int, default = 42)
|
||||
p.add_argument("--memory-mode", default = "balanced", choices = list(_MODE_TO_POLICY))
|
||||
p.add_argument("--out-image", default = "outputs/sdcpp_verify/sdcpp_smoke.png")
|
||||
p.add_argument("--timeout", type = float, default = 1800.0)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
binary = args.binary or find_sd_cpp_binary()
|
||||
engine = SdCppEngine(binary = binary)
|
||||
print(f"binary: {engine.binary}", flush = True)
|
||||
print(f"available: {engine.is_available()}", flush = True)
|
||||
print(f"version: {engine.version()}", flush = True)
|
||||
if not engine.is_available():
|
||||
print("ERROR: sd-cli not found (set --binary / SD_CLI_PATH / UNSLOTH_SD_CPP_PATH).", flush = True)
|
||||
return 2
|
||||
|
||||
files = SdCppModelFiles(
|
||||
diffusion_model = args.diffusion_model, vae = args.vae,
|
||||
clip_l = args.clip_l, t5xxl = args.t5xxl, llm = args.llm, qwen2vl = args.qwen2vl,
|
||||
)
|
||||
params = SdCppGenParams(
|
||||
prompt = args.prompt, negative_prompt = args.negative_prompt,
|
||||
width = args.width, height = args.height, steps = args.steps,
|
||||
cfg_scale = args.cfg_scale, seed = args.seed,
|
||||
)
|
||||
policy = _MODE_TO_POLICY[args.memory_mode]
|
||||
off = offload_flags(policy)
|
||||
print(f"memory: {args.memory_mode} -> policy={policy} -> flags={off}", flush = True)
|
||||
|
||||
out = Path(args.out_image)
|
||||
t0 = time.time()
|
||||
result = engine.generate(
|
||||
files, params, output_path = str(out), offload = off, verbose = True,
|
||||
timeout = args.timeout, on_log = lambda ln: print(f" [sd] {ln}", flush = True),
|
||||
)
|
||||
dt = time.time() - t0
|
||||
size_kb = result.stat().st_size / 1024 if result.is_file() else 0
|
||||
print(f"\nOK: generated {result} ({size_kb:.0f} KB) in {dt:.1f}s", flush = True)
|
||||
print("SD-CPP-SMOKE-OK", flush = True)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
201
studio/backend/core/inference/sd_cpp_args.py
Normal file
201
studio/backend/core/inference/sd_cpp_args.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Pure command builder for the ``sd-cli`` (stable-diffusion.cpp) native engine.
|
||||
|
||||
No torch / diffusers / subprocess here: everything is a pure function of its
|
||||
arguments so the whole argv construction can be unit-tested without a GPU, the
|
||||
binary, or any model files. The engine (``sd_cpp_engine.py``) calls
|
||||
``build_sd_cpp_command`` and runs the result.
|
||||
|
||||
stable-diffusion.cpp consumes the *same* split GGUF assets Studio already
|
||||
curates for the diffusers path: a transformer (``--diffusion-model``), a VAE
|
||||
(``--vae``), and one or more text encoders, wired to the family-specific flag
|
||||
(Z-Image's Qwen3 -> ``--llm``, Qwen-Image's Qwen2-VL -> ``--qwen2vl``,
|
||||
FLUX.1's CLIP-L + T5 -> ``--clip_l`` / ``--t5xxl``). The memory policy the
|
||||
diffusers planner picks (``none`` / ``group`` / ``model`` / ``sequential``)
|
||||
maps to sd.cpp's own offload flags here, so one user-facing knob drives both
|
||||
engines.
|
||||
|
||||
Ref (flag names): stable-diffusion.cpp ``examples/cli`` and ``docs/z_image.md``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from core.inference.diffusion_memory import (
|
||||
OFFLOAD_GROUP,
|
||||
OFFLOAD_MODEL,
|
||||
OFFLOAD_NONE,
|
||||
OFFLOAD_SEQUENTIAL,
|
||||
)
|
||||
|
||||
# Per-family text-encoder wiring: ordered (sd-cli flag, human label) pairs in the
|
||||
# order the text encoders are supplied. Z-Image uses a single Qwen3 LLM; FLUX.1
|
||||
# pairs CLIP-L with a T5; Qwen-Image uses Qwen2-VL; FLUX.2-klein uses a Qwen3 LLM.
|
||||
# Keyed by ``DiffusionFamily.name`` so the engine stays data-driven and the
|
||||
# family registry (diffusion_families.py) need not import sd.cpp specifics.
|
||||
_TE_FLAGS_BY_FAMILY: dict[str, tuple[str, ...]] = {
|
||||
"z-image": ("--llm",),
|
||||
"flux.2-klein": ("--llm",),
|
||||
"qwen-image": ("--qwen2vl",),
|
||||
"flux.1": ("--clip_l", "--t5xxl"),
|
||||
}
|
||||
|
||||
|
||||
def text_encoder_flags_for_family(family_name: str) -> tuple[str, ...]:
|
||||
"""sd-cli text-encoder flags for ``family_name`` (empty if unknown)."""
|
||||
return _TE_FLAGS_BY_FAMILY.get(family_name, ())
|
||||
|
||||
|
||||
# sd-cli's image-generation mode. Current stable-diffusion.cpp uses ``img_gen``
|
||||
# (the older spelling was ``txt2img``); img2img is the same mode with --init-img,
|
||||
# so this one token covers both text- and image-conditioned generation.
|
||||
DEFAULT_MODE = "img_gen"
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class SdCppModelFiles:
|
||||
"""Resolved on-disk paths for one diffusion checkpoint's components.
|
||||
|
||||
``diffusion_model`` (the transformer) is required; the rest are optional so a
|
||||
single-file checkpoint that bundles the VAE / encoders still builds a valid
|
||||
command. Any field left None is simply omitted from the argv.
|
||||
"""
|
||||
|
||||
diffusion_model: str
|
||||
vae: Optional[str] = None
|
||||
clip_l: Optional[str] = None
|
||||
clip_g: Optional[str] = None
|
||||
t5xxl: Optional[str] = None
|
||||
llm: Optional[str] = None
|
||||
qwen2vl: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass(frozen = True)
|
||||
class SdCppGenParams:
|
||||
"""Generation parameters, mapped 1:1 onto sd-cli's sampling flags."""
|
||||
|
||||
prompt: str
|
||||
negative_prompt: Optional[str] = None
|
||||
width: int = 1024
|
||||
height: int = 1024
|
||||
steps: Optional[int] = None
|
||||
cfg_scale: Optional[float] = None
|
||||
guidance: Optional[float] = None
|
||||
seed: Optional[int] = None
|
||||
sampling_method: Optional[str] = None
|
||||
batch_count: int = 1
|
||||
|
||||
|
||||
def offload_flags(
|
||||
policy: str,
|
||||
*,
|
||||
vae_tiling: bool = False,
|
||||
diffusion_fa: bool = False,
|
||||
) -> list[str]:
|
||||
"""Translate a diffusers memory *policy* into sd-cli offload flags.
|
||||
|
||||
- ``none``: weights stay resident, no offload flags.
|
||||
- ``group`` (balanced): stream the model through VRAM (``--offload-to-cpu``);
|
||||
flash attention shrinks the activation peak.
|
||||
- ``model`` / ``sequential`` (low-VRAM): offload everything, also push the
|
||||
CLIP/VAE to CPU and tile the VAE so the decode fits a tiny budget.
|
||||
|
||||
``vae_tiling`` / ``diffusion_fa`` force those flags on regardless of policy.
|
||||
"""
|
||||
flags: list[str] = []
|
||||
fa = diffusion_fa
|
||||
tile = vae_tiling
|
||||
if policy in (OFFLOAD_GROUP, OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL):
|
||||
flags.append("--offload-to-cpu")
|
||||
fa = True
|
||||
if policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL):
|
||||
flags.append("--clip-on-cpu")
|
||||
flags.append("--vae-on-cpu")
|
||||
tile = True
|
||||
if fa:
|
||||
flags.append("--diffusion-fa")
|
||||
if tile:
|
||||
flags.append("--vae-tiling")
|
||||
# Stable order, de-duplicated (a forced flag can coincide with a policy one).
|
||||
seen: set[str] = set()
|
||||
out: list[str] = []
|
||||
for f in flags:
|
||||
if f not in seen:
|
||||
seen.add(f)
|
||||
out.append(f)
|
||||
return out
|
||||
|
||||
|
||||
def build_sd_cpp_command(
|
||||
binary: str,
|
||||
files: SdCppModelFiles,
|
||||
params: SdCppGenParams,
|
||||
*,
|
||||
output_path: str,
|
||||
offload: Optional[list[str]] = None,
|
||||
threads: Optional[int] = None,
|
||||
verbose: bool = False,
|
||||
extra_args: Optional[list[str]] = None,
|
||||
) -> list[str]:
|
||||
"""Build the full ``sd-cli`` argv for one text-to-image generation.
|
||||
|
||||
Required model/IO flags come first (so a failure points at the obvious
|
||||
place), then sampling params (only those that are set), then offload flags,
|
||||
then any caller ``extra_args`` last so sd.cpp's last-wins parser lets a power
|
||||
user override anything Studio set.
|
||||
"""
|
||||
if not files.diffusion_model:
|
||||
raise ValueError("diffusion_model path is required")
|
||||
if not str(params.prompt).strip():
|
||||
raise ValueError("prompt is required")
|
||||
|
||||
cmd: list[str] = [binary, "--mode", DEFAULT_MODE, "--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]
|
||||
|
||||
cmd += ["--prompt", params.prompt]
|
||||
if params.negative_prompt:
|
||||
cmd += ["--negative-prompt", params.negative_prompt]
|
||||
cmd += ["--width", str(int(params.width)), "--height", str(int(params.height))]
|
||||
if params.steps is not None:
|
||||
cmd += ["--steps", str(int(params.steps))]
|
||||
if params.cfg_scale is not None:
|
||||
cmd += ["--cfg-scale", _fmt_float(params.cfg_scale)]
|
||||
if params.guidance is not None:
|
||||
cmd += ["--guidance", _fmt_float(params.guidance)]
|
||||
if params.sampling_method:
|
||||
cmd += ["--sampling-method", str(params.sampling_method)]
|
||||
if params.seed is not None:
|
||||
cmd += ["--seed", str(int(params.seed))]
|
||||
if params.batch_count and params.batch_count != 1:
|
||||
cmd += ["--batch-count", str(int(params.batch_count))]
|
||||
|
||||
cmd += ["--output", output_path]
|
||||
if threads is not None:
|
||||
cmd += ["--threads", str(int(threads))]
|
||||
if offload:
|
||||
cmd += list(offload)
|
||||
if verbose:
|
||||
cmd += ["-v"]
|
||||
if extra_args:
|
||||
cmd += list(extra_args)
|
||||
return cmd
|
||||
|
||||
|
||||
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)."""
|
||||
f = float(value)
|
||||
return str(int(f)) if f.is_integer() else repr(f)
|
||||
289
studio/backend/core/inference/sd_cpp_engine.py
Normal file
289
studio/backend/core/inference/sd_cpp_engine.py
Normal file
|
|
@ -0,0 +1,289 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Native stable-diffusion.cpp (``sd-cli``) engine for the diffusion backend.
|
||||
|
||||
This mirrors the chat backend's llama.cpp shell-out: locate a prebuilt / built
|
||||
binary, then run it as a one-shot subprocess. It is the CPU / Apple-Silicon
|
||||
(and low-VRAM) tier of the two-engine strategy -- the diffusers path stays the
|
||||
default on CUDA / ROCm / XPU, this covers the hardware diffusers serves poorly.
|
||||
|
||||
Kept deliberately thin:
|
||||
* ``find_sd_cpp_binary()`` -- env override -> ``~/.unsloth/stable-diffusion.cpp``
|
||||
build layouts -> in-tree build -> PATH. Same precedence as the llama finder.
|
||||
* ``SdCppEngine`` -- ``is_available`` / ``version`` probe + a ``generate`` that
|
||||
builds the argv (``sd_cpp_args``), runs ``sd-cli``, and returns the PNG path.
|
||||
* ``select_diffusion_engine(...)`` -- the routing decision (which backend gets
|
||||
diffusers vs. sd.cpp), a pure function so the device layer can call it.
|
||||
|
||||
Everything heavy (the subprocess, the binary) is reached only inside ``generate``
|
||||
/ ``version`` so importing this module is free and unit tests stay hermetic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
from core.inference.sd_cpp_args import (
|
||||
SdCppGenParams,
|
||||
SdCppModelFiles,
|
||||
build_sd_cpp_command,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Binary name: sd-cli (sd-cli.exe on Windows). The stable-diffusion.cpp CMake
|
||||
# target is ``sd-cli``; older builds shipped ``sd`` -- both are probed on PATH.
|
||||
_BINARY_STEM = "sd-cli"
|
||||
_LEGACY_STEM = "sd"
|
||||
|
||||
|
||||
def _binary_name(stem: str) -> str:
|
||||
return f"{stem}.exe" if sys.platform == "win32" else stem
|
||||
|
||||
|
||||
def _lib_path_var() -> str:
|
||||
"""The platform's shared-library search env var."""
|
||||
if sys.platform == "darwin":
|
||||
return "DYLD_LIBRARY_PATH"
|
||||
if sys.platform == "win32":
|
||||
return "PATH"
|
||||
return "LD_LIBRARY_PATH"
|
||||
|
||||
|
||||
def runtime_env(binary: str, base_env: Optional[dict[str, str]] = None) -> dict[str, str]:
|
||||
"""Environment that lets ``binary`` find its bundled shared libraries.
|
||||
|
||||
The prebuilt archives ship ``libstable-diffusion.so`` (/ ``.dylib`` / DLLs)
|
||||
next to ``sd-cli``, so prepend the binary's own directory to the platform
|
||||
library path. A locally-built binary that is already linked finds its libs
|
||||
regardless, so this is harmless there.
|
||||
"""
|
||||
env = dict(os.environ if base_env is None else base_env)
|
||||
var = _lib_path_var()
|
||||
bindir = str(Path(binary).resolve().parent)
|
||||
existing = env.get(var, "")
|
||||
env[var] = bindir + (os.pathsep + existing if existing else "")
|
||||
return env
|
||||
|
||||
|
||||
def _layout_candidates(root: Path) -> list[Path]:
|
||||
"""sd-cli 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)
|
||||
cands = [
|
||||
root / "build" / "bin" / name,
|
||||
root / "build" / "bin" / "Release" / name,
|
||||
root / "bin" / name,
|
||||
root / name,
|
||||
]
|
||||
return cands
|
||||
|
||||
|
||||
def find_sd_cpp_binary() -> Optional[str]:
|
||||
"""Locate the ``sd-cli`` binary, or 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.
|
||||
2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir.
|
||||
3. ``~/.unsloth/stable-diffusion.cpp`` build layouts (the installer target).
|
||||
4. ``./stable-diffusion.cpp`` in-tree build (developer checkout).
|
||||
5. ``sd-cli`` (then legacy ``sd``) on PATH.
|
||||
"""
|
||||
|
||||
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")
|
||||
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)))
|
||||
if hit:
|
||||
return hit
|
||||
|
||||
# 3. Default install root (sibling of ~/.unsloth/llama.cpp).
|
||||
hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp"))
|
||||
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"))
|
||||
if hit:
|
||||
return hit
|
||||
except (OSError, IndexError):
|
||||
pass
|
||||
|
||||
# 5. PATH.
|
||||
for stem in (_BINARY_STEM, _LEGACY_STEM):
|
||||
on_path = shutil.which(stem)
|
||||
if on_path:
|
||||
return on_path
|
||||
return None
|
||||
|
||||
|
||||
class SdCppEngine:
|
||||
"""A thin handle over a located ``sd-cli`` binary.
|
||||
|
||||
Construct it (cheap -- just resolves the path), check ``is_available``, then
|
||||
call ``generate``. Holds no process: each generation is an independent
|
||||
one-shot ``sd-cli`` run, so there is nothing to leak or clean up.
|
||||
"""
|
||||
|
||||
def __init__(self, binary: Optional[str] = None) -> None:
|
||||
self.binary = binary or find_sd_cpp_binary()
|
||||
self._version: Optional[str] = None
|
||||
|
||||
def is_available(self) -> bool:
|
||||
return bool(self.binary) and Path(self.binary).is_file()
|
||||
|
||||
def version(self, *, timeout: float = 10.0) -> Optional[str]:
|
||||
"""First line of ``sd-cli --version``, cached. None if it can't run."""
|
||||
if not self.is_available():
|
||||
return None
|
||||
if self._version is not None:
|
||||
return self._version
|
||||
try:
|
||||
res = subprocess.run(
|
||||
[self.binary, "--version"],
|
||||
capture_output = True,
|
||||
text = True,
|
||||
errors = "replace",
|
||||
timeout = timeout,
|
||||
check = False,
|
||||
env = runtime_env(self.binary),
|
||||
)
|
||||
text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip()
|
||||
self._version = text.splitlines()[0] if text else ""
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
self._version = ""
|
||||
return self._version
|
||||
|
||||
def generate(
|
||||
self,
|
||||
files: SdCppModelFiles,
|
||||
params: SdCppGenParams,
|
||||
*,
|
||||
output_path: str,
|
||||
offload: Optional[list[str]] = None,
|
||||
threads: Optional[int] = None,
|
||||
verbose: bool = False,
|
||||
extra_args: Optional[list[str]] = None,
|
||||
timeout: Optional[float] = 1800.0,
|
||||
env: Optional[dict[str, str]] = None,
|
||||
on_log: Optional[Callable[[str], None]] = None,
|
||||
) -> Path:
|
||||
"""Run one ``sd-cli`` generation; return the written image path.
|
||||
|
||||
Raises ``RuntimeError`` if the binary is missing, the process exits
|
||||
nonzero, or no output file is produced. ``on_log`` (if given) receives
|
||||
each line of sd-cli's progress output as it arrives.
|
||||
"""
|
||||
if not self.is_available():
|
||||
raise RuntimeError(
|
||||
"sd-cli (stable-diffusion.cpp) binary not found. Build it or set "
|
||||
"SD_CLI_PATH / UNSLOTH_SD_CPP_PATH."
|
||||
)
|
||||
out = Path(output_path)
|
||||
out.parent.mkdir(parents = True, exist_ok = True)
|
||||
cmd = build_sd_cpp_command(
|
||||
self.binary, files, params,
|
||||
output_path = str(out), offload = offload, threads = threads,
|
||||
verbose = verbose, extra_args = extra_args,
|
||||
)
|
||||
base = dict(os.environ)
|
||||
if env:
|
||||
base.update(env)
|
||||
run_env = runtime_env(self.binary, base)
|
||||
logger.info("sd-cli generate: %s", " ".join(cmd))
|
||||
|
||||
t0 = time.time()
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.STDOUT,
|
||||
text = True,
|
||||
errors = "replace",
|
||||
env = run_env,
|
||||
)
|
||||
tail: list[str] = []
|
||||
try:
|
||||
assert proc.stdout is not None
|
||||
for line in proc.stdout:
|
||||
line = line.rstrip("\n")
|
||||
tail.append(line)
|
||||
if len(tail) > 40:
|
||||
tail.pop(0)
|
||||
if on_log is not None:
|
||||
on_log(line)
|
||||
ret = proc.wait(timeout = timeout)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
raise RuntimeError(f"sd-cli timed out after {timeout}s")
|
||||
finally:
|
||||
if proc.poll() is None:
|
||||
proc.kill()
|
||||
|
||||
if ret != 0:
|
||||
raise RuntimeError(
|
||||
f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])
|
||||
)
|
||||
if not out.is_file():
|
||||
raise RuntimeError(
|
||||
f"sd-cli reported success but no image at {out}. Last output:\n"
|
||||
+ "\n".join(tail[-12:])
|
||||
)
|
||||
logger.info("sd-cli generate ok in %.1fs -> %s", time.time() - t0, out)
|
||||
return out
|
||||
|
||||
|
||||
# ── engine routing ──────────────────────────────────────────────────────────
|
||||
|
||||
ENGINE_DIFFUSERS = "diffusers"
|
||||
ENGINE_SD_CPP = "sd_cpp"
|
||||
|
||||
# Backends diffusers serves well with GPU acceleration. Everything else (CPU,
|
||||
# Apple MPS) is where the native engine earns its place.
|
||||
_GPU_BACKENDS = frozenset({"cuda", "rocm", "xpu"})
|
||||
|
||||
|
||||
def select_diffusion_engine(
|
||||
backend: str,
|
||||
*,
|
||||
native_available: bool,
|
||||
prefer_native: bool = False,
|
||||
) -> str:
|
||||
"""Choose the engine for a resolved device ``backend``.
|
||||
|
||||
- ``prefer_native`` + an available binary always wins (a user can force the
|
||||
native engine even on a CUDA box, e.g. to fit a tiny VRAM budget).
|
||||
- CPU / MPS route to sd.cpp when the binary is available, else fall back to
|
||||
diffusers (which still runs there, just slowly).
|
||||
- CUDA / ROCm / XPU stay on diffusers.
|
||||
"""
|
||||
if prefer_native and native_available:
|
||||
return ENGINE_SD_CPP
|
||||
if backend not in _GPU_BACKENDS and native_available:
|
||||
return ENGINE_SD_CPP
|
||||
return ENGINE_DIFFUSERS
|
||||
156
studio/backend/tests/test_sd_cpp_args.py
Normal file
156
studio/backend/tests/test_sd_cpp_args.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the sd-cli command builder (``sd_cpp_args.py``).
|
||||
|
||||
Pure: no torch, no subprocess, no files. Just argv construction and the
|
||||
policy -> offload-flag / family -> text-encoder-flag mappings.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference.diffusion_memory import (
|
||||
OFFLOAD_GROUP,
|
||||
OFFLOAD_MODEL,
|
||||
OFFLOAD_NONE,
|
||||
OFFLOAD_SEQUENTIAL,
|
||||
)
|
||||
from core.inference.sd_cpp_args import (
|
||||
SdCppGenParams,
|
||||
SdCppModelFiles,
|
||||
build_sd_cpp_command,
|
||||
offload_flags,
|
||||
text_encoder_flags_for_family,
|
||||
)
|
||||
|
||||
|
||||
def _pair(cmd: list[str], flag: str):
|
||||
"""Value following ``flag`` in ``cmd``, or None if the flag is absent."""
|
||||
return cmd[cmd.index(flag) + 1] if flag in cmd else None
|
||||
|
||||
|
||||
# ── family text-encoder wiring ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_te_flags_by_family():
|
||||
assert text_encoder_flags_for_family("z-image") == ("--llm",)
|
||||
assert text_encoder_flags_for_family("qwen-image") == ("--qwen2vl",)
|
||||
assert text_encoder_flags_for_family("flux.1") == ("--clip_l", "--t5xxl")
|
||||
assert text_encoder_flags_for_family("flux.2-klein") == ("--llm",)
|
||||
assert text_encoder_flags_for_family("unknown") == ()
|
||||
|
||||
|
||||
# ── offload policy -> sd-cli flags ──────────────────────────────────────────
|
||||
|
||||
|
||||
def test_offload_none_is_empty():
|
||||
assert offload_flags(OFFLOAD_NONE) == []
|
||||
|
||||
|
||||
def test_offload_group_streams_with_flash_attention():
|
||||
flags = offload_flags(OFFLOAD_GROUP)
|
||||
assert "--offload-to-cpu" in flags
|
||||
assert "--diffusion-fa" in flags
|
||||
# group keeps CLIP/VAE resident -> no per-component cpu flags
|
||||
assert "--clip-on-cpu" not in flags
|
||||
assert "--vae-on-cpu" not in flags
|
||||
|
||||
|
||||
def test_offload_model_pushes_everything_to_cpu_and_tiles():
|
||||
flags = offload_flags(OFFLOAD_MODEL)
|
||||
for expected in ("--offload-to-cpu", "--clip-on-cpu", "--vae-on-cpu", "--vae-tiling", "--diffusion-fa"):
|
||||
assert expected in flags
|
||||
# sequential maps the same as model
|
||||
assert offload_flags(OFFLOAD_SEQUENTIAL) == flags
|
||||
|
||||
|
||||
def test_offload_forced_flags_dedup():
|
||||
# vae_tiling/diffusion_fa forced on with a policy that already sets them
|
||||
flags = offload_flags(OFFLOAD_MODEL, vae_tiling = True, diffusion_fa = True)
|
||||
assert flags.count("--vae-tiling") == 1
|
||||
assert flags.count("--diffusion-fa") == 1
|
||||
# forced on top of a no-offload policy
|
||||
none_forced = offload_flags(OFFLOAD_NONE, vae_tiling = True, diffusion_fa = True)
|
||||
assert none_forced == ["--diffusion-fa", "--vae-tiling"]
|
||||
|
||||
|
||||
# ── full command construction ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_build_zimage_command_minimal():
|
||||
files = SdCppModelFiles(
|
||||
diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/qwen3.gguf",
|
||||
)
|
||||
params = SdCppGenParams(prompt = "a cat", width = 512, height = 768, steps = 8, cfg_scale = 1.0, seed = 42)
|
||||
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/out/x.png")
|
||||
|
||||
assert cmd[0] == "/bin/sd-cli"
|
||||
assert _pair(cmd, "--mode") == "img_gen"
|
||||
assert _pair(cmd, "--diffusion-model") == "/m/z.gguf"
|
||||
assert _pair(cmd, "--vae") == "/m/ae.sft"
|
||||
assert _pair(cmd, "--llm") == "/m/qwen3.gguf"
|
||||
assert _pair(cmd, "--prompt") == "a cat"
|
||||
assert _pair(cmd, "--width") == "512"
|
||||
assert _pair(cmd, "--height") == "768"
|
||||
assert _pair(cmd, "--steps") == "8"
|
||||
assert _pair(cmd, "--cfg-scale") == "1" # 1.0 -> "1"
|
||||
assert _pair(cmd, "--seed") == "42"
|
||||
assert _pair(cmd, "--output") == "/out/x.png"
|
||||
# unset encoders are omitted
|
||||
assert "--t5xxl" not in cmd
|
||||
assert "--qwen2vl" not in cmd
|
||||
|
||||
|
||||
def test_build_flux1_dual_text_encoders():
|
||||
files = SdCppModelFiles(
|
||||
diffusion_model = "/m/flux.gguf", vae = "/m/ae.sft",
|
||||
clip_l = "/m/clip_l.sft", t5xxl = "/m/t5.gguf",
|
||||
)
|
||||
params = SdCppGenParams(prompt = "x", guidance = 3.5)
|
||||
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
|
||||
assert _pair(cmd, "--clip_l") == "/m/clip_l.sft"
|
||||
assert _pair(cmd, "--t5xxl") == "/m/t5.gguf"
|
||||
assert _pair(cmd, "--guidance") == "3.5"
|
||||
assert "--llm" not in cmd
|
||||
|
||||
|
||||
def test_build_appends_offload_and_extra_args_last():
|
||||
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
|
||||
params = SdCppGenParams(prompt = "x")
|
||||
off = offload_flags(OFFLOAD_GROUP)
|
||||
cmd = build_sd_cpp_command(
|
||||
"/bin/sd-cli", files, params, output_path = "/o.png",
|
||||
offload = off, threads = 8, verbose = True, extra_args = ["--rng", "cuda"],
|
||||
)
|
||||
assert "--offload-to-cpu" in cmd
|
||||
assert _pair(cmd, "--threads") == "8"
|
||||
assert "-v" in cmd
|
||||
# extra args come after everything Studio set (last-wins for power users)
|
||||
assert cmd[-2:] == ["--rng", "cuda"]
|
||||
|
||||
|
||||
def test_build_negative_prompt_and_batch():
|
||||
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
|
||||
params = SdCppGenParams(prompt = "x", negative_prompt = "blurry", batch_count = 3)
|
||||
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
|
||||
assert _pair(cmd, "--negative-prompt") == "blurry"
|
||||
assert _pair(cmd, "--batch-count") == "3"
|
||||
|
||||
|
||||
def test_build_omits_unset_optional_params():
|
||||
files = SdCppModelFiles(diffusion_model = "/m/z.gguf")
|
||||
params = SdCppGenParams(prompt = "x") # no steps/cfg/seed/sampler
|
||||
cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png")
|
||||
for flag in ("--steps", "--cfg-scale", "--guidance", "--seed", "--sampling-method", "--batch-count", "--threads", "-v"):
|
||||
assert flag not in cmd
|
||||
|
||||
|
||||
def test_build_requires_diffusion_model_and_prompt():
|
||||
with pytest.raises(ValueError):
|
||||
build_sd_cpp_command("/bin/sd-cli", SdCppModelFiles(diffusion_model = ""),
|
||||
SdCppGenParams(prompt = "x"), output_path = "/o.png")
|
||||
with pytest.raises(ValueError):
|
||||
build_sd_cpp_command("/bin/sd-cli", SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
||||
SdCppGenParams(prompt = " "), output_path = "/o.png")
|
||||
230
studio/backend/tests/test_sd_cpp_engine.py
Normal file
230
studio/backend/tests/test_sd_cpp_engine.py
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the sd-cli engine + routing (``sd_cpp_engine.py``).
|
||||
|
||||
Hermetic: the binary finder is driven against a tmp filesystem, and ``generate``
|
||||
runs a fake ``subprocess.Popen`` that emits canned lines and writes the output
|
||||
PNG -- no real ``sd-cli``, no GPU.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from core.inference import sd_cpp_engine as eng
|
||||
from core.inference.sd_cpp_engine import (
|
||||
ENGINE_DIFFUSERS,
|
||||
ENGINE_SD_CPP,
|
||||
SdCppEngine,
|
||||
find_sd_cpp_binary,
|
||||
runtime_env,
|
||||
select_diffusion_engine,
|
||||
)
|
||||
from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles
|
||||
|
||||
|
||||
# ── binary discovery ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _clear_env(monkeypatch):
|
||||
monkeypatch.delenv("SD_CLI_PATH", raising = False)
|
||||
monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False)
|
||||
|
||||
|
||||
def test_find_prefers_sd_cli_path_env(tmp_path, monkeypatch):
|
||||
_clear_env(monkeypatch)
|
||||
binary = tmp_path / "sd-cli"
|
||||
binary.write_text("#!/bin/sh\n")
|
||||
monkeypatch.setenv("SD_CLI_PATH", str(binary))
|
||||
# even with PATH empty, the direct env wins
|
||||
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
||||
assert find_sd_cpp_binary() == str(binary)
|
||||
|
||||
|
||||
def test_find_custom_install_dir_build_layout(tmp_path, monkeypatch):
|
||||
_clear_env(monkeypatch)
|
||||
root = tmp_path / "sdcpp"
|
||||
built = root / "build" / "bin" / "sd-cli"
|
||||
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_cpp_binary() == str(built)
|
||||
|
||||
|
||||
def test_find_falls_back_to_path(tmp_path, monkeypatch):
|
||||
_clear_env(monkeypatch)
|
||||
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
||||
monkeypatch.setattr(eng.shutil, "which", lambda stem: "/usr/bin/sd-cli" if stem == "sd-cli" else None)
|
||||
assert find_sd_cpp_binary() == "/usr/bin/sd-cli"
|
||||
|
||||
|
||||
def test_find_returns_none_when_absent(tmp_path, monkeypatch):
|
||||
_clear_env(monkeypatch)
|
||||
monkeypatch.setattr(eng.Path, "home", staticmethod(lambda: tmp_path / "nohome"))
|
||||
monkeypatch.setattr(eng.shutil, "which", lambda *_a: None)
|
||||
assert find_sd_cpp_binary() is None
|
||||
|
||||
|
||||
# ── availability / version ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_engine_unavailable_when_no_binary():
|
||||
e = SdCppEngine(binary = None)
|
||||
assert e.is_available() is False
|
||||
assert e.version() is None
|
||||
|
||||
|
||||
def test_engine_version_parsed_and_cached(tmp_path, monkeypatch):
|
||||
binary = tmp_path / "sd-cli"
|
||||
binary.write_text("x")
|
||||
e = SdCppEngine(binary = str(binary))
|
||||
calls = {"n": 0}
|
||||
|
||||
def _fake_run(*_a, **_k):
|
||||
calls["n"] += 1
|
||||
return types.SimpleNamespace(stdout = "stable-diffusion.cpp version master-721\n", stderr = "")
|
||||
|
||||
monkeypatch.setattr(eng.subprocess, "run", _fake_run)
|
||||
assert e.version() == "stable-diffusion.cpp version master-721"
|
||||
assert e.version() == "stable-diffusion.cpp version master-721"
|
||||
assert calls["n"] == 1 # cached after the first probe
|
||||
|
||||
|
||||
# ── runtime env (bundled shared libs) ───────────────────────────────────────
|
||||
|
||||
|
||||
def test_runtime_env_prepends_binary_dir_to_lib_path():
|
||||
var = eng._lib_path_var()
|
||||
env = runtime_env("/opt/sdcpp/bin/sd-cli", {var: "/existing"})
|
||||
first = env[var].split(os.pathsep)[0]
|
||||
assert first == "/opt/sdcpp/bin"
|
||||
assert "/existing" in env[var]
|
||||
|
||||
|
||||
def test_runtime_env_handles_missing_lib_path():
|
||||
var = eng._lib_path_var()
|
||||
env = runtime_env("/opt/sdcpp/bin/sd-cli", {})
|
||||
assert env[var] == "/opt/sdcpp/bin"
|
||||
|
||||
|
||||
# ── generate (fake subprocess) ──────────────────────────────────────────────
|
||||
|
||||
|
||||
class _FakePopen:
|
||||
"""Stand-in for subprocess.Popen: streams ``lines`` then writes ``out_file``
|
||||
(unless ``write`` is False) and exits with ``returncode``."""
|
||||
|
||||
captured_cmd: list[str] = []
|
||||
captured_env: dict = {}
|
||||
|
||||
def __init__(self, cmd, *, lines, returncode, out_file, write, env = None):
|
||||
type(self).captured_cmd = list(cmd)
|
||||
type(self).captured_env = dict(env or {})
|
||||
self._lines = list(lines)
|
||||
self.returncode = returncode
|
||||
self._out_file = out_file
|
||||
self._write = write
|
||||
|
||||
@property
|
||||
def stdout(self):
|
||||
return iter(self._lines)
|
||||
|
||||
def wait(self, timeout = None):
|
||||
if self._write:
|
||||
Path(self._out_file).write_bytes(b"\x89PNG\r\n")
|
||||
return self.returncode
|
||||
|
||||
def poll(self):
|
||||
return self.returncode
|
||||
|
||||
def kill(self):
|
||||
pass
|
||||
|
||||
|
||||
def _patch_popen(monkeypatch, *, lines, returncode, out_file, write = True):
|
||||
def _factory(cmd, **kw):
|
||||
return _FakePopen(cmd, lines = lines, returncode = returncode, out_file = out_file,
|
||||
write = write, env = kw.get("env"))
|
||||
|
||||
monkeypatch.setattr(eng.subprocess, "Popen", _factory)
|
||||
|
||||
|
||||
def _engine(tmp_path):
|
||||
binary = tmp_path / "sd-cli"
|
||||
binary.write_text("x")
|
||||
return SdCppEngine(binary = str(binary))
|
||||
|
||||
|
||||
def test_generate_success_returns_path_and_collects_logs(tmp_path, monkeypatch):
|
||||
e = _engine(tmp_path)
|
||||
out = tmp_path / "img.png"
|
||||
_patch_popen(monkeypatch, lines = ["loading model", "step 1/8", "done"], returncode = 0, out_file = out)
|
||||
seen: list[str] = []
|
||||
files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf")
|
||||
params = SdCppGenParams(prompt = "a cat", steps = 8, seed = 1)
|
||||
|
||||
result = e.generate(files, params, output_path = str(out), on_log = seen.append)
|
||||
|
||||
assert result == out and out.is_file()
|
||||
assert seen == ["loading model", "step 1/8", "done"]
|
||||
# the real argv was built and handed to Popen
|
||||
assert "--diffusion-model" in _FakePopen.captured_cmd
|
||||
assert str(out) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--output") + 1]
|
||||
# the subprocess env carries the binary's dir on the library path
|
||||
var = eng._lib_path_var()
|
||||
assert str(Path(e.binary).resolve().parent) in _FakePopen.captured_env.get(var, "")
|
||||
|
||||
|
||||
def test_generate_raises_on_nonzero_exit(tmp_path, monkeypatch):
|
||||
e = _engine(tmp_path)
|
||||
out = tmp_path / "img.png"
|
||||
_patch_popen(monkeypatch, lines = ["boom: bad gguf"], returncode = 1, out_file = out, write = False)
|
||||
with pytest.raises(RuntimeError, match = "exited 1"):
|
||||
e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
||||
SdCppGenParams(prompt = "x"), output_path = str(out))
|
||||
|
||||
|
||||
def test_generate_raises_when_no_output_despite_success(tmp_path, monkeypatch):
|
||||
e = _engine(tmp_path)
|
||||
out = tmp_path / "img.png"
|
||||
_patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out, write = False)
|
||||
with pytest.raises(RuntimeError, match = "no image"):
|
||||
e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
||||
SdCppGenParams(prompt = "x"), output_path = str(out))
|
||||
|
||||
|
||||
def test_generate_raises_when_binary_missing():
|
||||
e = SdCppEngine(binary = None)
|
||||
with pytest.raises(RuntimeError, match = "not found"):
|
||||
e.generate(SdCppModelFiles(diffusion_model = "/m/z.gguf"),
|
||||
SdCppGenParams(prompt = "x"), output_path = "/tmp/x.png")
|
||||
|
||||
|
||||
# ── engine routing ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_routing_gpu_backends_use_diffusers():
|
||||
for backend in ("cuda", "rocm", "xpu"):
|
||||
assert select_diffusion_engine(backend, native_available = True) == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
def test_routing_cpu_and_mps_use_native_when_available():
|
||||
assert select_diffusion_engine("cpu", native_available = True) == ENGINE_SD_CPP
|
||||
assert select_diffusion_engine("mps", native_available = True) == ENGINE_SD_CPP
|
||||
|
||||
|
||||
def test_routing_cpu_falls_back_to_diffusers_without_binary():
|
||||
assert select_diffusion_engine("cpu", native_available = False) == ENGINE_DIFFUSERS
|
||||
|
||||
|
||||
def test_routing_prefer_native_overrides_gpu():
|
||||
assert select_diffusion_engine("cuda", native_available = True, prefer_native = True) == ENGINE_SD_CPP
|
||||
# but only if a binary is actually available
|
||||
assert select_diffusion_engine("cuda", native_available = False, prefer_native = True) == ENGINE_DIFFUSERS
|
||||
107
studio/backend/tests/test_sd_cpp_install.py
Normal file
107
studio/backend/tests/test_sd_cpp_install.py
Normal file
|
|
@ -0,0 +1,107 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Unit tests for the prebuilt sd-cli asset resolver (``install_sd_cpp_prebuilt``).
|
||||
|
||||
Pure: the host -> release-asset matrix is exercised against a fixed asset list
|
||||
(a real stable-diffusion.cpp release), no network. The installer lives under
|
||||
``studio/`` (not ``studio/backend``), so the test puts that dir on the path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
_STUDIO = Path(__file__).resolve().parents[2]
|
||||
if str(_STUDIO) not in sys.path:
|
||||
sys.path.insert(0, str(_STUDIO))
|
||||
|
||||
from install_sd_cpp_prebuilt import default_install_dir, resolve_release_asset # noqa: E402
|
||||
|
||||
# A real stable-diffusion.cpp latest-release asset list.
|
||||
_ASSETS = [
|
||||
"cudart-sd-bin-win-cu12-x64.zip",
|
||||
"sd-master-8caa3f9-bin-Darwin-macOS-15.7.7-arm64.zip",
|
||||
"sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-rocm-7.13.0.zip",
|
||||
"sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-rocm-7.2.1.zip",
|
||||
"sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-vulkan.zip",
|
||||
"sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64.zip",
|
||||
"sd-master-8caa3f9-bin-win-avx-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-avx2-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-avx512-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-cuda12-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-noavx-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-rocm-7.13.0-x64.zip",
|
||||
"sd-master-8caa3f9-bin-win-vulkan-x64.zip",
|
||||
]
|
||||
|
||||
|
||||
def _resolve(system, machine, accelerator = "auto"):
|
||||
return resolve_release_asset(_ASSETS, system = system, machine = machine, accelerator = accelerator)
|
||||
|
||||
|
||||
# ── macOS (the key Apple-Silicon target) ────────────────────────────────────
|
||||
|
||||
|
||||
def test_macos_arm64_picks_darwin_arm64():
|
||||
assert _resolve("Darwin", "arm64") == "sd-master-8caa3f9-bin-Darwin-macOS-15.7.7-arm64.zip"
|
||||
# aarch64 spelling resolves the same
|
||||
assert _resolve("Darwin", "aarch64").startswith("sd-master") and "arm64" in _resolve("Darwin", "aarch64")
|
||||
|
||||
|
||||
def test_macos_intel_has_no_prebuilt():
|
||||
# only an arm64 Darwin asset exists -> Intel Macs must build from source
|
||||
assert _resolve("Darwin", "x86_64") is None
|
||||
|
||||
|
||||
# ── Linux (CPU is the default tier) ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_linux_x86_64_auto_picks_plain_cpu_build():
|
||||
# the plain x86_64 zip, NOT a rocm/vulkan one
|
||||
assert _resolve("Linux", "x86_64") == "sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64.zip"
|
||||
|
||||
|
||||
def test_linux_vulkan_and_rocm_select_accelerator_builds():
|
||||
assert _resolve("Linux", "x86_64", "vulkan") == "sd-master-8caa3f9-bin-Linux-Ubuntu-24.04-x86_64-vulkan.zip"
|
||||
assert "rocm" in _resolve("Linux", "x86_64", "rocm")
|
||||
|
||||
|
||||
def test_linux_arm64_has_no_prebuilt():
|
||||
assert _resolve("Linux", "aarch64") is None
|
||||
|
||||
|
||||
# ── Windows ─────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_windows_auto_picks_avx2():
|
||||
assert _resolve("Windows", "AMD64") == "sd-master-8caa3f9-bin-win-avx2-x64.zip"
|
||||
|
||||
|
||||
def test_windows_cuda_picks_cuda12():
|
||||
assert _resolve("Windows", "AMD64", "cuda") == "sd-master-8caa3f9-bin-win-cuda12-x64.zip"
|
||||
|
||||
|
||||
def test_windows_vulkan_picks_vulkan():
|
||||
assert _resolve("Windows", "AMD64", "vulkan") == "sd-master-8caa3f9-bin-win-vulkan-x64.zip"
|
||||
|
||||
|
||||
# ── cudart helper archive is never chosen as the engine ─────────────────────
|
||||
|
||||
|
||||
def test_cudart_runtime_archive_never_selected():
|
||||
for accel in ("auto", "cuda", "vulkan", "rocm"):
|
||||
chosen = _resolve("Windows", "AMD64", accel)
|
||||
assert chosen is None or not chosen.startswith("cudart")
|
||||
|
||||
|
||||
# ── install dir ─────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_default_install_dir_is_sibling_of_llama(monkeypatch):
|
||||
monkeypatch.delenv("UNSLOTH_STUDIO_HOME", raising = False)
|
||||
monkeypatch.delenv("STUDIO_HOME", raising = False)
|
||||
d = default_install_dir()
|
||||
assert d.name == "stable-diffusion.cpp"
|
||||
assert d.parent.name == ".unsloth"
|
||||
201
studio/install_sd_cpp_prebuilt.py
Normal file
201
studio/install_sd_cpp_prebuilt.py
Normal file
|
|
@ -0,0 +1,201 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Install a prebuilt ``sd-cli`` (stable-diffusion.cpp) for the native diffusion
|
||||
engine.
|
||||
|
||||
The chat backend ships a prebuilt llama-server; this is the diffusion analogue,
|
||||
kept deliberately small. stable-diffusion.cpp publishes per-platform release
|
||||
zips (macOS-arm64/Metal, Linux x86_64 CPU, plus Vulkan / ROCm / Windows
|
||||
variants), so on the Phase-4 targets (Apple Silicon and CPU) there is nothing to
|
||||
compile: resolve the right asset, download, extract into
|
||||
``~/.unsloth/stable-diffusion.cpp``, and the engine's finder picks it up.
|
||||
|
||||
``resolve_release_asset`` -- the host -> asset choice -- is a pure function so the
|
||||
matching matrix is unit-tested without any network. CUDA / ROCm / XPU hosts stay
|
||||
on diffusers and never need this; it exists for the engines diffusers serves
|
||||
poorly.
|
||||
|
||||
Usage:
|
||||
python studio/install_sd_cpp_prebuilt.py # auto-detect host
|
||||
python studio/install_sd_cpp_prebuilt.py --accelerator vulkan
|
||||
python studio/install_sd_cpp_prebuilt.py --print-asset # resolve only
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import stat
|
||||
import sys
|
||||
import urllib.request
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from typing import Optional, Sequence
|
||||
|
||||
REPO = "leejet/stable-diffusion.cpp"
|
||||
RELEASES_API = f"https://api.github.com/repos/{REPO}/releases/latest"
|
||||
|
||||
# accelerator -> the token that must appear in a Linux/Windows asset name.
|
||||
_LINUX_ACCEL_TOKEN = {"rocm": "rocm", "vulkan": "vulkan"}
|
||||
_WINDOWS_ACCEL_TOKEN = {"cuda": "cuda12", "vulkan": "vulkan", "rocm": "rocm", "cpu": "avx2", "auto": "avx2"}
|
||||
# Tokens that mark an accelerator-specific Linux build; "auto"/"cpu" want none of them.
|
||||
_LINUX_ACCEL_MARKERS = ("rocm", "vulkan", "cuda", "sycl", "musa")
|
||||
|
||||
_ARCH_TOKENS = {
|
||||
"x86_64": ("x86_64", "x64", "amd64"),
|
||||
"amd64": ("x86_64", "x64", "amd64"),
|
||||
"arm64": ("arm64", "aarch64"),
|
||||
"aarch64": ("arm64", "aarch64"),
|
||||
}
|
||||
|
||||
|
||||
def _arch_tokens(machine: str) -> tuple[str, ...]:
|
||||
return _ARCH_TOKENS.get(machine.lower(), (machine.lower(),))
|
||||
|
||||
|
||||
def resolve_release_asset(
|
||||
asset_names: Sequence[str],
|
||||
*,
|
||||
system: str,
|
||||
machine: str,
|
||||
accelerator: str = "auto",
|
||||
) -> Optional[str]:
|
||||
"""Pick the best release asset for a host, or None if none matches.
|
||||
|
||||
``system`` / ``machine`` are ``platform.system()`` / ``platform.machine()``
|
||||
values; ``accelerator`` is ``auto`` (CPU/Metal default), ``vulkan``,
|
||||
``rocm``, or ``cuda`` (Windows only). Pure -- the caller passes the release's
|
||||
asset name list.
|
||||
"""
|
||||
system = system.lower()
|
||||
accel = accelerator.lower()
|
||||
arch = _arch_tokens(machine)
|
||||
zips = [a for a in asset_names if a.lower().endswith(".zip") and not a.lower().startswith("cudart")]
|
||||
|
||||
if system == "darwin":
|
||||
pool = [a for a in zips if ("darwin" in a.lower() or "macos" in a.lower())
|
||||
and any(t in a.lower() for t in arch)]
|
||||
return pool[0] if pool else None
|
||||
|
||||
if system == "windows":
|
||||
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)
|
||||
|
||||
# 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()]
|
||||
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
|
||||
|
||||
|
||||
def _fetch_latest_release(*, token: Optional[str] = None, timeout: float = 30.0) -> dict:
|
||||
"""GET the latest-release JSON from GitHub (token optional, lifts rate limit)."""
|
||||
req = urllib.request.Request(RELEASES_API, headers = {"Accept": "application/vnd.github+json"})
|
||||
token = token or os.environ.get("GH_TOKEN") or os.environ.get("GITHUB_TOKEN")
|
||||
if token:
|
||||
req.add_header("Authorization", f"Bearer {token}")
|
||||
with urllib.request.urlopen(req, timeout = timeout) as resp: # noqa: S310 (fixed https host)
|
||||
return json.loads(resp.read().decode("utf-8"))
|
||||
|
||||
|
||||
def default_install_dir() -> Path:
|
||||
"""``~/.unsloth/stable-diffusion.cpp`` (or under ``UNSLOTH_STUDIO_HOME`` /
|
||||
``STUDIO_HOME`` if set), the sibling of the llama.cpp install the finder
|
||||
probes."""
|
||||
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
|
||||
base = Path(home).parent if home else Path.home() / ".unsloth"
|
||||
return base / "stable-diffusion.cpp"
|
||||
|
||||
|
||||
def _make_executable(path: Path) -> None:
|
||||
mode = path.stat().st_mode
|
||||
path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
|
||||
|
||||
|
||||
def _locate_sd_cli(root: Path) -> Optional[Path]:
|
||||
name = "sd-cli.exe" if sys.platform == "win32" else "sd-cli"
|
||||
for p in root.rglob(name):
|
||||
if p.is_file():
|
||||
return p
|
||||
return None
|
||||
|
||||
|
||||
def install(
|
||||
*,
|
||||
install_dir: Optional[Path] = None,
|
||||
accelerator: str = "auto",
|
||||
token: Optional[str] = None,
|
||||
) -> 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``.
|
||||
"""
|
||||
target = install_dir or default_install_dir()
|
||||
release = _fetch_latest_release(token = token)
|
||||
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:
|
||||
raise RuntimeError(
|
||||
f"No prebuilt sd-cli for {platform.system()}/{platform.machine()} "
|
||||
f"(accelerator={accelerator}). Build from source: "
|
||||
f"https://github.com/{REPO}"
|
||||
)
|
||||
url = next(a["browser_download_url"] for a in release["assets"] if a["name"] == chosen)
|
||||
target.mkdir(parents = True, exist_ok = True)
|
||||
archive = target / chosen
|
||||
print(f"downloading {chosen} -> {archive}", flush = True)
|
||||
urllib.request.urlretrieve(url, archive) # noqa: S310 (github release URL)
|
||||
print("extracting ...", flush = True)
|
||||
with zipfile.ZipFile(archive) as zf:
|
||||
zf.extractall(target)
|
||||
archive.unlink(missing_ok = True)
|
||||
sd_cli = _locate_sd_cli(target)
|
||||
if not sd_cli:
|
||||
raise RuntimeError(f"archive {chosen} contained no sd-cli binary")
|
||||
if sys.platform != "win32":
|
||||
_make_executable(sd_cli)
|
||||
print(f"installed sd-cli -> {sd_cli}", flush = True)
|
||||
return sd_cli
|
||||
|
||||
|
||||
def main(argv: Optional[list[str]] = None) -> int:
|
||||
p = argparse.ArgumentParser(description = "Install a prebuilt sd-cli (stable-diffusion.cpp).")
|
||||
p.add_argument("--accelerator", default = "auto", choices = ["auto", "cpu", "vulkan", "rocm", "cuda"])
|
||||
p.add_argument("--install-dir", default = None)
|
||||
p.add_argument("--print-asset", action = "store_true", help = "resolve + print the asset, don't download")
|
||||
args = p.parse_args(argv)
|
||||
|
||||
if args.print_asset:
|
||||
release = _fetch_latest_release()
|
||||
names = [a["name"] for a in release.get("assets", [])]
|
||||
chosen = resolve_release_asset(
|
||||
names, system = platform.system(), machine = platform.machine(), accelerator = args.accelerator,
|
||||
)
|
||||
print(chosen or "(no matching prebuilt; build from source)")
|
||||
return 0 if chosen else 2
|
||||
|
||||
try:
|
||||
install(
|
||||
install_dir = Path(args.install_dir) if args.install_dir else None,
|
||||
accelerator = args.accelerator,
|
||||
)
|
||||
except RuntimeError as exc:
|
||||
print(f"error: {exc}", file = sys.stderr)
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Loading…
Add table
Add a link
Reference in a new issue