diff --git a/scripts/compare_engines.py b/scripts/compare_engines.py index 48e50ac981..dba173f55a 100644 --- a/scripts/compare_engines.py +++ b/scripts/compare_engines.py @@ -45,13 +45,22 @@ def bench_pytorch(repo, gguf, resolutions, steps, seed, iters): raise RuntimeError(backend.load_progress()) 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) + 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)) @@ -74,18 +83,27 @@ def bench_sdcpp(binary, gguf, vae, llm, resolutions, steps, seed, iters): 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) + 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) + 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"]) @@ -93,7 +111,10 @@ def bench_sdcpp(binary, gguf, vae, llm, resolutions, steps, seed, iters): 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) + print( + f" sdcpp speed={tag:12s} {res}px compute={med_c}s wall={med_w:.3f}s", + flush = True, + ) return rows @@ -102,9 +123,17 @@ def main(argv = None) -> int: 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 = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/flux_vae/ae.safetensors") - p.add_argument("--llm", default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/qwen3_te/Qwen3-4B-Instruct-2507-Q4_K_M.gguf") + p.add_argument( + "--sd-gguf", default = None, help = "local gguf for sd.cpp (default: same as pytorch via cache)" + ) + p.add_argument( + "--vae", + default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/flux_vae/ae.safetensors", + ) + p.add_argument( + "--llm", + default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/qwen3_te/Qwen3-4B-Instruct-2507-Q4_K_M.gguf", + ) p.add_argument("--resolutions", default = "512,1024") p.add_argument("--steps", type = int, default = 8) p.add_argument("--seed", type = int, default = 42) @@ -121,7 +150,9 @@ def main(argv = None) -> int: 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) + 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) diff --git a/scripts/leverage_probe.py b/scripts/leverage_probe.py index 287d005242..756fe25e34 100644 --- a/scripts/leverage_probe.py +++ b/scripts/leverage_probe.py @@ -39,7 +39,9 @@ def _load(): 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", + torch_dtype = torch.bfloat16, + config = BASE, + subfolder = "transformer", ) pipe = diffusers.ZImagePipeline.from_pretrained(BASE, torch_dtype = torch.bfloat16, transformer = t) pipe.to("cuda") @@ -52,8 +54,14 @@ def _gen(pipe, steps, seed, res): 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] + 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 @@ -74,9 +82,16 @@ def main(argv = None) -> int: ic.coordinate_descent_tuning = True pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) - def run(tag, *, compile = False, cdt = False, fbc = None): + 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() @@ -110,7 +125,10 @@ def main(argv = None) -> int: 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) + 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) diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index 447825c361..2548c0f0dc 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -155,9 +155,7 @@ def apply_speed_optims( # 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, max_autotune = mode == SPEED_MAX - ) + 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. @@ -181,7 +179,12 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False -def _compile_repeated_blocks(pipe: Any, logger: Any, *, max_autotune: bool = False) -> 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):