From 692d3c197545ec8f7955c4c6a682cfb3cc92b68e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:43:56 -0700 Subject: [PATCH] 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"), )