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-25 11:30:37 +00:00
commit 7f6e69dd5c
6 changed files with 157 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

131
scripts/diffusion_bench.py
View file

@ -62,7 +62,9 @@ def _git_commit() -> Optional[str]:
out = subprocess.run(
["git", "rev-parse", "HEAD"],
cwd = str(Path(__file__).resolve().parent),
capture_output = True, text = True, timeout = 10,
capture_output = True,
text = True,
timeout = 10,
)
return out.stdout.strip() or None if out.returncode == 0 else None
except Exception:
@ -85,40 +87,34 @@ def _is_cuda(device: Optional[str]) -> bool:
def _cuda_reset_peak() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
def _cuda_sync() -> None:
import torch
if torch.cuda.is_available():
torch.cuda.synchronize()
def _cuda_peak_alloc() -> Optional[int]:
import torch
return int(torch.cuda.max_memory_allocated()) if torch.cuda.is_available() else None
def _cuda_peak_reserved() -> Optional[int]:
import torch
return int(torch.cuda.max_memory_reserved()) if torch.cuda.is_available() else None
def _cuda_alloc() -> Optional[int]:
import torch
return int(torch.cuda.memory_allocated()) if torch.cuda.is_available() else None
def _gpu_name() -> Optional[str]:
try:
import torch
if torch.cuda.is_available():
return torch.cuda.get_device_name(0)
except Exception:
@ -130,13 +126,11 @@ def _versions() -> dict[str, Optional[str]]:
out: dict[str, Optional[str]] = {"torch": None, "diffusers": None}
try:
import torch
out["torch"] = torch.__version__
except Exception:
pass
try:
import diffusers
out["diffusers"] = diffusers.__version__
except Exception:
pass
@ -255,7 +249,9 @@ def _run(args: argparse.Namespace) -> dict[str, Any]:
"latencies_s": [round(x, 4) for x in latencies],
"median_latency_s": round(_percentile(latencies, 50), 4),
"p90_latency_s": round(_percentile(latencies, 90), 4),
"images_per_sec": round((args.batch_size * len(latencies)) / total, 4) if total > 0 else None,
"images_per_sec": round((args.batch_size * len(latencies)) / total, 4)
if total > 0
else None,
"peak_vram_bytes": _cuda_peak_alloc(),
}
@ -282,10 +278,17 @@ 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,
"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,
},
}
@ -301,8 +304,10 @@ def _write_baseline(args: argparse.Namespace) -> int:
metrics = _run(args)
metrics["accuracy"] = {
"reference_png": str(ref_png),
"width": args.width, "height": args.height,
"steps": args.steps, "seed": args.seed,
"width": args.width,
"height": args.height,
"steps": args.steps,
"seed": args.seed,
"dtype": (metrics["env"]["status"] or {}).get("dtype"),
}
@ -312,10 +317,13 @@ def _write_baseline(args: argparse.Namespace) -> int:
print(f" json: {baseline_path}", flush = True)
print(f" reference: {ref_png}", flush = True)
print(f" load: {metrics['load']}", flush = True)
print(f" generate: median={metrics['generate'].get('median_latency_s')}s "
f"p90={metrics['generate'].get('p90_latency_s')}s "
f"img/s={metrics['generate'].get('images_per_sec')} "
f"peak_vram={metrics['generate'].get('peak_vram_bytes')}", flush = True)
print(
f" generate: median={metrics['generate'].get('median_latency_s')}s "
f"p90={metrics['generate'].get('p90_latency_s')}s "
f"img/s={metrics['generate'].get('images_per_sec')} "
f"peak_vram={metrics['generate'].get('peak_vram_bytes')}",
flush = True,
)
return 0
@ -362,14 +370,22 @@ def _compare(args: argparse.Namespace) -> int:
print("\n=== REGRESSION REPORT ===", flush = True)
print(f" {'metric':<22}{'baseline':>16}{'current':>16}{'delta':>12}", flush = True)
print(f" {'median_latency_s':<22}{base_median:>16.4f}{cur_median:>16.4f}{latency_reg * 100:>11.1f}%", flush = True)
print(
f" {'median_latency_s':<22}{base_median:>16.4f}{cur_median:>16.4f}{latency_reg * 100:>11.1f}%",
flush = True,
)
if base_peak and cur_peak:
print(f" {'peak_vram_MB':<22}{base_peak / 1e6:>16.1f}{cur_peak / 1e6:>16.1f}{vram_reg * 100:>11.1f}%", flush = True)
print(
f" {'peak_vram_MB':<22}{base_peak / 1e6:>16.1f}{cur_peak / 1e6:>16.1f}{vram_reg * 100:>11.1f}%",
flush = True,
)
print(f" {'psnr_dB(vs ref)':<22}{'-':>16}{psnr:>16.2f}{'':>12}", flush = True)
failures = []
if latency_reg > args.max_latency_regression:
failures.append(f"latency +{latency_reg * 100:.1f}% > {args.max_latency_regression * 100:.0f}%")
failures.append(
f"latency +{latency_reg * 100:.1f}% > {args.max_latency_regression * 100:.0f}%"
)
if base_peak and cur_peak and vram_reg > args.max_vram_regression:
failures.append(f"peak VRAM +{vram_reg * 100:.1f}% > {args.max_vram_regression * 100:.0f}%")
if not math.isnan(psnr) and psnr < args.min_psnr:
@ -390,15 +406,21 @@ def _build_parser() -> argparse.ArgumentParser:
description = "Benchmark + regression guard 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("--gguf", default = "z-image-turbo-Q4_K_M.gguf",
help = "transformer GGUF filename inside --model")
p.add_argument(
"--model", default = "unsloth/Z-Image-Turbo-GGUF", help = "GGUF repo id or local path"
)
p.add_argument(
"--gguf",
default = "z-image-turbo-Q4_K_M.gguf",
help = "transformer GGUF filename inside --model",
)
p.add_argument("--base-repo", default = None, help = "override the diffusers base repo")
p.add_argument("--family-override", default = None, help = "force a diffusion family")
p.add_argument("--prompt",
default = "A cozy reading nook by a rain-streaked window, warm lamplight, "
"a cat asleep on a stack of books, highly detailed")
p.add_argument(
"--prompt",
default = "A cozy reading nook by a rain-streaked window, warm lamplight, "
"a cat asleep on a stack of books, highly detailed",
)
p.add_argument("--width", type = int, default = 1024)
p.add_argument("--height", type = int, default = 1024)
p.add_argument("--steps", type = int, default = 9)
@ -407,20 +429,41 @@ 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("--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(
"--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

5
studio/backend/core/inference/diffusion.py
View file

@ -466,6 +466,7 @@ class DiffusionBackend:
batch_size: int = 1,
) -> dict[str, Any]:
import torch
# A per-generation cancel Event: unload()/a superseding load set THIS event
# (registered under _lock below) to abort just this denoise. _generate_lock
# serialises generations and is the only lock the denoise holds, so a slow
@ -521,7 +522,9 @@ class DiffusionBackend:
gen.step = step_index + 1
if gen.first_step_at == 0.0:
gen.first_step_at = now
gen.eta_seconds = _estimate_eta(gen.total_steps, gen.step, gen.first_step_at, now)
gen.eta_seconds = _estimate_eta(
gen.total_steps, gen.step, gen.first_step_at, now
)
# Preempt a long denoise on unload/eviction or a superseding load:
# diffusers checks pipe._interrupt and stops after the current step.
if cancel.is_set():

1
studio/backend/core/inference/diffusion_device.py
View file

@ -109,7 +109,6 @@ def diffusion_device_target_from_torch_device(
if device == "cuda":
try:
import torch
is_rocm = bool(getattr(getattr(torch, "version", None), "hip", None))
except Exception:
is_rocm = False

51
studio/backend/tests/test_diffusion_backend.py
View file

@ -565,8 +565,13 @@ def test_generate_lock_split_keeps_status_and_unload_responsive(fake_runtime):
fam = detect_family("unsloth/Z-Image-GGUF")
backend._state = _LoadState(
pipe = _BlockingPipe(), family = fam, repo_id = "r", base_repo = "b",
device = "cpu", dtype = "float32", cpu_offload = False,
pipe = _BlockingPipe(),
family = fam,
repo_id = "r",
base_repo = "b",
device = "cpu",
dtype = "float32",
cpu_offload = False,
)
out: dict = {}
@ -610,7 +615,13 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime):
self._interrupt = False
self.steps_run = 0
def __call__(self, *, callback_on_step_end = None, num_inference_steps = 8, **kwargs):
def __call__(
self,
*,
callback_on_step_end = None,
num_inference_steps = 8,
**kwargs,
):
for i in range(num_inference_steps):
if self._interrupt: # diffusers' interrupt protocol
break
@ -625,8 +636,13 @@ def test_callback_cancellation_interrupts_denoise(fake_runtime):
pipe = _SteppingPipe()
fam = detect_family("unsloth/Z-Image-GGUF")
backend._state = _LoadState(
pipe = pipe, family = fam, repo_id = "r", base_repo = "b",
device = "cpu", dtype = "float32", cpu_offload = False,
pipe = pipe,
family = fam,
repo_id = "r",
base_repo = "b",
device = "cpu",
dtype = "float32",
cpu_offload = False,
)
out: dict = {}
@ -658,18 +674,22 @@ def test_validate_load_request(tmp_path):
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF")
with pytest.raises(ValueError, match = "family"):
backend.validate_load_request("meta/Llama-3", gguf_filename = "q.gguf")
assert backend.validate_load_request(
"unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf"
).name == "z-image"
assert (
backend.validate_load_request("unsloth/Z-Image-Turbo-GGUF", gguf_filename = "q.gguf").name
== "z-image"
)
# A local path with a missing child fails here (before any GPU/network work).
with pytest.raises(FileNotFoundError):
backend.validate_load_request(
str(tmp_path), gguf_filename = "missing.gguf", family_override = "z-image"
)
(tmp_path / "m.gguf").write_bytes(b"x")
assert backend.validate_load_request(
str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image"
).name == "z-image"
assert (
backend.validate_load_request(
str(tmp_path), gguf_filename = "m.gguf", family_override = "z-image"
).name
== "z-image"
)
# A path-shaped repo_id that does not exist is rejected here (it would otherwise
# be treated as remote, evict chat, and only fail in the background load).
with pytest.raises(FileNotFoundError):
@ -698,8 +718,13 @@ def test_replacement_load_waits_for_inflight_generation(fake_runtime, tmp_path):
fam = detect_family("unsloth/Z-Image-GGUF")
backend._state = _LoadState(
pipe = _BlockingPipe(), family = fam, repo_id = "r", base_repo = "b",
device = "cpu", dtype = "float32", cpu_offload = False,
pipe = _BlockingPipe(),
family = fam,
repo_id = "r",
base_repo = "b",
device = "cpu",
dtype = "float32",
cpu_offload = False,
)
gen_out: dict = {}

27
studio/backend/tests/test_diffusion_device.py
View file

@ -65,13 +65,13 @@ class _FakeTensor:
def _make_torch(
*,
cuda_available: bool = False,
capability=(8, 0),
capability = (8, 0),
capability_raises: bool = False,
bf16_supported: bool = False,
hip=None,
hip = None,
mps_available: bool = False,
mps_probe: str = "pass", # "pass" | "raise" | "nonfinite"
xpu_available=None, # None -> no xpu attr; True/False -> present
xpu_available = None, # None -> no xpu attr; True/False -> present
xpu_bf16: bool = False,
) -> types.ModuleType:
torch = types.ModuleType("torch")
@ -110,7 +110,14 @@ def _make_torch(
return torch
def _install(monkeypatch, torch, *, studio_device=None, is_rocm=False, hardware_fails=False):
def _install(
monkeypatch,
torch,
*,
studio_device = None,
is_rocm = False,
hardware_fails = False,
):
"""Install the fake torch and either a fake or failing `utils.hardware`."""
monkeypatch.setitem(sys.modules, "torch", torch)
if hardware_fails:
@ -139,7 +146,11 @@ def test_cuda_ampere_bf16(monkeypatch):
_install(monkeypatch, torch, studio_device = "cuda")
t = dd.resolve_diffusion_device_target()
assert (t.device, t.dtype, t.backend, t.vendor) == ("cuda", BF16, "cuda", "nvidia")
assert t.supports_model_cpu_offload and t.supports_default_torch_compile and t.supports_pinned_transfer
assert (
t.supports_model_cpu_offload
and t.supports_default_torch_compile
and t.supports_pinned_transfer
)
def test_cuda_pre_ampere_fp16(monkeypatch):
@ -185,7 +196,11 @@ def test_xpu_bf16(monkeypatch):
_install(monkeypatch, torch, studio_device = "xpu")
t = dd.resolve_diffusion_device_target()
assert (t.device, t.backend, t.vendor, t.dtype) == ("xpu", "xpu", "intel", BF16)
assert t.supports_model_cpu_offload and not t.supports_default_torch_compile and not t.supports_pinned_transfer
assert (
t.supports_model_cpu_offload
and not t.supports_default_torch_compile
and not t.supports_pinned_transfer
)
def test_xpu_without_bf16_fp16(monkeypatch):

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

@ -29,7 +29,13 @@ class _FakeBackend:
def is_loaded(self) -> bool:
return self.loaded
def validate_load_request(self, model_path, *, gguf_filename = None, family_override = None):
def validate_load_request(
self,
model_path,
*,
gguf_filename = None,
family_override = None,
):
# Mirror the real backend's cheap validation so the route's
# validate-before-evict ordering is exercised.
from core.inference.diffusion_families import detect_family