From 96940d87b7b94a88980acac38c1ce879282e4666 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:03:53 -0700 Subject: [PATCH] Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac (#6679) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/diffusion_bench.py | 9 + scripts/sd_cpp_smoke.py | 133 +++++++ studio/backend/core/inference/diffusion.py | 158 +++----- .../core/inference/diffusion_device.py | 68 +++- .../core/inference/diffusion_families.py | 4 +- .../core/inference/diffusion_memory.py | 138 +++++-- .../core/inference/diffusion_precision.py | 3 +- .../backend/core/inference/diffusion_speed.py | 26 +- studio/backend/core/inference/sd_cpp_args.py | 201 ++++++++++ .../backend/core/inference/sd_cpp_engine.py | 329 ++++++++++++++++ studio/backend/models/inference.py | 4 +- studio/backend/routes/inference.py | 1 + .../backend/tests/test_diffusion_backend.py | 241 +++--------- studio/backend/tests/test_diffusion_device.py | 39 ++ studio/backend/tests/test_diffusion_memory.py | 83 +++- .../backend/tests/test_diffusion_precision.py | 69 +--- studio/backend/tests/test_diffusion_speed.py | 41 +- studio/backend/tests/test_sd_cpp_args.py | 189 ++++++++++ studio/backend/tests/test_sd_cpp_engine.py | 354 ++++++++++++++++++ studio/backend/tests/test_sd_cpp_install.py | 166 ++++++++ studio/install_sd_cpp_prebuilt.py | 278 ++++++++++++++ 21 files changed, 2102 insertions(+), 432 deletions(-) create mode 100644 scripts/sd_cpp_smoke.py create mode 100644 studio/backend/core/inference/sd_cpp_args.py create mode 100644 studio/backend/core/inference/sd_cpp_engine.py create mode 100644 studio/backend/tests/test_sd_cpp_args.py create mode 100644 studio/backend/tests/test_sd_cpp_engine.py create mode 100644 studio/backend/tests/test_sd_cpp_install.py create mode 100644 studio/install_sd_cpp_prebuilt.py diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 0a4c66697c..dc49fe756b 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -81,6 +81,10 @@ def _percentile(values: list[float], pct: float) -> float: return ordered[rank] +def _is_cuda(device: Optional[str]) -> bool: + return bool(device) and device.split(":", 1)[0] == "cuda" + + def _cuda_reset_peak() -> None: import torch if torch.cuda.is_available(): @@ -210,6 +214,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: base_repo = args.base_repo, family_override = args.family_override, hf_token = os.environ.get("HF_TOKEN"), + cpu_offload = args.cpu_offload, memory_mode = args.memory_mode, text_encoder_quant = args.text_encoder_quant, ) @@ -290,6 +295,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "seed": args.seed, "batch_size": args.batch_size, "memory_mode": args.memory_mode, + "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, }, } @@ -450,6 +456,9 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["fp8", "nvfp4"], help = "quantise the companion text encoder (fp8 or nvfp4)", ) + p.add_argument( + "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload" + ) p.add_argument( "--write-baseline", metavar = "PATH", diff --git a/scripts/sd_cpp_smoke.py b/scripts/sd_cpp_smoke.py new file mode 100644 index 0000000000..4398fa6cf3 --- /dev/null +++ b/scripts/sd_cpp_smoke.py @@ -0,0 +1,133 @@ +# 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/diffusion.py b/studio/backend/core/inference/diffusion.py index 9cb0563f77..8ad5ffa250 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -16,7 +16,7 @@ from __future__ import annotations import inspect import threading import time -from dataclasses import dataclass, replace +from dataclasses import dataclass from pathlib import Path from typing import Any, Optional @@ -31,17 +31,20 @@ from .diffusion_families import ( ) from .diffusion_device import ( DiffusionDeviceTarget, + diffusion_device_target_from_torch_device, resolve_diffusion_device_target, ) from .diffusion_memory import ( OFFLOAD_NONE, apply_memory_plan, + estimate_gguf_dense_mib, estimate_image_runtime_mib, file_size_mib, + infer_gguf_quant_label, plan_diffusion_memory, snapshot_device_memory, ) -from .diffusion_speed import SPEED_OFF, apply_speed_optims, restore_tf32 +from .diffusion_speed import SPEED_OFF, apply_speed_optims from .diffusion_precision import quantize_text_encoders logger = get_logger(__name__) @@ -149,20 +152,28 @@ class DiffusionBackend: def is_loaded(self) -> bool: return self._state is not None - def _resolve_device_target(self, fam: Optional[DiffusionFamily]) -> DiffusionDeviceTarget: - """The device target for the current host with the family fp16 guard applied: - promote float16 -> float32 for fp16-incompatible families (Z-Image), whose - activations overflow float16 and render a black image. The capability flags - carry through unchanged; only the dtype moves.""" + def _pick_device_and_dtype(self) -> tuple[str, Any]: + """(device, dtype) for the current host. Thin wrapper over the device + policy module, kept as a method so tests can still monkeypatch it.""" target = resolve_diffusion_device_target() - effective = _resolve_diffusion_compute_dtype(fam, target.dtype) - if effective is target.dtype: - return target - logger.warning( - "diffusion.dtype_promoted: family=%s float16 -> float32 (fp16-incompatible)", - getattr(fam, "name", None), - ) - return replace(target, dtype = effective) + return target.device, target.dtype + + def _resolve_device_target(self, fam: Optional[DiffusionFamily]) -> DiffusionDeviceTarget: + """The device target with the family fp16 guard applied. + + Routes through _pick_device_and_dtype() (so a monkeypatched override still + drives the result), then promotes float16 -> float32 for fp16-incompatible + families (Z-Image), rebuilding the target so dtype + capability flags stay + consistent with the effective dtype. + """ + device, dtype = self._pick_device_and_dtype() + effective = _resolve_diffusion_compute_dtype(fam, dtype) + if effective is not dtype: + logger.warning( + "diffusion.dtype_promoted: family=%s float16 -> float32 (fp16-incompatible)", + getattr(fam, "name", None), + ) + return diffusion_device_target_from_torch_device(device, effective) def _resolve_gguf_path(self, repo_id: str, gguf_filename: str, hf_token: Optional[str]) -> str: local_root = Path(repo_id).expanduser() @@ -263,6 +274,7 @@ class DiffusionBackend: base_repo: Optional[str] = None, family_override: Optional[str] = None, hf_token: Optional[str] = None, + cpu_offload: bool = False, memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, @@ -293,6 +305,7 @@ class DiffusionBackend: base_repo = base_repo, family_override = family_override, hf_token = hf_token, + cpu_offload = cpu_offload, memory_mode = memory_mode, speed_mode = speed_mode, text_encoder_quant = text_encoder_quant, @@ -415,47 +428,6 @@ class DiffusionBackend: return 0 # repo not in cache yet return total - @staticmethod - def _cache_bytes_loaded(repo_id: str) -> int: - """Cached bytes of only the base-repo files the GGUF pipeline actually loads. - - Walks the snapshot dir (filename -> blob symlinks) and keeps the files - ``_base_file_downloaded`` accepts, so the companion-memory estimate counts just - the VAE + text encoders. The raw ``_cache_bytes`` blob sum would also include any - dense ``transformer/`` shards left in the cache by a previous diffusers run, which - the GGUF path never loads -- inflating companion memory and pushing the auto plan - toward needless offload/tiling. Returns 0 when the repo isn't cached or has no - snapshot layout; the planner then treats companion size as unknown.""" - from huggingface_hub import constants - - snaps = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "snapshots" - total = 0 - counted: set = set() - try: - snap_dirs = list(snaps.iterdir()) - except OSError: - return 0 # repo not in cache yet - if not snap_dirs: - return 0 - for snap in snap_dirs: - try: - for f in snap.rglob("*"): - if not f.is_file(): - continue - if not _base_file_downloaded(f.relative_to(snap).as_posix()): - continue - try: - blob = f.resolve() - if blob in counted: # same blob shared across snapshots - continue - counted.add(blob) - total += blob.stat().st_size - except OSError: - continue - except OSError: - continue - return total - # ── Synchronous load / generate / unload ─────────────────────────────── def load_pipeline( @@ -466,6 +438,7 @@ class DiffusionBackend: base_repo: Optional[str] = None, family_override: Optional[str] = None, hf_token: Optional[str] = None, + cpu_offload: bool = False, memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, @@ -488,18 +461,12 @@ class DiffusionBackend: # The cancel makes that wait ~one step (or the rest of the denoise for a # pipeline that ignores the step callback). with self._lock: - # Check the token BEFORE cancelling: a superseded/cancelled worker (its - # token already bumped by unload or a newer load) must not abort the - # current model's unrelated in-flight generation on its way out. - if _load_token is not None and _load_token != self._load_token: - raise RuntimeError("Diffusion load was cancelled.") if self._active_generate_cancel is not None: self._active_generate_cancel.set() with self._generate_lock: with self._lock: - # Re-check under the lock (we dropped it to wait on _generate_lock): an - # unload/eviction or a newer load may have superseded this one while we - # waited, and the slow VRAM-heavy build below must not run if so. + # Bail before the (slow, VRAM-heavy) build if an unload/eviction or a + # newer load superseded this one while we were resolving/downloading. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") @@ -533,6 +500,7 @@ class DiffusionBackend: speed_applied = apply_speed_optims( pipe, target, + is_gguf = bool(gguf_filename), family = fam, speed_mode = speed_mode or SPEED_OFF, logger = logger, @@ -550,8 +518,11 @@ class DiffusionBackend: # estimated resident size (transformer GGUF dequantised + the # companion text-encoder / VAE already cached for `base`), then # apply it. Computed here, after the build but before placement, - # because the weights are still on CPU so free VRAM is the real budget. - plan = self._plan_memory(target, gguf_path, gguf_filename, base, fam, memory_mode) + # because the weights are still on CPU so free VRAM is the real + # budget. `cpu_offload=True` stays an explicit override. + plan = self._plan_memory( + target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload + ) # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it # may fall back to whole-module offload, and tiling is a no-op on a # pipeline with no tiling control), so status stays honest. @@ -594,28 +565,24 @@ class DiffusionBackend: base: str, fam: DiffusionFamily, memory_mode: Optional[str], + cpu_offload: bool, ): """Build the memory plan for this load: snapshot free device memory and estimate the model's resident footprint, then let the planner pick an offload policy + VAE memory savers. Kept on the backend so the cached base repo (companion text-encoder / VAE) feeds the size estimate.""" device_memory = snapshot_device_memory(target) - # diffusers keeps the GGUF transformer PACKED on the device (uint8) and - # dequantises each layer transiently during the forward, so its resident - # footprint is ~the on-disk file size, NOT the dequantised dense size. (The - # transient per-layer dequant + activations are covered by the planner's - # runtime headroom + base overhead.) Measured: a 2463 MiB Q4 GGUF sits resident - # at ~2482 MiB. - transformer_resident = file_size_mib(gguf_path) - # The companion components (VAE + text encoders) load dense, near their on-disk - # size; sum only the base-repo files the GGUF pipeline actually loads (NOT any - # leftover dense transformer/ shards), so the plan isn't skewed toward offload. - companion = self._cache_bytes_loaded(base) + transformer_dense = estimate_gguf_dense_mib( + file_size_mib(gguf_path), infer_gguf_quant_label(gguf_filename) + ) + # The companion components (VAE + text encoders) load near their on-disk + # size; sum whatever the prefetch already placed in the base-repo cache. + companion = self._cache_bytes(base) companion_mib = int(companion // (1024 * 1024)) if companion else None model_dense_mib = None - if transformer_resident is not None: - model_dense_mib = transformer_resident + (companion_mib or 0) - runtime_headroom = estimate_image_runtime_mib(family = fam.name) + if transformer_dense is not None: + model_dense_mib = transformer_dense + (companion_mib or 0) + runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name) return plan_diffusion_memory( target = target, device_memory = device_memory, @@ -623,6 +590,7 @@ class DiffusionBackend: companion_dense_mib = companion_mib, runtime_headroom_mib = runtime_headroom, requested_mode = memory_mode, + explicit_offload = cpu_offload, ) def generate( @@ -751,32 +719,22 @@ class DiffusionBackend: # Abort an in-flight download so unload/an eviction returns promptly instead # of waiting it out (the download runs without _lock and checks this event). self._cancel_event.set() - # Signal an in-flight denoise, then WAIT on _generate_lock for it to exit - # before freeing _state. The GPU arbiter releases diffusion ownership and - # hands the card to chat the instant this returns, so chat must not start - # allocating VRAM while the old pipeline is still resident — that transient - # overlap is an OOM risk. The cancel bounds the wait to ~one step (the same - # bound the replacement-load path in load_pipeline already relies on). with self._lock: + # Abort an in-flight denoise too by setting ITS cancel event, so the step + # callback stops it. unload does NOT take _generate_lock — it must return + # promptly; the running generate keeps its own pipe reference, so freeing + # _state here can't crash it, and its VRAM is reclaimed when it returns + # (within ~one step thanks to the cancel). if self._active_generate_cancel is not None: self._active_generate_cancel.set() - with self._generate_lock: - with self._lock: - self._unload_locked() - # Cancel any in-flight load (its worker checks this token before - # committing) and drop the marker so the next load starts clean. - self._load_token += 1 - self._loading = None + self._unload_locked() + # Cancel any in-flight load (its worker checks this token before + # committing) and drop the marker so the next load starts clean. + self._load_token += 1 + self._loading = None return self.status() def _unload_locked(self) -> None: - # Restore TF32 before the no-state early return: a speed_mode="max" load flips - # process-global TF32 (in apply_speed_optims) and can then fail before _state is - # committed, leaving the flag on with _state still None. This is the universal - # handoff gate (explicit unload, an arbiter eviction for chat, and the start of - # the next load all route here), so it's the only place that reliably puts TF32 - # back before the next tenant runs. No-op when no max load ever flipped it. - restore_tf32() state = self._state if state is None: return diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index 8a49f1248d..0c874e4b66 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -30,6 +30,21 @@ class DiffusionDeviceTarget: supports_default_torch_compile: bool supports_pinned_transfer: bool + @property + def is_cuda_torch_device(self) -> bool: + return self.device == "cuda" + + def as_public_dict(self) -> dict[str, Any]: + return { + "device": self.device, + "dtype": str(self.dtype).replace("torch.", ""), + "backend": self.backend, + "vendor": self.vendor, + "supports_model_cpu_offload": self.supports_model_cpu_offload, + "supports_default_torch_compile": self.supports_default_torch_compile, + "supports_pinned_transfer": self.supports_pinned_transfer, + } + def _studio_device_is(studio_device: Any, device_type: Any, name: str) -> bool: """True if ``studio_device`` equals ``DeviceType.`` (when that member exists).""" @@ -82,6 +97,53 @@ def resolve_diffusion_device_target() -> DiffusionDeviceTarget: return _mps_or_cpu_target(torch) +def diffusion_device_target_from_torch_device( + torch_device: str, dtype: Any +) -> DiffusionDeviceTarget: + """Reconstruct a target from a (device, dtype) pair. + + Keeps the ``_pick_device_and_dtype`` shim / monkeypatch path working: a caller + that overrides the (device, dtype) tuple can still recover the capability flags. + """ + device = str(torch_device).split(":", 1)[0] + if device == "cuda": + try: + import torch + is_rocm = bool(getattr(getattr(torch, "version", None), "hip", None)) + except Exception: + is_rocm = False + return DiffusionDeviceTarget( + device = "cuda", + dtype = dtype, + backend = "rocm" if is_rocm else "cuda", + vendor = "amd" if is_rocm else "nvidia", + supports_model_cpu_offload = True, + supports_default_torch_compile = not is_rocm, + supports_pinned_transfer = True, + ) + if device == "xpu": + return DiffusionDeviceTarget( + device = "xpu", + dtype = dtype, + backend = "xpu", + vendor = "intel", + supports_model_cpu_offload = True, + supports_default_torch_compile = False, + supports_pinned_transfer = False, + ) + if device == "mps": + return DiffusionDeviceTarget( + device = "mps", + dtype = dtype, + backend = "mps", + vendor = "apple", + supports_model_cpu_offload = False, + supports_default_torch_compile = False, + supports_pinned_transfer = False, + ) + return _cpu_target(torch = None, dtype = dtype) + + def _cuda_or_rocm_target(torch: Any, *, is_rocm: bool) -> DiffusionDeviceTarget: if is_rocm: # ROCm (AMD) does not have NVIDIA's pre-Ampere bf16-emulation quirk, so @@ -175,10 +237,12 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget: return _cpu_target(torch) -def _cpu_target(torch: Any) -> DiffusionDeviceTarget: +def _cpu_target(torch: Any, dtype: Any = None) -> DiffusionDeviceTarget: + if dtype is None: + dtype = torch.float32 return DiffusionDeviceTarget( device = "cpu", - dtype = torch.float32, + dtype = dtype, backend = "cpu", vendor = None, supports_model_cpu_offload = False, diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 22c546ac95..4fecf950ad 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -36,8 +36,8 @@ class DiffusionFamily: # promotes a resolved float16 to float32 for these at load time. fp16_incompatible: bool = False # False for families whose denoiser block doesn't compile cleanly with - # regional torch.compile (Z-Image). The GGUF transformer compiles cleanly, so - # this gates the opt-in speed-mode compile for every family. + # regional torch.compile (Z-Image). Only consulted on the non-GGUF path; the + # GGUF transformer is never compiled regardless. supports_torch_compile: bool = True diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 19d1558f8e..ca90e11c73 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -53,6 +53,8 @@ OFFLOAD_SEQUENTIAL = "sequential" # = lower peak VRAM, more host<->device traffic. One is the lowest-VRAM setting. DEFAULT_GROUP_BLOCKS = 1 +DEFAULT_IMAGE_WIDTH = 1024 +DEFAULT_IMAGE_HEIGHT = 1024 # A flat allowance for the pipeline's fixed costs (scheduler, embeddings, the # CUDA context, fragmentation) on top of the model weights and per-step runtime. DEFAULT_BASE_OVERHEAD_MIB = 2048 @@ -61,10 +63,8 @@ DEFAULT_BASE_OVERHEAD_MIB = 2048 def normalize_memory_mode(value: Optional[str]) -> Optional[str]: """Lower/strip a requested mode, accepting dashes; None passes through. - Raises ValueError for an unsupported mode. The route already rejects bad values - at the Pydantic Literal boundary (422, before any GPU work); this is a defense-in- - depth guard for direct / script callers (bench, quality harness) that bypass it, - and runs on the load thread.""" + Raises ValueError for an unsupported mode so a bad request can be rejected + cheaply (the route surfaces it as a 4xx before any GPU work).""" if value is None: return None normalized = str(value).strip().lower().replace("-", "_") @@ -93,6 +93,15 @@ class DeviceMemory: def is_unified(self) -> bool: return self.memory_kind in ("unified_memory", "system_memory") + def as_public_dict(self) -> dict[str, Any]: + return { + "backend": self.backend, + "device": self.device, + "memory_kind": self.memory_kind, + "free_mib": self.free_mib, + "total_mib": self.total_mib, + } + @dataclass(frozen = True) class MemoryPlan: @@ -103,8 +112,24 @@ class MemoryPlan: vae_tiling: bool vae_slicing: bool device_memory: DeviceMemory + estimates: dict[str, Optional[int]] reasons: tuple[str, ...] = () + @property + def engages_offload(self) -> bool: + return self.offload_policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL) + + def as_public_dict(self) -> dict[str, Any]: + return { + "requested_mode": self.requested_mode, + "offload_policy": self.offload_policy, + "vae_tiling": self.vae_tiling, + "vae_slicing": self.vae_slicing, + "device_memory": self.device_memory.as_public_dict(), + "estimates": dict(self.estimates), + "reasons": list(self.reasons), + } + # ── hardware snapshot ───────────────────────────────────────────────────────── @@ -193,17 +218,73 @@ def file_size_mib(path: Any) -> Optional[int]: return None -def estimate_image_runtime_mib(*, family: Optional[str] = None) -> int: - """Per-call activation / latent headroom for an image generation (at the - default ~1MP resolution the planner budgets for). Distilled / turbo models - (few steps, no CFG) need less; editing pipelines need more.""" +def infer_gguf_quant_label(filename: Optional[str]) -> Optional[str]: + """Pull a quant tag (Q4_K_M, Q8_0, BF16, ...) out of a GGUF filename.""" + if not filename: + return None + from pathlib import Path + + stem = Path(filename).name + if stem.lower().endswith(".gguf"): + stem = stem[:-5] + parts = [p.upper() for p in stem.replace("-", "_").split("_") if p] + for index, part in enumerate(parts): + if part in ("BF16", "F16", "FP16", "FP8", "Q8", "Q6", "Q5", "Q4", "Q3", "Q2"): + suffix = parts[index + 1 :] + # Quant names carry either a K-family suffix (Q4_K_M) or a legacy + # numeric one (Q8_0, Q5_1); keep up to two suffix tokens. + if suffix and suffix[0] in ("K", "M", "S", "L", "XS", "XXS", "0", "1"): + return "_".join([part] + suffix[:2]) + return part + if part.startswith("IQ") or part.startswith("UD"): + return "_".join(parts[index : index + 3]) + return None + + +def estimate_gguf_dense_mib(storage_mib: Optional[int], quant: Optional[str]) -> Optional[int]: + """Approximate the dequantised (device) size of a GGUF from its on-disk size + and quant label. The compute dtype is bf16/fp16, so a 4-bit file roughly + quadruples once unpacked; higher-bit quants expand less.""" + if storage_mib is None: + return None + q = (quant or "").upper() + if any(t in q for t in ("BF16", "F16", "FP16")): + return storage_mib + if "FP8" in q or "Q8" in q: + return int(storage_mib * 2.0) + if "Q6" in q: + return int(storage_mib * 2.8) + if "Q5" in q: + return int(storage_mib * 3.3) + if "Q4" in q or "IQ4" in q or "UD" in q: + return int(storage_mib * 4.0) + if "Q3" in q or "IQ3" in q: + return int(storage_mib * 5.3) + if "Q2" in q or "Q1" in q or "IQ2" in q or "IQ1" in q: + return int(storage_mib * 8.0) + return int(storage_mib * 4.0) # unknown: assume 4-bit-ish + + +def estimate_image_runtime_mib( + *, + width: Optional[int], + height: Optional[int], + batch_size: int = 1, + family: Optional[str] = None, +) -> int: + """Per-call activation / latent headroom for an image generation, scaled by + pixel area and batch. Distilled / turbo models (few steps, no CFG) need less.""" + w = max(64, int(width or DEFAULT_IMAGE_WIDTH)) + h = max(64, int(height or DEFAULT_IMAGE_HEIGHT)) + batch = max(1, int(batch_size or 1)) + pixel_scale = (w * h * batch) / float(DEFAULT_IMAGE_WIDTH * DEFAULT_IMAGE_HEIGHT) fam = (family or "").lower() multiplier = 1.0 if "edit" in fam: multiplier *= 1.35 if "turbo" in fam or "distilled" in fam or "schnell" in fam: multiplier *= 0.85 - return max(1024, int(8192 * multiplier)) + return max(1024, int(8192 * max(0.25, pixel_scale) * multiplier)) def _safe_device_budget_mib(memory: DeviceMemory) -> Optional[int]: @@ -242,13 +323,15 @@ def plan_diffusion_memory( companion_dense_mib: Optional[int] = None, base_overhead_mib: int = DEFAULT_BASE_OVERHEAD_MIB, requested_mode: Optional[str] = None, + explicit_offload: bool = False, ) -> MemoryPlan: """Pick an offload policy + VAE memory savers for the current load. ``model_dense_mib`` is the estimated resident device size of all weights (transformer + companion text-encoder / VAE); ``companion_dense_mib`` is just the companions, which stay resident under streamed (group) offload while the - transformer is streamed block by block. + transformer is streamed block by block. ``explicit_offload`` is the back-compat + ``cpu_offload=True`` request: it forces whole-module offload. Policy meanings, ordered by measured speed/VRAM tradeoff: none - everything resident: fastest, highest VRAM. @@ -263,6 +346,15 @@ def plan_diffusion_memory( # The resident floor under group offload: companions stay, the transformer streams. group_floor = _sum_required(companion_dense_mib, runtime_headroom_mib, base_overhead_mib) reasons: list[str] = [] + estimates: dict[str, Optional[int]] = { + "safe_device_budget_mib": budget, + "model_dense_mib": model_dense_mib, + "companion_dense_mib": companion_dense_mib, + "runtime_headroom_mib": runtime_headroom_mib, + "base_overhead_mib": base_overhead_mib, + "resident_required_mib": required, + "group_floor_mib": group_floor, + } def _group_fits() -> bool: # Group offload only helps if the resident remainder (companions) fits; when @@ -304,25 +396,23 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("companions exceed budget; whole-module offload of every component") - # VAE TILING decodes in spatially-overlapping chunks and blends the seams, - # capping the decode-time VRAM spike at high resolution -- but it is LOSSY above - # the tile threshold (diffusers blends overlapping tiles), so only engage it when - # weights are being offloaded (the device is already tight) or there is no spare - # device pool to offload to -- unified / system memory, i.e. MPS, CPU, and - # integrated / unified-memory CUDA (which is_unified covers; a backend-string - # check would miss it). On a roomy discrete GPU it stays off so output is - # bit-identical. - tile = policy != OFFLOAD_NONE or device_memory.is_unified - # VAE SLICING decodes one batch element at a time and concatenates (no blending): - # bit-identical to a non-sliced decode, and diffusers no-ops it below batch>1, so - # it's free at the batch=1 size every request uses today. Always on -- it never - # changes output and a future batched caller gets the memory saving for free. + if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified: + policy = OFFLOAD_MODEL + reasons.append("explicit cpu_offload overrides resident placement") + + # VAE tiling/slicing decode the image in chunks, capping the decode-time spike + # that often dominates peak VRAM at high resolution. Turn it on whenever weights + # are being offloaded (the device is already tight) or the backend has no spare + # device pool (MPS/CPU). On a roomy discrete GPU it stays off so output is + # bit-identical to a plain resident run. + tile = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu") return MemoryPlan( requested_mode = mode, offload_policy = policy, vae_tiling = tile, - vae_slicing = True, + vae_slicing = tile, device_memory = device_memory, + estimates = estimates, reasons = tuple(reasons), ) diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 9b8dba28e3..030b6da7f2 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -37,8 +37,7 @@ _TEXT_ENCODER_ATTRS = ("text_encoder", "text_encoder_2", "text_encoder_3") def normalize_te_quant(value: Optional[str]) -> Optional[str]: """Lower/strip a requested text-encoder quant; None / "" / "none" -> None. - Raises ValueError for an unsupported value. The route rejects bad values at the - Pydantic Literal boundary; this guard covers direct / script callers.""" + Raises ValueError for an unsupported value so a bad request is rejected cheaply.""" if value is None: return None normalized = str(value).strip().lower().replace("-", "_") diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 60421ee2e4..e01ba48b63 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -10,15 +10,14 @@ lossless-to-near-lossless speedups in the order the diffusers guides recommend off - nothing (default). default - lossless: channels_last VAE memory format + regional torch.compile of - the denoiser's repeated block WHERE eligible (bf16, CUDA, and a - compile-friendly family). + the denoiser's repeated block WHERE eligible (non-GGUF, bf16, CUDA, and + a compile-friendly family). max - default plus near-lossless TF32 matmul and fused QKV projections. -Regional compile is gated only by family compatibility (Z-Image's block is flagged -off) and a bf16 / CUDA target. A GGUF transformer compiles cleanly on diffusers' -native dequant path -- verified empirically: fullgraph, zero graph breaks, ~1.35-1.44x -faster with lower peak VRAM (the per-op dequant is pure tensor ops inductor fuses) -- -so it is NOT excluded. torch is imported lazily. +Regional compile is gated off for the GGUF transformer (it dequantises per-op and +doesn't compile cleanly) and for families flagged not compile-friendly (Z-Image), so +on today's GGUF path only channels_last / TF32 engage; the compile path activates +automatically once a non-GGUF bf16 transformer is loaded. torch is imported lazily. """ from __future__ import annotations @@ -45,12 +44,14 @@ def normalize_speed_mode(value: Optional[str]) -> str: return normalized -def compile_eligible(target: Any, *, family: Any) -> bool: +def compile_eligible(target: Any, *, is_gguf: bool, family: Any) -> bool: """Whether the denoiser's repeated block should be regionally compiled. - Only on CUDA (incl. ROCm via supports_default_torch_compile), for a bf16 - transformer, on a compile-friendly family. A GGUF transformer IS eligible -- it - compiles cleanly on diffusers' native dequant path (verified ~1.4x faster).""" + Only on CUDA (incl. ROCm via supports_default_torch_compile), for a non-GGUF + bf16 transformer, on a compile-friendly family. The GGUF transformer is never + compiled (it dequantises per-op).""" + if is_gguf: + return False if not bool(getattr(target, "supports_default_torch_compile", False)): return False if not bool(getattr(family, "supports_torch_compile", True)): @@ -70,6 +71,7 @@ def apply_speed_optims( pipe: Any, target: Any, *, + is_gguf: bool, family: Any, speed_mode: str = SPEED_OFF, logger: Any = None, @@ -92,7 +94,7 @@ def apply_speed_optims( applied["channels_last"] = _vae_channels_last(pipe, logger) # Lossless-ish: regional compile of the repeated denoiser block, where eligible. - if compile_eligible(target, family = family): + if compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks(pipe, logger) if mode == SPEED_MAX: 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..80078ef6fc --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -0,0 +1,329 @@ +# 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 threading +import time +from pathlib import Path +from typing import Callable, Optional + +from core.inference.sd_cpp_args import ( + SdCppGenParams, + SdCppModelFiles, + build_sd_cpp_command, +) +from utils.process_lifetime import child_popen_kwargs + +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. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME exactly like the + # installer's default_install_dir(), so a binary installed under a custom Studio root + # is found without also having to set UNSLOTH_SD_CPP_PATH. + studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") + default_base = Path(studio_home).parent if studio_home else Path.home() / ".unsloth" + hit = _first_file(_layout_candidates(default_base / "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, + # Bind the child to the backend's lifetime (PR_SET_PDEATHSIG on Linux), so a + # long sd-cli denoise is SIGKILLed if the backend dies instead of orphaning. + **child_popen_kwargs(), + ) + # Drain stdout on a background thread and wait on the PROCESS, not the stream: + # iterating proc.stdout directly blocks until the stream closes, so a sd-cli that + # hangs without producing output (or closing stdout) would never reach proc.wait + # and the timeout would be silently bypassed. With the reader on its own thread the + # main thread always enforces the wall-clock timeout and kills a hung process. + tail: list[str] = [] + + def _drain() -> None: + 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) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + try: + ret = proc.wait(timeout = timeout) + except subprocess.TimeoutExpired: + proc.kill() + reader.join(timeout = 5.0) + raise RuntimeError(f"sd-cli timed out after {timeout}s") + finally: + if proc.poll() is None: + proc.kill() + # Reap the SIGKILLed child, or it lingers as a zombie until this process + # exits (the cancel/timeout branches kill without waiting otherwise). + try: + proc.wait(timeout = 5.0) + except Exception: # noqa: BLE001 + pass + # The process has exited; let the reader finish draining the buffered output. + reader.join(timeout = 5.0) + + if ret != 0: + raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) + if out.is_file(): + produced: Optional[Path] = out + else: + # For batch_count > 1, stable-diffusion.cpp's save_results() writes + # "_" (base_0.png, base_1.png, ...) rather than the + # literal --output path, so the single-path check above misses them. + # Fall back to the numbered siblings and return the first. + batch = sorted(out.parent.glob(f"{out.stem}_*{out.suffix}")) + produced = batch[0] if batch else None + if produced is None: + 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, produced) + return produced + + +# ── 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/models/inference.py b/studio/backend/models/inference.py index 49c9865f0d..fff68c090a 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1699,11 +1699,13 @@ class DiffusionLoadRequest(BaseModel): None, description = "Force a family when it can't be inferred from the repo id" ) hf_token: Optional[str] = Field(None, description = "HuggingFace token for gated repos") + cpu_offload: bool = Field(False, description = "Enable model CPU offload to fit low-VRAM cards") memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field( None, description = "Memory policy: auto (measured), fast (resident), balanced " "(stream the transformer, near-resident speed, moderate VRAM " - "cut), low_vram (offload every component, lowest VRAM, slower).", + "cut), low_vram (offload every component, lowest VRAM, slower). " + "Overrides cpu_offload when set.", ) speed_mode: Optional[Literal["off", "default", "max"]] = Field( None, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ed6819cdea..1aa4be56b3 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10327,6 +10327,7 @@ async def load_diffusion_model( base_repo = request.base_repo, family_override = request.family_override, hf_token = request.hf_token, + cpu_offload = request.cpu_offload, memory_mode = request.memory_mode, speed_mode = request.speed_mode, text_encoder_quant = request.text_encoder_quant, diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index b77becc404..c7c4736502 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -26,7 +26,6 @@ from core.inference.diffusion_families import ( resolve_base_repo, resolve_local_gguf_child, ) -from core.inference.diffusion_device import DiffusionDeviceTarget # Pure family helpers @@ -50,13 +49,9 @@ def test_detect_family_from_repo_id(): # Qwen-Image guides via true_cfg_scale, not guidance_scale. assert detect_family("unsloth/Qwen-Image-2512-GGUF").cfg_kwarg == "true_cfg_scale" assert detect_family("unsloth/Z-Image-GGUF").cfg_kwarg == "guidance_scale" - # Image-editing checkpoints are rejected (text-to-image backend only): the - # edit keyword is matched as a whole id segment, so an "edit" that's only a - # substring of a normal word ("Edition") still loads. + # Image-editing checkpoints are rejected (text-to-image backend only). assert detect_family("unsloth/Qwen-Image-Edit-2511-GGUF") is None assert detect_family("unsloth/FLUX.1-Kontext-dev-GGUF") is None - assert detect_family("unsloth/Qwen-Image-Inpainting-GGUF") is None - assert detect_family("unsloth/Z-Image-Edition-GGUF").name == "z-image" assert detect_family("meta-llama/Llama-3-8B") is None @@ -84,11 +79,6 @@ def test_resolve_local_gguf_child(tmp_path): resolve_local_gguf_child(tmp_path, "..\\secret.gguf") with pytest.raises(FileNotFoundError): resolve_local_gguf_child(tmp_path, "missing.gguf") - # A directory named like the gguf exists but isn't loadable; reject it here so - # the preflight doesn't pass and evict chat for a pick from_single_file can't load. - (tmp_path / "dir.gguf").mkdir() - with pytest.raises(FileNotFoundError): - resolve_local_gguf_child(tmp_path, "dir.gguf") def test_resolve_local_gguf_child_blocks_symlink_escape(tmp_path): @@ -226,12 +216,6 @@ def fake_runtime(monkeypatch): # The backend imports clear_gpu_cache by reference; no-op it so unload doesn't # run real hardware detection against the stubbed torch. monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: None) - # Fake the hardware layer too: resolve_diffusion_device_target() consults - # utils.hardware.get_device(), so on a real XPU/ROCm box the host would leak - # through the stubbed torch and the device tests would resolve a non-CUDA - # target. None -> fall through to the (stubbed, monkeypatchable) torch probe. - monkeypatch.setattr("utils.hardware.get_device", lambda: None, raising = False) - monkeypatch.setattr("utils.hardware.hardware.IS_ROCM", False, raising = False) _FakePipeline.last = {} _FakeTransformer.last = {} yield @@ -286,6 +270,20 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert backend.is_loaded is False +def test_cpu_offload_ignored_off_cuda(fake_runtime, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + family_override = "z-image", + base_repo = "base/repo", + cpu_offload = True, + ) + # No CUDA in the stub, so offload is not engaged. + assert status["cpu_offload"] is False + + def test_low_vram_ignored_off_cuda(fake_runtime, tmp_path): (tmp_path / "model.gguf").write_bytes(b"x") backend = DiffusionBackend() @@ -338,9 +336,7 @@ def test_load_without_gguf_raises(): def test_load_unknown_family_raises(): backend = DiffusionBackend() - # load_pipeline validates via validate_load_request, which rejects an - # undetectable family before any GPU/network work. - with pytest.raises(ValueError, match = "family"): + with pytest.raises(ValueError): backend.load_pipeline("some/unrecognised-repo", gguf_filename = "x.gguf") @@ -474,11 +470,9 @@ def test_pick_dtype_bf16_only_on_ampere(fake_runtime, monkeypatch): backend = DiffusionBackend() monkeypatch.setattr(torch.cuda, "is_available", lambda: True, raising = False) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0), raising = False) - t = backend._resolve_device_target(None) - assert (t.device, t.dtype) == ("cuda", torch.bfloat16) + assert backend._pick_device_and_dtype() == ("cuda", torch.bfloat16) monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (7, 5), raising = False) - t = backend._resolve_device_target(None) - assert (t.device, t.dtype) == ("cuda", torch.float16) + assert backend._pick_device_and_dtype() == ("cuda", torch.float16) def test_unload_sets_cancel_event(fake_runtime): @@ -582,7 +576,7 @@ def test_load_promotes_fp16_to_fp32_for_zimage_only(fake_runtime, monkeypatch, t # Lock split + mid-denoise cancellation -def test_generate_lock_split_keeps_status_responsive_and_unload_waits(fake_runtime): +def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime): import threading backend = DiffusionBackend() @@ -616,28 +610,21 @@ def test_generate_lock_split_keeps_status_responsive_and_unload_waits(fake_runti t = threading.Thread(target = _run) t.start() - assert started.wait(5) # the denoise is in flight, holding _generate_lock + assert started.wait(5) # the denoise is in flight, holding only _generate_lock - # status() / generate_progress() read _state lock-free, so they must NOT block - # behind the denoise. + # status() / generate_progress() must NOT block behind the denoise. assert backend.status()["loaded"] is True assert backend.generate_progress()["active"] is True - # unload() signals THIS generation's cancel, then WAITS on _generate_lock for it - # to exit before freeing _state -- so when the GPU arbiter hands the card to chat - # the moment unload returns, the old pipeline is already gone (no dual-allocation - # OOM). Run it on a thread to observe that it blocks behind the live denoise. - unload_done = threading.Event() - threading.Thread(target = lambda: (backend.unload(), unload_done.set()), daemon = True).start() - assert not unload_done.wait(0.5) # blocked behind the live denoise - assert backend._state is not None # not freed while the pipeline is still live + # unload() must return promptly (it does not wait on _generate_lock) and signal + # THIS in-flight generation's cancel event. + backend.unload() assert backend._active_generate_cancel is not None - assert backend._active_generate_cancel.is_set() # cancel was signalled up front - - release.set() # denoise returns and releases _generate_lock - t.join(5) - assert unload_done.wait(5) # only now does unload free _state and return + assert backend._active_generate_cancel.is_set() assert backend.status()["loaded"] is False + + release.set() + t.join(5) # The cancelled generation raised rather than returning a now-evicted image. assert "exc" in out and "cancelled" in str(out["exc"]).lower() @@ -737,18 +724,6 @@ def test_validate_load_request(tmp_path): gguf_filename = "m.gguf", family_override = "z-image", ) - # Windows-shaped local paths (backslash separator / .\ ..\ prefixes) are also caught - # here, so a mistyped local pick on Windows can't be mistaken for a remote HF repo. - for missing in (r".\models\z-image", r"..\z-image", r"models\z-image", r"C:\models\z"): - with pytest.raises(FileNotFoundError): - backend.validate_load_request( - missing, gguf_filename = "m.gguf", family_override = "z-image" - ) - # A bare "org/name" HF id is still treated as remote (not rejected as a local path). - assert ( - backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name - == "z-image" - ) def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path): @@ -832,21 +807,9 @@ def test_load_reports_memory_plan_fields_on_cpu(fake_runtime, tmp_path): def _force_cuda_target(backend, monkeypatch): - """Drive the loader down the CUDA (offload-capable) path under the stub by - overriding the device resolver with a fixed CUDA target.""" + """Drive the loader down the CUDA (offload-capable) path under the stub.""" torch = sys.modules["torch"] - cuda_target = DiffusionDeviceTarget( - device = "cuda", - dtype = torch.bfloat16, - backend = "cuda", - vendor = "nvidia", - supports_model_cpu_offload = True, - supports_default_torch_compile = True, - supports_pinned_transfer = True, - ) - monkeypatch.setattr( - "core.inference.diffusion.resolve_diffusion_device_target", lambda: cuda_target - ) + monkeypatch.setattr(backend, "_pick_device_and_dtype", lambda: ("cuda", torch.bfloat16)) def test_load_memory_mode_balanced_streams_or_falls_back(fake_runtime, tmp_path, monkeypatch): @@ -879,6 +842,20 @@ def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, assert pipe.offloaded is True and pipe.moved_to is None # offload owns placement +def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( + fake_runtime, tmp_path, monkeypatch +): + # cpu_offload=True with no mode: auto would stay resident (budget unknown under + # the stub), but the explicit flag forces whole-module offload. + (tmp_path / "m.gguf").write_bytes(b"x") + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + status = backend.load_pipeline( + str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", cpu_offload = True + ) + assert status["offload_policy"] == "model" and status["cpu_offload"] is True + + def test_load_speed_mode_threads_and_defaults_off(fake_runtime, tmp_path): # No speed_mode -> off, no optimisations engaged (the bit-identical default). (tmp_path / "m.gguf").write_bytes(b"x") @@ -903,37 +880,6 @@ def test_load_speed_mode_threads_and_defaults_off(fake_runtime, tmp_path): assert status3["text_encoder_quant"] is None -def test_failed_max_load_restores_tf32_on_unload(monkeypatch): - # A speed_mode="max" load flips process-global TF32 (in apply_speed_optims) and can - # then fail before _state is committed (e.g. an OOM during placement), leaving the - # flag on with _state still None. _unload_locked is the universal handoff gate, so it - # must restore TF32 even on this no-state path or the next tenant (chat, or a default - # load) silently inherits it. Regression: the restore used to sit below the - # `if state is None: return` early return and never ran here. - from core.inference import diffusion_speed - - torch = types.ModuleType("torch") - torch.backends = types.SimpleNamespace( - cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)), - cudnn = types.SimpleNamespace(allow_tf32 = False), - ) - monkeypatch.setitem(sys.modules, "torch", torch) - monkeypatch.setattr(diffusion_speed, "_tf32_prev", None) - - # Simulate the max load flipping TF32 on, then failing before committing _state. - diffusion_speed._enable_tf32(None) - assert torch.backends.cuda.matmul.allow_tf32 is True - assert diffusion_speed._tf32_prev == (False, False) - - backend = DiffusionBackend() - assert backend._state is None - backend._unload_locked() - - assert torch.backends.cuda.matmul.allow_tf32 is False - assert torch.backends.cudnn.allow_tf32 is False - assert diffusion_speed._tf32_prev is None - - def test_load_fast_mode_stays_resident_on_cuda(fake_runtime, tmp_path, monkeypatch): (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() @@ -943,102 +889,3 @@ def test_load_fast_mode_stays_resident_on_cuda(fake_runtime, tmp_path, monkeypat ) assert status["offload_policy"] == "none" and status["cpu_offload"] is False assert backend._state.pipe.moved_to == "cuda" - - -def test_superseded_load_does_not_cancel_unrelated_generation(fake_runtime, tmp_path): - # A stale worker (its _load_token already bumped by unload/a newer load) must bail on - # the token check BEFORE touching the current model's in-flight generation -- otherwise - # it would abort an unrelated denoise on its way out. - import threading - - backend = DiffusionBackend() - started = threading.Event() - release = threading.Event() - - class _BlockingPipe: - def __call__(self, **kwargs): - started.set() - release.wait(5) - return types.SimpleNamespace(images = [_FakeImage()]) - - fam = detect_family("unsloth/Z-Image-GGUF") - backend._state = _LoadState( - pipe = _BlockingPipe(), - family = fam, - repo_id = "current", - base_repo = "b", - device = "cpu", - dtype = "float32", - cpu_offload = False, - ) - backend._load_token = 7 # the "current" token - - gen_out: dict = {} - - def _gen(): - try: - backend.generate(prompt = "p", steps = 4) - except Exception as exc: # noqa: BLE001 - gen_out["exc"] = exc - - gt = threading.Thread(target = _gen) - gt.start() - assert started.wait(5) # generation in flight - cancel_event = backend._active_generate_cancel - assert cancel_event is not None - - # A superseded load arrives with a STALE token -> must raise without cancelling. - (tmp_path / "m.gguf").write_bytes(b"x") - with pytest.raises(RuntimeError, match = "cancelled"): - backend.load_pipeline( - str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", _load_token = 1 - ) - # The in-flight generation's cancel must NOT have been signalled by the stale worker. - assert not cancel_event.is_set() - - release.set() - gt.join(5) - # The generation completed normally (not cancelled). - assert "exc" not in gen_out - # The stale load did not replace the current model. - assert backend.status()["repo_id"] == "current" - - -def test_detect_family_from_local_gguf_filename(tmp_path): - # A direct local .gguf pick splits into (parent dir, basename); the family - # keyword can live only in the filename, so validate must scan it too. - (tmp_path / "z-image-turbo-Q4_K_M.gguf").write_bytes(b"w") - backend = DiffusionBackend() - fam = backend.validate_load_request(str(tmp_path), gguf_filename = "z-image-turbo-Q4_K_M.gguf") - assert fam.name == "z-image" - # An edit-checkpoint filename is still rejected even via the filename path. - (tmp_path / "FLUX.1-Kontext-dev-Q4.gguf").write_bytes(b"w") - with pytest.raises(ValueError): - backend.validate_load_request(str(tmp_path), gguf_filename = "FLUX.1-Kontext-dev-Q4.gguf") - # A bare parent dir whose name DOES carry the keyword still works (unchanged). - assert backend._detect_family_for_pick("unsloth/Z-Image-GGUF", "x.gguf", None).name == "z-image" - - -def test_run_load_does_not_stamp_superseded_progress(fake_runtime, monkeypatch): - # A worker whose load is superseded mid-resolve must not stamp its progress - # (base_repo / expected_bytes) onto the new load's _LoadingState. - backend = DiffusionBackend() - backend._loading = _LoadingState(repo_id = "unsloth/Z-Image-Turbo-GGUF", base_repo = "seed") - backend._load_token = 5 - monkeypatch.setattr("core.inference.diffusion._hf_base_model", lambda *a, **k: None) - - def supersede_then_estimate(*a, **k): - backend._load_token = 6 # a newer begin_load bumped the token mid-resolve - return (99999, []) - - monkeypatch.setattr( - DiffusionBackend, "_estimate_download_bytes", staticmethod(supersede_then_estimate) - ) - monkeypatch.setattr(DiffusionBackend, "_prefetch_files", lambda self, *a, **k: None) - monkeypatch.setattr(DiffusionBackend, "load_pipeline", lambda self, **k: None) - - backend._run_load( - repo_id = "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "m.gguf", base_repo = None, _load_token = 5 - ) - assert backend._loading.expected_bytes == 0 - assert backend._loading.base_repo == "seed" diff --git a/studio/backend/tests/test_diffusion_device.py b/studio/backend/tests/test_diffusion_device.py index 2179b16735..2eb9b3a822 100644 --- a/studio/backend/tests/test_diffusion_device.py +++ b/studio/backend/tests/test_diffusion_device.py @@ -13,6 +13,8 @@ from __future__ import annotations import sys import types +import pytest + from core.inference import diffusion_device as dd @@ -283,3 +285,40 @@ def test_fallback_cpu(monkeypatch): _install(monkeypatch, torch, hardware_fails = True) t = dd.resolve_diffusion_device_target() assert t.device == "cpu" and t.dtype == FP32 + + +# ── from-torch-device reconstruction + public dict ──────────────────── + + +def test_from_torch_device_cuda(monkeypatch): + torch = _make_torch() + monkeypatch.setitem(sys.modules, "torch", torch) + t = dd.diffusion_device_target_from_torch_device("cuda:0", FP32) + assert (t.device, t.backend, t.vendor, t.dtype) == ("cuda", "cuda", "nvidia", FP32) + assert t.is_cuda_torch_device + + +def test_from_torch_device_mps_and_cpu(monkeypatch): + torch = _make_torch() + monkeypatch.setitem(sys.modules, "torch", torch) + mps = dd.diffusion_device_target_from_torch_device("mps", FP16) + assert mps.device == "mps" and not mps.supports_model_cpu_offload + cpu = dd.diffusion_device_target_from_torch_device("cpu", FP32) + assert cpu.device == "cpu" and cpu.vendor is None + + +@pytest.mark.parametrize( + "dtype,expected", [(BF16, "bfloat16"), (FP16, "float16"), (FP32, "float32")] +) +def test_public_dict_dtype_string(dtype, expected): + t = dd.DiffusionDeviceTarget( + device = "cuda", + dtype = dtype, + backend = "cuda", + vendor = "nvidia", + supports_model_cpu_offload = True, + supports_default_torch_compile = True, + supports_pinned_transfer = True, + ) + d = t.as_public_dict() + assert d["dtype"] == expected and "torch." not in d["dtype"] diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 6ac1abb755..858e87ffa1 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -15,6 +15,8 @@ import types import pytest from core.inference.diffusion_memory import ( + DEFAULT_IMAGE_HEIGHT, + DEFAULT_IMAGE_WIDTH, MEMORY_MODE_BALANCED, MEMORY_MODE_FAST, MEMORY_MODE_LOW_VRAM, @@ -25,7 +27,9 @@ from core.inference.diffusion_memory import ( DeviceMemory, MemoryPlan, apply_memory_plan, + estimate_gguf_dense_mib, estimate_image_runtime_mib, + infer_gguf_quant_label, normalize_memory_mode, plan_diffusion_memory, snapshot_device_memory, @@ -63,16 +67,43 @@ def test_normalize_memory_mode_accepts_and_rejects(): normalize_memory_mode("ultra") -# ── size estimates ──────────────────────────────────────────────────────────── +# ── filename / size estimates ───────────────────────────────────────────────── -def test_estimate_image_runtime_scales_by_family(): - base = estimate_image_runtime_mib(family = "z-image") - # Distilled / turbo families get a discount; editing pipelines need more. - turbo = estimate_image_runtime_mib(family = "z-image-turbo") - edit = estimate_image_runtime_mib(family = "flux-kontext-edit") - assert turbo < base < edit - assert turbo >= 1024 # never below the floor +@pytest.mark.parametrize( + "filename,expected", + [ + ("z-image-turbo-Q4_K_M.gguf", "Q4_K_M"), + ("flux1-dev-Q8_0.gguf", "Q8_0"), + ("model-BF16.gguf", "BF16"), + ("qwen-image-IQ4_XS.gguf", "IQ4_XS"), + ("no-quant-here.gguf", None), + (None, None), + ], +) +def test_infer_gguf_quant_label(filename, expected): + assert infer_gguf_quant_label(filename) == expected + + +def test_estimate_gguf_dense_mib_expansion(): + # 4-bit roughly quadruples once dequantised to bf16; F16 is already dense. + assert estimate_gguf_dense_mib(1000, "Q4_K_M") == 4000 + assert estimate_gguf_dense_mib(1000, "Q8_0") == 2000 + assert estimate_gguf_dense_mib(1000, "BF16") == 1000 + assert estimate_gguf_dense_mib(None, "Q4_K_M") is None + # Unknown quant falls back to the conservative 4-bit-ish factor. + assert estimate_gguf_dense_mib(1000, None) == 4000 + + +def test_estimate_image_runtime_scales_with_pixels_and_family(): + base = estimate_image_runtime_mib(width = DEFAULT_IMAGE_WIDTH, height = DEFAULT_IMAGE_HEIGHT) + bigger = estimate_image_runtime_mib(width = 2048, height = 2048) + assert bigger > base + # Distilled / turbo families get a discount. + turbo = estimate_image_runtime_mib( + width = DEFAULT_IMAGE_WIDTH, height = DEFAULT_IMAGE_HEIGHT, family = "z-image-turbo" + ) + assert turbo < base # ── planner: device classes ─────────────────────────────────────────────────── @@ -110,9 +141,6 @@ def test_unified_cuda_skips_offload_even_if_offload_capable(): runtime_headroom_mib = 4000, ) assert plan.offload_policy == OFFLOAD_NONE - # ... but it's still memory-tight, so VAE tiling must engage (keyed on is_unified, - # which a backend-string check would have missed for unified-memory CUDA). - assert plan.vae_tiling is True # ── planner: auto budget tiers on a discrete GPU ────────────────────────────── @@ -127,8 +155,7 @@ def test_auto_resident_when_roomy(): runtime_headroom_mib = 4000, ) assert plan.offload_policy == OFFLOAD_NONE - assert plan.vae_tiling is False # roomy -> no lossy tiling - assert plan.vae_slicing is True # slicing is always on (lossless, free at batch=1) + assert plan.vae_tiling is False and plan.vae_slicing is False # roomy -> no tiling def test_auto_model_offload_on_tight_fit(): @@ -194,7 +221,7 @@ def test_auto_stays_resident_when_budget_unknown(): assert any("unknown" in r for r in plan.reasons) -# ── planner: explicit modes ─────────────────────────────────────────────────── +# ── planner: explicit modes + cpu_offload override ──────────────────────────── def test_explicit_modes_force_policy_regardless_of_budget(): @@ -242,6 +269,30 @@ def test_fast_falls_back_to_model_offload_when_it_does_not_fit(): assert plan.offload_policy == OFFLOAD_MODEL +def test_explicit_cpu_offload_overrides_resident_auto_choice(): + # Roomy GPU -> auto would stay resident, but cpu_offload=True forces offload. + plan = plan_diffusion_memory( + target = _target(), + device_memory = _discrete(80000), + model_dense_mib = 4000, + runtime_headroom_mib = 2000, + explicit_offload = True, + ) + assert plan.offload_policy == OFFLOAD_MODEL + assert any("explicit cpu_offload" in r for r in plan.reasons) + + +def test_explicit_cpu_offload_ignored_on_cpu_target(): + plan = plan_diffusion_memory( + target = _target(device = "cpu", backend = "cpu", supports_offload = False), + device_memory = DeviceMemory("cpu", "cpu", "system_memory", 8000, 16000), + model_dense_mib = 4000, + runtime_headroom_mib = 2000, + explicit_offload = True, + ) + assert plan.offload_policy == OFFLOAD_NONE + + # ── snapshot ────────────────────────────────────────────────────────────────── @@ -328,14 +379,14 @@ def _manual_plan(policy, *, tiling): vae_tiling = tiling, vae_slicing = tiling, device_memory = _discrete(4000, 8000), + estimates = {}, ) def test_apply_none_places_resident(): pipe = _RecordingPipe() effective, tiled = apply_memory_plan(pipe, _plan(OFFLOAD_NONE, tiling = False), device = "cuda") - # Resident placement, no lossy tiling, but slicing is always on (lossless / free). - assert pipe.calls == ["vae_slicing", "to:cuda"] + assert pipe.calls == ["to:cuda"] # no tiling on a roomy resident run assert effective == OFFLOAD_NONE and tiled is False diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index f879ad99fc..072a43e5f8 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -43,9 +43,10 @@ def _stub_torch( torch.float16 = "float16" if with_fp8: torch.float8_e4m3fn = "float8_e4m3fn" - torch.cuda = types.SimpleNamespace(get_device_capability = lambda *a: cc) - # _cast_fp8 skips nn.Embedding modules from layerwise casting. + # _cast_fp8 skips nn.Embedding tables (skip_modules_classes) to keep prompt + # tokens full precision, so the stub torch must expose torch.nn.Embedding. torch.nn = types.SimpleNamespace(Embedding = type("Embedding", (), {})) + torch.cuda = types.SimpleNamespace(get_device_capability = lambda *a: cc) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -67,20 +68,6 @@ def _stub_casters(monkeypatch, recorder): monkeypatch.setitem(sys.modules, "torchao.prototype.mx_formats", mx) -def _stub_fp8_capture(monkeypatch): - # Stub the fp8 caster to record the skip_modules_pattern passed for each encoder. - captured: dict = {} - hooks = types.ModuleType("diffusers.hooks") - casting = types.ModuleType("diffusers.hooks.layerwise_casting") - casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",) - hooks.apply_layerwise_casting = lambda module, **kw: captured.__setitem__( - id(module), kw["skip_modules_pattern"] - ) - monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks) - monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting) - return captured - - # ── normalisation ───────────────────────────────────────────────────────────── @@ -133,56 +120,6 @@ def test_quantize_fp8_casts_all_encoders(monkeypatch): assert recorder == [("fp8", te1), ("fp8", te3)] -def test_fp8_skips_encoder_keep_in_fp32_modules(monkeypatch): - # T5 keeps "wo" in fp32: its gated FF reads wo.weight.dtype and casts activations to - # match BEFORE calling wo, which races diffusers' forward-time upcast hook and crashes - # generation (fp8 input vs bf16 weight). _cast_fp8 must add the encoder's own - # _keep_in_fp32_modules to the layerwise-casting skip patterns; encoders without such a - # list (CLIP, Qwen) get nothing extra skipped. - _stub_torch(monkeypatch) - captured = _stub_fp8_capture(monkeypatch) - - t5 = types.SimpleNamespace(_keep_in_fp32_modules = ["wo"]) - qwen = types.SimpleNamespace(_keep_in_fp32_modules = None) - pipe = types.SimpleNamespace(text_encoder = t5, text_encoder_2 = qwen) - quantize_text_encoders(pipe, _target(), mode = "fp8") - - assert "wo" in captured[id(t5)] and "norm" in captured[id(t5)] - assert "wo" not in captured[id(qwen)] and "norm" in captured[id(qwen)] - - -def test_fp8_skips_tied_output_embedding(monkeypatch): - # A CausalLM encoder (FLUX.2's Qwen3) ties lm_head.weight to the input embedding. - # lm_head is nn.Linear, so layerwise casting would fp8 it and drag the shared - # embedding tensor to fp8, making the embedding emit fp8 activations that crash the - # first norm. _cast_fp8 must skip the tied output projection by name; an untied one - # (distinct weight tensors) is left to quantise normally. - _stub_torch(monkeypatch) - captured = _stub_fp8_capture(monkeypatch) - - shared = object() # the one tensor lm_head and embed_tokens share - emb = types.SimpleNamespace(weight = shared) - head = types.SimpleNamespace(weight = shared) - tied = types.SimpleNamespace( - _keep_in_fp32_modules = None, - get_input_embeddings = lambda: emb, - get_output_embeddings = lambda: head, - named_modules = lambda: [("model.embed_tokens", emb), ("lm_head", head)], - ) - u_emb, u_head = types.SimpleNamespace(weight = object()), types.SimpleNamespace(weight = object()) - untied = types.SimpleNamespace( - _keep_in_fp32_modules = None, - get_input_embeddings = lambda: u_emb, - get_output_embeddings = lambda: u_head, - named_modules = lambda: [("lm_head", u_head)], - ) - pipe = types.SimpleNamespace(text_encoder = tied, text_encoder_2 = untied) - quantize_text_encoders(pipe, _target(), mode = "fp8") - - assert r"^lm_head$" in captured[id(tied)] - assert r"^lm_head$" not in captured[id(untied)] - - def test_quantize_nvfp4_uses_torchao(monkeypatch): _stub_torch(monkeypatch, cc = (10, 0)) recorder: list = [] diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 9af23d330f..ade37482ba 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -67,16 +67,18 @@ def test_normalize_speed_mode(): # ── compile gating ──────────────────────────────────────────────────────────── -def test_compile_eligible_requires_bf16_cuda_friendly(monkeypatch): +def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch): _stub_torch(monkeypatch) - # The happy path: bf16, CUDA, compile-friendly family (GGUF is eligible too). - assert compile_eligible(_target(), family = _family()) is True + # The happy path: non-GGUF, bf16, CUDA, compile-friendly family. + assert compile_eligible(_target(), is_gguf = False, family = _family()) is True + # GGUF is never compiled. + assert compile_eligible(_target(), is_gguf = True, family = _family()) is False # fp16 (non-bf16) is excluded. - assert compile_eligible(_target(dtype = "float16"), family = _family()) is False + assert compile_eligible(_target(dtype = "float16"), is_gguf = False, family = _family()) is False # A family flagged not compile-friendly (Z-Image) is excluded. - assert compile_eligible(_target(), family = _family(compile_ok = False)) is False + assert compile_eligible(_target(), is_gguf = False, family = _family(compile_ok = False)) is False # No compile support (e.g. ROCm/XPU/MPS) is excluded. - assert compile_eligible(_target(compile_ok = False), family = _family()) is False + assert compile_eligible(_target(compile_ok = False), is_gguf = False, family = _family()) is False # ── applier ─────────────────────────────────────────────────────────────────── @@ -111,7 +113,9 @@ class _Pipe: def test_speed_off_applies_nothing(monkeypatch): _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = True) - applied = apply_speed_optims(pipe, _target(), family = _family(), speed_mode = SPEED_OFF) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF + ) assert applied == {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False} assert pipe.vae.mem_format is None and pipe.compiled is False @@ -119,17 +123,31 @@ def test_speed_off_applies_nothing(monkeypatch): def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) - applied = apply_speed_optims(pipe, _target(), family = _family(), speed_mode = SPEED_DEFAULT) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_DEFAULT + ) assert applied["channels_last"] is True and pipe.vae.mem_format == torch.channels_last assert applied["compiled"] is True and pipe.compiled is True # default does not flip TF32 or fuse QKV. assert applied["tf32"] is False and applied["fused_qkv"] is False +def test_speed_default_skips_compile_for_gguf(monkeypatch): + _stub_torch(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT + ) + assert applied["channels_last"] is True # lossless layout still applies + assert applied["compiled"] is False and pipe.compiled is False # GGUF never compiles + + def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = True) - applied = apply_speed_optims(pipe, _target(), family = _family(), speed_mode = SPEED_MAX) + applied = apply_speed_optims( + pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX + ) assert applied["tf32"] is True and torch.backends.cuda.matmul.allow_tf32 is True assert applied["fused_qkv"] is True and pipe.fused is True @@ -140,6 +158,7 @@ def test_speed_max_tf32_only_on_cuda(monkeypatch): applied = apply_speed_optims( pipe, _target(device = "mps", compile_ok = False), + is_gguf = True, family = _family(), speed_mode = SPEED_MAX, ) @@ -150,5 +169,7 @@ def test_apply_tolerates_missing_optims(monkeypatch): _stub_torch(monkeypatch) # A bare pipe (no vae.to, no compile, no fuse) must not crash. bare = types.SimpleNamespace(vae = None, transformer = types.SimpleNamespace()) - applied = apply_speed_optims(bare, _target(), family = _family(), speed_mode = SPEED_MAX) + applied = apply_speed_optims( + bare, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_MAX + ) assert applied["channels_last"] is False and applied["fused_qkv"] is False 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..99e758e22e --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -0,0 +1,189 @@ +# 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..ca6e875ea6 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -0,0 +1,354 @@ +# 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(monkeypatch): + # Hermetic: force discovery to find nothing so a real sd-cli installed on the host + # (e.g. ~/.unsloth) can't leak in and make binary=None resolve to a real binary. + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) + 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_collects_batch_output_paths(tmp_path, monkeypatch): + # batch_count > 1: stable-diffusion.cpp writes "_" + # (img_0.png, img_1.png, ...) rather than the literal --output path, so the + # single-path check must fall back to the numbered siblings. + e = _engine(tmp_path) + out = tmp_path / "img.png" + + def _factory(cmd, **kw): + # Emulate batch save_results(): write the numbered files, NOT the literal path. + (tmp_path / "img_0.png").write_bytes(b"\x89PNG\r\n") + (tmp_path / "img_1.png").write_bytes(b"\x89PNG\r\n") + return _FakePopen( + cmd, lines = ["done"], returncode = 0, out_file = out, write = False, env = kw.get("env") + ) + + monkeypatch.setattr(eng.subprocess, "Popen", _factory) + result = e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x", batch_count = 2), + output_path = str(out), + ) + assert result == tmp_path / "img_0.png" and result.is_file() + assert not out.exists() # the literal --output path was never written + + +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", + ) + + +def test_generate_times_out_even_when_stdout_blocks(tmp_path, monkeypatch): + # A sd-cli that hangs WITHOUT closing stdout must still hit the wall-clock timeout: + # the reader drains on a thread while the main thread waits on the PROCESS, so the + # timeout can no longer be bypassed by an unending stdout stream. + import subprocess as _sp + import threading as _threading + + released = _threading.Event() + + class _Block: + def __iter__(self): + return self + + def __next__(self): + # Models a hung stream that only ends once the process is killed. + if not released.wait(5.0): + raise AssertionError("stdout was never released by kill()") + raise StopIteration + + class _HangingPopen: + def __init__(self, cmd, **kw): + self.killed = False + + @property + def stdout(self): + return _Block() + + def wait(self, timeout = None): + raise _sp.TimeoutExpired(cmd = "sd-cli", timeout = timeout) + + def poll(self): + return 0 if self.killed else None + + def kill(self): + self.killed = True + released.set() # killing closes the pipe, so the reader unblocks + + holder: dict = {} + + def _factory(cmd, **kw): + holder["proc"] = _HangingPopen(cmd, **kw) + return holder["proc"] + + monkeypatch.setattr(eng.subprocess, "Popen", _factory) + + e = _engine(tmp_path) + with pytest.raises(RuntimeError, match = "timed out"): + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(tmp_path / "o.png"), + timeout = 0.01, + ) + assert holder["proc"].killed is True + + +# ── 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..b009cd37a6 --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_install.py @@ -0,0 +1,166 @@ +# 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)) + +import zipfile # noqa: E402 + +import pytest # noqa: E402 + +from install_sd_cpp_prebuilt import ( # noqa: E402 + _safe_extractall, + default_install_dir, + resolve_release_asset, +) + +# 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" + + +# ── safe extraction (Zip-Slip guard) ───────────────────────────────────────── + + +def test_safe_extractall_rejects_path_traversal(tmp_path): + target = tmp_path / "install" + target.mkdir() + archive = tmp_path / "evil.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("sd-cli", b"ok") + zf.writestr("../escape.txt", b"pwned") # escapes the install dir + with zipfile.ZipFile(archive) as zf: + with pytest.raises(RuntimeError, match = "unsafe path"): + _safe_extractall(zf, target) + assert not (tmp_path / "escape.txt").exists() + + +def test_safe_extractall_extracts_normal_members(tmp_path): + target = tmp_path / "install" + target.mkdir() + archive = tmp_path / "ok.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("build/bin/sd-cli", b"ok") + with zipfile.ZipFile(archive) as zf: + _safe_extractall(zf, target) + assert (target / "build" / "bin" / "sd-cli").read_bytes() == b"ok" + + +def test_find_sd_cpp_binary_honors_studio_home(tmp_path, monkeypatch): + # A binary installed under a custom Studio root must be discovered without also + # setting UNSLOTH_SD_CPP_PATH (matches default_install_dir's env handling). + from core.inference import sd_cpp_engine as eng + + monkeypatch.delenv("SD_CLI_PATH", raising = False) + monkeypatch.delenv("UNSLOTH_SD_CPP_PATH", raising = False) + studio_home = tmp_path / "studio_root" + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home)) + binary = tmp_path / "stable-diffusion.cpp" / "build" / "bin" / "sd-cli" + binary.parent.mkdir(parents = True) + binary.write_bytes(b"x") + assert eng.find_sd_cpp_binary() == str(binary) diff --git a/studio/install_sd_cpp_prebuilt.py b/studio/install_sd_cpp_prebuilt.py new file mode 100644 index 0000000000..4b2b36f95a --- /dev/null +++ b/studio/install_sd_cpp_prebuilt.py @@ -0,0 +1,278 @@ +# 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 _download( + url: str, + dest: Path, + *, + timeout: float = 300.0, +) -> None: + """Stream ``url`` to ``dest`` with an explicit timeout. ``urlretrieve`` takes no + timeout and can hang forever on a stalled socket.""" + import shutil + + req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-sd-cpp-installer"}) + with urllib.request.urlopen(req, timeout = timeout) as resp, open(dest, "wb") as f: # noqa: S310 + shutil.copyfileobj(resp, f) + + +def _safe_extractall(zf: zipfile.ZipFile, target: Path) -> None: + """``extractall`` with a per-member containment check, so an archive carrying an + absolute path or a ``..`` entry can't write outside ``target`` (Zip-Slip).""" + base = target.resolve() + for member in zf.infolist(): + dest = (base / member.filename).resolve() + if dest != base and base not in dest.parents: + raise RuntimeError(f"unsafe path in archive: {member.filename!r}") + zf.extractall(target) + + +def _maybe_fetch_windows_cudart(release: dict, chosen: str, target: Path) -> None: + """On Windows + a CUDA build, also fetch the separate CUDA-runtime DLL archive. + + Upstream ships the runtime as ``cudart-sd-...-win-cu12-...zip`` (which + ``resolve_release_asset`` filters out); without those DLLs ``sd-cli.exe`` cannot start + on a machine that does not already have the CUDA runtime installed.""" + if platform.system().lower() != "windows" or "cuda" not in chosen.lower(): + return + cudart = next( + ( + a + for a in release.get("assets", []) + if a["name"].lower().startswith("cudart") and "win" in a["name"].lower() + ), + None, + ) + if cudart is None: + return + dest = target / cudart["name"] + print(f"downloading CUDA runtime {cudart['name']} ...", flush = True) + try: + _download(cudart["browser_download_url"], dest) + with zipfile.ZipFile(dest) as zf: + _safe_extractall(zf, target) + finally: + dest.unlink(missing_ok = True) + + +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) + _download(url, archive) + print("extracting ...", flush = True) + with zipfile.ZipFile(archive) as zf: + _safe_extractall(zf, target) + archive.unlink(missing_ok = True) + # Windows CUDA builds need the separately-published cudart runtime DLLs. + _maybe_fetch_windows_cudart(release, chosen, target) + 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).expanduser() 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())