From 7f6e69dd5cc8e3d134b2a13878792b4a4fbe3c04 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 25 Jun 2026 11:30:37 +0000 Subject: [PATCH] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/diffusion_bench.py | 131 ++++++++++++------ studio/backend/core/inference/diffusion.py | 5 +- .../core/inference/diffusion_device.py | 1 - .../backend/tests/test_diffusion_backend.py | 51 +++++-- studio/backend/tests/test_diffusion_device.py | 27 +++- studio/backend/tests/test_diffusion_routes.py | 8 +- 6 files changed, 157 insertions(+), 66 deletions(-) diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index de8a9a570c..4ee86d2be3 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -62,7 +62,9 @@ def _git_commit() -> Optional[str]: out = subprocess.run( ["git", "rev-parse", "HEAD"], cwd = str(Path(__file__).resolve().parent), - capture_output = True, text = True, timeout = 10, + capture_output = True, + text = True, + timeout = 10, ) return out.stdout.strip() or None if out.returncode == 0 else None except Exception: @@ -85,40 +87,34 @@ def _is_cuda(device: Optional[str]) -> bool: def _cuda_reset_peak() -> None: import torch - if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() def _cuda_sync() -> None: import torch - if torch.cuda.is_available(): torch.cuda.synchronize() def _cuda_peak_alloc() -> Optional[int]: import torch - return int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None def _cuda_peak_reserved() -> Optional[int]: import torch - return int(torch.cuda.max_memory_reserved()) if torch.cuda.is_available() else None def _cuda_alloc() -> Optional[int]: import torch - return int(torch.cuda.memory_allocated()) if torch.cuda.is_available() else None def _gpu_name() -> Optional[str]: try: import torch - if torch.cuda.is_available(): return torch.cuda.get_device_name(0) except Exception: @@ -130,13 +126,11 @@ def _versions() -> dict[str, Optional[str]]: out: dict[str, Optional[str]] = {"torch": None, "diffusers": None} try: import torch - out["torch"] = torch.__version__ except Exception: pass try: import diffusers - out["diffusers"] = diffusers.__version__ except Exception: pass @@ -255,7 +249,9 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "latencies_s": [round(x, 4) for x in latencies], "median_latency_s": round(_percentile(latencies, 50), 4), "p90_latency_s": round(_percentile(latencies, 90), 4), - "images_per_sec": round((args.batch_size * len(latencies)) / total, 4) if total > 0 else None, + "images_per_sec": round((args.batch_size * len(latencies)) / total, 4) + if total > 0 + else None, "peak_vram_bytes": _cuda_peak_alloc(), } @@ -282,10 +278,17 @@ def _run(args: argparse.Namespace) -> dict[str, Any]: "load": load_metrics, "generate": gen_metrics, "config": { - "model": args.model, "gguf": args.gguf, "base_repo": args.base_repo, - "family_override": args.family_override, "prompt": args.prompt, - "width": args.width, "height": args.height, "steps": args.steps, - "guidance": args.guidance, "seed": args.seed, "batch_size": args.batch_size, + "model": args.model, + "gguf": args.gguf, + "base_repo": args.base_repo, + "family_override": args.family_override, + "prompt": args.prompt, + "width": args.width, + "height": args.height, + "steps": args.steps, + "guidance": args.guidance, + "seed": args.seed, + "batch_size": args.batch_size, }, } @@ -301,8 +304,10 @@ def _write_baseline(args: argparse.Namespace) -> int: metrics = _run(args) metrics["accuracy"] = { "reference_png": str(ref_png), - "width": args.width, "height": args.height, - "steps": args.steps, "seed": args.seed, + "width": args.width, + "height": args.height, + "steps": args.steps, + "seed": args.seed, "dtype": (metrics["env"]["status"] or {}).get("dtype"), } @@ -312,10 +317,13 @@ def _write_baseline(args: argparse.Namespace) -> int: print(f" json: {baseline_path}", flush = True) print(f" reference: {ref_png}", flush = True) print(f" load: {metrics['load']}", flush = True) - print(f" generate: median={metrics['generate'].get('median_latency_s')}s " - f"p90={metrics['generate'].get('p90_latency_s')}s " - f"img/s={metrics['generate'].get('images_per_sec')} " - f"peak_vram={metrics['generate'].get('peak_vram_bytes')}", flush = True) + print( + f" generate: median={metrics['generate'].get('median_latency_s')}s " + f"p90={metrics['generate'].get('p90_latency_s')}s " + f"img/s={metrics['generate'].get('images_per_sec')} " + f"peak_vram={metrics['generate'].get('peak_vram_bytes')}", + flush = True, + ) return 0 @@ -362,14 +370,22 @@ def _compare(args: argparse.Namespace) -> int: print("\n=== REGRESSION REPORT ===", flush = True) print(f" {'metric':<22}{'baseline':>16}{'current':>16}{'delta':>12}", flush = True) - print(f" {'median_latency_s':<22}{base_median:>16.4f}{cur_median:>16.4f}{latency_reg * 100:>11.1f}%", flush = True) + print( + f" {'median_latency_s':<22}{base_median:>16.4f}{cur_median:>16.4f}{latency_reg * 100:>11.1f}%", + flush = True, + ) if base_peak and cur_peak: - print(f" {'peak_vram_MB':<22}{base_peak / 1e6:>16.1f}{cur_peak / 1e6:>16.1f}{vram_reg * 100:>11.1f}%", flush = True) + print( + f" {'peak_vram_MB':<22}{base_peak / 1e6:>16.1f}{cur_peak / 1e6:>16.1f}{vram_reg * 100:>11.1f}%", + flush = True, + ) print(f" {'psnr_dB(vs ref)':<22}{'-':>16}{psnr:>16.2f}{'':>12}", flush = True) failures = [] if latency_reg > args.max_latency_regression: - failures.append(f"latency +{latency_reg * 100:.1f}% > {args.max_latency_regression * 100:.0f}%") + failures.append( + f"latency +{latency_reg * 100:.1f}% > {args.max_latency_regression * 100:.0f}%" + ) if base_peak and cur_peak and vram_reg > args.max_vram_regression: failures.append(f"peak VRAM +{vram_reg * 100:.1f}% > {args.max_vram_regression * 100:.0f}%") if not math.isnan(psnr) and psnr < args.min_psnr: @@ -390,15 +406,21 @@ def _build_parser() -> argparse.ArgumentParser: description = "Benchmark + regression guard for the Studio diffusion backend.", formatter_class = argparse.ArgumentDefaultsHelpFormatter, ) - p.add_argument("--model", default = "unsloth/Z-Image-Turbo-GGUF", - help = "GGUF repo id or local path") - p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf", - help = "transformer GGUF filename inside --model") + p.add_argument( + "--model", default = "unsloth/Z-Image-Turbo-GGUF", help = "GGUF repo id or local path" + ) + p.add_argument( + "--gguf", + default = "z-image-turbo-Q4_K_M.gguf", + help = "transformer GGUF filename inside --model", + ) p.add_argument("--base-repo", default = None, help = "override the diffusers base repo") p.add_argument("--family-override", default = None, help = "force a diffusion family") - p.add_argument("--prompt", - default = "A cozy reading nook by a rain-streaked window, warm lamplight, " - "a cat asleep on a stack of books, highly detailed") + p.add_argument( + "--prompt", + default = "A cozy reading nook by a rain-streaked window, warm lamplight, " + "a cat asleep on a stack of books, highly detailed", + ) p.add_argument("--width", type = int, default = 1024) p.add_argument("--height", type = int, default = 1024) p.add_argument("--steps", type = int, default = 9) @@ -407,20 +429,41 @@ def _build_parser() -> argparse.ArgumentParser: p.add_argument("--batch-size", type = int, default = 1) p.add_argument("--warmup", type = int, default = 1, help = "discarded warmup generations") p.add_argument("--iters", type = int, default = 3, help = "measured generations") - p.add_argument("--write-baseline", metavar = "PATH", default = None, - help = "run once and save metrics JSON + reference.png") - p.add_argument("--compare", metavar = "PATH", default = None, - help = "run again and diff against a baseline JSON") - p.add_argument("--max-latency-regression", type = float, default = 0.10, - help = "fail if median latency rises by more than this fraction") - p.add_argument("--max-vram-regression", type = float, default = 0.10, - help = "fail if peak generation VRAM rises by more than this fraction") - p.add_argument("--min-psnr", type = float, default = 35.0, - help = "fail if the fixed-seed image PSNR vs reference drops below this") - p.add_argument("--force-compare", action = "store_true", - help = "compare even when GPU/device/dtype differ from the baseline") - p.add_argument("--out-dir", default = "outputs/diffusion_bench", - help = "where compare.png is written") + p.add_argument( + "--write-baseline", + metavar = "PATH", + default = None, + help = "run once and save metrics JSON + reference.png", + ) + p.add_argument( + "--compare", metavar = "PATH", default = None, help = "run again and diff against a baseline JSON" + ) + p.add_argument( + "--max-latency-regression", + type = float, + default = 0.10, + help = "fail if median latency rises by more than this fraction", + ) + p.add_argument( + "--max-vram-regression", + type = float, + default = 0.10, + help = "fail if peak generation VRAM rises by more than this fraction", + ) + p.add_argument( + "--min-psnr", + type = float, + default = 35.0, + help = "fail if the fixed-seed image PSNR vs reference drops below this", + ) + p.add_argument( + "--force-compare", + action = "store_true", + help = "compare even when GPU/device/dtype differ from the baseline", + ) + p.add_argument( + "--out-dir", default = "outputs/diffusion_bench", help = "where compare.png is written" + ) return p diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 6cfdfd4911..c4dee53e5e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -466,6 +466,7 @@ class DiffusionBackend: batch_size: int = 1, ) -> dict[str, Any]: import torch + # A per-generation cancel Event: unload()/a superseding load set THIS event # (registered under _lock below) to abort just this denoise. _generate_lock # serialises generations and is the only lock the denoise holds, so a slow @@ -521,7 +522,9 @@ class DiffusionBackend: gen.step = step_index + 1 if gen.first_step_at == 0.0: gen.first_step_at = now - gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now) + gen.eta_seconds = _estimate_eta( + gen.total_steps, gen.step, gen.first_step_at, now + ) # Preempt a long denoise on unload/eviction or a superseding load: # diffusers checks pipe._interrupt and stops after the current step. if cancel.is_set(): diff --git a/studio/backend/core/inference/diffusion_device.py b/studio/backend/core/inference/diffusion_device.py index f72cdf464e..0c874e4b66 100644 --- a/studio/backend/core/inference/diffusion_device.py +++ b/studio/backend/core/inference/diffusion_device.py @@ -109,7 +109,6 @@ def diffusion_device_target_from_torch_device( if device == "cuda": try: import torch - is_rocm = bool(getattr(getattr(torch, "version", None), "hip", None)) except Exception: is_rocm = False diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 133a24351c..73888f450b 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -565,8 +565,13 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime): fam = detect_family("unsloth/Z-Image-GGUF") backend._state = _LoadState( - pipe = _BlockingPipe(), family = fam, repo_id = "r", base_repo = "b", - device = "cpu", dtype = "float32", cpu_offload = False, + pipe = _BlockingPipe(), + family = fam, + repo_id = "r", + base_repo = "b", + device = "cpu", + dtype = "float32", + cpu_offload = False, ) out: dict = {} @@ -610,7 +615,13 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime): self._interrupt = False self.steps_run = 0 - def __call__(self, *, callback_on_step_end = None, num_inference_steps = 8, **kwargs): + def __call__( + self, + *, + callback_on_step_end = None, + num_inference_steps = 8, + **kwargs, + ): for i in range(num_inference_steps): if self._interrupt: # diffusers' interrupt protocol break @@ -625,8 +636,13 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime): pipe = _SteppingPipe() fam = detect_family("unsloth/Z-Image-GGUF") backend._state = _LoadState( - pipe = pipe, family = fam, repo_id = "r", base_repo = "b", - device = "cpu", dtype = "float32", cpu_offload = False, + pipe = pipe, + family = fam, + repo_id = "r", + base_repo = "b", + device = "cpu", + dtype = "float32", + cpu_offload = False, ) out: dict = {} @@ -658,18 +674,22 @@ def test_validate_load_request(tmp_path): backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF") with pytest.raises(ValueError, match = "family"): backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf") - assert backend.validate_load_request( - "unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf" - ).name == "z-image" + assert ( + backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name + == "z-image" + ) # A local path with a missing child fails here (before any GPU/network work). with pytest.raises(FileNotFoundError): backend.validate_load_request( str(tmp_path), gguf_filename = "missing.gguf", family_override = "z-image" ) (tmp_path / "m.gguf").write_bytes(b"x") - assert backend.validate_load_request( - str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image" - ).name == "z-image" + assert ( + backend.validate_load_request( + str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image" + ).name + == "z-image" + ) # A path-shaped repo_id that does not exist is rejected here (it would otherwise # be treated as remote, evict chat, and only fail in the background load). with pytest.raises(FileNotFoundError): @@ -698,8 +718,13 @@ def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path): fam = detect_family("unsloth/Z-Image-GGUF") backend._state = _LoadState( - pipe = _BlockingPipe(), family = fam, repo_id = "r", base_repo = "b", - device = "cpu", dtype = "float32", cpu_offload = False, + pipe = _BlockingPipe(), + family = fam, + repo_id = "r", + base_repo = "b", + device = "cpu", + dtype = "float32", + cpu_offload = False, ) gen_out: dict = {} diff --git a/studio/backend/tests/test_diffusion_device.py b/studio/backend/tests/test_diffusion_device.py index 112cc3f397..2eb9b3a822 100644 --- a/studio/backend/tests/test_diffusion_device.py +++ b/studio/backend/tests/test_diffusion_device.py @@ -65,13 +65,13 @@ class _FakeTensor: def _make_torch( *, cuda_available: bool = False, - capability=(8, 0), + capability = (8, 0), capability_raises: bool = False, bf16_supported: bool = False, - hip=None, + hip = None, mps_available: bool = False, mps_probe: str = "pass", # "pass" | "raise" | "nonfinite" - xpu_available=None, # None -> no xpu attr; True/False -> present + xpu_available = None, # None -> no xpu attr; True/False -> present xpu_bf16: bool = False, ) -> types.ModuleType: torch = types.ModuleType("torch") @@ -110,7 +110,14 @@ def _make_torch( return torch -def _install(monkeypatch, torch, *, studio_device=None, is_rocm=False, hardware_fails=False): +def _install( + monkeypatch, + torch, + *, + studio_device = None, + is_rocm = False, + hardware_fails = False, +): """Install the fake torch and either a fake or failing `utils.hardware`.""" monkeypatch.setitem(sys.modules, "torch", torch) if hardware_fails: @@ -139,7 +146,11 @@ def test_cuda_ampere_bf16(monkeypatch): _install(monkeypatch, torch, studio_device = "cuda") t = dd.resolve_diffusion_device_target() assert (t.device, t.dtype, t.backend, t.vendor) == ("cuda", BF16, "cuda", "nvidia") - assert t.supports_model_cpu_offload and t.supports_default_torch_compile and t.supports_pinned_transfer + assert ( + t.supports_model_cpu_offload + and t.supports_default_torch_compile + and t.supports_pinned_transfer + ) def test_cuda_pre_ampere_fp16(monkeypatch): @@ -185,7 +196,11 @@ def test_xpu_bf16(monkeypatch): _install(monkeypatch, torch, studio_device = "xpu") t = dd.resolve_diffusion_device_target() assert (t.device, t.backend, t.vendor, t.dtype) == ("xpu", "xpu", "intel", BF16) - assert t.supports_model_cpu_offload and not t.supports_default_torch_compile and not t.supports_pinned_transfer + assert ( + t.supports_model_cpu_offload + and not t.supports_default_torch_compile + and not t.supports_pinned_transfer + ) def test_xpu_without_bf16_fp16(monkeypatch): diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index 23650494b0..632e2d1f6c 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -29,7 +29,13 @@ class _FakeBackend: def is_loaded(self) -> bool: return self.loaded - def validate_load_request(self, model_path, *, gguf_filename = None, family_override = None): + def validate_load_request( + self, + model_path, + *, + gguf_filename = None, + family_override = None, + ): # Mirror the real backend's cheap validation so the route's # validate-before-evict ordering is exercised. from core.inference.diffusion_families import detect_family