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

Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac

Browse source 
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:
Daniel Han 2026-06-25 15:49:39 +00:00
commit a7b8f825da
7 changed files with 1298 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

201
studio/install_sd_cpp_prebuilt.py Normal file
View 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())