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/backend/core/inference/sd_cpp_args.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
"""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
View 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