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 03:19:59 +00:00
commit 141cb3ae10
7 changed files with 133 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

77
scripts/compile_probe.py
View file

@ -40,8 +40,12 @@ def _gen(pipe, prompt, *, steps, seed, width, height, guidance):
torch.cuda.synchronize()
t0 = time.time()
image = pipe(
prompt = prompt, width = width, height = height,
num_inference_steps = steps, guidance_scale = guidance, generator = gen,
prompt = prompt,
width = width,
height = height,
num_inference_steps = steps,
guidance_scale = guidance,
generator = gen,
).images[0]
torch.cuda.synchronize()
return image, time.time() - t0
@ -54,14 +58,21 @@ def main(argv = None) -> int:
p.add_argument("--base-repo", default = "Tongyi-MAI/Z-Image-Turbo")
p.add_argument("--transformer-class", default = "ZImageTransformer2DModel")
p.add_argument("--pipeline-class", default = "ZImagePipeline")
p.add_argument("--prompt", default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed")
p.add_argument(
"--prompt",
default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed",
)
p.add_argument("--steps", type = int, default = 8)
p.add_argument("--seed", type = int, default = 42)
p.add_argument("--width", type = int, default = 1024)
p.add_argument("--height", type = int, default = 1024)
p.add_argument("--guidance", type = float, default = 0.0)
p.add_argument("--mode", default = "default", help = "compile mode: default | max-autotune-no-cudagraphs")
p.add_argument("--dynamic", action = "store_true", help = "dynamic=True (default False here for speed)")
p.add_argument(
"--mode", default = "default", help = "compile mode: default | max-autotune-no-cudagraphs"
)
p.add_argument(
"--dynamic", action = "store_true", help = "dynamic=True (default False here for speed)"
)
p.add_argument("--out-dir", default = "outputs/compile_probe")
args = p.parse_args(argv)
@ -80,7 +91,9 @@ def main(argv = None) -> int:
transformer = transformer_cls.from_single_file(
gguf_path,
quantization_config = diffusers.GGUFQuantizationConfig(compute_dtype = dtype),
torch_dtype = dtype, config = args.base_repo, subfolder = "transformer",
torch_dtype = dtype,
config = args.base_repo,
subfolder = "transformer",
)
pipeline_cls = getattr(diffusers, args.pipeline_class)
pipe = pipeline_cls.from_pretrained(args.base_repo, torch_dtype = dtype, transformer = transformer)
@ -88,8 +101,24 @@ def main(argv = None) -> int:
print("pipeline loaded on cuda", flush = True)
# warm the eager path once (allocator / cudnn), then time eager.
_gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance)
eager_img, eager_t = _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance)
_gen(
pipe,
args.prompt,
steps = args.steps,
seed = args.seed,
width = args.width,
height = args.height,
guidance = args.guidance,
)
eager_img, eager_t = _gen(
pipe,
args.prompt,
steps = args.steps,
seed = args.seed,
width = args.width,
height = args.height,
guidance = args.guidance,
)
eager_img.save(out / "eager.png")
eager_arr = np.array(eager_img)
print(f"EAGER: {eager_t:.2f}s/gen", flush = True)
@ -106,7 +135,10 @@ def main(argv = None) -> int:
try:
t0 = time.time()
fn(**compile_kwargs)
print(f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)", flush = True)
print(
f" compile_repeated_blocks() returned in {time.time()-t0:.1f}s (compilation is lazy)",
flush = True,
)
except Exception as exc: # noqa: BLE001
print(f"RESULT: compile_repeated_blocks RAISED: {type(exc).__name__}: {exc}", flush = True)
return 1
@ -114,13 +146,29 @@ def main(argv = None) -> int:
# first compiled gen triggers the actual compilation (untimed warmup).
try:
t0 = time.time()
_gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance)
_gen(
pipe,
args.prompt,
steps = args.steps,
seed = args.seed,
width = args.width,
height = args.height,
guidance = args.guidance,
)
print(f" first compiled gen (compilation) took {time.time()-t0:.1f}s", flush = True)
except Exception as exc: # noqa: BLE001
print(f"RESULT: first compiled generation RAISED: {type(exc).__name__}: {exc}", flush = True)
return 2
comp_img, comp_t = _gen(pipe, args.prompt, steps = args.steps, seed = args.seed, width = args.width, height = args.height, guidance = args.guidance)
comp_img, comp_t = _gen(
pipe,
args.prompt,
steps = args.steps,
seed = args.seed,
width = args.width,
height = args.height,
guidance = args.guidance,
)
comp_img.save(out / "compiled.png")
psnr = _psnr(eager_arr, np.array(comp_img))
@ -129,8 +177,11 @@ def main(argv = None) -> int:
print(f" eager: {eager_t:.2f}s/gen", flush = True)
print(f" compiled: {comp_t:.2f}s/gen ({speedup:+.1f}% vs eager)", flush = True)
print(f" PSNR(compiled vs eager): {psnr:.1f} dB", flush = True)
print(f" verdict: {'COMPILE-WORKS' if psnr >= 30 else 'COMPILE-DIVERGES'} "
f"{'FASTER' if comp_t < eager_t else 'NOT-FASTER'}", flush = True)
print(
f" verdict: {'COMPILE-WORKS' if psnr >= 30 else 'COMPILE-DIVERGES'} "
f"{'FASTER' if comp_t < eager_t else 'NOT-FASTER'}",
flush = True,
)
return 0

2
scripts/diffusion_bench.py
View file

@ -448,7 +448,7 @@ def _build_parser() -> argparse.ArgumentParser:
default = None,
choices = ["off", "default", "max"],
help = "speed profile: off is bit-identical; default adds compile + "
"cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV",
"cudnn.benchmark (near-lossless); max also adds TF32 + fused QKV",
)
p.add_argument(
"--text-encoder-quant",

58
scripts/perf_verify.py
View file

@ -42,7 +42,10 @@ def main(argv = None) -> int:
p = argparse.ArgumentParser()
p.add_argument("--model", default = "unsloth/Z-Image-Turbo-GGUF")
p.add_argument("--gguf", default = "z-image-turbo-Q4_K_M.gguf")
p.add_argument("--prompt", default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed")
p.add_argument(
"--prompt",
default = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed",
)
p.add_argument("--steps", type = int, default = 8)
p.add_argument("--seed", type = int, default = 42)
p.add_argument("--width", type = int, default = 1024)
@ -62,8 +65,11 @@ def main(argv = None) -> int:
def load(mode_speed = None, mode_mem = None):
backend.begin_load(
args.model, gguf_filename = args.gguf, hf_token = token,
speed_mode = mode_speed, memory_mode = mode_mem,
args.model,
gguf_filename = args.gguf,
hf_token = token,
speed_mode = mode_speed,
memory_mode = mode_mem,
)
deadline = time.time() + 2400
while time.time() < deadline:
@ -79,13 +85,25 @@ def main(argv = None) -> int:
torch.cuda.synchronize()
t0 = time.time()
img = backend.generate(
prompt = args.prompt, width = args.width, height = args.height,
steps = args.steps, guidance = 0.0, seed = args.seed, batch_size = 1,
prompt = args.prompt,
width = args.width,
height = args.height,
steps = args.steps,
guidance = 0.0,
seed = args.seed,
batch_size = 1,
)["images"][0]
torch.cuda.synchronize()
return img, time.time() - t0
def timed(mode_speed, *, warmup, iters, mem = None, tag = ""):
def timed(
mode_speed,
*,
warmup,
iters,
mem = None,
tag = "",
):
st = load(mode_speed, mem)
for _ in range(warmup):
gen()
@ -97,29 +115,41 @@ def main(argv = None) -> int:
img.save(out / f"{tag}.png")
backend.unload()
med = sorted(lats)[len(lats) // 2]
print(f" [{tag}] speed={mode_speed} mem={mem} optims={st.get('speed_optims')} "
f"tiling={st.get('vae_tiling')} median={med:.3f}s", flush = True)
print(
f" [{tag}] speed={mode_speed} mem={mem} optims={st.get('speed_optims')} "
f"tiling={st.get('vae_tiling')} median={med:.3f}s",
flush = True,
)
return np.array(img), med
print("== 1. speed: off vs default ==", flush = True)
off_img, off_t = timed("off", warmup = 1, iters = 3, tag = "off")
def_img, def_t = timed("default", warmup = 1, iters = 3, tag = "default")
print(f" PSNR(default vs off) = {_psnr(off_img, def_img):.1f} dB", flush = True)
print(f" speedup: off {off_t:.3f}s -> default {def_t:.3f}s "
f"({(off_t-def_t)/off_t*100:+.1f}%)", flush = True)
print(
f" speedup: off {off_t:.3f}s -> default {def_t:.3f}s "
f"({(off_t-def_t)/off_t*100:+.1f}%)",
flush = True,
)
print("== 2. TF32-leak fix: max then off must be byte-identical ==", flush = True)
timed("max", warmup = 0, iters = 1, tag = "max") # flips + should restore globals
off2_img, _ = timed("off", warmup = 0, iters = 1, tag = "off2")
leak_psnr = _psnr(off_img, off2_img)
print(f" PSNR(off-after-max vs off) = {leak_psnr:.1f} dB "
f"({'OK byte-identical' if leak_psnr == float('inf') else 'LEAK! globals not restored'})", flush = True)
print(
f" PSNR(off-after-max vs off) = {leak_psnr:.1f} dB "
f"({'OK byte-identical' if leak_psnr == float('inf') else 'LEAK! globals not restored'})",
flush = True,
)
print("== 3. balanced is bit-identical (tiling off) ==", flush = True)
bal_img, bal_t = timed("off", warmup = 0, iters = 1, mem = "balanced", tag = "balanced")
bal_psnr = _psnr(off_img, bal_img)
print(f" PSNR(balanced vs off) = {bal_psnr:.1f} dB "
f"({'OK bit-identical' if bal_psnr == float('inf') else 'differs'})", flush = True)
print(
f" PSNR(balanced vs off) = {bal_psnr:.1f} dB "
f"({'OK bit-identical' if bal_psnr == float('inf') else 'differs'})",
flush = True,
)
ok = (leak_psnr == float("inf")) and (def_t < off_t) and (_psnr(off_img, def_img) >= 30)
print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True)

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

@ -45,7 +45,6 @@ def snapshot_backend_flags() -> Optional[dict]:
caller can restore them on unload. None if torch is unavailable."""
try:
import torch
return {
"matmul_tf32": bool(torch.backends.cuda.matmul.allow_tf32),
"cudnn_tf32": bool(torch.backends.cudnn.allow_tf32),
@ -134,8 +133,11 @@ def apply_speed_optims(
BEFORE placement / offload. Returns which optimisations actually engaged. Every
step is best-effort: a pipeline that doesn't support one is simply skipped."""
applied = {
"channels_last": False, "cudnn_benchmark": False,
"tf32": False, "fused_qkv": False, "compiled": False,
"channels_last": False,
"cudnn_benchmark": False,
"tf32": False,
"fused_qkv": False,
"compiled": False,
}
mode = normalize_speed_mode(speed_mode)
if mode == SPEED_OFF:
@ -192,7 +194,6 @@ def _compile_repeated_blocks(pipe: Any, logger: Any) -> bool:
def _enable_cudnn_benchmark(logger: Any) -> bool:
try:
import torch
torch.backends.cudnn.benchmark = True
return True
except Exception as exc: # noqa: BLE001 — optimisation only

4
studio/backend/core/inference/sd_cpp_args.py
View file

@ -140,9 +140,7 @@ def native_speed_flags(speed_mode: Optional[str]) -> list[str]:
return ["--diffusion-fa"]
if mode == NATIVE_SPEED_MAX:
return ["--diffusion-fa", "--diffusion-conv-direct"]
raise ValueError(
f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'"
)
raise ValueError(f"native speed_mode must be one of {NATIVE_SPEED_MODES}, got '{speed_mode}'")
def offload_flags(

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

@ -212,9 +212,14 @@ class SdCppEngine:
speed = [f for f in native_speed_flags(native_speed) if f not in offload]
merged_extra = speed + list(extra_args or [])
cmd = build_sd_cpp_command(
self._require_binary(), files, params,
output_path = str(self._prepare_out(output_path)), offload = offload,
threads = threads, verbose = verbose, extra_args = merged_extra,
self._require_binary(),
files,
params,
output_path = str(self._prepare_out(output_path)),
offload = offload,
threads = threads,
verbose = verbose,
extra_args = merged_extra,
)
return self._run(cmd, output_path, timeout = timeout, env = env, on_log = on_log)

14
studio/backend/tests/test_diffusion_speed.py
View file

@ -152,8 +152,11 @@ def test_speed_off_applies_nothing(monkeypatch):
pipe, _target(), is_gguf = False, family = _family(), speed_mode = SPEED_OFF
)
assert applied == {
"channels_last": False, "cudnn_benchmark": False,
"tf32": False, "fused_qkv": False, "compiled": False,
"channels_last": False,
"cudnn_benchmark": False,
"tf32": False,
"fused_qkv": False,
"compiled": False,
}
assert pipe.vae.mem_format is None and pipe.compiled is False
# off must not touch any process-wide flag (bit-identical reference path).
@ -188,8 +191,11 @@ def test_speed_default_cudnn_benchmark_only_on_cuda(monkeypatch):
_stub_torch(monkeypatch)
pipe = _Pipe(with_compile = True)
applied = apply_speed_optims(
pipe, _target(device = "mps", compile_ok = False), is_gguf = True,
family = _family(), speed_mode = SPEED_DEFAULT,
pipe,
_target(device = "mps", compile_ok = False),
is_gguf = True,
family = _family(),
speed_mode = SPEED_DEFAULT,
)
assert applied["cudnn_benchmark"] is False # not CUDA -> no autotune flip