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 05:04:39 +00:00
commit ff60263834
3 changed files with 72 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

53
scripts/compare_engines.py
View file

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

28
scripts/leverage_probe.py
View file

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

11
studio/backend/core/inference/diffusion_speed.py
View file

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