From 6b9b1c72d3a5d3241faba3f9c2024d61bc712b68 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 1 Jul 2026 11:31:50 -0700 Subject: [PATCH] Studio diffusion (Phase 7): accuracy-preserving speed pass (2.2x via GGUF compile) (#6690) * Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict Phase 1 of porting the richer diffusion stack onto the image-generation backend. - Add a compartmentalized device/dtype policy module (diffusion_device.py) resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or fp32, never a silent fp16 that renders a black image. - Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved float16 to float32 for those families so they do not produce black images. - Split the backend locks: a generation holds only _generate_lock, so status, unload, and a new load are never blocked by a long denoise. Add per-generation cancellation via callback_on_step_end so an eviction or a superseding load preempts a running generation; a replacement load waits for it to stop before allocating, so two pipelines never sit in VRAM at once. - Validate a load request before the GPU handoff so an unloadable pick never evicts a working chat model, and reject missing local paths up front. - Add CPU-only tests for the device policy, dtype guard, lock split and cancellation, and validate-before-evict, plus a GPU benchmark/regression script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR against a saved reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and VAE tiling/slicing from measured free device memory vs the model's estimated resident footprint, then applies it to the built pipeline. auto stays resident when the model fits (byte-identical to the prior resident path), and falls to whole-module offload when tight; fast/balanced/low_vram are explicit overrides. Sequential submodule offload is unreliable for GGUF transformers on diffusers 0.38, so it falls back to whole-module offload and status reports the policy actually engaged. Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM 47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost. 73 prior + 35 new CPU tests pass. * Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level, use_stream) that keeps the transformer flowing through the GPU a few blocks at a time while the text encoder / VAE stay resident, and fix VAE tiling to drive the VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so status never overstates either, and group falls back to whole-module offload when the transformer can't be streamed. Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 -> 2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names now match that tradeoff: balanced = stream the transformer, low_vram = offload every component. auto picks group when the companions fit resident, else model. 112 CPU tests pass. * Studio diffusion (Phase 5): image quality-vs-quant accuracy harness Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold prompt + seed fixed, render a grid with a reference quant (default BF16), then render each candidate quant and measure drift from the reference. Records mean PSNR + SSIM (pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity (transformers, --clip), plus file size, latency, and peak VRAM, then prints a quality-vs-cost table and recommends the smallest quant within a quality budget. --selftest validates the metrics on synthetic images with no GPU or model. Verified on Z-Image (B200): the table degrades monotonically with quant size (Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat (~0.34) -- quantization erodes fine detail far more than prompt adherence. * Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32) Add a speed_mode knob (off by default, so the render path stays bit-identical): default applies channels_last VAE + regional torch.compile of the denoiser's repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional compile is gated off for the GGUF transformer (dequantises per-op) and for families flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image), so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed optims run before placement/offload, per the diffusers composition order. status now reports speed_mode + the optims actually engaged. Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last', 'tf32'], compile correctly skipped for GGUF; generation works in every mode. 121 CPU tests pass. * Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3) storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16 compute dtype while normalisations and embeddings stay full precision. Applied before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder dense). status reports which encoders were cast. Verified on Z-Image (B200, balanced/group mode where the encoder stays resident): generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload) at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off by default and documented as such, with the Phase 5 harness to size the cost. 127 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob) Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant (fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8 stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run before placement; status reports the mode actually engaged. This is the lean realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the 3045-line port. Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default); size it per model with the Phase 5 quality harness. diffusion_bench gains --text-encoder-quant. 129 CPU tests pass. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm / XPU; this covers the hardware diffusers serves poorly, consuming the same split GGUF assets Studio already curates. - sd_cpp_args.py: pure sd-cli command builder. Maps the family to its text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1 CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential) to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu / --vae-tiling / --diffusion-fa), so one user knob drives both engines. - sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary() with the same precedence as the llama finder (env override, then the Studio install root, then in-tree, then PATH), an is_available/version probe, and a one-shot subprocess generate that streams progress and returns the PNG. runtime_env() prepends the binary's directory to the platform library path so a prebuilt's bundled libstable-diffusion.so resolves. select_diffusion_engine() is the pure routing decision (GPU backends to diffusers, CPU/MPS to native when present). - install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt (macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the Studio install root. resolve_release_asset() is a pure, unit-tested host-to-asset matrix. - scripts/sd_cpp_smoke.py: end-to-end native generation harness. Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine, routing, runtime env, and the installer resolver. Full diffusion suite 166 passing. Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine Builds on Phase 4's native stable-diffusion.cpp engine, extending it from text-to-image to the wider feature surface, since sd.cpp supports all of these through the binary already. Pure command-builder additions plus one engine method, so the txt2img path is unchanged. - sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img + strength make a run img2img, adding mask makes it inpaint, ref_images drives FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and lora_dir + the prompt syntax select LoRAs. New SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run mode (input image + esrgan model, no prompt / text encoders). - sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so generate() (now carrying the conditioning flags) and a new upscale() reuse the same streaming / error / output-check path. - scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img / --strength / --upscale-model / --upscale-repeats. Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the upscale builder and its validation, and the engine's img2img + upscale paths. Full diffusion suite 176 passing. Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale (512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing coherent images. Video and the diffusers-path feature wiring are deferred. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): accuracy-preserving speed pass Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy bug and a dead-on-arrival speed path; this fixes both and adds the lossless / near-lossless wins, all measured on a B200. Correctness: - TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32 process-wide and never restored them, so a later `off` load silently inherited TF32 and was no longer bit-identical. Added snapshot_backend_flags / restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer runs and restored on unload. Verified: load max -> unload -> load off is now byte-identical (PSNR inf) to a fresh off. - sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and only checked the timeout after EOF, so a child stuck in model load / GPU init with no output ignored the timeout. Drained stdout on a reader thread with a wall-clock deadline. Added a silent-hang regression test. Speed (diffusers path), near-lossless, opt-in tiers: - Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks compiles and runs ~2.2x faster on the GGUF Z-Image transformer on torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move output quality. Gate relaxed; default tier delivers it. - cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs). - torch.inference_mode() around the pipeline call (lossless, strictly faster than the no_grad diffusers uses internally). Memory path: - VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers; the balanced (group) tier keeps exact slicing only, so it is now bit-identical to the resident image (verified PSNR inf) and slightly faster. - Group offload adds non_blocking + record_stream on the CUDA stream path to overlap each block's H2D copy with compute (lossless; gated on the installed diffusers signature so older versions still work). Native (sd.cpp) path: - native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a near-lossless CUDA win that was previously only added on offload tiers; max also -> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so it is never auto-on. Engine generate() merges it, de-duped against offload flags. Default profile: a GGUF model with no explicit speed_mode now resolves to the `default` profile (resolve_speed_mode), since compile's perturbation sits below the quantisation noise floor and so does not reduce quality versus the dense reference; out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is always honored, so the byte-identical path remains one flag away and is the regression reference. Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/ perf_verify.py (the B200 verification above), and diffusion_bench.py gains --speed-mode so the speed tiers are benchmarkable. Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore, GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags + the engine de-dup, and the sd-cli silent-hang timeout. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks The opt-in `max` speed tier now compiles the repeated block with mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode: Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are deliberately avoided -- both crash on the regionally-compiled block (its static output buffer is overwritten across denoise steps), measured. Adds two reproducible benchmarks used to validate the optimization research: - scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head. - scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes. Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen; coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune); FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load - snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the whole snapshot. restore_backend_flags restores each flag independently so one failure can't leave the others leaked process-wide. - load_pipeline restores the flags (and clears the GPU cache) when the build fails after apply_speed_optims mutated the process-wide flags but before _state captured them for unload to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and contaminated later off generations. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs without producing output (or without closing stdout) would never reach proc.wait and the wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the PROCESS, so the main thread always enforces the timeout and kills a hung process (which closes the pipe and ends the reader). Add a test that times out even when stdout blocks, and make the no-binary test hermetic so a host-installed sd-cli can't leak in. * Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts - diffusion_memory: when group offload is unavailable and the plan falls back to whole-module offload, enable VAE tiling (the group plan left it off, but the fallback is the low-VRAM path where the decode spike can OOM). Covers both the group and sequential fallback branches. - perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a balanced bit-identity regression actually fails the check instead of exiting 0. - compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and the load-progress poll has a 30 min deadline instead of looping forever on a hang. - test for the group->model fallback enabling VAE tiling. * Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening - install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket); extract through a per-member containment check (Zip-Slip guard); expanduser the --install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch the separately-published cudart runtime DLL archive so sd-cli.exe can start. - sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie. - tests: Zip-Slip rejection, normal extraction, studio-home discovery. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes the numbered files _ (base_0.png, base_1.png, ...) instead of the literal --output path. SdCppEngine.generate checked only the literal path, so a batch generation would exit 0 and then raise 'no image' (or return a stale file). generate now returns the literal path when present and otherwise falls back to the numbered siblings; single-image behavior is unchanged. Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected without error. * Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats Codex review on the native engine arg builder: - build_sd_cpp_command emitted --width/--height unconditionally, so an img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of the input. width/height are now Optional (None = unset): an image-conditioned run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives the size from the input image (set_width_and_height_if_unset); a plain txt2img run with unset dims keeps the prior 1024x1024 default; explicit dims are always honored. width/height are read only by the builder, so the type change is local. - build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...) that silently swallowed repeats=0 into sd-cli's default of one pass, turning an explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError and emits the flag for any explicit value != 1. Tests: img2img unset dims omit width/height (init_img and ref_images), explicit dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at the default. (Two pre-existing binary-discovery tests fail only because a real sd-cli is installed in this dev environment; unrelated to this change.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- scripts/compare_engines.py | 174 ++++++++++++++++ scripts/compile_probe.py | 189 ++++++++++++++++++ scripts/diffusion_bench.py | 9 + scripts/leverage_probe.py | 146 ++++++++++++++ scripts/perf_verify.py | 165 +++++++++++++++ studio/backend/core/inference/diffusion.py | 135 ++++++++----- .../core/inference/diffusion_families.py | 8 +- .../core/inference/diffusion_memory.py | 72 ++++--- .../backend/core/inference/diffusion_speed.py | 150 ++++++++++++-- studio/backend/core/inference/sd_cpp_args.py | 28 +++ .../backend/core/inference/sd_cpp_engine.py | 60 +++++- .../backend/tests/test_diffusion_backend.py | 16 +- studio/backend/tests/test_diffusion_memory.py | 15 ++ studio/backend/tests/test_diffusion_speed.py | 131 ++++++++++-- studio/backend/tests/test_sd_cpp_args.py | 11 + studio/backend/tests/test_sd_cpp_engine.py | 80 ++++++++ 16 files changed, 1261 insertions(+), 128 deletions(-) create mode 100644 scripts/compare_engines.py create mode 100644 scripts/compile_probe.py create mode 100644 scripts/leverage_probe.py create mode 100644 scripts/perf_verify.py diff --git a/scripts/compare_engines.py b/scripts/compare_engines.py new file mode 100644 index 0000000000..57abddfd3a --- /dev/null +++ b/scripts/compare_engines.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Head-to-head: PyTorch (diffusers GGUF) vs native stable-diffusion.cpp. + +Same Z-Image GGUF transformer, same VAE + text encoder, same resolution / steps / +seed, both resident (no CPU offload) on the same GPU. Reports per-engine compute +latency (model already loaded) so the denoise + VAE + TE work is compared fairly; +for sd.cpp it also reports the one-shot wall time (compute + the per-call model +reload, which a persistent sd-server would remove). + +PyTorch runs first (load / warmup / median), is unloaded, then sd.cpp runs. +""" + +from __future__ import annotations + +import argparse +import re +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend" +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" +_DONE_RE = re.compile(r"generate_image completed in ([0-9.]+)s") + + +def _median(xs): + return sorted(xs)[len(xs) // 2] + + +def bench_pytorch(repo, gguf, resolutions, steps, seed, iters): + import torch + from core.inference.diffusion import DiffusionBackend + + rows = [] + backend = DiffusionBackend() + for speed in ("off", "default"): + backend.begin_load(repo, gguf_filename = gguf, speed_mode = speed) + deadline = time.time() + 1800 # 30 min: a stuck download/load must not hang forever + while backend.load_progress().get("phase") != "ready": + prog = backend.load_progress() + if prog.get("phase") == "error": + raise RuntimeError(prog) + if time.time() > deadline: + raise TimeoutError(f"load timed out (last progress: {prog})") + time.sleep(0.5) + for res in resolutions: + + def gen(): + torch.cuda.synchronize() + t0 = time.time() + backend.generate( + prompt = PROMPT, + width = res, + height = res, + steps = steps, + guidance = 0.0, + seed = seed, + batch_size = 1, + ) + torch.cuda.synchronize() + return time.time() - t0 + + gen() # warmup (compiles for `default`) + med = _median([gen() for _ in range(iters)]) + rows.append(("pytorch", speed, res, med, None)) + print(f" pytorch speed={speed:7s} {res}px compute={med:.3f}s", flush = True) + backend.unload() + return rows + + +def bench_sdcpp(binary, gguf, vae, llm, resolutions, steps, seed, iters): + from core.inference.sd_cpp_args import SdCppGenParams, SdCppModelFiles + from core.inference.sd_cpp_engine import SdCppEngine + + engine = SdCppEngine(binary = binary) + if not engine.is_available(): + print(" sd.cpp binary not available; skipping", flush = True) + return [] + files = SdCppModelFiles(diffusion_model = gguf, vae = vae, llm = llm) + rows = [] + out_dir = Path("outputs/compare_engines") + out_dir.mkdir(parents = True, exist_ok = True) + for native in (None, "default"): # resident-no-fa vs resident+--diffusion-fa + for res in resolutions: + params = SdCppGenParams( + prompt = PROMPT, width = res, height = res, steps = steps, cfg_scale = 1.0, seed = seed + ) + computes, walls = [], [] + for _ in range(iters): + captured = {"c": None} + + def _log(ln): + m = _DONE_RE.search(ln) + if m: + captured["c"] = float(m.group(1)) + + t0 = time.time() + engine.generate( + files, + params, + output_path = str(out_dir / f"sd_{native}_{res}.png"), + offload = [], + native_speed = native, + on_log = _log, + ) + walls.append(time.time() - t0) + if captured["c"] is not None: + computes.append(captured["c"]) + med_c = _median(computes) if computes else None + med_w = _median(walls) + tag = "default(+fa)" if native == "default" else "off" + rows.append(("sdcpp", tag, res, med_c, med_w)) + print( + f" sdcpp speed={tag:12s} {res}px compute={med_c}s wall={med_w:.3f}s", + flush = True, + ) + return rows + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf-name", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument("--sd-binary", default = None) + p.add_argument( + "--sd-gguf", default = None, help = "local gguf for sd.cpp (default: same as pytorch via cache)" + ) + p.add_argument( + "--vae", + default = None, + help = "VAE safetensors for sd.cpp (required when benchmarking the sd.cpp engine)", + ) + p.add_argument( + "--llm", + default = None, + help = "text-encoder GGUF for sd.cpp (required when benchmarking the sd.cpp engine)", + ) + p.add_argument("--resolutions", default = "512,1024") + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--iters", type = int, default = 3) + args = p.parse_args(argv) + + from huggingface_hub import hf_hub_download + from core.inference.sd_cpp_engine import find_sd_cpp_binary + + resolutions = [int(x) for x in args.resolutions.split(",")] + sd_gguf = args.sd_gguf or hf_hub_download(args.repo, args.gguf_name) + binary = args.sd_binary or find_sd_cpp_binary() + + print("== PyTorch (diffusers GGUF) ==", flush = True) + pt = bench_pytorch(args.repo, args.gguf_name, resolutions, args.steps, args.seed, args.iters) + print("== stable-diffusion.cpp (native) ==", flush = True) + sd = bench_sdcpp( + binary, sd_gguf, args.vae, args.llm, resolutions, args.steps, args.seed, args.iters + ) + + print("\n==== COMPARISON (Z-Image-Turbo Q4, fixed seed, resident) ====", flush = True) + print(f"{'engine':9s} {'config':13s} {'res':>5s} {'compute_s':>10s} {'wall_s':>8s}", flush = True) + for eng, cfg, res, c, w in pt + sd: + cs = f"{c:.3f}" if c is not None else "n/a" + ws = f"{w:.3f}" if w is not None else "-" + print(f"{eng:9s} {cfg:13s} {res:5d} {cs:>10s} {ws:>8s}", flush = True) + print("COMPARE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/compile_probe.py b/scripts/compile_probe.py new file mode 100644 index 0000000000..5ff16ef763 --- /dev/null +++ b/scripts/compile_probe.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe: does regional ``torch.compile`` work on the GGUF diffusion transformer? + +The speed layer gates ``compile_repeated_blocks`` OFF for GGUF (it dequantises +per-op). Since the backend is GGUF-only, that makes regional compile dead on +every shipping model. This probe loads a GGUF transformer exactly as +``diffusion.py`` does, runs an eager generation, then compiles the repeated +denoiser block and runs the same seed again, reporting: whether compile raised, +per-generation latency eager vs compiled, and PSNR(compiled vs eager). If compile +is clean and PSNR is high, the gate can be relaxed for this family. + +Run on one CUDA GPU. Read-only w.r.t. the backend (does not import the gate). +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + + +def _psnr(a: "np.ndarray", b: "np.ndarray") -> float: + a = a.astype(np.float64) + b = b.astype(np.float64) + mse = float(np.mean((a - b) ** 2)) + if mse == 0.0: + return float("inf") + return float(10.0 * np.log10((255.0**2) / mse)) + + +def _gen(pipe, prompt, *, steps, seed, width, height, guidance): + import torch + + gen = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + image = pipe( + prompt = prompt, + width = width, + height = height, + num_inference_steps = steps, + guidance_scale = guidance, + generator = gen, + ).images[0] + torch.cuda.synchronize() + return image, time.time() - t0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument("--base-repo", default = "Tongyi-MAI/Z-Image-Turbo") + p.add_argument("--transformer-class", default = "ZImageTransformer2DModel") + p.add_argument("--pipeline-class", default = "ZImagePipeline") + p.add_argument( + "--prompt", + default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed", + ) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--width", type = int, default = 1024) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--guidance", type = float, default = 0.0) + p.add_argument( + "--mode", default = "default", help = "compile mode: default | max-autotune-no-cudagraphs" + ) + p.add_argument( + "--dynamic", action = "store_true", help = "dynamic=True (default False here for speed)" + ) + p.add_argument("--out-dir", default = "outputs/compile_probe") + args = p.parse_args(argv) + + import torch + import diffusers + from huggingface_hub import hf_hub_download + + out = Path(args.out_dir) + out.mkdir(parents = True, exist_ok = True) + dtype = torch.bfloat16 + + gguf_path = hf_hub_download(args.repo, args.gguf) + print(f"gguf: {gguf_path}", flush = True) + + transformer_cls = getattr(diffusers, args.transformer_class) + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, + config = args.base_repo, + subfolder = "transformer", + ) + pipeline_cls = getattr(diffusers, args.pipeline_class) + pipe = pipeline_cls.from_pretrained(args.base_repo, torch_dtype = dtype, transformer = transformer) + pipe.to("cuda") + print("pipeline loaded on cuda", flush = True) + + # warm the eager path once (allocator / cudnn), then time eager. + _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + eager_img, eager_t = _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + eager_img.save(out / "eager.png") + eager_arr = np.array(eager_img) + print(f"EAGER: {eager_t:.2f}s/gen", flush = True) + + # compile the repeated denoiser block. + fn = getattr(pipe.transformer, "compile_repeated_blocks", None) + if not callable(fn): + print("RESULT: transformer has no compile_repeated_blocks -> N/A", flush = True) + return 3 + compile_kwargs = {"fullgraph": True, "dynamic": bool(args.dynamic)} + if args.mode and args.mode != "default": + compile_kwargs["mode"] = args.mode + print(f"compiling repeated blocks: {compile_kwargs} ...", flush = True) + try: + t0 = time.time() + fn(**compile_kwargs) + print( + f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: compile_repeated_blocks RAISED: {type(exc).__name__}: {exc}", flush = True) + return 1 + + # first compiled gen triggers the actual compilation (untimed warmup). + try: + t0 = time.time() + _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + print(f" first compiled gen (compilation) took {time.time()-t0:.1f}s", flush = True) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: first compiled generation RAISED: {type(exc).__name__}: {exc}", flush = True) + return 2 + + comp_img, comp_t = _gen( + pipe, + args.prompt, + steps = args.steps, + seed = args.seed, + width = args.width, + height = args.height, + guidance = args.guidance, + ) + comp_img.save(out / "compiled.png") + psnr = _psnr(eager_arr, np.array(comp_img)) + + speedup = (eager_t - comp_t) / eager_t * 100.0 + print("\n==== COMPILE PROBE RESULT ====", flush = True) + print(f" eager: {eager_t:.2f}s/gen", flush = True) + print(f" compiled: {comp_t:.2f}s/gen ({speedup:+.1f}% vs eager)", flush = True) + print(f" PSNR(compiled vs eager): {psnr:.1f} dB", flush = True) + print( + f" verdict: {'COMPILE-WORKS' if psnr >= 30 else 'COMPILE-DIVERGES'} " + f"{'FASTER' if comp_t < eager_t else 'NOT-FASTER'}", + flush = True, + ) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index dc49fe756b..89f59a688a 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -216,6 +216,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: hf_token = os.environ.get("HF_TOKEN"), cpu_offload = args.cpu_offload, memory_mode = args.memory_mode, + speed_mode = args.speed_mode, text_encoder_quant = args.text_encoder_quant, ) _wait_for_load(backend) @@ -295,6 +296,7 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "seed": args.seed, "batch_size": args.batch_size, "memory_mode": args.memory_mode, + "speed_mode": args.speed_mode, "cpu_offload": args.cpu_offload, "text_encoder_quant": args.text_encoder_quant, }, @@ -450,6 +452,13 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["auto", "fast", "balanced", "low_vram"], help = "memory policy (default: backend auto)", ) + p.add_argument( + "--speed-mode", + default = None, + choices = ["off", "default", "max"], + help = "speed profile: off is bit-identical; default adds compile + " + "cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV", + ) p.add_argument( "--text-encoder-quant", default = None, diff --git a/scripts/leverage_probe.py b/scripts/leverage_probe.py new file mode 100644 index 0000000000..756fe25e34 --- /dev/null +++ b/scripts/leverage_probe.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe two candidate levers from the optimization research, on the real GGUF path: + + * `coordinate_descent_tuning` (Inductor) -- lossless extra kernel autotuning. + * FirstBlockCache (diffusers `apply_first_block_cache`) -- step-skip cache, lossy, + evaluated at a low (8) step count where its ceiling is lower. + +Each config is a fresh pipeline load (so Inductor config / compile artifacts don't +cross-contaminate). Reports latency + PSNR vs the eager reference. Run on one CUDA GPU. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +REPO = "unsloth/Z-Image-Turbo-GGUF" +GGUF = "z-image-turbo-Q4_K_M.gguf" +BASE = "Tongyi-MAI/Z-Image-Turbo" +PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed" + + +def _psnr(a, b): + mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)) + return float("inf") if mse == 0 else float(10 * np.log10(255.0**2 / mse)) + + +def _load(): + import torch + import diffusers + from huggingface_hub import hf_hub_download + + t = diffusers.ZImageTransformer2DModel.from_single_file( + hf_hub_download(REPO, GGUF), + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = torch.bfloat16), + torch_dtype = torch.bfloat16, + config = BASE, + subfolder = "transformer", + ) + pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) + pipe.to("cuda") + return pipe + + +def _gen(pipe, steps, seed, res): + import torch + + g = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + img = pipe( + prompt = PROMPT, + width = res, + height = res, + num_inference_steps = steps, + guidance_scale = 0.0, + generator = g, + ).images[0] + torch.cuda.synchronize() + return img, time.time() - t0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + args = p.parse_args(argv) + steps, res, seed = args.steps, args.res, args.seed + + import torch + + def compile_blocks(pipe, *, cdt = False): + if cdt: + import torch._inductor.config as ic + ic.coordinate_descent_tuning = True + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) + + def run( + tag, + *, + compile = False, + cdt = False, + fbc = None, + ): + # reset inductor config between runs + import torch._inductor.config as ic + + ic.coordinate_descent_tuning = False + torch.compiler.reset() + pipe = _load() + if fbc is not None: + from diffusers.hooks import FirstBlockCacheConfig, apply_first_block_cache + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbc)) + if compile: + compile_blocks(pipe, cdt = cdt) + _gen(pipe, steps, seed, res) # warmup / compilation + else: + _gen(pipe, steps, seed, res) # allocator warmup + img, dt = _gen(pipe, steps, seed, res) + del pipe + torch.cuda.empty_cache() + return tag, np.array(img), dt + + results = [] + print(f"== leverage probe (Z-Image Q4_K_M, {res}px, {steps} steps) ==", flush = True) + _, eager, eager_t = run("eager") + print(f" eager: {eager_t:.3f}s", flush = True) + results.append(("eager", eager_t, 0.0)) + + for tag, kw in [ + ("compile(default)", dict(compile = True)), + ("compile+coord_desc", dict(compile = True, cdt = True)), + ("fbc0.12+compile", dict(compile = True, fbc = 0.12)), + ("fbc0.20+compile", dict(compile = True, fbc = 0.20)), + ("fbc0.20(no compile)", dict(fbc = 0.20)), + ]: + try: + t, img, dt = run(tag, **kw) + ps = _psnr(eager, img) + results.append((tag, dt, ps)) + print( + f" {tag:22s} {dt:.3f}s ({(eager_t-dt)/eager_t*100:+.0f}% vs eager) PSNR={ps:.1f} dB", + flush = True, + ) + except Exception as exc: # noqa: BLE001 + print(f" {tag:22s} FAILED: {type(exc).__name__}: {str(exc)[:140]}", flush = True) + + print("\n==== SUMMARY ====", flush = True) + for tag, dt, ps in results: + sp = f"{(results[0][1]-dt)/results[0][1]*100:+.0f}%" if tag != "eager" else "ref" + pss = f"{ps:.1f}dB" if ps else "ref" + print(f" {tag:24s} {dt:.3f}s {sp:>6s} {pss:>8s}", flush = True) + print("LEVERAGE-PROBE-DONE", flush = True) + return 0 + + +if __name__ == "__main__": + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend")) + sys.exit(main()) diff --git a/scripts/perf_verify.py b/scripts/perf_verify.py new file mode 100644 index 0000000000..d3c8330239 --- /dev/null +++ b/scripts/perf_verify.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""GPU verification for the diffusion performance pass (Phase 7). + +Drives the real ``DiffusionBackend`` through several loads in one process and +checks, at a fixed seed: + + 1. speed: ``default`` (compile + cudnn.benchmark + channels_last) vs ``off`` + -- expect a large denoise speedup at high PSNR (near-lossless). + 2. the TF32-leak fix: load ``max`` (flips global TF32 / cudnn.benchmark), unload, + then load ``off`` -- the ``off`` image must be byte-identical (PSNR inf) to a + fresh ``off`` baseline, proving the globals were restored on unload. + 3. ``balanced`` is now bit-identical: with VAE tiling restricted to the low tiers, + streamed (group) offload should match the resident image (PSNR inf). + +Run on one CUDA GPU with the GGUF + base repo cached. +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + +_BACKEND_ROOT = Path(__file__).resolve().parent.parent / "studio" / "backend" +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +def _psnr(a: "np.ndarray", b: "np.ndarray") -> float: + a = a.astype(np.float64) + b = b.astype(np.float64) + mse = float(np.mean((a - b) ** 2)) + return float("inf") if mse == 0.0 else float(10.0 * np.log10((255.0**2) / mse)) + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--model", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument( + "--prompt", + default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed", + ) + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--width", type = int, default = 1024) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--out-dir", default = "outputs/perf_verify") + args = p.parse_args(argv) + + import os + + import torch + from core.inference.diffusion import DiffusionBackend + + out = Path(args.out_dir) + out.mkdir(parents = True, exist_ok = True) + backend = DiffusionBackend() + token = os.environ.get("HF_TOKEN") + + def load(mode_speed = None, mode_mem = None): + backend.begin_load( + args.model, + gguf_filename = args.gguf, + hf_token = token, + speed_mode = mode_speed, + memory_mode = mode_mem, + ) + deadline = time.time() + 2400 + while time.time() < deadline: + ph = backend.load_progress().get("phase") + if ph == "ready": + return backend.status() + if ph == "error": + raise RuntimeError(f"load error: {backend.load_progress()}") + time.sleep(0.5) + raise RuntimeError("load timed out") + + def gen(): + torch.cuda.synchronize() + t0 = time.time() + img = backend.generate( + prompt = args.prompt, + width = args.width, + height = args.height, + steps = args.steps, + guidance = 0.0, + seed = args.seed, + batch_size = 1, + )["images"][0] + torch.cuda.synchronize() + return img, time.time() - t0 + + def timed( + mode_speed, + *, + warmup, + iters, + mem = None, + tag = "", + ): + st = load(mode_speed, mem) + for _ in range(warmup): + gen() + lats = [] + img = None + for _ in range(iters): + img, dt = gen() + lats.append(dt) + img.save(out / f"{tag}.png") + backend.unload() + med = sorted(lats)[len(lats) // 2] + print( + f" [{tag}] speed={mode_speed} mem={mem} optims={st.get('speed_optims')} " + f"tiling={st.get('vae_tiling')} median={med:.3f}s", + flush = True, + ) + return np.array(img), med + + print("== 1. speed: off vs default ==", flush = True) + off_img, off_t = timed("off", warmup = 1, iters = 3, tag = "off") + def_img, def_t = timed("default", warmup = 1, iters = 3, tag = "default") + print(f" PSNR(default vs off) = {_psnr(off_img, def_img):.1f} dB", flush = True) + print( + f" speedup: off {off_t:.3f}s -> default {def_t:.3f}s " + f"({(off_t-def_t)/off_t*100:+.1f}%)", + flush = True, + ) + + print("== 2. TF32-leak fix: max then off must be byte-identical ==", flush = True) + timed("max", warmup = 0, iters = 1, tag = "max") # flips + should restore globals + off2_img, _ = timed("off", warmup = 0, iters = 1, tag = "off2") + leak_psnr = _psnr(off_img, off2_img) + print( + f" PSNR(off-after-max vs off) = {leak_psnr:.1f} dB " + f"({'OK byte-identical' if leak_psnr == float('inf') else 'LEAK! globals not restored'})", + flush = True, + ) + + print("== 3. balanced is bit-identical (tiling off) ==", flush = True) + bal_img, bal_t = timed("off", warmup = 0, iters = 1, mem = "balanced", tag = "balanced") + bal_psnr = _psnr(off_img, bal_img) + print( + f" PSNR(balanced vs off) = {bal_psnr:.1f} dB " + f"({'OK bit-identical' if bal_psnr == float('inf') else 'differs'})", + flush = True, + ) + + ok = ( + (leak_psnr == float("inf")) + and (bal_psnr == float("inf")) # check 3: balanced must be bit-identical to off + and (def_t < off_t) + and (_psnr(off_img, def_img) >= 30) + ) + print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8ad5ffa250..9500f2ed9f 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -44,7 +44,13 @@ from .diffusion_memory import ( plan_diffusion_memory, snapshot_device_memory, ) -from .diffusion_speed import SPEED_OFF, apply_speed_optims +from .diffusion_speed import ( + SPEED_OFF, + apply_speed_optims, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, +) from .diffusion_precision import quantize_text_encoders logger = get_logger(__name__) @@ -69,6 +75,10 @@ class _LoadState: # The opt-in speed profile (Phase 3). speed_mode: str = SPEED_OFF speed_optims: tuple = () + # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the + # speed layer mutated them, restored on unload so a later `off` load is not + # contaminated by this one's globals. None when nothing was changed. + backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None @@ -495,56 +505,77 @@ class DiffusionBackend: pipeline_cls = getattr(diffusers, fam.pipeline_class) pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs) + # Resolve the effective speed mode: GGUF models default to the + # near-lossless `default` profile (compile is ~2.2x and sits below + # the quant noise floor), dense models stay bit-identical `off`. An + # explicit speed_mode (incl. "off") is honored verbatim. + effective_speed = resolve_speed_mode(speed_mode, is_gguf = bool(gguf_filename)) # Opt-in speed optims run BEFORE placement (channels_last / compile - # must precede CPU offload). Off by default -> bit-identical output. - speed_applied = apply_speed_optims( - pipe, - target, - is_gguf = bool(gguf_filename), - family = fam, - speed_mode = speed_mode or SPEED_OFF, - logger = logger, - ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), - # also before placement so the offload hooks move the smaller weights. - te_quant = quantize_text_encoders( - pipe, - target, - mode = text_encoder_quant, - logger = logger, - ) + # must precede CPU offload). Snapshot the process-wide backend flags + # first so unload can restore them: TF32 / cudnn.benchmark are global, + # and a later `off` load must not inherit this load's settings. + backend_flags_before = snapshot_backend_flags() + # apply_speed_optims mutates PROCESS-WIDE flags (TF32 / cudnn.benchmark); + # they are only restored via _LoadState.backend_flags_before on unload. If + # the build fails after this but before _state commits (e.g. an OOM in + # apply_memory_plan / pipe.to), nothing would restore them and a later `off` + # generation would be contaminated, so restore on any non-committed exit. + committed = False + try: + speed_applied = apply_speed_optims( + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = effective_speed, + logger = logger, + ) + # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), + # also before placement so the offload hooks move the smaller weights. + te_quant = quantize_text_encoders( + pipe, + target, + mode = text_encoder_quant, + logger = logger, + ) - # Decide placement from MEASURED free device memory vs the model's - # estimated resident size (transformer GGUF dequantised + the - # companion text-encoder / VAE already cached for `base`), then - # apply it. Computed here, after the build but before placement, - # because the weights are still on CPU so free VRAM is the real - # budget. `cpu_offload=True` stays an explicit override. - plan = self._plan_memory( - target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload - ) - # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it - # may fall back to whole-module offload, and tiling is a no-op on a - # pipeline with no tiling control), so status stays honest. - effective_policy, effective_tiling = apply_memory_plan( - pipe, plan, device = device, logger = logger - ) + # Decide placement from MEASURED free device memory vs the model's + # estimated resident size (transformer GGUF dequantised + the + # companion text-encoder / VAE already cached for `base`), then + # apply it. Computed here, after the build but before placement, + # because the weights are still on CPU so free VRAM is the real + # budget. `cpu_offload=True` stays an explicit override. + plan = self._plan_memory( + target, gguf_path, gguf_filename, base, fam, memory_mode, cpu_offload + ) + # apply_memory_plan returns the (policy, tiling) ACTUALLY engaged (it + # may fall back to whole-module offload, and tiling is a no-op on a + # pipeline with no tiling control), so status stays honest. + effective_policy, effective_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) - self._state = _LoadState( - pipe = pipe, - family = fam, - repo_id = repo_id, - base_repo = base, - device = device, - dtype = str(dtype).replace("torch.", ""), - cpu_offload = effective_policy != OFFLOAD_NONE, - offload_policy = effective_policy, - vae_tiling = effective_tiling, - memory_mode = plan.requested_mode, - speed_mode = (speed_mode or SPEED_OFF), - speed_optims = tuple(k for k, v in speed_applied.items() if v), - text_encoder_quant = te_quant, - ) + self._state = _LoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + cpu_offload = effective_policy != OFFLOAD_NONE, + offload_policy = effective_policy, + vae_tiling = effective_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(k for k, v in speed_applied.items() if v), + backend_flags_before = backend_flags_before, + text_encoder_quant = te_quant, + ) + committed = True + finally: + if not committed: + restore_backend_flags(backend_flags_before) + clear_gpu_cache() logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", @@ -679,7 +710,10 @@ class DiffusionBackend: self._gen = gen try: - images = state.pipe(**kwargs).images + # inference_mode is strictly faster than the no_grad diffusers + # uses internally and numerically identical for inference. + with torch.inference_mode(): + images = state.pipe(**kwargs).images finally: self._gen = None # A cancelled denoise returns early with a partial/garbage image; @@ -738,6 +772,9 @@ class DiffusionBackend: state = self._state if state is None: return + # Restore the process-wide backend flags (TF32 / cudnn.benchmark) this load + # may have flipped, so the next `off` load is bit-identical again. + restore_backend_flags(state.backend_flags_before) self._state = None del state clear_gpu_cache() diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 4fecf950ad..4a3dd616ac 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -35,9 +35,9 @@ class DiffusionFamily: # (~6.5e4) and produce inf -> NaN latents -> a black image. The backend # promotes a resolved float16 to float32 for these at load time. fp16_incompatible: bool = False - # False for families whose denoiser block doesn't compile cleanly with - # regional torch.compile (Z-Image). Only consulted on the non-GGUF path; the - # GGUF transformer is never compiled regardless. + # Set False only for a family whose denoiser block does not compile cleanly with + # regional torch.compile. Now consulted on the GGUF path too (compile runs on the + # GGUF transformer); all current families compile, so this stays True. supports_torch_compile: bool = True @@ -81,8 +81,6 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("zimage", "z_image"), # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, - # Z-Image's denoiser block is excluded from regional torch.compile. - supports_torch_compile = False, ), ) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index ca90e11c73..9172ef7cf8 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -400,17 +400,21 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("explicit cpu_offload overrides resident placement") - # VAE tiling/slicing decode the image in chunks, capping the decode-time spike - # that often dominates peak VRAM at high resolution. Turn it on whenever weights - # are being offloaded (the device is already tight) or the backend has no spare - # device pool (MPS/CPU). On a roomy discrete GPU it stays off so output is - # bit-identical to a plain resident run. - tile = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu") + # VAE savers cap the decode-time spike that dominates peak VRAM at high res. + # Slicing (decode a batch one image at a time) is EXACT, so enable it on any + # offload tier / non-discrete backend. Tiling (spatial chunks) is only bit- + # identical for a single tile (<=1MP), so restrict it to the lowest tiers where + # the VAE itself is offloaded (model / sequential) or there is no spare device + # pool (MPS / CPU). Under group offload the transformer streams but the VAE stays + # resident and fits, so it keeps exact full-image decode -> balanced is both + # faster and bit-identical. On a roomy discrete GPU both stay off. + any_offload = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu") + tile = policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL) or device_memory.backend in ("mps", "cpu") return MemoryPlan( requested_mode = mode, offload_policy = policy, vae_tiling = tile, - vae_slicing = tile, + vae_slicing = any_offload, device_memory = device_memory, estimates = estimates, reasons = tuple(reasons), @@ -442,12 +446,22 @@ def apply_memory_plan( if plan.vae_slicing: _enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger) + def _fallback_to_model_offload() -> None: + # Group offload keeps the VAE resident, so the GROUP plan set vae_tiling=False. + # When group offload is unavailable and we drop to whole-module offload, the card + # is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE + # tiling on now (if not already engaged) to cap it. + nonlocal tiling_engaged + pipe.enable_model_cpu_offload() + if not tiling_engaged: + tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger) + policy = plan.offload_policy if policy == OFFLOAD_MODEL: pipe.enable_model_cpu_offload() elif policy == OFFLOAD_GROUP: if not _apply_group_offload(pipe, device, logger): - pipe.enable_model_cpu_offload() + _fallback_to_model_offload() policy = OFFLOAD_MODEL elif policy == OFFLOAD_SEQUENTIAL: try: @@ -459,7 +473,7 @@ def apply_memory_plan( "falling back to whole-module offload", exc, ) - pipe.enable_model_cpu_offload() + _fallback_to_model_offload() policy = OFFLOAD_MODEL else: pipe.to(device) @@ -491,32 +505,40 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: if transformer is None: return False try: + import inspect + import torch from diffusers.hooks import apply_group_offloading onload = torch.device(device) use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA - # Place the smaller components resident FIRST: moving them onto the device is - # the only step here that can OOM on a tight GPU. Doing it before - # apply_group_offloading means a failure leaves NO group-offload hooks on the - # transformer, so the caller's whole-module-offload fallback gets a clean - # pipeline -- diffusers refuses enable_model_cpu_offload() while group hooks - # are attached, which would otherwise turn the fallback into a hard crash. + gkwargs: dict[str, Any] = { + "onload_device": onload, + "offload_device": torch.device("cpu"), + "offload_type": "block_level", + "num_blocks_per_group": DEFAULT_GROUP_BLOCKS, + "use_stream": use_stream, + } + # On the CUDA stream path, overlap each block's host->device copy with + # compute: non_blocking issues the copy asynchronously and record_stream + # defers the free until the copy's stream is done. Lossless (only transfer + # scheduling changes). Safe for the group tier specifically, where the + # companions stay resident; gated on the installed signature so an older + # diffusers that lacks these kwargs still works (no hard fallback). + if use_stream: + _params = inspect.signature(apply_group_offloading).parameters + if "non_blocking" in _params: + gkwargs["non_blocking"] = True + if "record_stream" in _params: + gkwargs["record_stream"] = True + apply_group_offloading(transformer, **gkwargs) + # Place the remaining (smaller) components resident; the streamed + # transformer manages its own placement via the offloading hooks. for name, comp in getattr(pipe, "components", {}).items(): if name == "transformer": continue if isinstance(comp, torch.nn.Module): comp.to(onload) - # Stream the transformer a few blocks at a time; it manages its own placement - # via the offloading hooks. - apply_group_offloading( - transformer, - onload_device = onload, - offload_device = torch.device("cpu"), - offload_type = "block_level", - num_blocks_per_group = DEFAULT_GROUP_BLOCKS, - use_stream = use_stream, - ) return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index e01ba48b63..9184f29f88 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -5,19 +5,29 @@ Off by default, so the default render path stays bit-identical to a plain run (the property the regression harness checks). When the operator opts in, this applies the -lossless-to-near-lossless speedups in the order the diffusers guides recommend -(channels_last -> regional compile, with TF32 / fused-QKV under "max"): +near-lossless speedups in the order the diffusers guides recommend +(channels_last + cudnn.benchmark -> regional compile, with TF32 / fused-QKV under +"max"): - off - nothing (default). - default - lossless: channels_last VAE memory format + regional torch.compile of - the denoiser's repeated block WHERE eligible (non-GGUF, bf16, CUDA, and - a compile-friendly family). + off - nothing (default; bit-identical reference). + default - near-lossless: channels_last VAE memory format + cudnn.benchmark conv + autotune + regional torch.compile of the denoiser's repeated block WHERE + eligible (bf16, CUDA, a compile-friendly family). Compile is the big win + (~2.3x denoise on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, + well above the Q4 quantisation noise floor, so it does not meaningfully + move output quality). max - default plus near-lossless TF32 matmul and fused QKV projections. -Regional compile is gated off for the GGUF transformer (it dequantises per-op and -doesn't compile cleanly) and for families flagged not compile-friendly (Z-Image), so -on today's GGUF path only channels_last / TF32 engage; the compile path activates -automatically once a non-GGUF bf16 transformer is loaded. torch is imported lazily. +Regional compile used to be gated off for the GGUF transformer, but it compiles and +runs faster on the current diffusers/torch (measured; the GGUF dequant ops stay +eager and the rest of the repeated block compiles), so the GGUF gate is removed; the +per-family ``supports_torch_compile`` flag and the bf16/CUDA checks still apply. + +The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so +``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the +prior values at load and restore them at unload, keeping a later ``off`` load +bit-identical instead of inheriting a previous ``max`` run's globals. torch is +imported lazily. """ from __future__ import annotations @@ -30,6 +40,54 @@ SPEED_MAX = "max" SPEED_MODES = (SPEED_OFF, SPEED_DEFAULT, SPEED_MAX) +def snapshot_backend_flags() -> Optional[dict]: + """Capture the process-wide torch backend flags this layer may mutate, so the + caller can restore them on unload. None if torch is unavailable. Each flag is read + defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still + captures the rest -- otherwise a single missing attribute would skip the whole + snapshot and a real mutated flag would leak.""" + try: + import torch + except Exception: # noqa: BLE001 — no torch -> nothing to snapshot/restore + return None + state: dict[str, bool] = {} + matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None) + if matmul is not None and hasattr(matmul, "allow_tf32"): + state["matmul_tf32"] = bool(matmul.allow_tf32) + cudnn = getattr(torch.backends, "cudnn", None) + if cudnn is not None: + if hasattr(cudnn, "allow_tf32"): + state["cudnn_tf32"] = bool(cudnn.allow_tf32) + if hasattr(cudnn, "benchmark"): + state["cudnn_benchmark"] = bool(cudnn.benchmark) + return state + + +def restore_backend_flags(state: Optional[dict]) -> None: + """Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each + flag is restored independently so one failure can't leave the others leaked.""" + if not state: + return + try: + import torch + except Exception: # noqa: BLE001 — no torch -> nothing to restore + return + + def _set(obj: Any, attr: str, key: str) -> None: + if obj is not None and key in state and hasattr(obj, attr): + try: + setattr(obj, attr, state[key]) + except Exception: # noqa: BLE001 — best-effort per-flag restore + pass + + _set( + getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32" + ) + cudnn = getattr(torch.backends, "cudnn", None) + _set(cudnn, "allow_tf32", "cudnn_tf32") + _set(cudnn, "benchmark", "cudnn_benchmark") + + def normalize_speed_mode(value: Optional[str]) -> str: """Lower/strip a requested speed mode (dashes ok); None / "" -> off.""" if value is None: @@ -44,14 +102,29 @@ def normalize_speed_mode(value: Optional[str]) -> str: return normalized +def resolve_speed_mode(value: Optional[str], *, is_gguf: bool) -> str: + """The effective speed mode when the caller leaves it UNSET (``None``). + + A GGUF model defaults to ``default``: regional compile is ~2.2x faster and its + numeric perturbation sits well below the quantisation noise floor (measured + PSNR ~37 dB compile-vs-eager versus ~21 dB Q4-vs-bf16), so it does not reduce + output quality relative to the dense reference. A dense (non-GGUF) model stays + ``off`` / bit-identical, since there compile would be the only source of drift. + An explicit value -- including ``"off"`` -- is always honored verbatim.""" + if value is None: + return SPEED_DEFAULT if is_gguf else SPEED_OFF + return normalize_speed_mode(value) + + def compile_eligible(target: Any, *, is_gguf: bool, family: Any) -> bool: """Whether the denoiser's repeated block should be regionally compiled. - Only on CUDA (incl. ROCm via supports_default_torch_compile), for a non-GGUF - bf16 transformer, on a compile-friendly family. The GGUF transformer is never - compiled (it dequantises per-op).""" - if is_gguf: - return False + Only on CUDA (incl. ROCm via supports_default_torch_compile), for a bf16 + transformer, on a compile-friendly family. ``is_gguf`` no longer disqualifies: + ``compile_repeated_blocks`` runs fine on the GGUF transformer (the per-op + dequant stays eager, the rest of the block compiles) and is ~2.3x faster, so it + is kept only for signature/logging compatibility.""" + del is_gguf # GGUF is compile-eligible now; param kept for call-site compat. if not bool(getattr(target, "supports_default_torch_compile", False)): return False if not bool(getattr(family, "supports_torch_compile", True)): @@ -79,7 +152,13 @@ def apply_speed_optims( """Apply the opt-in speed optimisations for ``speed_mode`` to a built pipeline, BEFORE placement / offload. Returns which optimisations actually engaged. Every step is best-effort: a pipeline that doesn't support one is simply skipped.""" - applied = {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False} + applied = { + "channels_last": False, + "cudnn_benchmark": False, + "tf32": False, + "fused_qkv": False, + "compiled": False, + } mode = normalize_speed_mode(speed_mode) # TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this # load isn't max, so a later default/off diffusion load -- or chat inference in the @@ -93,9 +172,16 @@ def apply_speed_optims( # Lossless: a channels-last VAE speeds up its convolutions with no numeric change. applied["channels_last"] = _vae_channels_last(pipe, logger) - # Lossless-ish: regional compile of the repeated denoiser block, where eligible. + # Near-lossless: let cuDNN autotune the fixed-shape VAE convs (CUDA only). It may + # pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win. + if getattr(target, "device", None) == "cuda": + applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) + + # Near-lossless and the largest win: regional compile of the repeated denoiser + # block, where eligible (now incl. the GGUF transformer). `max` opts into + # max-autotune (longer compile, autotuned kernels). if compile_eligible(target, is_gguf = is_gguf, family = family): - applied["compiled"] = _compile_repeated_blocks(pipe, logger) + applied["compiled"] = _compile_repeated_blocks(pipe, logger, max_autotune = mode == SPEED_MAX) if mode == SPEED_MAX: # Near-lossless: TF32 matmul (CUDA only) trades a few mantissa bits for speed. @@ -119,19 +205,43 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False -def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool: +def _compile_repeated_blocks( + pipe: Any, + logger: Any, + *, + max_autotune: bool = False, +) -> bool: transformer = getattr(pipe, "transformer", None) fn = getattr(transformer, "compile_repeated_blocks", None) if not callable(fn): return False + # default: mode="default" + dynamic=True -- fast cold start, robust to resolution + # changes (no recompile). max: mode="max-autotune-no-cudagraphs" + dynamic=False -- + # Triton autotuning for a few % more on GEMM/conv-heavy models, at a much longer + # compile and a recompile per new resolution. The CUDA-graph modes (reduce-overhead + # / max-autotune) are deliberately NOT used: they crash on the regionally-compiled + # block because its static output buffer is overwritten across denoise steps. + kwargs: dict[str, Any] = {"fullgraph": True, "dynamic": not max_autotune} + if max_autotune: + kwargs["mode"] = "max-autotune-no-cudagraphs" try: - fn(fullgraph = True, dynamic = True) + fn(**kwargs) return True except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) return False +def _enable_cudnn_benchmark(logger: Any) -> bool: + try: + import torch + torch.backends.cudnn.benchmark = True + return True + except Exception as exc: # noqa: BLE001 — optimisation only + _warn(logger, "cudnn_benchmark", exc) + return False + + # The TF32 flag values from before the first max load flipped them, so a later # non-max load / unload can put the process back exactly as it found it (rather than # forcing a hardcoded default that might clobber another component's choice). diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 44f57fb989..8069e8f610 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -118,6 +118,34 @@ class SdCppUpscaleParams: tile_size: Optional[int] = None +# Native (sd.cpp) speed profiles, the engine-side analogue of diffusion_speed's +# modes. off: nothing (default). default: --diffusion-fa (flash attention; upstream +# reports it usually speeds CUDA and cuts attention memory, near-lossless). max: also +# --diffusion-conv-direct (direct conv; helps some backends, but measured +45% on +# CUDA here, so it stays opt-in/experimental, never auto-on for CUDA). +NATIVE_SPEED_OFF = "off" +NATIVE_SPEED_DEFAULT = "default" +NATIVE_SPEED_MAX = "max" +NATIVE_SPEED_MODES = (NATIVE_SPEED_OFF, NATIVE_SPEED_DEFAULT, NATIVE_SPEED_MAX) + + +def native_speed_flags(speed_mode: Optional[str]) -> list[str]: + """sd-cli speed flags for a native speed mode (empty for off / None). + + These are separate from the offload flags: ``--diffusion-fa`` is a speed/memory + win in its own right, not tied to whether weights are offloaded. De-duplicated + against offload flags at the call site (offload already adds ``--diffusion-fa``). + """ + mode = (speed_mode or NATIVE_SPEED_OFF).strip().lower() + if mode in ("", NATIVE_SPEED_OFF): + return [] + if mode == NATIVE_SPEED_DEFAULT: + return ["--diffusion-fa"] + if mode == NATIVE_SPEED_MAX: + return ["--diffusion-fa", "--diffusion-conv-direct"] + raise ValueError(f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'") + + def offload_flags( policy: str, *, diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index b1daafce72..186ea51631 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -24,9 +24,11 @@ from __future__ import annotations import logging import os +import queue import shutil import subprocess import sys +import threading import time from pathlib import Path from typing import Callable, Optional @@ -37,6 +39,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, ) logger = logging.getLogger(__name__) @@ -189,6 +192,7 @@ class SdCppEngine: *, output_path: str, offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, threads: Optional[int] = None, verbose: bool = False, extra_args: Optional[list[str]] = None, @@ -198,10 +202,15 @@ class SdCppEngine: ) -> Path: """Run one ``sd-cli`` generation; return the written image path. - Raises ``RuntimeError`` if the binary is missing, the process exits - nonzero, or no output file is produced. ``on_log`` (if given) receives - each line of sd-cli's progress output as it arrives. + ``native_speed`` ("default"/"max") adds sd.cpp's own speed flags + (``--diffusion-fa`` etc.), de-duplicated against the offload flags that may + already include them. Raises ``RuntimeError`` if the binary is missing, the + process exits nonzero, or no output file is produced. ``on_log`` (if given) + receives each line of sd-cli's progress output as it arrives. """ + offload = list(offload or []) + speed = [f for f in native_speed_flags(native_speed) if f not in offload] + merged_extra = speed + list(extra_args or []) cmd = build_sd_cpp_command( self._require_binary(), files, @@ -210,7 +219,7 @@ class SdCppEngine: offload = offload, threads = threads, verbose = verbose, - extra_args = extra_args, + extra_args = merged_extra, ) return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) @@ -281,20 +290,49 @@ class SdCppEngine: errors = "replace", env = run_env, ) + # Drain stdout on a reader thread so the timeout is enforced even when the + # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain + # `for line in proc.stdout` blocks until EOF, so proc.wait(timeout) would + # never be reached. The reader pushes lines (then a None sentinel at EOF) to a + # queue the main loop polls against a wall-clock deadline. tail: list[str] = [] + line_q: "queue.Queue[Optional[str]]" = queue.Queue() + + def _drain() -> None: + try: + assert proc.stdout is not None + for raw in proc.stdout: + line_q.put(raw.rstrip("\n")) + finally: + line_q.put(None) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + + deadline = None if timeout is None else time.monotonic() + float(timeout) + stdout_done = False try: - assert proc.stdout is not None - for line in proc.stdout: - line = line.rstrip("\n") + while True: + if deadline is not None and time.monotonic() >= deadline and proc.poll() is None: + proc.kill() + raise RuntimeError(f"sd-cli timed out after {timeout}s") + try: + line = line_q.get(timeout = 0.1) + except queue.Empty: + if proc.poll() is not None and stdout_done: + break + continue + if line is None: + stdout_done = True + if proc.poll() is not None: + break + continue tail.append(line) if len(tail) > 40: tail.pop(0) if on_log is not None: on_log(line) - ret = proc.wait(timeout = timeout) - except subprocess.TimeoutExpired: - proc.kill() - raise RuntimeError(f"sd-cli timed out after {timeout}s") + ret = proc.wait(timeout = 5.0) finally: if proc.poll() is None: proc.kill() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index c7c4736502..fe4cabe8b5 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -10,6 +10,7 @@ GPU, weights, or network access is needed (sub-second, CI-friendly). from __future__ import annotations +import contextlib import sys import types @@ -202,6 +203,8 @@ def fake_runtime(monkeypatch): torch.Generator = _FakeGenerator torch.cuda = types.SimpleNamespace(is_available = lambda: False) torch.backends = types.SimpleNamespace(mps = None) + # generate() wraps the pipe call in torch.inference_mode(); a no-op CM here. + torch.inference_mode = lambda: contextlib.nullcontext() diffusers = types.ModuleType("diffusers") diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) @@ -856,12 +859,19 @@ def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( assert status["offload_policy"] == "model" and status["cpu_offload"] is True -def test_load_speed_mode_threads_and_defaults_off(fake_runtime, tmp_path): - # No speed_mode -> off, no optimisations engaged (the bit-identical default). +def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path): + # No speed_mode on a GGUF model -> auto `default` (near-lossless, compile sits + # below the quant noise floor). compile itself only engages on CUDA, so on this + # CPU stub no optim need engage, but the resolved mode is `default`. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") - assert status["speed_mode"] == "off" and status["speed_optims"] == [] + assert status["speed_mode"] == "default" + # An explicit "off" opts back into the bit-identical path (engages nothing). + status_off = backend.load_pipeline( + str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "off" + ) + assert status_off["speed_mode"] == "off" and status_off["speed_optims"] == [] # An explicit speed_mode threads through to status (engaged optims are GPU-verified). status2 = backend.load_pipeline( str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "max" diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 858e87ffa1..bc09ffbd0f 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -185,6 +185,9 @@ def test_auto_group_offload_when_transformer_overflows_but_companions_fit(): base_overhead_mib = 1000, ) assert plan.offload_policy == OFFLOAD_GROUP + # Group keeps the VAE resident, so it uses exact slicing but NOT lossy tiling + # -> balanced stays bit-identical while still capping the offload footprint. + assert plan.vae_slicing is True and plan.vae_tiling is False def test_auto_model_offload_when_companions_exceed_budget(): @@ -432,6 +435,18 @@ def test_apply_group_falls_back_to_model_without_transformer(): assert effective == OFFLOAD_MODEL and "model_offload" in pipe.calls +def test_apply_group_fallback_enables_vae_tiling(): + # A balanced/group plan keeps the VAE resident (tiling off); when group offload can't + # engage and we drop to whole-module offload, the applier must turn VAE tiling ON to + # cap the decode-time spike on what is now a low-VRAM path. + plan = _plan(OFFLOAD_GROUP, tiling = True) + assert plan.vae_tiling is False # group plan leaves tiling off by design + pipe = _RecordingPipe() # no .transformer -> group offload falls back to model + effective, tiled = apply_memory_plan(pipe, plan, device = "cuda") + assert effective == OFFLOAD_MODEL + assert tiled is True and "vae_tiling" in pipe.calls + + def test_apply_sequential_offload(): pipe = _RecordingPipe() effective, _ = apply_memory_plan( diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index ade37482ba..b0ca10af4c 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -21,6 +21,9 @@ from core.inference.diffusion_speed import ( apply_speed_optims, compile_eligible, normalize_speed_mode, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, ) @@ -47,7 +50,7 @@ def _stub_torch(monkeypatch): torch.channels_last = "channels_last" torch.backends = types.SimpleNamespace( cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)), - cudnn = types.SimpleNamespace(allow_tf32 = False), + cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False), ) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -64,23 +67,91 @@ def test_normalize_speed_mode(): normalize_speed_mode("ludicrous") +def test_resolve_speed_mode_gguf_auto_default(): + # Unset (None) -> default for GGUF (near-lossless), off for dense. + assert resolve_speed_mode(None, is_gguf = True) == SPEED_DEFAULT + assert resolve_speed_mode(None, is_gguf = False) == SPEED_OFF + # An explicit value is honored verbatim, including an explicit opt-out to off. + assert resolve_speed_mode("off", is_gguf = True) == SPEED_OFF + assert resolve_speed_mode("max", is_gguf = True) == SPEED_MAX + assert resolve_speed_mode("max", is_gguf = False) == SPEED_MAX + + # ── compile gating ──────────────────────────────────────────────────────────── -def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch): +def test_compile_eligible_requires_bf16_cuda_friendly(monkeypatch): _stub_torch(monkeypatch) - # The happy path: non-GGUF, bf16, CUDA, compile-friendly family. + # The happy path: bf16, CUDA, compile-friendly family. assert compile_eligible(_target(), is_gguf = False, family = _family()) is True - # GGUF is never compiled. - assert compile_eligible(_target(), is_gguf = True, family = _family()) is False + # GGUF is now compile-eligible too (measured ~2.3x, PSNR ~37 dB vs eager). + assert compile_eligible(_target(), is_gguf = True, family = _family()) is True # fp16 (non-bf16) is excluded. assert compile_eligible(_target(dtype = "float16"), is_gguf = False, family = _family()) is False - # A family flagged not compile-friendly (Z-Image) is excluded. + # A family flagged not compile-friendly is excluded. assert compile_eligible(_target(), is_gguf = False, family = _family(compile_ok = False)) is False - # No compile support (e.g. ROCm/XPU/MPS) is excluded. + # No compile support (e.g. XPU/MPS) is excluded. assert compile_eligible(_target(compile_ok = False), is_gguf = False, family = _family()) is False +# ── backend-flag snapshot / restore (TF32 / cudnn.benchmark leak guard) ──────── + + +def test_snapshot_restore_backend_flags(monkeypatch): + torch = _stub_torch(monkeypatch) + snap = snapshot_backend_flags() + assert snap == {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False} + # An opt-in max run flips the globals on... + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + # ...and restore puts them back, so a later `off` load is bit-identical again. + restore_backend_flags(snap) + assert torch.backends.cuda.matmul.allow_tf32 is False + assert torch.backends.cudnn.allow_tf32 is False + assert torch.backends.cudnn.benchmark is False + + +def test_restore_backend_flags_tolerates_none(): + restore_backend_flags(None) # no torch needed, no-op + + +def test_snapshot_partial_when_some_backends_missing(monkeypatch): + # A build/platform without cuda.matmul (e.g. CPU/MPS) must still snapshot + restore the + # flags it does have, rather than skipping the whole snapshot on one missing attribute. + torch = types.ModuleType("torch") + torch.backends = types.SimpleNamespace( + cuda = types.SimpleNamespace(), # no .matmul + cudnn = types.SimpleNamespace(benchmark = True), # no .allow_tf32 + ) + monkeypatch.setitem(sys.modules, "torch", torch) + snap = snapshot_backend_flags() + assert snap == {"cudnn_benchmark": True} + torch.backends.cudnn.benchmark = False + restore_backend_flags(snap) + assert torch.backends.cudnn.benchmark is True + + +def test_restore_is_independent_per_flag(monkeypatch): + # A read-only / failing attribute must not abort restoring the remaining flags. + torch = _stub_torch(monkeypatch) + + class _NoMatmulSet: + @property + def allow_tf32(self): + return False + + @allow_tf32.setter + def allow_tf32(self, value): + raise RuntimeError("read-only on this build") + + torch.backends.cuda.matmul = _NoMatmulSet() + snap = {"matmul_tf32": False, "cudnn_tf32": False, "cudnn_benchmark": False} + torch.backends.cudnn.benchmark = True + restore_backend_flags(snap) # matmul setter raises, cudnn still restored + assert torch.backends.cudnn.benchmark is False + + # ── applier ─────────────────────────────────────────────────────────────────── @@ -103,24 +174,33 @@ class _Pipe: def _vae_to(self, *, memory_format): self.vae.mem_format = memory_format - def _compile(self, *, fullgraph, dynamic): + def _compile(self, **kwargs): self.compiled = True + self.compile_kwargs = kwargs def _fuse(self): self.fused = True def test_speed_off_applies_nothing(monkeypatch): - _stub_torch(monkeypatch) + torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True, with_fuse = True) applied = apply_speed_optims( pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF ) - assert applied == {"channels_last": False, "tf32": False, "fused_qkv": False, "compiled": False} + assert applied == { + "channels_last": False, + "cudnn_benchmark": False, + "tf32": False, + "fused_qkv": False, + "compiled": False, + } assert pipe.vae.mem_format is None and pipe.compiled is False + # off must not touch any process-wide flag (bit-identical reference path). + assert torch.backends.cudnn.benchmark is False -def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch): +def test_speed_default_channels_last_compile_and_cudnn_benchmark(monkeypatch): torch = _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( @@ -128,18 +208,36 @@ def test_speed_default_channels_last_and_compile_when_eligible(monkeypatch): ) assert applied["channels_last"] is True and pipe.vae.mem_format == torch.channels_last assert applied["compiled"] is True and pipe.compiled is True - # default does not flip TF32 or fuse QKV. + # default compiles with dynamic=True and no autotune mode (fast cold start, + # resolution-robust, sidesteps the CUDA-graph crash). + assert pipe.compile_kwargs == {"fullgraph": True, "dynamic": True} + # default also autotunes the VAE convs but does NOT flip TF32 or fuse QKV. + assert applied["cudnn_benchmark"] is True and torch.backends.cudnn.benchmark is True assert applied["tf32"] is False and applied["fused_qkv"] is False -def test_speed_default_skips_compile_for_gguf(monkeypatch): +def test_speed_default_compiles_gguf(monkeypatch): _stub_torch(monkeypatch) pipe = _Pipe(with_compile = True) applied = apply_speed_optims( pipe, _target(), is_gguf = True, family = _family(), speed_mode = SPEED_DEFAULT ) - assert applied["channels_last"] is True # lossless layout still applies - assert applied["compiled"] is False and pipe.compiled is False # GGUF never compiles + assert applied["channels_last"] is True + # GGUF now compiles (the big near-lossless win). + assert applied["compiled"] is True and pipe.compiled is True + + +def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch): + _stub_torch(monkeypatch) + pipe = _Pipe(with_compile = True) + applied = apply_speed_optims( + pipe, + _target(device = "mps", compile_ok = False), + is_gguf = True, + family = _family(), + speed_mode = SPEED_DEFAULT, + ) + assert applied["cudnn_benchmark"] is False # not CUDA -> no autotune flip def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): @@ -150,6 +248,9 @@ def test_speed_max_enables_tf32_and_fused_qkv(monkeypatch): ) assert applied["tf32"] is True and torch.backends.cuda.matmul.allow_tf32 is True assert applied["fused_qkv"] is True and pipe.fused is True + # max opts into autotuned kernels (static shapes); CUDA-graph modes are avoided. + assert pipe.compile_kwargs["mode"] == "max-autotune-no-cudagraphs" + assert pipe.compile_kwargs["dynamic"] is False def test_speed_max_tf32_only_on_cuda(monkeypatch): diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index c33a4a01f9..200672aa26 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -23,6 +23,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, offload_flags, text_encoder_flags_for_family, ) @@ -47,6 +48,16 @@ def test_te_flags_by_family(): # ── offload policy -> sd-cli flags ────────────────────────────────────────── +def test_native_speed_flags(): + assert native_speed_flags(None) == [] + assert native_speed_flags("off") == [] + assert native_speed_flags("") == [] + assert native_speed_flags("default") == ["--diffusion-fa"] + assert native_speed_flags("max") == ["--diffusion-fa", "--diffusion-conv-direct"] + with pytest.raises(ValueError): + native_speed_flags("ludicrous") + + def test_offload_none_is_empty(): assert offload_flags(OFFLOAD_NONE) == [] diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 7aeb537033..a879252761 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -12,6 +12,7 @@ from __future__ import annotations import os import sys +import time import types from pathlib import Path @@ -242,6 +243,55 @@ def test_generate_raises_when_binary_missing(): ) +class _HangingPopen: + """A child that runs but never prints and never exits -- the case a plain + `for line in stdout` would block on forever, ignoring the timeout.""" + + def __init__(self, cmd, **_kw): + self._alive = True + + class _Blocking: + def __init__(self, owner): + self.owner = owner + + def __iter__(self): + return self + + def __next__(self): + while self.owner._alive: + time.sleep(0.01) + raise StopIteration + + @property + def stdout(self): + return self._Blocking(self) + + def poll(self): + return None if self._alive else -9 + + def wait(self, timeout = None): + self._alive = False + return -9 + + def kill(self): + self._alive = False + + +def test_generate_times_out_on_silent_hang(tmp_path, monkeypatch): + e = _engine(tmp_path) + monkeypatch.setattr(eng.subprocess, "Popen", lambda cmd, **kw: _HangingPopen(cmd, **kw)) + t0 = time.time() + with pytest.raises(RuntimeError, match = "timed out"): + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(tmp_path / "x.png"), + timeout = 0.3, + ) + # The timeout is enforced promptly (not blocked until stdout EOF). + assert time.time() - t0 < 5.0 + + def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "img.png" @@ -257,6 +307,36 @@ def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1] +def test_generate_native_speed_dedupes_against_offload(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + # offload already adds --diffusion-fa; native_speed="default" would add it again. + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = ["--offload-to-cpu", "--diffusion-fa"], + native_speed = "default", + ) + # --diffusion-fa appears exactly once (de-duped), not twice. + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + +def test_generate_native_speed_adds_flag_when_not_offloaded(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = [], # fast/resident tier: no offload, but speed flag still applies + native_speed = "default", + ) + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "big.png"