diff --git a/scripts/sd_cpp_smoke.py b/scripts/sd_cpp_smoke.py new file mode 100644 index 0000000000..8693481d2b --- /dev/null +++ b/scripts/sd_cpp_smoke.py @@ -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()) diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py new file mode 100644 index 0000000000..876cd7e557 --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -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) diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py new file mode 100644 index 0000000000..cf645a91c5 --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -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: /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 diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py new file mode 100644 index 0000000000..9d86a4cb8b --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -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") diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py new file mode 100644 index 0000000000..8a20c73963 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -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 diff --git a/studio/backend/tests/test_sd_cpp_install.py b/studio/backend/tests/test_sd_cpp_install.py new file mode 100644 index 0000000000..ae5fc4e5d3 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -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" diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py new file mode 100644 index 0000000000..2156995356 --- /dev/null +++ b/studio/install_sd_cpp_prebuilt.py @@ -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())