ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-06-26 12:36:45 +00:00
commit ca52807680
4 changed files with 88 additions and 44 deletions
Download patch file Download diff file Expand all files Collapse all files

107
scripts/fbcache_flux_probe.py
View file

@ -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

9
studio/backend/core/inference/diffusion_cache.py
View file

@ -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}"

4
studio/backend/models/inference.py
View file

@ -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")

12
studio/backend/tests/test_diffusion_routes.py
View file

@ -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