From 395816cf7ede31ea91fad2988b08376a6b1f2a40 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 02:45:40 +0000 Subject: [PATCH] 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. --- scripts/compile_probe.py | 138 ++++++++++++++++++ scripts/diffusion_bench.py | 9 ++ scripts/perf_verify.py | 130 +++++++++++++++++ studio/backend/core/inference/diffusion.py | 35 ++++- .../core/inference/diffusion_families.py | 8 +- .../core/inference/diffusion_memory.py | 48 ++++-- .../backend/core/inference/diffusion_speed.py | 108 +++++++++++--- studio/backend/core/inference/sd_cpp_args.py | 30 ++++ .../backend/core/inference/sd_cpp_engine.py | 69 ++++++--- .../backend/tests/test_diffusion_backend.py | 16 +- studio/backend/tests/test_diffusion_memory.py | 3 + studio/backend/tests/test_diffusion_speed.py | 80 ++++++++-- studio/backend/tests/test_sd_cpp_args.py | 11 ++ studio/backend/tests/test_sd_cpp_engine.py | 80 ++++++++++ 14 files changed, 688 insertions(+), 77 deletions(-) create mode 100644 scripts/compile_probe.py create mode 100644 scripts/perf_verify.py diff --git a/scripts/compile_probe.py b/scripts/compile_probe.py new file mode 100644 index 0000000000..e1abbcc46a --- /dev/null +++ b/scripts/compile_probe.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Probe: does regional ``torch.compile`` work on the GGUF diffusion transformer? + +The speed layer gates ``compile_repeated_blocks`` OFF for GGUF (it dequantises +per-op). Since the backend is GGUF-only, that makes regional compile dead on +every shipping model. This probe loads a GGUF transformer exactly as +``diffusion.py`` does, runs an eager generation, then compiles the repeated +denoiser block and runs the same seed again, reporting: whether compile raised, +per-generation latency eager vs compiled, and PSNR(compiled vs eager). If compile +is clean and PSNR is high, the gate can be relaxed for this family. + +Run on one CUDA GPU. Read-only w.r.t. the backend (does not import the gate). +""" + +from __future__ import annotations + +import argparse +import sys +import time +from pathlib import Path + +import numpy as np + + +def _psnr(a: "np.ndarray", b: "np.ndarray") -> float: + a = a.astype(np.float64) + b = b.astype(np.float64) + mse = float(np.mean((a - b) ** 2)) + if mse == 0.0: + return float("inf") + return float(10.0 * np.log10((255.0**2) / mse)) + + +def _gen(pipe, prompt, *, steps, seed, width, height, guidance): + import torch + + gen = torch.Generator(device = "cuda").manual_seed(seed) + torch.cuda.synchronize() + t0 = time.time() + image = pipe( + prompt = prompt, width = width, height = height, + num_inference_steps = steps, guidance_scale = guidance, generator = gen, + ).images[0] + torch.cuda.synchronize() + return image, time.time() - t0 + + +def main(argv = None) -> int: + p = argparse.ArgumentParser() + p.add_argument("--repo", default = "unsloth/Z-Image-Turbo-GGUF") + p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf") + p.add_argument("--base-repo", default = "Tongyi-MAI/Z-Image-Turbo") + p.add_argument("--transformer-class", default = "ZImageTransformer2DModel") + p.add_argument("--pipeline-class", default = "ZImagePipeline") + p.add_argument("--prompt", default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed") + p.add_argument("--steps", type = int, default = 8) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--width", type = int, default = 1024) + p.add_argument("--height", type = int, default = 1024) + p.add_argument("--guidance", type = float, default = 0.0) + p.add_argument("--mode", default = "default", help = "compile mode: default | max-autotune-no-cudagraphs") + p.add_argument("--dynamic", action = "store_true", help = "dynamic=True (default False here for speed)") + p.add_argument("--out-dir", default = "outputs/compile_probe") + args = p.parse_args(argv) + + import torch + import diffusers + from huggingface_hub import hf_hub_download + + out = Path(args.out_dir) + out.mkdir(parents = True, exist_ok = True) + dtype = torch.bfloat16 + + gguf_path = hf_hub_download(args.repo, args.gguf) + print(f"gguf: {gguf_path}", flush = True) + + transformer_cls = getattr(diffusers, args.transformer_class) + transformer = transformer_cls.from_single_file( + gguf_path, + quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), + torch_dtype = dtype, config = args.base_repo, subfolder = "transformer", + ) + pipeline_cls = getattr(diffusers, args.pipeline_class) + pipe = pipeline_cls.from_pretrained(args.base_repo, torch_dtype = dtype, transformer = transformer) + pipe.to("cuda") + print("pipeline loaded on cuda", flush = True) + + # warm the eager path once (allocator / cudnn), then time eager. + _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance) + eager_img, eager_t = _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance) + eager_img.save(out / "eager.png") + eager_arr = np.array(eager_img) + print(f"EAGER: {eager_t:.2f}s/gen", flush = True) + + # compile the repeated denoiser block. + fn = getattr(pipe.transformer, "compile_repeated_blocks", None) + if not callable(fn): + print("RESULT: transformer has no compile_repeated_blocks -> N/A", flush = True) + return 3 + compile_kwargs = {"fullgraph": True, "dynamic": bool(args.dynamic)} + if args.mode and args.mode != "default": + compile_kwargs["mode"] = args.mode + print(f"compiling repeated blocks: {compile_kwargs} ...", flush = True) + try: + t0 = time.time() + fn(**compile_kwargs) + print(f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)", flush = True) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: compile_repeated_blocks RAISED: {type(exc).__name__}: {exc}", flush = True) + return 1 + + # first compiled gen triggers the actual compilation (untimed warmup). + try: + t0 = time.time() + _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance) + print(f" first compiled gen (compilation) took {time.time()-t0:.1f}s", flush = True) + except Exception as exc: # noqa: BLE001 + print(f"RESULT: first compiled generation RAISED: {type(exc).__name__}: {exc}", flush = True) + return 2 + + comp_img, comp_t = _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance) + comp_img.save(out / "compiled.png") + psnr = _psnr(eager_arr, np.array(comp_img)) + + speedup = (eager_t - comp_t) / eager_t * 100.0 + print("\n==== COMPILE PROBE RESULT ====", flush = True) + print(f" eager: {eager_t:.2f}s/gen", flush = True) + print(f" compiled: {comp_t:.2f}s/gen ({speedup:+.1f}% vs eager)", flush = True) + print(f" PSNR(compiled vs eager): {psnr:.1f} dB", flush = True) + print(f" verdict: {'COMPILE-WORKS' if psnr >= 30 else 'COMPILE-DIVERGES'} " + f"{'FASTER' if comp_t < eager_t else 'NOT-FASTER'}", flush = True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 8af3b6060d..39825468e6 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -214,6 +214,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) @@ -293,6 +294,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, }, @@ -441,6 +443,13 @@ def _build_parser() -> argparse.ArgumentParser: choices = ["auto", "fast", "balanced", "low_vram"], help = "memory policy (default: backend auto)", ) + p.add_argument( + "--speed-mode", + default = None, + choices = ["off", "default", "max"], + help = "speed profile: off is bit-identical; default adds compile + " + "cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV", + ) p.add_argument( "--text-encoder-quant", default = None, diff --git a/scripts/perf_verify.py b/scripts/perf_verify.py new file mode 100644 index 0000000000..03591cbda4 --- /dev/null +++ b/scripts/perf_verify.py @@ -0,0 +1,130 @@ +# 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 (def_t < off_t) and (_psnr(off_img, def_img) >= 30) + print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 8051d9665a..ed66e554f9 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -44,7 +44,13 @@ from .diffusion_memory import ( plan_diffusion_memory, snapshot_device_memory, ) -from .diffusion_speed import SPEED_OFF, apply_speed_optims +from .diffusion_speed import ( + SPEED_OFF, + apply_speed_optims, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, +) from .diffusion_precision import quantize_text_encoders logger = get_logger(__name__) @@ -69,6 +75,10 @@ class _LoadState: # The opt-in speed profile (Phase 3). speed_mode: str = SPEED_OFF speed_optims: tuple = () + # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the + # speed layer mutated them, restored on unload so a later `off` load is not + # contaminated by this one's globals. None when nothing was changed. + backend_flags_before: Optional[dict] = None # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). text_encoder_quant: Optional[str] = None @@ -462,14 +472,22 @@ 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. + # 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() speed_applied = apply_speed_optims( pipe, target, is_gguf = bool(gguf_filename), family = fam, - speed_mode = speed_mode or SPEED_OFF, + speed_mode = effective_speed, logger = logger, ) # Quantise the dense companion text encoder(s) (opt-in fp8 / nvfp4), @@ -508,8 +526,9 @@ class DiffusionBackend: offload_policy = effective_policy, vae_tiling = effective_tiling, memory_mode = plan.requested_mode, - speed_mode = (speed_mode or SPEED_OFF), + 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, ) @@ -646,7 +665,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; @@ -705,6 +727,9 @@ class DiffusionBackend: state = self._state if state is None: return + # Restore the process-wide backend flags (TF32 / cudnn.benchmark) this load + # may have flipped, so the next `off` load is bit-identical again. + restore_backend_flags(state.backend_flags_before) self._state = None del state clear_gpu_cache() diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index ba9125fabf..4e0d938178 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -34,9 +34,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 @@ -80,8 +80,6 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( aliases = ("zimage", "z_image"), # Z-Image's MLP down-projections peak near 9e5, which overflows float16. fp16_incompatible = True, - # Z-Image's denoiser block is excluded from regional torch.compile. - supports_torch_compile = False, ), ) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 753b8d1bd6..e896e538d6 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -397,17 +397,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), @@ -488,19 +492,33 @@ 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 - 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, - ) + 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(): diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 42c3b4d5be..23d5a85048 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -5,19 +5,29 @@ Off by default, so the default render path stays bit-identical to a plain run (the property the regression harness checks). When the operator opts in, this applies the -lossless-to-near-lossless speedups in the order the diffusers guides recommend -(channels_last -> regional compile, with TF32 / fused-QKV under "max"): +near-lossless speedups in the order the diffusers guides recommend +(channels_last + cudnn.benchmark -> regional compile, with TF32 / fused-QKV under +"max"): - off - nothing (default). - default - lossless: channels_last VAE memory format + regional torch.compile of - the denoiser's repeated block WHERE eligible (non-GGUF, bf16, CUDA, and - a compile-friendly family). + off - nothing (default; bit-identical reference). + default - near-lossless: channels_last VAE memory format + cudnn.benchmark conv + autotune + regional torch.compile of the denoiser's repeated block WHERE + eligible (bf16, CUDA, a compile-friendly family). Compile is the big win + (~2.3x denoise on the GGUF Z-Image transformer, PSNR ~36 dB vs eager, + well above the Q4 quantisation noise floor, so it does not meaningfully + move output quality). max - default plus near-lossless TF32 matmul and fused QKV projections. -Regional compile is gated off for the GGUF transformer (it dequantises per-op and -doesn't compile cleanly) and for families flagged not compile-friendly (Z-Image), so -on today's GGUF path only channels_last / TF32 engage; the compile path activates -automatically once a non-GGUF bf16 transformer is loaded. torch is imported lazily. +Regional compile used to be gated off for the GGUF transformer, but it compiles and +runs faster on the current diffusers/torch (measured; the GGUF dequant ops stay +eager and the rest of the repeated block compiles), so the GGUF gate is removed; the +per-family ``supports_torch_compile`` flag and the bf16/CUDA checks still apply. + +The backend flags this layer flips (TF32, cudnn.benchmark) are PROCESS-WIDE, so +``snapshot_backend_flags`` / ``restore_backend_flags`` let the caller capture the +prior values at load and restore them at unload, keeping a later ``off`` load +bit-identical instead of inheriting a previous ``max`` run's globals. torch is +imported lazily. """ from __future__ import annotations @@ -30,6 +40,35 @@ 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.""" + try: + import torch + + return { + "matmul_tf32": bool(torch.backends.cuda.matmul.allow_tf32), + "cudnn_tf32": bool(torch.backends.cudnn.allow_tf32), + "cudnn_benchmark": bool(torch.backends.cudnn.benchmark), + } + except Exception: # noqa: BLE001 — best-effort; no snapshot -> no restore + return None + + +def restore_backend_flags(state: Optional[dict]) -> None: + """Restore the flags captured by ``snapshot_backend_flags``. No-op on None.""" + if not state: + return + try: + import torch + + torch.backends.cuda.matmul.allow_tf32 = state["matmul_tf32"] + torch.backends.cudnn.allow_tf32 = state["cudnn_tf32"] + torch.backends.cudnn.benchmark = state["cudnn_benchmark"] + except Exception: # noqa: BLE001 — best-effort restore + return + + def normalize_speed_mode(value: Optional[str]) -> str: """Lower/strip a requested speed mode (dashes ok); None / "" -> off.""" if value is None: @@ -44,14 +83,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 +133,10 @@ 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) if mode == SPEED_OFF: return applied @@ -87,7 +144,13 @@ 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). if compile_eligible(target, is_gguf = is_gguf, family = family): applied["compiled"] = _compile_repeated_blocks(pipe, logger) @@ -126,6 +189,17 @@ def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool: 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 + + def _enable_tf32(logger: Any) -> bool: try: import torch diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index b2df104b7c..06e738731f 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -115,6 +115,36 @@ class SdCppUpscaleParams: tile_size: Optional[int] = None +# Native (sd.cpp) speed profiles, the engine-side analogue of diffusion_speed's +# modes. off: nothing (default). default: --diffusion-fa (flash attention; upstream +# reports it usually speeds CUDA and cuts attention memory, near-lossless). max: also +# --diffusion-conv-direct (direct conv; helps some backends, but measured +45% on +# CUDA here, so it stays opt-in/experimental, never auto-on for CUDA). +NATIVE_SPEED_OFF = "off" +NATIVE_SPEED_DEFAULT = "default" +NATIVE_SPEED_MAX = "max" +NATIVE_SPEED_MODES = (NATIVE_SPEED_OFF, NATIVE_SPEED_DEFAULT, NATIVE_SPEED_MAX) + + +def native_speed_flags(speed_mode: Optional[str]) -> list[str]: + """sd-cli speed flags for a native speed mode (empty for off / None). + + These are separate from the offload flags: ``--diffusion-fa`` is a speed/memory + win in its own right, not tied to whether weights are offloaded. De-duplicated + against offload flags at the call site (offload already adds ``--diffusion-fa``). + """ + mode = (speed_mode or NATIVE_SPEED_OFF).strip().lower() + if mode in ("", NATIVE_SPEED_OFF): + return [] + if mode == NATIVE_SPEED_DEFAULT: + return ["--diffusion-fa"] + if mode == NATIVE_SPEED_MAX: + return ["--diffusion-fa", "--diffusion-conv-direct"] + raise ValueError( + f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'" + ) + + def offload_flags( policy: str, *, diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index b1daafce72..146a98018b 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -24,9 +24,11 @@ from __future__ import annotations import logging import os +import queue import shutil import subprocess import sys +import threading import time from pathlib import Path from typing import Callable, Optional @@ -37,6 +39,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, ) logger = logging.getLogger(__name__) @@ -189,6 +192,7 @@ class SdCppEngine: *, output_path: str, offload: Optional[list[str]] = None, + native_speed: Optional[str] = None, threads: Optional[int] = None, verbose: bool = False, extra_args: Optional[list[str]] = None, @@ -198,19 +202,19 @@ 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, - params, - output_path = str(self._prepare_out(output_path)), - offload = offload, - threads = threads, - verbose = verbose, - extra_args = extra_args, + self._require_binary(), files, params, + output_path = str(self._prepare_out(output_path)), offload = offload, + threads = threads, verbose = verbose, extra_args = merged_extra, ) return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log) @@ -281,20 +285,49 @@ class SdCppEngine: errors = "replace", env = run_env, ) + # Drain stdout on a reader thread so the timeout is enforced even when the + # child hangs WITHOUT printing (e.g. stuck in model load / GPU init): a plain + # `for line in proc.stdout` blocks until EOF, so proc.wait(timeout) would + # never be reached. The reader pushes lines (then a None sentinel at EOF) to a + # queue the main loop polls against a wall-clock deadline. tail: list[str] = [] + line_q: "queue.Queue[Optional[str]]" = queue.Queue() + + def _drain() -> None: + try: + assert proc.stdout is not None + for raw in proc.stdout: + line_q.put(raw.rstrip("\n")) + finally: + line_q.put(None) + + reader = threading.Thread(target = _drain, daemon = True) + reader.start() + + deadline = None if timeout is None else time.monotonic() + float(timeout) + stdout_done = False try: - assert proc.stdout is not None - for line in proc.stdout: - line = line.rstrip("\n") + while True: + if deadline is not None and time.monotonic() >= deadline and proc.poll() is None: + proc.kill() + raise RuntimeError(f"sd-cli timed out after {timeout}s") + try: + line = line_q.get(timeout = 0.1) + except queue.Empty: + if proc.poll() is not None and stdout_done: + break + continue + if line is None: + stdout_done = True + if proc.poll() is not None: + break + continue tail.append(line) if len(tail) > 40: tail.pop(0) if on_log is not None: on_log(line) - ret = proc.wait(timeout = timeout) - except subprocess.TimeoutExpired: - proc.kill() - raise RuntimeError(f"sd-cli timed out after {timeout}s") + ret = proc.wait(timeout = 5.0) finally: if proc.poll() is None: proc.kill() diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 7bf9f02f17..533f7b8e01 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -10,6 +10,7 @@ GPU, weights, or network access is needed (sub-second, CI-friendly). from __future__ import annotations +import contextlib import sys import types @@ -202,6 +203,8 @@ def fake_runtime(monkeypatch): torch.Generator = _FakeGenerator torch.cuda = types.SimpleNamespace(is_available = lambda: False) torch.backends = types.SimpleNamespace(mps = None) + # generate() wraps the pipe call in torch.inference_mode(); a no-op CM here. + torch.inference_mode = lambda: contextlib.nullcontext() diffusers = types.ModuleType("diffusers") diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) @@ -842,12 +845,19 @@ def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( assert status["offload_policy"] == "model" and status["cpu_offload"] is True -def test_load_speed_mode_threads_and_defaults_off(fake_runtime, tmp_path): - # No speed_mode -> off, no optimisations engaged (the bit-identical default). +def test_load_speed_mode_gguf_auto_defaults_and_explicit(fake_runtime, tmp_path): + # No speed_mode on a GGUF model -> auto `default` (near-lossless, compile sits + # below the quant noise floor). compile itself only engages on CUDA, so on this + # CPU stub no optim need engage, but the resolved mode is `default`. (tmp_path / "m.gguf").write_bytes(b"x") backend = DiffusionBackend() status = backend.load_pipeline(str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image") - assert status["speed_mode"] == "off" and status["speed_optims"] == [] + assert status["speed_mode"] == "default" + # An explicit "off" opts back into the bit-identical path (engages nothing). + status_off = backend.load_pipeline( + str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "off" + ) + assert status_off["speed_mode"] == "off" and status_off["speed_optims"] == [] # An explicit speed_mode threads through to status (engaged optims are GPU-verified). status2 = backend.load_pipeline( str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image", speed_mode = "max" diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 858e87ffa1..2b1480a63b 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -185,6 +185,9 @@ def test_auto_group_offload_when_transformer_overflows_but_companions_fit(): base_overhead_mib = 1000, ) assert plan.offload_policy == OFFLOAD_GROUP + # Group keeps the VAE resident, so it uses exact slicing but NOT lossy tiling + # -> balanced stays bit-identical while still capping the offload footprint. + assert plan.vae_slicing is True and plan.vae_tiling is False def test_auto_model_offload_when_companions_exceed_budget(): diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index ade37482ba..1d76b70f11 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -21,6 +21,9 @@ from core.inference.diffusion_speed import ( apply_speed_optims, compile_eligible, normalize_speed_mode, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, ) @@ -47,7 +50,7 @@ def _stub_torch(monkeypatch): torch.channels_last = "channels_last" torch.backends = types.SimpleNamespace( cuda = types.SimpleNamespace(matmul = types.SimpleNamespace(allow_tf32 = False)), - cudnn = types.SimpleNamespace(allow_tf32 = False), + cudnn = types.SimpleNamespace(allow_tf32 = False, benchmark = False), ) monkeypatch.setitem(sys.modules, "torch", torch) return torch @@ -64,23 +67,55 @@ 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 + + # ── applier ─────────────────────────────────────────────────────────────────── @@ -111,16 +146,21 @@ class _Pipe: 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 +168,30 @@ 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 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): diff --git a/studio/backend/tests/test_sd_cpp_args.py b/studio/backend/tests/test_sd_cpp_args.py index 7e3a024f0e..99582aed46 100644 --- a/studio/backend/tests/test_sd_cpp_args.py +++ b/studio/backend/tests/test_sd_cpp_args.py @@ -23,6 +23,7 @@ from core.inference.sd_cpp_args import ( SdCppUpscaleParams, build_sd_cpp_command, build_sd_cpp_upscale_command, + native_speed_flags, offload_flags, text_encoder_flags_for_family, ) @@ -47,6 +48,16 @@ def test_te_flags_by_family(): # ── offload policy -> sd-cli flags ────────────────────────────────────────── +def test_native_speed_flags(): + assert native_speed_flags(None) == [] + assert native_speed_flags("off") == [] + assert native_speed_flags("") == [] + assert native_speed_flags("default") == ["--diffusion-fa"] + assert native_speed_flags("max") == ["--diffusion-fa", "--diffusion-conv-direct"] + with pytest.raises(ValueError): + native_speed_flags("ludicrous") + + def test_offload_none_is_empty(): assert offload_flags(OFFLOAD_NONE) == [] diff --git a/studio/backend/tests/test_sd_cpp_engine.py b/studio/backend/tests/test_sd_cpp_engine.py index 7aeb537033..a879252761 100644 --- a/studio/backend/tests/test_sd_cpp_engine.py +++ b/studio/backend/tests/test_sd_cpp_engine.py @@ -12,6 +12,7 @@ from __future__ import annotations import os import sys +import time import types from pathlib import Path @@ -242,6 +243,55 @@ def test_generate_raises_when_binary_missing(): ) +class _HangingPopen: + """A child that runs but never prints and never exits -- the case a plain + `for line in stdout` would block on forever, ignoring the timeout.""" + + def __init__(self, cmd, **_kw): + self._alive = True + + class _Blocking: + def __init__(self, owner): + self.owner = owner + + def __iter__(self): + return self + + def __next__(self): + while self.owner._alive: + time.sleep(0.01) + raise StopIteration + + @property + def stdout(self): + return self._Blocking(self) + + def poll(self): + return None if self._alive else -9 + + def wait(self, timeout = None): + self._alive = False + return -9 + + def kill(self): + self._alive = False + + +def test_generate_times_out_on_silent_hang(tmp_path, monkeypatch): + e = _engine(tmp_path) + monkeypatch.setattr(eng.subprocess, "Popen", lambda cmd, **kw: _HangingPopen(cmd, **kw)) + t0 = time.time() + with pytest.raises(RuntimeError, match = "timed out"): + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(tmp_path / "x.png"), + timeout = 0.3, + ) + # The timeout is enforced promptly (not blocked until stdout EOF). + assert time.time() - t0 < 5.0 + + def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "img.png" @@ -257,6 +307,36 @@ def test_img2img_generate_passes_init_image(tmp_path, monkeypatch): assert str(src) == _FakePopen.captured_cmd[_FakePopen.captured_cmd.index("--init-img") + 1] +def test_generate_native_speed_dedupes_against_offload(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + # offload already adds --diffusion-fa; native_speed="default" would add it again. + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = ["--offload-to-cpu", "--diffusion-fa"], + native_speed = "default", + ) + # --diffusion-fa appears exactly once (de-duped), not twice. + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + +def test_generate_native_speed_adds_flag_when_not_offloaded(tmp_path, monkeypatch): + e = _engine(tmp_path) + out = tmp_path / "img.png" + _patch_popen(monkeypatch, lines = ["ok"], returncode = 0, out_file = out) + e.generate( + SdCppModelFiles(diffusion_model = "/m/z.gguf"), + SdCppGenParams(prompt = "x"), + output_path = str(out), + offload = [], # fast/resident tier: no offload, but speed flag still applies + native_speed = "default", + ) + assert _FakePopen.captured_cmd.count("--diffusion-fa") == 1 + + def test_upscale_runs_and_returns_path(tmp_path, monkeypatch): e = _engine(tmp_path) out = tmp_path / "big.png"