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 <lora:name:weight> 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 <stem>_<idx><suffix> (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>
This commit is contained in:
Daniel Han
GitHub
parent
2d09508951
commit
6b9b1c72d3
16 changed files with 1262 additions and 129 deletions
174
scripts/compare_engines.py
Normal file
174
scripts/compare_engines.py
Normal file
|
|
@ -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())
|
||||
189
scripts/compile_probe.py
Normal file
189
scripts/compile_probe.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
146
scripts/leverage_probe.py
Normal file
146
scripts/leverage_probe.py
Normal file
|
|
@ -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())
|
||||
165
scripts/perf_verify.py
Normal file
165
scripts/perf_verify.py
Normal file
|
|
@ -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())
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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) == []
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue