From 96940d87b7b94a88980acac38c1ce879282e4666 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:03:53 -0700 Subject: [PATCH 01/16] 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()) From 2d09508951c41b1e170b35ee8ea6ed3f80544b2b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:18:38 -0700 Subject: [PATCH 02/16] Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine (#6680) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/sd_cpp_smoke.py | 51 +++++- studio/backend/core/inference/sd_cpp_args.py | 104 ++++++++++++- .../backend/core/inference/sd_cpp_engine.py | 120 +++++++------- studio/backend/tests/test_sd_cpp_args.py | 147 ++++++++++++++++++ studio/backend/tests/test_sd_cpp_engine.py | 117 +++++--------- 5 files changed, 397 insertions(+), 142 deletions(-) diff --git a/scripts/sd_cpp_smoke.py b/scripts/sd_cpp_smoke.py index 4398fa6cf3..9deafb7e4f 100644 --- a/scripts/sd_cpp_smoke.py +++ b/scripts/sd_cpp_smoke.py @@ -41,6 +41,7 @@ from core.inference.diffusion_memory import ( # noqa: E402 from core.inference.sd_cpp_args import ( # noqa: E402 SdCppGenParams, SdCppModelFiles, + SdCppUpscaleParams, offload_flags, ) from core.inference.sd_cpp_engine import SdCppEngine, find_sd_cpp_binary # noqa: E402 @@ -55,9 +56,15 @@ _MODE_TO_POLICY = { def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(description = "Native sd-cli engine smoke test.") + p.add_argument("--task", default = "txt2img", choices = ["txt2img", "img2img", "upscale"]) 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("--diffusion-model", default = None) + # img2img + upscale inputs + p.add_argument("--init-img", default = None) + p.add_argument("--strength", type = float, default = 0.6) + p.add_argument("--upscale-model", default = None) + p.add_argument("--upscale-repeats", type = int, default = 1) p.add_argument("--vae", default = None) p.add_argument("--clip_l", default = None) p.add_argument("--t5xxl", default = None) @@ -90,6 +97,36 @@ def main(argv: list[str] | None = None) -> int: ) return 2 + out = Path(args.out_image) + + if args.task == "upscale": + if not args.init_img or not args.upscale_model: + print("ERROR: upscale needs --init-img and --upscale-model.", flush = True) + return 2 + t0 = time.time() + result = engine.upscale( + SdCppUpscaleParams( + input_image = args.init_img, + upscale_model = args.upscale_model, + repeats = args.upscale_repeats, + ), + output_path = str(out), + verbose = True, + timeout = args.timeout, + on_log = lambda ln: print(f" [sd] {ln}", flush = True), + ) + dt = time.time() - t0 + print( + f"\nOK: upscaled {result} ({result.stat().st_size/1024:.0f} KB) in {dt:.1f}s", + flush = True, + ) + print("SD-CPP-SMOKE-OK", flush = True) + return 0 + + if not args.diffusion_model: + print("ERROR: --diffusion-model is required for txt2img / img2img.", flush = True) + return 2 + files = SdCppModelFiles( diffusion_model = args.diffusion_model, vae = args.vae, @@ -98,6 +135,7 @@ def main(argv: list[str] | None = None) -> int: llm = args.llm, qwen2vl = args.qwen2vl, ) + is_img2img = args.task == "img2img" params = SdCppGenParams( prompt = args.prompt, negative_prompt = args.negative_prompt, @@ -106,12 +144,21 @@ def main(argv: list[str] | None = None) -> int: steps = args.steps, cfg_scale = args.cfg_scale, seed = args.seed, + init_img = args.init_img if is_img2img else None, + strength = args.strength if is_img2img else None, ) + if is_img2img and not args.init_img: + print("ERROR: img2img needs --init-img.", flush = True) + return 2 policy = _MODE_TO_POLICY[args.memory_mode] off = offload_flags(policy) + print( + f"task: {args.task}" + + (f" (init={args.init_img}, strength={args.strength})" if is_img2img else ""), + flush = True, + ) print(f"memory: {args.memory_mode} -> policy={policy} -> flags={off}", flush = True) - out = Path(args.out_image) t0 = time.time() result = engine.generate( files, diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 876cd7e557..44f57fb989 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -76,18 +76,46 @@ class SdCppModelFiles: @dataclass(frozen = True) class SdCppGenParams: - """Generation parameters, mapped 1:1 onto sd-cli's sampling flags.""" + """Generation parameters, mapped 1:1 onto sd-cli's sampling flags. + + The image-conditioning fields cover the img_gen variants: ``init_img`` + + ``strength`` make it img2img, adding ``mask`` makes it inpaint, and + ``ref_images`` drives FLUX Kontext / Qwen-Image-Edit style editing. ``lora_dir`` + points sd-cli at a LoRA directory; the LoRAs themselves are selected with + ```` tags inside ``prompt`` (sd.cpp's own syntax). + """ prompt: str negative_prompt: Optional[str] = None - width: int = 1024 - height: int = 1024 + # None = "unset": an image-conditioned run (img2img/inpaint/edit) then lets + # sd.cpp derive the size from the input image instead of forcing a resize; a + # plain txt2img run with unset dims falls back to 1024x1024 (see the builder). + width: Optional[int] = None + height: Optional[int] = None 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 + # image-to-image / inpaint / edit + init_img: Optional[str] = None + strength: Optional[float] = None + mask: Optional[str] = None + ref_images: tuple[str, ...] = () + # LoRA + lora_dir: Optional[str] = None + lora_apply_mode: Optional[str] = None + + +@dataclass(frozen = True) +class SdCppUpscaleParams: + """Inputs for sd-cli's ESRGAN upscale mode (a separate run mode).""" + + input_image: str + upscale_model: str + repeats: int = 1 + tile_size: Optional[int] = None def offload_flags( @@ -168,7 +196,31 @@ def build_sd_cpp_command( 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))] + # img2img / inpaint / edit conditioning (img_gen mode with an input image). + if params.init_img: + cmd += ["--init-img", params.init_img] + if params.strength is not None: + cmd += ["--strength", _fmt_float(params.strength)] + if params.mask: + cmd += ["--mask", params.mask] + for ref in params.ref_images: + cmd += ["--ref-image", ref] + # LoRA: the directory to scan; individual LoRAs are tags in prompt. + if params.lora_dir: + cmd += ["--lora-model-dir", params.lora_dir] + if params.lora_apply_mode: + cmd += ["--lora-apply-mode", params.lora_apply_mode] + # Emit explicit dims when given. For an image-conditioned run (img2img / + # inpaint / edit) that leaves them unset, omit the flags so sd.cpp derives the + # size from the input image (set_width_and_height_if_unset) rather than forcing + # a 1024x1024 resize/crop of the source. A plain txt2img run with unset dims + # keeps the prior 1024 default. + if params.width is not None or params.height is not None: + w = int(params.width) if params.width is not None else 1024 + h = int(params.height) if params.height is not None else 1024 + cmd += ["--width", str(w), "--height", str(h)] + elif not (params.init_img or params.ref_images): + cmd += ["--width", "1024", "--height", "1024"] if params.steps is not None: cmd += ["--steps", str(int(params.steps))] if params.cfg_scale is not None: @@ -194,6 +246,50 @@ def build_sd_cpp_command( return cmd +def build_sd_cpp_upscale_command( + binary: str, + params: SdCppUpscaleParams, + *, + output_path: str, + verbose: bool = False, + extra_args: Optional[list[str]] = None, +) -> list[str]: + """Build the ``sd-cli --mode upscale`` argv (ESRGAN super-resolution). + + Upscale is a distinct run mode: it takes an input image and an ESRGAN model, + no prompt or text encoders. ``repeats`` runs the upscaler N times (each pass + is a fixed scale factor for the model). + """ + if not params.input_image: + raise ValueError("input_image is required for upscale") + if not params.upscale_model: + raise ValueError("upscale_model is required for upscale") + # A truthiness guard below would silently swallow repeats=0 and fall back to + # sd-cli's default of one pass, turning an explicit no-op into a real upscale. + # Reject it (and negatives) so the caller's intent isn't quietly changed. + if params.repeats < 1: + raise ValueError("repeats must be >= 1 for upscale") + cmd: list[str] = [ + binary, + "--mode", + "upscale", + "--init-img", + params.input_image, + "--upscale-model", + params.upscale_model, + ] + if params.repeats != 1: + cmd += ["--upscale-repeats", str(int(params.repeats))] + if params.tile_size is not None: + cmd += ["--upscale-tile-size", str(int(params.tile_size))] + cmd += ["--output", output_path] + 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).""" diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 80078ef6fc..b1daafce72 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -27,7 +27,6 @@ import os import shutil import subprocess import sys -import threading import time from pathlib import Path from typing import Callable, Optional @@ -35,9 +34,10 @@ from typing import Callable, Optional from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, + SdCppUpscaleParams, build_sd_cpp_command, + build_sd_cpp_upscale_command, ) -from utils.process_lifetime import child_popen_kwargs logger = logging.getLogger(__name__) @@ -123,12 +123,8 @@ def find_sd_cpp_binary() -> Optional[str]: 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")) + # 3. Default install root (sibling of ~/.unsloth/llama.cpp). + hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp")) if hit: return hit @@ -206,28 +202,75 @@ class SdCppEngine: 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, + self._require_binary(), files, params, - output_path = str(out), + output_path = str(self._prepare_out(output_path)), offload = offload, threads = threads, verbose = verbose, extra_args = extra_args, ) + return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + + def upscale( + self, + params: "SdCppUpscaleParams", + *, + output_path: str, + 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: + """Upscale an image with an ESRGAN model; return the written path.""" + cmd = build_sd_cpp_upscale_command( + self._require_binary(), + params, + output_path = str(self._prepare_out(output_path)), + verbose = verbose, + extra_args = extra_args, + ) + return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + + # ── internals ───────────────────────────────────────────────────────────── + + def _require_binary(self) -> str: + 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." + ) + return self.binary # type: ignore[return-value] + + @staticmethod + def _prepare_out(output_path: str) -> Path: + out = Path(output_path) + out.parent.mkdir(parents = True, exist_ok = True) + return out + + def _run( + self, + cmd: list[str], + output_path: str, + *, + timeout: Optional[float], + env: Optional[dict[str, str]], + on_log: Optional[Callable[[str], None]], + ) -> Path: + """Run an sd-cli argv, stream output, and return the produced image path. + + Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output. + Shared by ``generate`` and ``upscale``. + """ + out = Path(output_path) base = dict(os.environ) if env: base.update(env) - run_env = runtime_env(self.binary, base) - logger.info("sd-cli generate: %s", " ".join(cmd)) + run_env = runtime_env(self._require_binary(), base) + logger.info("sd-cli run: %s", " ".join(cmd)) t0 = time.time() proc = subprocess.Popen( @@ -237,18 +280,9 @@ class SdCppEngine: 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: + try: assert proc.stdout is not None for line in proc.stdout: line = line.rstrip("\n") @@ -257,45 +291,23 @@ class SdCppEngine: 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: + if not out.is_file(): raise RuntimeError( f"sd-cli reported success but no image at {out}. Last output:\n" + "\n".join(tail[-12:]) ) - logger.info("sd-cli generate ok in %.1fs -> %s", time.time() - t0, produced) - return produced + logger.info("sd-cli run ok in %.1fs -> %s", time.time() - t0, out) + return out # ── engine routing ────────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 99e758e22e..c33a4a01f9 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -20,7 +20,9 @@ from core.inference.diffusion_memory import ( from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, + SdCppUpscaleParams, build_sd_cpp_command, + build_sd_cpp_upscale_command, offload_flags, text_encoder_flags_for_family, ) @@ -187,3 +189,148 @@ def test_build_requires_diffusion_model_and_prompt(): SdCppGenParams(prompt = " "), output_path = "/o.png", ) + + +# ── img2img / inpaint / edit / LoRA (Phase 6) ─────────────────────────────── + + +def test_build_img2img_adds_init_and_strength(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf", vae = "/m/ae.sft", llm = "/m/q.gguf") + params = SdCppGenParams(prompt = "make it autumn", init_img = "/in/src.png", strength = 0.6) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--init-img") == "/in/src.png" + assert _pair(cmd, "--strength") == "0.6" + assert _pair(cmd, "--mode") == "img_gen" # img2img is still img_gen mode + + +def test_build_inpaint_adds_mask(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + params = SdCppGenParams(prompt = "x", init_img = "/in/src.png", mask = "/in/mask.png", strength = 0.8) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--mask") == "/in/mask.png" + assert _pair(cmd, "--init-img") == "/in/src.png" + + +def test_build_edit_repeats_ref_image(): + files = SdCppModelFiles(diffusion_model = "/m/flux.gguf") + params = SdCppGenParams(prompt = "add a hat", ref_images = ("/r/a.png", "/r/b.png")) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + # each ref image gets its own --ref-image flag + idxs = [i for i, t in enumerate(cmd) if t == "--ref-image"] + assert len(idxs) == 2 + assert [cmd[i + 1] for i in idxs] == ["/r/a.png", "/r/b.png"] + + +def test_img2img_unset_dims_lets_sdcpp_derive_from_source(): + # img2img/inpaint/edit with dims left unset must NOT force --width/--height, + # so sd.cpp derives the size from the input image instead of resizing it to 1024. + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_command( + "/bin/sd-cli", + files, + SdCppGenParams(prompt = "x", init_img = "/in/src.png"), + output_path = "/o.png", + ) + assert "--width" not in cmd and "--height" not in cmd + # an edit (ref-image) run derives its size too + cmd2 = build_sd_cpp_command( + "/bin/sd-cli", + files, + SdCppGenParams(prompt = "x", ref_images = ("/r/a.png",)), + output_path = "/o.png", + ) + assert "--width" not in cmd2 and "--height" not in cmd2 + + +def test_img2img_explicit_dims_are_emitted(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_command( + "/bin/sd-cli", + files, + SdCppGenParams(prompt = "x", init_img = "/in/src.png", width = 768, height = 512), + output_path = "/o.png", + ) + assert _pair(cmd, "--width") == "768" and _pair(cmd, "--height") == "512" + + +def test_txt2img_unset_dims_keep_1024_default(): + # A plain txt2img run with no dims keeps the prior 1024x1024 default. + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_command( + "/bin/sd-cli", files, SdCppGenParams(prompt = "x"), output_path = "/o.png" + ) + assert _pair(cmd, "--width") == "1024" and _pair(cmd, "--height") == "1024" + + +def test_build_lora_dir_and_apply_mode(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + params = SdCppGenParams( + prompt = "a portrait ", + lora_dir = "/loras", + lora_apply_mode = "at_runtime", + ) + cmd = build_sd_cpp_command("/bin/sd-cli", files, params, output_path = "/o.png") + assert _pair(cmd, "--lora-model-dir") == "/loras" + assert _pair(cmd, "--lora-apply-mode") == "at_runtime" + # the tag rides in the prompt unchanged + assert _pair(cmd, "--prompt") == "a portrait " + + +def test_txt2img_omits_image_conditioning_flags(): + files = SdCppModelFiles(diffusion_model = "/m/z.gguf") + cmd = build_sd_cpp_command( + "/bin/sd-cli", files, SdCppGenParams(prompt = "x"), output_path = "/o.png" + ) + for flag in ("--init-img", "--strength", "--mask", "--ref-image", "--lora-model-dir"): + assert flag not in cmd + + +# ── upscale mode ──────────────────────────────────────────────────────────── + + +def test_build_upscale_command(): + params = SdCppUpscaleParams( + input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2 + ) + cmd = build_sd_cpp_upscale_command("/bin/sd-cli", params, output_path = "/out/big.png") + assert _pair(cmd, "--mode") == "upscale" + assert _pair(cmd, "--init-img") == "/in/small.png" + assert _pair(cmd, "--upscale-model") == "/m/esrgan.pth" + assert _pair(cmd, "--upscale-repeats") == "2" + assert _pair(cmd, "--output") == "/out/big.png" + # no prompt / text-encoder flags in upscale mode + assert "--prompt" not in cmd and "--llm" not in cmd + + +def test_build_upscale_rejects_non_positive_repeats(): + # repeats=0 must not be silently swallowed into sd-cli's default of one pass. + with pytest.raises(ValueError, match = "repeats"): + build_sd_cpp_upscale_command( + "/bin/sd-cli", + SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth", repeats = 0), + output_path = "/o.png", + ) + + +def test_build_upscale_default_repeats_omits_flag(): + cmd = build_sd_cpp_upscale_command( + "/bin/sd-cli", + SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), # repeats=1 + output_path = "/o.png", + ) + assert "--upscale-repeats" not in cmd + + +def test_build_upscale_requires_input_and_model(): + with pytest.raises(ValueError): + build_sd_cpp_upscale_command( + "/bin/sd-cli", + SdCppUpscaleParams(input_image = "", upscale_model = "/m/e.pth"), + output_path = "/o.png", + ) + with pytest.raises(ValueError): + build_sd_cpp_upscale_command( + "/bin/sd-cli", + SdCppUpscaleParams(input_image = "/i.png", upscale_model = ""), + output_path = "/o.png", + ) diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index ca6e875ea6..7aeb537033 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -26,7 +26,7 @@ from core.inference.sd_cpp_engine import ( runtime_env, select_diffusion_engine, ) -from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, SdCppUpscaleParams # ── binary discovery ──────────────────────────────────────────────────────── @@ -77,10 +77,7 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): # ── 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) +def test_engine_unavailable_when_no_binary(): e = SdCppEngine(binary = None) assert e.is_available() is False assert e.version() is None @@ -211,31 +208,6 @@ def test_generate_success_returns_path_and_collects_logs(tmp_path, monkeypatch): 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" @@ -270,60 +242,41 @@ def test_generate_raises_when_binary_missing(): ) -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) - +def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): 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, + out = tmp_path / "img.png" + src = tmp_path / "src.png" + src.write_bytes(b"\x89PNG\r\n") + _patch_popen(monkeypatch, lines = ["img2img"], returncode = 0, out_file = out) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x", init_img = str(src), strength = 0.5), + output_path = str(out), + ) + assert "--init-img" in _FakePopen.captured_cmd + assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1] + + +def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "big.png" + _patch_popen(monkeypatch, lines = ["upscaling", "done"], returncode = 0, out_file = out) + result = e.upscale( + SdCppUpscaleParams(input_image = "/in/small.png", upscale_model = "/m/esrgan.pth", repeats = 2), + output_path = str(out), + ) + assert result == out and out.is_file() + assert _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--mode") + 1] == "upscale" + assert "--upscale-model" in _FakePopen.captured_cmd + + +def test_upscale_raises_when_binary_missing(): + e = SdCppEngine(binary = None) + with pytest.raises(RuntimeError, match = "not found"): + e.upscale( + SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), + output_path = "/tmp/x.png", ) - assert holder["proc"].killed is True # ── engine routing ────────────────────────────────────────────────────────── From 6b9b1c72d3a5d3241faba3f9c2024d61bc712b68 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:31:50 -0700 Subject: [PATCH 03/16] Studio diffusion (Phase 7): accuracy-preserving speed pass (2.2x via GGUF compile) (#6690) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/compare_engines.py | 174 ++++++++++++++++ scripts/compile_probe.py | 189 ++++++++++++++++++ scripts/diffusion_bench.py | 9 + scripts/leverage_probe.py | 146 ++++++++++++++ scripts/perf_verify.py | 165 +++++++++++++++ studio/backend/core/inference/diffusion.py | 135 ++++++++----- .../core/inference/diffusion_families.py | 8 +- .../core/inference/diffusion_memory.py | 72 ++++--- .../backend/core/inference/diffusion_speed.py | 150 ++++++++++++-- studio/backend/core/inference/sd_cpp_args.py | 28 +++ .../backend/core/inference/sd_cpp_engine.py | 60 +++++- .../backend/tests/test_diffusion_backend.py | 16 +- studio/backend/tests/test_diffusion_memory.py | 15 ++ studio/backend/tests/test_diffusion_speed.py | 131 ++++++++++-- studio/backend/tests/test_sd_cpp_args.py | 11 + studio/backend/tests/test_sd_cpp_engine.py | 80 ++++++++ 16 files changed, 1261 insertions(+), 128 deletions(-) create mode 100644 scripts/compare_engines.py create mode 100644 scripts/compile_probe.py create mode 100644 scripts/leverage_probe.py create mode 100644 scripts/perf_verify.py diff --git a/scripts/compare_engines.py b/scripts/compare_engines.py new file mode 100644 index 0000000000..57abddfd3a --- /dev/null +++ b/scripts/compare_engines.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Head-to-head: PyTorch (diffusers GGUF) vs native stable-diffusion.cpp. + +Same Z-Image GGUF transformer, same VAE + text encoder, same resolution / steps / +seed, both resident (no CPU offload) on the same GPU. Reports per-engine compute +latency (model already loaded) so the denoise + VAE + TE work is compared fairly; +for sd.cpp it also reports the one-shot wall time (compute + the per-call model +reload, which a persistent sd-server would remove). + +PyTorch runs first (load / warmup / median), is unloaded, then sd.cpp runs. +""" + +from __future__ import annotations + +import argparse +import re +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)) + +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +_DONE_RE = re.compile(r"generate_image completed in ([0-9.]+)s") + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def bench_pytorch(repo, gguf, resolutions, steps, seed, iters): + import torch + from core.inference.diffusion import DiffusionBackend + + rows = [] + backend = DiffusionBackend() + for speed in ("off", "default"): + backend.begin_load(repo, gguf_filename = gguf, speed_mode = speed) + deadline = time.time() + 1800 # 30 min: a stuck download/load must not hang forever + while backend.load_progress().get("phase") != "ready": + prog = backend.load_progress() + if prog.get("phase") == "error": + raise RuntimeError(prog) + if time.time() > deadline: + raise TimeoutError(f"load timed out (last progress: {prog})") + time.sleep(0.5) + for res in resolutions: + + def gen(): + torch.cuda.synchronize() + t0 = time.time() + backend.generate( + prompt = PROMPT, + width = res, + height = res, + steps = steps, + guidance = 0.0, + seed = seed, + batch_size = 1, + ) + torch.cuda.synchronize() + return time.time() - t0 + + gen() # warmup (compiles for `default`) + med = _median([gen() for _ in range(iters)]) + rows.append(("pytorch", speed, res, med, None)) + print(f" pytorch speed={speed:7s} {res}px compute={med:.3f}s", flush = True) + backend.unload() + return rows + + +def bench_sdcpp(binary, gguf, vae, llm, resolutions, steps, seed, iters): + from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles + from core.inference.sd_cpp_engine import SdCppEngine + + engine = SdCppEngine(binary = binary) + if not engine.is_available(): + print(" sd.cpp binary not available; skipping", flush = True) + return [] + files = SdCppModelFiles(diffusion_model = gguf, vae = vae, llm = llm) + rows = [] + out_dir = Path("outputs/compare_engines") + out_dir.mkdir(parents = True, exist_ok = True) + for native in (None, "default"): # resident-no-fa vs resident+--diffusion-fa + for res in resolutions: + params = SdCppGenParams( + prompt = PROMPT, width = res, height = res, steps = steps, cfg_scale = 1.0, seed = seed + ) + computes, walls = [], [] + for _ in range(iters): + captured = {"c": None} + + def _log(ln): + m = _DONE_RE.search(ln) + if m: + captured["c"] = float(m.group(1)) + + t0 = time.time() + engine.generate( + files, + params, + output_path = str(out_dir / f"sd_{native}_{res}.png"), + offload = [], + native_speed = native, + on_log = _log, + ) + walls.append(time.time() - t0) + if captured["c"] is not None: + computes.append(captured["c"]) + med_c = _median(computes) if computes else None + med_w = _median(walls) + tag = "default(+fa)" if native == "default" else "off" + rows.append(("sdcpp", tag, res, med_c, med_w)) + print( + f" sdcpp speed={tag:12s} {res}px compute={med_c}s wall={med_w:.3f}s", + flush = True, + ) + return rows + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf-name", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument("--sd-binary", default = None) + p.add_argument( + "--sd-gguf", default = None, help = "local gguf for sd.cpp (default: same as pytorch via cache)" + ) + p.add_argument( + "--vae", + default = None, + help = "VAE safetensors for sd.cpp (required when benchmarking the sd.cpp engine)", + ) + p.add_argument( + "--llm", + default = None, + help = "text-encoder GGUF for sd.cpp (required when benchmarking the sd.cpp engine)", + ) + p.add_argument("--resolutions", default = "512,1024") + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + args = p.parse_args(argv) + + from huggingface_hub import hf_hub_download + from core.inference.sd_cpp_engine import find_sd_cpp_binary + + resolutions = [int(x) for x in args.resolutions.split(",")] + sd_gguf = args.sd_gguf or hf_hub_download(args.repo, args.gguf_name) + binary = args.sd_binary or find_sd_cpp_binary() + + print("== PyTorch (diffusers GGUF) ==", flush = True) + pt = bench_pytorch(args.repo, args.gguf_name, resolutions, args.steps, args.seed, args.iters) + print("== stable-diffusion.cpp (native) ==", flush = True) + sd = bench_sdcpp( + binary, sd_gguf, args.vae, args.llm, resolutions, args.steps, args.seed, args.iters + ) + + print("\n==== COMPARISON (Z-Image-Turbo Q4, fixed seed, resident) ====", flush = True) + print(f"{'engine':9s} {'config':13s} {'res':>5s} {'compute_s':>10s} {'wall_s':>8s}", flush = True) + for eng, cfg, res, c, w in pt + sd: + cs = f"{c:.3f}" if c is not None else "n/a" + ws = f"{w:.3f}" if w is not None else "-" + print(f"{eng:9s} {cfg:13s} {res:5d} {cs:>10s} {ws:>8s}", flush = True) + print("COMPARE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/compile_probe.py b/scripts/compile_probe.py new file mode 100644 index 0000000000..5ff16ef763 --- /dev/null +++ b/scripts/compile_probe.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 + +"""Probe: does regional ``torch.compile`` work on the GGUF diffusion transformer? + +The speed layer gates ``compile_repeated_blocks`` OFF for GGUF (it dequantises +per-op). Since the backend is GGUF-only, that makes regional compile dead on +every shipping model. This probe loads a GGUF transformer exactly as +``diffusion.py`` does, runs an eager generation, then compiles the repeated +denoiser block and runs the same seed again, reporting: whether compile raised, +per-generation latency eager vs compiled, and PSNR(compiled vs eager). If compile +is clean and PSNR is high, the gate can be relaxed for this family. + +Run on one CUDA GPU. Read-only w.r.t. the backend (does not import the gate). +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + + +def _psnr(a: "np.ndarray", b: "np.ndarray") -> float: + a = a.astype(np.float64) + b = b.astype(np.float64) + mse = float(np.mean((a - b) ** 2)) + if mse == 0.0: + return float("inf") + return float(10.0 * np.log10((255.0**2) / mse)) + + +def _gen(pipe, prompt, *, steps, seed, width, height, guidance): + import torch + + gen = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + image = pipe( + prompt = prompt, + width = width, + height = height, + num_inference_steps = steps, + guidance_scale = guidance, + generator = gen, + ).images[0] + torch.cuda.synchronize() + return image, time.time() - t0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument("--base-repo", default = "Tongyi-MAI/Z-Image-Turbo") + p.add_argument("--transformer-class", default = "ZImageTransformer2DModel") + p.add_argument("--pipeline-class", default = "ZImagePipeline") + p.add_argument( + "--prompt", + default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed", + ) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--width", type = int, default = 1024) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--guidance", type = float, default = 0.0) + p.add_argument( + "--mode", default = "default", help = "compile mode: default | max-autotune-no-cudagraphs" + ) + p.add_argument( + "--dynamic", action = "store_true", help = "dynamic=True (default False here for speed)" + ) + p.add_argument("--out-dir", default = "outputs/compile_probe") + args = p.parse_args(argv) + + import torch + import diffusers + from huggingface_hub import hf_hub_download + + out = Path(args.out_dir) + out.mkdir(parents = True, exist_ok = True) + dtype = torch.bfloat16 + + gguf_path = hf_hub_download(args.repo, args.gguf) + print(f"gguf: {gguf_path}", flush = True) + + transformer_cls = getattr(diffusers, args.transformer_class) + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = args.base_repo, + subfolder = "transformer", + ) + pipeline_cls = getattr(diffusers, args.pipeline_class) + pipe = pipeline_cls.from_pretrained(args.base_repo, torch_dtype = dtype, transformer = transformer) + pipe.to("cuda") + print("pipeline loaded on cuda", flush = True) + + # warm the eager path once (allocator / cudnn), then time eager. + _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + eager_img, eager_t = _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + eager_img.save(out / "eager.png") + eager_arr = np.array(eager_img) + print(f"EAGER: {eager_t:.2f}s/gen", flush = True) + + # compile the repeated denoiser block. + fn = getattr(pipe.transformer, "compile_repeated_blocks", None) + if not callable(fn): + print("RESULT: transformer has no compile_repeated_blocks -> N/A", flush = True) + return 3 + compile_kwargs = {"fullgraph": True, "dynamic": bool(args.dynamic)} + if args.mode and args.mode != "default": + compile_kwargs["mode"] = args.mode + print(f"compiling repeated blocks: {compile_kwargs} ...", flush = True) + try: + t0 = time.time() + fn(**compile_kwargs) + print( + f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: compile_repeated_blocks RAISED: {type(exc).__name__}: {exc}", flush = True) + return 1 + + # first compiled gen triggers the actual compilation (untimed warmup). + try: + t0 = time.time() + _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + print(f" first compiled gen (compilation) took {time.time()-t0:.1f}s", flush = True) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: first compiled generation RAISED: {type(exc).__name__}: {exc}", flush = True) + return 2 + + comp_img, comp_t = _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + comp_img.save(out / "compiled.png") + psnr = _psnr(eager_arr, np.array(comp_img)) + + speedup = (eager_t - comp_t) / eager_t * 100.0 + print("\n==== COMPILE PROBE RESULT ====", flush = True) + print(f" eager: {eager_t:.2f}s/gen", flush = True) + print(f" compiled: {comp_t:.2f}s/gen ({speedup:+.1f}% vs eager)", flush = True) + print(f" PSNR(compiled vs eager): {psnr:.1f} dB", flush = True) + print( + f" verdict: {'COMPILE-WORKS' if psnr >= 30 else 'COMPILE-DIVERGES'} " + f"{'FASTER' if comp_t < eager_t else 'NOT-FASTER'}", + flush = True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index dc49fe756b..89f59a688a 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -216,6 +216,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: hf_token = os.environ.get("HF_TOKEN"), cpu_offload = args.cpu_offload, memory_mode = args.memory_mode, + speed_mode = args.speed_mode, text_encoder_quant = args.text_encoder_quant, ) _wait_for_load(backend) @@ -295,6 +296,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "seed": args.seed, "batch_size": args.batch_size, "memory_mode": args.memory_mode, + "speed_mode": args.speed_mode, "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, }, @@ -450,6 +452,13 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["auto", "fast", "balanced", "low_vram"], help = "memory policy (default: backend auto)", ) + p.add_argument( + "--speed-mode", + default = None, + choices = ["off", "default", "max"], + help = "speed profile: off is bit-identical; default adds compile + " + "cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV", + ) p.add_argument( "--text-encoder-quant", default = None, diff --git a/scripts/leverage_probe.py b/scripts/leverage_probe.py new file mode 100644 index 0000000000..756fe25e34 --- /dev/null +++ b/scripts/leverage_probe.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe two candidate levers from the optimization research, on the real GGUF path: + + * `coordinate_descent_tuning` (Inductor) -- lossless extra kernel autotuning. + * FirstBlockCache (diffusers `apply_first_block_cache`) -- step-skip cache, lossy, + evaluated at a low (8) step count where its ceiling is lower. + +Each config is a fresh pipeline load (so Inductor config / compile artifacts don't +cross-contaminate). Reports latency + PSNR vs the eager reference. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +REPO = "unsloth/Z-Image-Turbo-GGUF" +GGUF = "z-image-turbo-Q4_K_M.gguf" +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +def _load(): + import torch + import diffusers + from huggingface_hub import hf_hub_download + + t = diffusers.ZImageTransformer2DModel.from_single_file( + hf_hub_download(REPO, GGUF), + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16), + torch_dtype = torch.bfloat16, + config = BASE, + subfolder = "transformer", + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + args = p.parse_args(argv) + steps, res, seed = args.steps, args.res, args.seed + + import torch + + def compile_blocks(pipe, *, cdt = False): + if cdt: + import torch._inductor.config as ic + ic.coordinate_descent_tuning = True + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + + def run( + tag, + *, + compile = False, + cdt = False, + fbc = None, + ): + # reset inductor config between runs + import torch._inductor.config as ic + + ic.coordinate_descent_tuning = False + torch.compiler.reset() + pipe = _load() + if fbc is not None: + from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbc)) + if compile: + compile_blocks(pipe, cdt = cdt) + _gen(pipe, steps, seed, res) # warmup / compilation + else: + _gen(pipe, steps, seed, res) # allocator warmup + img, dt = _gen(pipe, steps, seed, res) + del pipe + torch.cuda.empty_cache() + return tag, np.array(img), dt + + results = [] + print(f"== leverage probe (Z-Image Q4_K_M, {res}px, {steps} steps) ==", flush = True) + _, eager, eager_t = run("eager") + print(f" eager: {eager_t:.3f}s", flush = True) + results.append(("eager", eager_t, 0.0)) + + for tag, kw in [ + ("compile(default)", dict(compile = True)), + ("compile+coord_desc", dict(compile = True, cdt = True)), + ("fbc0.12+compile", dict(compile = True, fbc = 0.12)), + ("fbc0.20+compile", dict(compile = True, fbc = 0.20)), + ("fbc0.20(no compile)", dict(fbc = 0.20)), + ]: + try: + t, img, dt = run(tag, **kw) + ps = _psnr(eager, img) + results.append((tag, dt, ps)) + print( + f" {tag:22s} {dt:.3f}s ({(eager_t-dt)/eager_t*100:+.0f}% vs eager) PSNR={ps:.1f} dB", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f" {tag:22s} FAILED: {type(exc).__name__}: {str(exc)[:140]}", flush = True) + + print("\n==== SUMMARY ====", flush = True) + for tag, dt, ps in results: + sp = f"{(results[0][1]-dt)/results[0][1]*100:+.0f}%" if tag != "eager" else "ref" + pss = f"{ps:.1f}dB" if ps else "ref" + print(f" {tag:24s} {dt:.3f}s {sp:>6s} {pss:>8s}", flush = True) + print("LEVERAGE-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/perf_verify.py b/scripts/perf_verify.py new file mode 100644 index 0000000000..d3c8330239 --- /dev/null +++ b/scripts/perf_verify.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GPU verification for the diffusion performance pass (Phase 7). + +Drives the real ``DiffusionBackend`` through several loads in one process and +checks, at a fixed seed: + + 1. speed: ``default`` (compile + cudnn.benchmark + channels_last) vs ``off`` + -- expect a large denoise speedup at high PSNR (near-lossless). + 2. the TF32-leak fix: load ``max`` (flips global TF32 / cudnn.benchmark), unload, + then load ``off`` -- the ``off`` image must be byte-identical (PSNR inf) to a + fresh ``off`` baseline, proving the globals were restored on unload. + 3. ``balanced`` is now bit-identical: with VAE tiling restricted to the low tiers, + streamed (group) offload should match the resident image (PSNR inf). + +Run on one CUDA GPU with the GGUF + base repo cached. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend" +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +def _psnr(a: "np.ndarray", b: "np.ndarray") -> float: + a = a.astype(np.float64) + b = b.astype(np.float64) + mse = float(np.mean((a - b) ** 2)) + return float("inf") if mse == 0.0 else float(10.0 * np.log10((255.0**2) / mse)) + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--model", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument( + "--prompt", + default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed", + ) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--width", type = int, default = 1024) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--out-dir", default = "outputs/perf_verify") + args = p.parse_args(argv) + + import os + + import torch + from core.inference.diffusion import DiffusionBackend + + out = Path(args.out_dir) + out.mkdir(parents = True, exist_ok = True) + backend = DiffusionBackend() + token = os.environ.get("HF_TOKEN") + + def load(mode_speed = None, mode_mem = None): + backend.begin_load( + args.model, + gguf_filename = args.gguf, + hf_token = token, + speed_mode = mode_speed, + memory_mode = mode_mem, + ) + deadline = time.time() + 2400 + while time.time() < deadline: + ph = backend.load_progress().get("phase") + if ph == "ready": + return backend.status() + if ph == "error": + raise RuntimeError(f"load error: {backend.load_progress()}") + time.sleep(0.5) + raise RuntimeError("load timed out") + + def gen(): + torch.cuda.synchronize() + t0 = time.time() + img = backend.generate( + prompt = args.prompt, + width = args.width, + height = args.height, + steps = args.steps, + guidance = 0.0, + seed = args.seed, + batch_size = 1, + )["images"][0] + torch.cuda.synchronize() + return img, time.time() - t0 + + def timed( + mode_speed, + *, + warmup, + iters, + mem = None, + tag = "", + ): + st = load(mode_speed, mem) + for _ in range(warmup): + gen() + lats = [] + img = None + for _ in range(iters): + img, dt = gen() + lats.append(dt) + img.save(out / f"{tag}.png") + backend.unload() + med = sorted(lats)[len(lats) // 2] + print( + f" [{tag}] speed={mode_speed} mem={mem} optims={st.get('speed_optims')} " + f"tiling={st.get('vae_tiling')} median={med:.3f}s", + flush = True, + ) + return np.array(img), med + + print("== 1. speed: off vs default ==", flush = True) + off_img, off_t = timed("off", warmup = 1, iters = 3, tag = "off") + def_img, def_t = timed("default", warmup = 1, iters = 3, tag = "default") + print(f" PSNR(default vs off) = {_psnr(off_img, def_img):.1f} dB", flush = True) + print( + f" speedup: off {off_t:.3f}s -> default {def_t:.3f}s " + f"({(off_t-def_t)/off_t*100:+.1f}%)", + flush = True, + ) + + print("== 2. TF32-leak fix: max then off must be byte-identical ==", flush = True) + timed("max", warmup = 0, iters = 1, tag = "max") # flips + should restore globals + off2_img, _ = timed("off", warmup = 0, iters = 1, tag = "off2") + leak_psnr = _psnr(off_img, off2_img) + print( + f" PSNR(off-after-max vs off) = {leak_psnr:.1f} dB " + f"({'OK byte-identical' if leak_psnr == float('inf') else 'LEAK! globals not restored'})", + flush = True, + ) + + print("== 3. balanced is bit-identical (tiling off) ==", flush = True) + bal_img, bal_t = timed("off", warmup = 0, iters = 1, mem = "balanced", tag = "balanced") + bal_psnr = _psnr(off_img, bal_img) + print( + f" PSNR(balanced vs off) = {bal_psnr:.1f} dB " + f"({'OK bit-identical' if bal_psnr == float('inf') else 'differs'})", + flush = True, + ) + + ok = ( + (leak_psnr == float("inf")) + and (bal_psnr == float("inf")) # check 3: balanced must be bit-identical to off + and (def_t < off_t) + and (_psnr(off_img, def_img) >= 30) + ) + print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8ad5ffa250..9500f2ed9f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -44,7 +44,13 @@ from .diffusion_memory import ( plan_diffusion_memory, snapshot_device_memory, ) -from .diffusion_speed import SPEED_OFF, apply_speed_optims +from .diffusion_speed import ( + SPEED_OFF, + apply_speed_optims, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, +) from .diffusion_precision import quantize_text_encoders logger = get_logger(__name__) @@ -69,6 +75,10 @@ class _LoadState: # The opt-in speed profile (Phase 3). speed_mode: str = SPEED_OFF speed_optims: tuple = () + # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the + # speed layer mutated them, restored on unload so a later `off` load is not + # contaminated by this one's globals. None when nothing was changed. + backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None @@ -495,56 +505,77 @@ class DiffusionBackend: pipeline_cls = getattr(diffusers, fam.pipeline_class) pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + # Resolve the effective speed mode: GGUF models default to the + # near-lossless `default` profile (compile is ~2.2x and sits below + # the quant noise floor), dense models stay bit-identical `off`. An + # explicit speed_mode (incl. "off") is honored verbatim. + effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) # Opt-in speed optims run BEFORE placement (channels_last / compile - # must precede CPU offload). Off by default -> bit-identical output. - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = speed_mode or SPEED_OFF, - logger = logger, - ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, - ) + # must precede CPU offload). Snapshot the process-wide backend flags + # first so unload can restore them: TF32 / cudnn.benchmark are global, + # and a later `off` load must not inherit this load's settings. + backend_flags_before = snapshot_backend_flags() + # apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark); + # they are only restored via _LoadState.backend_flags_before on unload. If + # the build fails after this but before _state commits (e.g. an OOM in + # apply_memory_plan / pipe.to), nothing would restore them and a later `off` + # generation would be contaminated, so restore on any non-committed exit. + committed = False + try: + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + logger = logger, + ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, + ) - # Decide placement from MEASURED free device memory vs the model's - # 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. `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. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Decide placement from MEASURED free device memory vs the model's + # 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. `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. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = (speed_mode or SPEED_OFF), - speed_optims = tuple(k for k, v in speed_applied.items() if v), - text_encoder_quant = te_quant, - ) + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + ) + committed = True + finally: + if not committed: + restore_backend_flags(backend_flags_before) + clear_gpu_cache() logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -679,7 +710,10 @@ class DiffusionBackend: self._gen = gen try: - images = state.pipe(**kwargs).images + # inference_mode is strictly faster than the no_grad diffusers + # uses internally and numerically identical for inference. + with torch.inference_mode(): + images = state.pipe(**kwargs).images finally: self._gen = None # A cancelled denoise returns early with a partial/garbage image; @@ -738,6 +772,9 @@ class DiffusionBackend: state = self._state if state is None: return + # Restore the process-wide backend flags (TF32 / cudnn.benchmark) this load + # may have flipped, so the next `off` load is bit-identical again. + restore_backend_flags(state.backend_flags_before) self._state = None del state clear_gpu_cache() diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 4fecf950ad..4a3dd616ac 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -35,9 +35,9 @@ class DiffusionFamily: # (~6.5e4) and produce inf -> NaN latents -> a black image. The backend # 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). Only consulted on the non-GGUF path; the - # GGUF transformer is never compiled regardless. + # Set False only for a family whose denoiser block does not compile cleanly with + # regional torch.compile. Now consulted on the GGUF path too (compile runs on the + # GGUF transformer); all current families compile, so this stays True. supports_torch_compile: bool = True @@ -81,8 +81,6 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("zimage", "z_image"), # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, - # Z-Image's denoiser block is excluded from regional torch.compile. - supports_torch_compile = False, ), ) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index ca90e11c73..9172ef7cf8 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -400,17 +400,21 @@ def plan_diffusion_memory( 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") + # VAE savers cap the decode-time spike that dominates peak VRAM at high res. + # Slicing (decode a batch one image at a time) is EXACT, so enable it on any + # offload tier / non-discrete backend. Tiling (spatial chunks) is only bit- + # identical for a single tile (<=1MP), so restrict it to the lowest tiers where + # the VAE itself is offloaded (model / sequential) or there is no spare device + # pool (MPS / CPU). Under group offload the transformer streams but the VAE stays + # resident and fits, so it keeps exact full-image decode -> balanced is both + # faster and bit-identical. On a roomy discrete GPU both stay off. + any_offload = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu") + tile = policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL) or device_memory.backend in ("mps", "cpu") return MemoryPlan( requested_mode = mode, offload_policy = policy, vae_tiling = tile, - vae_slicing = tile, + vae_slicing = any_offload, device_memory = device_memory, estimates = estimates, reasons = tuple(reasons), @@ -442,12 +446,22 @@ def apply_memory_plan( if plan.vae_slicing: _enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger) + def _fallback_to_model_offload() -> None: + # Group offload keeps the VAE resident, so the GROUP plan set vae_tiling=False. + # When group offload is unavailable and we drop to whole-module offload, the card + # is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE + # tiling on now (if not already engaged) to cap it. + nonlocal tiling_engaged + pipe.enable_model_cpu_offload() + if not tiling_engaged: + tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger) + policy = plan.offload_policy if policy == OFFLOAD_MODEL: pipe.enable_model_cpu_offload() elif policy == OFFLOAD_GROUP: if not _apply_group_offload(pipe, device, logger): - pipe.enable_model_cpu_offload() + _fallback_to_model_offload() policy = OFFLOAD_MODEL elif policy == OFFLOAD_SEQUENTIAL: try: @@ -459,7 +473,7 @@ def apply_memory_plan( "falling back to whole-module offload", exc, ) - pipe.enable_model_cpu_offload() + _fallback_to_model_offload() policy = OFFLOAD_MODEL else: pipe.to(device) @@ -491,32 +505,40 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: if transformer is None: return False try: + import inspect + import torch from diffusers.hooks import apply_group_offloading onload = torch.device(device) use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA - # Place the smaller components resident FIRST: moving them onto the device is - # the only step here that can OOM on a tight GPU. Doing it before - # apply_group_offloading means a failure leaves NO group-offload hooks on the - # transformer, so the caller's whole-module-offload fallback gets a clean - # pipeline -- diffusers refuses enable_model_cpu_offload() while group hooks - # are attached, which would otherwise turn the fallback into a hard crash. + gkwargs: dict[str, Any] = { + "onload_device": onload, + "offload_device": torch.device("cpu"), + "offload_type": "block_level", + "num_blocks_per_group": DEFAULT_GROUP_BLOCKS, + "use_stream": use_stream, + } + # On the CUDA stream path, overlap each block's host->device copy with + # compute: non_blocking issues the copy asynchronously and record_stream + # defers the free until the copy's stream is done. Lossless (only transfer + # scheduling changes). Safe for the group tier specifically, where the + # companions stay resident; gated on the installed signature so an older + # diffusers that lacks these kwargs still works (no hard fallback). + if use_stream: + _params = inspect.signature(apply_group_offloading).parameters + if "non_blocking" in _params: + gkwargs["non_blocking"] = True + if "record_stream" in _params: + gkwargs["record_stream"] = True + apply_group_offloading(transformer, **gkwargs) + # Place the remaining (smaller) components resident; the streamed + # transformer manages its own placement via the offloading hooks. for name, comp in getattr(pipe, "components", {}).items(): if name == "transformer": continue if isinstance(comp, torch.nn.Module): comp.to(onload) - # Stream the transformer a few blocks at a time; it manages its own placement - # via the offloading hooks. - apply_group_offloading( - transformer, - onload_device = onload, - offload_device = torch.device("cpu"), - offload_type = "block_level", - num_blocks_per_group = DEFAULT_GROUP_BLOCKS, - use_stream = use_stream, - ) return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index e01ba48b63..9184f29f88 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -5,19 +5,29 @@ Off by default, so the default render path stays bit-identical to a plain run (the property the regression harness checks). When the operator opts in, this applies the -lossless-to-near-lossless speedups in the order the diffusers guides recommend -(channels_last -> regional compile, with TF32 / fused-QKV under "max"): +near-lossless speedups in the order the diffusers guides recommend +(channels_last + cudnn.benchmark -> regional compile, with TF32 / fused-QKV under +"max"): - off - nothing (default). - default - lossless: channels_last VAE memory format + regional torch.compile of - the denoiser's repeated block WHERE eligible (non-GGUF, bf16, CUDA, and - a compile-friendly family). + off - nothing (default; bit-identical reference). + default - near-lossless: channels_last VAE memory format + cudnn.benchmark conv + autotune + regional torch.compile of the denoiser's repeated block WHERE + eligible (bf16, CUDA, a compile-friendly family). Compile is the big win + (~2.3x denoise on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, + well above the Q4 quantisation noise floor, so it does not meaningfully + move output quality). max - default plus near-lossless TF32 matmul and fused QKV projections. -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. +Regional compile used to be gated off for the GGUF transformer, but it compiles and +runs faster on the current diffusers/torch (measured; the GGUF dequant ops stay +eager and the rest of the repeated block compiles), so the GGUF gate is removed; the +per-family ``supports_torch_compile`` flag and the bf16/CUDA checks still apply. + +The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so +``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the +prior values at load and restore them at unload, keeping a later ``off`` load +bit-identical instead of inheriting a previous ``max`` run's globals. torch is +imported lazily. """ from __future__ import annotations @@ -30,6 +40,54 @@ SPEED_MAX = "max" SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX) +def snapshot_backend_flags() -> Optional[dict]: + """Capture the process-wide torch backend flags this layer may mutate, so the + caller can restore them on unload. None if torch is unavailable. Each flag is read + defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still + captures the rest -- otherwise a single missing attribute would skip the whole + snapshot and a real mutated flag would leak.""" + try: + import torch + except Exception: # noqa: BLE001 — no torch -> nothing to snapshot/restore + return None + state: dict[str, bool] = {} + matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None) + if matmul is not None and hasattr(matmul, "allow_tf32"): + state["matmul_tf32"] = bool(matmul.allow_tf32) + cudnn = getattr(torch.backends, "cudnn", None) + if cudnn is not None: + if hasattr(cudnn, "allow_tf32"): + state["cudnn_tf32"] = bool(cudnn.allow_tf32) + if hasattr(cudnn, "benchmark"): + state["cudnn_benchmark"] = bool(cudnn.benchmark) + return state + + +def restore_backend_flags(state: Optional[dict]) -> None: + """Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each + flag is restored independently so one failure can't leave the others leaked.""" + if not state: + return + try: + import torch + except Exception: # noqa: BLE001 — no torch -> nothing to restore + return + + def _set(obj: Any, attr: str, key: str) -> None: + if obj is not None and key in state and hasattr(obj, attr): + try: + setattr(obj, attr, state[key]) + except Exception: # noqa: BLE001 — best-effort per-flag restore + pass + + _set( + getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32" + ) + cudnn = getattr(torch.backends, "cudnn", None) + _set(cudnn, "allow_tf32", "cudnn_tf32") + _set(cudnn, "benchmark", "cudnn_benchmark") + + def normalize_speed_mode(value: Optional[str]) -> str: """Lower/strip a requested speed mode (dashes ok); None / "" -> off.""" if value is None: @@ -44,14 +102,29 @@ def normalize_speed_mode(value: Optional[str]) -> str: return normalized +def resolve_speed_mode(value: Optional[str], *, is_gguf: bool) -> str: + """The effective speed mode when the caller leaves it UNSET (``None``). + + A GGUF model defaults to ``default``: regional compile is ~2.2x faster and its + numeric perturbation sits well below the quantisation noise floor (measured + PSNR ~37 dB compile-vs-eager versus ~21 dB Q4-vs-bf16), so it does not reduce + output quality relative to the dense reference. A dense (non-GGUF) model stays + ``off`` / bit-identical, since there compile would be the only source of drift. + An explicit value -- including ``"off"`` -- is always honored verbatim.""" + if value is None: + return SPEED_DEFAULT if is_gguf else SPEED_OFF + return normalize_speed_mode(value) + + 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 non-GGUF - bf16 transformer, on a compile-friendly family. The GGUF transformer is never - compiled (it dequantises per-op).""" - if is_gguf: - return False + Only on CUDA (incl. ROCm via supports_default_torch_compile), for a bf16 + transformer, on a compile-friendly family. ``is_gguf`` no longer disqualifies: + ``compile_repeated_blocks`` runs fine on the GGUF transformer (the per-op + dequant stays eager, the rest of the block compiles) and is ~2.3x faster, so it + is kept only for signature/logging compatibility.""" + del is_gguf # GGUF is compile-eligible now; param kept for call-site compat. if not bool(getattr(target, "supports_default_torch_compile", False)): return False if not bool(getattr(family, "supports_torch_compile", True)): @@ -79,7 +152,13 @@ def apply_speed_optims( """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, BEFORE placement / offload. Returns which optimisations actually engaged. Every step is best-effort: a pipeline that doesn't support one is simply skipped.""" - applied = {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False} + applied = { + "channels_last": False, + "cudnn_benchmark": False, + "tf32": False, + "fused_qkv": False, + "compiled": False, + } mode = normalize_speed_mode(speed_mode) # TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this # load isn't max, so a later default/off diffusion load -- or chat inference in the @@ -93,9 +172,16 @@ def apply_speed_optims( # Lossless: a channels-last VAE speeds up its convolutions with no numeric change. applied["channels_last"] = _vae_channels_last(pipe, logger) - # Lossless-ish: regional compile of the repeated denoiser block, where eligible. + # Near-lossless: let cuDNN autotune the fixed-shape VAE convs (CUDA only). It may + # pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win. + if getattr(target, "device", None) == "cuda": + applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) + + # Near-lossless and the largest win: regional compile of the repeated denoiser + # block, where eligible (now incl. the GGUF transformer). `max` opts into + # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): - applied["compiled"] = _compile_repeated_blocks(pipe, logger) + applied["compiled"] = _compile_repeated_blocks(pipe, logger, max_autotune = mode == SPEED_MAX) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. @@ -119,19 +205,43 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False -def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool: +def _compile_repeated_blocks( + pipe: Any, + logger: Any, + *, + max_autotune: bool = False, +) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) if not callable(fn): return False + # default: mode="default" + dynamic=True -- fast cold start, robust to resolution + # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False -- + # Triton autotuning for a few % more on GEMM/conv-heavy models, at a much longer + # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead + # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled + # block because its static output buffer is overwritten across denoise steps. + kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune} + if max_autotune: + kwargs["mode"] = "max-autotune-no-cudagraphs" try: - fn(fullgraph = True, dynamic = True) + fn(**kwargs) return True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) return False +def _enable_cudnn_benchmark(logger: Any) -> bool: + try: + import torch + torch.backends.cudnn.benchmark = True + return True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "cudnn_benchmark", exc) + return False + + # The TF32 flag values from before the first max load flipped them, so a later # non-max load / unload can put the process back exactly as it found it (rather than # forcing a hardcoded default that might clobber another component's choice). diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 44f57fb989..8069e8f610 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -118,6 +118,34 @@ class SdCppUpscaleParams: tile_size: Optional[int] = None +# Native (sd.cpp) speed profiles, the engine-side analogue of diffusion_speed's +# modes. off: nothing (default). default: --diffusion-fa (flash attention; upstream +# reports it usually speeds CUDA and cuts attention memory, near-lossless). max: also +# --diffusion-conv-direct (direct conv; helps some backends, but measured +45% on +# CUDA here, so it stays opt-in/experimental, never auto-on for CUDA). +NATIVE_SPEED_OFF = "off" +NATIVE_SPEED_DEFAULT = "default" +NATIVE_SPEED_MAX = "max" +NATIVE_SPEED_MODES = (NATIVE_SPEED_OFF, NATIVE_SPEED_DEFAULT, NATIVE_SPEED_MAX) + + +def native_speed_flags(speed_mode: Optional[str]) -> list[str]: + """sd-cli speed flags for a native speed mode (empty for off / None). + + These are separate from the offload flags: ``--diffusion-fa`` is a speed/memory + win in its own right, not tied to whether weights are offloaded. De-duplicated + against offload flags at the call site (offload already adds ``--diffusion-fa``). + """ + mode = (speed_mode or NATIVE_SPEED_OFF).strip().lower() + if mode in ("", NATIVE_SPEED_OFF): + return [] + if mode == NATIVE_SPEED_DEFAULT: + return ["--diffusion-fa"] + if mode == NATIVE_SPEED_MAX: + return ["--diffusion-fa", "--diffusion-conv-direct"] + raise ValueError(f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'") + + def offload_flags( policy: str, *, diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index b1daafce72..186ea51631 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -24,9 +24,11 @@ from __future__ import annotations import logging import os +import queue import shutil import subprocess import sys +import threading import time from pathlib import Path from typing import Callable, Optional @@ -37,6 +39,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, ) logger = logging.getLogger(__name__) @@ -189,6 +192,7 @@ class SdCppEngine: *, output_path: str, offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, threads: Optional[int] = None, verbose: bool = False, extra_args: Optional[list[str]] = None, @@ -198,10 +202,15 @@ class SdCppEngine: ) -> 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. + ``native_speed`` ("default"/"max") adds sd.cpp's own speed flags + (``--diffusion-fa`` etc.), de-duplicated against the offload flags that may + already include them. 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. """ + offload = list(offload or []) + speed = [f for f in native_speed_flags(native_speed) if f not in offload] + merged_extra = speed + list(extra_args or []) cmd = build_sd_cpp_command( self._require_binary(), files, @@ -210,7 +219,7 @@ class SdCppEngine: offload = offload, threads = threads, verbose = verbose, - extra_args = extra_args, + extra_args = merged_extra, ) return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) @@ -281,20 +290,49 @@ class SdCppEngine: errors = "replace", env = run_env, ) + # Drain stdout on a reader thread so the timeout is enforced even when the + # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain + # `for line in proc.stdout` blocks until EOF, so proc.wait(timeout) would + # never be reached. The reader pushes lines (then a None sentinel at EOF) to a + # queue the main loop polls against a wall-clock deadline. tail: list[str] = [] + line_q: "queue.Queue[Optional[str]]" = queue.Queue() + + def _drain() -> None: + try: + assert proc.stdout is not None + for raw in proc.stdout: + line_q.put(raw.rstrip("\n")) + finally: + line_q.put(None) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + + deadline = None if timeout is None else time.monotonic() + float(timeout) + stdout_done = False try: - assert proc.stdout is not None - for line in proc.stdout: - line = line.rstrip("\n") + while True: + if deadline is not None and time.monotonic() >= deadline and proc.poll() is None: + proc.kill() + raise RuntimeError(f"sd-cli timed out after {timeout}s") + try: + line = line_q.get(timeout = 0.1) + except queue.Empty: + if proc.poll() is not None and stdout_done: + break + continue + if line is None: + stdout_done = True + if proc.poll() is not None: + break + continue tail.append(line) if len(tail) > 40: tail.pop(0) if on_log is not None: on_log(line) - ret = proc.wait(timeout = timeout) - except subprocess.TimeoutExpired: - proc.kill() - raise RuntimeError(f"sd-cli timed out after {timeout}s") + ret = proc.wait(timeout = 5.0) finally: if proc.poll() is None: proc.kill() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index c7c4736502..fe4cabe8b5 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -10,6 +10,7 @@ GPU, weights, or network access is needed (sub-second, CI-friendly). from __future__ import annotations +import contextlib import sys import types @@ -202,6 +203,8 @@ def fake_runtime(monkeypatch): torch.Generator = _FakeGenerator torch.cuda = types.SimpleNamespace(is_available = lambda: False) torch.backends = types.SimpleNamespace(mps = None) + # generate() wraps the pipe call in torch.inference_mode(); a no-op CM here. + torch.inference_mode = lambda: contextlib.nullcontext() diffusers = types.ModuleType("diffusers") diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) @@ -856,12 +859,19 @@ def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( 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). +def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path): + # No speed_mode on a GGUF model -> auto `default` (near-lossless, compile sits + # below the quant noise floor). compile itself only engages on CUDA, so on this + # CPU stub no optim need engage, but the resolved mode is `default`. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") - assert status["speed_mode"] == "off" and status["speed_optims"] == [] + assert status["speed_mode"] == "default" + # An explicit "off" opts back into the bit-identical path (engages nothing). + status_off = backend.load_pipeline( + str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "off" + ) + assert status_off["speed_mode"] == "off" and status_off["speed_optims"] == [] # An explicit speed_mode threads through to status (engaged optims are GPU-verified). status2 = backend.load_pipeline( str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "max" diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 858e87ffa1..bc09ffbd0f 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -185,6 +185,9 @@ def test_auto_group_offload_when_transformer_overflows_but_companions_fit(): base_overhead_mib = 1000, ) assert plan.offload_policy == OFFLOAD_GROUP + # Group keeps the VAE resident, so it uses exact slicing but NOT lossy tiling + # -> balanced stays bit-identical while still capping the offload footprint. + assert plan.vae_slicing is True and plan.vae_tiling is False def test_auto_model_offload_when_companions_exceed_budget(): @@ -432,6 +435,18 @@ def test_apply_group_falls_back_to_model_without_transformer(): assert effective == OFFLOAD_MODEL and "model_offload" in pipe.calls +def test_apply_group_fallback_enables_vae_tiling(): + # A balanced/group plan keeps the VAE resident (tiling off); when group offload can't + # engage and we drop to whole-module offload, the applier must turn VAE tiling ON to + # cap the decode-time spike on what is now a low-VRAM path. + plan = _plan(OFFLOAD_GROUP, tiling = True) + assert plan.vae_tiling is False # group plan leaves tiling off by design + pipe = _RecordingPipe() # no .transformer -> group offload falls back to model + effective, tiled = apply_memory_plan(pipe, plan, device = "cuda") + assert effective == OFFLOAD_MODEL + assert tiled is True and "vae_tiling" in pipe.calls + + def test_apply_sequential_offload(): pipe = _RecordingPipe() effective, _ = apply_memory_plan( diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index ade37482ba..b0ca10af4c 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -21,6 +21,9 @@ from core.inference.diffusion_speed import ( apply_speed_optims, compile_eligible, normalize_speed_mode, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, ) @@ -47,7 +50,7 @@ def _stub_torch(monkeypatch): torch.channels_last = "channels_last" torch.backends = types.SimpleNamespace( cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)), - cudnn = types.SimpleNamespace(allow_tf32 = False), + cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False), ) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -64,23 +67,91 @@ def test_normalize_speed_mode(): normalize_speed_mode("ludicrous") +def test_resolve_speed_mode_gguf_auto_default(): + # Unset (None) -> default for GGUF (near-lossless), off for dense. + assert resolve_speed_mode(None, is_gguf = True) == SPEED_DEFAULT + assert resolve_speed_mode(None, is_gguf = False) == SPEED_OFF + # An explicit value is honored verbatim, including an explicit opt-out to off. + assert resolve_speed_mode("off", is_gguf = True) == SPEED_OFF + assert resolve_speed_mode("max", is_gguf = True) == SPEED_MAX + assert resolve_speed_mode("max", is_gguf = False) == SPEED_MAX + + # ── compile gating ──────────────────────────────────────────────────────────── -def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch): +def test_compile_eligible_requires_bf16_cuda_friendly(monkeypatch): _stub_torch(monkeypatch) - # The happy path: non-GGUF, bf16, CUDA, compile-friendly family. + # The happy path: 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 + # GGUF is now compile-eligible too (measured ~2.3x, PSNR ~37 dB vs eager). + assert compile_eligible(_target(), is_gguf = True, family = _family()) is True # fp16 (non-bf16) is excluded. assert compile_eligible(_target(dtype = "float16"), is_gguf = False, family = _family()) is False - # A family flagged not compile-friendly (Z-Image) is excluded. + # A family flagged not compile-friendly is excluded. assert compile_eligible(_target(), is_gguf = False, family = _family(compile_ok = False)) is False - # No compile support (e.g. ROCm/XPU/MPS) is excluded. + # No compile support (e.g. XPU/MPS) is excluded. assert compile_eligible(_target(compile_ok = False), is_gguf = False, family = _family()) is False +# ── backend-flag snapshot / restore (TF32 / cudnn.benchmark leak guard) ──────── + + +def test_snapshot_restore_backend_flags(monkeypatch): + torch = _stub_torch(monkeypatch) + snap = snapshot_backend_flags() + assert snap == {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False} + # An opt-in max run flips the globals on... + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + # ...and restore puts them back, so a later `off` load is bit-identical again. + restore_backend_flags(snap) + assert torch.backends.cuda.matmul.allow_tf32 is False + assert torch.backends.cudnn.allow_tf32 is False + assert torch.backends.cudnn.benchmark is False + + +def test_restore_backend_flags_tolerates_none(): + restore_backend_flags(None) # no torch needed, no-op + + +def test_snapshot_partial_when_some_backends_missing(monkeypatch): + # A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the + # flags it does have, rather than skipping the whole snapshot on one missing attribute. + torch = types.ModuleType("torch") + torch.backends = types.SimpleNamespace( + cuda = types.SimpleNamespace(), # no .matmul + cudnn = types.SimpleNamespace(benchmark = True), # no .allow_tf32 + ) + monkeypatch.setitem(sys.modules, "torch", torch) + snap = snapshot_backend_flags() + assert snap == {"cudnn_benchmark": True} + torch.backends.cudnn.benchmark = False + restore_backend_flags(snap) + assert torch.backends.cudnn.benchmark is True + + +def test_restore_is_independent_per_flag(monkeypatch): + # A read-only / failing attribute must not abort restoring the remaining flags. + torch = _stub_torch(monkeypatch) + + class _NoMatmulSet: + @property + def allow_tf32(self): + return False + + @allow_tf32.setter + def allow_tf32(self, value): + raise RuntimeError("read-only on this build") + + torch.backends.cuda.matmul = _NoMatmulSet() + snap = {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False} + torch.backends.cudnn.benchmark = True + restore_backend_flags(snap) # matmul setter raises, cudnn still restored + assert torch.backends.cudnn.benchmark is False + + # ── applier ─────────────────────────────────────────────────────────────────── @@ -103,24 +174,33 @@ class _Pipe: def _vae_to(self, *, memory_format): self.vae.mem_format = memory_format - def _compile(self, *, fullgraph, dynamic): + def _compile(self, **kwargs): self.compiled = True + self.compile_kwargs = kwargs def _fuse(self): self.fused = True def test_speed_off_applies_nothing(monkeypatch): - _stub_torch(monkeypatch) + torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = True) 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 applied == { + "channels_last": False, + "cudnn_benchmark": False, + "tf32": False, + "fused_qkv": False, + "compiled": False, + } assert pipe.vae.mem_format is None and pipe.compiled is False + # off must not touch any process-wide flag (bit-identical reference path). + assert torch.backends.cudnn.benchmark is False -def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch): +def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( @@ -128,18 +208,36 @@ def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch): ) 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. + # default compiles with dynamic=True and no autotune mode (fast cold start, + # resolution-robust, sidesteps the CUDA-graph crash). + assert pipe.compile_kwargs == {"fullgraph": True, "dynamic": True} + # default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV. + assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True assert applied["tf32"] is False and applied["fused_qkv"] is False -def test_speed_default_skips_compile_for_gguf(monkeypatch): +def test_speed_default_compiles_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 + assert applied["channels_last"] is True + # GGUF now compiles (the big near-lossless win). + assert applied["compiled"] is True and pipe.compiled is True + + +def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch): + _stub_torch(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, + _target(device = "mps", compile_ok = False), + is_gguf = True, + family = _family(), + speed_mode = SPEED_DEFAULT, + ) + assert applied["cudnn_benchmark"] is False # not CUDA -> no autotune flip def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): @@ -150,6 +248,9 @@ def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): ) 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 + # max opts into autotuned kernels (static shapes); CUDA-graph modes are avoided. + assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs" + assert pipe.compile_kwargs["dynamic"] is False def test_speed_max_tf32_only_on_cuda(monkeypatch): diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index c33a4a01f9..200672aa26 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -23,6 +23,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, offload_flags, text_encoder_flags_for_family, ) @@ -47,6 +48,16 @@ def test_te_flags_by_family(): # ── offload policy -> sd-cli flags ────────────────────────────────────────── +def test_native_speed_flags(): + assert native_speed_flags(None) == [] + assert native_speed_flags("off") == [] + assert native_speed_flags("") == [] + assert native_speed_flags("default") == ["--diffusion-fa"] + assert native_speed_flags("max") == ["--diffusion-fa", "--diffusion-conv-direct"] + with pytest.raises(ValueError): + native_speed_flags("ludicrous") + + def test_offload_none_is_empty(): assert offload_flags(OFFLOAD_NONE) == [] diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 7aeb537033..a879252761 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -12,6 +12,7 @@ from __future__ import annotations import os import sys +import time import types from pathlib import Path @@ -242,6 +243,55 @@ def test_generate_raises_when_binary_missing(): ) +class _HangingPopen: + """A child that runs but never prints and never exits -- the case a plain + `for line in stdout` would block on forever, ignoring the timeout.""" + + def __init__(self, cmd, **_kw): + self._alive = True + + class _Blocking: + def __init__(self, owner): + self.owner = owner + + def __iter__(self): + return self + + def __next__(self): + while self.owner._alive: + time.sleep(0.01) + raise StopIteration + + @property + def stdout(self): + return self._Blocking(self) + + def poll(self): + return None if self._alive else -9 + + def wait(self, timeout = None): + self._alive = False + return -9 + + def kill(self): + self._alive = False + + +def test_generate_times_out_on_silent_hang(tmp_path, monkeypatch): + e = _engine(tmp_path) + monkeypatch.setattr(eng.subprocess, "Popen", lambda cmd, **kw: _HangingPopen(cmd, **kw)) + t0 = time.time() + with pytest.raises(RuntimeError, match = "timed out"): + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(tmp_path / "x.png"), + timeout = 0.3, + ) + # The timeout is enforced promptly (not blocked until stdout EOF). + assert time.time() - t0 < 5.0 + + def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "img.png" @@ -257,6 +307,36 @@ def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1] +def test_generate_native_speed_dedupes_against_offload(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + # offload already adds --diffusion-fa; native_speed="default" would add it again. + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = ["--offload-to-cpu", "--diffusion-fa"], + native_speed = "default", + ) + # --diffusion-fa appears exactly once (de-duped), not twice. + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + +def test_generate_native_speed_adds_flag_when_not_offloaded(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = [], # fast/resident tier: no offload, but speed flag still applies + native_speed = "default", + ) + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "big.png" From f43910a4a4bdff028d63ee7b56c1482a1a8d9388 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:36:43 -0700 Subject: [PATCH 04/16] Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) (#6694) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/diffusion_bench.py | 21 ++ scripts/fp8_overflow_check.py | 115 ++++++ scripts/nvfp4_probe.py | 186 ++++++++++ scripts/nvfp4_t211_probe.py | 301 +++++++++++++++ scripts/quant_probe.py | 318 ++++++++++++++++ scripts/sparse_accum_probe.py | 251 +++++++++++++ studio/backend/core/inference/diffusion.py | 246 +++++++++---- .../inference/diffusion_transformer_quant.py | 347 ++++++++++++++++++ studio/backend/models/inference.py | 22 ++ studio/backend/routes/inference.py | 2 + .../backend/tests/test_diffusion_backend.py | 128 +++++++ studio/backend/tests/test_diffusion_routes.py | 38 ++ .../tests/test_diffusion_transformer_quant.py | 342 +++++++++++++++++ 13 files changed, 2243 insertions(+), 74 deletions(-) create mode 100644 scripts/fp8_overflow_check.py create mode 100644 scripts/nvfp4_probe.py create mode 100644 scripts/nvfp4_t211_probe.py create mode 100644 scripts/quant_probe.py create mode 100644 scripts/sparse_accum_probe.py create mode 100644 studio/backend/core/inference/diffusion_transformer_quant.py create mode 100644 studio/backend/tests/test_diffusion_transformer_quant.py diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 89f59a688a..43a92ac8ee 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -218,6 +218,10 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: memory_mode = args.memory_mode, speed_mode = args.speed_mode, text_encoder_quant = args.text_encoder_quant, + transformer_quant = args.transformer_quant, + transformer_quant_fast_accum = {"auto": None, "on": True, "off": False}[ + args.fp8_fast_accum + ], ) _wait_for_load(backend) _cuda_sync() @@ -299,6 +303,8 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "speed_mode": args.speed_mode, "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, + "transformer_quant": args.transformer_quant, + "fp8_fast_accum": args.fp8_fast_accum, }, } @@ -465,6 +471,21 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["fp8", "nvfp4"], help = "quantise the companion text encoder (fp8 or nvfp4)", ) + p.add_argument( + "--transformer-quant", + default = None, + choices = ["auto", "int8", "fp8", "nvfp4", "mxfp8"], + help = "opt-in fast transformer: load the DENSE bf16 transformer and torchao-" + "quantise it onto the low-precision tensor cores (faster than GGUF, higher " + "VRAM). auto picks per GPU; falls back to GGUF if unsupported / no VRAM", + ) + p.add_argument( + "--fp8-fast-accum", + default = "auto", + choices = ["auto", "on", "off"], + help = "fp8 accumulate: auto picks by GPU class (fast on consumer, precise on " + "data-center); on/off force it", + ) p.add_argument( "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload" ) diff --git a/scripts/fp8_overflow_check.py b/scripts/fp8_overflow_check.py new file mode 100644 index 0000000000..157b9cce30 --- /dev/null +++ b/scripts/fp8_overflow_check.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empirically check that fp8 dynamic quant with fast accumulation does not overflow. + +Hooks every quantised Linear's output during a real Z-Image generation and reports the +global max |output| and any non-finite (Inf/NaN) count, for use_fast_accum True vs False. +The concern fast_accum raises is accumulation *precision*, not overflow (the accumulator +stays FP32-range and torchao's dynamic per-row scale keeps FP8 inputs <= 448); this proves +it on the real model, including Z-Image's large (~9e5) activation peaks. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _run(fast_accum, steps, res, seed, mf): + import torch + import torch.nn as nn + from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig + from torchao.float8 import Float8MMConfig + + pipe = _load_dense() + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + quantize_( + pipe.transformer, + Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = fast_accum) + ), + filter_fn = filt, + ) + + stats = {"max_abs": 0.0, "nonfinite": 0, "hooked": 0} + + def hook(mod, inp, out): + t = out[0] if isinstance(out, tuple) else out + if not torch.is_tensor(t): + return + finite = torch.isfinite(t) + nf = int((~finite).sum().item()) + stats["nonfinite"] += nf + m = float(t[finite].abs().max().item()) if finite.any() else float("inf") + if m > stats["max_abs"]: + stats["max_abs"] = m + + # Hook the quantised linears (where an fp8-accumulation overflow would surface). + # Run EAGER: forward hooks don't trace through torch.compile, and the fp8 fast-accum + # accumulation is identical compiled or eager -- compile only changes scheduling. + for m in pipe.transformer.modules(): + if isinstance(m, nn.Linear): + m.register_forward_hook(hook) + stats["hooked"] += 1 + + g = torch.Generator(device = "cuda").manual_seed(seed) + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + import numpy as np + + arr = np.array(img) + img_finite = bool(np.isfinite(arr).all()) + del pipe + torch.cuda.empty_cache() + return stats, img_finite + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 4) + p.add_argument("--res", type = int, default = 512) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--min-feat", type = int, default = 512) + args = p.parse_args(argv) + + print(f"== fp8 overflow check (Z-Image dense, {args.res}px, {args.steps} steps) ==", flush = True) + for fast in (True, False): + stats, img_finite = _run(fast, args.steps, args.res, args.seed, args.min_feat) + print( + f" fast_accum={str(fast):5s} hooked_linears={stats['hooked']:3d} " + f"max|linear_out|={stats['max_abs']:.1f} nonfinite_elems={stats['nonfinite']} " + f"image_all_finite={img_finite}", + flush = True, + ) + print("FP8-OVERFLOW-CHECK-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/nvfp4_probe.py b/scripts/nvfp4_probe.py new file mode 100644 index 0000000000..902cf8a2d2 --- /dev/null +++ b/scripts/nvfp4_probe.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe NVFP4 via torchao with use_triton_kernel=False (no MSLK) on the real dense +Z-Image transformer: is it a genuine FP4-tensor-core speedup over fp8, and is quality +in-bar? Reference for LPIPS is dense bf16 eager. Run on one CUDA (Blackwell) GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import torch, lpips + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + p.add_argument("--min-feat", type = int, default = 512) + p.add_argument("--out-dir", default = None, help = "image output dir (default: repo outputs/)") + args = p.parse_args(argv) + steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat + import torch + import torch.nn as nn + + global OUT + if args.out_dir: + OUT = Path(args.out_dir).expanduser() + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + cfg = None, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + if cfg is not None: + from torchao.quantization import quantize_ + quantize_(pipe.transformer, cfg, filter_fn = filt) + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print( + f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True + ) + _gen(pipe, steps, seed, res) # warmup / compile + dts, img = [], None + for _ in range(args.iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + + print(f"== nvfp4 probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True) + bref, ref, _ = run("bf16_eager", cfg = None, compile = False) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + rows = [("bf16_eager", bref, float("inf"), 0.0, None)] + + specs = [ + ("bf16_compile", None, True), + ("fp8_compile", FP8(), True), + ("nvfp4_notriton_compile", NV(use_triton_kernel = False), True), + ("nvfp4_notriton_eager", NV(use_triton_kernel = False), False), + ] + for tag, cfg, comp in specs: + try: + med, arr, gp = run(tag, cfg = cfg, compile = comp) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None) + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:24s} FAILED") + continue + vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + print("NVFP4-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/nvfp4_t211_probe.py b/scripts/nvfp4_t211_probe.py new file mode 100644 index 0000000000..e69e01b2de --- /dev/null +++ b/scripts/nvfp4_t211_probe.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""torch>=2.11 NVFP4 probe. Three parts: + + A. diagnostics -- torch/torchao versions, cpp-extension load state, device. + B. GEMM micro -- isolated per-linear forward latency (bf16 / fp8 / nvfp4-cutlass / + nvfp4-triton) at Z-Image-like shapes, to measure raw FP4 + tensor-core throughput free of pipeline overhead. + C. end-to-end -- real dense Z-Image transformer, latency + LPIPS + PSNR + VRAM, + reference = dense bf16 eager. + +Run on one CUDA (Blackwell) GPU. This is the experiment that decides whether NVFP4 +becomes a genuine speedup once torch>=2.11 + torchao's CUTLASS FP4 GEMM is present.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images") + + +# ----------------------------------------------------------------------------- diag +def diagnostics() -> None: + import torch + import torchao + + print("== A. diagnostics ==", flush = True) + print(f" torch {torch.__version__}", flush = True) + print(f" torchao {torchao.__version__}", flush = True) + print(f" cuda {torch.version.cuda}", flush = True) + if torch.cuda.is_available(): + print( + f" device {torch.cuda.get_device_name(0)} sm{torch.cuda.get_device_capability(0)}", + flush = True, + ) + print(f" torch.ops.torchao present: {hasattr(torch.ops, 'torchao')}", flush = True) + print( + f" fp4 primitives: e2m1={hasattr(torch, 'float4_e2m1fn_x2')} " + f"e8m0={hasattr(torch, 'float8_e8m0fnu')} _scaled_mm={hasattr(torch, '_scaled_mm')}", + flush = True, + ) + # torchao prints "Skipping import of cpp extensions ..." to stderr at import on torch<2.11. + # On 2.11 that line is absent -> the CUTLASS FP4 GEMM extension is live. + print( + " (no 'Skipping import of cpp extensions' line above => cpp/CUTLASS ext loaded)", + flush = True, + ) + + +# ----------------------------------------------------------------------------- micro +def _configs(): + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + return { + "bf16": None, + "fp8": FP8(), + "nvfp4_cutlass": NV(use_triton_kernel = False), + "nvfp4_triton": NV(use_triton_kernel = True), + } + + +def _bench_linear(K, N, M, cfg, iters, compile_): + import torch + import torch.nn as nn + from torchao.quantization import quantize_ + + torch.compiler.reset() + torch.cuda.empty_cache() + m = nn.Sequential(nn.Linear(K, N, bias = False)).cuda().to(torch.bfloat16) + if cfg is not None: + quantize_(m, cfg) + fn = torch.compile(m, fullgraph = True, dynamic = False) if compile_ else m + x = torch.randn(M, K, device = "cuda", dtype = torch.bfloat16) + with torch.no_grad(): + for _ in range(3): # warmup / compile + fn(x) + torch.cuda.synchronize() + dts = [] + for _ in range(iters): + t0 = time.perf_counter() + fn(x) + torch.cuda.synchronize() + dts.append(time.perf_counter() - t0) + del m, fn, x + torch.cuda.empty_cache() + med = sorted(dts)[len(dts) // 2] + tflops = 2.0 * M * K * N / med / 1e12 + return med, tflops + + +def micro(M, iters, compile_): + print(f"\n== B. GEMM micro (M={M}, compile={compile_}, iters={iters}) ==", flush = True) + # (K, N): qkv-ish, mlp-up, mlp-down for a ~3072-dim DiT + shapes = [(3072, 3072), (3072, 12288), (12288, 3072)] + cfgs = _configs() + for K, N in shapes: + print(f" shape K={K} N={N}:", flush = True) + base_ms = None + fp8_ms = None + for name, cfg in cfgs.items(): + try: + med, tfl = _bench_linear(K, N, M, cfg, iters, compile_) + ms = med * 1e3 + if name == "bf16": + base_ms = ms + if name == "fp8": + fp8_ms = ms + vs_bf16 = f"{base_ms/ms:.2f}x" if base_ms else "-" + vs_fp8 = f"{fp8_ms/ms:.2f}x" if fp8_ms else "-" + print( + f" {name:16s} {ms:7.3f} ms {tfl:7.1f} TFLOPS vs_bf16={vs_bf16:>6s} vs_fp8={vs_fp8:>6s}", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f" {name:16s} FAILED: {type(exc).__name__}: {str(exc)[:120]}", flush = True) + + +# ----------------------------------------------------------------------------- e2e +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import diffusers + import torch + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def e2e(steps, res, seed, iters, mf): + import torch + import torch.nn as nn + + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + cfg = None, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + if cfg is not None: + from torchao.quantization import quantize_ + quantize_(pipe.transformer, cfg, filter_fn = filt) + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print( + f" [{tag}] compile failed: {type(exc).__name__}: {str(exc)[:90]}", flush = True + ) + _gen(pipe, steps, seed, res) # warmup / compile + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig as NV + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig as FP8 + + print( + f"\n== C. end-to-end (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", flush = True + ) + bref, ref, _ = run("bf16_eager", cfg = None, compile = False) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + rows = [("bf16_eager", bref, float("inf"), 0.0, None)] + + specs = [ + ("bf16_compile", None, True), + ("fp8_compile", FP8(), True), + ("nvfp4_cutlass_compile", NV(use_triton_kernel = False), True), + ("nvfp4_triton_compile", NV(use_triton_kernel = True), True), + ] + for tag, cfg, comp in specs: + try: + med, arr, gp = run(tag, cfg = cfg, compile = comp) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:24s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:24s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + fp8 = next((r[1] for r in rows if r[0] == "fp8_compile" and r[1]), None) + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:24s} FAILED") + continue + vs_fp8 = f"{fp8/med:.2f}x" if fp8 else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:24s} {med:.3f}s vs_fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + p.add_argument("--micro-M", type = int, default = 4096) + p.add_argument("--min-feat", type = int, default = 512) + p.add_argument("--only", choices = ["diag", "micro", "e2e", "all"], default = "all") + args = p.parse_args(argv) + + diagnostics() + if args.only in ("micro", "all"): + micro(args.micro_M, args.iters, compile_ = True) + if args.only in ("e2e", "all"): + e2e(args.steps, args.res, args.seed, args.iters, args.min_feat) + print("NVFP4-T211-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/quant_probe.py b/scripts/quant_probe.py new file mode 100644 index 0000000000..8675fb75bf --- /dev/null +++ b/scripts/quant_probe.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Empirical quant probe: torchao int8/fp8/fp4 dynamic quant vs GGUF+compile. + +Question this answers: GGUF stores the Z-Image DiT at 4-bit but dequantizes to bf16 +per matmul, so it runs at bf16 tensor-core rate. Can a low-precision *tensor-core* +path (int8dq on any Ampere+, fp8dq on Ada+, NVFP4/MXFP8 on Blackwell), loaded from +the dense bf16 transformer, beat GGUF+compile on speed while staying inside the +quality bar -- and how does its quality compare to GGUF's own 4-bit loss? + +Reference for all quality numbers is the DENSE bf16 EAGER image (the best this model +can do). Each config is a fresh pipeline (no compile/quant cross-contamination). +Reports median latency, PSNR + LPIPS vs reference, and peak VRAM. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +REPO = "unsloth/Z-Image-Turbo-GGUF" +GGUF = "z-image-turbo-Q4_K_M.gguf" +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LPIPS = {"fn": None} + + +def _lpips(ref_arr, arr): + """Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable. + + Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a + resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM.""" + try: + import torch + import lpips + + if _LPIPS["fn"] is None: + _LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval() + + def t(x): + return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 + + with torch.no_grad(): + return float(_LPIPS["fn"](t(ref_arr), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips unavailable: {type(exc).__name__}: {str(exc)[:80]})", flush = True) + return None + + +def _load_dense(): + import torch + import diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _load_gguf(): + import torch + import diffusers + from huggingface_hub import hf_hub_download + + t = diffusers.ZImageTransformer2DModel.from_single_file( + hf_hub_download(REPO, GGUF), + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16), + torch_dtype = torch.bfloat16, + config = BASE, + subfolder = "transformer", + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _quant_config(name): + """Return a torchao config instance for `name`, or raise to mark FAILED.""" + from torchao.quantization import ( + Int8WeightOnlyConfig, + Int8DynamicActivationInt8WeightConfig, + Float8DynamicActivationFloat8WeightConfig, + ) + + if name == "int8wo": + return Int8WeightOnlyConfig() + if name == "int8dq": + return Int8DynamicActivationInt8WeightConfig() + if name == "fp8dq": + return Float8DynamicActivationFloat8WeightConfig() + if name == "nvfp4": + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig + return NVFP4DynamicActivationNVFP4WeightConfig() + if name == "mxfp8": + from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig + try: + import torch + return MXDynamicActivationMXWeightConfig( + activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn + ) + except (TypeError, AttributeError): + return MXDynamicActivationMXWeightConfig() + raise ValueError(name) + + +def _make_filter_fn(min_features): + """Keep only the FLOP-heavy linears: nn.Linear with both in/out >= min_features. + The int8 dynamic path uses torch._int_mm (needs activation M>16), and the tiny + timestep/pooled projections (in_features=256) run at M=1 and crash it -- skip them.""" + import torch.nn as nn + + def filter_fn(module, fqn = ""): + return ( + isinstance(module, nn.Linear) + and getattr(module, "in_features", 0) >= min_features + and getattr(module, "out_features", 0) >= min_features + ) + + return filter_fn + + +def _apply_quant(pipe, name, log, min_features): + import torch.nn as nn + from torchao.quantization import quantize_ + + cfg = _quant_config(name) + total = sum(1 for m in pipe.transformer.modules() if isinstance(m, nn.Linear)) + filt = _make_filter_fn(min_features) + q = sum(1 for n, m in pipe.transformer.named_modules() if filt(m, n)) + quantize_(pipe.transformer, cfg, filter_fn = filt) + log(f" quantized transformer with {name} ({q}/{total} linears >= {min_features} feat)") + + +def _compile(pipe, log): + fn = getattr(pipe.transformer, "compile_repeated_blocks", None) + if not callable(fn): + return False + for kw in ({"fullgraph": True, "dynamic": True}, {"dynamic": True}, {}): + try: + fn(**kw) + log(f" compiled repeated blocks {kw}") + return True + except Exception as exc: # noqa: BLE001 + log(f" compile {kw} failed: {type(exc).__name__}: {str(exc)[:90]}") + return False + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + p.add_argument( + "--min-feat", + type = int, + default = 512, + help = "only quantize Linear with in&out features >= this (int8 _int_mm needs M>16)", + ) + p.add_argument( + "--configs", + default = "bf16,bf16_c,gguf_c,int8dq_c,fp8dq_c,nvfp4_c,mxfp8_c,int8wo_c", + help = "comma list; suffix _c = +compile", + ) + args = p.parse_args(argv) + steps, res, seed, iters = args.steps, args.res, args.seed, args.iters + + import torch + + OUT.mkdir(parents = True, exist_ok = True) + + def run( + tag, + *, + source, + quant = None, + compile = False, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() if source == "dense" else _load_gguf() + load_peak = torch.cuda.max_memory_allocated() / 1e9 + if quant is not None: + _apply_quant(pipe, quant, print_, args.min_feat) + if compile: + _compile(pipe, print_) + _gen(pipe, steps, seed, res) # warmup / compilation + else: + _gen(pipe, steps, seed, res) # allocator warmup + torch.cuda.reset_peak_memory_stats() + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gen_peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return tag, _median(dts), arr, load_peak, gen_peak + + print_ = lambda s: print(s, flush = True) # noqa: E731 + + # config table: tag -> (source, quant, compile) + table = { + "bf16": ("dense", None, False), + "bf16_c": ("dense", None, True), + "gguf_c": ("gguf", None, True), + "int8wo_c": ("dense", "int8wo", True), + "int8dq_c": ("dense", "int8dq", True), + "fp8dq_c": ("dense", "fp8dq", True), + "nvfp4_c": ("dense", "nvfp4", True), + "mxfp8_c": ("dense", "mxfp8", True), + } + want = [c.strip() for c in args.configs.split(",") if c.strip()] + + print(f"== quant probe (Z-Image-Turbo, {res}px, {steps} steps, seed {seed}) ==", flush = True) + ref_arr = None + rows = [] + for tag in want: + if tag not in table: + print(f" {tag}: unknown config, skipping", flush = True) + continue + source, quant, compile = table[tag] + print(f"-- {tag} (source={source} quant={quant} compile={compile}) --", flush = True) + try: + _, med, arr, lp, gp = run(tag, source = source, quant = quant, compile = compile) + except Exception as exc: # noqa: BLE001 + import traceback + + print(f" {tag:10s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + traceback.print_exc() + rows.append((tag, None, None, None, None, None)) + continue + if ref_arr is None and tag == "bf16": + ref_arr = arr + psnr = _psnr(ref_arr, arr) if ref_arr is not None else None + lpips_v = ( + _lpips(ref_arr, arr) + if (ref_arr is not None and tag != "bf16") + else (0.0 if tag == "bf16" else None) + ) + rows.append((tag, med, psnr, lpips_v, lp, gp)) + ps = f"{psnr:.1f}dB" if psnr is not None else "n/a" + lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a" + print( + f" {tag:10s} {med:.3f}s PSNR={ps:>7s} LPIPS={lps:>6s} loadVRAM={lp:.1f}G genVRAM={gp:.1f}G", + flush = True, + ) + + base = next((r[1] for r in rows if r[0] == "bf16" and r[1]), None) + gguf = next((r[1] for r in rows if r[0] == "gguf_c" and r[1]), None) + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + print( + f"{'config':10s} {'sec':>7s} {'vs_bf16':>8s} {'vs_gguf':>8s} {'PSNR':>8s} {'LPIPS':>7s} {'loadG':>6s} {'genG':>6s}", + flush = True, + ) + for tag, med, psnr, lpips_v, lp, gp in rows: + if med is None: + print(f"{tag:10s} {'FAILED':>7s}", flush = True) + continue + vb = f"{base/med:.2f}x" if base else "-" + vg = f"{gguf/med:.2f}x" if gguf else "-" + ps = ( + f"{psnr:.1f}" + if psnr is not None and psnr != float("inf") + else ("inf" if psnr == float("inf") else "n/a") + ) + lps = f"{lpips_v:.3f}" if lpips_v is not None else "n/a" + print( + f"{tag:10s} {med:>7.3f} {vb:>8s} {vg:>8s} {ps:>8s} {lps:>7s} {lp:>6.1f} {gp:>6.1f}", + flush = True, + ) + print("QUANT-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/sparse_accum_probe.py b/scripts/sparse_accum_probe.py new file mode 100644 index 0000000000..4f91f905b6 --- /dev/null +++ b/scripts/sparse_accum_probe.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe two consumer-GPU-motivated levers on the real dense Z-Image transformer: + + * fp8 fast_accum on/off -- on consumer Blackwell, fp8 with FP16 accumulate is ~2x + fp8 with FP32 accumulate (838 vs 419 TFLOPS). torchao defaults use_fast_accum=True, + so this confirms we are already on the fast path and quantifies it (muted on a B200, + which is not nerfed, but the knob still moves latency). + * 2:4 semi-structured sparsity -- doubles tensor-core rate in theory. Two blockers to + test empirically: (a) QUALITY -- inference-only 2:4 magnitude-pruning drops 50% of + weights with no fine-tune; (b) it does NOT compose with torch.compile, so the real + sparse path runs eager. We measure sparse-no-compile speed vs our fp8+compile + baseline (the bar it must beat) and the LPIPS of 2:4 pruning. + +Reference for quality is the dense bf16 eager image. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images") + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import torch, lpips + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load_dense(): + import torch, diffusers + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _big_linears(transformer, min_feat = 512): + import torch.nn as nn + return [ + m + for m in transformer.modules() + if isinstance(m, nn.Linear) and m.in_features >= min_feat and m.out_features >= min_feat + ] + + +def _prune_24_(transformer, min_feat = 512): + """In-place 2:4 magnitude prune (zero the 2 smallest of every 4 along in_features) + of the FLOP-heavy linears. Dense format -> measures the QUALITY of 2:4 with no kernel.""" + import torch + + n = 0 + for lin in _big_linears(transformer, min_feat): + w = lin.weight.data + o, i = w.shape + if i % 4: + continue + g = w.view(o, i // 4, 4) + idx = g.abs().argsort(dim = -1)[..., :2] + g.scatter_(-1, idx, 0.0) + n += 1 + return n + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + p.add_argument("--min-feat", type = int, default = 512) + args = p.parse_args(argv) + steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat + import torch + + OUT.mkdir(parents = True, exist_ok = True) + + def filt(mod, fqn = ""): + import torch.nn as nn + return isinstance(mod, nn.Linear) and mod.in_features >= mf and mod.out_features >= mf + + def run( + tag, + *, + quant = None, + fast_accum = True, + prune = False, + real_sparse = False, + compile = True, + ): + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load_dense() + note = "" + if prune or real_sparse: + n = _prune_24_(pipe.transformer, mf) + note += f" pruned24={n}" + if real_sparse: + from torchao.sparsity import sparsify_, semi_sparse_weight + sparsify_(pipe.transformer, semi_sparse_weight(), filter_fn = filt) + note += " +semi_sparse" + if quant == "fp8": + from torchao.quantization import quantize_, Float8DynamicActivationFloat8WeightConfig + from torchao.float8 import Float8MMConfig + + cfg = Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = fast_accum) + ) + quantize_(pipe.transformer, cfg, filter_fn = filt) + note += f" fp8(fast_accum={fast_accum})" + if compile: + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + note += f" [compile FAILED {type(exc).__name__}]" + print(f" [{tag}]{note}", flush = True) + _gen(pipe, steps, seed, res) # warmup / compile + dts = [] + img = None + for _ in range(args.iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + gp = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, gp + + print( + f"== sparse/accum probe (Z-Image dense, {res}px, {steps} steps, min_feat={mf}) ==", + flush = True, + ) + rows = [] + # quality reference: dense bf16 eager (no compile, no quant) + bref, ref, _ = run("bf16_eager", compile = False) + rows.append(("bf16_eager", bref, float("inf"), 0.0, None)) + print(f" bf16 eager ref: {bref:.3f}s", flush = True) + + specs = [ + ("bf16_compile", dict()), + ("fp8_fastT_c", dict(quant = "fp8", fast_accum = True)), + ("fp8_fastF_c", dict(quant = "fp8", fast_accum = False)), + ( + "fake24_fp8_c", + dict(quant = "fp8", fast_accum = True, prune = True), + ), # quality of 2:4+fp8 (fake=no kernel) + ( + "real24_nocompile", + dict(real_sparse = True, compile = False), + ), # sparse SPEED (no quant, no compile) + ( + "real24_compile_try", + dict(real_sparse = True, compile = True), + ), # does sparse survive compile? + ] + for tag, kw in specs: + try: + med, arr, gp = run(tag, **kw) + ps, lp = _psnr(ref, arr), _lpips(ref, arr) + rows.append((tag, med, ps, lp, gp)) + print( + f" {tag:18s} {med:.3f}s ({bref/med:.2f}x vs eager) PSNR={ps:.1f} LPIPS={lp} VRAM={gp:.1f}G", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" {tag:18s} FAILED: {type(exc).__name__}: {str(exc)[:160]}", flush = True) + rows.append((tag, None, None, None, None)) + + print("\n==== SUMMARY (ref = bf16 dense eager) ====", flush = True) + base = next((r[1] for r in rows if r[0] == "fp8_fastT_c" and r[1]), None) + for tag, med, ps, lp, gp in rows: + if med is None: + print(f" {tag:18s} FAILED") + continue + vs_eager = f"{bref/med:.2f}x" + vs_fp8 = f"{base/med:.2f}x" if base else "-" + psv = "inf" if ps == float("inf") else f"{ps:.1f}" + lpv = ( + "ref" + if (lp == 0.0 and tag == "bf16_eager") + else (f"{lp:.3f}" if lp is not None else "n/a") + ) + print( + f" {tag:18s} {med:.3f}s eager:{vs_eager:>6s} fp8:{vs_fp8:>6s} PSNR={psv:>5s} LPIPS={lpv:>6s}", + flush = True, + ) + print("SPARSE-ACCUM-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 9500f2ed9f..3cac5629ac 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -45,6 +45,7 @@ from .diffusion_memory import ( snapshot_device_memory, ) from .diffusion_speed import ( + SPEED_DEFAULT, SPEED_OFF, apply_speed_optims, resolve_speed_mode, @@ -52,6 +53,11 @@ from .diffusion_speed import ( snapshot_backend_flags, ) from .diffusion_precision import quantize_text_encoders +from .diffusion_transformer_quant import ( + dense_transformer_supported, + normalize_transformer_quant, + quantize_transformer, +) logger = get_logger(__name__) @@ -81,6 +87,9 @@ class _LoadState: backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None + # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" + # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. + transformer_quant: Optional[str] = None @dataclass @@ -288,6 +297,8 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -319,6 +330,8 @@ class DiffusionBackend: memory_mode = memory_mode, speed_mode = speed_mode, text_encoder_quant = text_encoder_quant, + transformer_quant = transformer_quant, + transformer_quant_fast_accum = transformer_quant_fast_accum, _load_token = token, ), daemon = True, @@ -452,6 +465,8 @@ class DiffusionBackend: memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -484,98 +499,146 @@ class DiffusionBackend: # checkpoints never sit in VRAM at once. self._unload_locked() - # Dequantise the GGUF transformer on-device; the VAE / text-encoder / - # scheduler come from the base diffusers repo (GGUF is transformer-only). gguf_path = self._resolve_gguf_path(repo_id, gguf_filename, hf_token) transformer_cls = getattr(diffusers, fam.transformer_class) - transformer = transformer_cls.from_single_file( - gguf_path, - quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), - torch_dtype = dtype, - config = base, - subfolder = "transformer", - # Forward the token: the config is fetched from the (possibly gated) - # base repo before from_pretrained gets a chance to authenticate. - token = hf_token, + pipeline_cls = getattr(diffusers, fam.pipeline_class) + + # Decide placement up front (the weights are still on CPU, so free VRAM is + # the real budget) -- this also doubles as the dense-quant preflight: the + # dense bf16 transformer must fit resident, so the fast path is offered only + # when the plan is `none`. + plan = self._plan_memory( + target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload ) - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} - if hf_token: - pipe_kwargs["token"] = hf_token - pipeline_cls = getattr(diffusers, fam.pipeline_class) - pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it + # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul + # dequant on both speed and quality, at the cost of a higher-memory dense + # load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch + # / scheme, OOM, partial quant) falls back to the GGUF build below. + pipe = None + transformer_quant_engaged = None + if ( + normalize_transformer_quant(transformer_quant) is not None + and dense_transformer_supported(target) + and plan.offload_policy == OFFLOAD_NONE + ): + try: + pipe, transformer_quant_engaged = self._load_dense_quant_pipeline( + transformer_cls, + pipeline_cls, + base, + device, + dtype, + hf_token, + target, + transformer_quant, + transformer_quant_fast_accum, + ) + except Exception as exc: # noqa: BLE001 — fall back to the GGUF build + logger.warning( + "diffusion.transformer_quant_fallback: %s (loading GGUF)", exc + ) + pipe = None + transformer_quant_engaged = None + clear_gpu_cache() + + if pipe is None: + # Default: dequantise the single-file GGUF transformer on-device; the + # VAE / text-encoder / scheduler come from the base diffusers repo + # (GGUF is transformer-only). + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = base, + subfolder = "transformer", + # Forward the token: the config is fetched from the (possibly gated) + # base repo before from_pretrained gets a chance to authenticate. + token = hf_token, + ) + + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) # Resolve the effective speed mode: GGUF models default to the # near-lossless `default` profile (compile is ~2.2x and sits below # the quant noise floor), dense models stay bit-identical `off`. An # explicit speed_mode (incl. "off") is honored verbatim. effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) + # A torchao-quantized dense transformer runs its matmuls through the + # regional torch.compile; UNcompiled (eager) it is ~30x slower and would + # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so + # force at least `default` (regional compile) whenever the quant engaged, + # or the opt-in "fast" path silently commits an eager, pathologically slow + # pipeline. + if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: + logger.info( + "diffusion.transformer_quant: forcing speed_mode=default " + "(quantized transformer must be compiled; eager is ~30x slower)" + ) + effective_speed = SPEED_DEFAULT # Opt-in speed optims run BEFORE placement (channels_last / compile # must precede CPU offload). Snapshot the process-wide backend flags # first so unload can restore them: TF32 / cudnn.benchmark are global, # and a later `off` load must not inherit this load's settings. backend_flags_before = snapshot_backend_flags() - # apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark); - # they are only restored via _LoadState.backend_flags_before on unload. If - # the build fails after this but before _state commits (e.g. an OOM in - # apply_memory_plan / pipe.to), nothing would restore them and a later `off` - # generation would be contaminated, so restore on any non-committed exit. - committed = False - try: - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = effective_speed, - logger = logger, - ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + logger = logger, + ) + if transformer_quant_engaged is not None and not speed_applied.get("compiled"): + # Promotion above could not engage compile (e.g. the family is not + # compile-friendly, or compile_repeated_blocks failed): the quantized + # transformer is now running eager, which is far slower than the GGUF + # path it replaced. Surface it loudly rather than hiding the regression. + logger.warning( + "diffusion.transformer_quant: %s engaged but the transformer is NOT " + "compiled; eager torchao quant is ~30x slower than GGUF here", + transformer_quant_engaged, ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, + ) - # Decide placement from MEASURED free device memory vs the model's - # 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. `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. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Apply the placement planned above (from MEASURED free device memory vs + # the model's estimated resident size). 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. The dense fast path already placed the pipe resident; + # for the `none` policy this is an idempotent re-placement. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - speed_optims = tuple(k for k, v in speed_applied.items() if v), - backend_flags_before = backend_flags_before, - text_encoder_quant = te_quant, - ) - committed = True - finally: - if not committed: - restore_backend_flags(backend_flags_before) - clear_gpu_cache() + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + transformer_quant = transformer_quant_engaged, + ) logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -588,6 +651,39 @@ class DiffusionBackend: ) return self.status() + def _load_dense_quant_pipeline( + self, + transformer_cls: Any, + pipeline_cls: Any, + base: str, + device: str, + dtype: Any, + hf_token: Optional[str], + target: DiffusionDeviceTarget, + mode: Optional[str], + fast_accum: Optional[bool] = None, + ) -> tuple[Any, str]: + """Build the opt-in fast pipeline: load the DENSE bf16 transformer from the base + repo (``subfolder="transformer"``), assemble the pipeline, place it on the device, + and torchao-quantise the transformer in place. Returns ``(pipe, engaged_scheme)``. + + Raises if the scheme is unsupported or quantisation fails, so ``load_pipeline`` + catches it and falls back to the GGUF build. Quantisation runs ON the device (the + dynamic int8 / fp8 / fp4 kernels need the weights on CUDA) and BEFORE the loader + compiles the repeated block, so the order is quantize -> compile -> placement.""" + transformer = transformer_cls.from_pretrained( + base, subfolder = "transformer", torch_dtype = dtype, token = hf_token + ) + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + if hf_token: + pipe_kwargs["token"] = hf_token + pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + pipe.to(device) + scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger) + if scheme is None: + raise RuntimeError("transformer quant unsupported for this device/scheme") + return pipe, scheme + def _plan_memory( self, target: DiffusionDeviceTarget, @@ -796,6 +892,7 @@ class DiffusionBackend: "speed_mode": None, "speed_optims": [], "text_encoder_quant": None, + "transformer_quant": None, } return { "loaded": True, @@ -811,6 +908,7 @@ class DiffusionBackend: "speed_mode": state.speed_mode, "speed_optims": list(state.speed_optims), "text_encoder_quant": state.text_encoder_quant, + "transformer_quant": state.transformer_quant, } diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py new file mode 100644 index 0000000000..7cb1074425 --- /dev/null +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -0,0 +1,347 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in low-precision quantisation of the diffusion DiT transformer. + +The default path loads the transformer as a single-file GGUF, which stores weights +4-bit but DEQUANTISES to bf16 on every matmul -- so it runs at bf16 tensor-core rate +and never touches the int8 / fp8 / fp4 tensor cores. It is a memory win that costs +speed. This module is the opt-in alternative: load the DENSE bf16 transformer from the +base repo and torchao-quantise it with a DYNAMIC-ACTIVATION scheme so the matmul runs +on the low-precision tensor cores. Measured on a B200 (Z-Image-Turbo, 1024px / 8 steps) +vs the GGUF+compile default (0.802s, LPIPS 0.083 vs dense bf16): fp8 dynamic 0.585s +(1.37x), int8 dynamic 0.603s (1.33x), both at LOWER LPIPS than GGUF -- faster AND a hair +more accurate, at the cost of a higher-memory dense load. So it is strictly opt-in; the +loader keeps GGUF as the low-memory default and the fallback. + +Scheme by architecture (``auto`` picks the best supported, best first): + nvfp4 / mxfp8 - Blackwell sm_100+ FP4 / MX tensor cores (biggest win; prototype). + fp8 - Ada / Hopper / Blackwell (sm_89+) fp8 tensor cores. + int8 - Ampere+ (sm_80+) int8 tensor cores -- the broadest-hardware lever. + +Every scheme needs ``torch.compile`` to realise the speedup (dynamic quant is ~30x +slower eager); the loader already compiles the repeated block AFTER this runs. torch / +torchao are imported lazily so the module stays importable in a no-torch runtime, and +every probe is best-effort: an unsupported scheme yields None and the caller loads GGUF. +""" + +from __future__ import annotations + +from typing import Any, Optional + +TQ_INT8 = "int8" +TQ_FP8 = "fp8" +TQ_NVFP4 = "nvfp4" +TQ_MXFP8 = "mxfp8" +TQ_AUTO = "auto" +TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8) +TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES + +# Skip linears whose in/out features are below this. The int8 dynamic path uses +# torch._int_mm, which requires the activation row count M > 16, and the DiT's tiny +# timestep / pooled / modulation projections run at M=1 and crash it. They are a +# negligible share of the FLOPs, so leaving them bf16 costs ~nothing (measured: +# 239/276 Z-Image linears quantised, full speedup) and keeps quality a touch higher. +DEFAULT_MIN_LINEAR_FEATURES = 512 + +# Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with +# the lower-precision schemes listed as fallbacks for that arch tier. On Blackwell, fp8 +# leads: measured on a B200, plain fp8 dynamic is both faster AND more accurate than the +# alternatives for the DiT's shapes. mxfp8's block scaling adds overhead without a speed +# win, so it sits below fp8. nvfp4 is intentionally below fp8 too: the FP4 tensor-core +# GEMM is real once torch>=2.11 + torchao's CUTLASS FP4 kernel is present (verified: a +# 16384^3 GEMM hits ~3826 TFLOPS, 1.37x fp8), but it only beats fp8 on very large GEMMs. +# At the DiT's actual shapes (hidden ~3072, MLP ~12288, M~4096) it is *slower* than fp8 +# (0.81x end-to-end on Z-Image 1024px) AND notably less accurate (LPIPS 0.166 vs fp8's +# 0.044), because FP4's per-forward quant overhead is not amortised and the format is +# coarser. So nvfp4 is kept as an explicit opt-in, never the auto pick for diffusion. +_AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( + ((10, 0), (TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8)), # Blackwell sm_100+ + ((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90 + ((8, 0), (TQ_INT8,)), # Ampere sm_80 / sm_86 +) + +# Cache of (scheme, device) -> bool so the quantise+matmul smoke test runs once. +_SMOKE_CACHE: dict[tuple[str, str], bool] = {} + +# Data-center GPU model tokens (un-nerfed FP32 accumulate). Matched as whole tokens of +# torch.cuda.get_device_name(), so the workstation "A4000" is not mistaken for the +# data-center "A40". Anything not here -- GeForce, workstation RTX, or an unknown name -- +# is treated as consumer-class (FP32-accumulate halved). See developer.nvidia.com/cuda/gpus. +_DATACENTER_GPU_TOKENS = frozenset( + { + "B200", + "B100", + "GB200", + "GB300", + "GB10", # Blackwell data center + "H200", + "H100", + "H800", + "H20", # Hopper data center + "A100", + "A800", + "A30", + "A40", + "A16", + "A10", + "A2", # Ampere data center + "L40", + "L40S", + "L4", + "L20", + "L2", # Ada data center + "V100", + "P100", + "P40", + "T4", # legacy data center + } +) + + +def _is_consumer_gpu(device: Any = None) -> bool: + """Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32 + accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center + HBM parts (recognised by name token) are not nerfed and return False, so they keep + the higher-precision default accumulate for free. Heuristic on the device name: a + GeForce / TITAN name is always consumer; a recognised data-center token is not; + anything else (workstation RTX, unknown) defaults to consumer -- the safe choice, + since fast accumulate is free on data-center and a win on consumer. Best-effort: + True on any probe failure.""" + try: + import re + + import torch + name = torch.cuda.get_device_name(device).upper() + except Exception: # noqa: BLE001 — no torch / no device -> assume consumer + return True + if "GEFORCE" in name or "TITAN" in name: + return True + tokens = set(re.split(r"[^A-Z0-9]+", name)) + return not (tokens & _DATACENTER_GPU_TOKENS) + + +def normalize_transformer_quant(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested transformer quant; None / "" / "none" / "off" -> None. + + 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("-", "_") + if not normalized or normalized in ("none", "off"): + return None + if normalized not in TQ_MODES: + raise ValueError( + f"Unsupported transformer_quant '{value}'. Use one of: {', '.join(TQ_MODES)}." + ) + return normalized + + +def dense_transformer_supported(target: Any) -> bool: + """Whether the dense-source quant path is usable for ``target``: a CUDA device with + a bf16 compute dtype (the only configuration any torchao dynamic scheme accelerates). + A cheap pre-check the loader runs before loading the (large) dense transformer.""" + if getattr(target, "device", None) != "cuda": + return False + try: + import torch + return getattr(target, "dtype", None) is torch.bfloat16 + except Exception: + return False + + +def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Optional[str]: + """The concrete scheme to apply, or None to fall back to GGUF. + + ``auto`` walks the per-arch ladder and returns the first scheme that passes a real + quantise+matmul smoke test, so on a box where the Blackwell fp4 / mx kernels are + unavailable it lands on fp8 / int8 with no error. An explicit scheme is honored only + if supported (else None -> GGUF), never silently swapped for a different one.""" + requested = normalize_transformer_quant(requested) + if requested is None or not dense_transformer_supported(target): + return None + device = str(getattr(target, "device", "cuda")) + if requested != TQ_AUTO: + return requested if _scheme_supported(requested, device) else None + cap = _capability() + if cap is None: + return None + for floor, schemes in _AUTO_LADDER: + if cap >= floor: + for scheme in schemes: + if _scheme_supported(scheme, device): + return scheme + return None + return None + + +def _capability() -> Optional[tuple[int, int]]: + try: + import torch + major, minor = torch.cuda.get_device_capability() + return (int(major), int(minor)) + except Exception: + return None + + +def _scheme_supported(scheme: str, device: str) -> bool: + """CUDA + (for fp8) the fp8 dtype + a cached quantise+matmul smoke test for ``scheme``.""" + try: + import torch + if not torch.cuda.is_available(): + return False + if scheme == TQ_FP8 and not hasattr(torch, "float8_e4m3fn"): + return False + except Exception: + return False + return _smoke_probe(scheme, device) + + +def _smoke_probe(scheme: str, device: str) -> bool: + """True iff a tiny Linear quantised with ``scheme`` runs one M=32 forward without + error. Cached per (scheme, device). This is what makes ``auto`` robust to a torch / + torchao build where a prototype (nvfp4 / mxfp8) kernel is unavailable: it fails here + and the ladder moves on, rather than crashing at the first real denoise step.""" + key = (scheme, device) + if key in _SMOKE_CACHE: + return _SMOKE_CACHE[key] + ok = False + try: + import torch + from torchao.quantization import quantize_ + + lin = torch.nn.Linear(512, 512, bias = False).to(device = device, dtype = torch.bfloat16) + quantize_(lin, _make_quant_config(scheme), filter_fn = make_filter_fn(0)) + x = torch.randn(32, 512, device = device, dtype = torch.bfloat16) + with torch.no_grad(): + lin(x) + torch.cuda.synchronize() + ok = True + except Exception: + ok = False + _SMOKE_CACHE[key] = ok + return ok + + +def _resolve_fast_accum(fast_accum: Optional[bool]) -> bool: + """The fp8 ``use_fast_accum`` to apply. ``None`` auto-detects by GPU class + (consumer / workstation -> fast; data-center -> precise); an explicit bool forces it.""" + return _is_consumer_gpu() if fast_accum is None else bool(fast_accum) + + +def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: + """The torchao dynamic-activation config for ``scheme`` (lazy import; prototype + import for the Blackwell fp4 / mx schemes is inside the branch that needs it). + + ``fast_accum`` applies to fp8 only: None auto-detects by GPU class, True/False force it.""" + from torchao.quantization import ( + Float8DynamicActivationFloat8WeightConfig, + Int8DynamicActivationInt8WeightConfig, + ) + + if scheme == TQ_INT8: + return Int8DynamicActivationInt8WeightConfig() + if scheme == TQ_FP8: + # Choose fp8 accumulate by GPU class (unless forced). On consumer / workstation + # cards (GDDR) the fp8 tensor cores run ~2x faster with FP16 (fast) accumulate + # than FP32 (e.g. ~838 vs ~419 TFLOPS on RTX 50xx), so fast accumulate is a real + # win there. Data-center HBM parts default to the higher-precision accumulate. + # fast accumulate is a precision (not overflow) tradeoff and stays below the fp8 + # quant noise floor (measured 0 non-finite even on Z-Image's ~1e6 activations). + try: + from torchao.float8 import Float8MMConfig + return Float8DynamicActivationFloat8WeightConfig( + mm_config = Float8MMConfig(use_fast_accum = _resolve_fast_accum(fast_accum)) + ) + except Exception: # noqa: BLE001 — older torchao without the explicit knob + return Float8DynamicActivationFloat8WeightConfig() + if scheme == TQ_NVFP4: + from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig + + # Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults + # use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the + # CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail + # and silently fall back to GGUF instead of using the FP4 tensor cores. + try: + return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False) + except TypeError: # older torchao without the knob + return NVFP4DynamicActivationNVFP4WeightConfig() + if scheme == TQ_MXFP8: + import torch + from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig + try: + return MXDynamicActivationMXWeightConfig( + activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn + ) + except (TypeError, AttributeError): + # TypeError: older torchao without the explicit dtype knobs. + # AttributeError: a torch build without torch.float8_e4m3fn. + return MXDynamicActivationMXWeightConfig() + raise ValueError(f"unknown transformer quant scheme '{scheme}'") + + +def make_filter_fn(min_features: int): + """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear + with both in/out features >= ``min_features``. Hides the (module, fqn) callback arity.""" + + def filter_fn(module: Any, fqn: str = "") -> bool: + try: + import torch + if not isinstance(module, torch.nn.Linear): + return False + except Exception: + return False + in_features = getattr(module, "in_features", None) + out_features = getattr(module, "out_features", None) + if in_features is None or out_features is None: + return False + return in_features >= min_features and out_features >= min_features + + return filter_fn + + +def quantize_transformer( + pipe: Any, + target: Any, + *, + mode: Optional[str], + min_features: int = DEFAULT_MIN_LINEAR_FEATURES, + fast_accum: Optional[bool] = None, + logger: Any = None, +) -> Optional[str]: + """Quantise ``pipe.transformer``'s FLOP-heavy linears in place with the arch-chosen + dynamic scheme. Returns the scheme actually engaged, or None when disabled / + unsupported / failed -- the caller then loads GGUF instead. Best-effort: it never + raises for an ordinary unsupported environment (a failure leaves the module dense). + + ``fast_accum`` (fp8 only) overrides the per-GPU-class accumulate choice: None + auto-detects (fast on consumer, precise on data-center), True/False force it.""" + scheme = select_transformer_quant_scheme(target, mode) + if scheme is None: + return None + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return None + try: + from torchao.quantization import quantize_ + + quantize_( + transformer, + _make_quant_config(scheme, fast_accum = fast_accum), + filter_fn = make_filter_fn(min_features), + ) + # Runtime-only marker (torchao tensors are not safetensors-serializable; this + # backend is inference-only, so this is purely diagnostic). + try: + transformer._unsloth_runtime_quant = scheme + except Exception: # noqa: BLE001 — marker is best-effort + pass + return scheme + except Exception as exc: # noqa: BLE001 — leave the transformer dense -> GGUF fallback + _warn(logger, scheme, exc) + return None + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.transformer_quant: %s failed: %s", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index fff68c090a..56ee571dae 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1720,6 +1720,23 @@ class DiffusionLoadRequest(BaseModel): "memory-vs-quality tradeoff (shifts fine detail), not free; " "pairs well with balanced mode.", ) + transformer_quant: Optional[Literal["auto", "int8", "fp8", "nvfp4", "mxfp8"]] = Field( + None, + description = "Opt-in fast transformer: load the DENSE bf16 transformer instead " + "of the GGUF and torchao-quantise it onto the low-precision tensor " + "cores (faster than GGUF's bf16-rate dequant, at higher VRAM). auto " + "picks the best for the GPU (Blackwell nvfp4/mxfp8, Ada/Hopper fp8, " + "Ampere int8); an explicit scheme forces it. Needs CUDA + bf16 + room " + "for the dense load; falls back to GGUF otherwise.", + ) + transformer_quant_fast_accum: Optional[bool] = Field( + None, + description = "fp8 only: FP8 matmul accumulate. null auto-detects by GPU class " + "(fast FP16 accumulate on consumer/workstation cards, where FP32 " + "accumulate is ~2x slower; precise FP32 accumulate on data-center " + "HBM cards, which are not nerfed). true/false force it. Negligible " + "quality effect (below the fp8 quant noise floor); no overflow risk.", + ) class DiffusionGenerateRequest(BaseModel): @@ -1830,3 +1847,8 @@ class DiffusionStatusResponse(BaseModel): text_encoder_quant: Optional[str] = Field( None, description = "Text-encoder quantisation engaged: fp8 | nvfp4 | null" ) + transformer_quant: Optional[str] = Field( + None, + description = "Transformer quant engaged on the dense fast path: int8 | fp8 | " + "nvfp4 | mxfp8 | null (null = the GGUF transformer was loaded)", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1aa4be56b3..13beee00e0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10331,6 +10331,8 @@ async def load_diffusion_model( memory_mode = request.memory_mode, speed_mode = request.speed_mode, text_encoder_quant = request.text_encoder_quant, + transformer_quant = request.transformer_quant, + transformer_quant_fast_accum = request.transformer_quant_fast_accum, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index fe4cabe8b5..2e714b217b 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -899,3 +899,131 @@ 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" + + +# ── transformer quant (opt-in dense fast path) ──────────────────────────────── + + +def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): + """Force the dense+quant branch hermetically: a supported dense source, a + from_pretrained on the fake transformer, and a quantizer that engages `scheme`. + Returns a dict recording the dense-loader / quantizer calls.""" + from core.inference import diffusion as dmod + + calls: dict = {"from_pretrained": 0, "quantize": 0, "quant_mode": None} + + @classmethod + def _from_pretrained(cls, base, **kwargs): + calls["from_pretrained"] += 1 + calls["fp_kwargs"] = {"base": base, **kwargs} + return object() + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + def _quantize(pipe, target, *, mode, **kw): + calls["quantize"] += 1 + calls["quant_mode"] = mode + return scheme + + monkeypatch.setattr(dmod, "quantize_transformer", _quantize) + return calls + + +def test_default_load_skips_dense_quant_path(fake_runtime, tmp_path, monkeypatch): + # With no transformer_quant flag the GGUF path is taken and the dense gate is + # never even consulted (short-circuit), so the default cannot regress. + from core.inference import diffusion as dmod + + monkeypatch.setattr( + dmod, + "dense_transformer_supported", + lambda *a, **k: pytest.fail("dense path must not run without the flag"), + ) + (tmp_path / "m.gguf").write_bytes(b"x") + backend = DiffusionBackend() + status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") + assert status["transformer_quant"] is None + assert _FakeTransformer.last["path"] # GGUF from_single_file was used + + +def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatch): + # transformer_quant + a CUDA resident plan -> load the DENSE transformer from the + # base repo, place it on the device, quantise it, and report the engaged scheme. + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + calls = _stub_dense_quant(monkeypatch, scheme = "fp8") + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert status["transformer_quant"] == "fp8" + # No speed_mode was given, but a quantized transformer is ~30x slower eager, so the + # backend promotes it to `default` (regional compile) instead of the dense `off`. + assert status["speed_mode"] == "default" + assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 + assert calls["quant_mode"] == "fp8" + assert calls["fp_kwargs"]["subfolder"] == "transformer" # dense transformer subfolder + # The GGUF single-file path was NOT used for the transformer. + assert _FakeTransformer.last == {} + # quantize ran on-device: the dense pipe was placed on cuda (before compile). + assert backend._state.pipe.moved_to == "cuda" + assert status["offload_policy"] == "none" + + +def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, monkeypatch): + # A dense/quant failure (here: quantize returns None -> unsupported) must fall back + # to the GGUF build, not error -- status reports no transformer_quant engaged. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + @classmethod + def _from_pretrained(cls, base, **kwargs): + return object() + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) + monkeypatch.setattr(dmod, "quantize_transformer", lambda pipe, target, **kw: None) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert status["loaded"] is True + assert status["transformer_quant"] is None # fell back + assert _FakeTransformer.last["path"] # GGUF from_single_file used + + +def test_transformer_quant_skipped_when_plan_offloads(fake_runtime, tmp_path, monkeypatch): + # The dense bf16 transformer only fits resident, so when the memory plan would + # offload (here low_vram) the fast path is skipped and GGUF loads instead -- the + # dense transformer is never even loaded. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + + @classmethod + def _fp_fail(cls, *a, **k): + pytest.fail("dense transformer must not load when the plan offloads") + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + memory_mode = "low_vram", + ) + assert status["transformer_quant"] is None + assert status["offload_policy"] == "model" + assert _FakeTransformer.last["path"] # GGUF path used diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index b0fb3836e6..137b218636 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -381,6 +381,44 @@ def test_memory_mode_threads_through_to_backend(client, monkeypatch): assert backend.last_load_kwargs.get("memory_mode") == "low_vram" +def test_transformer_quant_threads_through_to_backend(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "auto"}, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_quant") == "auto" + + +def test_transformer_quant_fast_accum_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_quant": "fp8", + "transformer_quant_fast_accum": False, + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False + + +def test_invalid_transformer_quant_returns_422_without_eviction(client): + # An unsupported transformer_quant is rejected by the request schema (Literal), so + # the GPU is never acquired and no chat model is evicted. + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_quant": "int2"}, + ) + assert resp.status_code == 422 + assert gpu_arbiter._owner is None + + def test_invalid_memory_mode_returns_422_without_eviction(client): # An unsupported memory_mode is rejected by the request schema (Literal), so the # GPU is never acquired and no chat model is evicted. diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py new file mode 100644 index 0000000000..d2adaa9f7c --- /dev/null +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -0,0 +1,342 @@ +# 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 transformer quantisation (``diffusion_transformer_quant.py``). + +Hermetic: torch + torchao are stubbed via ``sys.modules``, and the per-scheme smoke +probe (``_scheme_supported`` / ``_smoke_probe``) is monkeypatched where the test cares +about the selection ladder rather than the GPU probe, so everything runs CPU-only. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +import core.inference.diffusion_transformer_quant as tq +from core.inference.diffusion_transformer_quant import ( + TQ_FP8, + TQ_INT8, + TQ_MXFP8, + TQ_NVFP4, + dense_transformer_supported, + make_filter_fn, + normalize_transformer_quant, + quantize_transformer, + select_transformer_quant_scheme, +) + + +def _target(*, device = "cuda", dtype = "bfloat16"): + return types.SimpleNamespace(device = device, dtype = dtype) + + +def _stub_torch( + monkeypatch, + *, + cc = (10, 0), + with_fp8 = True, + cuda_available = True, +): + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.float16 = "float16" + if with_fp8: + torch.float8_e4m3fn = "float8_e4m3fn" + torch.cuda = types.SimpleNamespace( + is_available = lambda: cuda_available, + get_device_capability = lambda *a: cc, + ) + monkeypatch.setitem(sys.modules, "torch", torch) + return torch + + +# ── normalisation ───────────────────────────────────────────────────────────── + + +def test_normalize_transformer_quant(): + assert normalize_transformer_quant(None) is None + assert normalize_transformer_quant("") is None + assert normalize_transformer_quant("none") is None + assert normalize_transformer_quant("off") is None + assert normalize_transformer_quant("AUTO") == "auto" + assert normalize_transformer_quant("INT8") == TQ_INT8 + assert normalize_transformer_quant("fp8") == TQ_FP8 + with pytest.raises(ValueError): + normalize_transformer_quant("int2") + + +# ── dense-source gate ─────────────────────────────────────────────────────────── + + +def test_dense_transformer_supported_requires_cuda_bf16(monkeypatch): + _stub_torch(monkeypatch) + assert dense_transformer_supported(_target()) is True + assert dense_transformer_supported(_target(device = "cpu")) is False + assert dense_transformer_supported(_target(dtype = "float16")) is False + + +# ── scheme selection ladder ───────────────────────────────────────────────────── + + +def _allow(monkeypatch, allowed): + """Force ``_scheme_supported`` to accept only ``allowed`` (simulates smoke results).""" + monkeypatch.setattr(tq, "_scheme_supported", lambda scheme, device: scheme in allowed) + + +def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch): + _stub_torch(monkeypatch, cc = (10, 0)) + # Even with every scheme available, auto picks fp8 on Blackwell: measured on a B200 + # (torch 2.11 + torchao CUTLASS FP4), fp8 is both faster and more accurate than nvfp4 + # for the DiT's shapes -- nvfp4's FP4 GEMM only wins on very large GEMMs, not here. + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + # fp8 unavailable: nvfp4 is the next pick (above mxfp8 / int8). + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_NVFP4 + # Only mxfp8 + int8 left -> mxfp8 (still above int8). + _allow(monkeypatch, {TQ_MXFP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_MXFP8 + # Only int8 usable -> int8. + _allow(monkeypatch, {TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_ada_hopper_prefers_fp8(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 9)) + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + _stub_torch(monkeypatch, cc = (9, 0)) # Hopper + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + +def test_auto_ampere_prefers_int8(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 0)) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) # fp8 cores absent on Ampere -> int8 only in ladder + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + _stub_torch(monkeypatch, cc = (8, 6)) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_pre_ampere_unsupported(monkeypatch): + _stub_torch(monkeypatch, cc = (7, 5)) # Turing: below the int8-dynamic floor + _allow(monkeypatch, {TQ_INT8, TQ_FP8}) + assert select_transformer_quant_scheme(_target(), "auto") is None + + +def test_explicit_scheme_honored_or_none(monkeypatch): + _stub_torch(monkeypatch, cc = (8, 0)) + _allow(monkeypatch, {TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "int8") == TQ_INT8 + # Explicit unsupported scheme is NOT silently downgraded -> None (-> GGUF fallback). + assert select_transformer_quant_scheme(_target(), "fp8") is None + assert select_transformer_quant_scheme(_target(), "nvfp4") is None + + +def test_select_none_when_disabled_or_non_cuda(monkeypatch): + _stub_torch(monkeypatch) + _allow(monkeypatch, {TQ_INT8, TQ_FP8, TQ_NVFP4}) + assert select_transformer_quant_scheme(_target(), None) is None + assert select_transformer_quant_scheme(_target(device = "cpu"), "auto") is None + + +# ── _scheme_supported / _smoke_probe ──────────────────────────────────────────── + + +def test_scheme_supported_shortcircuits(monkeypatch): + # No CUDA -> False without running the smoke probe. + _stub_torch(monkeypatch, cuda_available = False) + monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run")) + assert tq._scheme_supported(TQ_INT8, "cuda") is False + # fp8 requested but the fp8 dtype is missing -> False before the probe. + _stub_torch(monkeypatch, with_fp8 = False) + monkeypatch.setattr(tq, "_smoke_probe", lambda *a: pytest.fail("probe should not run")) + assert tq._scheme_supported(TQ_FP8, "cuda") is False + + +def test_smoke_probe_caches_and_tolerates_failure(monkeypatch): + tq._SMOKE_CACHE.clear() + calls = {"n": 0} + + class _Lin: + def __init__(self, *a, **k): + pass + + def to(self, **k): + return self + + torch = types.ModuleType("torch") + torch.bfloat16 = "bfloat16" + torch.nn = types.SimpleNamespace(Linear = _Lin) + torch.randn = lambda *a, **k: object() + torch.no_grad = lambda: __import__("contextlib").nullcontext() + torch.cuda = types.SimpleNamespace(is_available = lambda: True, synchronize = lambda: None) + monkeypatch.setitem(sys.modules, "torch", torch) + + tqz = types.ModuleType("torchao.quantization") + + def _quantize_ok( + module, + config, + filter_fn = None, + ): + calls["n"] += 1 + + tqz.quantize_ = _quantize_ok + tqz.Int8DynamicActivationInt8WeightConfig = lambda: "int8cfg" + tqz.Float8DynamicActivationFloat8WeightConfig = lambda: "fp8cfg" + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + # _Lin is callable? No -> the forward lin(x) would fail. Make instances callable. + _Lin.__call__ = lambda self, x: x + + assert tq._smoke_probe(TQ_INT8, "cuda") is True + assert tq._smoke_probe(TQ_INT8, "cuda") is True # cached, no second quantize_ + assert calls["n"] == 1 + + # A scheme whose quantize_ raises -> probe False (and cached). + tq._SMOKE_CACHE.clear() + + def _quantize_boom( + module, + config, + filter_fn = None, + ): + raise RuntimeError("kernel unavailable") + + tqz.quantize_ = _quantize_boom + assert tq._smoke_probe(TQ_FP8, "cuda") is False + + +# ── consumer-vs-datacenter detection (fp8 fast-accumulate gate) ────────────────── + + +def _stub_device_name(monkeypatch, name): + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(get_device_name = lambda device = None: name) + monkeypatch.setitem(sys.modules, "torch", torch) + + +@pytest.mark.parametrize( + "name", + [ + "NVIDIA GeForce RTX 5090", + "NVIDIA GeForce RTX 4090", + "NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40 + "NVIDIA RTX 6000 Ada Generation", + "NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC) + ], +) +def test_is_consumer_gpu_true(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is True + + +@pytest.mark.parametrize( + "name", + [ + "NVIDIA B200", + "NVIDIA H100 80GB HBM3", + "NVIDIA A100-SXM4-80GB", + "NVIDIA A40", # data-center Ampere (distinct token from RTX A4000) + "NVIDIA L40S", + "NVIDIA L4", + "Tesla V100-SXM2-16GB", + ], +) +def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name): + _stub_device_name(monkeypatch, name) + assert tq._is_consumer_gpu() is False + + +def test_is_consumer_gpu_defaults_true_on_probe_failure(monkeypatch): + # No torch / no device name available -> assume consumer (safe: fast accum is free + # on data center and a win on consumer). + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace() # no get_device_name + monkeypatch.setitem(sys.modules, "torch", torch) + assert tq._is_consumer_gpu() is True + + +# ── filter ────────────────────────────────────────────────────────────────────── + + +def test_make_filter_fn(monkeypatch): + class _Lin: + def __init__(self, i, o): + self.in_features, self.out_features = i, o + + torch = types.ModuleType("torch") + torch.nn = types.SimpleNamespace(Linear = _Lin) + monkeypatch.setitem(sys.modules, "torch", torch) + + keep = make_filter_fn(512) + assert keep(_Lin(1024, 4096), "blocks.0.attn.to_q") is True + assert keep(_Lin(256, 4096), "time_proj") is False # small in_features -> skip + assert keep(_Lin(4096, 256), "out_proj") is False # small out_features -> skip + assert keep(object(), "not_linear") is False # non-Linear -> skip + assert keep(types.SimpleNamespace(), "no_attrs") is False + + +# ── apply ─────────────────────────────────────────────────────────────────────── + + +def test_resolve_fast_accum(monkeypatch): + # None auto-detects by GPU class; an explicit bool forces it. + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: True) + assert tq._resolve_fast_accum(None) is True + monkeypatch.setattr(tq, "_is_consumer_gpu", lambda *a: False) + assert tq._resolve_fast_accum(None) is False + assert tq._resolve_fast_accum(True) is True # forced on (e.g. on a data-center card) + assert tq._resolve_fast_accum(False) is False # forced off (e.g. on a consumer card) + + +def test_quantize_transformer_applies_and_marks(monkeypatch): + monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_FP8) + seen: dict = {} + + def _mk(scheme, fast_accum = None): + seen["scheme"], seen["fast_accum"] = scheme, fast_accum + return f"{scheme}cfg" + + monkeypatch.setattr(tq, "_make_quant_config", _mk) + recorder: list = [] + tqz = types.ModuleType("torchao.quantization") + tqz.quantize_ = lambda module, config, filter_fn = None: recorder.append( + (module, config, filter_fn) + ) + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + + transformer = types.SimpleNamespace() + pipe = types.SimpleNamespace(transformer = transformer) + assert quantize_transformer(pipe, _target(), mode = "fp8", fast_accum = False) == TQ_FP8 + assert len(recorder) == 1 and recorder[0][0] is transformer and recorder[0][1] == "fp8cfg" + assert callable(recorder[0][2]) # a filter_fn was passed + assert transformer._unsloth_runtime_quant == TQ_FP8 # diagnostic marker set + assert seen["fast_accum"] is False # the override is forwarded into the config + + +def test_quantize_transformer_none_when_unsupported(monkeypatch): + monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: None) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert quantize_transformer(pipe, _target(), mode = "auto") is None + + +def test_quantize_transformer_tolerates_failure(monkeypatch): + monkeypatch.setattr(tq, "select_transformer_quant_scheme", lambda target, mode: TQ_INT8) + monkeypatch.setattr(tq, "_make_quant_config", lambda scheme: "cfg") + tqz = types.ModuleType("torchao.quantization") + + def _boom( + module, + config, + filter_fn = None, + ): + raise RuntimeError("partial quant failure") + + tqz.quantize_ = _boom + monkeypatch.setitem(sys.modules, "torchao.quantization", tqz) + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + # A quantise failure returns None (caller falls back to GGUF), never raises. + assert quantize_transformer(pipe, _target(), mode = "int8") is None From ddb6084ac4f272227b19d74308f8bc91024130cd Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:37:53 -0700 Subject: [PATCH 05/16] Studio diffusion (Phase 9): pre-quantized transformer loading (#6700) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/build_prequant_checkpoint.py | 132 +++++++ scripts/prequant_probe.py | 196 ++++++++++ scripts/verify_prequant_backend.py | 152 ++++++++ studio/backend/core/inference/diffusion.py | 80 +++- .../core/inference/diffusion_families.py | 14 + .../core/inference/diffusion_prequant.py | 287 +++++++++++++++ studio/backend/models/inference.py | 13 + studio/backend/routes/inference.py | 1 + .../backend/tests/test_diffusion_backend.py | 75 ++++ .../backend/tests/test_diffusion_prequant.py | 348 ++++++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 29 ++ 11 files changed, 1317 insertions(+), 10 deletions(-) create mode 100644 scripts/build_prequant_checkpoint.py create mode 100644 scripts/prequant_probe.py create mode 100644 scripts/verify_prequant_backend.py create mode 100644 studio/backend/core/inference/diffusion_prequant.py create mode 100644 studio/backend/tests/test_diffusion_prequant.py diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py new file mode 100644 index 0000000000..18dd064650 --- /dev/null +++ b/scripts/build_prequant_checkpoint.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Build a pre-quantized transformer checkpoint for the Studio diffusion fast path. + +Quantise a model's dense bf16 DiT transformer ONCE and save the quantized state dict, so +the backend can load the already-quantized weights at runtime (meta-init + +load_state_dict(assign=True)) instead of materialising the dense bf16 on the GPU. That +drops the transformer GPU load peak ~2x and the download ~2x for fp8 (measured on Z-Image: +12.9 -> 6.3 GB peak, 12 -> 6.28 GB on disk), with bit-identical output -- it is the exact +same torchao config + min_features filter the runtime path uses, applied ahead of time. + +Run on one CUDA (Blackwell / Ada / Hopper) GPU. fp8 works on torch 2.9+; the FP4/MX schemes +need the newer kernels (see scripts/nvfp4_t211_probe.py). + + python scripts/build_prequant_checkpoint.py \ + --base Tongyi-MAI/Z-Image-Turbo --family z-image --scheme fp8 \ + --out outputs/quant_research/prequant_fp8/transformer_fp8.pt [--upload-repo ORG/REPO] +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend" + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument( + "--base", required = True, help = "diffusers base repo (carries the transformer subfolder)" + ) + p.add_argument("--family", required = True, help = "diffusion family name/alias (e.g. z-image)") + p.add_argument("--scheme", required = True, help = "quant scheme: int8 | fp8 | nvfp4 | mxfp8") + p.add_argument("--out", required = True, help = "output .pt path for the checkpoint") + p.add_argument("--min-features", type = int, default = 512) + p.add_argument("--dtype", default = "bfloat16", choices = ["bfloat16"]) + p.add_argument("--hf-token", default = None) + p.add_argument( + "--upload-repo", default = None, help = "optional HF repo id to upload the checkpoint to" + ) + p.add_argument("--upload-revision", default = None) + args = p.parse_args(argv) + + sys.path.insert(0, str(BACKEND)) + import torch + import torchao + import diffusers + + from core.inference.diffusion_families import detect_family + from core.inference.diffusion_prequant import PREQUANT_FORMAT, prequant_filename + + # Reuse the runtime quant factory + filter so offline == runtime (the LPIPS-0 invariant). + from core.inference.diffusion_transformer_quant import ( + TQ_SCHEMES, + _make_quant_config, + make_filter_fn, + ) + from torchao.quantization import quantize_ + + scheme = args.scheme.strip().lower() + if scheme not in TQ_SCHEMES: + print(f"error: --scheme must be one of {TQ_SCHEMES} (not 'auto')", flush = True) + return 2 + fam = detect_family(args.base, override = args.family) + if fam is None: + print(f"error: unknown family '{args.family}'", flush = True) + return 2 + transformer_cls = getattr(diffusers, fam.transformer_class) + + print(f"== build prequant ({fam.name}/{scheme}, min_feat={args.min_features}) ==", flush = True) + print(f" loading dense transformer from {args.base} (subfolder=transformer) ...", flush = True) + t0 = time.time() + transformer = transformer_cls.from_pretrained( + args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token + ).to("cuda") + print(f" quantising in place ({scheme}) ...", flush = True) + quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features)) + + # Move the state dict to CPU for a portable, GPU-free artifact. + state_dict = { + k: (v.detach().to("cpu") if hasattr(v, "detach") else v) + for k, v in transformer.state_dict().items() + } + ckpt = { + "format": PREQUANT_FORMAT, + "metadata": { + "base_model_id": args.base, + "family": fam.name, + "scheme": scheme, + "min_features": args.min_features, + "torch_dtype": args.dtype, + "quant_backend": "torchao", + "transformer_class": fam.transformer_class, + "torch_version": torch.__version__, + "torchao_version": getattr(torchao, "__version__", "?"), + "diffusers_version": diffusers.__version__, + }, + "state_dict": state_dict, + } + + out = Path(args.out) + out.parent.mkdir(parents = True, exist_ok = True) + torch.save(ckpt, out) + size_gb = out.stat().st_size / 1e9 + print(f" saved {out} ({size_gb:.2f} GB) in {time.time() - t0:.0f}s", flush = True) + print(f" metadata: {ckpt['metadata']}", flush = True) + + if args.upload_repo: + from huggingface_hub import HfApi + + dest = prequant_filename(scheme) + print(f" uploading -> {args.upload_repo}:{dest} ...", flush = True) + api = HfApi(token = args.hf_token) + api.create_repo(args.upload_repo, exist_ok = True) + api.upload_file( + path_or_fileobj = str(out), + path_in_repo = dest, + repo_id = args.upload_repo, + revision = args.upload_revision, + ) + print(f" uploaded {dest} to {args.upload_repo}", flush = True) + + print("BUILD-PREQUANT-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/prequant_probe.py b/scripts/prequant_probe.py new file mode 100644 index 0000000000..cbcbc245fe --- /dev/null +++ b/scripts/prequant_probe.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Does a *pre-quantized* checkpoint fix the dense-quant load-VRAM spike? + +The current fast-transformer path materialises the dense bf16 transformer on the GPU +and quantises it in place -> ~2x the GGUF load peak. This probe checks the fix: quantise +once, ``torch.save`` the quantized state dict, then load it onto an empty (meta) model +with ``load_state_dict(assign=True)`` so the bf16 never touches the GPU. + +Modes (run each in its own process so peak VRAM is clean): + build -- load dense bf16, quantize_ fp8, torch.save the state dict + on-disk size. + baseline -- current path: from_pretrained bf16 -> quantize_ on GPU. Report load peak + gen. + prequant -- meta-init -> load_state_dict(saved, assign=True) -> cuda. Report load peak + gen. + +Run on one CUDA (Blackwell) GPU. Reference image for LPIPS is the baseline path.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +ROOT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" +CKPT = ROOT / "prequant_fp8" / "transformer_fp8_state.pt" +OUT = ROOT / "prequant_images" +MIN_FEAT = 512 + + +def _filt(mod, fqn = ""): + import torch.nn as nn + return ( + isinstance(mod, nn.Linear) and mod.in_features >= MIN_FEAT and mod.out_features >= MIN_FEAT + ) + + +def _fp8_cfg(): + from torchao.quantization import Float8DynamicActivationFloat8WeightConfig + return Float8DynamicActivationFloat8WeightConfig() + + +def _build(): + import torch + import diffusers + from torchao.quantization import quantize_ + + torch.cuda.reset_peak_memory_stats() + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ).to("cuda") + quantize_(t, _fp8_cfg(), filter_fn = _filt) + CKPT.parent.mkdir(parents = True, exist_ok = True) + sd = t.state_dict() + # move to cpu for a portable, gpu-free checkpoint + sd = {k: (v.detach().to("cpu") if hasattr(v, "detach") else v) for k, v in sd.items()} + torch.save(sd, CKPT) + sz = CKPT.stat().st_size / 1e9 + peak = torch.cuda.max_memory_allocated() / 1e9 + print( + f"[build] saved {CKPT.name} on-disk={sz:.2f} GB build_gpu_peak={peak:.1f} GB", flush = True + ) + return 0 + + +def _make_pipe_from_transformer(t): + import diffusers + import torch + + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _baseline(steps, seed, res): + import torch + import diffusers + from torchao.quantization import quantize_ + + torch.cuda.reset_peak_memory_stats() + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ).to("cuda") + quantize_(t, _fp8_cfg(), filter_fn = _filt) + load_peak = torch.cuda.max_memory_allocated() / 1e9 + pipe = _make_pipe_from_transformer(t) + img, dt = _gen(pipe, steps, seed, res) # warmup + img, dt = _gen(pipe, steps, seed, res) + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / "baseline.png") + print(f"[baseline] transformer_load_gpu_peak={load_peak:.1f} GB gen={dt:.3f}s", flush = True) + return 0 + + +def _prequant(steps, seed, res): + import torch + import diffusers + from accelerate import init_empty_weights + + if not CKPT.exists(): + print(f"[prequant] missing checkpoint {CKPT}; run --mode build first", flush = True) + return 1 + torch.cuda.reset_peak_memory_stats() + cfg = diffusers.ZImageTransformer2DModel.load_config(BASE, subfolder = "transformer") + with init_empty_weights(): + t = diffusers.ZImageTransformer2DModel.from_config(cfg) + sd = torch.load(CKPT, weights_only = False, map_location = "cpu") + missing, unexpected = t.load_state_dict(sd, strict = False, assign = True) + # any param/buffer still on meta (e.g. non-persistent buffers) -> materialise on cuda + leftover = [n for n, p in t.named_parameters() if p.is_meta] + [ + n for n, b in t.named_buffers() if b.is_meta + ] + if leftover: + print( + f"[prequant] {len(leftover)} meta leftovers (non-persistent buffers): {leftover[:4]}", + flush = True, + ) + t = t.to_empty(device = "cuda") # fallback path; re-loads sd below + t.load_state_dict(sd, strict = False, assign = True) + t = t.to(torch.bfloat16).to("cuda") + load_peak = torch.cuda.max_memory_allocated() / 1e9 + print( + f"[prequant] missing={len(missing)} unexpected={len(unexpected)} " + f"transformer_load_gpu_peak={load_peak:.1f} GB", + flush = True, + ) + pipe = _make_pipe_from_transformer(t) + img, dt = _gen(pipe, steps, seed, res) # warmup + img, dt = _gen(pipe, steps, seed, res) + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / "prequant.png") + # LPIPS vs baseline if present + bpath = OUT / "baseline.png" + lp = None + if bpath.exists(): + try: + import lpips + from PIL import Image + + fn = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def tt(p): + a = np.array(Image.open(p).convert("RGB")) + return ( + torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 + ).cuda() + + with torch.no_grad(): + lp = float(fn(tt(bpath), tt(OUT / "prequant.png")).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + print(f"[prequant] gen={dt:.3f}s LPIPS_vs_baseline={lp}", flush = True) + return 0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--mode", choices = ["build", "baseline", "prequant"], required = True) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + args = p.parse_args(argv) + if args.mode == "build": + return _build() + if args.mode == "baseline": + return _baseline(args.steps, args.seed, args.res) + return _prequant(args.steps, args.seed, args.res) + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + rc = main() + print("PREQUANT-PROBE-DONE", flush = True) + sys.exit(rc) diff --git a/scripts/verify_prequant_backend.py b/scripts/verify_prequant_backend.py new file mode 100644 index 0000000000..e160597599 --- /dev/null +++ b/scripts/verify_prequant_backend.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GPU verification of the Phase 9 pre-quantized load path through the real backend code. + +Exercises the actual product functions (``load_prequantized_transformer`` and the runtime +``quantize_transformer``), not a reimplementation: + + prequant -- load the checkpoint built by build_prequant_checkpoint.py via the real + ``load_prequantized_transformer`` (meta-init + assign), measure GPU load peak, + generate. + runtime -- the existing path: from_pretrained dense bf16 -> ``quantize_transformer`` on + device, measure GPU load peak, generate (the LPIPS reference). + +Asserts the prequant load peak is far below the dense one and the images match (LPIPS ~0). +Run each mode in its own process for a clean peak. One CUDA GPU.""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from pathlib import Path + +import numpy as np + +_REPO = Path(__file__).resolve().parent.parent +_RESEARCH = _REPO / "outputs" / "quant_research" +BACKEND = _REPO / "studio" / "backend" +BASE = "Tongyi-MAI/Z-Image-Turbo" +CKPT = os.environ.get("PREQUANT_CKPT", str(_RESEARCH / "prequant_fp8" / "transformer_fp8.pt")) +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_images"))) + +logging.basicConfig(level = logging.INFO, format = "%(message)s") +LOGGER = logging.getLogger("verify_prequant") + + +def _target(dtype): + import types + return types.SimpleNamespace(device = "cuda", dtype = dtype) + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _lpips(ref, arr): + try: + import lpips, torch + + fn = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(fn(t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def run(mode, steps, seed, res): + sys.path.insert(0, str(BACKEND)) + import torch + import diffusers + from core.inference.diffusion_prequant import PrequantSource, load_prequantized_transformer + from core.inference.diffusion_transformer_quant import quantize_transformer + + OUT.mkdir(parents = True, exist_ok = True) + transformer_cls = diffusers.ZImageTransformer2DModel + torch.cuda.reset_peak_memory_stats() + torch.cuda.empty_cache() + + if mode == "prequant": + source = PrequantSource(kind = "path", location = CKPT, filename = None) + transformer = load_prequantized_transformer( + transformer_cls, + BASE, + source, + device = "cuda", + dtype = torch.bfloat16, + hf_token = None, + scheme = "fp8", + logger = LOGGER, + ) + if transformer is None: + print("prequant load FAILED (returned None)", flush = True) + return 1 + pipe = diffusers.ZImagePipeline.from_pretrained( + BASE, torch_dtype = torch.bfloat16, transformer = transformer + ) + pipe.to("cuda") + load_peak = torch.cuda.max_memory_allocated() / 1e9 + marker = getattr(transformer, "_unsloth_runtime_quant", None) + print(f"[prequant] load_gpu_peak={load_peak:.1f} GB marker={marker}", flush = True) + else: # runtime + transformer = transformer_cls.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ).to("cuda") + pipe = diffusers.ZImagePipeline.from_pretrained( + BASE, torch_dtype = torch.bfloat16, transformer = transformer + ) + pipe.to("cuda") + scheme = quantize_transformer(pipe, _target(torch.bfloat16), mode = "fp8", logger = LOGGER) + load_peak = torch.cuda.max_memory_allocated() / 1e9 + print(f"[runtime] engaged={scheme} load_gpu_peak={load_peak:.1f} GB", flush = True) + + img, dt = _gen(pipe, steps, seed, res) # warmup + img, dt = _gen(pipe, steps, seed, res) + img.save(OUT / f"{mode}.png") + print(f"[{mode}] gen={dt:.3f}s saved {mode}.png", flush = True) + + ref_path = OUT / "runtime.png" + if mode == "prequant" and ref_path.exists(): + from PIL import Image + lp = _lpips(np.array(Image.open(ref_path).convert("RGB")), np.array(img)) + print(f"[prequant] LPIPS_vs_runtime={lp}", flush = True) + return 0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--mode", choices = ["prequant", "runtime"], required = True) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + args = p.parse_args(argv) + rc = run(args.mode, args.steps, args.seed, args.res) + print("VERIFY-PREQUANT-DONE", flush = True) + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 3cac5629ac..aac4069833 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -53,10 +53,16 @@ from .diffusion_speed import ( snapshot_backend_flags, ) from .diffusion_precision import quantize_text_encoders +from .diffusion_prequant import ( + load_prequantized_transformer, + resolve_prequant_source, +) from .diffusion_transformer_quant import ( + DEFAULT_MIN_LINEAR_FEATURES, dense_transformer_supported, normalize_transformer_quant, quantize_transformer, + select_transformer_quant_scheme, ) logger = get_logger(__name__) @@ -299,6 +305,7 @@ class DiffusionBackend: text_encoder_quant: Optional[str] = None, transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, + transformer_prequant_path: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -332,6 +339,7 @@ class DiffusionBackend: text_encoder_quant = text_encoder_quant, transformer_quant = transformer_quant, transformer_quant_fast_accum = transformer_quant_fast_accum, + transformer_prequant_path = transformer_prequant_path, _load_token = token, ), daemon = True, @@ -467,6 +475,7 @@ class DiffusionBackend: text_encoder_quant: Optional[str] = None, transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, + transformer_prequant_path: Optional[str] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -534,6 +543,8 @@ class DiffusionBackend: target, transformer_quant, transformer_quant_fast_accum, + fam = fam, + prequant_path = transformer_prequant_path, ) except Exception as exc: # noqa: BLE001 — fall back to the GGUF build logger.warning( @@ -662,27 +673,76 @@ class DiffusionBackend: target: DiffusionDeviceTarget, mode: Optional[str], fast_accum: Optional[bool] = None, + *, + fam: Optional[DiffusionFamily] = None, + prequant_path: Optional[str] = None, ) -> tuple[Any, str]: - """Build the opt-in fast pipeline: load the DENSE bf16 transformer from the base - repo (``subfolder="transformer"``), assemble the pipeline, place it on the device, - and torchao-quantise the transformer in place. Returns ``(pipe, engaged_scheme)``. + """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``. + + Two ways to get the quantized transformer, in order: + + 1. Pre-quantized: if a checkpoint is configured for the chosen scheme (an explicit + ``prequant_path`` or the family's hosted repo), load the already-quantized + weights onto the meta device and assign them in -- the dense bf16 never lands on + the GPU, so the load peak is ~half and the download is smaller. + 2. Dense + quantise (fallback): load the DENSE bf16 transformer from the base repo, + place it on the device, and torchao-quantise it in place. Raises if the scheme is unsupported or quantisation fails, so ``load_pipeline`` - catches it and falls back to the GGUF build. Quantisation runs ON the device (the - dynamic int8 / fp8 / fp4 kernels need the weights on CUDA) and BEFORE the loader - compiles the repeated block, so the order is quantize -> compile -> placement.""" + catches it and falls back to the GGUF build. Quantisation runs ON the device and + BEFORE the loader compiles the repeated block, so the order stays quantize -> + compile -> placement.""" + # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme. + scheme = select_transformer_quant_scheme(target, mode) + if scheme is not None and fam is not None: + source = resolve_prequant_source(fam, scheme, path_override = prequant_path) + if source is not None: + transformer = load_prequantized_transformer( + transformer_cls, + base, + source, + device = device, + dtype = dtype, + hf_token = hf_token, + scheme = scheme, + # Reject a checkpoint built with a different Linear filter than the + # dense path uses, so the prequant and runtime-quant models match. + min_features = DEFAULT_MIN_LINEAR_FEATURES, + logger = logger, + ) + if transformer is not None: + pipe = self._assemble_pipe( + pipeline_cls, base, transformer, dtype, hf_token, device + ) + return pipe, scheme + + # 2. Fallback: materialise the dense bf16 transformer and quantise it on-device. transformer = transformer_cls.from_pretrained( base, subfolder = "transformer", torch_dtype = dtype, token = hf_token ) + pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device) + scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger) + if scheme is None: + raise RuntimeError("transformer quant unsupported for this device/scheme") + return pipe, scheme + + @staticmethod + def _assemble_pipe( + pipeline_cls: Any, + base: str, + transformer: Any, + dtype: Any, + hf_token: Optional[str], + device: str, + ) -> Any: + """Assemble the diffusers pipeline around ``transformer`` and place it on ``device`` + (a no-op for an already-placed pre-quantized transformer; it moves the companions).""" pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} if hf_token: pipe_kwargs["token"] = hf_token pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) pipe.to(device) - scheme = quantize_transformer(pipe, target, mode = mode, fast_accum = fast_accum, logger = logger) - if scheme is None: - raise RuntimeError("transformer quant unsupported for this device/scheme") - return pipe, scheme + return pipe def _plan_memory( self, diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 4a3dd616ac..d5068cbf49 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -39,6 +39,12 @@ class DiffusionFamily: # regional torch.compile. Now consulted on the GGUF path too (compile runs on the # GGUF transformer); all current families compile, so this stays True. supports_torch_compile: bool = True + # Optional pre-quantized transformer checkpoints, as (scheme, repo_id) pairs (a + # hashable mapping). When the fast transformer_quant path resolves a scheme with a + # hosted checkpoint, the loader fetches the already-quantized weights instead of + # materialising the dense bf16 transformer on the GPU (much lower load VRAM + a + # smaller download). Empty until checkpoints are hosted -> behaviour is unchanged. + prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple) # Keyed by architecture, not per model variant: a checkpoint's specific base repo @@ -123,6 +129,14 @@ def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: return base or fam.base_repo +def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: + """The hosted pre-quantized transformer repo for ``scheme`` in this family, or None.""" + for entry_scheme, repo_id in fam.prequant_repos: + if entry_scheme == scheme: + return repo_id + return None + + def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: """Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes. diff --git a/studio/backend/core/inference/diffusion_prequant.py b/studio/backend/core/inference/diffusion_prequant.py new file mode 100644 index 0000000000..207f305220 --- /dev/null +++ b/studio/backend/core/inference/diffusion_prequant.py @@ -0,0 +1,287 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Load a *pre-quantized* transformer instead of quantising a dense one on the GPU. + +The opt-in fast transformer_quant path (see ``diffusion_transformer_quant.py``) loads +the dense bf16 transformer and torchao-``quantize_``s it in place. That materialises the +full bf16 weights on the GPU before quantising, so the load peak is ~2x the GGUF's and it +pulls the full bf16 download. When a transformer has already been quantised once and saved +(``scripts/build_prequant_checkpoint.py``), this module loads those weights directly: + + 1. build the transformer skeleton on the ``meta`` device (no storage) via + ``accelerate.init_empty_weights`` + ``from_config``; + 2. ``load_state_dict(assign=True)`` the quantized state dict (the torchao weight subclass + tensors are assigned in, not copied), so the dense bf16 never touches the GPU; + 3. move to the device. + +Measured (B200, Z-Image fp8): transformer GPU load peak 12.9 -> 6.3 GB, download 12 -> +6.28 GB, output bit-identical (LPIPS 0.0). The checkpoint carries the exact same scheme + +``min_features`` as the runtime path, so the result is identical to quantising on the fly. + +Best-effort and lazily imported throughout: a missing / mismatched / unreadable checkpoint +returns None and the caller falls back to the dense-quantise path (and then to GGUF). All +behaviour is gated on a configured source -- with nothing configured this module is inert. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Optional + +# torch.save dict layout this module reads (and the build script writes). Bumped if the +# on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded. +PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1" + +# Loading a checkpoint ends in ``torch.load(weights_only=False)``, which executes arbitrary +# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted, +# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field +# of a load request -- i.e. an authenticated API caller naming an arbitrary local file. +# Unpickling that is remote code execution, so a request-supplied path is unpickled ONLY when +# it resolves inside an operator-configured ALLOWLIST of directories. A bare on/off toggle is +# deliberately NOT accepted as a wildcard: enabling local checkpoints for one trusted +# directory must never also permit unpickling any other path a request happens to name. The +# trusted hosted-repo path is unaffected. +ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH" + +_PREQUANT_TOGGLE_TOKENS = {"1", "true", "yes", "on", "0", "false", "no", "off"} + + +def _allowed_prequant_roots() -> list: + """Operator-allowlisted directories whose pre-quant checkpoints may be unpickled. + + Set ``UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH`` to one or more directories (separated by + ``os.pathsep``). A bare truthy/falsey toggle is ignored on purpose -- it must name a + directory, so there is no "allow everything" mode.""" + import os + + raw = (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip() + if not raw: + return [] + roots = [] + for part in raw.split(os.pathsep): + part = part.strip() + if not part or part.lower() in _PREQUANT_TOGGLE_TOKENS: + continue # a bare on/off value is not a directory -> never a wildcard allow + try: + roots.append(os.path.realpath(os.path.expanduser(part))) + except Exception: # noqa: BLE001 — a bad entry is simply not allowlisted + continue + return roots + + +def _local_prequant_path_allowed(path: str) -> bool: + """True only when ``path`` resolves inside an operator-allowlisted directory; an + arbitrary request-supplied path is never unpickled. ``realpath`` first so a symlink + cannot point an allowlisted name at a file outside the allowed roots.""" + import os + + roots = _allowed_prequant_roots() + if not roots: + return False + try: + real = os.path.realpath(os.path.expanduser(path)) + except Exception: # noqa: BLE001 + return False + return any(real == r or real.startswith(r + os.sep) for r in roots) + + +@dataclass(frozen = True) +class PrequantSource: + """Where a pre-quantized transformer checkpoint lives. ``kind`` is "path" (a local + file) or "repo" (a Hub repo id in ``location`` + ``filename`` inside it).""" + + kind: str + location: str + filename: Optional[str] = None + + +def prequant_filename(scheme: str) -> str: + """The conventional checkpoint filename for ``scheme`` inside a Hub repo.""" + return f"transformer_{scheme}.pt" + + +def resolve_prequant_source( + fam: Any, + scheme: str, + *, + path_override: Optional[str] = None, +) -> Optional[PrequantSource]: + """Resolve where the pre-quantized checkpoint for ``(fam, scheme)`` should come from. + + Priority: (1) an explicit local ``path_override`` (testing / power users); (2) the + family's hosted repo for ``scheme``; (3) None -> no pre-quant, caller quantises dense. + Pure: no IO, no torch -- it only decides the source, the loader fetches it. + """ + override = (path_override or "").strip() + if override: + return PrequantSource(kind = "path", location = override, filename = None) + try: + from .diffusion_families import family_prequant_repo + repo_id = family_prequant_repo(fam, scheme) + except Exception: # noqa: BLE001 — a bad family object must not break the load + repo_id = None + if repo_id: + return PrequantSource(kind = "repo", location = repo_id, filename = prequant_filename(scheme)) + return None + + +def load_prequantized_transformer( + transformer_cls: Any, + base: str, + source: PrequantSource, + *, + device: str, + dtype: Any, + hf_token: Optional[str] = None, + scheme: str, + min_features: Optional[int] = None, + logger: Any = None, +) -> Optional[Any]: + """Load the pre-quantized transformer described by ``source`` onto ``device``. + + Returns the placed, already-quantized transformer, or None on any problem (missing / + mismatched / unreadable checkpoint, or a meta-init the class does not support) so the + caller falls back to the dense-quantise path. Best-effort: never raises for an + ordinary unavailable artifact. + """ + try: + # weights_only=False (required below) executes pickle code, so a caller-supplied + # local path is unpickled ONLY when it resolves inside an operator-allowlisted + # directory. The hosted family repo is first-party and always allowed. + if source.kind == "path" and not _local_prequant_path_allowed(source.location): + _warn( + logger, + f"{scheme}:path", + RuntimeError( + "request-supplied local pre-quant path refused (unpickling an arbitrary " + f"file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV} to an allowlisted " + "directory containing trusted checkpoints to permit it", + ), + ) + return None + + path = _resolve_checkpoint_path(source, hf_token) + if path is None: + return None + + import torch + + # torchao weight subclasses are not safetensors-serializable, so the checkpoint is + # a torch.save pickle. weights_only=False is required to rebuild those subclasses. + # The local-path branch is gated above; the repo branch is a first-party artifact. + ckpt = torch.load(path, weights_only = False, map_location = "cpu") + if not _validate_checkpoint(ckpt, scheme, base, logger, min_features = min_features): + return None + state_dict = ckpt["state_dict"] + + config = transformer_cls.load_config(base, subfolder = "transformer", token = hf_token) + from accelerate import init_empty_weights + + with init_empty_weights(): + transformer = transformer_cls.from_config(config) + # assign=True swaps in the loaded (quantized) tensors rather than copying into the + # meta tensors (a copy into meta is a no-op); strict=True since the saved state + # dict is the full state dict of the same class (non-persistent buffers excluded). + transformer.load_state_dict(state_dict, strict = True, assign = True) + if _has_meta_tensors(transformer): + # A class with non-persistent buffers (computed in __init__, absent from the + # state dict) leaves those on meta. Rebuild on CPU so the buffers hold their + # real values, then re-assign the quantized weights. The dense bf16 lives in + # CPU RAM only -- the GPU still receives just the quantized footprint. + transformer = transformer_cls.from_config(config) + transformer.load_state_dict(state_dict, strict = True, assign = True) + + transformer = transformer.to(device) + try: # diagnostic marker, mirrors the runtime-quant path + transformer._unsloth_runtime_quant = scheme + except Exception: # noqa: BLE001 — marker is best-effort + pass + if logger is not None: + logger.info( + "diffusion.prequant: loaded %s checkpoint (%s) onto %s", + scheme, + source.kind, + device, + ) + return transformer + except Exception as exc: # noqa: BLE001 — fall back to the dense-quantise path + _warn(logger, f"{scheme}:{source.kind}", exc) + return None + + +def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) -> Optional[str]: + """The local file path for ``source``, downloading from the Hub if needed; None if absent.""" + if source.kind == "path": + import os + return source.location if os.path.isfile(source.location) else None + if source.kind == "repo": + from huggingface_hub import hf_hub_download + return hf_hub_download(repo_id = source.location, filename = source.filename, token = hf_token) + return None + + +def _validate_checkpoint( + ckpt: Any, + scheme: str, + base: str, + logger: Any, + min_features: Optional[int] = None, +) -> bool: + """Reject a checkpoint that is the wrong format / scheme / base model / filter. + + ``min_features`` (when given) is the runtime Linear-feature threshold: a checkpoint + built with a different ``--min-features`` quantises a different set of Linear layers, + so ``load_state_dict(assign=True)`` would silently install a model that does not match + what the dense path produces while status still reports the requested scheme. Reject it.""" + if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT: + _warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format")) + return False + if "state_dict" not in ckpt: + _warn(logger, scheme, ValueError("pre-quant checkpoint has no state_dict")) + return False + meta = ckpt.get("metadata") or {} + if meta.get("scheme") != scheme: + _warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}")) + return False + ckpt_base = meta.get("base_model_id") + if ckpt_base and base and not _same_base_model(ckpt_base, base): + _warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}")) + return False + if min_features is not None: + ckpt_min = meta.get("min_features") + if ckpt_min is not None and int(ckpt_min) != int(min_features): + _warn( + logger, + scheme, + ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"), + ) + return False + return True + + +def _same_base_model(a: str, b: str) -> bool: + """Tolerant compare of two base-model ids: an exact match, or the same final + path/repo segment (so a local path or a fork id matches the canonical repo, e.g. + ``/models/Z-Image-Turbo`` vs ``Tongyi-MAI/Z-Image-Turbo``).""" + + def _tail(x: str) -> str: + return x.replace("\\", "/").rstrip("/").split("/")[-1].lower() + + return a == b or _tail(a) == _tail(b) + + +def _has_meta_tensors(module: Any) -> bool: + """True if any parameter or buffer is still on the meta device after loading.""" + from itertools import chain + try: + return any( + getattr(t, "is_meta", False) for t in chain(module.parameters(), module.buffers()) + ) + except Exception: # noqa: BLE001 + return False + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.prequant: %s failed: %s", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 56ee571dae..08d309f257 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1737,6 +1737,19 @@ class DiffusionLoadRequest(BaseModel): "HBM cards, which are not nerfed). true/false force it. Negligible " "quality effect (below the fp8 quant noise floor); no overflow risk.", ) + transformer_prequant_path: Optional[str] = Field( + None, + description = "Local path to a pre-quantized transformer checkpoint (built by " + "scripts/build_prequant_checkpoint.py) for the requested transformer_quant " + "scheme. Loads the already-quantized weights with the dense bf16 never on the " + "GPU (~half the load VRAM and a smaller download). null uses the family's hosted " + "checkpoint if configured, else quantises the dense transformer at load time. " + "Loading a local path unpickles the file (arbitrary code execution), so it is " + "ignored unless the path resolves inside a directory the operator allowlisted " + "via UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH (one or more directories, separated by " + "the OS path separator). A bare on/off value such as '1' is deliberately not " + "accepted -- it must name an allowed directory.", + ) class DiffusionGenerateRequest(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 13beee00e0..ec5382e190 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10333,6 +10333,7 @@ async def load_diffusion_model( text_encoder_quant = request.text_encoder_quant, transformer_quant = request.transformer_quant, transformer_quant_fast_accum = request.transformer_quant_fast_accum, + transformer_prequant_path = request.transformer_prequant_path, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 2e714b217b..f8d0172173 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -920,6 +920,10 @@ def _stub_dense_quant(monkeypatch, *, scheme = "fp8"): monkeypatch.setattr(_FakeTransformer, "from_pretrained", _from_pretrained, raising = False) monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + # Resolve the scheme without the real GPU smoke probe, and configure no pre-quant + # checkpoint so the dense materialise+quantise branch is the one exercised. + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: scheme) + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) def _quantize(pipe, target, *, mode, **kw): calls["quantize"] += 1 @@ -974,6 +978,77 @@ def test_transformer_quant_dense_path_engaged(fake_runtime, tmp_path, monkeypatc assert status["offload_policy"] == "none" +def test_transformer_quant_prequant_path_engaged(fake_runtime, tmp_path, monkeypatch): + # A configured pre-quant checkpoint -> load the already-quantized transformer directly; + # the dense from_pretrained and the on-device quantize_transformer are NOT used. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) + prequant_obj = object() + loaded: dict = {"n": 0} + + def _load_prequant(transformer_cls, base, source, **kw): + loaded["n"] += 1 + loaded["scheme"] = kw.get("scheme") + return prequant_obj + + monkeypatch.setattr(dmod, "load_prequantized_transformer", _load_prequant) + + @classmethod + def _fp_fail(cls, *a, **k): + pytest.fail("dense from_pretrained must not run when a prequant checkpoint loads") + + monkeypatch.setattr(_FakeTransformer, "from_pretrained", _fp_fail, raising = False) + monkeypatch.setattr( + dmod, + "quantize_transformer", + lambda *a, **k: pytest.fail("quantize_transformer must not run on the prequant path"), + ) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + transformer_prequant_path = str(tmp_path / "zimage_fp8.pt"), + ) + assert status["transformer_quant"] == "fp8" + assert loaded["n"] == 1 and loaded["scheme"] == "fp8" + # The pre-quantized transformer object was assembled into the pipeline... + assert _FakePipeline.last.get("transformer") is prequant_obj + # ...and the GGUF single-file path was not used. + assert _FakeTransformer.last == {} + + +def test_transformer_quant_prequant_load_fails_falls_back_to_dense( + fake_runtime, tmp_path, monkeypatch +): + # A configured prequant source whose load returns None must fall back to the dense + # materialise+quantise path (not straight to GGUF), preserving the fast mode. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + calls = _stub_dense_quant(monkeypatch, scheme = "fp8") + # Override the no-prequant default: a source resolves, but its load fails. + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) + monkeypatch.setattr(dmod, "load_prequantized_transformer", lambda *a, **k: None) + (tmp_path / "m.gguf").write_bytes(b"x") + status = backend.load_pipeline( + str(tmp_path), + gguf_filename = "m.gguf", + family_override = "z-image", + transformer_quant = "fp8", + ) + assert status["transformer_quant"] == "fp8" + assert calls["from_pretrained"] == 1 and calls["quantize"] == 1 # dense path ran + assert _FakeTransformer.last == {} # GGUF not used + + def test_transformer_quant_falls_back_to_gguf_on_failure(fake_runtime, tmp_path, monkeypatch): # A dense/quant failure (here: quantize returns None -> unsupported) must fall back # to the GGUF build, not error -- status reports no transformer_quant engaged. diff --git a/studio/backend/tests/test_diffusion_prequant.py b/studio/backend/tests/test_diffusion_prequant.py new file mode 100644 index 0000000000..a233f6225f --- /dev/null +++ b/studio/backend/tests/test_diffusion_prequant.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for the pre-quantized transformer load path. + +torch / accelerate are stubbed via ``sys.modules`` (the module under test imports them +lazily), and ``transformer_cls`` is a fake that records calls -- so the resolver, the +meta-init + ``load_state_dict(assign=True)`` flow, and the validation/fallback behaviour +are all exercised without CUDA, torchao, or a real diffusers model. +""" + +from __future__ import annotations + +import contextlib +import sys +import types + +import core.inference.diffusion_prequant as pq +from core.inference.diffusion_families import DiffusionFamily +from core.inference.diffusion_prequant import ( + PREQUANT_FORMAT, + PrequantSource, + load_prequantized_transformer, + resolve_prequant_source, +) + + +# ── resolve_prequant_source ────────────────────────────────────────────────────── +def _fam(prequant_repos = ()): + return DiffusionFamily( + name = "z-image", + pipeline_class = "ZImagePipeline", + transformer_class = "ZImageTransformer2DModel", + base_repo = "Tongyi-MAI/Z-Image-Turbo", + prequant_repos = prequant_repos, + ) + + +def test_resolve_path_override_wins(): + fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),)) + src = resolve_prequant_source(fam, "fp8", path_override = "/tmp/local.pt") + assert src == PrequantSource(kind = "path", location = "/tmp/local.pt", filename = None) + + +def test_resolve_family_repo_by_scheme(): + fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"), ("int8", "org/hosted-int8"))) + src = resolve_prequant_source(fam, "int8") + assert src.kind == "repo" and src.location == "org/hosted-int8" + assert src.filename == "transformer_int8.pt" + + +def test_resolve_wrong_scheme_is_none(): + fam = _fam(prequant_repos = (("fp8", "org/hosted-fp8"),)) + assert resolve_prequant_source(fam, "int8") is None + + +def test_resolve_nothing_configured_is_none(): + assert resolve_prequant_source(_fam(), "fp8") is None + assert resolve_prequant_source(_fam(), "fp8", path_override = "") is None + + +# ── load_prequantized_transformer ──────────────────────────────────────────────── +class _FakeTransformer: + calls: dict = {} + + def __init__(self): + self.assigned = None + self.moved = None + + @classmethod + def load_config(cls, base, **kw): + cls.calls["load_config"] = {"base": base, **kw} + return {"cfg": True} + + @classmethod + def from_config(cls, config): + cls.calls["from_config"] = config + return cls() + + @classmethod + def from_pretrained(cls, *a, **k): # the dense path -- must never run here + cls.calls["from_pretrained"] = True + raise AssertionError("from_pretrained must not be called on the prequant path") + + def load_state_dict( + self, + sd, + strict = True, + assign = False, + ): + _FakeTransformer.calls["load_state_dict"] = {"strict": strict, "assign": assign} + self.assigned = sd + + def parameters(self): + return [] + + def buffers(self): + return [] + + def to(self, device): + self.moved = device + return self + + +def _stub_torch_accelerate( + monkeypatch, + ckpt, + *, + load_raises = False, +): + torch = types.ModuleType("torch") + + def _load( + path, + weights_only = False, + map_location = None, + ): + if load_raises: + raise RuntimeError("corrupt checkpoint") + return ckpt + + torch.load = _load + monkeypatch.setitem(sys.modules, "torch", torch) + + accelerate = types.ModuleType("accelerate") + accelerate.init_empty_weights = lambda: contextlib.nullcontext() + monkeypatch.setitem(sys.modules, "accelerate", accelerate) + + +def _good_ckpt(scheme = "fp8", base = "Tongyi-MAI/Z-Image-Turbo"): + return { + "format": PREQUANT_FORMAT, + "metadata": {"scheme": scheme, "base_model_id": base}, + "state_dict": {"weight": object()}, + } + + +def _load( + monkeypatch, + tmp_path, + ckpt, + *, + scheme = "fp8", + load_raises = False, + exists = True, + allow_local = True, +): + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, ckpt, load_raises = load_raises) + # The local-path branch is opt-in via a directory ALLOWLIST (it unpickles an arbitrary + # file); these tests exercise the load mechanics, so allowlist tmp_path (where ckpt.pt + # lives) unless a test is checking the gate. + if allow_local: + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path)) + else: + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + path = tmp_path / "ckpt.pt" + if exists: + path.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(path), filename = None) + return load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = scheme, + logger = None, + ) + + +def test_load_meta_init_and_assign(monkeypatch, tmp_path): + t = _load(monkeypatch, tmp_path, _good_ckpt()) + assert t is not None + # meta-init path was used, not the dense from_pretrained. + assert "from_config" in _FakeTransformer.calls + assert "from_pretrained" not in _FakeTransformer.calls + # assign=True is the whole point (copy into meta is a no-op). + assert _FakeTransformer.calls["load_state_dict"] == {"strict": True, "assign": True} + assert t.moved == "cuda" + assert t._unsloth_runtime_quant == "fp8" + + +def test_load_missing_file_is_none(monkeypatch, tmp_path): + assert _load(monkeypatch, tmp_path, _good_ckpt(), exists = False) is None + + +def test_load_torch_load_raises_is_none(monkeypatch, tmp_path): + assert _load(monkeypatch, tmp_path, _good_ckpt(), load_raises = True) is None + + +def test_load_format_mismatch_is_none(monkeypatch, tmp_path): + bad = _good_ckpt() + bad["format"] = "something_else" + assert _load(monkeypatch, tmp_path, bad) is None + + +def test_load_scheme_mismatch_is_none(monkeypatch, tmp_path): + # checkpoint built for int8, but fp8 was requested. + assert _load(monkeypatch, tmp_path, _good_ckpt(scheme = "int8"), scheme = "fp8") is None + + +def test_load_base_mismatch_is_none(monkeypatch, tmp_path): + assert _load(monkeypatch, tmp_path, _good_ckpt(base = "other/model")) is None + + +# ── local-path opt-in gate (RCE guard) ─────────────────────────────────────────── +def test_load_local_path_refused_by_default(monkeypatch, tmp_path): + # A valid checkpoint at a real file is still refused: torch.load must never run on a + # request-supplied path without the operator opt-in. + called = {"load": False} + + def _explode(*a, **k): + called["load"] = True + raise AssertionError("torch.load must not run on a refused local path") + + torch = types.ModuleType("torch") + torch.load = _explode + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + + path = tmp_path / "ckpt.pt" + path.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(path), filename = None) + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is None + assert called["load"] is False + + +def test_load_local_path_allowed_with_optin(monkeypatch, tmp_path): + assert _load(monkeypatch, tmp_path, _good_ckpt(), allow_local = True) is not None + + +def test_load_repo_source_allowed_without_optin(monkeypatch, tmp_path): + # The hosted-repo branch is first-party and trusted: it loads with no opt-in env set. + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, _good_ckpt()) + monkeypatch.delenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + + downloaded = tmp_path / "transformer_fp8.pt" + downloaded.write_bytes(b"x") + hub = types.ModuleType("huggingface_hub") + hub.hf_hub_download = lambda repo_id, filename, token = None: str(downloaded) + monkeypatch.setitem(sys.modules, "huggingface_hub", hub) + + source = PrequantSource(kind = "repo", location = "org/hosted-fp8", filename = "transformer_fp8.pt") + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is not None + + +def test_load_local_path_outside_allowlist_refused(monkeypatch, tmp_path): + # Even with the opt-in set, a path OUTSIDE every allowlisted directory must not be + # unpickled: enabling one trusted dir is not a wildcard for arbitrary request paths. + called = {"load": False} + + def _explode(*a, **k): + called["load"] = True + raise AssertionError("torch.load must not run on a path outside the allowlist") + + torch = types.ModuleType("torch") + torch.load = _explode + monkeypatch.setitem(sys.modules, "torch", torch) + + allowed = tmp_path / "allowed" + allowed.mkdir() + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(allowed)) + + outside = tmp_path / "evil.pt" # a real file, but outside the allowlisted dir + outside.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(outside), filename = None) + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is None + assert called["load"] is False + + +def test_load_min_features_mismatch_is_none(monkeypatch, tmp_path): + # A checkpoint built with a different --min-features quantises a different Linear set, + # so it must be rejected when the runtime threshold is supplied. + ckpt = _good_ckpt() + ckpt["metadata"]["min_features"] = 256 # built with 256, runtime asks for 512 + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, ckpt) + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path)) + path = tmp_path / "ckpt.pt" + path.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(path), filename = None) + result = load_prequantized_transformer( + _FakeTransformer, + "Tongyi-MAI/Z-Image-Turbo", + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + min_features = 512, + logger = None, + ) + assert result is None + + +def test_load_base_fork_tail_matches(monkeypatch, tmp_path): + # A local path / fork id with the same final segment as the canonical base is accepted. + ckpt = _good_ckpt(base = "Tongyi-MAI/Z-Image-Turbo") + _FakeTransformer.calls = {} + _stub_torch_accelerate(monkeypatch, ckpt) + monkeypatch.setenv(pq.ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path)) + path = tmp_path / "ckpt.pt" + path.write_bytes(b"x") + source = PrequantSource(kind = "path", location = str(path), filename = None) + result = load_prequantized_transformer( + _FakeTransformer, + "/local/models/Z-Image-Turbo", # different prefix, same tail + source, + device = "cuda", + dtype = "bfloat16", + hf_token = None, + scheme = "fp8", + logger = None, + ) + assert result is not None diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 137b218636..839d325443 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -408,6 +408,35 @@ def test_transformer_quant_fast_accum_threads_through(client, monkeypatch): assert backend.last_load_kwargs.get("transformer_quant_fast_accum") is False +def test_transformer_prequant_path_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_quant": "fp8", + "transformer_prequant_path": "/data/zimage_fp8.pt", + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_prequant_path") == "/data/zimage_fp8.pt" + + +def test_prequant_path_doc_describes_allowlist_not_toggle(): + # The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a + # directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots + # drops bare on/off tokens), so operators following the doc don't get every + # request silently refused. + from models.inference import DiffusionLoadRequest + + desc = DiffusionLoadRequest.model_fields["transformer_prequant_path"].description + assert "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH" in desc + assert "=1" not in desc + assert "allowlist" in desc.lower() or "director" in desc.lower() + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted. From 2ffe6fbeaba07a443e236e91717de08f53cad943 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:38:53 -0700 Subject: [PATCH 06/16] Studio diffusion (Phase 10): attention-backend selection (#6701) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/perf_levers_probe.py | 238 +++++++++++++++++ studio/backend/core/inference/diffusion.py | 25 ++ .../core/inference/diffusion_attention.py | 245 ++++++++++++++++++ studio/backend/models/inference.py | 29 +++ studio/backend/routes/inference.py | 1 + .../backend/tests/test_diffusion_attention.py | 225 ++++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 23 ++ 7 files changed, 786 insertions(+) create mode 100644 scripts/perf_levers_probe.py create mode 100644 studio/backend/core/inference/diffusion_attention.py create mode 100644 studio/backend/tests/test_diffusion_attention.py diff --git a/scripts/perf_levers_probe.py b/scripts/perf_levers_probe.py new file mode 100644 index 0000000000..3964ab46ad --- /dev/null +++ b/scripts/perf_levers_probe.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Measure the next-phase diffusion levers on the real model, vs today's compiled baseline. + +Variants (Z-Image dense bf16, regional compile = the shipped "default" speed profile): + baseline -- channels_last + compile_repeated_blocks (reference image) + inductor_flags -- + the lossless inductor autotune flags (conv_1x1_as_mm, + coordinate_descent_tuning(+all_dirs), epilogue_fusion=False) + attn_cudnn -- + set_attention_backend("_native_cudnn") (exact) + attn_flash4 -- + set_attention_backend("flash_4_hub") (exact, SM100) + attn_sage -- + set_attention_backend("sage") (INT8 QK, quantized) + fbcache -- + First-Block-Cache (threshold 0.12) (few-step headroom test) + +Reports median latency, vs-baseline speedup, peak VRAM, and LPIPS vs baseline. One CUDA GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "perf_levers_images" + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + + # Keep the metric model on CPU: caching it on CUDA leaves it resident across + # variants, and each run resets peak-memory stats, so its VRAM would be charged + # to (and reduce headroom for) every later variant's measurement. + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval() + + def t(x): + return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0 + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _set_inductor_flags(): + import torch._inductor.config as ic + + ic.conv_1x1_as_mm = True + ic.coordinate_descent_tuning = True + ic.coordinate_descent_check_all_directions = True + ic.epilogue_fusion = False + try: + ic.force_fuse_int_mm_with_mul = True + except Exception: # noqa: BLE001 + pass + + +def _reset_inductor_flags(): + import torch._inductor.config as ic + + ic.conv_1x1_as_mm = False + ic.coordinate_descent_tuning = False + ic.coordinate_descent_check_all_directions = False + ic.epilogue_fusion = True + # Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into + # every later compiled row and the attention/fbcache measurements stop being isolated. + try: + ic.force_fuse_int_mm_with_mul = False + except Exception: # noqa: BLE001 + pass + + +def _load(): + import diffusers + import torch + + t = diffusers.ZImageTransformer2DModel.from_pretrained( + BASE, subfolder = "transformer", torch_dtype = torch.bfloat16 + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + try: + pipe.vae.to(memory_format = torch.channels_last) + except Exception: # noqa: BLE001 + pass + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def run( + tag, + steps, + seed, + res, + iters, + *, + attn = None, + fbcache = None, + inductor = False, +): + import torch + + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + _reset_inductor_flags() + if inductor: + _set_inductor_flags() + pipe = _load() + note = "" + if attn is not None: + try: + pipe.transformer.set_attention_backend(attn) + except Exception as exc: # noqa: BLE001 + note = f"attn({attn})={type(exc).__name__}:{str(exc)[:60]}" + print(f" [{tag}] {note}", flush = True) + del pipe # free the resident pipe so a skipped variant doesn't leak VRAM + torch.cuda.empty_cache() + return None + if fbcache is not None: + try: + from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbcache)) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] fbcache={type(exc).__name__}:{str(exc)[:60]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None + try: + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] compile={type(exc).__name__}:{str(exc)[:60]}", flush = True) + try: + _gen(pipe, steps, seed, res) # warmup / compile + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" [{tag}] FAILED first gen: {type(exc).__name__}:{str(exc)[:80]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res) + dts.append(dt) + peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, peak + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + args = p.parse_args(argv) + s, r, seed, it = args.steps, args.res, args.seed, args.iters + + print(f"== perf levers (Z-Image dense, {r}px, {s} steps) ==", flush = True) + base = run("baseline", s, seed, r, it) + if base is None: + print("baseline FAILED", flush = True) + return 1 + bmed, ref, bpeak = base + print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush = True) + rows = [("baseline", bmed, bpeak, 0.0)] + + variants = [ + ("inductor_flags", dict(inductor = True)), + ("attn_cudnn", dict(attn = "_native_cudnn")), + ("attn_flash4", dict(attn = "flash_4_hub")), + ("attn_sage", dict(attn = "sage")), + ("attn_sage_inductor", dict(attn = "sage", inductor = True)), + ("fbcache_0p12", dict(fbcache = 0.12)), + ] + for tag, kw in variants: + out = run(tag, s, seed, r, it, **kw) + if out is None: + rows.append((tag, None, None, None)) + continue + med, arr, peak = out + lp = _lpips(ref, arr) + rows.append((tag, med, peak, lp)) + spd = f"{bmed/med:.2f}x" if med else "-" + print(f" {tag:20s} {med:.3f}s ({spd} vs base) peak={peak:.1f}G LPIPS={lp}", flush = True) + + print("\n==== SUMMARY (ref = baseline compile) ====", flush = True) + for tag, med, peak, lp in rows: + if med is None: + print(f" {tag:20s} FAILED") + continue + spd = f"{bmed/med:.2f}x" if med else "-" + lpv = "ref" if (tag == "baseline") else (f"{lp:.3f}" if lp is not None else "n/a") + print(f" {tag:20s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush = True) + print("PERF-LEVERS-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index aac4069833..b61badb50b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -52,6 +52,10 @@ from .diffusion_speed import ( restore_backend_flags, snapshot_backend_flags, ) +from .diffusion_attention import ( + apply_attention_backend, + select_attention_backend, +) from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( load_prequantized_transformer, @@ -96,6 +100,9 @@ class _LoadState: # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. transformer_quant: Optional[str] = None + # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or + # None for the default SDPA. Set before compile; orthogonal to the weight quant. + attention_backend: Optional[str] = None @dataclass @@ -306,6 +313,7 @@ class DiffusionBackend: transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -340,6 +348,7 @@ class DiffusionBackend: transformer_quant = transformer_quant, transformer_quant_fast_accum = transformer_quant_fast_accum, transformer_prequant_path = transformer_prequant_path, + attention_backend = attention_backend, _load_token = token, ), daemon = True, @@ -476,6 +485,7 @@ class DiffusionBackend: transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -596,6 +606,18 @@ class DiffusionBackend: # first so unload can restore them: TF32 / cudnn.benchmark are global, # and a later `off` load must not inherit this load's settings. backend_flags_before = snapshot_backend_flags() + # Pick the attention kernel BEFORE compile (compile traces attention). auto + # upgrades to cuDNN fused attention on NVIDIA when a speed profile is active + # (~1.18x, near-lossless); an explicit backend is honored, falling back to + # the diffusers default if its kernel is unavailable. Orthogonal to the + # weight quant -- it speeds the QK/PV matmuls torchao does not touch. + attention_engaged = apply_attention_backend( + pipe, + select_attention_backend( + target, attention_backend, speed_active = effective_speed != SPEED_OFF + ), + logger = logger, + ) speed_applied = apply_speed_optims( pipe, target, @@ -649,6 +671,7 @@ class DiffusionBackend: backend_flags_before = backend_flags_before, text_encoder_quant = te_quant, transformer_quant = transformer_quant_engaged, + attention_backend = attention_engaged, ) logger.info( @@ -953,6 +976,7 @@ class DiffusionBackend: "speed_optims": [], "text_encoder_quant": None, "transformer_quant": None, + "attention_backend": None, } return { "loaded": True, @@ -969,6 +993,7 @@ class DiffusionBackend: "speed_optims": list(state.speed_optims), "text_encoder_quant": state.text_encoder_quant, "transformer_quant": state.transformer_quant, + "attention_backend": state.attention_backend, } diff --git a/studio/backend/core/inference/diffusion_attention.py b/studio/backend/core/inference/diffusion_attention.py new file mode 100644 index 0000000000..e652b0068c --- /dev/null +++ b/studio/backend/core/inference/diffusion_attention.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Select the diffusion transformer's attention backend. + +diffusers exposes a unified ``transformer.set_attention_backend(name)`` dispatcher that +swaps the scaled-dot-product-attention kernel, validating hardware/package requirements at +set time and otherwise leaving the default (``native`` = ``F.scaled_dot_product_attention``). +Attention is memory-bandwidth bound, so a better kernel is a real end-to-end win that is +orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never +touches) and composes with torch.compile. + + auto - the best *exact* (non-quantized) backend for the device. On NVIDIA CUDA that is + cuDNN's fused attention (``_native_cudnn``), measured ~1.18x end-to-end on a B200 + with LPIPS ~0.004 vs the default (below the compile/quant noise floor). On + AMD/Intel/Apple/CPU it stays ``native`` (the dispatcher already routes those). + ``auto`` only upgrades when a speed profile is active, so ``speed_mode=off`` stays + bit-identical. + native - force the default SDPA (bit-identical reference). + cudnn - cuDNN fused attention (exact; NVIDIA). + flash / flash3 / flash4 - FlashAttention 2 / 3 (Hopper) / 4 (SM100); exact, kernel-gated. + sage - SageAttention (INT8 QK); quantized, a small quality cost, consumer-friendly. + xformers / aiter - memory-efficient (NVIDIA) / AITER (AMD ROCm). + +Best-effort: an unavailable backend (missing kernel / wrong arch) is caught and the load +falls back to the diffusers default rather than failing. torch/diffusers imported lazily. +""" + +from __future__ import annotations + +from typing import Any, Optional + +ATTN_AUTO = "auto" +ATTN_NATIVE = "native" + +# User-facing alias -> the diffusers dispatcher backend name. +_ALIASES: dict[str, str] = { + "native": "native", + "sdpa": "native", + "cudnn": "_native_cudnn", + "flash": "flash", + "flash2": "flash", + "flash3": "_flash_3_hub", + "flash4": "flash_4_hub", + "sage": "sage", + "xformers": "xformers", + "aiter": "aiter", +} +ATTN_ALIASES = (ATTN_AUTO,) + tuple(dict.fromkeys(_ALIASES)) + + +def normalize_attention_backend(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested attention backend; None / "" / "auto" -> "auto". + + Raises ValueError for an unsupported alias so a bad request is rejected cheaply.""" + if value is None: + return ATTN_AUTO + normalized = str(value).strip().lower() + if not normalized: + return ATTN_AUTO + if normalized not in ATTN_ALIASES: + raise ValueError( + f"Unsupported attention_backend '{value}'. Use one of: {', '.join(ATTN_ALIASES)}." + ) + return normalized + + +# Backends diffusers validates only by *package* at set time (``_check_attention_backend_ +# requirements`` checks the ``kernels`` install, not the GPU), but whose kernels need a +# specific CUDA arch at run time -- so an explicit request on the wrong card loads/sets fine +# and then crashes mid-generation. Gate them up front by a (min, max-exclusive) compute +# capability range. FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so it +# needs an upper bound: an explicit flash3 on a B200 (SM100) must drop to native instead of +# setting fine then crashing at generation. FlashAttention 4 is Blackwell+ (no upper bound). +_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = { + "_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only + "flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+ +} + + +def _cuda_capability() -> Optional[tuple[int, int]]: + """(major, minor) compute capability of the active CUDA device, or None if unknown.""" + try: + import torch + if not torch.cuda.is_available(): + return None + return tuple(torch.cuda.get_device_capability()) # type: ignore[return-value] + except Exception: # noqa: BLE001 + return None + + +def _backend_arch_supported(backend: str) -> bool: + """False only when ``backend`` needs a CUDA arch outside this device's supported range. + + Unknown capability (no CUDA / detection failure) returns True so we never block on a + guess -- diffusers' own set-time check still guards the package, and a genuine run-time + failure falls back to native.""" + bounds = _ARCH_CAPABILITY.get(backend) + if bounds is None: + return True + have = _cuda_capability() + if have is None: + return True + low, high = bounds + return have >= low and (high is None or have < high) + + +def _is_cuda_nvidia(target: Any) -> bool: + """CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies.""" + if getattr(target, "device", None) != "cuda": + return False + try: + import torch + return getattr(torch.version, "hip", None) is None + except Exception: # noqa: BLE001 + return False + + +def select_attention_backend( + target: Any, requested: Optional[str], *, speed_active: bool +) -> Optional[str]: + """The dispatcher backend name to apply, or None to leave the diffusers default. + + An explicit alias is honored verbatim (apply falls back if its kernel is unavailable). + ``auto`` upgrades to cuDNN on NVIDIA CUDA only when a speed profile is active (so + ``off`` stays bit-identical); everywhere else it returns None (native default).""" + alias = normalize_attention_backend(requested) + if alias != ATTN_AUTO: + backend = _ALIASES[alias] + if backend == "native": + return None + # An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine + # then crash mid-generation, so drop it to the native default up front. + if not _backend_arch_supported(backend): + return None + # cuDNN fused SDPA needs Ampere+ (SM80); diffusers accepts it on pre-SM80 cards + # (T4/V100) then fails at the first generation, so apply the same gate to an + # explicit cuDNN request as the auto path already does. + if backend == "_native_cudnn" and not _cudnn_attention_supported(): + return None + return backend + # auto + if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported(): + return "_native_cudnn" + return None + + +def _cudnn_attention_supported() -> bool: + """cuDNN fused SDPA needs Ampere+ (SM80). On pre-SM80 NVIDIA cards (T4 SM75 / + V100 SM70) diffusers accepts ``_native_cudnn`` at set time but the kernel fails at + the first generation, so gate the auto-cuDNN upgrade on capability. Unknown + capability allows it (diffusers' set-time check + the run-time fallback still guard).""" + have = _cuda_capability() + return have is None or have >= (8, 0) + + +def apply_attention_backend( + pipe: Any, + backend: Optional[str], + *, + logger: Any = None, +) -> Optional[str]: + """Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher. + + Returns the backend actually engaged, or None when left at the native default (either + because ``backend`` was None or because the requested kernel was unavailable -> graceful + fallback, never a load failure). + + diffusers keeps a *process-wide* active attention backend that ``set_attention_backend`` + also updates, and a fresh transformer's processors follow it (their ``_attention_backend`` + defaults to None). So a load that wants native must restore it explicitly: otherwise it + silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile), + breaking the bit-identical/``off`` guarantee. Best-effort throughout.""" + transformer = getattr(pipe, "transformer", None) + fn = getattr(transformer, "set_attention_backend", None) + if not callable(fn): + return None + if backend is not None: + try: + fn(backend) + # set_attention_backend also pins the backend in diffusers' process-wide + # registry. This transformer's own processors keep it locally (their + # _attention_backend is now explicit), so reset the global default back to + # native -- otherwise a later component whose processors are unconfigured + # (backend None) silently inherits this kernel. + _reset_global_backend_to_native(logger) + if logger is not None: + logger.info("diffusion.attention: backend=%s", backend) + return backend + except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below + _warn(logger, backend, exc) + # No backend requested, or the requested one failed: pin the native default so a stale + # process-wide backend from a previous load can't leak into this one. + _restore_native_backend(fn, logger) + return None + + +def _active_attention_backend() -> Optional[str]: + """The diffusers process-wide active attention backend name, or None if undeterminable.""" + try: + from diffusers.models.attention_dispatch import _AttentionBackendRegistry + + # get_active_backend() returns a (AttentionBackendName, fn) tuple (or None), so + # take element 0 and read its .value (e.g. "native"); reading .value off the + # tuple itself would yield a junk string that never compares equal to a name. + active = _AttentionBackendRegistry.get_active_backend() + if active is None: + return None + name = active[0] if isinstance(active, tuple) else active + return getattr(name, "value", str(name)) + except Exception: # noqa: BLE001 + return None + + +def _reset_global_backend_to_native(logger: Any) -> None: + """Reset diffusers' process-wide active attention backend to native after a + successful per-transformer set, so a later component whose processors are + unconfigured (backend None) does not inherit this transformer's kernel. The + transformer's own processors keep the backend just set. Best-effort and silent: + if the diffusers internals move, the prior (leaking) behavior is unchanged.""" + if _active_attention_backend() == ATTN_NATIVE: + return + try: + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + except Exception: # noqa: BLE001 — best-effort; leave the global as-is on any change + pass + + +def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None: + """Force the native default when the global active backend isn't already native.""" + if _active_attention_backend() == ATTN_NATIVE: + return # already native -> avoid redundant work and an extra dispatcher warning + try: + set_backend_fn(ATTN_NATIVE) + except Exception as exc: # noqa: BLE001 — best-effort restore + _warn(logger, ATTN_NATIVE, exc) + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 08d309f257..402c931da1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1750,6 +1750,30 @@ class DiffusionLoadRequest(BaseModel): "the OS path separator). A bare on/off value such as '1' is deliberately not " "accepted -- it must name an allowed directory.", ) + attention_backend: Optional[ + Literal[ + "auto", + "native", + "sdpa", + "cudnn", + "flash", + "flash2", + "flash3", + "flash4", + "sage", + "xformers", + "aiter", + ] + ] = Field( + None, + description = "Attention kernel via the diffusers dispatcher. auto picks the best " + "exact backend for the device (cuDNN fused attention on NVIDIA, ~1.18x and " + "near-lossless, when a speed profile is active; native SDPA elsewhere and when " + "speed=off). native (alias sdpa) forces default SDPA; cudnn/flash/flash3/flash4 are exact " + "(kernel/arch-gated); sage is INT8 attention (a small quality cost, consumer " + "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An " + "unavailable kernel falls back to the default.", + ) class DiffusionGenerateRequest(BaseModel): @@ -1865,3 +1889,8 @@ class DiffusionStatusResponse(BaseModel): description = "Transformer quant engaged on the dense fast path: int8 | fp8 | " "nvfp4 | mxfp8 | null (null = the GGUF transformer was loaded)", ) + attention_backend: Optional[str] = Field( + None, + description = "Attention backend engaged via the diffusers dispatcher (e.g. " + "_native_cudnn), or null for the default SDPA", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ec5382e190..c89a0cc7eb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10334,6 +10334,7 @@ async def load_diffusion_model( transformer_quant = request.transformer_quant, transformer_quant_fast_accum = request.transformer_quant_fast_accum, transformer_prequant_path = request.transformer_prequant_path, + attention_backend = request.attention_backend, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_attention.py b/studio/backend/tests/test_diffusion_attention.py new file mode 100644 index 0000000000..3e43aab8a9 --- /dev/null +++ b/studio/backend/tests/test_diffusion_attention.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for attention-backend selection. No torch/diffusers needed: +``_is_cuda_nvidia`` is monkeypatched for the policy tests, and the apply path uses a fake +transformer that records / raises on ``set_attention_backend``. +""" + +from __future__ import annotations + +import types + +import pytest + +import core.inference.diffusion_attention as att +from core.inference.diffusion_attention import ( + ATTN_AUTO, + apply_attention_backend, + normalize_attention_backend, + select_attention_backend, +) + + +def _target(device = "cuda"): + return types.SimpleNamespace(device = device) + + +# ── normalize ──────────────────────────────────────────────────────────────────── +def test_normalize_defaults_and_aliases(): + assert normalize_attention_backend(None) == ATTN_AUTO + assert normalize_attention_backend("") == ATTN_AUTO + assert normalize_attention_backend("auto") == ATTN_AUTO + assert normalize_attention_backend("CuDNN") == "cudnn" + assert normalize_attention_backend("FLASH3") == "flash3" + assert normalize_attention_backend("sdpa") == "sdpa" + + +def test_normalize_rejects_unknown(): + with pytest.raises(ValueError): + normalize_attention_backend("bogus") + # dashes are no longer silently rewritten to underscores -> a dashed alias is rejected. + with pytest.raises(ValueError): + normalize_attention_backend("flash-3") + + +def test_sdpa_alias_maps_to_native(): + # sdpa is an alias for native -> nothing to set on the dispatcher. + assert select_attention_backend(_target(), "sdpa", speed_active = True) is None + + +# ── select policy ───────────────────────────────────────────────────────────────── +def test_auto_upgrades_to_cudnn_on_nvidia_when_speed_active(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) # Ampere+: cuDNN ok + assert select_attention_backend(_target(), "auto", speed_active = True) == "_native_cudnn" + + +def test_auto_does_not_pin_cudnn_below_sm80(monkeypatch): + # cuDNN fused SDPA fails at run time on pre-SM80 (T4 SM75 / V100 SM70); auto must stay + # on the native default there rather than pin a backend that crashes on first generation. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) # Turing T4 + assert select_attention_backend(_target(), "auto", speed_active = True) is None + + +def test_auto_stays_native_when_speed_off(monkeypatch): + # off must stay bit-identical -> no backend change even on NVIDIA. + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: True) + assert select_attention_backend(_target(), "auto", speed_active = False) is None + + +def test_auto_stays_native_off_nvidia(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) + assert select_attention_backend(_target(device = "mps"), "auto", speed_active = True) is None + + +def test_explicit_backend_honored_regardless_of_speed(monkeypatch): + monkeypatch.setattr(att, "_is_cuda_nvidia", lambda target: False) + # Pin a high capability so the arch-gated flash4 isn't dropped by the runtime check. + monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0)) + assert select_attention_backend(_target(), "sage", speed_active = False) == "sage" + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" + + +def test_explicit_native_returns_none(): + # native is the default -> nothing to set. + assert select_attention_backend(_target(), "native", speed_active = True) is None + + +# ── arch gating (flash3/flash4 need a specific CUDA capability) ───────────────────── +def test_flash3_dropped_below_hopper(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 9)) # Ada / consumer + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + + +def test_flash4_dropped_below_blackwell(monkeypatch): + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) # Hopper, but FA4 needs SM100 + assert select_attention_backend(_target(), "flash4", speed_active = False) is None + # flash3 still allowed on Hopper. + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +def test_arch_gate_does_not_block_when_capability_unknown(monkeypatch): + # Unknown capability (e.g. no CUDA) must not block -> diffusers' set-time check still guards. + monkeypatch.setattr(att, "_cuda_capability", lambda: None) + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + + +def test_flash3_dropped_on_blackwell(monkeypatch): + # FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel: an explicit + # flash3 on a B200 (SM100) must drop to native rather than set fine then crash. + monkeypatch.setattr(att, "_cuda_capability", lambda: (10, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) is None + # FA4 is still honored on Blackwell. + assert select_attention_backend(_target(), "flash4", speed_active = False) == "flash_4_hub" + # flash3 is allowed exactly on Hopper SM90. + monkeypatch.setattr(att, "_cuda_capability", lambda: (9, 0)) + assert select_attention_backend(_target(), "flash3", speed_active = False) == "_flash_3_hub" + + +def test_explicit_cudnn_dropped_below_sm80(monkeypatch): + # An explicit cuDNN request on pre-Ampere (T4 SM75 / V100 SM70) must drop to native, + # not set fine and crash at first generation -- the same gate the auto path applies. + monkeypatch.setattr(att, "_cuda_capability", lambda: (7, 5)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) is None + # Ampere+ still honors it. + monkeypatch.setattr(att, "_cuda_capability", lambda: (8, 0)) + assert select_attention_backend(_target(), "cudnn", speed_active = False) == "_native_cudnn" + + +# ── apply ───────────────────────────────────────────────────────────────────────── +class _FakeTransformer: + def __init__(self, *, fail = False): + self.fail = fail + self.set_to = None + + def set_attention_backend(self, name): + if self.fail: + raise RuntimeError(f"{name} kernel unavailable") + self.set_to = name + + +def _pipe(transformer): + return types.SimpleNamespace(transformer = transformer) + + +def test_apply_none_leaves_native_when_global_already_native(monkeypatch): + # Global already native -> no redundant set call, returns None. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to is None + + +def test_apply_none_restores_native_when_global_polluted(monkeypatch): + # A previous load pinned cuDNN process-wide; a native load must reset it so it can't + # silently inherit cuDNN (the bit-identical/off guarantee). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + t = _FakeTransformer() + assert apply_attention_backend(_pipe(t), None) is None + assert t.set_to == "native" + + +def test_apply_sets_backend(): + t = _FakeTransformer() + engaged = apply_attention_backend(_pipe(t), "_native_cudnn") + assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" + + +def test_apply_falls_back_on_unavailable_kernel(monkeypatch): + # an unavailable kernel must not fail the load -> returns None (diffusers default). + monkeypatch.setattr(att, "_active_attention_backend", lambda: "native") + t = _FakeTransformer(fail = True) + assert apply_attention_backend(_pipe(t), "sage") is None + + +def test_apply_failed_kernel_restores_native_when_polluted(monkeypatch): + # Requested kernel fails AND the global is polluted: restore native before returning. + monkeypatch.setattr(att, "_active_attention_backend", lambda: "_native_cudnn") + + class _FailOnceTransformer: + def __init__(self): + self.calls = [] + + def set_attention_backend(self, name): + self.calls.append(name) + if name != "native": + raise RuntimeError(f"{name} kernel unavailable") + + t = _FailOnceTransformer() + assert apply_attention_backend(_pipe(t), "sage") is None + assert t.calls == ["sage", "native"] + + +def test_apply_handles_missing_method(): + pipe = types.SimpleNamespace(transformer = types.SimpleNamespace()) + assert apply_attention_backend(pipe, "_native_cudnn") is None + + +def test_apply_resets_global_registry_after_success(monkeypatch): + # After a successful per-transformer set, the process-wide registry must be reset to + # native so a later component (unconfigured processors) can't inherit this kernel -- + # while the transformer's own backend stays the engaged one. + called = {"reset": False} + monkeypatch.setattr( + att, "_reset_global_backend_to_native", lambda logger: called.__setitem__("reset", True) + ) + t = _FakeTransformer() + engaged = apply_attention_backend(_pipe(t), "_native_cudnn") + assert engaged == "_native_cudnn" and t.set_to == "_native_cudnn" + assert called["reset"] is True + + +def test_active_attention_backend_reads_tuple_return(): + # get_active_backend() returns a (AttentionBackendName, fn) tuple; the helper must read + # the name's .value, not stringify the tuple (which never compares equal to a name). + pytest.importorskip("diffusers") + from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + ) + + _AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE) + assert att._active_attention_backend() == "native" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 839d325443..dc34e60cef 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -424,6 +424,29 @@ def test_transformer_prequant_path_threads_through(client, monkeypatch): assert backend.last_load_kwargs.get("transformer_prequant_path") == "/data/zimage_fp8.pt" +def test_attention_backend_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "attention_backend": "cudnn", + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("attention_backend") == "cudnn" + + +def test_invalid_attention_backend_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "attention_backend": "bogus"}, + ) + assert resp.status_code == 422 + + def test_prequant_path_doc_describes_allowlist_not_toggle(): # The field help must match the code: UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH is a # directory allowlist, not a =1 toggle (diffusion_prequant._allowed_prequant_roots From 33df9d7b8d5660aeb5bd65591dff1c4950bea4b3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:39:50 -0700 Subject: [PATCH 07/16] Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder (#6702) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../inference/diffusion_transformer_quant.py | 47 +++++++++++++---- .../tests/test_diffusion_transformer_quant.py | 52 +++++++++++++++++-- 2 files changed, 86 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index 7cb1074425..cb6873a661 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -55,6 +55,11 @@ DEFAULT_MIN_LINEAR_FEATURES = 512 # (0.81x end-to-end on Z-Image 1024px) AND notably less accurate (LPIPS 0.166 vs fp8's # 0.044), because FP4's per-forward quant overhead is not amortised and the format is # coarser. So nvfp4 is kept as an explicit opt-in, never the auto pick for diffusion. +# +# This order is the DATA-CENTER preference. On a consumer / workstation GPU it is reordered +# to put int8 first (see ``_prefer_consumer_scheme``): consumer cards halve fp8/fp16 FP32- +# accumulate throughput, while int8 runs full-rate (int32 accumulate is not nerfed), so int8 +# is as fast or faster than fp8 on every consumer NVIDIA / AMD / Intel part. _AUTO_LADDER: tuple[tuple[tuple[int, int], tuple[str, ...]], ...] = ( ((10, 0), (TQ_FP8, TQ_NVFP4, TQ_MXFP8, TQ_INT8)), # Blackwell sm_100+ ((8, 9), (TQ_FP8, TQ_INT8)), # Ada sm_89 / Hopper sm_90 @@ -72,13 +77,15 @@ _DATACENTER_GPU_TOKENS = frozenset( { "B200", "B100", + "B300", # Blackwell Ultra data center "GB200", "GB300", "GB10", # Blackwell data center "H200", "H100", "H800", - "H20", # Hopper data center + "H20", + "GH200", # Grace-Hopper superchip (data center) "A100", "A800", "A30", @@ -99,15 +106,22 @@ _DATACENTER_GPU_TOKENS = frozenset( ) +# Professional parts the rest of the backend treats as datacenter-class (see llama_cpp.py +# _DATACENTER_GPU_RE, which applies the same FP32-accum tuning to them). Matched as phrases +# because the marker spans tokens ("RTX PRO 6000", "RTX 6000 ADA"), so they must not be +# misread as consumer (which would put int8 ahead of fp8 and pick fast accumulate). +_PROFESSIONAL_GPU_MARKERS = ("RTX PRO 6000", "RTX 6000 ADA") + + def _is_consumer_gpu(device: Any = None) -> bool: - """Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32 - accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center - HBM parts (recognised by name token) are not nerfed and return False, so they keep - the higher-precision default accumulate for free. Heuristic on the device name: a - GeForce / TITAN name is always consumer; a recognised data-center token is not; - anything else (workstation RTX, unknown) defaults to consumer -- the safe choice, - since fast accumulate is free on data-center and a win on consumer. Best-effort: - True on any probe failure.""" + """Whether the active GPU is consumer-class (GDDR), where fp8 FP32 accumulate is + throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center HBM parts and + professional parts (recognised by name) are not nerfed and return False, so they keep + the higher-precision default accumulate and fp8 first. Heuristic on the device name: a + GeForce / TITAN name is always consumer; a recognised data-center token or professional + marker is not; anything else (unknown) defaults to consumer -- the safe choice, since + fast accumulate is free on data-center and a win on consumer. Best-effort: True on any + probe failure.""" try: import re @@ -117,6 +131,8 @@ def _is_consumer_gpu(device: Any = None) -> bool: return True if "GEFORCE" in name or "TITAN" in name: return True + if any(marker in name for marker in _PROFESSIONAL_GPU_MARKERS): + return False tokens = set(re.split(r"[^A-Z0-9]+", name)) return not (tokens & _DATACENTER_GPU_TOKENS) @@ -168,13 +184,24 @@ def select_transformer_quant_scheme(target: Any, requested: Optional[str]) -> Op return None for floor, schemes in _AUTO_LADDER: if cap >= floor: - for scheme in schemes: + for scheme in _prefer_consumer_scheme(schemes, device): if _scheme_supported(scheme, device): return scheme return None return None +def _prefer_consumer_scheme(schemes: tuple[str, ...], device: Any) -> tuple[str, ...]: + """Reorder an arch tier's schemes for the GPU class. On a consumer / workstation card + move int8 to the front: consumer parts halve fp8/fp16 FP32-accumulate throughput, while + int8 runs at full rate (int32 accumulate is not nerfed), so int8 is as fast or faster + than fp8 on every consumer NVIDIA / AMD / Intel GPU (and the only path on pre-Ada + consumer without fp8 tensor cores). Data-center HBM parts keep fp8 first.""" + if TQ_INT8 in schemes and schemes[0] != TQ_INT8 and _is_consumer_gpu(device): + return (TQ_INT8,) + tuple(s for s in schemes if s != TQ_INT8) + return schemes + + def _capability() -> Optional[tuple[int, int]]: try: import torch diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index d2adaa9f7c..d527723b2e 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -39,6 +39,7 @@ def _stub_torch( cc = (10, 0), with_fp8 = True, cuda_available = True, + device_name = "NVIDIA B200", ): torch = types.ModuleType("torch") torch.bfloat16 = "bfloat16" @@ -48,6 +49,9 @@ def _stub_torch( torch.cuda = types.SimpleNamespace( is_available = lambda: cuda_available, get_device_capability = lambda *a: cc, + # data-center name by default so the ladder tests get the data-center order; + # consumer tests pass a GeForce name (or monkeypatch _is_consumer_gpu). + get_device_name = lambda *a: device_name, ) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -104,11 +108,49 @@ def test_auto_blackwell_prefers_fp8_then_falls_back(monkeypatch): assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 +def test_auto_consumer_blackwell_prefers_int8(monkeypatch): + # Consumer Blackwell (RTX 50xx): fp8 FP32-accumulate is throughput-halved while int8 is + # full-rate, so auto prefers int8 even though fp8 is available (the data-center default). + _stub_torch(monkeypatch, cc = (10, 0), device_name = "NVIDIA GeForce RTX 5090") + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + # int8 unavailable -> falls back to the rest of the tier (fp8 next). + _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + +def test_auto_consumer_ada_prefers_int8(monkeypatch): + # Consumer Ada (RTX 4090): int8 runs ~2x fp8's nerfed FP32-accumulate rate. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA GeForce RTX 4090") + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_workstation_unknown_prefers_int8(monkeypatch): + # Unknown / workstation name -> treated as consumer (the safe default) -> int8 first. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA RTX A5000") + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_INT8 + + +def test_auto_professional_rtx_prefers_fp8(monkeypatch): + # Professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) are classified datacenter + # by the rest of the backend, so auto keeps fp8 first (not int8) -- matching llama_cpp. + for device_name, cc in ( + ("NVIDIA RTX PRO 6000 Blackwell Server Edition", (10, 0)), + ("NVIDIA RTX 6000 Ada Generation", (8, 9)), + ): + _stub_torch(monkeypatch, cc = cc, device_name = device_name) + _allow(monkeypatch, {TQ_FP8, TQ_INT8}) + assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 + + def test_auto_ada_hopper_prefers_fp8(monkeypatch): - _stub_torch(monkeypatch, cc = (8, 9)) + # Data-center Ada (L40S) / Hopper (H100): not nerfed -> fp8 first. + _stub_torch(monkeypatch, cc = (8, 9), device_name = "NVIDIA L40S") _allow(monkeypatch, {TQ_NVFP4, TQ_MXFP8, TQ_FP8, TQ_INT8}) assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 - _stub_torch(monkeypatch, cc = (9, 0)) # Hopper + _stub_torch(monkeypatch, cc = (9, 0), device_name = "NVIDIA H100 80GB HBM3") # Hopper assert select_transformer_quant_scheme(_target(), "auto") == TQ_FP8 @@ -224,7 +266,7 @@ def _stub_device_name(monkeypatch, name): "NVIDIA GeForce RTX 5090", "NVIDIA GeForce RTX 4090", "NVIDIA RTX A4000", # workstation: A4000 token, NOT the data-center A40 - "NVIDIA RTX 6000 Ada Generation", + "NVIDIA RTX A5000", # workstation: A5000 token, not professional/datacenter "NVIDIA Some Future Card 9000", # unknown -> default consumer (fast accum is free on DC) ], ) @@ -237,12 +279,16 @@ def test_is_consumer_gpu_true(monkeypatch, name): "name", [ "NVIDIA B200", + "NVIDIA B300", # Blackwell Ultra (matches llama_cpp datacenter regex) + "NVIDIA GH200 480GB", # Grace-Hopper superchip (was misread as consumer) "NVIDIA H100 80GB HBM3", "NVIDIA A100-SXM4-80GB", "NVIDIA A40", # data-center Ampere (distinct token from RTX A4000) "NVIDIA L40S", "NVIDIA L4", "Tesla V100-SXM2-16GB", + "NVIDIA RTX PRO 6000 Blackwell Server Edition", # professional -> datacenter-class + "NVIDIA RTX 6000 Ada Generation", # professional -> datacenter-class ], ) def test_is_consumer_gpu_false_for_datacenter(monkeypatch, name): From e16fd32b5006fabb40b0fd6c0c1c4e21c43e9920 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:40:52 -0700 Subject: [PATCH 08/16] Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT (#6703) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/fbcache_flux_probe.py | 180 ++++++++++++++++++ studio/backend/core/inference/diffusion.py | 27 +++ .../backend/core/inference/diffusion_cache.py | 110 +++++++++++ .../backend/core/inference/diffusion_speed.py | 13 +- studio/backend/models/inference.py | 18 ++ studio/backend/routes/inference.py | 2 + studio/backend/tests/test_diffusion_cache.py | 149 +++++++++++++++ studio/backend/tests/test_diffusion_routes.py | 41 ++++ 8 files changed, 538 insertions(+), 2 deletions(-) create mode 100644 scripts/fbcache_flux_probe.py create mode 100644 studio/backend/core/inference/diffusion_cache.py create mode 100644 studio/backend/tests/test_diffusion_cache.py diff --git a/scripts/fbcache_flux_probe.py b/scripts/fbcache_flux_probe.py new file mode 100644 index 0000000000..c0462f1947 --- /dev/null +++ b/scripts/fbcache_flux_probe.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Validate First-Block-Cache (FBCache) on a MANY-step DiT (Flux.1-dev), vs the compiled +baseline. FBCache reuses the transformer tail across denoise steps when the first block's +residual barely changes -- a real speedup only when there are enough steps (it is why it is +gated OFF for few-step distilled models like Z-Image-Turbo). Reports median latency, +speedup, peak VRAM, and LPIPS vs the no-cache baseline. One CUDA GPU.""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +BASE = "black-forest-labs/FLUX.1-dev" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "fbcache_flux_images" + + +_LP = {"fn": None} + + +def _lpips(ref, arr): + try: + import lpips + import torch + + if _LP["fn"] is None: + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() + + def t(x): + return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() + + with torch.no_grad(): + return float(_LP["fn"](t(ref), t(arr)).item()) + except Exception as exc: # noqa: BLE001 + print(f" (lpips: {type(exc).__name__})", flush = True) + return None + + +def _load(): + import os + import diffusers + import torch + + pipe = diffusers.FluxPipeline.from_pretrained( + BASE, torch_dtype = torch.bfloat16, token = os.environ.get("HF_TOKEN") + ) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res, guidance): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = guidance, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def run( + tag, + steps, + seed, + res, + guidance, + iters, + *, + threshold = None, + compile_ = True, +): + import torch + + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + pipe = _load() + if threshold is not None: + from diffusers import FirstBlockCacheConfig + try: + pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold = threshold)) + except Exception as exc: # noqa: BLE001 + from diffusers.hooks import apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold)) + if compile_: + # FBCache's per-step decision is a graph break, so a cached run must compile with + # fullgraph=False (mirroring the production path); fullgraph=True would fail the + # warmup compile and the row would silently fall back to an eager cached run, + # producing misleading speedup numbers. + fullgraph = threshold is None + try: + pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True) + except Exception as exc: # noqa: BLE001 + print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush = True) + try: + _gen(pipe, steps, seed, res, guidance) # warmup / compile + except Exception as exc: # noqa: BLE001 + import traceback + + traceback.print_exc() + print(f" [{tag}] FAILED: {type(exc).__name__}: {str(exc)[:100]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None + dts, img = [], None + for _ in range(iters): + img, dt = _gen(pipe, steps, seed, res, guidance) + dts.append(dt) + peak = torch.cuda.max_memory_allocated() / 1e9 + arr = np.array(img) + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() + return _median(dts), arr, peak + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 28) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--guidance", type = float, default = 3.5) + p.add_argument("--iters", type = int, default = 2) + args = p.parse_args(argv) + s, r, seed, gd, it = args.steps, args.res, args.seed, args.guidance, args.iters + + print(f"== FBCache on Flux.1-dev ({r}px, {s} steps, guidance {gd}) ==", flush = True) + base = run("baseline", s, seed, r, gd, it) + if base is None: + print("baseline FAILED", flush = True) + return 1 + bmed, ref, bpeak = base + print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush = True) + rows = [("baseline", bmed, bpeak, 0.0)] + for thr in (0.08, 0.12, 0.20): + out = run(f"fbcache_{thr}", s, seed, r, gd, it, threshold = thr) + if out is None: + rows.append((f"fbcache_{thr}", None, None, None)) + continue + med, arr, peak = out + lp = _lpips(ref, arr) + rows.append((f"fbcache_{thr}", med, peak, lp)) + print( + f" fbcache_{thr}: {med:.3f}s ({bmed/med:.2f}x) peak={peak:.1f}G LPIPS={lp}", flush = True + ) + + print("\n==== SUMMARY (Flux.1-dev, ref = no-cache compile) ====", flush = True) + for tag, med, peak, lp in rows: + if med is None: + print(f" {tag:16s} FAILED") + continue + spd = f"{bmed/med:.2f}x" + lpv = "ref" if tag == "baseline" else (f"{lp:.3f}" if lp is not None else "n/a") + print(f" {tag:16s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush = True) + print("FBCACHE-FLUX-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index b61badb50b..cba3ba2af5 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -56,6 +56,7 @@ from .diffusion_attention import ( apply_attention_backend, select_attention_backend, ) +from .diffusion_cache import apply_step_cache from .diffusion_precision import quantize_text_encoders from .diffusion_prequant import ( load_prequantized_transformer, @@ -103,6 +104,8 @@ class _LoadState: # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or # None for the default SDPA. Set before compile; orthogonal to the weight quant. attention_backend: Optional[str] = None + # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. + transformer_cache: Optional[str] = None @dataclass @@ -314,6 +317,8 @@ class DiffusionBackend: transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" fam = self.validate_load_request( @@ -349,6 +354,8 @@ class DiffusionBackend: transformer_quant_fast_accum = transformer_quant_fast_accum, transformer_prequant_path = transformer_prequant_path, attention_backend = attention_backend, + transformer_cache = transformer_cache, + transformer_cache_threshold = transformer_cache_threshold, _load_token = token, ), daemon = True, @@ -486,6 +493,8 @@ class DiffusionBackend: transformer_quant_fast_accum: Optional[bool] = None, transformer_prequant_path: Optional[str] = None, attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, _load_token: Optional[int] = None, ) -> dict[str, Any]: # Validate first (cheap, no torch/diffusers) so a direct call with a bad @@ -618,12 +627,27 @@ class DiffusionBackend: ), logger = logger, ) + # Opt-in step caching (First-Block-Cache), also before compile. OFF by + # default; for many-step models it reuses the transformer tail across steps + # (~1.4x on Flux at LPIPS ~0.08). When engaged, compile must drop fullgraph + # (the cache's per-step decision is a graph break), so pass it through. + cache_engaged = apply_step_cache( + pipe, + mode = transformer_cache, + threshold = transformer_cache_threshold, + # GGUF transformers are quantized too (the default Studio path), so the + # cache needs the higher quantized threshold to still trigger -- not just + # the dense-quant fast path. + quant_active = transformer_quant_engaged is not None or bool(gguf_filename), + logger = logger, + ) speed_applied = apply_speed_optims( pipe, target, is_gguf = bool(gguf_filename), family = fam, speed_mode = effective_speed, + cache_active = cache_engaged is not None, logger = logger, ) if transformer_quant_engaged is not None and not speed_applied.get("compiled"): @@ -672,6 +696,7 @@ class DiffusionBackend: text_encoder_quant = te_quant, transformer_quant = transformer_quant_engaged, attention_backend = attention_engaged, + transformer_cache = cache_engaged, ) logger.info( @@ -977,6 +1002,7 @@ class DiffusionBackend: "text_encoder_quant": None, "transformer_quant": None, "attention_backend": None, + "transformer_cache": None, } return { "loaded": True, @@ -994,6 +1020,7 @@ class DiffusionBackend: "text_encoder_quant": state.text_encoder_quant, "transformer_quant": state.transformer_quant, "attention_backend": state.attention_backend, + "transformer_cache": state.transformer_cache, } diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py new file mode 100644 index 0000000000..b7a2b7ac45 --- /dev/null +++ b/studio/backend/core/inference/diffusion_cache.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Opt-in step caching for the diffusion transformer (First-Block-Cache). + +Across denoising steps a DiT's output changes little once the trajectory settles, so most of +the transformer can be reused. First-Block-Cache (FBCache) computes the first block, and if +its residual barely changed from the previous step (within ``threshold``) it skips the +remaining blocks and reuses their cached output. diffusers ships it natively +(``transformer.enable_cache(FirstBlockCacheConfig(...))`` for CacheMixin models, or the +standalone ``apply_first_block_cache`` hook). + +Measured on Flux.1-dev (28 steps, 1024px, B200): ~1.4x on top of torch.compile (2.83 -> +2.03 s) at LPIPS ~0.08 vs the no-cache output -- deep inside the speed-for-quality bar. + +OFF by default and a deliberate per-load opt-in, because the win scales with step count: a +few-step distilled model (e.g. Z-Image-Turbo at ~8 steps) has almost no headroom and a +single skipped step is a large fraction of the trajectory, so caching is for many-step +models (Flux / Qwen-Image). It composes with torch.compile only with ``fullgraph=False`` +(the cache's compiler-disabled decision is a graph break), which the speed layer switches to +automatically when a cache is engaged. Best-effort: an incompatible model (e.g. a transformer +whose block signature the hook does not recognise) is caught and the load proceeds uncached. +torch / diffusers imported lazily. +""" + +from __future__ import annotations + +from typing import Any, Optional + +TC_OFF = "off" +TC_FBCACHE = "fbcache" +TC_MODES = (TC_FBCACHE,) + +# FBCache residual thresholds: higher skips more steps (faster, lower quality). The dense +# bf16 default; a quantised transformer shifts the residual distribution, so it needs a +# higher threshold for the cache to trigger at all (per ParaAttention's fp8 guidance). +DEFAULT_FBCACHE_THRESHOLD = 0.08 +QUANT_FBCACHE_THRESHOLD = 0.12 + + +def normalize_transformer_cache(value: Optional[str]) -> Optional[str]: + """Lower/strip a requested cache mode; None / "" / "none" / "off" -> None (disabled). + + 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("-", "_") + if not normalized or normalized in ("none", "off"): + return None + if normalized not in TC_MODES: + raise ValueError( + f"Unsupported transformer_cache '{value}'. Use one of: off, {', '.join(TC_MODES)}." + ) + return normalized + + +def apply_step_cache( + pipe: Any, + *, + mode: Optional[str], + threshold: Optional[float] = None, + quant_active: bool = False, + logger: Any = None, +) -> Optional[str]: + """Engage step caching on ``pipe.transformer``. Returns the mode actually engaged, or + None when disabled / unsupported (the load then runs uncached). ``threshold`` overrides + the default; ``quant_active`` raises the default so the cache still triggers on a + quantised transformer. Best-effort: never raises for an incompatible model.""" + mode = normalize_transformer_cache(mode) + if mode is None: + return None + transformer = getattr(pipe, "transformer", None) + if transformer is None: + return None + thr = ( + threshold + if threshold is not None + else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD) + ) + # Only engage via the transformer's native enable_cache (the diffusers CacheMixin path). + # That mixin is present exactly when the pipeline wraps the transformer call in a + # cache_context, which the First-Block-Cache hook requires at run time. The lower-level + # apply_first_block_cache hook would install on a non-CacheMixin transformer too (e.g. + # Z-Image), but its pipeline opens no cache_context, so the first generation would crash + # inside the hook -- so a model without enable_cache runs uncached per the best-effort + # contract instead of being reported as cached and then failing. + enable_cache = getattr(transformer, "enable_cache", None) + if not callable(enable_cache): + _warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)")) + return None + try: + from diffusers import FirstBlockCacheConfig + + config = FirstBlockCacheConfig(threshold = thr) + enable_cache(config) + try: + transformer._unsloth_step_cache = f"{mode}@{thr}" + except Exception: # noqa: BLE001 — marker is best-effort + pass + if logger is not None: + logger.info("diffusion.cache: %s engaged (threshold=%s)", mode, thr) + return mode + except Exception as exc: # noqa: BLE001 — incompatible model -> run uncached + _warn(logger, mode, exc) + return None + + +def _warn(logger: Any, what: str, exc: Exception) -> None: + if logger is not None: + logger.warning("diffusion.cache: %s unavailable (%s); running uncached", what, exc) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 9184f29f88..208cb8b513 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -147,6 +147,7 @@ def apply_speed_optims( is_gguf: bool, family: Any, speed_mode: str = SPEED_OFF, + cache_active: bool = False, logger: Any = None, ) -> dict[str, bool]: """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, @@ -181,7 +182,9 @@ def apply_speed_optims( # block, where eligible (now incl. the GGUF transformer). `max` opts into # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): - applied["compiled"] = _compile_repeated_blocks(pipe, logger, max_autotune = mode == SPEED_MAX) + applied["compiled"] = _compile_repeated_blocks( + pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active + ) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. @@ -210,6 +213,7 @@ def _compile_repeated_blocks( logger: Any, *, max_autotune: bool = False, + cache_active: bool = False, ) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) @@ -221,7 +225,12 @@ def _compile_repeated_blocks( # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled # block because its static output buffer is overwritten across denoise steps. - kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune} + # + # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is + # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip + # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the + # block still compiles. + kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune} if max_autotune: kwargs["mode"] = "max-autotune-no-cudagraphs" try: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 402c931da1..21e9760aed 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1774,6 +1774,23 @@ class DiffusionLoadRequest(BaseModel): "friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An " "unavailable kernel falls back to the default.", ) + transformer_cache: Optional[Literal["off", "fbcache"]] = Field( + None, + description = "Opt-in step caching (off by default). fbcache = First-Block-Cache: " + "reuse the transformer tail across denoise steps when the first block's residual " + "barely changes (~1.4x on Flux 28-step at LPIPS ~0.08). For MANY-step models " + "(Flux / Qwen-Image); leave off for few-step distilled models (e.g. Z-Image-Turbo), " + "which have no caching headroom. Composes with compile (drops fullgraph " + "automatically); incompatible models run uncached.", + ) + transformer_cache_threshold: Optional[float] = Field( + None, + ge = 0.0, + le = 1.0, + description = "FBCache residual threshold (higher = skips more steps = faster, lower " + "quality). null auto-picks 0.08 (0.12 when the transformer is quantised, which " + "shifts the residual distribution).", + ) class DiffusionGenerateRequest(BaseModel): @@ -1894,3 +1911,4 @@ class DiffusionStatusResponse(BaseModel): description = "Attention backend engaged via the diffusers dispatcher (e.g. " "_native_cudnn), or null for the default SDPA", ) + transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c89a0cc7eb..38dd2d1f20 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10335,6 +10335,8 @@ async def load_diffusion_model( transformer_quant_fast_accum = request.transformer_quant_fast_accum, transformer_prequant_path = request.transformer_prequant_path, attention_backend = request.attention_backend, + transformer_cache = request.transformer_cache, + transformer_cache_threshold = request.transformer_cache_threshold, ) return DiffusionStatusResponse(**status_dict) except (ValueError, FileNotFoundError) as exc: diff --git a/studio/backend/tests/test_diffusion_cache.py b/studio/backend/tests/test_diffusion_cache.py new file mode 100644 index 0000000000..62071d9aa6 --- /dev/null +++ b/studio/backend/tests/test_diffusion_cache.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for opt-in step caching (First-Block-Cache). + +``diffusers`` is stubbed via ``sys.modules`` (the module under test imports +``FirstBlockCacheConfig`` lazily), and the pipeline is a fake that records the engaged config. +So normalisation, the CacheMixin (``enable_cache``) gating, threshold selection, and the +best-effort failure handling are all exercised without torch or a real diffusers model. +""" + +from __future__ import annotations + +import sys +import types + +import pytest + +from core.inference.diffusion_cache import ( + DEFAULT_FBCACHE_THRESHOLD, + QUANT_FBCACHE_THRESHOLD, + TC_FBCACHE, + apply_step_cache, + normalize_transformer_cache, +) + + +# ── normalize_transformer_cache ──────────────────────────────────────────────────── +def test_normalize_disabled_values_are_none(): + for value in (None, "", " ", "none", "off", "OFF", "None"): + assert normalize_transformer_cache(value) is None + + +def test_normalize_fbcache_and_casing(): + assert normalize_transformer_cache("fbcache") == TC_FBCACHE + assert normalize_transformer_cache("FBCache") == TC_FBCACHE + assert normalize_transformer_cache(" fbcache ") == TC_FBCACHE + + +def test_normalize_rejects_unknown(): + with pytest.raises(ValueError): + normalize_transformer_cache("deepcache") + + +# ── apply_step_cache ─────────────────────────────────────────────────────────────── +class _Config: + def __init__(self, threshold): + self.threshold = threshold + + +class _MixinTransformer: + """A CacheMixin-style transformer: exposes ``enable_cache``.""" + + def __init__(self, *, fail = False): + self.fail = fail + self.enabled_with = None + + def enable_cache(self, config): + if self.fail: + raise RuntimeError("block signature not recognised") + self.enabled_with = config + + +class _NonCacheMixinTransformer: + """A transformer with no ``enable_cache`` (not a CacheMixin) -> must run uncached. + + Its pipeline opens no ``cache_context``, so installing FBCache would crash at generation; + the load runs uncached instead (e.g. Z-Image).""" + + +def _pipe(transformer): + return types.SimpleNamespace(transformer = transformer) + + +def _stub_diffusers(monkeypatch, *, hook_recorder = None): + diffusers = types.ModuleType("diffusers") + diffusers.FirstBlockCacheConfig = _Config + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + + hooks = types.ModuleType("diffusers.hooks") + + def _apply_first_block_cache(transformer, config): + if hook_recorder is not None: + hook_recorder["transformer"] = transformer + hook_recorder["config"] = config + + hooks.apply_first_block_cache = _apply_first_block_cache + monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks) + + +def test_disabled_mode_is_noop(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = None) is None + assert apply_step_cache(_pipe(t), mode = "off") is None + assert t.enabled_with is None + + +def test_enable_cache_path_default_threshold(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + engaged = apply_step_cache(_pipe(t), mode = "fbcache") + assert engaged == TC_FBCACHE + assert t.enabled_with.threshold == DEFAULT_FBCACHE_THRESHOLD + assert t._unsloth_step_cache == f"fbcache@{DEFAULT_FBCACHE_THRESHOLD}" + + +def test_quant_active_raises_default_threshold(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", quant_active = True) + assert t.enabled_with.threshold == QUANT_FBCACHE_THRESHOLD + + +def test_explicit_threshold_overrides_quant(monkeypatch): + _stub_diffusers(monkeypatch) + t = _MixinTransformer() + apply_step_cache(_pipe(t), mode = "fbcache", threshold = 0.2, quant_active = True) + assert t.enabled_with.threshold == 0.2 + + +def test_non_cachemixin_runs_uncached(monkeypatch): + # A transformer without enable_cache (e.g. Z-Image) must NOT install the standalone hook + # -- its pipeline opens no cache_context, so it runs uncached instead of crashing at gen. + rec: dict = {} + _stub_diffusers(monkeypatch, hook_recorder = rec) + t = _NonCacheMixinTransformer() + assert apply_step_cache(_pipe(t), mode = "fbcache") is None + assert rec == {} # the standalone hook was never called + + +def test_incompatible_model_runs_uncached(monkeypatch): + # enable_cache raising (e.g. unrecognised block signature) must not fail the load. + _stub_diffusers(monkeypatch) + t = _MixinTransformer(fail = True) + assert apply_step_cache(_pipe(t), mode = "fbcache") is None + + +def test_missing_transformer_is_none(monkeypatch): + _stub_diffusers(monkeypatch) + pipe = types.SimpleNamespace(transformer = None) + assert apply_step_cache(pipe, mode = "fbcache") is None + + +def test_diffusers_unavailable_runs_uncached(monkeypatch): + # no diffusers import -> best-effort returns None, load proceeds uncached. + monkeypatch.setitem(sys.modules, "diffusers", None) + t = _MixinTransformer() + assert apply_step_cache(_pipe(t), mode = "fbcache") is None diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index dc34e60cef..d03c48b477 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -460,6 +460,47 @@ def test_prequant_path_doc_describes_allowlist_not_toggle(): assert "allowlist" in desc.lower() or "director" in desc.lower() +def test_transformer_cache_threads_through(client, monkeypatch): + backend = _FakeBackend() + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache": "fbcache", + "transformer_cache_threshold": 0.1, + }, + ) + assert resp.status_code == 200 + assert backend.last_load_kwargs.get("transformer_cache") == "fbcache" + assert backend.last_load_kwargs.get("transformer_cache_threshold") == 0.1 + + +def test_invalid_transformer_cache_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache": "deepcache", + }, + ) + assert resp.status_code == 422 + + +def test_out_of_range_cache_threshold_returns_422(client): + resp = client.post( + "/api/inference/images/load", + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache_threshold": 1.5, + }, + ) + assert resp.status_code == 422 + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted. From 6209497fd1d019b85eba05788ce0eb5357710e66 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:41:49 -0700 Subject: [PATCH 09/16] Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) (#6716) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/build_prequant_checkpoint.py | 13 ++- scripts/int8_linear_probe.py | 88 +++++++++++++++++++ .../inference/diffusion_transformer_quant.py | 60 ++++++++++--- .../tests/test_diffusion_transformer_quant.py | 61 +++++++++++++ 4 files changed, 211 insertions(+), 11 deletions(-) create mode 100644 scripts/int8_linear_probe.py diff --git a/scripts/build_prequant_checkpoint.py b/scripts/build_prequant_checkpoint.py index 18dd064650..366822de94 100644 --- a/scripts/build_prequant_checkpoint.py +++ b/scripts/build_prequant_checkpoint.py @@ -57,6 +57,7 @@ def main(argv = None) -> int: from core.inference.diffusion_transformer_quant import ( TQ_SCHEMES, _make_quant_config, + exclude_tokens_for_scheme, make_filter_fn, ) from torchao.quantization import quantize_ @@ -78,7 +79,17 @@ def main(argv = None) -> int: args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token ).to("cuda") print(f" quantising in place ({scheme}) ...", flush = True) - quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features)) + # Mirror the runtime path EXACTLY (the offline == runtime, LPIPS-0 invariant): for int8 also + # skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint + # bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on + # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns (). + quantize_( + transformer, + _make_quant_config(scheme), + filter_fn = make_filter_fn( + args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme) + ), + ) # Move the state dict to CPU for a portable, GPU-free artifact. state_dict = { diff --git a/scripts/int8_linear_probe.py b/scripts/int8_linear_probe.py new file mode 100644 index 0000000000..5cef1b68e7 --- /dev/null +++ b/scripts/int8_linear_probe.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Examine which Linear layers the int8 dense-quant filter would select, to find the large M=1 +modulation/embedder projections that crash torch._int_mm (M>16). Loads each transformer on the +META device (no weights, no GPU) from its base-repo config, lists nn.Linear fqn/in/out, and marks +those that pass min_features=512. CPU-only, fast.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + +# (label, transformer_class, base_repo) +MODELS = [ + ("flux.1-dev", "FluxTransformer2DModel", "black-forest-labs/FLUX.1-dev"), + ("qwen-image", "QwenImageTransformer2DModel", "Qwen/Qwen-Image"), + ("z-image", "ZImageTransformer2DModel", "Tongyi-MAI/Z-Image-Turbo"), + ("flux.2-klein-4b", "Flux2Transformer2DModel", "black-forest-labs/FLUX.2-klein-4B"), +] +MIN = 512 + + +def main() -> int: + import diffusers + import torch + from accelerate import init_empty_weights + + tok = os.environ.get("HF_TOKEN") + for label, cls_name, base in MODELS: + cls = getattr(diffusers, cls_name, None) + if cls is None: + print(f"\n### {label}: {cls_name} NOT in diffusers") + continue + try: + cfg = cls.load_config(base, subfolder = "transformer", token = tok) + with init_empty_weights(): + model = cls.from_config(cfg) + except Exception as e: # noqa: BLE001 + print(f"\n### {label}: load failed {type(e).__name__}: {e}") + continue + lins = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)] + selected = [(n, m) for n, m in lins if m.in_features >= MIN and m.out_features >= MIN] + print(f"\n### {label}: {len(lins)} Linear, {len(selected)} pass min_features={MIN}") + # Heuristic: a modulation/embedder Linear is one OUTSIDE the repeated transformer blocks, + # i.e. its fqn does not contain a numeric block index, OR out==k*in (k>=3) AdaLN shape. + sus = [] + for n, m in selected: + depth_idx = any(p.isdigit() for p in n.split(".")) + ratio = m.out_features / m.in_features if m.in_features else 0 + tag = [] + if not depth_idx: + tag.append("NO-BLOCK-IDX") + if ratio >= 3: + tag.append(f"out={ratio:.0f}xin") + if any( + t in n.lower() + for t in ("norm", "embed", "time", "guidance", "modulation", "adaln", "cond") + ): + tag.append("NAME") + if tag: + sus.append((n, m.in_features, m.out_features, ",".join(tag))) + # Print the distinct fqn shapes (collapse block indices to {i}) + import re + + seen = {} + for n, i, o, tag in sus: + key = re.sub(r"\.\d+\.", ".{i}.", n) + seen.setdefault((key, i, o, tag), 0) + seen[(key, i, o, tag)] += 1 + print(f" SUSPECT (M=1 risk) distinct patterns:") + for (key, i, o, tag), cnt in sorted(seen.items()): + print(f" [{cnt:>3}x] {key:55s} {i:>6}->{o:<6} [{tag}]") + # Also show a few non-suspect selected names for contrast (the real FLOP linears) + good = [ + n + for n, m in selected + if (n, m.in_features, m.out_features) not in {(s[0], s[1], s[2]) for s in sus} + ][:6] + print(f" kept-for-int8 examples: {good}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py index cb6873a661..9e2e56a2a5 100644 --- a/studio/backend/core/inference/diffusion_transformer_quant.py +++ b/studio/backend/core/inference/diffusion_transformer_quant.py @@ -37,13 +37,42 @@ TQ_AUTO = "auto" TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8) TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES -# Skip linears whose in/out features are below this. The int8 dynamic path uses -# torch._int_mm, which requires the activation row count M > 16, and the DiT's tiny -# timestep / pooled / modulation projections run at M=1 and crash it. They are a -# negligible share of the FLOPs, so leaving them bf16 costs ~nothing (measured: -# 239/276 Z-Image linears quantised, full speedup) and keeps quality a touch higher. +# Skip linears whose in/out features are below this. A small share of the FLOPs, so leaving +# them bf16 costs ~nothing and keeps quality a touch higher. DEFAULT_MIN_LINEAR_FEATURES = 512 +# int8-ONLY name exclusions. The int8 dynamic path uses torch._int_mm, which requires the +# activation row count M > 16. A DiT's AdaLN *modulation* projections and its timestep / +# guidance / pooled-text *conditioning embedders* are computed once from the [batch, dim] +# conditioning vector (M = batch = 1), not per token -- so they crash _int_mm even though +# their feature dims are large (e.g. Flux's norm1.linear 3072->18432, Qwen's img_mod.1, Flux.2's +# *_modulation.linear). min_features does NOT catch them (they are big), so int8 also skips any +# Linear whose fqn matches one of these tokens. They are a negligible share of the FLOPs (run at +# M=1, once per block), so int8 keeps the full speedup on the attention/FFN layers (M = seq). +# fp8 / nvfp4 / mxfp8 use scaled_mm (no M>16 limit) and quantise these layers fine, so the +# exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) +# are deliberately NOT excluded. +_INT8_EXCLUDE_NAME_TOKENS = ( + "norm", # AdaLN modulation .linear (norm1 / norm1_context / norm / norm_out) + "_mod", # Qwen img_mod / txt_mod + "modulation", # Flux.2 double/single_stream_modulation + "timestep_embed", + "guidance_embed", + "time_text_embed", # Flux/Qwen time_text_embed.* (pooled-text + timestep); NOT context_embedder + "pooled", +) + + +def exclude_tokens_for_scheme(scheme: str) -> tuple[str, ...]: + """Name tokens to exclude from quantisation for ``scheme``. int8 (torch._int_mm, M>16) + skips the M=1 modulation / conditioning-embedder projections (see _INT8_EXCLUDE_NAME_TOKENS); + every other scheme uses scaled_mm (no M limit) and excludes nothing. Shared by the runtime + quantise path and the offline prequant-checkpoint builder so the two never drift -- an int8 + checkpoint built offline must skip exactly the layers the runtime path skips, or it bakes the + M=1 projections as int8 and crashes at the first denoise step on Flux / Qwen.""" + return _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else () + + # Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with # the lower-precision schemes listed as fallbacks for that arch tier. On Blackwell, fp8 # leads: measured on a B200, plain fp8 dynamic is both faster AND more accurate than the @@ -307,9 +336,11 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any: raise ValueError(f"unknown transformer quant scheme '{scheme}'") -def make_filter_fn(min_features: int): - """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear - with both in/out features >= ``min_features``. Hides the (module, fqn) callback arity.""" +def make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ()): + """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear with both + in/out features >= ``min_features`` AND whose fully-qualified name contains none of + ``exclude_name_tokens`` (used by int8 to skip the M=1 modulation / conditioning-embedder + projections that crash ``torch._int_mm``). Hides the (module, fqn) callback arity.""" def filter_fn(module: Any, fqn: str = "") -> bool: try: @@ -322,7 +353,13 @@ def make_filter_fn(min_features: int): out_features = getattr(module, "out_features", None) if in_features is None or out_features is None: return False - return in_features >= min_features and out_features >= min_features + if in_features < min_features or out_features < min_features: + return False + if exclude_name_tokens: + name = fqn.lower() if fqn else "" + if any(tok in name for tok in exclude_name_tokens): + return False + return True return filter_fn @@ -352,10 +389,13 @@ def quantize_transformer( try: from torchao.quantization import quantize_ + # int8 (torch._int_mm, M>16) additionally skips the M=1 modulation / conditioning-embedder + # projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything. + exclude = exclude_tokens_for_scheme(scheme) quantize_( transformer, _make_quant_config(scheme, fast_accum = fast_accum), - filter_fn = make_filter_fn(min_features), + filter_fn = make_filter_fn(min_features, exclude_name_tokens = exclude), ) # Runtime-only marker (torchao tensors are not safetensors-serializable; this # backend is inference-only, so this is purely diagnostic). diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index d527723b2e..a418367e98 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -325,6 +325,67 @@ def test_make_filter_fn(monkeypatch): assert keep(types.SimpleNamespace(), "no_attrs") is False +def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch): + # The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections + # (they crash torch._int_mm's M>16), while keeping the attention / FFN compute layers and + # the sequence embedders. fp8 (no exclusion) keeps everything. + from core.inference.diffusion_transformer_quant import _INT8_EXCLUDE_NAME_TOKENS + + class _Lin: + def __init__(self, i, o): + self.in_features, self.out_features = i, o + + torch = types.ModuleType("torch") + torch.nn = types.SimpleNamespace(Linear = _Lin) + monkeypatch.setitem(sys.modules, "torch", torch) + + keep = make_filter_fn(512, exclude_name_tokens = _INT8_EXCLUDE_NAME_TOKENS) + big = lambda: _Lin(3072, 18432) # noqa: E731 — large enough to pass min_features + # Excluded (M=1 modulation / conditioning embedders), despite large features: + for fqn in ( + "transformer_blocks.0.norm1.linear", + "transformer_blocks.0.norm1_context.linear", + "single_transformer_blocks.0.norm.linear", + "norm_out.linear", + "transformer_blocks.0.img_mod.1", + "transformer_blocks.0.txt_mod.1", + "double_stream_modulation_img.linear", + "time_text_embed.timestep_embedder.linear_2", + "time_text_embed.guidance_embedder.linear_2", + "time_guidance_embed.timestep_embedder.linear_2", + ): + assert keep(big(), fqn) is False, fqn + # Kept (M=seq compute layers + sequence embedders), NOT matched by the modulation tokens: + for fqn in ( + "transformer_blocks.0.attn.to_q", + "transformer_blocks.0.ff.net.0.proj", + "single_transformer_blocks.0.proj_mlp", + "single_transformer_blocks.0.attn.to_qkv_mlp_proj", + "context_embedder", # "context" contains "text" -> must NOT be excluded + "txt_in", + ): + assert keep(big(), fqn) is True, fqn + # Without the exclusion (fp8 path), the modulation layer is kept. + assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True + # A None / empty fqn must not crash the exclusion check (defensive against the callback + # passing no name); with no name nothing matches the exclusion tokens -> kept. + assert keep(big(), None) is True + assert keep(big(), "") is True + + +def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder(): + # The runtime quantiser and the offline prequant builder must apply the SAME int8 + # exclusion, or an int8 prequant artifact quantises the M=1 modulation/embedder linears + # and reintroduces the torch._int_mm crash. int8 gets the exclusion; others get none. + from core.inference.diffusion_transformer_quant import ( + _INT8_EXCLUDE_NAME_TOKENS, + exclude_tokens_for_scheme, + ) + assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8): + assert exclude_tokens_for_scheme(scheme) == () + + # ── apply ─────────────────────────────────────────────────────────────────────── From 465c238787a05bdf53593610c5719116ca0ae03e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:42:51 -0700 Subject: [PATCH 10/16] Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) (#6717) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../tests/test_diffusion_transformer_quant.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py index a418367e98..4953380412 100644 --- a/studio/backend/tests/test_diffusion_transformer_quant.py +++ b/studio/backend/tests/test_diffusion_transformer_quant.py @@ -386,6 +386,22 @@ def test_exclude_tokens_for_scheme_shared_by_runtime_and_builder(): assert exclude_tokens_for_scheme(scheme) == () +def test_exclude_tokens_for_scheme(): + # The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline + # prequant-checkpoint builder, so an int8 checkpoint built ahead of time skips exactly the + # layers the runtime path skips (offline == runtime). int8 excludes the M=1 modulation / + # embedder tokens; every scaled_mm scheme excludes nothing. + from core.inference.diffusion_transformer_quant import ( + _INT8_EXCLUDE_NAME_TOKENS, + exclude_tokens_for_scheme, + ) + + assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS + assert exclude_tokens_for_scheme(TQ_FP8) == () + assert exclude_tokens_for_scheme(TQ_NVFP4) == () + assert exclude_tokens_for_scheme(TQ_MXFP8) == () + + # ── apply ─────────────────────────────────────────────────────────────────────── From 692d3c197545ec8f7955c4c6a682cfb3cc92b68e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:43:56 -0700 Subject: [PATCH 11/16] Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine (#6724) * 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 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source) Add an opt-in transformer_quant mode that loads the dense bf16 transformer and torchao-quantises it onto the low-precision tensor cores, instead of the GGUF transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's 0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor. GGUF+compile stays the low-memory default and the fallback. The mode is gated on CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB); any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe (Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips the tiny projections that crash int8's torch._int_mm. New module mirrors diffusion_precision.py; quant runs before compile before placement. 184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py is the standalone torchao lever probe. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838 TFLOPS with FP16 accumulate; bf16 209), so: - fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on; pinning it guards consumer cards against a default change. On B200 it is identical speed and slightly better quality (LPIPS 0.050 vs 0.091). - the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate). 2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py): 2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile (our main ~2x). Documented as a dead end, not shipped. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image generation and reports max-abs + non-finite counts for use_fast_accum True vs False. Confirms fast accumulation is an accumulation-precision knob, not an overflow one: across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16), 0 non-finite elements and identical max-abs for both modes. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40; GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it. Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s), no overflow, quality below the quant noise floor. So the default leans to accuracy on data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force) lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off). 187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override threading). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9 scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer. Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with use_triton_kernel=False (the default triton path needs the missing MSLK library), but only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4 tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS, 2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166 vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): pre-quantized transformer loading The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs 13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline (scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True) the quantized weights, so the dense bf16 never touches the GPU. Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same torchao config + min_features filter the runtime path uses, applied ahead of time. New core/inference/diffusion_prequant.py (resolve_prequant_source + load_prequantized_transformer, best-effort, lazy imports). diffusion.py _load_dense_quant_pipeline tries the pre-quant source first and falls back to the dense materialise+quantise path, then to GGUF, so the default is unchanged. DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for the resolver, the meta-init+assign loader, and the backend branch selection + fallbacks; GPU verification via scripts/verify_prequant_backend.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): attention-backend selection Add a selectable attention kernel via the diffusers set_attention_backend dispatcher. Attention is memory-bandwidth bound, so a better kernel is an end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV matmuls torchao never touches) and composes with torch.compile. auto picks the best exact backend for the device: cuDNN fused attention (_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default (below the compile/quant noise floor); native SDPA elsewhere and when speed=off (so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/ xformers/aiter are honored, and an unavailable kernel falls back to the default rather than failing the load. New core/inference/diffusion_attention.py (normalize + per-device select + apply, best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline; attention_backend threads through begin_load / load_pipeline / status like the other load knobs. New request field attention_backend + status field. Hermetic CPU tests for normalize / select policy / apply fallback, plus route threading + 422. Measured via scripts/perf_levers_probe.py. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the arch tier to put int8 first on a consumer/workstation GPU (detected by the existing _is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first. Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation (-> int8) and data-center Ada/Hopper/Blackwell (-> fp8). * Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across denoise steps a DiT's output settles, so once the first block's residual barely changes the remaining blocks are skipped and their cached output reused. diffusers ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the standalone apply_first_block_cache hook as a fallback). Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 -> 2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar. OFF by default and a per-load opt-in: the win scales with step count, so it is for many-step models (Flux / Qwen-Image) and pointless for few-step distilled models (e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction of the trajectory. It composes with regional compile only with fullgraph=False (the cache's per-step decision is a torch.compiler.disable graph break), which the speed layer now switches to automatically when a cache is engaged. Best-effort: a model whose block signature the hook does not recognise is caught and the load proceeds uncached. - new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache (enable_cache / apply_first_block_cache fallback; threshold auto-raised for a quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import). - diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph when a cache is engaged. - diffusion.py: apply_step_cache before compile; thread transformer_cache / transformer_cache_threshold through begin_load -> load_pipeline and report the engaged mode in status(). - models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and transformer_cache_threshold request fields, engaged mode in the status response. - hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold selection, and best-effort failure handling, plus route threading + validation. - scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM / LPIPS vs the compiled no-cache baseline). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears) The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with 'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2 *_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing min_features filter did not exclude them. Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed, guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers (M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and quantises these layers fine, so the exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note 'context_embedder' contains the substring 'text_embed', which is why the token is the specific 'time_text_embed', not 'text_embed'. Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s (5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already working) are unchanged. - diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes exclude_name_tokens; quantize_transformer passes it for int8 only. - hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps attention / FFN / sequence-embedder linears), while fp8 keeps them. - scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's Linear layers and derive the exclusion list. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears) The prequant-checkpoint builder applied the dense quant filter without the int8-only M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8 checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now produces a working checkpoint on every supported model, giving int8 (the consumer-preferred scheme) the same ~2x load-VRAM and download reduction fp8 already had. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the native stable-diffusion.cpp engine instead of diffusers, with diffusers as the guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM. - diffusion_engine_router: centralised engine selection (built on the existing select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason. - sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface backed by sd-cli, with lazy binary install, registry-driven asset fetch, step-progress parsing, and cancellation. - diffusion_families: per-family single-file VAE + text-encoder asset mapping. - sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled). - routes/inference + gpu_arbiter: drive the active engine via the router; the API now reports the active engine and any fallback reason. - tests for the backend, router, route selection, and cancellation. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler Address review feedback on #6724: - engine router: unload the engine being deactivated on a switch, so the old model is not left resident-but-unreachable (the evictor only targets the active engine). - generate route: sd.cpp execution errors (nonzero exit / timeout / missing output) now map to 500, not 409 (which only means not-loaded / cancelled). - native batch: return per-image seeds and persist the actual seed for each image so every batch image is reproducible. - Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the stable-diffusion.cpp docs; other families keep sd-cli defaults. - honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload offload flags on the native load instead of hardcoding them off. - fail the load when the sd-cli binary is present but not runnable (version() now returns None on exec error / nonzero exit). - size estimate: only treat the transformer asset as a possible local path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in load_prequantized_transformer ends in torch.load(weights_only=False), which executes arbitrary code from the pickle. The transformer_prequant_path load-request field reached that unpickle for any local file an authenticated caller named, so a request could trigger remote code execution. Refuse the source.kind=='path' branch unless the operator sets UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted and unaffected. Document the requirement on the API field and add gate tests. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa - apply_attention_backend now restores the native default when no backend is requested or a kernel fails. diffusers keeps a process-wide active attention backend that set_attention_backend updates, and a fresh transformer's processors follow it, so a load that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile load pinned, breaking the bit-identical/off guarantee. - select_attention_backend drops flash3/flash4 up front when the CUDA capability is below Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit request on the wrong card set fine then crashed mid-generation; it now falls back to native. - Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a valid alias of native) is accepted instead of 422-rejected by Pydantic. - Drop the dead replace('-','_') normalization (no alias uses dashes/underscores). - perf_levers_probe.py output dir is now relative to the script, not a hardcoded path. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF - apply_step_cache now engages only via the transformer's native enable_cache (the diffusers CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model now runs uncached per the best-effort contract. - GGUF transformers are quantized (the default Studio load path), so they now use the higher quantized FBCache threshold when the caller leaves it unset, instead of the dense default that could keep the cache from triggering. - fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so fullgraph=True failed warmup and silently measured an eager cached run); output dir is now relative to the script, not a hardcoded path. * Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder _is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate. Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0). * Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just TypeError on older torchao), which would break the mxfp8 config helper instead of falling back to the default. Catch both so the fallback is robust. quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds CUDA memory during the per-row VRAM probe; output dir relative to the script. * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [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 14): guard the int8 exclusion filter against a None fqn The filter callback can be invoked without a module name, so fqn.lower() would raise AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens, so the linear is kept) instead of crashing the quantise pass. * Studio diffusion (Phase 16) review fixes: native engine robustness - sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64); a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds aliased to the same image. Random seeds stay 53-bit (JS-safe). - sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub fall back to anonymous instead of failing auth on a blank token. - sd_cpp_backend: a superseding load now cancels the in-flight generation, so the old sd-cli can no longer return/persist an image from the previous model. - diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a slow 10+ GB free / CUDA sync does not block engine selection. - diffusion_engine_router: probe sd-cli runnability (version()) before committing to native, so a present-but-unrunnable binary falls back to diffusers at selection. - diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a CPU-only install can still reach the native sd.cpp engine instead of failing load. - tests updated for the runnability probe + a not-runnable fallback case. * Studio diffusion (Phase 9) review fixes: prequant safety + validation - SECURITY: a request-supplied local pre-quant path is now unpickled only when it resolves inside an operator-configured ALLOWLIST of directories (UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in, once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on any path a load request named (arbitrary code execution). realpath() blocks symlink escapes; a bare on/off toggle is no longer a wildcard. - Validate the checkpoint's min_features against the runtime Linear filter, so a checkpoint that quantised a different layer set is rejected instead of silently loading a model that mismatches the dense path while reporting the same scheme. - Tolerant base_model_id compare (exact or same final path/repo segment), so a local path or fork of the canonical base is accepted instead of falling back to dense. - _has_meta_tensors uses any(chain(...)) (no intermediate lists). - prequant verify/probe scripts use repo-relative paths (+ env overrides), not the author's absolute /mnt paths. - tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path - diffusion: a torchao-quantized transformer is committed only compiled. A dense model resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF it replaced), so when transformer_quant engaged and speed resolved to off, promote to default (regional compile); warn loudly if compile still does not engage. - diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK); otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF. - nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path). - test asserts the eager-quant -> default-compile promotion. * Studio diffusion (Phase 10) review fixes: attention gating + probe isolation - diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation, so auto now stays on native SDPA there. - diffusion_attention: _active_attention_backend handles get_active_backend() returning an enum/None (not a tuple); the old unpack always raised and was swallowed, so the native-restore short-circuit never fired. - perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul so the inductor_flags variant doesn't leak into later compiled rows. - tests for the SM80 cuDNN gate. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc Codex review: the transformer_prequant_path field description still told operators to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the prior security fix made that variable a directory allowlist -- _allowed_prequant_roots deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator following the documented =1 would have every transformer_prequant_path request silently refused. The description now states it must name one or more allowlisted directories and that a bare on/off value is not accepted. Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does not say =1, and describes an allowlist/directory (guards against doc drift). * Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset Codex review on attention-backend selection: - Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first generation with no fallback. select_attention_backend now applies _cudnn_attention_supported() to an explicit cuDNN request too. - flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a (min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+. - apply_attention_backend's success path left diffusers' process-wide active backend pinned to the kernel it set; a later component whose processors are unconfigured (backend None) would inherit it. It now resets the global registry to native after a successful per-transformer set (the transformer keeps its own backend), best-effort. Also fixed _active_attention_backend: get_active_backend() returns a (name, fn) tuple, so the prior code stringified the tuple and never matched a name, defeating the native-restore short-circuit. Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on SM90; global registry reset after a successful set; _active_attention_backend reads the tuple return. * Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300 (Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So _is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they are treated as data-center class and keep the intended fp8-first behavior. Test: extends the datacenter parametrize with 'NVIDIA B300' and 'NVIDIA GH200 480GB' (now _is_consumer_gpu False). * Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder Codex review: the M=1 modulation/embedder exclusion was wired only into the dense runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint quantised the AdaLN modulation and conditioning-embedder linears, and loading it via transformer_prequant_path (the load path only loads already-quantised tensors, it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes for the runtime path. Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 -> the M=1 exclusion, every other scheme -> none) and use it in both the runtime quantiser and the builder, so a prequant artifact's quantised-layer set always matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion). Test: int8_exclude_name_tokens returns the exclusion for int8 and () for fp8/nvfp4/mxfp8. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race Codex review on the native-engine routing: - The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat) unconditionally after engine selection. A native sd.cpp load on a pure-CPU host never touches the GPU, so that needlessly tore down the resident chat model. The handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it. - sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when _run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so the setting was unverifiable. status now derives them from state.offload_flags (still 'none' on CPU, where the flags are empty). - _run_load committed the new state without cancelling/waiting on a generation that started during the (slow) asset download, so a stale sd-cli run against the OLD model could finish afterward and persist an image from the previous model once the new load reported ready. The commit now signals the in-flight cancel and waits on _generate_lock before swapping _state (taken only at commit, so the download never serialises against generation), mirroring the diffusers load path. Tests: CPU native load skips the arbiter while a GPU native load takes it; status reports offload active when flags are set; _run_load cancels and waits for an in-flight generation before committing. * Studio diffusion (Phase 14) review round 2: align helper name with the stack Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the identical helper already present higher in the diffusion stack (Phase 16). The helper definition, the runtime quantiser call, and the offline builder are now byte-identical to that version, so the two branches no longer introduce a divergent name for the same single-source-of-truth and the stack merges without a conflict on this fix. No behavior change. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../core/inference/diffusion_device.py | 18 +- .../core/inference/diffusion_engine_router.py | 166 +++++ .../core/inference/diffusion_families.py | 57 ++ studio/backend/core/inference/gpu_arbiter.py | 6 +- .../backend/core/inference/sd_cpp_backend.py | 661 ++++++++++++++++++ .../backend/core/inference/sd_cpp_engine.py | 80 ++- studio/backend/models/inference.py | 5 + studio/backend/routes/inference.py | 86 ++- .../tests/test_diffusion_engine_router.py | 157 +++++ studio/backend/tests/test_diffusion_routes.py | 121 ++++ studio/backend/tests/test_sd_cpp_backend.py | 310 ++++++++ studio/backend/tests/test_sd_cpp_engine.py | 14 +- 12 files changed, 1632 insertions(+), 49 deletions(-) create mode 100644 studio/backend/core/inference/diffusion_engine_router.py create mode 100644 studio/backend/core/inference/sd_cpp_backend.py create mode 100644 studio/backend/tests/test_diffusion_engine_router.py create mode 100644 studio/backend/tests/test_sd_cpp_backend.py diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index 0c874e4b66..2d2cc11343 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -59,8 +59,24 @@ def resolve_diffusion_device_target() -> DiffusionDeviceTarget: (CUDA -> XPU -> MPS -> CPU). On Apple Silicon Studio reports MLX/CPU when its product backend is gated on the ``mlx`` package, but diffusers runs on PyTorch's MPS backend, so those cases still fall through to the MPS probe. + + Torch is optional here: on a CPU-only install without PyTorch the native + stable-diffusion.cpp engine still runs (it shells out to sd-cli), so a missing + torch reports a torch-free CPU target instead of crashing the whole + ``/images/load`` before the engine router can select the native backend. """ - import torch + try: + import torch + except Exception: + return DiffusionDeviceTarget( + device = "cpu", + dtype = None, + backend = "cpu", + vendor = None, + supports_model_cpu_offload = False, + supports_default_torch_compile = False, + supports_pinned_transfer = False, + ) try: from utils.hardware import DeviceType, get_device diff --git a/studio/backend/core/inference/diffusion_engine_router.py b/studio/backend/core/inference/diffusion_engine_router.py new file mode 100644 index 0000000000..c8db7f81d2 --- /dev/null +++ b/studio/backend/core/inference/diffusion_engine_router.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 + +"""Selects the diffusion engine (diffusers vs native sd.cpp) for the live route. + +The image routes drive one engine at a time. On a CUDA/ROCm/XPU GPU that engine is +the diffusers ``DiffusionBackend`` (the default, and the only path with the torchao +fast-quant / compile stack). With no usable GPU (CPU, or Apple MPS when explicitly +enabled) it is the native ``SdCppDiffusionBackend``, which is faster and far lighter +on RAM there. The choice is made once at load time and remembered, so ``generate`` / +``unload`` / ``status`` / progress all act on the same engine the load committed to. + +Selection is centralised here and built on the existing pure ``select_diffusion_engine`` +decision; this module adds the policy around it (env opt-out, MPS gating, per-family +native-asset support, lazy binary availability) and records why a fallback happened. + +Env knobs (one canonical interpretation each): + UNSLOTH_DIFFUSION_ENGINE=auto|diffusers|sd_cpp force an engine (auto = decide) + UNSLOTH_DIFFUSION_SD_CPP=auto|0|1 enable/disable the native route + UNSLOTH_DIFFUSION_SD_CPP_MPS=0|1 allow native on Apple MPS (default off) + UNSLOTH_DIFFUSION_SD_CPP_INSTALL=auto|0|1 allow lazy binary install (in sd_cpp_backend) +""" + +from __future__ import annotations + +import os +import threading +from typing import Any, Optional + +from core.inference.diffusion_device import resolve_diffusion_device_target +from core.inference.diffusion_families import DiffusionFamily, family_sd_cpp_supported +from core.inference.sd_cpp_backend import _install_allowed, ensure_sd_cpp_binary +from core.inference.sd_cpp_engine import ( + ENGINE_DIFFUSERS, + ENGINE_SD_CPP, + SdCppEngine, + select_diffusion_engine, +) +from loggers import get_logger + +logger = get_logger(__name__) + +_DISABLE_TOKENS = frozenset({"0", "off", "false", "no"}) +_ENABLE_TOKENS = frozenset({"1", "on", "true", "yes"}) + +# The engine the current (or most recent) load committed to, and why a non-native +# choice was made. Mutated only under _lock during selection. +_lock = threading.Lock() +_active_engine_name: str = ENGINE_DIFFUSERS +_fallback_reason: Optional[str] = None + + +def _engine_config() -> tuple[str, str, bool]: + forced = os.environ.get("UNSLOTH_DIFFUSION_ENGINE", "auto").strip().lower() + sd_cpp = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP", "auto").strip().lower() + mps = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP_MPS", "0").strip().lower() in _ENABLE_TOKENS + return forced, sd_cpp, mps + + +def get_active_diffusion_engine() -> Any: + """The engine object the active selection points at (defaults to diffusers).""" + if _active_engine_name == ENGINE_SD_CPP: + from core.inference.sd_cpp_backend import get_sd_cpp_backend + return get_sd_cpp_backend() + from core.inference.diffusion import get_diffusion_backend + + return get_diffusion_backend() + + +def active_engine_name() -> str: + return _active_engine_name + + +def _activate(name: str, reason: Optional[str]) -> Any: + global _active_engine_name, _fallback_reason + # Switching engines: unload the one being deactivated first, or its model + # stays resident but unreachable (the arbiter evictor only targets the active + # engine), leaking 10+ GB and defeating the chat<->diffusion handoff. The unload + # itself (freeing 10+ GB / syncing CUDA) is slow, so resolve the engine under the + # lock but run unload() OUTSIDE it -- holding _lock across a slow unload would + # block every other selection caller. + engine_to_unload = None + old_name = None + with _lock: + if name != _active_engine_name: + engine_to_unload = get_active_diffusion_engine() + old_name = _active_engine_name + _active_engine_name = name + _fallback_reason = reason if name == ENGINE_DIFFUSERS else None + if engine_to_unload is not None: + try: + engine_to_unload.unload() + except Exception as exc: # noqa: BLE001 -- best-effort; never block the switch + logger.warning("failed to unload previous engine %s: %s", old_name, exc) + if name == ENGINE_SD_CPP: + logger.info("diffusion engine: sd_cpp") + else: + logger.info("diffusion engine: diffusers (%s)", reason or "selected") + return get_active_diffusion_engine() + + +def select_and_activate_engine(fam: DiffusionFamily, *, hf_token: Optional[str] = None) -> Any: + """Pick + activate the engine for loading ``fam`` on this host; return the engine. + + Falls back to diffusers (recording a reason) whenever the native route is + disabled, the device has a usable GPU, MPS is not enabled, the family has no + native asset mapping, or the sd-cli binary is unavailable -- always BEFORE the + slow load begins, so a fallback never strands a half-native load. + """ + forced, sd_cpp_pref, mps_enabled = _engine_config() + + if forced == ENGINE_DIFFUSERS: + return _activate(ENGINE_DIFFUSERS, "forced (UNSLOTH_DIFFUSION_ENGINE=diffusers)") + + prefer_native = forced == ENGINE_SD_CPP + if sd_cpp_pref in _DISABLE_TOKENS and not prefer_native: + return _activate(ENGINE_DIFFUSERS, "native engine disabled (UNSLOTH_DIFFUSION_SD_CPP=0)") + + target = resolve_diffusion_device_target() + backend = target.backend + # Policy: CPU is always native-eligible; MPS only when explicitly enabled; a GPU + # backend (cuda/rocm/xpu) never is, unless the user force-selects sd_cpp. + policy_eligible = backend == "cpu" or (backend == "mps" and mps_enabled) or prefer_native + fam_ok = family_sd_cpp_supported(fam) + + binary = None + if policy_eligible and fam_ok: + binary = ensure_sd_cpp_binary(allow_install = _install_allowed()) + # Probe runnability here, before committing the route to native: a present but + # non-runnable binary (wrong arch, missing shared libs, no execute bit) would + # otherwise pass as available and only fail inside the background load, instead + # of falling back to diffusers now. + if binary and SdCppEngine(binary = binary).version() is None: + logger.warning("sd-cli at %s is present but not runnable; using diffusers", binary) + binary = None + + native_available = bool(binary) and policy_eligible and fam_ok + choice = select_diffusion_engine( + backend, native_available = native_available, prefer_native = prefer_native + ) + if choice == ENGINE_SD_CPP: + return _activate(ENGINE_SD_CPP, None) + + # Explain the diffusers choice for status/telemetry. + if not policy_eligible: + reason = f"GPU backend '{backend}' uses diffusers" + elif not fam_ok: + reason = f"family '{fam.name}' has no native sd.cpp asset mapping" + elif not binary: + reason = "sd-cli binary unavailable" + else: + reason = "diffusers selected" + return _activate(ENGINE_DIFFUSERS, reason) + + +def annotate_status(status: dict[str, Any]) -> dict[str, Any]: + """Tag a backend status dict with the active engine + any fallback reason.""" + out = dict(status) + out["engine"] = _active_engine_name + out["fallback_reason"] = _fallback_reason + return out + + +def active_status() -> dict[str, Any]: + """The active engine's status, annotated with which engine + any fallback reason.""" + return annotate_status(get_active_diffusion_engine().status()) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index d5068cbf49..3afbe9b495 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -45,6 +45,26 @@ class DiffusionFamily: # materialising the dense bf16 transformer on the GPU (much lower load VRAM + a # smaller download). Empty until checkpoints are hosted -> behaviour is unchanged. prequant_repos: tuple[tuple[str, str], ...] = field(default_factory = tuple) + # Native (stable-diffusion.cpp) single-file assets, used only when the no-GPU + # sd.cpp engine is selected (CPU / Apple). The transformer GGUF is shared with + # the diffusers path; sd-cli additionally needs a single-file VAE and text + # encoder(s), because the diffusers base repo ships those sharded and sd-cli + # cannot read that layout. Each asset is a hashable (repo_id, filename) the + # backend fetches with hf_hub_download. ``sd_cpp_text_encoders`` carries a + # trailing SdCppModelFiles field name (clip_l / t5xxl / llm / qwen2vl / clip_g) + # so the backend maps each file onto the right sd-cli flag. Empty -> the family + # has no native mapping and the sd.cpp route falls back to diffusers. + sd_cpp_vae: Optional[tuple[str, str]] = None + # VAE latent-format override for sd-cli (--vae-format): "flux2" for the FLUX.2 + # autoencoder, None (auto) otherwise. + sd_cpp_vae_format: Optional[str] = None + sd_cpp_text_encoders: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) + # Family-specific sd-cli sampler settings, applied on the native path so the + # output matches the model's supported invocation (e.g. Qwen-Image needs + # --sampling-method euler --flow-shift 3 per stable-diffusion.cpp's docs). None + # leaves sd-cli's defaults (correct for the distilled flux/z-image families). + sd_cpp_sampling_method: Optional[str] = None + sd_cpp_flow_shift: Optional[float] = None # Keyed by architecture, not per model variant: a checkpoint's specific base repo @@ -60,6 +80,11 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "FluxTransformer2DModel", base_repo = "black-forest-labs/FLUX.1-schnell", aliases = ("flux1", "flux-1"), + sd_cpp_vae = ("black-forest-labs/FLUX.1-schnell", "ae.safetensors"), + sd_cpp_text_encoders = ( + ("comfyanonymous/flux_text_encoders", "clip_l.safetensors", "clip_l"), + ("comfyanonymous/flux_text_encoders", "t5xxl_fp16.safetensors", "t5xxl"), + ), ), # FLUX.2-klein is a distinct pipeline (Flux2KleinPipeline) with a Qwen3 text # encoder, not the Mistral-based Flux2Pipeline; it must precede a generic @@ -70,6 +95,14 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( transformer_class = "Flux2Transformer2DModel", base_repo = "black-forest-labs/FLUX.2-klein-4B", aliases = ("flux2-klein",), + # FLUX.2 uses a distinct 32-channel autoencoder; sd-cli needs the latent + # format override. The single-file VAE ships in Comfy-Org/flux2-dev (the + # klein-4B repo only has a sharded diffusers VAE). Shares Qwen3-4B with z-image. + sd_cpp_vae = ("Comfy-Org/flux2-dev", "split_files/vae/flux2-vae.safetensors"), + sd_cpp_vae_format = "flux2", + sd_cpp_text_encoders = ( + ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"), + ), ), DiffusionFamily( name = "qwen-image", @@ -78,6 +111,19 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( base_repo = "Qwen/Qwen-Image", cfg_kwarg = "true_cfg_scale", aliases = ("qwen_image", "qwenimage"), + sd_cpp_vae = ("Comfy-Org/Qwen-Image_ComfyUI", "split_files/vae/qwen_image_vae.safetensors"), + # The Qwen2.5-VL text encoder as a Q4_K_M GGUF keeps the CPU RAM win (the + # bf16 safetensors encoder is ~15 GB). sd-cli's --qwen2vl is an alias of --llm. + sd_cpp_text_encoders = ( + ( + "unsloth/Qwen2.5-VL-7B-Instruct-GGUF", + "Qwen2.5-VL-7B-Instruct-Q4_K_M.gguf", + "qwen2vl", + ), + ), + # Qwen-Image's supported sd.cpp invocation (docs/qwen_image.md). + sd_cpp_sampling_method = "euler", + sd_cpp_flow_shift = 3.0, ), DiffusionFamily( name = "z-image", @@ -87,6 +133,10 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("zimage", "z_image"), # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, + sd_cpp_vae = ("Comfy-Org/z_image_turbo", "split_files/vae/ae.safetensors"), + sd_cpp_text_encoders = ( + ("Comfy-Org/z_image_turbo", "split_files/text_encoders/qwen_3_4b.safetensors", "llm"), + ), ), ) @@ -137,6 +187,13 @@ def family_prequant_repo(fam: DiffusionFamily, scheme: str) -> Optional[str]: return None +def family_sd_cpp_supported(fam: DiffusionFamily) -> bool: + """True when the family has the single-file VAE + text-encoder mapping the + native sd.cpp engine needs. A family without it can only run on diffusers, so + the no-GPU route falls back rather than routing to sd-cli.""" + return bool(fam.sd_cpp_vae and fam.sd_cpp_text_encoders) + + def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: """Resolve ``gguf_filename`` to a file under ``repo_root``, rejecting escapes. diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index daa0ef4e94..a6c79bf1e8 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -55,8 +55,10 @@ def _evict_chat() -> None: def _evict_diffusion() -> None: - from core.inference.diffusion import get_diffusion_backend - get_diffusion_backend().unload() + # Unload whichever engine the router has active (diffusers or native sd.cpp), so a + # chat acquire frees the right one. + from core.inference.diffusion_engine_router import get_active_diffusion_engine + get_active_diffusion_engine().unload() # Patchable in tests via monkeypatch.setitem. diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py new file mode 100644 index 0000000000..4d0bbcbfa7 --- /dev/null +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -0,0 +1,661 @@ +# 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 diffusion backend (the no-GPU tier). + +``SdCppDiffusionBackend`` presents the SAME public surface the image routes use on +the diffusers ``DiffusionBackend`` (``begin_load`` / ``load_progress`` / ``generate`` +/ ``generate_progress`` / ``unload`` / ``status``), but is backed by the ``sd-cli`` +subprocess (``SdCppEngine``) instead of an in-process diffusers pipeline. The engine +router (``diffusion_engine_router.py``) selects this backend only when no usable +CUDA/ROCm/XPU GPU is present, where it is measurably faster and far lighter on RAM +than diffusers (see outputs/sdcpp_cpu). + +It reuses the transformer GGUF the diffusers path already downloads and additionally +fetches the per-family single-file VAE + text encoders declared in +``diffusion_families`` (sd-cli cannot read the sharded diffusers components). The +binary is installed lazily on first use; if it is unavailable or the family has no +native mapping, the router falls back to diffusers, so this backend is only ever +asked to run requests it can serve. + +Import-light on purpose: no torch / diffusers here, so selecting it on a CPU box +does not drag the heavy GPU stack into the process. +""" + +from __future__ import annotations + +import logging +import os +import re +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional + +from core.inference.diffusion_device import resolve_diffusion_device_target +from core.inference.diffusion_families import ( + DiffusionFamily, + detect_family, + family_sd_cpp_supported, + resolve_base_repo, + resolve_local_gguf_child, +) +from core.inference.diffusion_memory import ( + OFFLOAD_GROUP, + OFFLOAD_MODEL, + OFFLOAD_NONE, + OFFLOAD_SEQUENTIAL, +) +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles, offload_flags +from core.inference.sd_cpp_engine import ( + SdCppCancelled, + SdCppEngine, + find_sd_cpp_binary, +) +from loggers import get_logger + +logger = get_logger(__name__) + +# A sampling-progress line like " 4/4" / "[ 12/ 28]" / "sampling: 50%|...| 14/28". +# We only trust a match whose denominator equals the requested step count, so an +# unrelated "1/100" elsewhere in the log can't move the bar. +_STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)") + +# Serialises the one-time binary install so concurrent first-loads don't race on the +# download / extract / chmod. +_install_lock = threading.Lock() + + +def ensure_sd_cpp_binary(*, allow_install: bool = True, accelerator: str = "cpu") -> Optional[str]: + """Path to a usable ``sd-cli`` binary, installing the prebuilt once if needed. + + Returns the binary path, or None when it is absent and cannot be installed + (install disabled, no network, unsupported platform). Never raises -- a None + return is the router's signal to fall back to diffusers. + """ + found = find_sd_cpp_binary() + if found: + return found + if not allow_install: + return None + with _install_lock: + # Re-check inside the lock: a concurrent first-load may have installed it. + found = find_sd_cpp_binary() + if found: + return found + try: + import sys + + studio_dir = Path(__file__).resolve().parents[3] # .../studio + if str(studio_dir) not in sys.path: + sys.path.insert(0, str(studio_dir)) + from install_sd_cpp_prebuilt import install as _install + except Exception as exc: # noqa: BLE001 -- import path / module issues are non-fatal + logger.warning("sd-cli installer import failed: %s", exc) + return None + try: + path = _install(accelerator = accelerator) + logger.info("sd-cli installed at %s", path) + return str(path) + except Exception as exc: # noqa: BLE001 -- download/extract failure -> fall back + logger.warning("sd-cli auto-install failed: %s", exc) + return None + + +@dataclass(frozen = True) +class _SdState: + """The loaded native checkpoint: resolved asset paths + run settings.""" + + repo_id: str + base_repo: str + family: DiffusionFamily + device: str + files: SdCppModelFiles + vae_format: Optional[str] = None + native_speed: str = "off" + offload_flags: tuple[str, ...] = () + threads: Optional[int] = None + sampling_method: Optional[str] = None + flow_shift: Optional[float] = None + + +def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: + """Map the diffusers memory knobs onto an sd-cli offload policy. Only meaningful + off-CPU (forced sd_cpp / MPS); on CPU everything is resident in RAM anyway.""" + mode = (memory_mode or "").strip().lower() + if mode == "low_vram": + return OFFLOAD_SEQUENTIAL + if mode == "balanced": + return OFFLOAD_GROUP + if cpu_offload and mode in ("", "auto"): + return OFFLOAD_MODEL + return OFFLOAD_NONE + + +def _native_speed_for(speed_mode: Optional[str]) -> str: + mode = (speed_mode or "off").strip().lower() + return mode if mode in ("default", "max") else "off" + + +@dataclass +class _SdLoading: + """An in-flight asset download, polled for progress.""" + + repo_id: str + base_repo: str + expected_bytes: int = 0 + downloaded_bytes: int = 0 + error: Optional[str] = None + + +@dataclass +class _SdGen: + """An in-flight generation, updated from parsed sd-cli progress lines.""" + + total_steps: int + step: int = 0 + first_step_at: float = 0.0 + eta_seconds: Optional[float] = None + + +def _estimate_eta(total_steps: int, step: int, first_step_at: float, now: float) -> Optional[float]: + steps_since_first = step - 1 + if not first_step_at or steps_since_first <= 0: + return None + per_step = (now - first_step_at) / steps_since_first + return max(0.0, (total_steps - step) * per_step) + + +def _map_guidance( + fam: DiffusionFamily, guidance: Optional[float] +) -> tuple[Optional[float], Optional[float]]: + """(cfg_scale, guidance) for sd-cli from the single diffusers ``guidance`` value. + + FLUX families take a distilled embedded ``--guidance``; everyone else uses real + classifier-free ``--cfg-scale``. A distilled 0/1 means CFG off (sd-cli's 1.0); a + value > 1 is real CFG. Mirrors the engine mapping validated in the CPU benchmark. + """ + if fam.name in ("flux.1", "flux.2-klein"): + return None, (float(guidance) if guidance is not None else None) + cfg = float(guidance) if (guidance is not None and guidance > 1.0) else 1.0 + return cfg, None + + +class SdCppDiffusionBackend: + """Native sd.cpp backend with the diffusers ``DiffusionBackend`` method surface.""" + + def __init__(self, engine: Optional[SdCppEngine] = None) -> None: + self._lock = threading.Lock() + self._generate_lock = threading.Lock() + self._engine = engine # resolved lazily on first load so import stays cheap + self._state: Optional[_SdState] = None + self._loading: Optional[_SdLoading] = None + self._load_token = 0 + self._cancel_event = threading.Event() + self._active_generate_cancel: Optional[threading.Event] = None + self._gen: Optional[_SdGen] = None + + @property + def is_loaded(self) -> bool: + return self._state is not None + + def _resolve_engine(self) -> SdCppEngine: + """The SdCppEngine, installing the binary on first use. Raises if unusable.""" + if self._engine is not None and self._engine.is_available(): + return self._engine + binary = ensure_sd_cpp_binary(allow_install = _install_allowed()) + if not binary: + raise RuntimeError("sd-cli (stable-diffusion.cpp) binary is unavailable.") + self._engine = SdCppEngine(binary = binary) + return self._engine + + # ── Background load + progress ───────────────────────────────────────── + + def begin_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + 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, + # diffusers-only knobs accepted (so the route calls both engines uniformly) + # and ignored -- sd.cpp has no torchao quant / SDPA dispatcher / fbcache. + text_encoder_quant: Optional[str] = None, + transformer_quant: Optional[str] = None, + transformer_quant_fast_accum: Optional[bool] = None, + transformer_prequant_path: Optional[str] = None, + attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, + ) -> dict[str, Any]: + """Validate, then fetch assets on a daemon thread. Returns at once.""" + # An empty / whitespace token is "no token": passing "" verbatim to HfApi / + # hf_hub_download is treated as an explicit (invalid) credential and breaks the + # anonymous fallback for public repos. + hf_token = hf_token.strip() if hf_token and hf_token.strip() else None + if not gguf_filename: + raise ValueError( + "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." + ) + fam = detect_family(repo_id, family_override) + if fam is None: + raise ValueError(f"Could not infer a diffusion family for '{repo_id}'.") + if not family_sd_cpp_supported(fam): + raise ValueError(f"Family '{fam.name}' has no native sd.cpp asset mapping.") + + base = resolve_base_repo(fam, base_repo) + with self._lock: + if self._loading is not None and self._loading.error is None: + raise RuntimeError("A diffusion load is already in progress.") + # A superseding load must stop any in-flight generation, or the old sd-cli + # keeps running against the previous model and can still return / persist an + # image after the new load has started (matches unload()'s cancel). + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + self._load_token += 1 + token = self._load_token + self._cancel_event.clear() + self._loading = _SdLoading(repo_id = repo_id, base_repo = base) + + threading.Thread( + target = self._run_load, + kwargs = dict( + repo_id = repo_id, + gguf_filename = gguf_filename, + base = base, + fam = fam, + hf_token = hf_token, + cpu_offload = cpu_offload, + memory_mode = memory_mode, + speed_mode = speed_mode, + _load_token = token, + ), + daemon = True, + ).start() + return self.status() + + def _run_load( + self, + *, + repo_id: str, + gguf_filename: str, + base: str, + fam: DiffusionFamily, + hf_token: Optional[str], + cpu_offload: bool = False, + memory_mode: Optional[str] = None, + speed_mode: Optional[str] = None, + _load_token: int, + ) -> None: + try: + # Ensure the binary up front so an install failure surfaces before the + # multi-GB asset pull (the router also pre-checks, but a forced reload here + # must not silently download then fail at generate). + engine = self._resolve_engine() + + assets = self._asset_specs(repo_id, gguf_filename, fam) + self._set_expected_bytes(assets, hf_token) + paths = self._fetch_assets(assets, hf_token) + + files = SdCppModelFiles( + diffusion_model = paths["diffusion_model"], + vae = paths.get("vae"), + clip_l = paths.get("clip_l"), + clip_g = paths.get("clip_g"), + t5xxl = paths.get("t5xxl"), + llm = paths.get("llm"), + qwen2vl = paths.get("qwen2vl"), + ) + device = resolve_diffusion_device_target().device + # Honor the requested speed everywhere; offload only off-CPU (forced + # sd_cpp / MPS), since on CPU sd-cli is resident in RAM and the offload + # flags are no-ops. + offload: tuple[str, ...] = () + if device != "cpu": + offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload))) + state = _SdState( + repo_id = repo_id, + base_repo = base, + family = fam, + device = device, + files = files, + vae_format = fam.sd_cpp_vae_format, + native_speed = _native_speed_for(speed_mode), + offload_flags = offload, + threads = None, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, + ) + # Probe the binary: version() returns None when the present binary cannot + # run (bad permissions / missing shared libs), so fail the load now rather + # than commit a "ready" state that crashes on the first generation. + if engine.version() is None: + raise RuntimeError("sd-cli binary is present but not runnable.") + # A generation that started during the (slow) asset download is still running + # against the OLD model. Abort it, then WAIT on _generate_lock for it to exit + # before publishing the new state -- otherwise that stale sd-cli run can finish + # afterward and persist an image from the previous model once this load reports + # ready (mirrors the diffusers load_pipeline commit). _generate_lock is taken + # only here, not during the download, so the long fetch never serialises against + # generation; the inner token re-check guards an unload/newer load arriving while + # we waited. + with self._lock: + if self._load_token != _load_token: + return # superseded / cancelled + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + with self._generate_lock: + with self._lock: + if self._load_token != _load_token: + return # superseded / cancelled while waiting + self._state = state + self._loading = None + except SdCppCancelled: + return + except Exception as exc: # noqa: BLE001 -- surfaced via load_progress + if self._load_token != _load_token: + return + logger.error("sd_cpp.load_failed: %s", exc) + with self._lock: + if self._load_token == _load_token and self._loading is not None: + self._loading.error = str(exc) + + def _asset_specs( + self, repo_id: str, gguf_filename: str, fam: DiffusionFamily + ) -> list[tuple[str, str, str]]: + """(repo, filename, kind) for every file sd-cli needs. ``kind`` is the + SdCppModelFiles field; the transformer reuses the diffusers GGUF.""" + specs: list[tuple[str, str, str]] = [(repo_id, gguf_filename, "diffusion_model")] + if fam.sd_cpp_vae: + specs.append((fam.sd_cpp_vae[0], fam.sd_cpp_vae[1], "vae")) + for terepo, tefile, kind in fam.sd_cpp_text_encoders: + specs.append((terepo, tefile, kind)) + return specs + + def _set_expected_bytes( + self, assets: list[tuple[str, str, str]], hf_token: Optional[str] + ) -> None: + """Best-effort total download size for the progress bar (0 if unknown).""" + total = 0 + try: + from huggingface_hub import HfApi + api = HfApi(token = hf_token) + for repo, fn, kind in assets: + # Only the transformer can be a local path; for the others ``repo`` is + # an HF id (a same-named local dir must not skip the size estimate). + if kind == "diffusion_model" and Path(repo).expanduser().exists(): + continue + try: + info = api.get_paths_info(repo, paths = [fn], expand = False) + for it in info: + total += int(getattr(it, "size", 0) or 0) + except Exception: # noqa: BLE001 -- one missing size is non-fatal + continue + except Exception: # noqa: BLE001 -- estimate is best-effort + total = 0 + loading = self._loading + if loading is not None: + loading.expected_bytes = total + + def _fetch_assets( + self, assets: list[tuple[str, str, str]], hf_token: Optional[str] + ) -> dict[str, str]: + """Download every asset (cancellable), returning kind -> local path.""" + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + paths: dict[str, str] = {} + for repo, fn, kind in assets: + if self._cancel_event.is_set(): + raise SdCppCancelled("load cancelled") + local_root = Path(repo).expanduser() + if kind == "diffusion_model" and local_root.exists(): + path = str(resolve_local_gguf_child(local_root, fn)) + else: + path = hf_hub_download_with_xet_fallback( + repo, fn, hf_token, cancel_event = self._cancel_event + ) + paths[kind] = path + with self._lock: + if self._loading is not None: + try: + self._loading.downloaded_bytes += os.path.getsize(path) + except OSError: + pass + return paths + + def load_progress(self) -> dict[str, Any]: + loading = self._loading + if loading is not None and loading.error: + return _progress("error", error = loading.error) + if loading is None: + return _progress("ready" if self._state is not None else None) + downloaded = loading.downloaded_bytes + expected = loading.expected_bytes + if expected > 0 and downloaded >= expected * 0.999: + return _progress("finalizing", min(downloaded, expected), expected, 1.0) + fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 + return _progress("downloading", downloaded, expected, fraction) + + # ── Generate ─────────────────────────────────────────────────────────── + + def generate( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: int = 1024, + height: int = 1024, + steps: int = 9, + guidance: float = 0.0, + seed: Optional[int] = None, + batch_size: int = 1, + ) -> dict[str, Any]: + import tempfile + + from PIL import Image + + cancel = threading.Event() + with self._generate_lock: + with self._lock: + state = self._state + if state is None: + raise RuntimeError("No diffusion model is loaded.") + self._active_generate_cancel = cancel + engine = self._resolve_engine() + try: + if seed is None: + seed = int.from_bytes(os.urandom(6), "big") & ((1 << 53) - 1) + else: + seed = int(seed) + cfg_scale, flux_guidance = _map_guidance(state.family, guidance) + extra_args: list[str] = [] + if state.vae_format: + extra_args += ["--vae-format", state.vae_format] + if state.flow_shift is not None: + extra_args += ["--flow-shift", repr(float(state.flow_shift))] + + self._gen = _SdGen(total_steps = int(steps)) + images = [] + seeds: list[int] = [] + with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: + for index in range(max(1, int(batch_size))): + if cancel.is_set(): + raise RuntimeError("Diffusion generation was cancelled.") + # Distinct seed per batch image (sd-cli is one image/run here), + # so a batch is reproducible image-by-image from the base seed. + # Mask to sd-cli's int64 range, NOT 53 bits: the request model and + # the diffusers backend both accept large explicit seeds, so a tight + # 2**53 mask would silently truncate them (2**53 -> 0) and collide + # distinct requested seeds onto the same image. Randomly-drawn seeds + # above are already 53-bit (JS-safe); explicit seeds pass through. + seed_i = (seed + index) & ((1 << 63) - 1) + out_path = str(Path(tmpdir) / f"img_{index}.png") + params = SdCppGenParams( + prompt = prompt, + negative_prompt = negative_prompt or None, + width = int(width), + height = int(height), + steps = int(steps), + cfg_scale = cfg_scale, + guidance = flux_guidance, + seed = seed_i, + sampling_method = state.sampling_method, + batch_count = 1, + ) + engine.generate( + state.files, + params, + output_path = out_path, + offload = list(state.offload_flags) or None, + native_speed = state.native_speed, + threads = state.threads, + extra_args = extra_args or None, + on_log = self._on_log, + cancel_event = cancel, + ) + with Image.open(out_path) as im: + images.append(im.copy()) + seeds.append(seed_i) + if cancel.is_set(): + raise RuntimeError("Diffusion generation was cancelled.") + # ``seeds`` is the per-image seed (each sd-cli run used seed+index), so + # the route can persist the real seed for every image in the batch. + return { + "images": images, + "seed": int(seed), + "seeds": seeds, + "repo_id": state.repo_id, + } + except SdCppCancelled as exc: + raise RuntimeError("Diffusion generation was cancelled.") from exc + finally: + self._gen = None + with self._lock: + if self._active_generate_cancel is cancel: + self._active_generate_cancel = None + + def _on_log(self, line: str) -> None: + gen = self._gen + if gen is None or gen.total_steps <= 0: + return + for a, b in _STEP_RE.findall(line): + if int(b) == gen.total_steps: + now = time.time() + gen.step = min(int(a), gen.total_steps) + if gen.first_step_at == 0.0: + gen.first_step_at = now + gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now) + + def generate_progress(self) -> dict[str, Any]: + gen = self._gen + if gen is None or gen.total_steps <= 0: + return { + "active": False, + "step": 0, + "total_steps": 0, + "fraction": 0.0, + "eta_seconds": None, + } + return { + "active": True, + "step": gen.step, + "total_steps": gen.total_steps, + "fraction": min(gen.step / gen.total_steps, 1.0), + "eta_seconds": gen.eta_seconds, + } + + # ── Unload / status ────────────────────────────────────────────────────── + + def unload(self) -> dict[str, Any]: + self._cancel_event.set() + with self._lock: + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + self._state = None + self._load_token += 1 + self._loading = None + return self.status() + + def status(self) -> dict[str, Any]: + state = self._state + if state is None: + return { + "loaded": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "cpu_offload": False, + "offload_policy": None, + "vae_tiling": False, + "memory_mode": None, + "speed_mode": None, + "speed_optims": [], + "text_encoder_quant": None, + "transformer_quant": None, + "attention_backend": None, + "transformer_cache": None, + "engine": "sd_cpp", + } + return { + "loaded": True, + "repo_id": state.repo_id, + "family": state.family.name, + "base_repo": state.base_repo, + "device": state.device, + "dtype": "gguf", + # Reflect the offload flags actually passed to sd-cli, so a balanced/low_vram + # (or cpu_offload) load is verifiable from status instead of always reading + # "none". On CPU _run_load leaves offload_flags empty (the flags are no-ops), + # so this correctly stays "none" there. + "cpu_offload": bool(state.offload_flags), + "offload_policy": "active" if state.offload_flags else "none", + "vae_tiling": False, + "memory_mode": None, + "speed_mode": state.native_speed, + "speed_optims": [], + "text_encoder_quant": None, + "transformer_quant": None, + "attention_backend": None, + "transformer_cache": None, + "engine": "sd_cpp", + } + + +def _install_allowed() -> bool: + """Whether lazy binary install is permitted (UNSLOTH_DIFFUSION_SD_CPP_INSTALL).""" + val = os.environ.get("UNSLOTH_DIFFUSION_SD_CPP_INSTALL", "auto").strip().lower() + return val not in ("0", "off", "false", "no") + + +def _progress( + phase: Optional[str], + bytes_downloaded: int = 0, + bytes_total: int = 0, + fraction: float = 0.0, + *, + error: Optional[str] = None, +) -> dict[str, Any]: + return { + "phase": phase, + "bytes_downloaded": bytes_downloaded, + "bytes_total": bytes_total, + "fraction": fraction, + "error": error, + } + + +_sd_cpp_backend: Optional[SdCppDiffusionBackend] = None + + +def get_sd_cpp_backend() -> SdCppDiffusionBackend: + global _sd_cpp_backend + if _sd_cpp_backend is None: + _sd_cpp_backend = SdCppDiffusionBackend() + return _sd_cpp_backend diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 186ea51631..6714151045 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -26,6 +26,7 @@ import logging import os import queue import shutil +import signal import subprocess import sys import threading @@ -50,6 +51,30 @@ _BINARY_STEM = "sd-cli" _LEGACY_STEM = "sd" +class SdCppCancelled(RuntimeError): + """A generation was cancelled via its ``cancel_event`` (unload / superseding load + / arbiter eviction). Distinct from a generation *failure* so the caller can keep + cancellation semantics (no diffusers fallback, no error surfaced as a crash).""" + + +def _terminate(proc: "subprocess.Popen") -> None: + """Hard-stop an sd-cli process (and any children). On POSIX the process is its + own session leader (``start_new_session``), so kill the whole group; otherwise + fall back to killing just the process.""" + if proc.poll() is not None: + return + try: + if os.name == "posix": + os.killpg(os.getpgid(proc.pid), signal.SIGKILL) + else: + proc.kill() + except Exception: # noqa: BLE001 -- killpg can miss (no pgid / already gone); fall back + try: + proc.kill() + except Exception: # noqa: BLE001 -- best-effort teardown + pass + + def _binary_name(stem: str) -> str: return f"{stem}.exe" if sys.platform == "win32" else stem @@ -164,7 +189,10 @@ class SdCppEngine: 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.""" + """First line of ``sd-cli --version``, cached on success. ``None`` when the + binary is absent OR present-but-unrunnable (exec error / nonzero exit, e.g. + missing shared libraries / bad permissions), so callers can fail a load early + instead of committing a "ready" state that crashes on first generation.""" if not self.is_available(): return None if self._version is not None: @@ -179,10 +207,12 @@ class SdCppEngine: 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 None + if res.returncode != 0: + return None + text = ((res.stdout or "") + "\n" + (res.stderr or "")).strip() + self._version = text.splitlines()[0] if text else "" return self._version def generate( @@ -199,6 +229,7 @@ class SdCppEngine: timeout: Optional[float] = 1800.0, env: Optional[dict[str, str]] = None, on_log: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, ) -> Path: """Run one ``sd-cli`` generation; return the written image path. @@ -206,7 +237,9 @@ class SdCppEngine: (``--diffusion-fa`` etc.), de-duplicated against the offload flags that may already include them. 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. + receives each line of sd-cli's progress output as it arrives. ``cancel_event`` + (if given) is polled while the child runs; when set, the process tree is + killed and ``SdCppCancelled`` is raised. """ offload = list(offload or []) speed = [f for f in native_speed_flags(native_speed) if f not in offload] @@ -221,7 +254,14 @@ class SdCppEngine: verbose = verbose, extra_args = merged_extra, ) - return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + return self._run( + cmd, + output_path, + timeout = timeout, + env = env, + on_log = on_log, + cancel_event = cancel_event, + ) def upscale( self, @@ -233,6 +273,7 @@ class SdCppEngine: timeout: Optional[float] = 1800.0, env: Optional[dict[str, str]] = None, on_log: Optional[Callable[[str], None]] = None, + cancel_event: Optional[threading.Event] = None, ) -> Path: """Upscale an image with an ESRGAN model; return the written path.""" cmd = build_sd_cpp_upscale_command( @@ -242,7 +283,14 @@ class SdCppEngine: verbose = verbose, extra_args = extra_args, ) - return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) + return self._run( + cmd, + output_path, + timeout = timeout, + env = env, + on_log = on_log, + cancel_event = cancel_event, + ) # ── internals ───────────────────────────────────────────────────────────── @@ -268,11 +316,13 @@ class SdCppEngine: timeout: Optional[float], env: Optional[dict[str, str]], on_log: Optional[Callable[[str], None]], + cancel_event: Optional[threading.Event] = None, ) -> Path: """Run an sd-cli argv, stream output, and return the produced image path. - Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output. - Shared by ``generate`` and ``upscale``. + Raises ``RuntimeError`` on nonzero exit, timeout, or a missing output, and + ``SdCppCancelled`` when ``cancel_event`` fires. Shared by ``generate`` and + ``upscale``. """ out = Path(output_path) base = dict(os.environ) @@ -289,6 +339,9 @@ class SdCppEngine: text = True, errors = "replace", env = run_env, + # Own session/process group so cancellation/timeout can kill the whole + # tree, not just the parent (POSIX only; harmless flag elsewhere). + start_new_session = (os.name == "posix"), ) # Drain stdout on a reader thread so the timeout is enforced even when the # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain @@ -313,8 +366,13 @@ class SdCppEngine: stdout_done = False try: while True: + # Cancellation (unload / superseding load / arbiter eviction): kill the + # process tree and signal the caller it was cancelled, not a failure. + if cancel_event is not None and cancel_event.is_set() and proc.poll() is None: + _terminate(proc) + raise SdCppCancelled("sd-cli generation was cancelled.") if deadline is not None and time.monotonic() >= deadline and proc.poll() is None: - proc.kill() + _terminate(proc) raise RuntimeError(f"sd-cli timed out after {timeout}s") try: line = line_q.get(timeout = 0.1) @@ -335,7 +393,7 @@ class SdCppEngine: ret = proc.wait(timeout = 5.0) finally: if proc.poll() is None: - proc.kill() + _terminate(proc) if ret != 0: raise RuntimeError(f"sd-cli exited {ret}. Last output:\n" + "\n".join(tail[-12:])) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 21e9760aed..2c1701f6ea 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1912,3 +1912,8 @@ class DiffusionStatusResponse(BaseModel): "_native_cudnn), or null for the default SDPA", ) transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") + engine: Optional[str] = Field(None, description = "Active diffusion engine: diffusers | sd_cpp") + fallback_reason: Optional[str] = Field( + None, + description = "Why diffusers was chosen over the native sd.cpp engine (null when none)", + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 38dd2d1f20..400dc239fc 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10299,15 +10299,22 @@ async def load_diffusion_model( request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject) ): from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_device import resolve_diffusion_device_target + from core.inference.diffusion_engine_router import ( + active_engine_name, + annotate_status, + select_and_activate_engine, + ) from core.inference.gpu_arbiter import acquire_for, DIFFUSION + from core.inference.sd_cpp_engine import ENGINE_SD_CPP from utils.native_path_leases import redact_native_paths backend = get_diffusion_backend() try: # Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, - # missing local GGUF) must not evict a working chat model and then 400. - # validate_load_request does local-path I/O, so run it off the event loop. - await asyncio.to_thread( + # missing local GGUF) must not evict a working chat model and then 400. The + # validated family also drives engine selection below. + fam = await asyncio.to_thread( backend.validate_load_request, request.model_path, gguf_filename = request.gguf_filename, @@ -10317,11 +10324,22 @@ async def load_diffusion_model( # compete with the training subprocess for VRAM. The chat path does the # same via _guard_chat_load_against_training; this is its image sibling. _guard_diffusion_load_against_training() - # Now take the GPU from the chat backend, then kick the (slow) load onto a - # background thread and return at once — the client polls images/load-progress. - await asyncio.to_thread(acquire_for, DIFFUSION) + # Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU), + # installing the sd-cli binary if needed -- all BEFORE evicting chat, so a + # native fallback never strands a half-loaded state. + engine = await asyncio.to_thread(select_and_activate_engine, fam, hf_token = request.hf_token) + # Take the GPU from the chat backend only when this load will actually use it. + # diffusers always does; a *force-native* sd.cpp load on a CUDA/XPU/MPS box does + # too. But a native sd.cpp load on a pure-CPU host never touches the GPU, so + # acquiring would evict the resident chat model for nothing -- skip the handoff. + device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) + needs_gpu = active_engine_name() != ENGINE_SD_CPP or device != "cpu" + if needs_gpu: + # Then kick the (slow) load onto a background thread and return at once -- + # the client polls images/load-progress. + await asyncio.to_thread(acquire_for, DIFFUSION) status_dict = await asyncio.to_thread( - backend.begin_load, + engine.begin_load, request.model_path, gguf_filename = request.gguf_filename, base_repo = request.base_repo, @@ -10338,7 +10356,7 @@ async def load_diffusion_model( transformer_cache = request.transformer_cache, transformer_cache_threshold = request.transformer_cache_threshold, ) - return DiffusionStatusResponse(**status_dict) + return DiffusionStatusResponse(**annotate_status(status_dict)) except (ValueError, FileNotFoundError) as exc: raise HTTPException(status_code = 400, detail = redact_native_paths(str(exc))) except RuntimeError as exc: @@ -10351,9 +10369,9 @@ async def generate_diffusion_image( request: DiffusionGenerateRequest, current_subject: str = Depends(get_current_subject) ): from core.inference import image_gallery - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import get_active_diffusion_engine - backend = get_diffusion_backend() + backend = get_active_diffusion_engine() try: result = await asyncio.to_thread( backend.generate, @@ -10367,26 +10385,32 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - if not backend.is_loaded: - # The only genuine client-state 409: nothing is loaded to generate with. - raise HTTPException(status_code = 409, detail = "No diffusion model is loaded.") - # A pipeline RuntimeError (CUDA OOM, shape/device) is a server failure; fall - # through to the sanitized 500 instead of echoing raw exception text (which - # would 409 an OOM as retryable and leak VRAM totals / tensor shapes). + # Only "no model loaded" / cancelled are client-state (409). The native + # sd.cpp engine also raises RuntimeError for execution failures (nonzero + # exit, timeout, missing output), which are server errors (500). + msg = str(exc) + if "No diffusion model is loaded" in msg or "cancelled" in msg.lower(): + raise HTTPException(status_code = 409, detail = msg) logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") except Exception as exc: logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") - # Persist each image with its full recipe embedded. The whole batch shares - # one seed (drawn sequentially from one generator) and one timestamp (the - # images are generated together), so their gallery order is stable on reload. + # Persist each image with its full recipe embedded. The diffusers batch shares + # one seed (drawn sequentially from one generator); the native sd.cpp batch uses a + # distinct seed per image and returns them in ``seeds`` so each is reproducible. created_at = time.time() + per_image_seeds = result.get("seeds") def _persist() -> list[dict]: records = [] for index, image in enumerate(result["images"]): + seed = ( + per_image_seeds[index] + if per_image_seeds and index < len(per_image_seeds) + else result["seed"] + ) records.append( image_gallery.save( image, @@ -10397,9 +10421,9 @@ async def generate_diffusion_image( "height": request.height, "steps": request.steps, "guidance": request.guidance, - "seed": result["seed"], - # Position within the batch: images here share a seed + timestamp, - # so the export filename needs this to stay unique. + "seed": seed, + # Position within the batch: shared timestamp, so the export + # filename needs this to stay unique. "batch_index": index, # The batch shares one seed, so reproducing image batch_index>0 # needs the original batch_size: persist it so restore can replay. @@ -10476,27 +10500,27 @@ async def clear_gallery_images(current_subject: str = Depends(get_current_subjec @studio_router.post("/images/unload", response_model = DiffusionStatusResponse) async def unload_diffusion_model(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend + from core.inference.diffusion_engine_router import annotate_status, get_active_diffusion_engine from core.inference.gpu_arbiter import release, DIFFUSION - status_dict = await asyncio.to_thread(get_diffusion_backend().unload) + status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload) release(DIFFUSION) - return DiffusionStatusResponse(**status_dict) + return DiffusionStatusResponse(**annotate_status(status_dict)) @studio_router.get("/images/status", response_model = DiffusionStatusResponse) async def diffusion_status(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionStatusResponse(**get_diffusion_backend().status()) + from core.inference.diffusion_engine_router import active_status + return DiffusionStatusResponse(**active_status()) @studio_router.get("/images/load-progress", response_model = DiffusionLoadProgressResponse) async def diffusion_load_progress(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionLoadProgressResponse(**get_diffusion_backend().load_progress()) + from core.inference.diffusion_engine_router import get_active_diffusion_engine + return DiffusionLoadProgressResponse(**get_active_diffusion_engine().load_progress()) @studio_router.get("/images/generate-progress", response_model = DiffusionGenerateProgressResponse) async def diffusion_generate_progress(current_subject: str = Depends(get_current_subject)): - from core.inference.diffusion import get_diffusion_backend - return DiffusionGenerateProgressResponse(**get_diffusion_backend().generate_progress()) + from core.inference.diffusion_engine_router import get_active_diffusion_engine + return DiffusionGenerateProgressResponse(**get_active_diffusion_engine().generate_progress()) diff --git a/studio/backend/tests/test_diffusion_engine_router.py b/studio/backend/tests/test_diffusion_engine_router.py new file mode 100644 index 0000000000..07d495757b --- /dev/null +++ b/studio/backend/tests/test_diffusion_engine_router.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the diffusion engine router (diffusers vs native sd.cpp selection).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from core.inference import diffusion_engine_router as r +from core.inference.diffusion_families import detect_family +from core.inference.sd_cpp_engine import ENGINE_DIFFUSERS, ENGINE_SD_CPP + +_ENVS = ( + "UNSLOTH_DIFFUSION_ENGINE", + "UNSLOTH_DIFFUSION_SD_CPP", + "UNSLOTH_DIFFUSION_SD_CPP_MPS", + "UNSLOTH_DIFFUSION_SD_CPP_INSTALL", +) + + +@pytest.fixture(autouse = True) +def _clean_env_and_state(monkeypatch): + for e in _ENVS: + monkeypatch.delenv(e, raising = False) + # A light status-capable stub so neither selection nor active_status() imports the + # heavy diffusers/sd.cpp backends; the active engine NAME comes from module state. + monkeypatch.setattr( + r, + "get_active_diffusion_engine", + lambda: SimpleNamespace(status = lambda: {"loaded": False, "repo_id": None}), + ) + yield + + +def _set_device(monkeypatch, backend): + monkeypatch.setattr( + r, + "resolve_diffusion_device_target", + lambda: SimpleNamespace(backend = backend, device = backend), + ) + + +def _set_binary(monkeypatch, path): + monkeypatch.setattr(r, "ensure_sd_cpp_binary", lambda **_: path) + + +def _set_runnable(monkeypatch, version = "sd-cli v0"): + """Stub the runnability probe so a stubbed binary path is treated as executable + (the router now probes ``SdCppEngine(...).version()`` before committing to native).""" + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: version)) + + +def _select(fam_name = "z-image"): + """Activate the engine for a family and return which engine was chosen.""" + r.select_and_activate_engine(detect_family(fam_name)) + return r.active_engine_name() + + +# ── core selection matrix ───────────────────────────────────────────────────── + + +def test_cpu_with_binary_and_supported_family_picks_sd_cpp(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + assert _select() == ENGINE_SD_CPP + assert r.active_engine_name() == ENGINE_SD_CPP + + +def test_present_but_not_runnable_binary_falls_back(monkeypatch): + # A binary that exists but cannot run (version() -> None) must fall back to + # diffusers at selection, not commit native and fail inside the load. + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setattr(r, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: None)) + assert _select() == ENGINE_DIFFUSERS + assert "binary unavailable" in (r.active_status()["fallback_reason"] or "") + + +@pytest.mark.parametrize("gpu", ["cuda", "rocm", "xpu"]) +def test_gpu_backends_use_diffusers(monkeypatch, gpu): + _set_device(monkeypatch, gpu) + _set_binary(monkeypatch, "/usr/bin/sd-cli") # even with a binary, GPU stays diffusers + assert _select() == ENGINE_DIFFUSERS + assert "uses diffusers" in (r.active_status()["fallback_reason"] or "") + + +def test_forced_diffusers_overrides_cpu(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "diffusers") + assert _select() == ENGINE_DIFFUSERS + assert "forced" in (r.active_status()["fallback_reason"] or "") + + +def test_sd_cpp_disabled_uses_diffusers(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP", "0") + assert _select() == ENGINE_DIFFUSERS + assert "disabled" in (r.active_status()["fallback_reason"] or "") + + +def test_mps_default_diffusers_but_optin_sd_cpp(monkeypatch): + _set_device(monkeypatch, "mps") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + # Default: MPS is not native-eligible -> diffusers. + assert _select() == ENGINE_DIFFUSERS + # Opt in: MPS routes to sd.cpp. + monkeypatch.setenv("UNSLOTH_DIFFUSION_SD_CPP_MPS", "1") + assert _select() == ENGINE_SD_CPP + + +def test_unsupported_family_falls_back(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + monkeypatch.setattr(r, "family_sd_cpp_supported", lambda fam: False) + assert _select() == ENGINE_DIFFUSERS + assert "no native sd.cpp asset mapping" in (r.active_status()["fallback_reason"] or "") + + +def test_missing_binary_falls_back(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) # install unavailable + assert _select() == ENGINE_DIFFUSERS + assert "binary unavailable" in (r.active_status()["fallback_reason"] or "") + + +def test_force_sd_cpp_on_gpu_when_binary_present(monkeypatch): + _set_device(monkeypatch, "cuda") + _set_binary(monkeypatch, "/usr/bin/sd-cli") + _set_runnable(monkeypatch) + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp") + assert _select() == ENGINE_SD_CPP + + +def test_force_sd_cpp_without_binary_falls_back(monkeypatch): + _set_device(monkeypatch, "cuda") + _set_binary(monkeypatch, None) + monkeypatch.setenv("UNSLOTH_DIFFUSION_ENGINE", "sd_cpp") + assert _select() == ENGINE_DIFFUSERS + + +# ── active_status annotation ────────────────────────────────────────────────── + + +def test_active_status_injects_engine_and_reason(monkeypatch): + _set_device(monkeypatch, "cpu") + _set_binary(monkeypatch, None) + _select() # -> diffusers fallback (no binary) + st = r.active_status() + assert st["engine"] == ENGINE_DIFFUSERS + assert st["fallback_reason"] and "binary unavailable" in st["fallback_reason"] diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d03c48b477..d0b6997d07 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -114,6 +114,26 @@ def _unloaded_status(): def client(monkeypatch, tmp_path): backend = _FakeBackend() monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: backend) + # Neutralise the engine router so the routes deterministically drive this fake + # (diffusers) backend regardless of the host's real device, and never attempt a + # native sd.cpp install/download. The router's selection logic is covered in + # test_diffusion_engine_router.py; one route-level sd_cpp test lives below. + import core.inference.diffusion_engine_router as engine_router + + # Delegate to whatever get_diffusion_backend currently returns, so per-test + # re-patches of the backend still flow through the routes. + monkeypatch.setattr( + engine_router, + "select_and_activate_engine", + lambda fam, **kw: diffusion_module.get_diffusion_backend(), + ) + monkeypatch.setattr( + engine_router, + "get_active_diffusion_engine", + lambda: diffusion_module.get_diffusion_backend(), + ) + monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers") + monkeypatch.setattr(engine_router, "_fallback_reason", None) # Isolate from the real GPU arbiter: reset ownership and stub the evictors so # the load route's acquire_for() never touches live backend singletons. monkeypatch.setattr(gpu_arbiter, "_owner", None) @@ -501,6 +521,62 @@ def test_out_of_range_cache_threshold_returns_422(client): assert resp.status_code == 422 +def test_load_routes_to_sd_cpp_on_cpu(monkeypatch, tmp_path): + """End-to-end through the REAL router: a CPU host with an available binary routes + the load to the native sd.cpp engine and the response reports engine=sd_cpp.""" + from types import SimpleNamespace + + import core.inference.diffusion_engine_router as engine_router + import core.inference.sd_cpp_backend as sd_backend + + for e in ( + "UNSLOTH_DIFFUSION_ENGINE", + "UNSLOTH_DIFFUSION_SD_CPP", + "UNSLOTH_DIFFUSION_SD_CPP_MPS", + "UNSLOTH_DIFFUSION_SD_CPP_INSTALL", + ): + monkeypatch.delenv(e, raising = False) + + validator = _FakeBackend() # supplies validate_load_request (and is the diffusers fallback) + monkeypatch.setattr(diffusion_module, "get_diffusion_backend", lambda: validator) + # Force the router's decision inputs: CPU device + an available binary. + monkeypatch.setattr( + engine_router, + "resolve_diffusion_device_target", + lambda: SimpleNamespace(backend = "cpu", device = "cpu"), + ) + monkeypatch.setattr(engine_router, "ensure_sd_cpp_binary", lambda **_: "/x/sd-cli") + # The router now probes runnability before committing to native; treat the stub + # binary as executable. + monkeypatch.setattr( + engine_router, "SdCppEngine", lambda **_: SimpleNamespace(version = lambda: "sd-cli v0") + ) + monkeypatch.setattr(engine_router, "_active_engine_name", "diffusers") + monkeypatch.setattr(engine_router, "_fallback_reason", None) + # The native backend the router will activate. + sd_fake = _FakeBackend() + monkeypatch.setattr(sd_backend, "get_sd_cpp_backend", lambda: sd_fake) + + monkeypatch.setattr(gpu_arbiter, "_owner", None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.CHAT, lambda: None) + monkeypatch.setitem(gpu_arbiter._EVICTORS, gpu_arbiter.DIFFUSION, lambda: None) + + app = FastAPI() + app.include_router(studio_router, prefix = "/api/inference") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + client = TestClient(app) + + resp = client.post( + "/api/inference/images/load", + json = {"model_path": "unsloth/Z-Image-Turbo-GGUF", "gguf_filename": "z.gguf"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["engine"] == "sd_cpp" + assert body["fallback_reason"] is None + assert sd_fake.loaded is True # the native engine actually received the load + + def test_invalid_transformer_quant_returns_422_without_eviction(client): # An unsupported transformer_quant is rejected by the request schema (Literal), so # the GPU is never acquired and no chat model is evicted. @@ -537,3 +613,48 @@ def test_in_progress_returns_409_after_validation_passes(client, monkeypatch): assert resp.status_code == 409 # Validation passed first, so the GPU WAS acquired before begin_load reported busy. assert gpu_arbiter._owner == gpu_arbiter.DIFFUSION + + +def _force_engine(monkeypatch, backend, *, engine_name, device): + """Pin engine selection + device so the load route's arbiter gating is deterministic.""" + import types as _types + + import core.inference.diffusion_device as devmod + import core.inference.diffusion_engine_router as router + + monkeypatch.setattr(router, "select_and_activate_engine", lambda fam, **kw: backend) + monkeypatch.setattr(router, "active_engine_name", lambda: engine_name) + monkeypatch.setattr( + devmod, "resolve_diffusion_device_target", lambda: _types.SimpleNamespace(device = device) + ) + acquired: list = [] + monkeypatch.setattr(gpu_arbiter, "acquire_for", lambda role: acquired.append(role)) + return acquired + + +def test_cpu_native_load_skips_gpu_arbiter(client, monkeypatch): + # A native sd.cpp load on a pure-CPU host never touches the GPU, so the route must NOT + # evict the resident chat model -- the arbiter handoff is skipped. + from core.inference.sd_cpp_engine import ENGINE_SD_CPP + + backend = diffusion_module.get_diffusion_backend() + acquired = _force_engine(monkeypatch, backend, engine_name = ENGINE_SD_CPP, device = "cpu") + resp = client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + assert resp.status_code == 200 + assert acquired == [] # no arbiter handoff for a CPU native load + + +def test_gpu_native_load_takes_arbiter(client, monkeypatch): + # A force-native sd.cpp load on a GPU box DOES use the GPU, so the arbiter is acquired + # (same as the always-GPU diffusers path). + from core.inference.sd_cpp_engine import ENGINE_SD_CPP + + backend = diffusion_module.get_diffusion_backend() + acquired = _force_engine(monkeypatch, backend, engine_name = ENGINE_SD_CPP, device = "cuda") + resp = client.post( + "/api/inference/images/load", json = {"model_path": "x/z-image", "gguf_filename": "q.gguf"} + ) + assert resp.status_code == 200 + assert acquired == [gpu_arbiter.DIFFUSION] diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py new file mode 100644 index 0000000000..79464dd6be --- /dev/null +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the native sd.cpp diffusion backend (the no-GPU engine).""" + +from __future__ import annotations + +import threading +import types + +import pytest +from PIL import Image + +from core.inference import sd_cpp_backend as bk +from core.inference.diffusion_families import detect_family +from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles +from core.inference.sd_cpp_backend import ( + SdCppDiffusionBackend, + _map_guidance, + ensure_sd_cpp_binary, +) +from core.inference.sd_cpp_engine import SdCppCancelled + + +class _FakeEngine: + """Stands in for SdCppEngine: writes a 1x1 PNG and records the args.""" + + def __init__( + self, + *, + fail = None, + cancel_on_call = False, + ): + self.calls = [] + self.fail = fail + self.cancel_on_call = cancel_on_call + + def is_available(self): + return True + + def version(self, **_): + return "fake sd-cli" + + def generate( + self, + files, + params, + *, + output_path, + cancel_event = None, + **kw, + ): + self.calls.append((files, params, output_path, kw)) + if self.cancel_on_call and cancel_event is not None: + cancel_event.set() + if self.fail is not None: + raise self.fail + if cancel_event is not None and cancel_event.is_set(): + raise SdCppCancelled("cancelled") + Image.new("RGB", (1, 1), (10, 20, 30)).save(output_path) + from pathlib import Path + + return Path(output_path) + + +def _loaded_backend(fam_name = "z-image", engine = None): + b = SdCppDiffusionBackend(engine = engine or _FakeEngine()) + fam = detect_family(fam_name) + b._state = bk._SdState( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + base_repo = fam.base_repo, + family = fam, + device = "cpu", + files = SdCppModelFiles( + diffusion_model = "/m/z.gguf", vae = "/m/vae.safetensors", llm = "/m/llm.safetensors" + ), + vae_format = fam.sd_cpp_vae_format, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, + ) + return b + + +# ── asset resolution ────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize( + "fam_name,expect_kinds", + [ + ("flux.1", {"diffusion_model", "vae", "clip_l", "t5xxl"}), + ("z-image", {"diffusion_model", "vae", "llm"}), + ("qwen-image", {"diffusion_model", "vae", "qwen2vl"}), + ("flux.2-klein", {"diffusion_model", "vae", "llm"}), + ], +) +def test_asset_specs_cover_required_files(fam_name, expect_kinds): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family(fam_name) + specs = b._asset_specs("unsloth/x-GGUF", "x-Q4_K_M.gguf", fam) + kinds = {kind for _, _, kind in specs} + assert kinds == expect_kinds + # Every spec has a non-empty repo + filename. + assert all(repo and fn for repo, fn, _ in specs) + # The transformer reuses the requested GGUF, not a registry file. + tr = [s for s in specs if s[2] == "diffusion_model"][0] + assert tr[0] == "unsloth/x-GGUF" and tr[1] == "x-Q4_K_M.gguf" + + +# ── guidance mapping ────────────────────────────────────────────────────────── + + +def test_map_guidance_flux_uses_distilled_guidance(): + cfg, g = _map_guidance(detect_family("flux.1"), 3.5) + assert cfg is None and g == 3.5 + + +def test_map_guidance_cfg_family_off_when_distilled(): + # qwen-image uses real CFG; a distilled 0 -> CFG off (1.0), a >1 value passes through. + assert _map_guidance(detect_family("qwen-image"), 0.0) == (1.0, None) + assert _map_guidance(detect_family("qwen-image"), 4.0) == (4.0, None) + + +# ── status ──────────────────────────────────────────────────────────────────── + + +def test_status_unloaded_reports_sd_cpp_engine(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + st = b.status() + assert st["loaded"] is False and st["engine"] == "sd_cpp" + + +def test_status_loaded_shape(): + b = _loaded_backend() + st = b.status() + assert st["loaded"] is True + assert st["engine"] == "sd_cpp" + assert st["family"] == "z-image" + assert st["device"] == "cpu" + # diffusers-only fields are present (route response parity) but null. + for k in ("transformer_quant", "attention_backend", "transformer_cache", "text_encoder_quant"): + assert st[k] is None + + +# ── generate ────────────────────────────────────────────────────────────────── + + +def test_generate_returns_images_and_seed(): + eng = _FakeEngine() + b = _loaded_backend(engine = eng) + out = b.generate(prompt = "a fox", width = 64, height = 64, steps = 8, seed = 123, batch_size = 2) + assert out["seed"] == 123 + assert out["repo_id"] == "unsloth/Z-Image-Turbo-GGUF" + assert len(out["images"]) == 2 + assert all(isinstance(im, Image.Image) for im in out["images"]) + # One sd-cli run per batch image, each a distinct seed from the base. + assert len(eng.calls) == 2 + seeds = [params.seed for _, params, _, _ in eng.calls] + assert seeds == [123, 124] + # The per-image seeds are returned so the route can persist each one. + assert out["seeds"] == [123, 124] + + +def test_generate_qwen_passes_sampling_args(): + eng = _FakeEngine() + b = _loaded_backend(fam_name = "qwen-image", engine = eng) + b.generate(prompt = "x", steps = 20, guidance = 4.0, seed = 1) + _, params, _, kw = eng.calls[0] + assert params.sampling_method == "euler" # Qwen's supported sd.cpp sampler + assert "--flow-shift" in (kw.get("extra_args") or []) + + +def test_generate_raises_when_not_loaded(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + with pytest.raises(RuntimeError, match = "No diffusion model is loaded"): + b.generate(prompt = "x") + + +def test_generate_passes_vae_format_for_flux2(): + eng = _FakeEngine() + b = _loaded_backend(fam_name = "flux.2-klein", engine = eng) + b.generate(prompt = "x", steps = 4, seed = 1) + _, _, _, kw = eng.calls[0] + assert kw.get("extra_args") == ["--vae-format", "flux2"] + + +def test_generate_cancellation_raises_cancelled_not_failure(): + # The engine cancels mid-run; the backend surfaces a cancellation, not a crash. + eng = _FakeEngine(cancel_on_call = True) + b = _loaded_backend(engine = eng) + with pytest.raises(RuntimeError, match = "cancelled"): + b.generate(prompt = "x", steps = 8, seed = 5) + + +def test_generate_progress_tracks_parsed_steps(): + b = _loaded_backend() + b._gen = bk._SdGen(total_steps = 8) + b._on_log(" sampling 4/8 done") + p = b.generate_progress() + assert p["active"] is True and p["step"] == 4 and p["total_steps"] == 8 + # A fraction with a different denominator must not move the bar. + b._on_log("loaded 1/3 tensors") + assert b.generate_progress()["step"] == 4 + + +# ── load validation + binary install ────────────────────────────────────────── + + +def test_begin_load_rejects_unsupported_family(monkeypatch): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + # A family with no native asset mapping must be rejected (router falls back). + monkeypatch.setattr(bk, "family_sd_cpp_supported", lambda fam: False) + with pytest.raises(ValueError, match = "no native sd.cpp asset mapping"): + b.begin_load("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "z.gguf") + + +def test_begin_load_requires_gguf_filename(): + b = SdCppDiffusionBackend(engine = _FakeEngine()) + with pytest.raises(ValueError, match = "gguf_filename is required"): + b.begin_load("unsloth/Z-Image-Turbo-GGUF") + + +def test_ensure_binary_returns_found(monkeypatch): + monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli") + assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli" + + +def test_ensure_binary_install_disabled_returns_none(monkeypatch): + monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: None) + assert ensure_sd_cpp_binary(allow_install = False) is None + + +def test_unload_clears_state_and_signals_cancel(): + cancel = threading.Event() + b = _loaded_backend() + b._active_generate_cancel = cancel + st = b.unload() + assert st["loaded"] is False + assert cancel.is_set() + assert b._cancel_event.is_set() + + +def test_status_reports_offload_when_flags_active(): + # status must reflect the offload flags actually passed to sd-cli, not always "none", + # so a balanced/low_vram (or cpu_offload) load is verifiable. + b = _loaded_backend() + # No flags (CPU default) -> none. + assert b.status()["offload_policy"] == "none" and b.status()["cpu_offload"] is False + # Flags present (off-CPU offload) -> reported active. + s = b._state + b._state = bk._SdState( + repo_id = s.repo_id, + base_repo = s.base_repo, + family = s.family, + device = "cuda", + files = s.files, + offload_flags = ("--vae-on-cpu", "--clip-on-cpu"), + ) + st = b.status() + assert st["cpu_offload"] is True and st["offload_policy"] == "active" + + +def test_run_load_cancels_and_waits_for_inflight_generation(monkeypatch): + # A generation that started during the asset download is still running against the OLD + # model. _run_load must cancel it AND wait on _generate_lock before committing the new + # state, or a stale sd-cli run finishes afterward and persists an image from the previous + # model once the new load reports ready. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family("z-image") + monkeypatch.setattr(b, "_asset_specs", lambda *a, **k: []) + monkeypatch.setattr(b, "_set_expected_bytes", lambda *a, **k: None) + monkeypatch.setattr( + b, + "_fetch_assets", + lambda *a, **k: {"diffusion_model": "/m/z.gguf", "vae": "/m/vae.sft", "llm": "/m/llm.sft"}, + ) + # Avoid importing torch from the worker thread (its first import deadlocks off the main + # thread -- a test artifact, not a production path); the device only needs to be CPU here. + monkeypatch.setattr( + bk, "resolve_diffusion_device_target", lambda: types.SimpleNamespace(device = "cpu") + ) + + b._load_token = 5 + cancel = threading.Event() + b._active_generate_cancel = cancel # a generation is "in flight" + + committed = threading.Event() + + def _load(): + b._run_load( + repo_id = "unsloth/Z-Image-Turbo-GGUF", + gguf_filename = "z.gguf", + base = fam.base_repo, + fam = fam, + hf_token = None, + _load_token = 5, + ) + committed.set() + + b._generate_lock.acquire() # simulate the live denoise holding _generate_lock + try: + threading.Thread(target = _load, daemon = True).start() + # The commit must block behind the live generation and not publish the new state, + # but must already have signalled the in-flight cancel. + assert not committed.wait(0.5) + assert b._state is None + assert cancel.is_set() + finally: + b._generate_lock.release() + assert committed.wait(5) # only now does the commit run + assert b._state is not None and b._state.repo_id == "unsloth/Z-Image-Turbo-GGUF" diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index a879252761..054574d905 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -78,7 +78,10 @@ def test_find_returns_none_when_absent(tmp_path, monkeypatch): # ── availability / version ────────────────────────────────────────────────── -def test_engine_unavailable_when_no_binary(): +def test_engine_unavailable_when_no_binary(monkeypatch): + # Force the "no binary anywhere" condition so the test is hermetic even on a host + # that happens to have sd-cli installed. + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) assert e.is_available() is False assert e.version() is None @@ -92,7 +95,9 @@ def test_engine_version_parsed_and_cached(tmp_path, monkeypatch): def _fake_run(*_a, **_k): calls["n"] += 1 - return types.SimpleNamespace(stdout = "stable-diffusion.cpp version master-721\n", stderr = "") + return types.SimpleNamespace( + stdout = "stable-diffusion.cpp version master-721\n", stderr = "", returncode = 0 + ) monkeypatch.setattr(eng.subprocess, "run", _fake_run) assert e.version() == "stable-diffusion.cpp version master-721" @@ -350,12 +355,13 @@ def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): assert "--upscale-model" in _FakePopen.captured_cmd -def test_upscale_raises_when_binary_missing(): +def test_upscale_raises_when_binary_missing(monkeypatch, tmp_path): + monkeypatch.setattr(eng, "find_sd_cpp_binary", lambda: None) e = SdCppEngine(binary = None) with pytest.raises(RuntimeError, match = "not found"): e.upscale( SdCppUpscaleParams(input_image = "/i.png", upscale_model = "/m/e.pth"), - output_path = "/tmp/x.png", + output_path = str(tmp_path / "x.png"), ) From d0c5cf6e073848caa646b8a633e04c42c0793fae Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 1 Jul 2026 16:58:39 -0300 Subject: [PATCH 12/16] Fix diffusion flag leak, sd-cli orphan, and native family fallback Six correctness fixes to the diffusion stack, found reviewing the merged phase PRs on this branch: - load_pipeline: restore the try/finally guard around the speed/quant/ placement span. A failure after apply_speed_optims (e.g. OOM in quant or the memory plan) left TF32/cudnn flags flipped process-wide and the half-built pipe resident in VRAM. Now restores the flags and frees VRAM on a failed load. - sd-cli Popen binds to the parent (PR_SET_PDEATHSIG via child_popen_kwargs, matching the llama.cpp sites), so a parent crash mid-generation can't orphan it holding VRAM/RAM. - Native begin_load uses the filename-fallback family detector the route validated with, so a local .gguf whose family keyword lives only in the basename no longer dead-ends 400 on a no-GPU host. - Generate error handler matches exact sentinel messages instead of the "cancelled" substring, fixing a 409 misroute and a raw sd-cli output leak. - find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME/STUDIO_HOME like the installer, so a custom Studio home resolves. - Drop the redundant _tf32_prev bookkeeping; snapshot/restore_backend_flags is now the single owner of the TF32/cudnn restore. The two client-state messages are now shared constants so the 409-vs-500 contract can't drift. Adds a regression test for each behavioral fix. --- studio/backend/core/inference/diffusion.py | 151 +++++++++--------- .../core/inference/diffusion_families.py | 24 +++ .../backend/core/inference/diffusion_speed.py | 44 +---- .../backend/core/inference/sd_cpp_backend.py | 17 +- .../backend/core/inference/sd_cpp_engine.py | 16 +- studio/backend/routes/inference.py | 16 +- .../backend/tests/test_diffusion_backend.py | 32 ++++ studio/backend/tests/test_diffusion_routes.py | 31 ++++ studio/backend/tests/test_sd_cpp_backend.py | 11 ++ 9 files changed, 215 insertions(+), 127 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index cba3ba2af5..6d728df035 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -24,8 +24,10 @@ from loggers import get_logger from utils.hardware import clear_gpu_cache from .diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, DiffusionFamily, - detect_family, + detect_family_for_pick, resolve_base_repo, resolve_local_gguf_child, ) @@ -245,22 +247,6 @@ class DiffusionBackend: base, rfilename, hf_token, cancel_event = self._cancel_event ) - @staticmethod - def _detect_family_for_pick( - repo_id: str, gguf_filename: Optional[str], family_override: Optional[str] - ) -> Optional[DiffusionFamily]: - """Detect the family from the repo id, falling back to the combined - path/filename for a direct local .gguf pick. The frontend splits such a - pick into (parent dir, basename), so the family keyword can live only in - the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent - directory carries none; scan it too when the directory alone is - undetectable. Only used as a fallback, so remote 'org/name' picks and - explicit overrides behave exactly as before.""" - fam = detect_family(repo_id, family_override) - if fam is None and gguf_filename and not family_override: - fam = detect_family(f"{repo_id}/{gguf_filename}", family_override) - return fam - def validate_load_request( self, repo_id: str, @@ -277,7 +263,7 @@ class DiffusionBackend: raise ValueError( "gguf_filename is required: this backend loads single-file GGUF checkpoints only." ) - fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override) + fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError( f"Could not infer a diffusion family for '{repo_id}'. Pass family_override (z-image)." @@ -368,7 +354,7 @@ class DiffusionBackend: # Resolve the base repo and estimate sizes on this thread (both network # calls) so begin_load returns instantly; the bar shows raw bytes until # the total lands. This is the only writer of _loading's fields here. - fam = self._detect_family_for_pick( + fam = detect_family_for_pick( kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override") ) base = _resolve_base_repo( @@ -641,63 +627,78 @@ class DiffusionBackend: quant_active = transformer_quant_engaged is not None or bool(gguf_filename), logger = logger, ) - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = effective_speed, - cache_active = cache_engaged is not None, - logger = logger, - ) - if transformer_quant_engaged is not None and not speed_applied.get("compiled"): - # Promotion above could not engage compile (e.g. the family is not - # compile-friendly, or compile_repeated_blocks failed): the quantized - # transformer is now running eager, which is far slower than the GGUF - # path it replaced. Surface it loudly rather than hiding the regression. - logger.warning( - "diffusion.transformer_quant: %s engaged but the transformer is NOT " - "compiled; eager torchao quant is ~30x slower than GGUF here", - transformer_quant_engaged, + # apply_speed_optims flips the process-global TF32 / cudnn.benchmark + # flags. If a later step here (text-encoder quant, memory plan) then + # raises -- e.g. OOM -- those flags would leak flipped and a subsequent + # `off` load would no longer be bit-identical. Restore the snapshot + # unless we reach the commit (unload restores on the happy path). + committed = False + try: + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + cache_active = cache_engaged is not None, + logger = logger, + ) + if transformer_quant_engaged is not None and not speed_applied.get("compiled"): + # Promotion above could not engage compile (e.g. the family is not + # compile-friendly, or compile_repeated_blocks failed): the quantized + # transformer is now running eager, which is far slower than the GGUF + # path it replaced. Surface it loudly rather than hiding the regression. + logger.warning( + "diffusion.transformer_quant: %s engaged but the transformer is NOT " + "compiled; eager torchao quant is ~30x slower than GGUF here", + transformer_quant_engaged, + ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, - ) - # Apply the placement planned above (from MEASURED free device memory vs - # the model's estimated resident size). 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. The dense fast path already placed the pipe resident; - # for the `none` policy this is an idempotent re-placement. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Apply the placement planned above (from MEASURED free device memory vs + # the model's estimated resident size). 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. The dense fast path already placed the pipe resident; + # for the `none` policy this is an idempotent re-placement. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = effective_speed, - speed_optims = tuple(k for k, v in speed_applied.items() if v), - backend_flags_before = backend_flags_before, - text_encoder_quant = te_quant, - transformer_quant = transformer_quant_engaged, - attention_backend = attention_engaged, - transformer_cache = cache_engaged, - ) + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + transformer_quant = transformer_quant_engaged, + attention_backend = attention_engaged, + transformer_cache = cache_engaged, + ) + committed = True + finally: + if not committed: + # Restore the flags AND free the half-built pipe's VRAM: the + # failed load never commits _state, so nothing else reclaims it + # until the next unload. + restore_backend_flags(backend_flags_before) + clear_gpu_cache() logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -854,7 +855,7 @@ class DiffusionBackend: with self._lock: state = self._state if state is None: - raise RuntimeError("No diffusion model is loaded.") + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. # A cancel that arrived before now either nulled _state (we raised # above) or targets an older generation, so nothing is lost. @@ -923,7 +924,7 @@ class DiffusionBackend: # A cancelled denoise returns early with a partial/garbage image; # don't hand it back to be persisted. if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Return the PIL images (not yet encoded): the route embeds each # image's recipe and persists it via the gallery. return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 3afbe9b495..3ac238e399 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -20,6 +20,14 @@ from pathlib import Path, PurePosixPath from typing import Optional +# Runtime->route contract: the RuntimeError messages a backend raises for +# client-recoverable generate states. The /images/generate route matches these +# EXACTLY to return 409 (vs a sanitized 500 for real failures), so both engines +# must raise them verbatim -- keep them named here, not as scattered literals. +DIFFUSION_NOT_LOADED_MSG = "No diffusion model is loaded." +DIFFUSION_CANCELLED_MSG = "Diffusion generation was cancelled." + + @dataclass(frozen = True) class DiffusionFamily: name: str @@ -173,6 +181,22 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff return None +def detect_family_for_pick( + repo_id: str, gguf_filename: Optional[str] = None, override: Optional[str] = None +) -> Optional[DiffusionFamily]: + """``detect_family``, falling back to the combined path/filename for a direct + local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename), + so the family keyword can live only in the filename (e.g. + ``/models/z-image-turbo-Q4_K_M.gguf``) while the parent directory carries none; + scan the combined string too when the directory alone is undetectable. Only a + fallback, so remote ``org/name`` picks and explicit overrides behave exactly as + ``detect_family``. Shared by both engines so validation and load can't diverge.""" + fam = detect_family(repo_id, override) + if fam is None and gguf_filename and not override: + fam = detect_family(f"{repo_id}/{gguf_filename}", override) + return fam + + def resolve_base_repo(fam: DiffusionFamily, base_repo: Optional[str]) -> str: """The companion diffusers repo: caller-supplied if given, else the family fallback.""" base = (base_repo or "").strip() diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 208cb8b513..5f1ef64445 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -161,12 +161,11 @@ def apply_speed_optims( "compiled": False, } mode = normalize_speed_mode(speed_mode) - # TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this - # load isn't max, so a later default/off diffusion load -- or chat inference in the - # same long-lived process -- doesn't silently inherit a prior max load's TF32 and - # lose the bit-identical default the regression harness checks. - if mode != SPEED_MAX: - _restore_tf32(logger) + # TF32 and cudnn.benchmark are the process-global flags this may flip (TF32 on max, + # cudnn.benchmark on any non-off CUDA load). The caller snapshots them before this + # call and restores on unload / failed load via snapshot_backend_flags / + # restore_backend_flags, so a later `off` load -- or chat inference in the same + # process -- never inherits them. We keep no separate bookkeeping here. if mode == SPEED_OFF: return applied @@ -251,22 +250,9 @@ def _enable_cudnn_benchmark(logger: Any) -> bool: return False -# The TF32 flag values from before the first max load flipped them, so a later -# non-max load / unload can put the process back exactly as it found it (rather than -# forcing a hardcoded default that might clobber another component's choice). -_tf32_prev: Optional[tuple[bool, bool]] = None - - def _enable_tf32(logger: Any) -> bool: - global _tf32_prev try: import torch - - if _tf32_prev is None: - _tf32_prev = ( - torch.backends.cuda.matmul.allow_tf32, - torch.backends.cudnn.allow_tf32, - ) torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True return True @@ -275,26 +261,6 @@ def _enable_tf32(logger: Any) -> bool: return False -def restore_tf32(logger: Any = None) -> None: - """Put the process-global TF32 flags back to their pre-max-load values. No-op if - a max load never set them. Called on a non-max load and on unload.""" - _restore_tf32(logger) - - -def _restore_tf32(logger: Any) -> None: - global _tf32_prev - if _tf32_prev is None: - return - try: - import torch - torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0] - torch.backends.cudnn.allow_tf32 = _tf32_prev[1] - except Exception as exc: # noqa: BLE001 — best-effort restore - _warn(logger, "tf32_restore", exc) - finally: - _tf32_prev = None - - def _fuse_qkv(pipe: Any, logger: Any) -> bool: for owner in (pipe, getattr(pipe, "transformer", None)): fn = getattr(owner, "fuse_qkv_projections", None) diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 4d0bbcbfa7..88661f2be1 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -35,8 +35,10 @@ from typing import Any, Optional from core.inference.diffusion_device import resolve_diffusion_device_target from core.inference.diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, DiffusionFamily, - detect_family, + detect_family_for_pick, family_sd_cpp_supported, resolve_base_repo, resolve_local_gguf_child, @@ -242,7 +244,10 @@ class SdCppDiffusionBackend: raise ValueError( "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." ) - fam = detect_family(repo_id, family_override) + # Use the filename-fallback detector the route validated with, so a local + # .gguf pick whose family keyword lives only in the basename doesn't pass + # validation and then dead-end here on a no-GPU (native-routed) host. + fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError(f"Could not infer a diffusion family for '{repo_id}'.") if not family_sd_cpp_supported(fam): @@ -464,7 +469,7 @@ class SdCppDiffusionBackend: with self._lock: state = self._state if state is None: - raise RuntimeError("No diffusion model is loaded.") + raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) self._active_generate_cancel = cancel engine = self._resolve_engine() try: @@ -485,7 +490,7 @@ class SdCppDiffusionBackend: with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: for index in range(max(1, int(batch_size))): if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # Distinct seed per batch image (sd-cli is one image/run here), # so a batch is reproducible image-by-image from the base seed. # Mask to sd-cli's int64 range, NOT 53 bits: the request model and @@ -522,7 +527,7 @@ class SdCppDiffusionBackend: images.append(im.copy()) seeds.append(seed_i) if cancel.is_set(): - raise RuntimeError("Diffusion generation was cancelled.") + raise RuntimeError(DIFFUSION_CANCELLED_MSG) # ``seeds`` is the per-image seed (each sd-cli run used seed+index), so # the route can persist the real seed for every image in the batch. return { @@ -532,7 +537,7 @@ class SdCppDiffusionBackend: "repo_id": state.repo_id, } except SdCppCancelled as exc: - raise RuntimeError("Diffusion generation was cancelled.") from exc + raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc finally: self._gen = None with self._lock: diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 6714151045..557f6fcb3e 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -34,6 +34,7 @@ import time from pathlib import Path from typing import Callable, Optional +from utils.process_lifetime import child_popen_kwargs from core.inference.sd_cpp_args import ( SdCppGenParams, SdCppModelFiles, @@ -125,7 +126,8 @@ def find_sd_cpp_binary() -> Optional[str]: 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). + 3. the installer target: ``/../stable-diffusion.cpp`` when + that env (or ``STUDIO_HOME``) is set, else ``~/.unsloth/stable-diffusion.cpp``. 4. ``./stable-diffusion.cpp`` in-tree build (developer checkout). 5. ``sd-cli`` (then legacy ``sd``) on PATH. """ @@ -151,8 +153,12 @@ def find_sd_cpp_binary() -> Optional[str]: if hit: return hit - # 3. Default install root (sibling of ~/.unsloth/llama.cpp). - hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp")) + # 3. Default install root: the installer's default_install_dir() -- a sibling of + # the llama.cpp install under UNSLOTH_STUDIO_HOME / STUDIO_HOME when set, else + # ~/.unsloth. Mirror that env resolution or a custom Studio home never resolves. + 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 @@ -342,6 +348,10 @@ class SdCppEngine: # Own session/process group so cancellation/timeout can kill the whole # tree, not just the parent (POSIX only; harmless flag elsewhere). start_new_session = (os.name == "posix"), + # Bind the child to the parent's lifetime (Linux PR_SET_PDEATHSIG), so a + # hard parent crash mid-generation can't orphan sd-cli holding VRAM/RAM -- + # matching every llama.cpp Popen site. Composes with start_new_session. + **child_popen_kwargs(), ) # Drain stdout on a reader thread so the timeout is enforced even when the # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 400dc239fc..e20a2f1a49 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10370,6 +10370,10 @@ async def generate_diffusion_image( ): from core.inference import image_gallery from core.inference.diffusion_engine_router import get_active_diffusion_engine + from core.inference.diffusion_families import ( + DIFFUSION_CANCELLED_MSG, + DIFFUSION_NOT_LOADED_MSG, + ) backend = get_active_diffusion_engine() try: @@ -10385,11 +10389,15 @@ async def generate_diffusion_image( batch_size = request.batch_size, ) except RuntimeError as exc: - # Only "no model loaded" / cancelled are client-state (409). The native - # sd.cpp engine also raises RuntimeError for execution failures (nonzero - # exit, timeout, missing output), which are server errors (500). + # Only "no model loaded" / user-cancelled are client-state (409); both engines + # raise these two EXACT messages. The native sd.cpp engine also raises + # RuntimeError for execution failures (nonzero exit, timeout, missing output) + # whose text can embed the raw sd-cli tail (local paths / argv) -- those are + # server errors (500) returned as a fixed literal, never echoed. Match the + # sentinels exactly, not as a substring, so an sd-cli failure that merely + # contains "cancelled" can't misroute to 409 and leak that output. msg = str(exc) - if "No diffusion model is loaded" in msg or "cancelled" in msg.lower(): + if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG): raise HTTPException(status_code = 409, detail = msg) logger.error("diffusion.generate_failed: %s", exc) raise HTTPException(status_code = 500, detail = "Image generation failed.") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f8d0172173..5ec8bb3954 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -307,6 +307,38 @@ def test_generate_without_load_raises(fake_runtime): backend.generate(prompt = "x") +def test_failed_load_restores_backend_flags(fake_runtime, tmp_path, monkeypatch): + # A failure AFTER apply_speed_optims (here an OOM in apply_memory_plan) must go + # through the load's try/finally and restore the process-global TF32 / cudnn flags, + # so a later `off` load is still bit-identical, and must not commit a partial state. + # Regression: a refactor dropped this guard, leaking the flags on a failed load. + (tmp_path / "model.gguf").write_bytes(b"x") + backend = DiffusionBackend() + + restored: list = [] + cleared: list = [] + monkeypatch.setattr( + "core.inference.diffusion.restore_backend_flags", lambda snap: restored.append(snap) + ) + monkeypatch.setattr("core.inference.diffusion.clear_gpu_cache", lambda: cleared.append(True)) + monkeypatch.setattr( + "core.inference.diffusion.apply_memory_plan", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("CUDA out of memory")), + ) + + with pytest.raises(RuntimeError, match = "out of memory"): + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + family_override = "z-image", + base_repo = "base/repo", + speed_mode = "max", + ) + assert restored, "restore_backend_flags was not called on the failed-load path" + assert cleared, "clear_gpu_cache was not called on the failed-load path (VRAM leak)" + assert backend._state is None and backend.is_loaded is False + + def test_resolve_base_repo_prefers_caller_then_hf_tag_then_fallback(monkeypatch): from core.inference import diffusion from core.inference.diffusion_families import detect_family diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d0b6997d07..a3c78a21aa 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -281,6 +281,37 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch): assert "CUDA" not in resp.json()["detail"] +def test_generate_execution_error_with_cancelled_substring_is_sanitized_500(client, monkeypatch): + # A native sd-cli execution failure whose raw tail merely CONTAINS "cancelled" + # must stay a sanitized 500, not misroute to 409 and echo that output (path/arg + # leak). Regression: the handler matched "cancelled" as a substring. + backend = diffusion_module.get_diffusion_backend() + backend.loaded = True + + def _fail(**kwargs): + raise RuntimeError("sd-cli exited 1. Last output:\nop cancelled at /home/u/models/x.gguf") + + monkeypatch.setattr(backend, "generate", _fail) + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 500 + assert resp.json()["detail"] == "Image generation failed." + assert "cancelled" not in resp.json()["detail"] and "models" not in resp.json()["detail"] + + +def test_generate_user_cancellation_returns_409(client, monkeypatch): + # The exact cancellation sentinel both engines raise is client-state (409). + backend = diffusion_module.get_diffusion_backend() + backend.loaded = True + + def _cancel(**kwargs): + raise RuntimeError("Diffusion generation was cancelled.") + + monkeypatch.setattr(backend, "generate", _cancel) + resp = client.post("/api/inference/images/generate", json = {"prompt": "p"}) + assert resp.status_code == 409 + assert resp.json()["detail"] == "Diffusion generation was cancelled." + + def test_load_unknown_family_returns_400(client, monkeypatch): def _raise(*a, **k): raise ValueError("'x/y' isn't a supported image-generation model. Supported: Z-Image.") diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 79464dd6be..84e1b37dd2 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -219,6 +219,17 @@ def test_begin_load_requires_gguf_filename(): b.begin_load("unsloth/Z-Image-Turbo-GGUF") +def test_begin_load_resolves_family_from_filename_only(monkeypatch): + # A local .gguf pick whose family keyword lives only in the basename (parent dir + # carries none) must resolve via the same filename fallback the route validated + # with -- not dead-end with "Could not infer" on a native (no-GPU) host. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + monkeypatch.setattr(b, "_run_load", lambda **kwargs: None) # skip the download thread + b.begin_load("/models/gguf-store", gguf_filename = "Z-Image-Turbo-Q4_K_M.gguf") + # Validation passed (no ValueError) and the family was inferred from the filename. + assert b._loading is not None and b._loading.repo_id == "/models/gguf-store" + + def test_ensure_binary_returns_found(monkeypatch): monkeypatch.setattr(bk, "find_sd_cpp_binary", lambda: "/usr/bin/sd-cli") assert ensure_sd_cpp_binary() == "/usr/bin/sd-cli" From c34020bfec21557fdba87e77ca039f697a24376c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 20:02:34 +0000 Subject: [PATCH 13/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion_families.py | 4 +++- studio/backend/core/inference/diffusion_speed.py | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 3ac238e399..7257c69e9e 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -182,7 +182,9 @@ def detect_family(repo_id: str, override: Optional[str] = None) -> Optional[Diff def detect_family_for_pick( - repo_id: str, gguf_filename: Optional[str] = None, override: Optional[str] = None + repo_id: str, + gguf_filename: Optional[str] = None, + override: Optional[str] = None, ) -> Optional[DiffusionFamily]: """``detect_family``, falling back to the combined path/filename for a direct local ``.gguf`` pick. The frontend splits such a pick into (parent dir, basename), diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 5f1ef64445..fba962f3ba 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -253,6 +253,7 @@ def _enable_cudnn_benchmark(logger: Any) -> bool: def _enable_tf32(logger: Any) -> bool: try: import torch + torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True return True From 8d16ef977b77d3a5e16654ab07eeb7db58368dd3 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:01:12 -0300 Subject: [PATCH 14/16] Fix diffusion GGUF memory over-estimate and torch.compile crashes Three chained bugs that made Z-Image (and other GGUF DiTs) crash at generation on anything but a huge, fully-idle GPU. Verified end to end on an RTX 6000 Ada: Q2_K now plans resident and generates a real 1024x1024 PNG on both the resident and forced-group-offload paths. - Memory planner over-estimated the GGUF transformer's resident size. diffusers keeps GGUF weights PACKED (uint8 GGUFParameter) and dequantises per-matmul transiently, so resident VRAM is ~= the on-disk size, not the unpacked bf16 size (measured: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). The old per-quant expansion (x8 for Q2) over-estimated ~7.6x, so a 3.6 GB model on a 48 GB-free card was judged a "tight fit" and forced into group offload. Replace the multiplier table with estimate_gguf_resident_mib = storage * 1.05 (matches diffusers' own get_memory_footprint of a loaded GGUF model). - torch.compile with fullgraph=True crashed under CPU offload: group/model/ sequential offload installs a @torch.compiler.disable'd ModuleGroup.onload_ hook, which graph-breaks. Drop fullgraph when offloading is planned, same as the existing step-cache case (fullgraph = not (cache_active or offload_active)). This mirrors diffusers' documented compile+offload guidance. - compile_repeated_blocks compiles one graph per distinct block shape, but Z-Image's "repeated" blocks are heterogeneous (~11 variants), above dynamo's default recompile_limit of 8, so a resident load hard-errored under fullgraph. Raise the limit (diffusers' documented fix for regional-compile recompilation). Confirmed force_parameter_static_shapes=False is the wrong lever: same variant count, ~6x slower compile. Also drops the now-dead infer_gguf_quant_label / gguf_filename plumbing and adds regression tests for the estimate and the offload fullgraph drop. --- studio/backend/core/inference/diffusion.py | 17 +++--- .../core/inference/diffusion_memory.py | 53 +++++-------------- .../backend/core/inference/diffusion_speed.py | 44 ++++++++++++--- studio/backend/tests/test_diffusion_memory.py | 35 ++++-------- studio/backend/tests/test_diffusion_speed.py | 18 +++++++ 5 files changed, 85 insertions(+), 82 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 6d728df035..ecf1a1dcf9 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -39,10 +39,9 @@ from .diffusion_device import ( from .diffusion_memory import ( OFFLOAD_NONE, apply_memory_plan, - estimate_gguf_dense_mib, + estimate_gguf_resident_mib, estimate_image_runtime_mib, file_size_mib, - infer_gguf_quant_label, plan_diffusion_memory, snapshot_device_memory, ) @@ -522,7 +521,7 @@ class DiffusionBackend: # dense bf16 transformer must fit resident, so the fast path is offered only # when the plan is `none`. plan = self._plan_memory( - target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload + target, gguf_path, base, fam, memory_mode, cpu_offload ) # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it @@ -641,6 +640,9 @@ class DiffusionBackend: family = fam, speed_mode = effective_speed, cache_active = cache_engaged is not None, + # The planned offload policy: group/model/sequential offload installs + # compiler-disabled onload hooks, so compile must drop fullgraph. + offload_active = plan.offload_policy != OFFLOAD_NONE, logger = logger, ) if transformer_quant_engaged is not None and not speed_applied.get("compiled"): @@ -797,7 +799,6 @@ class DiffusionBackend: self, target: DiffusionDeviceTarget, gguf_path: str, - gguf_filename: Optional[str], base: str, fam: DiffusionFamily, memory_mode: Optional[str], @@ -808,16 +809,14 @@ class DiffusionBackend: 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) - transformer_dense = estimate_gguf_dense_mib( - file_size_mib(gguf_path), infer_gguf_quant_label(gguf_filename) - ) + transformer_resident = estimate_gguf_resident_mib(file_size_mib(gguf_path)) # 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_dense is not None: - model_dense_mib = transformer_dense + (companion_mib or 0) + if transformer_resident is not None: + model_dense_mib = transformer_resident + (companion_mib or 0) runtime_headroom = estimate_image_runtime_mib(width = None, height = None, family = fam.name) return plan_diffusion_memory( target = target, diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 9172ef7cf8..b2c1e25f11 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -218,51 +218,22 @@ def file_size_mib(path: Any) -> Optional[int]: return None -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 +def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]: + """Approximate the RESIDENT device size of a GGUF transformer loaded through + diffusers' ``GGUFQuantizationConfig``. - 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 + The weights stay PACKED on the device as quantised bytes (``GGUFParameter`` / + uint8); ``GGUFLinear.forward`` dequantises each weight to the bf16 compute dtype + transiently for its matmul and frees it immediately, so the persistent footprint + is ~= the on-disk tensor size, NOT the unpacked bf16 size. Measured on + Z-Image-Turbo: Q2_K 3.64 GiB -> 3.68 GiB, Q8_0 7.22 GiB -> 7.25 GiB resident. + The transient per-op dequant is covered by the separate runtime headroom. - -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.""" + (The prior per-quant expansion assumed a full unpack that never happens on this + path; it over-estimated e.g. Q2 ~7.6x, forcing needless offload.)""" 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 + return int(storage_mib * 1.05) # small margin for allocator + bf16 norms/biases def estimate_image_runtime_mib( diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index fba962f3ba..7c4a9018d8 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -148,11 +148,17 @@ def apply_speed_optims( family: Any, speed_mode: str = SPEED_OFF, cache_active: bool = False, + offload_active: bool = False, logger: Any = None, ) -> dict[str, bool]: """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, BEFORE placement / offload. Returns which optimisations actually engaged. Every - step is best-effort: a pipeline that doesn't support one is simply skipped.""" + step is best-effort: a pipeline that doesn't support one is simply skipped. + + ``offload_active`` is the planned offload policy != none: group/model/sequential + offloading installs ``@torch.compiler.disable``d onload hooks, so the compile must + drop ``fullgraph`` (same reason as an active step cache) or it crashes at the first + denoise step.""" applied = { "channels_last": False, "cudnn_benchmark": False, @@ -182,7 +188,11 @@ def apply_speed_optims( # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks( - pipe, logger, max_autotune = mode == SPEED_MAX, cache_active = cache_active + pipe, + logger, + max_autotune = mode == SPEED_MAX, + cache_active = cache_active, + offload_active = offload_active, ) if mode == SPEED_MAX: @@ -213,6 +223,7 @@ def _compile_repeated_blocks( *, max_autotune: bool = False, cache_active: bool = False, + offload_active: bool = False, ) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) @@ -225,14 +236,33 @@ def _compile_repeated_blocks( # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled # block because its static output buffer is overwritten across denoise steps. # - # fullgraph drops to False when a step cache is engaged: FBCache's per-step decision is - # ``@torch.compiler.disable``d, i.e. a graph break, which fullgraph=True rejects ("Skip - # inlining torch.compiler.disable()d function"). The break is cheap and the rest of the - # block still compiles. - kwargs: dict[str, Any] = {"fullgraph": not cache_active, "dynamic": not max_autotune} + # fullgraph drops to False when a step cache OR CPU offloading is engaged: both insert + # an ``@torch.compiler.disable``d function into the forward -- FBCache's per-step + # decision, and group/model/sequential offload's ``ModuleGroup.onload_`` streaming hook + # -- i.e. a graph break, which fullgraph=True rejects ("Skip inlining + # torch.compiler.disable()d function"). The break is cheap and the rest of the block + # still compiles. + kwargs: dict[str, Any] = { + "fullgraph": not (cache_active or offload_active), + "dynamic": not max_autotune, + } if max_autotune: kwargs["mode"] = "max-autotune-no-cudagraphs" try: + import torch + # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block + # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default + # recompile_limit of 8. Once the limit is hit a resident load hard-errors under + # fullgraph (and an offload/cache load silently drops the overflow blocks to eager), + # so raise it well past that (64) for headroom on larger heterogeneous DiTs. This is + # diffusers' own documented fix for regional-compile recompilation (their guide bumps + # cache_size_limit). Deliberately NOT force_parameter_static_shapes=False: it doesn't + # cut the variant count here and makes each compile ~6x slower (24s -> 143s cold). + dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) + if dynamo_cfg is not None: + for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver + if hasattr(dynamo_cfg, _limit_attr): + setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) fn(**kwargs) return True except Exception as exc: # noqa: BLE001 — optimisation only diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index bc09ffbd0f..cfdddb051e 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -27,9 +27,8 @@ from core.inference.diffusion_memory import ( DeviceMemory, MemoryPlan, apply_memory_plan, - estimate_gguf_dense_mib, + estimate_gguf_resident_mib, estimate_image_runtime_mib, - infer_gguf_quant_label, normalize_memory_mode, plan_diffusion_memory, snapshot_device_memory, @@ -70,29 +69,15 @@ def test_normalize_memory_mode_accepts_and_rejects(): # ── filename / size estimates ───────────────────────────────────────────────── -@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_gguf_resident_mib_matches_packed_size(): + # GGUF weights stay packed (uint8) on-device; diffusers dequantises per-matmul + # transiently, so the resident footprint ~= the on-disk size regardless of quant + # level (measured on Z-Image-Turbo: Q2_K 3.64->3.68 GiB, Q8_0 7.22->7.25 GiB). A + # small margin covers allocator overhead. The prior per-quant expansion over- + # estimated (Q2 ~7.6x) and forced needless offload on a roomy card. + assert estimate_gguf_resident_mib(1000) == 1050 + assert estimate_gguf_resident_mib(7220) == 7581 + assert estimate_gguf_resident_mib(None) is None def test_estimate_image_runtime_scales_with_pixels_and_family(): diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index b0ca10af4c..1f64b31b7a 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -216,6 +216,24 @@ def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): assert applied["tf32"] is False and applied["fused_qkv"] is False +def test_offload_active_drops_fullgraph(monkeypatch): + # Group/model/sequential offload installs a torch.compiler.disable'd onload hook; + # compiling with fullgraph=True then crashes at the first denoise step. Same reason + # as an active step cache -> fullgraph must drop to False when offload is planned. + _stub_torch(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, + _target(), + is_gguf = True, + family = _family(), + speed_mode = SPEED_DEFAULT, + offload_active = True, + ) + assert applied["compiled"] is True + assert pipe.compile_kwargs["fullgraph"] is False + + def test_speed_default_compiles_gguf(monkeypatch): _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) From a34d48d1fd498d73180797f399f0c3c678f9e903 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 21:05:52 +0000 Subject: [PATCH 15/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/diffusion.py | 4 +--- studio/backend/core/inference/diffusion_speed.py | 1 + 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index ecf1a1dcf9..894e3ce692 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -520,9 +520,7 @@ class DiffusionBackend: # the real budget) -- this also doubles as the dense-quant preflight: the # dense bf16 transformer must fit resident, so the fast path is offered only # when the plan is `none`. - plan = self._plan_memory( - target, gguf_path, base, fam, memory_mode, cpu_offload - ) + plan = self._plan_memory(target, gguf_path, base, fam, memory_mode, cpu_offload) # Opt-in fast path: load the DENSE bf16 transformer and torchao-quantise it # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 7c4a9018d8..99f506450a 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -250,6 +250,7 @@ def _compile_repeated_blocks( kwargs["mode"] = "max-autotune-no-cudagraphs" try: import torch + # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block # shape through compile_repeated_blocks; Z-Image needs ~11, above dynamo's default # recompile_limit of 8. Once the limit is hit a resident load hard-errors under From 68cad1bfbf890f09a241c42e9786c512bb324346 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 23:40:14 +0000 Subject: [PATCH 16/16] Remove stray async task scratch outputs committed by mistake --- .../async_task_output_0y2ic6.md | 34 ------------- .../async_task_output_24q0ob.md | 1 - .../async_task_output_2f8woi.md | 49 ------------------- .../async_task_output_3g25g3.md | 1 - .../async_task_output_3tedr8.md | 5 -- .../async_task_output_3uk9gw.md | 1 - .../async_task_output_44q27y.md | 8 --- .../async_task_output_49yim9.md | 1 - .../async_task_output_4fbmks.md | 1 - .../async_task_output_52h9ka.md | 2 - .../async_task_output_579ywh.md | 1 - .../async_task_output_5jwx8v.md | 8 --- .../async_task_output_67uiso.md | 1 - .../async_task_output_6emlbv.md | 1 - .../async_task_output_6l13hm.md | 1 - .../async_task_output_6mjfa2.md | 1 - .../async_task_output_6q6kut.md | 1 - .../async_task_output_7kn8en.md | 2 - .../async_task_output_83h6pm.md | 1 - .../async_task_output_8gk6gm.md | 2 - .../async_task_output_92i5k2.md | 1 - .../async_task_output_9dgu6d.md | 1 - .../async_task_output_9hresq.md | 1 - .../async_task_output_a478kh.md | 1 - .../async_task_output_asplkl.md | 22 --------- .../async_task_output_baeb1y.md | 5 -- .../async_task_output_cothbh.md | 1 - .../async_task_output_d1c766.md | 1 - .../async_task_output_d211gb.md | 4 -- .../async_task_output_d616s6.md | 2 - .../async_task_output_de1ekr.md | 1 - .../async_task_output_dhusjp.md | 1 - .../async_task_output_djnp8g.md | 1 - .../async_task_output_ed9jtn.md | 1 - .../async_task_output_f8uf97.md | 5 -- .../async_task_output_flqnxc.md | 1 - .../async_task_output_ftur3l.md | 1 - .../async_task_output_g7yinb.md | 6 --- .../async_task_output_ggz4m7.md | 1 - .../async_task_output_gs83o9.md | 1 - .../async_task_output_hbajbw.md | 1 - .../async_task_output_hh8yan.md | 1 - .../async_task_output_hm7vhm.md | 1 - .../async_task_output_hu6k45.md | 1 - .../async_task_output_irlty4.md | 1 - .../async_task_output_iyj2u3.md | 1 - .../async_task_output_j302lq.md | 1 - .../async_task_output_j6ykxj.md | 1 - .../async_task_output_jm0l3y.md | 1 - .../async_task_output_klpaw5.md | 1 - .../async_task_output_ktq93m.md | 1 - .../async_task_output_l026u1.md | 1 - .../async_task_output_l4hdlx.md | 1 - .../async_task_output_ldy0ne.md | 1 - .../async_task_output_m9h6x3.md | 2 - .../async_task_output_mwjqhd.md | 1 - .../async_task_output_na1kau.md | 1 - .../async_task_output_nagyj1.md | 1 - .../async_task_output_nc9gvk.md | 1 - .../async_task_output_nwjxf4.md | 1 - .../async_task_output_o5wtl3.md | 1 - .../async_task_output_o77ymr.md | 1 - .../async_task_output_p8crwl.md | 1 - .../async_task_output_p8rq7l.md | 1 - .../async_task_output_peiz71.md | 1 - .../async_task_output_pnabd9.md | 1 - .../async_task_output_ptrx41.md | 1 - .../async_task_output_pvflyj.md | 1 - .../async_task_output_pysgn3.md | 4 -- .../async_task_output_pzxpq2.md | 2 - .../async_task_output_q479i8.md | 1 - .../async_task_output_qz0t3d.md | 1 - .../async_task_output_rbj2zk.md | 7 --- .../async_task_output_s9q7qs.md | 1 - .../async_task_output_t0my78.md | 8 --- .../async_task_output_v3swmu.md | 1 - .../async_task_output_v6m49g.md | 1 - .../async_task_output_vbfgbz.md | 1 - .../async_task_output_viqlxl.md | 1 - .../async_task_output_vn6a5d.md | 4 -- .../async_task_output_vudd2x.md | 2 - .../async_task_output_w0h2kg.md | 9 ---- .../async_task_output_whtxy2.md | 1 - .../async_task_output_xcofo0.md | 2 - .../async_task_output_xds3jp.md | 1 - .../async_task_output_y2psj5.md | 1 - .../async_task_output_yz7170.md | 1 - .../async_task_output_zi1igk.md | 1 - .../async_task_output_zsrnxt.md | 1 - .../async_task_output_1csyea.md | 1 - .../async_task_output_560grk.md | 1 - .../async_task_output_6vw4lz.md | 1 - .../async_task_output_7p2hgc.md | 1 - .../async_task_output_bnb814.md | 1 - .../async_task_output_cymhgv.md | 3 -- .../async_task_output_d8ngtg.md | 1 - .../async_task_output_edpp2u.md | 6 --- .../async_task_output_g3wn5w.md | 1 - .../async_task_output_heqa1v.md | 1 - .../async_task_output_iedlnv.md | 24 --------- .../async_task_output_jlf3lc.md | 1 - .../async_task_output_jqe936.md | 1 - .../async_task_output_kc50ea.md | 1 - .../async_task_output_oh38q4.md | 1 - .../async_task_output_ot9we3.md | 1 - .../async_task_output_phzu13.md | 11 ----- .../async_task_output_qh9doa.md | 2 - .../async_task_output_s15ff5.md | 1 - 108 files changed, 320 deletions(-) delete mode 100644 studio/backend/async_task_outputs/async_task_output_0y2ic6.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_24q0ob.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_2f8woi.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_3g25g3.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_3tedr8.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_3uk9gw.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_44q27y.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_49yim9.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_4fbmks.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_52h9ka.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_579ywh.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_5jwx8v.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_67uiso.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_6emlbv.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_6l13hm.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_6mjfa2.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_6q6kut.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_7kn8en.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_83h6pm.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_8gk6gm.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_92i5k2.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_9dgu6d.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_9hresq.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_a478kh.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_asplkl.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_baeb1y.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_cothbh.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_d1c766.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_d211gb.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_d616s6.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_de1ekr.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_dhusjp.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_djnp8g.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ed9jtn.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_f8uf97.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_flqnxc.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ftur3l.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_g7yinb.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ggz4m7.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_gs83o9.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_hbajbw.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_hh8yan.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_hm7vhm.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_hu6k45.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_irlty4.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_iyj2u3.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_j302lq.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_j6ykxj.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_jm0l3y.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_klpaw5.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ktq93m.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_l026u1.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_l4hdlx.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ldy0ne.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_m9h6x3.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_mwjqhd.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_na1kau.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_nagyj1.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_nc9gvk.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_nwjxf4.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_o5wtl3.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_o77ymr.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_p8crwl.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_p8rq7l.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_peiz71.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_pnabd9.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_ptrx41.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_pvflyj.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_pysgn3.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_pzxpq2.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_q479i8.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_qz0t3d.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_rbj2zk.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_s9q7qs.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_t0my78.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_v3swmu.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_v6m49g.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_vbfgbz.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_viqlxl.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_vn6a5d.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_vudd2x.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_w0h2kg.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_whtxy2.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_xcofo0.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_xds3jp.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_y2psj5.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_yz7170.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_zi1igk.md delete mode 100644 studio/backend/async_task_outputs/async_task_output_zsrnxt.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_1csyea.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_560grk.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_6vw4lz.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_7p2hgc.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_bnb814.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_cymhgv.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_d8ngtg.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_edpp2u.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_g3wn5w.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_heqa1v.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_iedlnv.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_jlf3lc.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_jqe936.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_kc50ea.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_oh38q4.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_ot9we3.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_phzu13.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_qh9doa.md delete mode 100644 studio/frontend/async_task_outputs/async_task_output_s15ff5.md diff --git a/studio/backend/async_task_outputs/async_task_output_0y2ic6.md b/studio/backend/async_task_outputs/async_task_output_0y2ic6.md deleted file mode 100644 index f1c7bc93ec..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_0y2ic6.md +++ /dev/null @@ -1,34 +0,0 @@ -- Decision: Used plan mode because `workflows/benchmarking_workflow.md` explicitly requires planning before benchmark execution. -- Decision: User selected full scope: `all 8 variants` and both `GGUF` and dense `safetensors` fp8/int8 paths. -- Decision: Chose a focused-comprehensive benchmark matrix of `44 configs` instead of an exhaustive `549-630-config` Cartesian product because the latter would require excessive downloads/GPU time and produce an unreadable table. -- Decision: Excluded `FLUX.2-klein-9B` from actual sweep after `snapshot_download` returned `GatedRepoError: 403 Client Error`; needs manual HF license acceptance. -- Decision: Treated failures as benchmark findings, not harness bugs: dense int8 small-M failures, `mxfp8` unsupported, Flux T5 fp8 text-encoder failure, and FBCache GGUF slowdown. -- Created task `#66`: `Benchmark all supported diffusion models across all optimization levers`. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/diffusion_bench.py`: added/plumbed benchmark flags including `--transformer-cache {off,fbcache}` and related metrics/config output. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/benchall_orchestrator.py`: encodes the model/config matrix, runs `diffusion_bench.py` across GPUs, supports `--list`, `--gpu`, `--variants`, and `--aggregate`. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/benchall_predownload.py` via shell heredoc during predownload setup. -- Overwrote `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/wobbly-jumping-narwhal.md` with approved benchmark plan. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/diffusion_benchall/SUMMARY.md` with benchmark findings and headline tables. -- Generated `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/diffusion_benchall/results.csv` from aggregation; contains `40 rows`. -- Commands run: `ls workflows/benchmarking_workflow.md unsloth/workflows/benchmarking_workflow.md`, `ls parallel_planner.py gpu_pool_runner.sh`, `ls unsloth/scripts/diffusion_bench.py unsloth/scripts`; found workflow, planner, runner, `diffusion_bench.py`, `diffusion_quality.py`. -- Command run: `python3 -u parallel_planner.py --output_file extra_plan.md --prompt ...`; background task completed and wrote `/mnt/disks/unslothai/ubuntu/workspace_81/async_task_outputs/extra_plan.md`. -- Command run: `python -c "import ast; ast.parse(open('scripts/diffusion_bench.py').read()); print('bench OK')"`; exit `0`, output `bench OK`. -- Command run: `python scripts/diffusion_bench.py --help ... grep ...`; verified `--transformer-cache {off,fbcache}` appears. -- Command run: HF access resolver using `huggingface_hub.HfApi`; found all 8 GGUF `Q4_K_M` filenames and initially `list_repo_files` access appeared OK. -- Resolved GGUF files include `z-image-turbo-Q4_K_M.gguf`, `z-image-Q4_K_M.gguf`, `qwen-image-2512-Q4_K_M.gguf`, `qwen-image-Q4_K_M.gguf`, `flux1-schnell-Q4_K_M.gguf`. -- Command run: `python scripts/benchall_orchestrator.py --list`; output showed `44 configs`, including `z-image-turbo (5)`, `z-image (5)`, `qwen-image-2512 (5)`, `qwen-image (6)`, `flux.1-schnell (4)`, `flux.1-dev (11)`, plus klein configs. -- Command run: dry-run `CUDA_VISIBLE_DEVICES=5 python -u scripts/benchall_orchestrator.py ...`; exit `0`; all `z-image-turbo` configs passed: `G0_eager`, `G1_default`, `D1_fp8`, `D2_int8`, `G1_max`. -- Command run: predownload all assets; exit `0` overall, but logged `[base ERR] black-forest-labs/FLUX.2-klein-9B: GatedRepoError: 403 Client Error`; other 7 models downloaded. -- Command run: metrics schema check on `outputs/diffusion_benchall/_dryrun/z-image-turbo/G1_default/metrics.json`; verified keys `env`, `load`, `generate`, `config`, `accuracy`. -- Command run: launched full sweep across GPUs `4`, `5`, `6`, `7` using `scripts/benchall_orchestrator.py`; background task completed exit `0`. -- Command run: first progress check showed failures for `flux.1-dev/D2_int8`, `qwen-image/D2_int8`, `qwen-image-2512/D2_int8`; later aggregation found `33` completed metrics and `7` failed configs. -- Failed configs: `flux.1-dev/D2_int8`, `flux.1-dev/D_mxfp8`, `flux.1-dev/G1_teq_fp8`, `qwen-image/D2_int8`, `flux.1-schnell/D2_int8`, `qwen-image-2512/D2_int8`, `z-image/G2_fbcache`. -- Error root cause: dense int8 on Flux/Qwen hits `torch._int_mm: self.size(0) needs to be greater than 16, but got 1`; small-M conditioning embedders should be excluded from int8. -- Error root cause: `mxfp8` failed with `CUBLAS_STATUS_NOT_SUPPORTED` on B200/cuBLAS stack. -- Error root cause: Flux text encoder fp8 failed with `normal_kernel_cuda not implemented for Float8_e4m3fn`. -- Finding: FBCache on GGUF is slower despite compile engaging; log showed `diffusion.cache: fbcache engaged (threshold=0.08)` and `speed_optims` included `compiled`; likely due to `fullgraph=False` plus GGUF per-op dequant graph breaks. -- Final measured headline: GGUF compile+cuDNN stack gives `1.9x-3.55x`; dense fp8 gives `2.3x-5.0x`; GGUF lowest VRAM, dense fp8 fastest. -- Example final table values: `Z-Image-Turbo` best `3.58x`; `Z-Image` best `3.87x`; `Qwen-Image-2512` best `4.89x`; `Qwen-Image` best `4.97x`; `FLUX.1-schnell` best `3.32x`; `FLUX.1-dev` best `4.47x`; `FLUX.2-klein-4B` best `2.31x`; `FLUX.2-klein-9B` blocked. -- Task `#66` marked completed. -- COMPLETED: plan approved, benchmark harness patched, orchestrator created, dry run passed, assets predownloaded where accessible, 7-model full sweep completed, results aggregated, summary written. -- PENDING: optional follow-up fix recommended: gate FBCache to dense-only/off for GGUF; optional fix for dense int8 should exclude small-M embedders; accept HF license for `black-forest-labs/FLUX.2-klein-9B` before rerunning that model. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_24q0ob.md b/studio/backend/async_task_outputs/async_task_output_24q0ob.md deleted file mode 100644 index 50a3c01fd3..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_24q0ob.md +++ /dev/null @@ -1 +0,0 @@ -- PENDING: fp8/int8 UI; auto-resize 16 \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_2f8woi.md b/studio/backend/async_task_outputs/async_task_output_2f8woi.md deleted file mode 100644 index 1ae52c73ea..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_2f8woi.md +++ /dev/null @@ -1,49 +0,0 @@ -- Decision: Launched 6 parallel research agents to avoid duplicated work and cover shipped optimizations plus five lever clusters: caching, attention/token merging, peak-memory reduction, `torch.compile`/Inductor, and quant/kernel/cross-platform. -- Decision: Prioritized attention-backend selection after research and B200 probing showed it was the only immediately validated speed win: `_native_cudnn` gave `1.18x` end-to-end with `LPIPS=0.004448794759809971`; Sage/Flash/FBCache were not usable on the available Z-Image setup. -- Decision: Split Inductor flags out of the main implementation because measured gain was negligible on B200/bf16: `inductor_flags 0.681s (1.01x vs base) peak=40.9G LPIPS=0.0047416952438652515`. -- Decision: Rejected `flash_4_hub` for now because installing `kernels` upgraded `huggingface-hub` to `1.21.0`, breaking diffusers `0.38.0` which needs `<1.0`; env was restored to `huggingface-hub==0.36.2`. -- Decision: Rejected SageAttention for B200 because `sageattention` v1.0.6 lacked usable Blackwell/`sm_100` kernels: `Sage Attention backend 'sage' is not usable`. -- Decision: Implemented Phase 11 consumer-aware quant ladder because consumer/workstation GPUs have nerfed FP8 FP32-accumulate while INT8 is full-rate; kept FP8 first on data-center GPUs. -- Decision: Rejected VAE tiling as Phase 12 for Z-Image inference peak memory after measuring worse/negligible peaks; denoise transformer activations dominate, not VAE decode. -- Decision: Deferred FBCache shipping because it fails for Z-Image even eager with `ValueError: Parameter 'hidden_states' not found in function signature but was requested.` and `KeyError: 'hidden_states'`; it needs validation on many-step Flux/Qwen. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/quant_research/03_NEXT_LEVERS.md`: synthesis of next levers and later updated with negative VAE/FBCache findings. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/perf_levers_probe.py`: B200 probe for baseline, Inductor flags, attention backends, and FBCache. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_attention.py`: normalizes attention backend requests, selects `auto -> _native_cudnn` on NVIDIA CUDA, supports opt-in `sage`/`flash`/etc., applies via `set_attention_backend`, gracefully falls back to native. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion.py`: threaded `attention_backend` through load flow, applied attention backend before speed compile, added `_LoadState.attention_backend`, and exposed it in status. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/models/inference.py`: added request and status model fields for `attention_backend`. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/routes/inference.py`: forwards `attention_backend` from API request to backend load call. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_attention.py`: hermetic tests for backend normalization/selection/application/fallback. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_routes.py`: added route coverage for threading and invalid enum handling. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_transformer_quant.py`: added consumer-aware `auto` reorder so consumer GPUs prefer `int8`; B200/data-center keeps `fp8` first. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_transformer_quant.py`: updated torch stub with device names, added 3 consumer ladder tests, fixed data-center names. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/temp/phase10_pr_body.md` and `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/temp/phase11_pr_body.md` for PR bodies. -- Command: checked diffusers APIs; key output: `diffusers 0.38.0 torch 2.9.1+cu128`, `ZImageTransformer2DModel has set_attention_backend: True`, cache configs present, Inductor flags present. -- Command: listed attention backends; key output included `flash`, `flash_4_hub`, `aiter`, `flex`, `native`, `_native_cudnn`, `sage`, `sage_hub`, `sage_varlen`, `xformers`. -- Command: `uv pip install sageattention kernels`; caused `huggingface-hub` incompatibility. -- Command: restored env with `uv pip uninstall kernels` and `uv pip install "huggingface-hub>=0.34.0,<1.0"`; verified `hub 0.36.2 diffusers 0.38.0`, `ZImage import OK`, `sageattention OK`. -- Command: ran B200 perf probe; final key output: `baseline 0.686s peak=22.6G`, `attn_cudnn 0.584s (1.18x vs base) peak=40.6G LPIPS=0.004448794759809971`. -- Command: created branch `diffusion-phase10-attention`; verified `has set_attention_backend: True`, `has reset_attention_backend: True`, `hip: None`. -- Error resolved: initial tests failed because test expected `"flash-3"` alias; actual normalization did not map it. Fixed test to use valid case. -- Command: `python -m pytest tests/ -q -k "diffusion"` after Phase 10; result `213 passed, 4993 deselected, 1 warning`. -- Command: B200 smoke for product attention functions; key output: `select(auto, speed_active=True) -> _native_cudnn`, `apply -> _native_cudnn`, `image finite: True shape: (1024, 1024, 3)`, `ATTN-SMOKE-OK`, `=== EXIT 0 ===`. -- Commit: `b92367554 Studio diffusion (Phase 10): attention-backend selection`. -- Command: pushed `diffusion-phase10-attention` and created PR `https://github.com/unslothai/unsloth/pull/6701` stacked on `diffusion-phase9-prequant`. -- Command: pre-commit.ci moved Phase 10 remote; rebased successfully; sync output `0 0`. -- Command: re-ran tests for PR #6701; result `213 passed, 4993 deselected, 1 warning in 15.14s`; PR state `{"base":"diffusion-phase9-prequant","mergeable":"MERGEABLE","number":6701,"state":"OPEN"}`. -- Command: created branch `diffusion-phase11-consumer-int8`. -- Command: Phase 11 tests: `python -m pytest tests/ -q -k "diffusion"`; result `216 passed` (tail showed `216 p`). -- Command: B200 non-regression for quant ladder; key output: `device: NVIDIA B200`, `_is_consumer_gpu('cuda'): False`, `reorder (B200): ('fp8', 'nvfp4', 'mxfp8', 'int8')`, `auto scheme on B200: fp8`. -- Commit: `764a3e1dc Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder`. -- Command: pushed `diffusion-phase11-consumer-int8` and created PR `https://github.com/unslothai/unsloth/pull/6702` stacked on `diffusion-phase10-attention`. -- Command: Phase 11 pre-commit wait showed `no movement`, sync `0 0`. -- Command: PR state check: #6701 `MERGEABLE`/`OPEN`, #6702 `MERGEABLE`/`OPEN`. -- Command: measured VAE decode peak with pipeline-level tiling; errors: `'ZImagePipeline' object has no attribute 'enable_vae_tiling'` and `'ZImagePipeline' object has no attribute 'enable_vae_slicing'`. -- Command: checked Z-Image VAE methods; key output: `VAE class: AutoencoderKL`, `vae.enable_tiling: True`, `vae.enable_slicing: True`, `pipe.enable_vae_tiling: False`, `pipe.vae_scale_factor: 8`. -- Command: measured VAE-level tiling; key output: `res 1024: no-tile peak=23.3G | vae-tile peak=30.6G | saved=-7.3G (-31%) | LPIPS=0.0000`; `res 1536: no-tile peak=33.8G | vae-tile peak=37.9G | saved=-4.1G (-12%) | LPIPS=0.0037`; `res 2048: no-tile peak=45.6G | vae-tile peak=45.3G | saved=0.3G (1%) | LPIPS=0.0037`; `=== EXIT 0 ===`. -- Command: probed FBCache compatibility; key output: `[cache_eager] FAIL ValueError: Parameter 'hidden_states' not found in function signature but was requested.`, `[cache+regional_nofullgraph] FAIL KeyError: 'hidden_states'`, `[cache+fullcompile_nofullgraph] FAIL KeyError: 'hidden_states'`, `FBCACHE-COMPAT-DONE`, `=== EXIT 0 ===`. -- Command still in flight at end: prefetch Flux.1-dev for FBCache validation, background ID `bbujpwspc`, log path pattern `logs/flux_prefetch_$(date +%Y%m%d_%H%M%S).log`. -- Completed: Phase 10 PR #6701 is open/mergeable, tests passing, B200 smoke passing. -- Completed: Phase 11 PR #6702 is open/mergeable, tests passing, B200 non-regression verified. -- Completed: VAE tiling investigated and rejected for Z-Image inference peak. -- Completed: FBCache investigated on Z-Image and found incompatible due to `hidden_states` signature. -- Pending: Wait for Flux.1-dev prefetch, then validate FBCache on compatible many-step model and potentially ship gated Phase 12 for Flux/Qwen only. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_3g25g3.md b/studio/backend/async_task_outputs/async_task_output_3g25g3.md deleted file mode 100644 index 6a2d828294..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_3g25g3.md +++ /dev/null @@ -1 +0,0 @@ -- #156 done: FLUX.2 reference; 99 pass \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_3tedr8.md b/studio/backend/async_task_outputs/async_task_output_3tedr8.md deleted file mode 100644 index 520190e628..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_3tedr8.md +++ /dev/null @@ -1,5 +0,0 @@ -- User directed to expand scope to every architecture: `Yes do all arches.` -- User also requested validation benchmarks across all architectures: `Then also do a benchmark to check all.` -- User specifically requested benchmarking the alternative buffer-based implementation: `Also benchmark the buffers approach` -- COMPLETED: captured new requirements only. -- PENDING: implement/test all architectures; run comparative benchmarks including buffer approach. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_3uk9gw.md b/studio/backend/async_task_outputs/async_task_output_3uk9gw.md deleted file mode 100644 index 0569aa5f6e..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_3uk9gw.md +++ /dev/null @@ -1 +0,0 @@ -- `upscale` added+verified; hold for push/mirror \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_44q27y.md b/studio/backend/async_task_outputs/async_task_output_44q27y.md deleted file mode 100644 index bd48a2c47d..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_44q27y.md +++ /dev/null @@ -1,8 +0,0 @@ -- Decision: implement faithful `addcmul` fusions for all four Flux block classes (`FluxTransformerBlock`, `FluxSingleTransformerBlock`, `Flux2TransformerBlock`, `Flux2SingleTransformerBlock`) using exact `diffusers` 0.38 forward bodies to avoid semantic drift. -- Decision: run two follow-up tasks: `Task #112` for Flux.1/Flux.2 patch implementation with body-drift guards/specs/allclose tests, and `Task #113` for isolation benchmarks across all four arch families plus focused global buffer on/off benchmark. -- Commands run from `/mnt/disks/unslothai/ubuntu/workspace_81`: `CUDA_VISIBLE_DEVICES="" UNSLOTH_ALLOW_CPU=1 python - <<'PY' 2>/dev/null` to inspect `transformer_flux2` and `transformer_flux` sources. -- Key outputs: fetched `Flux2TransformerBlock.forward`, `Flux2SingleTransformerBlock.forward`, `FluxTransformerBlock.forward`, `FluxSingleTransformerBlock.forward`, relevant `__init__` signatures, and `Flux2Modulation` source. -- Finding: `Flux2Modulation.split` returns modulation tensors shaped `[B,1,dim]`, so test inputs can rely on clean broadcasting. -- Error: `API Error: Connection closed mid-response. The response above may be incomplete.` Work was interrupted before edits. -- Completed: source inspection and task creation. -- Pending: edit `diffusion_arch_patches.py`, add specs/guards/tests, run allclose checks and benchmarks. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_49yim9.md b/studio/backend/async_task_outputs/async_task_output_49yim9.md deleted file mode 100644 index 106f8313a8..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_49yim9.md +++ /dev/null @@ -1 +0,0 @@ -- Backend PR split blocked; plan updated. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_4fbmks.md b/studio/backend/async_task_outputs/async_task_output_4fbmks.md deleted file mode 100644 index 97ed70e280..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_4fbmks.md +++ /dev/null @@ -1 +0,0 @@ -- Done: default=compile-dequant; max=compile \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_52h9ka.md b/studio/backend/async_task_outputs/async_task_output_52h9ka.md deleted file mode 100644 index ef17b67b80..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_52h9ka.md +++ /dev/null @@ -1,2 +0,0 @@ -- Done: right-sidebar move #25, Advanced discoverability, `/16` auto-resize verified. -- Pending: push, backend path, mirror, cancel loop. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_579ywh.md b/studio/backend/async_task_outputs/async_task_output_579ywh.md deleted file mode 100644 index 2bf458f678..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_579ywh.md +++ /dev/null @@ -1 +0,0 @@ -Continue the user’s directive to make a PR or push, scoped as “Complete PR, new branch” per the answered scoping question. The last completed tool was the full diffusion pre-commit pytest run, logging to `/mnt/disks/unslothai/ubuntu/workspace_81/logs/pretest_230339.log`; do not rerun it. Next, inspect the test result/log as needed, then proceed with the planned commit/push flow: push the Phase-16 base branch to `oobabooga`, create the new complete PR branch, commit only the intended frontend/installer/backend changes while excluding logs/temp/plans/junk, push it, and open the PR against that base. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_5jwx8v.md b/studio/backend/async_task_outputs/async_task_output_5jwx8v.md deleted file mode 100644 index 121363abbf..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_5jwx8v.md +++ /dev/null @@ -1,8 +0,0 @@ -- User requested benchmarking performance gains across all PRs for every supported model format: `GGUFs` and `safetensors`. -- Verbatim directive to preserve: `And for all odels that are supported - GGUFs or just safetensors - across all PRs - what are the performance gains - benchmark all` -- Decision implied: scope expands from isolated PR review to cross-PR, cross-model benchmark comparison. -- Files created/edited: none in this span. -- Commands run: none in this span. -- Errors encountered: none; note typo `odels` appears in user directive. -- COMPLETED: request captured. -- PENDING: identify supported models/formats, enumerate relevant PRs, run benchmarks, report performance gains. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_67uiso.md b/studio/backend/async_task_outputs/async_task_output_67uiso.md deleted file mode 100644 index 67c816ec74..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_67uiso.md +++ /dev/null @@ -1 +0,0 @@ -- Done: PR synced; no pending \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_6emlbv.md b/studio/backend/async_task_outputs/async_task_output_6emlbv.md deleted file mode 100644 index 79bd452c65..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_6emlbv.md +++ /dev/null @@ -1 +0,0 @@ -- PR `#6694` opened; tests pass. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_6l13hm.md b/studio/backend/async_task_outputs/async_task_output_6l13hm.md deleted file mode 100644 index ad0ae35beb..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_6l13hm.md +++ /dev/null @@ -1 +0,0 @@ -- Asked check `https://github.com/city96/ComfyUI-GGUF` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_6mjfa2.md b/studio/backend/async_task_outputs/async_task_output_6mjfa2.md deleted file mode 100644 index 0f582f54af..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_6mjfa2.md +++ /dev/null @@ -1 +0,0 @@ -- Faster inference; `5 Opus subagents` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_6q6kut.md b/studio/backend/async_task_outputs/async_task_output_6q6kut.md deleted file mode 100644 index 7ecc105658..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_6q6kut.md +++ /dev/null @@ -1 +0,0 @@ -- Asked: `non GGUFs`, `fp8, int8, bf16` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_7kn8en.md b/studio/backend/async_task_outputs/async_task_output_7kn8en.md deleted file mode 100644 index b111aa9728..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_7kn8en.md +++ /dev/null @@ -1,2 +0,0 @@ -- User asked: `Also for https://github.com/unslothai/unsloth/pull/6658 - which models are supported?` -- PENDING: inspect PR `6658` and determine supported models. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_83h6pm.md b/studio/backend/async_task_outputs/async_task_output_83h6pm.md deleted file mode 100644 index 8c95fcb0a4..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_83h6pm.md +++ /dev/null @@ -1 +0,0 @@ -- `FLUX.2-klein` inpaint done; extend excluded \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_8gk6gm.md b/studio/backend/async_task_outputs/async_task_output_8gk6gm.md deleted file mode 100644 index fcd244c5d2..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_8gk6gm.md +++ /dev/null @@ -1,2 +0,0 @@ -- Done: HF data -- Pending: PR push \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_92i5k2.md b/studio/backend/async_task_outputs/async_task_output_92i5k2.md deleted file mode 100644 index 83bf8c7262..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_92i5k2.md +++ /dev/null @@ -1 +0,0 @@ -- Phase12 done: PR `#6703` diff --git a/studio/backend/async_task_outputs/async_task_output_9dgu6d.md b/studio/backend/async_task_outputs/async_task_output_9dgu6d.md deleted file mode 100644 index 11421b08c7..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_9dgu6d.md +++ /dev/null @@ -1 +0,0 @@ -- Rebased, pushed PR `#6675`; 127 pass \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_9hresq.md b/studio/backend/async_task_outputs/async_task_output_9hresq.md deleted file mode 100644 index f8c3551d05..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_9hresq.md +++ /dev/null @@ -1 +0,0 @@ -- quant built; 146 pass; GPU pending \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_a478kh.md b/studio/backend/async_task_outputs/async_task_output_a478kh.md deleted file mode 100644 index 537ced780b..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_a478kh.md +++ /dev/null @@ -1 +0,0 @@ -- No actions; context-summary request only \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_asplkl.md b/studio/backend/async_task_outputs/async_task_output_asplkl.md deleted file mode 100644 index 71ca6d679f..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_asplkl.md +++ /dev/null @@ -1,22 +0,0 @@ -- Decision: implemented all remaining diffusion arch addcmul patches because tests/benchmarks required all four families covered and bit-correct: `qwen-image`, `z-image`, `flux.1`, `flux.2-klein`. -- Decision: kept source/signature drift guards around arch patches because patch safety depends on matching the expected diffusers forward bodies. -- Decision: treated initial `flux.1-schnell` benchmark `-20.5%` as invalid because logs showed `offload=model` and latency was ~10x normal; reran contaminated cases sequentially on free GPU 5 resident. -- Decision: concluded global weight buffer is neutral because qwen was `+0.2%`, flux.1 repeated from `+11.0%` to `-2.4%`, and allocator profile had no allocation retries/steady reserved VRAM. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_arch_patches.py`: added all family/block patches and updated module docstring to state all arches are implemented. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_arch_patches.py`: added allclose tests for four flux blocks and updated install-count assertion. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/buffer_bench.py`: dedicated global-buffer on/off benchmark. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/outputs/arch_patch/SUMMARY.md`: final benchmark summary and interpretation. -- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/test_diffusion_arch_patches.py -q 2>&1 | tail -16`; result: `.......... [100%]`, 10 arch tests passed with warning. -- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/ -q -k "diffusion" 2>&1 | tail -4`; result: `300 passed, 1 skipped, 5012 deselected, 1 warning in 13.25s`. -- Launched benchmarks: `scripts/arch_patch_bench.py --model flux.1-schnell --gpu 4`, `--model flux.2-klein-4b --gpu 5`, `scripts/buffer_bench.py --model qwen-image --gpu 6`, `--model flux.1-schnell --gpu 7`. -- Error: result collection shell test produced `/bin/bash: line 23: [: 0 0: integer expression expected`; non-blocking, later results were collected another way. -- Error: GPUs 4/6/7 were occupied by external `python3` processes using ~158GB each; caused offload contamination. Did not kill external PIDs. -- Ran `kill 325007` to stop own offloaded qwen buffer bench; output `killed my qwen buffer bench (offloaded)`. -- Ran clean sequential re-runs on GPU 5: `scripts/arch_patch_bench.py --model flux.1-schnell --gpu 5`, `--model qwen-image --gpu 5`, `scripts/buffer_bench.py --model flux.1-schnell --gpu 5`, `--model qwen-image --gpu 5`; all 8 loads reported `offload=none`. -- Final resident arch benchmark results: `z-image-turbo off=1.554s on=1.557s -0.2% PSNR=32.1`; `flux.2-klein-4b off=0.626s on=0.623s +0.4% PSNR=33.6`; `flux.1-schnell off=1.417s on=1.405s +0.8% PSNR=30.2`; `qwen-image off=9.508s on=9.667s -1.6% PSNR=45.2`. -- Final resident buffer benchmark results: `qwen-image off=9.477s on=9.462s +0.2%`; `flux.1-schnell off=1.569/on=1.413 +11.0%` then repeat `off=1.415s on=1.450s -2.4% VRAM 18.8->18.9 PSNR=inf`. -- Error: repeat flux.1 buffer pretty-print wrapper had quote-nesting bug and printed inline Python source; underlying `buffer_bench` still wrote `outputs/buffer_bench/flux.1-schnell/result.json`, so resolved by reading result JSON. -- Ran final verification: `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && CUDA_VISIBLE_DEVICES="" python -m pytest tests/ -q -k "diffusion" 2>&1 | tail -3`; result: `300 passed, 1 skipped, 5012 deselected, 1 warning in 14.20s`. -- Task status: completed task `#112` and `#113`. -- COMPLETED: all arch patches implemented, flux allclose tests added, install-count assertion updated, buffer benchmark script added, clean benchmarks collected, summary written, final diffusion test suite green. -- PENDING: none noted. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_baeb1y.md b/studio/backend/async_task_outputs/async_task_output_baeb1y.md deleted file mode 100644 index f182fd9bb5..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_baeb1y.md +++ /dev/null @@ -1,5 +0,0 @@ -- Decision: continue active Stop-hook work on diffusion efficiency/memory without pausing, per directive `Keep working on making diffusion more efficent, use less memory and not reduce accuracy...`. -- Must use cwd only: `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/temp/temporary_qui4gujp`; never exit it. -- Goal constraints: optimize diffusion load/inference peak memory and speed across Mac/Windows/Linux, NVIDIA/AMD/Intel/CPU; support GGUF and safetensors; avoid large accuracy loss, `25% reduction is fine`. -- Research targets preserved: `torch._int_mm`, `torchao/sparsity`, SDNQ, int8 fused, FBGEMM, TorchInductor config, Diffusers optimization docs, FP8/NVFP4/MXFP8, TensorRT, FlashAttention diffusion. -- Existing PR context to consider: `#6658 #6670 #6675 #6679 #6680 #6690 #6694 #6700 #5872`. diff --git a/studio/backend/async_task_outputs/async_task_output_cothbh.md b/studio/backend/async_task_outputs/async_task_output_cothbh.md deleted file mode 100644 index bc419cf449..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_cothbh.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest benchmark/fix task from the completed state: the int8 dense quant crash investigation, fix, verification, and stacked PR are already done. Do not rerun the completed background benchmarks, pytest runs, branch push, or PR creation; PR #6716 exists, includes the pre-commit.ci autofix, and quant tests passed after syncing. Next tool action is to proceed only with the next separate benchmark failure if the broader effort continues, starting by inspecting the relevant code/logs for the T5 text-encoder fp8 issue rather than touching the completed int8 PR work. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_d1c766.md b/studio/backend/async_task_outputs/async_task_output_d1c766.md deleted file mode 100644 index 9f7a532071..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_d1c766.md +++ /dev/null @@ -1 +0,0 @@ -- NVFP4 Phase2C done; pushed `dbb029256` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_d211gb.md b/studio/backend/async_task_outputs/async_task_output_d211gb.md deleted file mode 100644 index a3f6e0b898..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_d211gb.md +++ /dev/null @@ -1,4 +0,0 @@ -- User requested testing safetensors inplace ops and `global` buffers. -- Desired defaults: `channels_last`, compiled dequant only, cuDNN if supported, inplace/eager opts. -- Max adds compile all + tf32. -- Check `https://github.com/unslothai/unsloth/blob/main/unsloth/models/llama.py#L1236`-`L1354`. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_d616s6.md b/studio/backend/async_task_outputs/async_task_output_d616s6.md deleted file mode 100644 index f72024acd8..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_d616s6.md +++ /dev/null @@ -1,2 +0,0 @@ -- Advanced moved to right card, open default; Chat-like toggle. -- `images-page.tsx`; `tsc`, build ok. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_de1ekr.md b/studio/backend/async_task_outputs/async_task_output_de1ekr.md deleted file mode 100644 index bc37983717..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_de1ekr.md +++ /dev/null @@ -1 +0,0 @@ -- Found/fixed `is_gguf` bug in `core/inference/diffusion.py`; non-GGUF fp8/int8 now route to regional `compiled`. diff --git a/studio/backend/async_task_outputs/async_task_output_dhusjp.md b/studio/backend/async_task_outputs/async_task_output_dhusjp.md deleted file mode 100644 index 715aed7bee..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_dhusjp.md +++ /dev/null @@ -1 +0,0 @@ -- Check NVFP4; test FP8/NVFP4 lowmem \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_djnp8g.md b/studio/backend/async_task_outputs/async_task_output_djnp8g.md deleted file mode 100644 index 438a77174a..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_djnp8g.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest directive to keep improving diffusion efficiency, memory use, and speed without unacceptable accuracy loss, specifically building on the validated Flux.1-dev FBCache Phase 12 work. The next concrete action is to continue wiring FBCache into the backend after the completed edit to `studio/backend/models/inference.py`: update the route/request plumbing as needed, finish status model exposure, then validate the new cache parameters and compile `fullgraph=False` behavior without re-running already completed Flux validation tools. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ed9jtn.md b/studio/backend/async_task_outputs/async_task_output_ed9jtn.md deleted file mode 100644 index 04eb6927ba..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ed9jtn.md +++ /dev/null @@ -1 +0,0 @@ -- Asked `NXFP4` via PyTorch/TorchAO \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_f8uf97.md b/studio/backend/async_task_outputs/async_task_output_f8uf97.md deleted file mode 100644 index fbc467182f..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_f8uf97.md +++ /dev/null @@ -1,5 +0,0 @@ -- Decision/request: make `Advanced` closed by default, because user wants it not open initially. -- Decision/request: when `Advanced` is expanded, keep the Advanced icon stationary; it should look like `Chat`, so the icon should not move. -- Reference URL provided verbatim: `https://huggingface.co/datasets/danielhanchen/screenshots/discussions/26` -- COMPLETED: no implementation in this span. -- PENDING: apply the UI behavior/icon alignment changes. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_flqnxc.md b/studio/backend/async_task_outputs/async_task_output_flqnxc.md deleted file mode 100644 index 2ec30e8f20..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_flqnxc.md +++ /dev/null @@ -1 +0,0 @@ -- `https://huggingface.co/datasets/danielhanchen/screenshots/discussions/25` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ftur3l.md b/studio/backend/async_task_outputs/async_task_output_ftur3l.md deleted file mode 100644 index b34b31a632..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ftur3l.md +++ /dev/null @@ -1 +0,0 @@ -- Load Studio `--secure`; ask sd.cpp \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_g7yinb.md b/studio/backend/async_task_outputs/async_task_output_g7yinb.md deleted file mode 100644 index f87185f3ae..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_g7yinb.md +++ /dev/null @@ -1,6 +0,0 @@ -- Decision: next lever should target inference-time peak memory/speed, because Phases 7-9 already covered compile, dense+torchao quant, consumer fast-accum gating, and pre-quantized loading/load VRAM. -- Goal preserved: push diffusion speed and `peak memory (both load and inference)` under ~25% accuracy cost across consumer/data-center GPUs and ideally Mac/AMD/Intel/CPU, for `GGUF + safetensors`. -- Mentioned candidate levers: step caching, attention backends, VAE tiling for decode peak memory. -- No files edited/created. -- No commands completed. -- Error: `API Error: Connection closed mid-response.` unresolved; assistant response may be incomplete. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ggz4m7.md b/studio/backend/async_task_outputs/async_task_output_ggz4m7.md deleted file mode 100644 index 62646a10f7..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ggz4m7.md +++ /dev/null @@ -1 +0,0 @@ -- PENDING: Z-Image, LPIPS, CPU cmp \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_gs83o9.md b/studio/backend/async_task_outputs/async_task_output_gs83o9.md deleted file mode 100644 index 2beee26459..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_gs83o9.md +++ /dev/null @@ -1 +0,0 @@ -- Waiting user decision. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_hbajbw.md b/studio/backend/async_task_outputs/async_task_output_hbajbw.md deleted file mode 100644 index d952820cab..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_hbajbw.md +++ /dev/null @@ -1 +0,0 @@ -- Studio launched: `--secure`, `:8890` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_hh8yan.md b/studio/backend/async_task_outputs/async_task_output_hh8yan.md deleted file mode 100644 index 5d392f8263..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_hh8yan.md +++ /dev/null @@ -1 +0,0 @@ -- COMPLETED; pending user go-ahead \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_hm7vhm.md b/studio/backend/async_task_outputs/async_task_output_hm7vhm.md deleted file mode 100644 index 1f2ed8aeed..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_hm7vhm.md +++ /dev/null @@ -1 +0,0 @@ -- PENDING: install `torch 2.11` `NVFP4` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_hu6k45.md b/studio/backend/async_task_outputs/async_task_output_hu6k45.md deleted file mode 100644 index 70f7bbe71a..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_hu6k45.md +++ /dev/null @@ -1 +0,0 @@ -- Pending: run `gh pr checks 6703`; use `CronList`/`CronDelete` only if checks done. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_irlty4.md b/studio/backend/async_task_outputs/async_task_output_irlty4.md deleted file mode 100644 index 98f1042c87..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_irlty4.md +++ /dev/null @@ -1 +0,0 @@ -- Studio up: `HTTP 200`; URL `https://courtesy-complicated-newman-citizens.trycloudflare.com` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_iyj2u3.md b/studio/backend/async_task_outputs/async_task_output_iyj2u3.md deleted file mode 100644 index 086f751253..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_iyj2u3.md +++ /dev/null @@ -1 +0,0 @@ -- fp8 auto; PR `6694` merged-ready. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_j302lq.md b/studio/backend/async_task_outputs/async_task_output_j302lq.md deleted file mode 100644 index 3a0001bdd3..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_j302lq.md +++ /dev/null @@ -1 +0,0 @@ -- Bench: 2.56x faster `1.83 s -> 0.71 s`; stack: compile+cudnn attention+channels_last+benchmark. FBCache extra for many-step only. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_j6ykxj.md b/studio/backend/async_task_outputs/async_task_output_j6ykxj.md deleted file mode 100644 index 9894461ab6..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_j6ykxj.md +++ /dev/null @@ -1 +0,0 @@ -- Complete; pending `push the PRs`/`set up the mirror` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_jm0l3y.md b/studio/backend/async_task_outputs/async_task_output_jm0l3y.md deleted file mode 100644 index e425cb6100..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_jm0l3y.md +++ /dev/null @@ -1 +0,0 @@ -- Decided fp8 fastest; GGUF smallest/default. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_klpaw5.md b/studio/backend/async_task_outputs/async_task_output_klpaw5.md deleted file mode 100644 index f4861b66b8..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_klpaw5.md +++ /dev/null @@ -1 +0,0 @@ -- PR-ready; pending `push the PRs` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ktq93m.md b/studio/backend/async_task_outputs/async_task_output_ktq93m.md deleted file mode 100644 index 88445f3962..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ktq93m.md +++ /dev/null @@ -1 +0,0 @@ -- No API errors; only benign ONNX warning \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_l026u1.md b/studio/backend/async_task_outputs/async_task_output_l026u1.md deleted file mode 100644 index 7708e7e86f..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_l026u1.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest directive: remove all global buffers since they looked useless, then ensure nothing else broke. The last completed action read the affected section of tests/test_diffusion_speed.py, so the next action is to edit that file to remove the weight_buffer stub/assertions without re-reading or rerunning the completed Read. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_l4hdlx.md b/studio/backend/async_task_outputs/async_task_output_l4hdlx.md deleted file mode 100644 index 4b073cd086..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_l4hdlx.md +++ /dev/null @@ -1 +0,0 @@ -- PENDING: user decide commit scope \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ldy0ne.md b/studio/backend/async_task_outputs/async_task_output_ldy0ne.md deleted file mode 100644 index b60abac3b3..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ldy0ne.md +++ /dev/null @@ -1 +0,0 @@ -- compile gives most win; conv opts ~noop \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_m9h6x3.md b/studio/backend/async_task_outputs/async_task_output_m9h6x3.md deleted file mode 100644 index 28e56d51be..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_m9h6x3.md +++ /dev/null @@ -1,2 +0,0 @@ -- Phase13 TE offload rejected: `+686%` latency for `-11%` peak; code reverted. -- `#6703` green: `42 pass`, `1 skipping`. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_mwjqhd.md b/studio/backend/async_task_outputs/async_task_output_mwjqhd.md deleted file mode 100644 index 79b2bcbacc..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_mwjqhd.md +++ /dev/null @@ -1 +0,0 @@ -- Diffusion: fp8 fastest/`auto`; caveats VRAM/shape \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_na1kau.md b/studio/backend/async_task_outputs/async_task_output_na1kau.md deleted file mode 100644 index 7abf640256..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_na1kau.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest directive: remove all global buffer machinery because the benchmarks showed it was useless, and ensure nothing else broke. The production module, diffusion comments, speed docstrings, and `test_diffusion_gguf_compile.py` were already edited; the last completed tool read `studio/backend/tests/test_diffusion_speed.py` and showed the remaining `weight_buffer` test stub/assertions. Next action is to edit `test_diffusion_speed.py` to remove `weight_buffer` from the stub and expectations, then grep for remaining buffer references and run the relevant diffusion tests. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_nagyj1.md b/studio/backend/async_task_outputs/async_task_output_nagyj1.md deleted file mode 100644 index add2b66bf9..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_nagyj1.md +++ /dev/null @@ -1 +0,0 @@ -- Pending: needs user go-ahead. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_nc9gvk.md b/studio/backend/async_task_outputs/async_task_output_nc9gvk.md deleted file mode 100644 index 5dd823b2ac..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_nc9gvk.md +++ /dev/null @@ -1 +0,0 @@ -- NVFP4 tried; Phase4 started \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_nwjxf4.md b/studio/backend/async_task_outputs/async_task_output_nwjxf4.md deleted file mode 100644 index 5634f36d20..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_nwjxf4.md +++ /dev/null @@ -1 +0,0 @@ -- User asked: `So FP8 is the fastest?` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_o5wtl3.md b/studio/backend/async_task_outputs/async_task_output_o5wtl3.md deleted file mode 100644 index cf875b16dc..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_o5wtl3.md +++ /dev/null @@ -1 +0,0 @@ -- Decision: keep `default dynamic=True`. diff --git a/studio/backend/async_task_outputs/async_task_output_o77ymr.md b/studio/backend/async_task_outputs/async_task_output_o77ymr.md deleted file mode 100644 index b267fbbe49..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_o77ymr.md +++ /dev/null @@ -1 +0,0 @@ -- Multi-ref FLUX.2 done; verified live. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_p8crwl.md b/studio/backend/async_task_outputs/async_task_output_p8crwl.md deleted file mode 100644 index 641c6d891d..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_p8crwl.md +++ /dev/null @@ -1 +0,0 @@ -- Use 5 Opus subagents; recheck links \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_p8rq7l.md b/studio/backend/async_task_outputs/async_task_output_p8rq7l.md deleted file mode 100644 index 846378cb67..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_p8rq7l.md +++ /dev/null @@ -1 +0,0 @@ -- Asked `How about memory usage and size` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_peiz71.md b/studio/backend/async_task_outputs/async_task_output_peiz71.md deleted file mode 100644 index 9b23c05cbe..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_peiz71.md +++ /dev/null @@ -1 +0,0 @@ -- `dynamic=True` default; global buffers \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_pnabd9.md b/studio/backend/async_task_outputs/async_task_output_pnabd9.md deleted file mode 100644 index 4462af1673..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_pnabd9.md +++ /dev/null @@ -1 +0,0 @@ -- Phase 7 pushed: `ede94176f`, green. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_ptrx41.md b/studio/backend/async_task_outputs/async_task_output_ptrx41.md deleted file mode 100644 index 770026eb3d..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_ptrx41.md +++ /dev/null @@ -1 +0,0 @@ -- User asks benchmark/settings recap \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_pvflyj.md b/studio/backend/async_task_outputs/async_task_output_pvflyj.md deleted file mode 100644 index a9a5b698eb..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_pvflyj.md +++ /dev/null @@ -1 +0,0 @@ -- Asked: `Search HuggingFace for any NVFP4 diffusion models and see if we can try them` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_pysgn3.md b/studio/backend/async_task_outputs/async_task_output_pysgn3.md deleted file mode 100644 index af50d169be..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_pysgn3.md +++ /dev/null @@ -1,4 +0,0 @@ -- Decision: discussion `#26` changes are complete because the UI is now closed by default and the Chat-style toggle is fixed. -- Completed: verified the discussion `#26` behavior; no further autonomous action was identified in this span. -- Pending: waiting for user go-ahead on remaining choices: `push the frontend + installer`, `backend path 1/2/3`, `set up the mirror`, `cancel the loop`, or continued UI feedback iteration. -- Live studio remains available at `https://incentive-topics-patrick-commissioners.trycloudflare.com` with credentials `unsloth` / `diffusion-Thrg7s3tr1-UoGYq`. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_pzxpq2.md b/studio/backend/async_task_outputs/async_task_output_pzxpq2.md deleted file mode 100644 index 32090b6d34..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_pzxpq2.md +++ /dev/null @@ -1,2 +0,0 @@ -- fixed UI+resize -- pending push/backend \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_q479i8.md b/studio/backend/async_task_outputs/async_task_output_q479i8.md deleted file mode 100644 index 7c5c1259f9..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_q479i8.md +++ /dev/null @@ -1 +0,0 @@ -Continue making diffusion more efficient with lower loading and inference peak memory, faster execution, and acceptable accuracy loss only within the user’s stated tolerance. The latest edit to inference.py is done; do not re-run it. Next, proceed from the completed status-response-model edit by inspecting the updated inference.py context only as needed, then continue implementing or validating the remaining quantization/memory/speed changes implied by the stop-hook feedback. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_qz0t3d.md b/studio/backend/async_task_outputs/async_task_output_qz0t3d.md deleted file mode 100644 index 492a77411b..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_qz0t3d.md +++ /dev/null @@ -1 +0,0 @@ -- Asked diffusers `int8 / float8` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_rbj2zk.md b/studio/backend/async_task_outputs/async_task_output_rbj2zk.md deleted file mode 100644 index a9e465d513..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_rbj2zk.md +++ /dev/null @@ -1,7 +0,0 @@ -- Decision: Used cached/local context first per `AGENTS.md`; no cache existed for PR `#6658`, so fetched live PR metadata with `gh`. -- Decision: Summarized PR `#6658` as foundational diffusion/image-generation support; detection is architecture-based so one family entry covers quants/variants. -- Command run: `cd /mnt/disks/unslothai/ubuntu/workspace_81 && echo "=== cached summary? ===" && ls -la data/pr_summaries/unslothai__unsloth/pr_6658.json 2>/dev/null || echo "(no cache)" && echo "=== PR meta ===" && gh pr view 6658 --repo unslothai/unsloth --json number,title,state,author,body ...`; key output: `(no cache)`, `#6658 [OPEN] Studio: add image generation with Unsloth GGUFs`, author `oobabooga`. -- Completed: Captured PR body summary: new Image tab with download menu, generation inputs, output image/gallery. -- Completed: Identified integrated families: `z-image`, `qwen-image`, `flux.1`, `flux.2-klein`; all Unsloth GGUF transformer-only with base diffusers components resolved via `base_model`. -- Completed: Noted excluded models/reasons: `ERNIE-Image`, `FLUX.2-dev`, editing models. -- Pending: none in this span. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_s9q7qs.md b/studio/backend/async_task_outputs/async_task_output_s9q7qs.md deleted file mode 100644 index 69b9e51eca..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_s9q7qs.md +++ /dev/null @@ -1 +0,0 @@ -- Done; pending user choice \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_t0my78.md b/studio/backend/async_task_outputs/async_task_output_t0my78.md deleted file mode 100644 index 63e72920d2..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_t0my78.md +++ /dev/null @@ -1,8 +0,0 @@ -- Decision: discussion #26 fixed by using one always-present top-bar `LayoutAlignRightIcon`; guarantees no icon swap or movement and matches Chat style. -- Decision: Advanced panel now closed by default so preview canvas loads full width. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend/src/features/images/images-page.tsx`: default Advanced state closed; replaced conditional gear/header close icons with a single fixed top-right `LayoutAlignRightIcon` toggle; removed redundant panel header close button. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/diffusion-workflows-pr-plan.md`: updated plan note for discussion #26. -- Ran `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend && npx tsc --noEmit ... && npm run build`; exit 0, key output `=== TSC OK ===`, `build exit 0 -> /mnt/disks/unslothai/ubuntu/workspace_81/logs/fe_build_222352.log`. -- Ran Playwright script `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_toggle.py`; exit 0, saved `toggle_closed.png` and `toggle_open.png`; output `ICON MOVED: False`, both boxes `{x:1558, y:11, width:34, height:34}`. -- Verified visually via `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/toggle_open.png` and `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/toggle_closed.png`. -- Completed: discussion #26 changes and verification. Pending: local frontend changes for PR 2 remain uncommitted. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_v3swmu.md b/studio/backend/async_task_outputs/async_task_output_v3swmu.md deleted file mode 100644 index e3cbe806cc..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_v3swmu.md +++ /dev/null @@ -1 +0,0 @@ -- Asked compile flags/CUDAGraphs \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_v6m49g.md b/studio/backend/async_task_outputs/async_task_output_v6m49g.md deleted file mode 100644 index 10940f22cf..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_v6m49g.md +++ /dev/null @@ -1 +0,0 @@ -- Done: adv open; `/16` resize \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_vbfgbz.md b/studio/backend/async_task_outputs/async_task_output_vbfgbz.md deleted file mode 100644 index 0405db9393..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_vbfgbz.md +++ /dev/null @@ -1 +0,0 @@ -- Done/no errors; pending go-ahead \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_viqlxl.md b/studio/backend/async_task_outputs/async_task_output_viqlxl.md deleted file mode 100644 index 2d4c6bf7d4..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_viqlxl.md +++ /dev/null @@ -1 +0,0 @@ -- FP8/INT8 PASS; all `/goal` done diff --git a/studio/backend/async_task_outputs/async_task_output_vn6a5d.md b/studio/backend/async_task_outputs/async_task_output_vn6a5d.md deleted file mode 100644 index fa0bf7662a..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_vn6a5d.md +++ /dev/null @@ -1,4 +0,0 @@ -- Decided pre-quantized torchao loading fixes fast-mode load VRAM; offline quant avoids dense bf16 GPU materialization. -- Created PR `https://github.com/unslothai/unsloth/pull/6700`, base `diffusion-phase8-quant`, branch `diffusion-phase9-prequant`, commit `b90f83346`. -- Added `diffusion_prequant.py`, builder/verify scripts, request field/tests. -- Verified: load peak `21.2 GB` -> `14.6 GB`, disk `6.28 GB`, `LPIPS 0.0`; `200` diffusion tests passed. diff --git a/studio/backend/async_task_outputs/async_task_output_vudd2x.md b/studio/backend/async_task_outputs/async_task_output_vudd2x.md deleted file mode 100644 index 751e263b51..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_vudd2x.md +++ /dev/null @@ -1,2 +0,0 @@ -- User asked: `Is there a way to fix the loading VRAM issue - pre-quantized?` -- Pending: answer likely needs repo/PR context; no commands or edits in this span. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_w0h2kg.md b/studio/backend/async_task_outputs/async_task_output_w0h2kg.md deleted file mode 100644 index b1cd283280..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_w0h2kg.md +++ /dev/null @@ -1,9 +0,0 @@ -- Approved plan saved at `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/wobbly-jumping-narwhal.md`. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_patch_backend.py`; lazy wraps `unsloth_zoo.temporary_patches.utils` `patch_function`/`can_safely_patch`/`restore_original`. -- Edited `diffusion_eager_patches.py`, `diffusion_gguf_compile.py`, `diffusion.py`, `tests/conftest.py`, `test_diffusion_gguf_compile.py`; CPU tests require `UNSLOTH_ALLOW_CPU=1`. -- Created `diffusion_arch_patches.py` + `test_diffusion_arch_patches.py`; qwen `_modulate` and z-image gated residuals use compile-safe `torch.addcmul`. -- Created `scripts/arch_patch_bench.py`. -- Tests: `77 passed, 1 skipped`; full diffusion: `296 passed, 1 skipped`. -- Bench: qwen parity default `5.952s`, max `2.741s`; arch isolation z-image `-0.2%`, qwen `+0.3%`. -- Alloc profile: `alloc_retries delta == 0`, reserved VRAM flat; global scratch buffers not useful. -- Pending: user asked whether to extend same addcmul recipe to flux.1/flux.2. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_whtxy2.md b/studio/backend/async_task_outputs/async_task_output_whtxy2.md deleted file mode 100644 index 33d12e7158..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_whtxy2.md +++ /dev/null @@ -1 +0,0 @@ -- Done: summarized `outputs/` benches \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_xcofo0.md b/studio/backend/async_task_outputs/async_task_output_xcofo0.md deleted file mode 100644 index ec7c368a0f..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_xcofo0.md +++ /dev/null @@ -1,2 +0,0 @@ -- DONE: build verified -- PENDING: choice \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_xds3jp.md b/studio/backend/async_task_outputs/async_task_output_xds3jp.md deleted file mode 100644 index e1645a50cb..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_xds3jp.md +++ /dev/null @@ -1 +0,0 @@ -- User asked: `So overall how much did increase the speed by?` \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_y2psj5.md b/studio/backend/async_task_outputs/async_task_output_y2psj5.md deleted file mode 100644 index 9962431fd7..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_y2psj5.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest directive: examine why int8 failed for some cases and fix if possible. The latest assistant state says this was completed: int8 failures were root-caused, fixed, verified on B200, tests passed, and PR #6716 was shipped stacked on #6703. No tool is in flight; the next action is to carry forward that completed status and, if the conversation continues, answer only any new user follow-up such as tackling the separate T5 fp8 issue. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_yz7170.md b/studio/backend/async_task_outputs/async_task_output_yz7170.md deleted file mode 100644 index 8f91db914e..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_yz7170.md +++ /dev/null @@ -1 +0,0 @@ -Continue the latest directive: “make a PR or push to another,” meaning push the completed diffusion image workflow work to the user’s fork and open the appropriate stacked PR. The full diffusion test run already completed, so do not re-run it just because of compaction; inspect its result/log if needed, then proceed with committing the work, pushing the base/work branches as described, and opening the PR against the fork branch. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_zi1igk.md b/studio/backend/async_task_outputs/async_task_output_zi1igk.md deleted file mode 100644 index cd0885926c..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_zi1igk.md +++ /dev/null @@ -1 +0,0 @@ -- Explained diffusion models/int8 UI gap. \ No newline at end of file diff --git a/studio/backend/async_task_outputs/async_task_output_zsrnxt.md b/studio/backend/async_task_outputs/async_task_output_zsrnxt.md deleted file mode 100644 index cff3739f42..0000000000 --- a/studio/backend/async_task_outputs/async_task_output_zsrnxt.md +++ /dev/null @@ -1 +0,0 @@ -- NVFP4 slow; added `scripts/nvfp4_probe.py` \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_1csyea.md b/studio/frontend/async_task_outputs/async_task_output_1csyea.md deleted file mode 100644 index a40c418b32..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_1csyea.md +++ /dev/null @@ -1 +0,0 @@ -- Studio live; pwd set; routes fixed \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_560grk.md b/studio/frontend/async_task_outputs/async_task_output_560grk.md deleted file mode 100644 index 531ef4548d..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_560grk.md +++ /dev/null @@ -1 +0,0 @@ -- Pending `/loop`: `Every 30 minutes` \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_6vw4lz.md b/studio/frontend/async_task_outputs/async_task_output_6vw4lz.md deleted file mode 100644 index 193984e94b..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_6vw4lz.md +++ /dev/null @@ -1 +0,0 @@ -- Continue work \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_7p2hgc.md b/studio/frontend/async_task_outputs/async_task_output_7p2hgc.md deleted file mode 100644 index 7e87657306..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_7p2hgc.md +++ /dev/null @@ -1 +0,0 @@ -- Started planner for diffusion UI; bg task `b1s4duoo4`; then user interrupted. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_bnb814.md b/studio/frontend/async_task_outputs/async_task_output_bnb814.md deleted file mode 100644 index 545387651e..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_bnb814.md +++ /dev/null @@ -1 +0,0 @@ -- Done: safetensors; tests 73 pass \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_cymhgv.md b/studio/frontend/async_task_outputs/async_task_output_cymhgv.md deleted file mode 100644 index 7da8bec177..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_cymhgv.md +++ /dev/null @@ -1,3 +0,0 @@ -- User directive received: `if the /goal is not yet achieved or there is an API error, say "Continue work"` -- No files edited, no commands run, no errors observed in this span. -- Pending: determine whether `/goal` is achieved; if not achieved or API error occurred, output `Continue work`. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_d8ngtg.md b/studio/frontend/async_task_outputs/async_task_output_d8ngtg.md deleted file mode 100644 index f0a3b252b4..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_d8ngtg.md +++ /dev/null @@ -1 +0,0 @@ -Continue work. Resume the current diffusion safetensors implementation task because the /goal is not yet achieved. The last completed action created `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_advanced.py`; next, run that Playwright screenshot script to capture the Advanced panel, then continue verification from its result without re-running the completed build or file creation. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_edpp2u.md b/studio/frontend/async_task_outputs/async_task_output_edpp2u.md deleted file mode 100644 index bea51ecfe6..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_edpp2u.md +++ /dev/null @@ -1,6 +0,0 @@ -- Fixed numeric spinner overlap in `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend/src/features/images/images-page.tsx`; hide all native spinners, widen `SliderField` number input. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/plans/diffusion-workflows-studio.md`. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_images_page.py`. -- Ran `npm run build`: exit 0. Playwright screenshots saved to `outputs/ui_shots/`; verified values no longer covered. -- Secure Studio relaunched: `https://perhaps-nick-valentine-park.trycloudflare.com`, login `unsloth` / `diffusion-Thrg7s3tr1-UoGYq`. -- Pending: tabbed workflow shell, Advanced Options, img2img. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_g3wn5w.md b/studio/frontend/async_task_outputs/async_task_output_g3wn5w.md deleted file mode 100644 index 9f3ac55aea..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_g3wn5w.md +++ /dev/null @@ -1 +0,0 @@ -- Comment; `unsloth/*` `safetensors` \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_heqa1v.md b/studio/frontend/async_task_outputs/async_task_output_heqa1v.md deleted file mode 100644 index d61513793b..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_heqa1v.md +++ /dev/null @@ -1 +0,0 @@ -Continue work on the current diffusion safetensors implementation task, specifically the Advanced Options panel: frontend fields are wired, task #147 is in progress, and `tsc` plus build already passed. The last completed tool action wrote `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_advanced.py`; do not re-run that write. Next, run that screenshot script to capture/verify the Advanced panel in the running studio, then inspect the result and continue verification or fixes as needed. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_iedlnv.md b/studio/frontend/async_task_outputs/async_task_output_iedlnv.md deleted file mode 100644 index 3e6ba8ca5a..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_iedlnv.md +++ /dev/null @@ -1,24 +0,0 @@ -- Decision: completed the frontend **Transform** tab as the next vertical slice because the img2img backend was already verified; UI now exposes Create/Transform workflow selection, gated by backend `status.workflows`. -- Decision: restarted the secure studio because the running process predated backend `workflows` support and kept Transform disabled; after restart, `workflows: ['txt2img', 'img2img', 'inpaint']` enabled the tab. -- Decision: added backend unit coverage for img2img to lock in `Pipeline.from_pipe` reuse and unsupported-family rejection. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend/src/features/images/api.ts`: added request/status typing for `init_image`, `mask_image`, `strength`, and workflow capability fields used by img2img UI. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend/src/features/images/images-page.tsx`: added `WORKFLOW_TABS`, workflow state, init image state, strength state, `ImageDropzone`, Create/Transform segmented tabs, Transform-only dropzone + Strength slider, capability gating, and threaded `init_image`/`strength` into generation. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/tests/test_diffusion_backend.py`: extended fakes and added img2img backend tests. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/scripts/shoot_transform_tab.py`: Playwright screenshot helper that loads a model, opens Transform, uploads a sample image, and saves `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/transform_tab.png`. -- Command run: `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/frontend && npx tsc --noEmit`; exit `0`, only npm warning `Unknown project config "min-release-age"`. -- Command run: `npm run build`; exit `0`, key output `✓ built in 1.95s`. -- Command run: `curl -s -m5 -o /dev/null -w "root=%{http_code}\n" http://127.0.0.1:8898/ ; ss -ltn | grep -q :8898 && echo listening || echo down`; key output `root=200`, `listening`. -- Command run: `/mnt/disks/unslothai/ubuntu/workspace_81/bin/python3 scripts/shoot_transform_tab.py`; first capture saved `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/transform_tab.png` but Transform stayed disabled because old backend was running. -- Error encountered: screenshot script initially logged repeated click waits ending with `element is not enabled`; resolved by restarting studio to pick up new backend `status.workflows`. -- Command run: `pkill -f "run.py --host 127.0.0.1 --port 8898"` and `pkill -f "unsloth studio --secure -p 8898"`; exit `144`; follow-up confirmed `(8898 free)`. -- Command run: relaunched secure studio with `UNSLOTH_STUDIO_HOME=/mnt/disks/unslothai/ubuntu/workspace_81/studio_home CUDA_VISIBLE_DEVICES=6,7 HF_HUB_DISABLE_XET=1 HF_TOKEN...`; background ID `bfrwp6n0d`. -- Command run: wait/check new studio; key output `https://function-shopping-grew-shoot.trycloudflare.com` and `root=200`. -- Command run: recaptured Transform tab with `/mnt/disks/unslothai/ubuntu/workspace_81/bin/python3 scripts/shoot_transform_tab.py`; key output `workflows: ['txt2img', 'img2img', 'inpaint']` and saved `/mnt/disks/unslothai/ubuntu/workspace_81/outputs/ui_shots/transform_tab.png`. -- Command run: live HTTP round-trip test through `http://127.0.0.1:8898`; key output `loaded: True workflows: ['txt2img', 'img2img', 'inpaint']`, `txt2img: 200`, `img2img HTTP: 200 images: 1`. -- Command run: `cd /mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend && /mnt/disks/unslothai/ubuntu/workspace_81/bin/python3 -m pytest tests/test_diffusion_backend.py tests/test_diffusion_routes.py -q`; key output `75 passed, 5 warnings in 21.40s`. -- Task update: marked task `#145` updated/completed for img2img vertical slice. -- Task update: updated task `#144` as partial: workflow tab strip and capability gating done; remaining tabs are Edit/Extend/Control/Enhance and Advanced Options work under task `147`. -- Completed: img2img full vertical slice verified across backend, frontend Transform tab, live HTTP, screenshot, and unit tests. -- Completed: running secure studio available at `https://function-shopping-grew-shoot.trycloudflare.com` with login `unsloth` / `diffusion-Thrg7s3tr1-UoGYq`. -- Pending: no commit or push performed. -- Pending: next planned work is **Edit tab** with inpaint mask canvas and instruction editing for Qwen-Image-Edit/FLUX Kontext, then Advanced Options panel for speed/compile/quant/memory knobs. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_jlf3lc.md b/studio/frontend/async_task_outputs/async_task_output_jlf3lc.md deleted file mode 100644 index 80fb824a49..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_jlf3lc.md +++ /dev/null @@ -1 +0,0 @@ -- Qwen layered unsupported; use non-layered \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_jqe936.md b/studio/frontend/async_task_outputs/async_task_output_jqe936.md deleted file mode 100644 index 6cb8c77002..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_jqe936.md +++ /dev/null @@ -1 +0,0 @@ -- `@codex review`+`/gemini review`x28 \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_kc50ea.md b/studio/frontend/async_task_outputs/async_task_output_kc50ea.md deleted file mode 100644 index f0e1a47872..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_kc50ea.md +++ /dev/null @@ -1 +0,0 @@ -- Backend WIP; frontend pending \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_oh38q4.md b/studio/frontend/async_task_outputs/async_task_output_oh38q4.md deleted file mode 100644 index 94d75cf9ec..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_oh38q4.md +++ /dev/null @@ -1 +0,0 @@ -- PENDING: Studio images UI + diffusion workflows. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_ot9we3.md b/studio/frontend/async_task_outputs/async_task_output_ot9we3.md deleted file mode 100644 index a249b8bfba..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_ot9we3.md +++ /dev/null @@ -1 +0,0 @@ -- Completed safetensors E2E; Studio up \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_phzu13.md b/studio/frontend/async_task_outputs/async_task_output_phzu13.md deleted file mode 100644 index cb33f7edf8..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_phzu13.md +++ /dev/null @@ -1,11 +0,0 @@ -- Decision: implemented first real workflow as `img2img` backend foundation for Transform/Edit/Extend; diffusers has `FluxImg2ImgPipeline`, `QwenImageImg2ImgPipeline`, `ZImageImg2ImgPipeline`, no `Flux2KleinImg2ImgPipeline`, so flux.2-klein is gated off. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion_families.py`: added `img2img_pipeline_class`/`inpaint_pipeline_class` slots and populated flux.1/qwen-image/z-image. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/diffusion.py`: added base64-to-PIL decode, auxiliary workflow pipeline cache via `Pipeline.from_pipe`, txt2img/img2img/inpaint branching, `status.workflows`, and VAE dtype realignment before image-conditioned calls. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/models/inference.py`: added `init_image`, `mask_image`, `strength`; updated status response `workflows`. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/routes/inference.py`: forwards new fields; maps bad image/unsupported workflow to 400. -- Edited `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/studio/backend/core/inference/sd_cpp_backend.py`: accepts new kwargs but rejects image-conditioned requests clearly. -- Created `/mnt/disks/unslothai/ubuntu/workspace_81/unsloth/scripts/verify_img2img.py` for GPU verification. -- Commands: AST parse succeeded for changed backend files; pytest `tests/test_diffusion_backend.py tests/test_diffusion_routes.py tests/test_diffusion_engine_router.py -q` ended `85 passed, 5 warnings in 20.82s`. -- Error resolved: img2img initially failed with VAE dtype mismatch (`input BFloat16, bias float`) after txt2img upcasted shared VAE; fixed by aligning VAE dtype to compute dtype before image-conditioned calls. -- Verification: B200 run with `CUDA_VISIBLE_DEVICES=5 HF_HUB_DISABLE_XET=1 HF_TOKEN=$HF_TOKEN` loaded `unsloth/Z-Image-Turbo-unsloth-bnb-4bit`; `status workflows: ['txt2img', 'img2img', 'inpaint']`; base 2.4s, img2img 1.7s, mean-abs-diff 15.3, std 59.2. -- Completed: backend img2img/inpaint plumbing verified. Pending: frontend Transform tab, backend unit test for img2img branch, then Edit/inpaint mask canvas and Advanced Options. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_qh9doa.md b/studio/frontend/async_task_outputs/async_task_output_qh9doa.md deleted file mode 100644 index a9f20ba337..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_qh9doa.md +++ /dev/null @@ -1,2 +0,0 @@ -- User directed: `No need to enter plan mode - just use the parallel planner, and you plan here.` -- Pending: continue without plan mode. \ No newline at end of file diff --git a/studio/frontend/async_task_outputs/async_task_output_s15ff5.md b/studio/frontend/async_task_outputs/async_task_output_s15ff5.md deleted file mode 100644 index fb7b2b7e9d..0000000000 --- a/studio/frontend/async_task_outputs/async_task_output_s15ff5.md +++ /dev/null @@ -1 +0,0 @@ -- Asked why `Image generation failed.` \ No newline at end of file