diff --git a/scripts/diffusion_bench.py b/scripts/diffusion_bench.py index d7381ed3a5..9e1a928362 100644 --- a/scripts/diffusion_bench.py +++ b/scripts/diffusion_bench.py @@ -280,11 +280,19 @@ 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, - "memory_mode": args.memory_mode, "cpu_offload": args.cpu_offload, + "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, + "memory_mode": args.memory_mode, + "cpu_offload": args.cpu_offload, }, } @@ -425,25 +433,50 @@ 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("--memory-mode", default = None, - choices = ["auto", "fast", "balanced", "low_vram"], - help = "memory policy (default: backend auto)") - p.add_argument("--cpu-offload", action = "store_true", - help = "legacy: force whole-module CPU offload") - 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( + "--memory-mode", + default = None, + choices = ["auto", "fast", "balanced", "low_vram"], + help = "memory policy (default: backend auto)", + ) + p.add_argument( + "--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload" + ) + 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/scripts/diffusion_quality.py b/scripts/diffusion_quality.py index fb233010b4..f4e54b9c0d 100644 --- a/scripts/diffusion_quality.py +++ b/scripts/diffusion_quality.py @@ -56,7 +56,6 @@ DEFAULT_PROMPTS = [ def _to_gray(img: Any) -> Any: import numpy as np - return np.asarray(img.convert("L"), dtype = np.float64) @@ -93,7 +92,11 @@ def _box_mean(x: Any, w: int) -> Any: return total / float(w * w) -def ssim(a_img: Any, b_img: Any, window: int = 7) -> float: +def ssim( + a_img: Any, + b_img: Any, + window: int = 7, +) -> float: """Mean structural similarity (luminance) over a uniform window; 1.0 when identical. Pure numpy box-window SSIM (Wang et al. constants), no skimage.""" a, b = _to_gray(a_img), _to_gray(b_img) @@ -133,9 +136,9 @@ class _Clip: return emb / emb.norm(dim = -1, keepdim = True) def _text_embed(self, text: str) -> Any: - inputs = self.proc( - text = [text], return_tensors = "pt", padding = True, truncation = True - ).to(self.device) + inputs = self.proc(text = [text], return_tensors = "pt", padding = True, truncation = True).to( + self.device + ) with self.torch.no_grad(): emb = self.model.get_text_features(**inputs) return emb / emb.norm(dim = -1, keepdim = True) @@ -153,7 +156,6 @@ class _Clip: def _cuda(call: str) -> Optional[int]: try: import torch - if not torch.cuda.is_available(): return None return int(getattr(torch.cuda, call)()) @@ -164,7 +166,6 @@ def _cuda(call: str) -> Optional[int]: def _cuda_reset_peak() -> None: try: import torch - if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() torch.cuda.synchronize() @@ -187,7 +188,6 @@ def _wait_for_load(backend: Any, timeout_s: int = 3600) -> None: def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]: try: from huggingface_hub import HfApi - info = HfApi().model_info(repo, files_metadata = True, token = os.environ.get("HF_TOKEN")) for s in info.siblings: if s.rfilename == filename and s.size: @@ -200,13 +200,18 @@ def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]: # ── one quant: load, render the grid, measure ──────────────────────────────── -def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Path) -> dict[str, Any]: +def _render_grid( + backend: Any, args: argparse.Namespace, gguf: str, out_dir: Path +) -> dict[str, Any]: """Load ``gguf`` and render one image per (prompt, seed); return images keyed by (prompt_index, seed) plus latency / VRAM metrics.""" _cuda_reset_peak() backend.begin_load( - args.model, gguf_filename = gguf, base_repo = args.base_repo, - family_override = args.family_override, hf_token = os.environ.get("HF_TOKEN"), + args.model, + gguf_filename = gguf, + base_repo = args.base_repo, + family_override = args.family_override, + hf_token = os.environ.get("HF_TOKEN"), memory_mode = args.memory_mode, ) _wait_for_load(backend) @@ -221,8 +226,13 @@ def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Pat for seed in args.seeds: t0 = time.time() result = backend.generate( - prompt = prompt, width = args.width, height = args.height, - steps = args.steps, guidance = args.guidance, seed = seed, batch_size = 1, + prompt = prompt, + width = args.width, + height = args.height, + steps = args.steps, + guidance = args.guidance, + seed = seed, + batch_size = 1, ) latencies.append(time.time() - t0) img = result["images"][0] @@ -243,7 +253,9 @@ def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Pat } -def _compare(grid: dict, ref_grid: dict, clip: Optional[_Clip], prompts: list[str]) -> dict[str, Any]: +def _compare( + grid: dict, ref_grid: dict, clip: Optional[_Clip], prompts: list[str] +) -> dict[str, Any]: psnrs, ssims, clip_txt, clip_sim = [], [], [], [] for key, img in grid["images"].items(): ref = ref_grid["images"].get(key) @@ -285,16 +297,20 @@ def _sweep(args: argparse.Namespace) -> int: rows: list[dict[str, Any]] = [] for gguf in quants: print(f"=== quant: {gguf} ===", flush = True) - grid = gguf == args.reference_quant and ref_grid or _render_grid(backend, args, gguf, out_dir) + grid = ( + gguf == args.reference_quant and ref_grid or _render_grid(backend, args, gguf, out_dir) + ) metrics = _compare(grid, ref_grid, clip, args.prompts) - rows.append({ - "quant": gguf, - "file_size_mib": grid["file_size_mib"], - "is_reference": gguf == args.reference_quant, - "median_latency_s": grid["median_latency_s"], - "peak_vram_mib": (grid["peak_vram_bytes"] or 0) // (1024 * 1024) or None, - **metrics, - }) + rows.append( + { + "quant": gguf, + "file_size_mib": grid["file_size_mib"], + "is_reference": gguf == args.reference_quant, + "median_latency_s": grid["median_latency_s"], + "peak_vram_mib": (grid["peak_vram_bytes"] or 0) // (1024 * 1024) or None, + **metrics, + } + ) print(f" {metrics}", flush = True) _write_outputs(args, out_dir, rows) @@ -304,17 +320,36 @@ def _sweep(args: argparse.Namespace) -> int: def _write_outputs(args: argparse.Namespace, out_dir: Path, rows: list[dict]) -> None: - (out_dir / "quality.json").write_text(json.dumps({ - "config": { - "model": args.model, "reference_quant": args.reference_quant, - "prompts": args.prompts, "seeds": args.seeds, "steps": args.steps, - "width": args.width, "height": args.height, "guidance": args.guidance, - "memory_mode": args.memory_mode, "clip": args.clip, - }, - "rows": rows, - }, indent = 2)) - fields = ["quant", "file_size_mib", "peak_vram_mib", "median_latency_s", - "mean_psnr", "mean_ssim", "mean_clip_text", "mean_clip_sim"] + (out_dir / "quality.json").write_text( + json.dumps( + { + "config": { + "model": args.model, + "reference_quant": args.reference_quant, + "prompts": args.prompts, + "seeds": args.seeds, + "steps": args.steps, + "width": args.width, + "height": args.height, + "guidance": args.guidance, + "memory_mode": args.memory_mode, + "clip": args.clip, + }, + "rows": rows, + }, + indent = 2, + ) + ) + fields = [ + "quant", + "file_size_mib", + "peak_vram_mib", + "median_latency_s", + "mean_psnr", + "mean_ssim", + "mean_clip_text", + "mean_clip_sim", + ] with (out_dir / "quality.csv").open("w", newline = "") as fh: writer = csv.DictWriter(fh, fieldnames = fields, extrasaction = "ignore") writer.writeheader() @@ -324,12 +359,17 @@ def _write_outputs(args: argparse.Namespace, out_dir: Path, rows: list[dict]) -> def _print_table(rows: list[dict]) -> None: - print("\n=== QUALITY vs QUANT (lower size/latency/VRAM better; higher PSNR/SSIM/CLIP better) ===", flush = True) + print( + "\n=== QUALITY vs QUANT (lower size/latency/VRAM better; higher PSNR/SSIM/CLIP better) ===", + flush = True, + ) hdr = f" {'quant':<28}{'size_MB':>9}{'vram_MB':>9}{'lat_s':>8}{'PSNR':>8}{'SSIM':>8}{'CLIPt':>8}{'CLIPs':>8}" print(hdr, flush = True) for r in rows: + def _f(v, fmt): return format(v, fmt) if isinstance(v, (int, float)) else "-" + psnr_str = "inf" if r.get("mean_psnr") == math.inf else _f(r.get("mean_psnr"), ".2f") print( f" {r['quant']:<28}{_f(r.get('file_size_mib'), '>9'):>9}" @@ -343,9 +383,11 @@ def _print_table(rows: list[dict]) -> None: def _recommend(args: argparse.Namespace, rows: list[dict]) -> None: # The smallest-on-disk non-reference quant that stays within the quality budget. passing = [ - r for r in rows + r + for r in rows if not r["is_reference"] - and r.get("mean_ssim") is not None and r["mean_ssim"] >= args.ssim_threshold + and r.get("mean_ssim") is not None + and r["mean_ssim"] >= args.ssim_threshold and (r.get("mean_psnr") is None or r["mean_psnr"] >= args.psnr_threshold) and r.get("file_size_mib") is not None ] @@ -355,9 +397,12 @@ def _recommend(args: argparse.Namespace, rows: list[dict]) -> None: print(" no candidate quant met the quality budget; keep the reference quant.", flush = True) return best = min(passing, key = lambda r: r["file_size_mib"]) - print(f" smallest quant within budget: {best['quant']} " - f"({best['file_size_mib']} MB, SSIM {best['mean_ssim']}, PSNR " - f"{'inf' if best['mean_psnr'] == math.inf else best['mean_psnr']})", flush = True) + print( + f" smallest quant within budget: {best['quant']} " + f"({best['file_size_mib']} MB, SSIM {best['mean_ssim']}, PSNR " + f"{'inf' if best['mean_psnr'] == math.inf else best['mean_psnr']})", + flush = True, + ) # ── self-test (CPU, no GPU/model) ───────────────────────────────────────────── @@ -371,17 +416,23 @@ def _selftest() -> int: base = rng.integers(0, 256, (128, 128, 3), dtype = np.uint8) a = Image.fromarray(base) b = Image.fromarray(base) # identical - noisy = Image.fromarray(np.clip(base.astype(int) + rng.integers(-40, 40, base.shape), 0, 255).astype(np.uint8)) + noisy = Image.fromarray( + np.clip(base.astype(int) + rng.integers(-40, 40, base.shape), 0, 255).astype(np.uint8) + ) checks = [] checks.append(("identical PSNR is inf", psnr(a, b) == math.inf)) checks.append(("identical SSIM ~ 1.0", abs(ssim(a, b) - 1.0) < 1e-9)) - checks.append(("noisy PSNR is finite + lower", math.isfinite(psnr(a, noisy)) and psnr(a, noisy) < 60)) + checks.append( + ("noisy PSNR is finite + lower", math.isfinite(psnr(a, noisy)) and psnr(a, noisy) < 60) + ) checks.append(("noisy SSIM < identical", ssim(a, noisy) < ssim(a, b))) checks.append(("shape mismatch -> 0", psnr(a, Image.fromarray(base[:64])) == 0.0)) # box mean of a constant field equals the constant const = np.full((32, 32), 7.0) - checks.append(("box mean of constant is constant", abs(_box_mean(const, 7).mean() - 7.0) < 1e-9)) + checks.append( + ("box mean of constant is constant", abs(_box_mean(const, 7).mean() - 7.0) < 1e-9) + ) ok = True for name, passed in checks: @@ -399,12 +450,24 @@ def _build_parser() -> argparse.ArgumentParser: description = "Image quality-vs-quant harness 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("--reference-quant", default = "z-image-turbo-BF16.gguf", - help = "high-fidelity reference GGUF filename") - p.add_argument("--quants", nargs = "*", default = [ - "z-image-turbo-Q8_0.gguf", "z-image-turbo-Q4_K_M.gguf", "z-image-turbo-Q2_K.gguf", - ], help = "candidate GGUF filenames to score against the reference") + p.add_argument( + "--model", default = "unsloth/Z-Image-Turbo-GGUF", help = "GGUF repo id or local path" + ) + p.add_argument( + "--reference-quant", + default = "z-image-turbo-BF16.gguf", + help = "high-fidelity reference GGUF filename", + ) + p.add_argument( + "--quants", + nargs = "*", + default = [ + "z-image-turbo-Q8_0.gguf", + "z-image-turbo-Q4_K_M.gguf", + "z-image-turbo-Q2_K.gguf", + ], + help = "candidate GGUF filenames to score against the reference", + ) p.add_argument("--base-repo", default = None) p.add_argument("--family-override", default = None) p.add_argument("--prompts", nargs = "*", default = DEFAULT_PROMPTS) @@ -415,10 +478,18 @@ def _build_parser() -> argparse.ArgumentParser: p.add_argument("--guidance", type = float, default = 0.0) p.add_argument("--memory-mode", default = None, choices = ["auto", "fast", "balanced", "low_vram"]) p.add_argument("--clip", action = "store_true", help = "also compute CLIP text + image scores") - p.add_argument("--psnr-threshold", type = float, default = 30.0, - help = "min mean PSNR (dB) vs reference for the recommendation") - p.add_argument("--ssim-threshold", type = float, default = 0.92, - help = "min mean SSIM vs reference for the recommendation") + p.add_argument( + "--psnr-threshold", + type = float, + default = 30.0, + help = "min mean PSNR (dB) vs reference for the recommendation", + ) + p.add_argument( + "--ssim-threshold", + type = float, + default = 0.92, + help = "min mean SSIM vs reference for the recommendation", + ) p.add_argument("--out-dir", default = "outputs/diffusion_quality") p.add_argument("--selftest", action = "store_true", help = "CPU metric sanity check; no GPU/model") return p diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index f1fafdf686..da0e0547bb 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -465,13 +465,20 @@ class DiffusionBackend: # 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, + pipe, + target, + is_gguf = bool(gguf_filename), + family = fam, + speed_mode = speed_mode or SPEED_OFF, + logger = logger, ) # Cast the dense companion text encoder(s) to fp8 storage (opt-in), # also before placement so the offload hooks move the smaller weights. fp8_cast = apply_fp8_text_encoder( - pipe, target, enable = text_encoder_fp8, logger = logger, + pipe, + target, + enable = text_encoder_fp8, + logger = logger, ) # Decide placement from MEASURED free device memory vs the model's @@ -508,7 +515,12 @@ class DiffusionBackend: logger.info( "diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s", - repo_id, base, device, effective_policy, effective_tiling, "; ".join(plan.reasons), + repo_id, + base, + device, + effective_policy, + effective_tiling, + "; ".join(plan.reasons), ) return self.status() diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 665c21fb9e..753b8d1bd6 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -177,7 +177,6 @@ def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]: def _xpu_memory() -> tuple[Optional[int], Optional[int]]: try: import torch - mem_get_info = getattr(getattr(torch, "xpu", None), "mem_get_info", None) if callable(mem_get_info): free, total = mem_get_info() @@ -191,7 +190,6 @@ def _system_memory_mib() -> tuple[Optional[int], Optional[int]]: """(total, available) host RAM in MiB, via psutil then POSIX sysconf.""" try: import psutil - vm = psutil.virtual_memory() return int(vm.total // (1024 * 1024)), int(vm.available // (1024 * 1024)) except Exception: @@ -212,7 +210,6 @@ def file_size_mib(path: Any) -> Optional[int]: """On-disk size of ``path`` in MiB, or None if it can't be stat'd.""" try: from pathlib import Path - return max(1, int(Path(path).expanduser().stat().st_size // (1024 * 1024))) except Exception: return None @@ -421,7 +418,11 @@ def plan_diffusion_memory( def apply_memory_plan( - pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None, + pipe: Any, + plan: MemoryPlan, + *, + device: str, + logger: Any = None, ) -> tuple[str, bool]: """Apply ``plan`` to a freshly built diffusers pipeline: enable the VAE memory savers then place / offload the weights. Exactly one placement call runs, so the @@ -452,7 +453,8 @@ def apply_memory_plan( if logger is not None: logger.warning( "diffusion.memory: sequential offload failed (%s); " - "falling back to whole-module offload", exc, + "falling back to whole-module offload", + exc, ) pipe.enable_model_cpu_offload() policy = OFFLOAD_MODEL @@ -511,6 +513,7 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: if logger is not None: logger.warning( "diffusion.memory: group offload failed (%s); falling back to " - "whole-module offload", exc, + "whole-module offload", + exc, ) return False diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index f354373ec2..de675e66cf 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -37,14 +37,17 @@ def fp8_text_encoder_supported(target: Any) -> bool: return False try: import torch - return getattr(target, "dtype", None) is torch.bfloat16 and hasattr(torch, "float8_e4m3fn") except Exception: return False def apply_fp8_text_encoder( - pipe: Any, target: Any, *, enable: bool, logger: Any = None, + pipe: Any, + target: Any, + *, + enable: bool, + logger: Any = None, ) -> list[str]: """Cast each text encoder's linear weights to fp8 storage (compute stays the target dtype). Returns the names of the encoders actually cast (empty when diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index a731c3124f..42c3b4d5be 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -62,7 +62,6 @@ def compile_eligible(target: Any, *, is_gguf: bool, family: Any) -> bool: def _is_bfloat16(dtype: Any) -> bool: try: import torch - return dtype is torch.bfloat16 except Exception: return str(dtype).endswith("bfloat16") @@ -107,7 +106,6 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False try: import torch - vae.to(memory_format = torch.channels_last) return True except Exception as exc: # noqa: BLE001 — optimisation only diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 24cc1ac513..fb309d7ff3 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1704,21 +1704,21 @@ class DiffusionLoadRequest(BaseModel): memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field( None, description = "Memory policy: auto (measured), fast (resident), balanced " - "(stream the transformer, near-resident speed, moderate VRAM " - "cut), low_vram (offload every component, lowest VRAM, slower). " - "Overrides cpu_offload when set.", + "(stream the transformer, near-resident speed, moderate VRAM " + "cut), low_vram (offload every component, lowest VRAM, slower). " + "Overrides cpu_offload when set.", ) speed_mode: Optional[Literal["off", "default", "max"]] = Field( None, description = "Opt-in speed optims (default off -> bit-identical output): " - "default (channels_last + regional torch.compile where eligible), " - "max (also TF32 + fused QKV).", + "default (channels_last + regional torch.compile where eligible), " + "max (also TF32 + fused QKV).", ) text_encoder_fp8: bool = Field( False, description = "Cast the companion text encoder(s) to fp8 storage (~2x smaller, " - "CUDA + bf16 only). A memory-vs-quality tradeoff (shifts fine " - "detail), not free; pairs well with balanced mode.", + "CUDA + bf16 only). A memory-vs-quality tradeoff (shifts fine " + "detail), not free; pairs well with balanced mode.", ) diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 39ffdcc9b9..6c82d528f9 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -828,7 +828,9 @@ def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path, assert pipe.offloaded is True and pipe.moved_to is None # offload owns placement -def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(fake_runtime, tmp_path, monkeypatch): +def test_load_explicit_cpu_offload_engages_model_offload_on_cuda( + fake_runtime, tmp_path, monkeypatch +): # cpu_offload=True with no mode: auto would stay resident (budget unknown under # the stub), but the explicit flag forces whole-module offload. (tmp_path / "m.gguf").write_bytes(b"x") diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index 357425c6c5..858e87ffa1 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -36,7 +36,12 @@ from core.inference.diffusion_memory import ( ) -def _target(*, device = "cuda", backend = "cuda", supports_offload = True): +def _target( + *, + device = "cuda", + backend = "cuda", + supports_offload = True, +): """A duck-typed stand-in for DiffusionDeviceTarget (only the fields the planner / snapshot read).""" return types.SimpleNamespace( @@ -221,18 +226,36 @@ def test_auto_stays_resident_when_budget_unknown(): def test_explicit_modes_force_policy_regardless_of_budget(): roomy = _discrete(80000) - assert plan_diffusion_memory( - target = _target(), device_memory = roomy, model_dense_mib = 1000, - runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_FAST, - ).offload_policy == OFFLOAD_NONE - assert plan_diffusion_memory( - target = _target(), device_memory = roomy, model_dense_mib = 1000, - runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_BALANCED, - ).offload_policy == OFFLOAD_GROUP - assert plan_diffusion_memory( - target = _target(), device_memory = roomy, model_dense_mib = 1000, - runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_LOW_VRAM, - ).offload_policy == OFFLOAD_MODEL + assert ( + plan_diffusion_memory( + target = _target(), + device_memory = roomy, + model_dense_mib = 1000, + runtime_headroom_mib = 1000, + requested_mode = MEMORY_MODE_FAST, + ).offload_policy + == OFFLOAD_NONE + ) + assert ( + plan_diffusion_memory( + target = _target(), + device_memory = roomy, + model_dense_mib = 1000, + runtime_headroom_mib = 1000, + requested_mode = MEMORY_MODE_BALANCED, + ).offload_policy + == OFFLOAD_GROUP + ) + assert ( + plan_diffusion_memory( + target = _target(), + device_memory = roomy, + model_dense_mib = 1000, + runtime_headroom_mib = 1000, + requested_mode = MEMORY_MODE_LOW_VRAM, + ).offload_policy + == OFFLOAD_MODEL + ) def test_fast_falls_back_to_model_offload_when_it_does_not_fit(): @@ -381,8 +404,10 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule(): class _VaeOnly: def __init__(self): self.vae = types.SimpleNamespace( - tiled = False, sliced = False, - enable_tiling = self._tile, enable_slicing = self._slice, + tiled = False, + sliced = False, + enable_tiling = self._tile, + enable_slicing = self._slice, ) def _tile(self): @@ -409,7 +434,9 @@ def test_apply_group_falls_back_to_model_without_transformer(): def test_apply_sequential_offload(): pipe = _RecordingPipe() - effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") + effective, _ = apply_memory_plan( + pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda" + ) assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls assert effective == OFFLOAD_SEQUENTIAL @@ -422,7 +449,9 @@ def test_apply_sequential_falls_back_to_model_offload_when_unsupported(): raise RuntimeError("sequential offload not supported for this transformer") pipe = _NoSeqPipe() - effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda") + effective, _ = apply_memory_plan( + pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda" + ) assert effective == OFFLOAD_MODEL assert "model_offload" in pipe.calls diff --git a/studio/backend/tests/test_diffusion_speed.py b/studio/backend/tests/test_diffusion_speed.py index 72a4f7899d..ade37482ba 100644 --- a/studio/backend/tests/test_diffusion_speed.py +++ b/studio/backend/tests/test_diffusion_speed.py @@ -24,9 +24,16 @@ from core.inference.diffusion_speed import ( ) -def _target(*, device = "cuda", dtype = "bfloat16", compile_ok = True): +def _target( + *, + device = "cuda", + dtype = "bfloat16", + compile_ok = True, +): return types.SimpleNamespace( - device = device, dtype = dtype, supports_default_torch_compile = compile_ok, + device = device, + dtype = dtype, + supports_default_torch_compile = compile_ok, ) @@ -78,7 +85,12 @@ def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch): class _Pipe: - def __init__(self, *, with_compile = False, with_fuse = False) -> None: + def __init__( + self, + *, + with_compile = False, + with_fuse = False, + ) -> None: self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to) self.transformer = types.SimpleNamespace() if with_compile: @@ -144,8 +156,11 @@ def test_speed_max_tf32_only_on_cuda(monkeypatch): _stub_torch(monkeypatch) pipe = _Pipe() applied = apply_speed_optims( - pipe, _target(device = "mps", compile_ok = False), is_gguf = True, - family = _family(), speed_mode = SPEED_MAX, + pipe, + _target(device = "mps", compile_ok = False), + is_gguf = True, + family = _family(), + speed_mode = SPEED_MAX, ) assert applied["tf32"] is False # not CUDA -> no TF32