diff --git a/scripts/fbcache_flux_probe.py b/scripts/fbcache_flux_probe.py index 2b41657dcf..31d96a1c40 100644 --- a/scripts/fbcache_flux_probe.py +++ b/scripts/fbcache_flux_probe.py @@ -28,8 +28,9 @@ def _lpips(ref, arr): try: import lpips import torch + if _LP["fn"] is None: - _LP["fn"] = lpips.LPIPS(net="alex", verbose=False).cuda().eval() + _LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval() def t(x): return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda() @@ -37,7 +38,7 @@ def _lpips(ref, arr): with torch.no_grad(): return float(_LP["fn"](t(ref), t(arr)).item()) except Exception as exc: # noqa: BLE001 - print(f" (lpips: {type(exc).__name__})", flush=True) + print(f" (lpips: {type(exc).__name__})", flush = True) return None @@ -45,17 +46,28 @@ def _load(): import os import diffusers import torch - pipe = diffusers.FluxPipeline.from_pretrained(BASE, torch_dtype=torch.bfloat16, token=os.environ.get("HF_TOKEN")) + + pipe = diffusers.FluxPipeline.from_pretrained( + BASE, torch_dtype = torch.bfloat16, token = os.environ.get("HF_TOKEN") + ) pipe.to("cuda") return pipe def _gen(pipe, steps, seed, res, guidance): import torch - 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=guidance, generator=g).images[0] + + 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 = guidance, + generator = g, + ).images[0] torch.cuda.synchronize() return img, time.time() - t0 @@ -64,72 +76,97 @@ def _median(xs): return sorted(xs)[len(xs) // 2] -def run(tag, steps, seed, res, guidance, iters, *, threshold=None, compile_=True): +def run( + tag, + steps, + seed, + res, + guidance, + iters, + *, + threshold = None, + compile_ = True, +): import torch - torch.compiler.reset(); torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats() + + torch.compiler.reset() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() pipe = _load() if threshold is not None: from diffusers import FirstBlockCacheConfig try: - pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=threshold)) + pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold = threshold)) except Exception as exc: # noqa: BLE001 from diffusers.hooks import apply_first_block_cache - apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=threshold)) + apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold)) if compile_: try: - pipe.transformer.compile_repeated_blocks(fullgraph=True, dynamic=True) + pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True) except Exception as exc: # noqa: BLE001 - print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush=True) + print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush = True) try: _gen(pipe, steps, seed, res, guidance) # warmup / compile except Exception as exc: # noqa: BLE001 - import traceback; traceback.print_exc() - print(f" [{tag}] FAILED: {type(exc).__name__}: {str(exc)[:100]}", flush=True) - del pipe; torch.cuda.empty_cache(); return None + import traceback + + traceback.print_exc() + print(f" [{tag}] FAILED: {type(exc).__name__}: {str(exc)[:100]}", flush = True) + del pipe + torch.cuda.empty_cache() + return None dts, img = [], None for _ in range(iters): - img, dt = _gen(pipe, steps, seed, res, guidance); dts.append(dt) + img, dt = _gen(pipe, steps, seed, res, guidance) + dts.append(dt) peak = torch.cuda.max_memory_allocated() / 1e9 arr = np.array(img) - OUT.mkdir(parents=True, exist_ok=True); img.save(OUT / f"{tag}.png") - del pipe; torch.cuda.empty_cache() + OUT.mkdir(parents = True, exist_ok = True) + img.save(OUT / f"{tag}.png") + del pipe + torch.cuda.empty_cache() return _median(dts), arr, peak -def main(argv=None) -> int: +def main(argv = None) -> int: p = argparse.ArgumentParser() - p.add_argument("--steps", type=int, default=28) - p.add_argument("--res", type=int, default=1024) - p.add_argument("--seed", type=int, default=42) - p.add_argument("--guidance", type=float, default=3.5) - p.add_argument("--iters", type=int, default=2) + p.add_argument("--steps", type = int, default = 28) + p.add_argument("--res", type = int, default = 1024) + p.add_argument("--seed", type = int, default = 42) + p.add_argument("--guidance", type = float, default = 3.5) + p.add_argument("--iters", type = int, default = 2) args = p.parse_args(argv) s, r, seed, gd, it = args.steps, args.res, args.seed, args.guidance, args.iters - print(f"== FBCache on Flux.1-dev ({r}px, {s} steps, guidance {gd}) ==", flush=True) + print(f"== FBCache on Flux.1-dev ({r}px, {s} steps, guidance {gd}) ==", flush = True) base = run("baseline", s, seed, r, gd, it) if base is None: - print("baseline FAILED", flush=True); return 1 + print("baseline FAILED", flush = True) + return 1 bmed, ref, bpeak = base - print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush=True) + print(f" baseline {bmed:.3f}s peak={bpeak:.1f}G", flush = True) rows = [("baseline", bmed, bpeak, 0.0)] for thr in (0.08, 0.12, 0.20): - out = run(f"fbcache_{thr}", s, seed, r, gd, it, threshold=thr) + out = run(f"fbcache_{thr}", s, seed, r, gd, it, threshold = thr) if out is None: - rows.append((f"fbcache_{thr}", None, None, None)); continue + rows.append((f"fbcache_{thr}", None, None, None)) + continue med, arr, peak = out lp = _lpips(ref, arr) rows.append((f"fbcache_{thr}", med, peak, lp)) - print(f" fbcache_{thr}: {med:.3f}s ({bmed/med:.2f}x) peak={peak:.1f}G LPIPS={lp}", flush=True) + print( + f" fbcache_{thr}: {med:.3f}s ({bmed/med:.2f}x) peak={peak:.1f}G LPIPS={lp}", flush = True + ) - print("\n==== SUMMARY (Flux.1-dev, ref = no-cache compile) ====", flush=True) + print("\n==== SUMMARY (Flux.1-dev, ref = no-cache compile) ====", flush = True) for tag, med, peak, lp in rows: if med is None: - print(f" {tag:16s} FAILED"); continue + print(f" {tag:16s} FAILED") + continue spd = f"{bmed/med:.2f}x" lpv = "ref" if tag == "baseline" else (f"{lp:.3f}" if lp is not None else "n/a") - print(f" {tag:16s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush=True) - print("FBCACHE-FLUX-DONE", flush=True) + print(f" {tag:16s} {med:.3f}s {spd:>6s} peak={peak:.1f}G LPIPS={lpv:>6s}", flush = True) + print("FBCACHE-FLUX-DONE", flush = True) return 0 diff --git a/studio/backend/core/inference/diffusion_cache.py b/studio/backend/core/inference/diffusion_cache.py index 2c67358ea9..4fd5815e16 100644 --- a/studio/backend/core/inference/diffusion_cache.py +++ b/studio/backend/core/inference/diffusion_cache.py @@ -72,19 +72,20 @@ def apply_step_cache( transformer = getattr(pipe, "transformer", None) if transformer is None: return None - thr = threshold if threshold is not None else ( - QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD + thr = ( + threshold + if threshold is not None + else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD) ) try: from diffusers import FirstBlockCacheConfig - config = FirstBlockCacheConfig(threshold=thr) + config = FirstBlockCacheConfig(threshold = thr) enable_cache = getattr(transformer, "enable_cache", None) if callable(enable_cache): enable_cache(config) else: from diffusers.hooks import apply_first_block_cache - apply_first_block_cache(transformer, config) try: transformer._unsloth_step_cache = f"{mode}@{thr}" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index eec529000a..3c34c47a45 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1903,6 +1903,4 @@ class DiffusionStatusResponse(BaseModel): description = "Attention backend engaged via the diffusers dispatcher (e.g. " "_native_cudnn), or null for the default SDPA", ) - transformer_cache: Optional[str] = Field( - None, description = "Step cache engaged: fbcache | null" - ) + transformer_cache: Optional[str] = Field(None, description = "Step cache engaged: fbcache | null") diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index c9c83009ee..4b572adcf7 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -402,7 +402,11 @@ def test_transformer_cache_threads_through(client, monkeypatch): def test_invalid_transformer_cache_returns_422(client): resp = client.post( "/api/inference/images/load", - json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache": "deepcache"}, + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache": "deepcache", + }, ) assert resp.status_code == 422 @@ -410,7 +414,11 @@ def test_invalid_transformer_cache_returns_422(client): def test_out_of_range_cache_threshold_returns_422(client): resp = client.post( "/api/inference/images/load", - json = {"model_path": "x/z-image", "gguf_filename": "q.gguf", "transformer_cache_threshold": 1.5}, + json = { + "model_path": "x/z-image", + "gguf_filename": "q.gguf", + "transformer_cache_threshold": 1.5, + }, ) assert resp.status_code == 422