From 141cb3ae1085f82808b3d4c683cfab5187fd6baf Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 26 Jun 2026 03:19:59 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/compile_probe.py | 77 +++++++++++++++---- scripts/diffusion_bench.py | 2 +- scripts/perf_verify.py | 58 ++++++++++---- .../backend/core/inference/diffusion_speed.py | 9 ++- studio/backend/core/inference/sd_cpp_args.py | 4 +- .../backend/core/inference/sd_cpp_engine.py | 11 ++- studio/backend/tests/test_diffusion_speed.py | 14 +++- 7 files changed, 133 insertions(+), 42 deletions(-) diff --git a/scripts/compile_probe.py b/scripts/compile_probe.py index e1abbcc46a..5ff16ef763 100644 --- a/scripts/compile_probe.py +++ b/scripts/compile_probe.py @@ -40,8 +40,12 @@ def _gen(pipe, prompt, *, steps, seed, width, height, guidance): torch.cuda.synchronize() t0 = time.time() image = pipe( - prompt = prompt, width = width, height = height, - num_inference_steps = steps, guidance_scale = guidance, generator = gen, + 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 @@ -54,14 +58,21 @@ def main(argv = None) -> int: 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( + "--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( + "--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) @@ -80,7 +91,9 @@ def main(argv = None) -> int: transformer = transformer_cls.from_single_file( gguf_path, quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype), - torch_dtype = dtype, config = args.base_repo, subfolder = "transformer", + 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) @@ -88,8 +101,24 @@ def main(argv = None) -> int: 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) + _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) @@ -106,7 +135,10 @@ def main(argv = None) -> int: try: t0 = time.time() fn(**compile_kwargs) - print(f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)", flush = True) + 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 @@ -114,13 +146,29 @@ def main(argv = None) -> int: # 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) + _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, 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)) @@ -129,8 +177,11 @@ def main(argv = None) -> int: 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) + 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 diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index 39825468e6..e8accf477d 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -448,7 +448,7 @@ def _build_parser() -> argparse.ArgumentParser: 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", + "cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV", ) p.add_argument( "--text-encoder-quant", diff --git a/scripts/perf_verify.py b/scripts/perf_verify.py index 03591cbda4..77f1361e0b 100644 --- a/scripts/perf_verify.py +++ b/scripts/perf_verify.py @@ -42,7 +42,10 @@ 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( + "--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) @@ -62,8 +65,11 @@ def main(argv = None) -> int: 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, + 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: @@ -79,13 +85,25 @@ def main(argv = None) -> int: 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, + 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 = ""): + def timed( + mode_speed, + *, + warmup, + iters, + mem = None, + tag = "", + ): st = load(mode_speed, mem) for _ in range(warmup): gen() @@ -97,29 +115,41 @@ def main(argv = None) -> int: 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) + 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( + 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( + 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) + 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) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 23d5a85048..4350b6aa9a 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -45,7 +45,6 @@ def snapshot_backend_flags() -> Optional[dict]: 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), @@ -134,8 +133,11 @@ def apply_speed_optims( 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, "cudnn_benchmark": False, - "tf32": False, "fused_qkv": False, "compiled": False, + "channels_last": False, + "cudnn_benchmark": False, + "tf32": False, + "fused_qkv": False, + "compiled": False, } mode = normalize_speed_mode(speed_mode) if mode == SPEED_OFF: @@ -192,7 +194,6 @@ def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool: 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 diff --git a/studio/backend/core/inference/sd_cpp_args.py b/studio/backend/core/inference/sd_cpp_args.py index 06e738731f..ea259b9000 100644 --- a/studio/backend/core/inference/sd_cpp_args.py +++ b/studio/backend/core/inference/sd_cpp_args.py @@ -140,9 +140,7 @@ def native_speed_flags(speed_mode: Optional[str]) -> list[str]: 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}'" - ) + raise ValueError(f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'") def offload_flags( diff --git a/studio/backend/core/inference/sd_cpp_engine.py b/studio/backend/core/inference/sd_cpp_engine.py index 146a98018b..186ea51631 100644 --- a/studio/backend/core/inference/sd_cpp_engine.py +++ b/studio/backend/core/inference/sd_cpp_engine.py @@ -212,9 +212,14 @@ class SdCppEngine: 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 = merged_extra, + 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) diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 1d76b70f11..51705442eb 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -152,8 +152,11 @@ def test_speed_off_applies_nothing(monkeypatch): pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF ) assert applied == { - "channels_last": False, "cudnn_benchmark": False, - "tf32": False, "fused_qkv": False, "compiled": False, + "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). @@ -188,8 +191,11 @@ 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, + 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