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 13:56:34 +00:00
commit 54eea13036
10 changed files with 287 additions and 121 deletions
Download patch file Download diff file Expand all files Collapse all files

81
scripts/diffusion_bench.py
View file

@ -280,11 +280,19 @@ 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,
"memory_mode": args.memory_mode, "cpu_offload": args.cpu_offload,
"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,
"memory_mode": args.memory_mode,
"cpu_offload": args.cpu_offload,
},
}
@ -425,25 +433,50 @@ 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("--memory-mode", default = None,
choices = ["auto", "fast", "balanced", "low_vram"],
help = "memory policy (default: backend auto)")
p.add_argument("--cpu-offload", action = "store_true",
help = "legacy: force whole-module CPU offload")
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(
"--memory-mode",
default = None,
choices = ["auto", "fast", "balanced", "low_vram"],
help = "memory policy (default: backend auto)",
)
p.add_argument(
"--cpu-offload", action = "store_true", help = "legacy: force whole-module CPU offload"
)
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

177
scripts/diffusion_quality.py
View file

@ -56,7 +56,6 @@ DEFAULT_PROMPTS = [
def _to_gray(img: Any) -> Any:
import numpy as np
return np.asarray(img.convert("L"), dtype = np.float64)
@ -93,7 +92,11 @@ def _box_mean(x: Any, w: int) -> Any:
return total / float(w * w)
def ssim(a_img: Any, b_img: Any, window: int = 7) -> float:
def ssim(
a_img: Any,
b_img: Any,
window: int = 7,
) -> float:
"""Mean structural similarity (luminance) over a uniform window; 1.0 when
identical. Pure numpy box-window SSIM (Wang et al. constants), no skimage."""
a, b = _to_gray(a_img), _to_gray(b_img)
@ -133,9 +136,9 @@ class _Clip:
return emb / emb.norm(dim = -1, keepdim = True)
def _text_embed(self, text: str) -> Any:
inputs = self.proc(
text = [text], return_tensors = "pt", padding = True, truncation = True
).to(self.device)
inputs = self.proc(text = [text], return_tensors = "pt", padding = True, truncation = True).to(
self.device
)
with self.torch.no_grad():
emb = self.model.get_text_features(**inputs)
return emb / emb.norm(dim = -1, keepdim = True)
@ -153,7 +156,6 @@ class _Clip:
def _cuda(call: str) -> Optional[int]:
try:
import torch
if not torch.cuda.is_available():
return None
return int(getattr(torch.cuda, call)())
@ -164,7 +166,6 @@ def _cuda(call: str) -> Optional[int]:
def _cuda_reset_peak() -> None:
try:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
@ -187,7 +188,6 @@ def _wait_for_load(backend: Any, timeout_s: int = 3600) -> None:
def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]:
try:
from huggingface_hub import HfApi
info = HfApi().model_info(repo, files_metadata = True, token = os.environ.get("HF_TOKEN"))
for s in info.siblings:
if s.rfilename == filename and s.size:
@ -200,13 +200,18 @@ def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]:
# ── one quant: load, render the grid, measure ────────────────────────────────
def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Path) -> dict[str, Any]:
def _render_grid(
backend: Any, args: argparse.Namespace, gguf: str, out_dir: Path
) -> dict[str, Any]:
"""Load ``gguf`` and render one image per (prompt, seed); return images keyed by
(prompt_index, seed) plus latency / VRAM metrics."""
_cuda_reset_peak()
backend.begin_load(
args.model, gguf_filename = gguf, base_repo = args.base_repo,
family_override = args.family_override, hf_token = os.environ.get("HF_TOKEN"),
args.model,
gguf_filename = gguf,
base_repo = args.base_repo,
family_override = args.family_override,
hf_token = os.environ.get("HF_TOKEN"),
memory_mode = args.memory_mode,
)
_wait_for_load(backend)
@ -221,8 +226,13 @@ def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Pat
for seed in args.seeds:
t0 = time.time()
result = backend.generate(
prompt = prompt, width = args.width, height = args.height,
steps = args.steps, guidance = args.guidance, seed = seed, batch_size = 1,
prompt = prompt,
width = args.width,
height = args.height,
steps = args.steps,
guidance = args.guidance,
seed = seed,
batch_size = 1,
)
latencies.append(time.time() - t0)
img = result["images"][0]
@ -243,7 +253,9 @@ def _render_grid(backend: Any, args: argparse.Namespace, gguf: str, out_dir: Pat
}
def _compare(grid: dict, ref_grid: dict, clip: Optional[_Clip], prompts: list[str]) -> dict[str, Any]:
def _compare(
grid: dict, ref_grid: dict, clip: Optional[_Clip], prompts: list[str]
) -> dict[str, Any]:
psnrs, ssims, clip_txt, clip_sim = [], [], [], []
for key, img in grid["images"].items():
ref = ref_grid["images"].get(key)
@ -285,16 +297,20 @@ def _sweep(args: argparse.Namespace) -> int:
rows: list[dict[str, Any]] = []
for gguf in quants:
print(f"=== quant: {gguf} ===", flush = True)
grid = gguf == args.reference_quant and ref_grid or _render_grid(backend, args, gguf, out_dir)
grid = (
gguf == args.reference_quant and ref_grid or _render_grid(backend, args, gguf, out_dir)
)
metrics = _compare(grid, ref_grid, clip, args.prompts)
rows.append({
"quant": gguf,
"file_size_mib": grid["file_size_mib"],
"is_reference": gguf == args.reference_quant,
"median_latency_s": grid["median_latency_s"],
"peak_vram_mib": (grid["peak_vram_bytes"] or 0) // (1024 * 1024) or None,
**metrics,
})
rows.append(
{
"quant": gguf,
"file_size_mib": grid["file_size_mib"],
"is_reference": gguf == args.reference_quant,
"median_latency_s": grid["median_latency_s"],
"peak_vram_mib": (grid["peak_vram_bytes"] or 0) // (1024 * 1024) or None,
**metrics,
}
)
print(f" {metrics}", flush = True)
_write_outputs(args, out_dir, rows)
@ -304,17 +320,36 @@ def _sweep(args: argparse.Namespace) -> int:
def _write_outputs(args: argparse.Namespace, out_dir: Path, rows: list[dict]) -> None:
(out_dir / "quality.json").write_text(json.dumps({
"config": {
"model": args.model, "reference_quant": args.reference_quant,
"prompts": args.prompts, "seeds": args.seeds, "steps": args.steps,
"width": args.width, "height": args.height, "guidance": args.guidance,
"memory_mode": args.memory_mode, "clip": args.clip,
},
"rows": rows,
}, indent = 2))
fields = ["quant", "file_size_mib", "peak_vram_mib", "median_latency_s",
"mean_psnr", "mean_ssim", "mean_clip_text", "mean_clip_sim"]
(out_dir / "quality.json").write_text(
json.dumps(
{
"config": {
"model": args.model,
"reference_quant": args.reference_quant,
"prompts": args.prompts,
"seeds": args.seeds,
"steps": args.steps,
"width": args.width,
"height": args.height,
"guidance": args.guidance,
"memory_mode": args.memory_mode,
"clip": args.clip,
},
"rows": rows,
},
indent = 2,
)
)
fields = [
"quant",
"file_size_mib",
"peak_vram_mib",
"median_latency_s",
"mean_psnr",
"mean_ssim",
"mean_clip_text",
"mean_clip_sim",
]
with (out_dir / "quality.csv").open("w", newline = "") as fh:
writer = csv.DictWriter(fh, fieldnames = fields, extrasaction = "ignore")
writer.writeheader()
@ -324,12 +359,17 @@ def _write_outputs(args: argparse.Namespace, out_dir: Path, rows: list[dict]) ->
def _print_table(rows: list[dict]) -> None:
print("\n=== QUALITY vs QUANT (lower size/latency/VRAM better; higher PSNR/SSIM/CLIP better) ===", flush = True)
print(
"\n=== QUALITY vs QUANT (lower size/latency/VRAM better; higher PSNR/SSIM/CLIP better) ===",
flush = True,
)
hdr = f" {'quant':<28}{'size_MB':>9}{'vram_MB':>9}{'lat_s':>8}{'PSNR':>8}{'SSIM':>8}{'CLIPt':>8}{'CLIPs':>8}"
print(hdr, flush = True)
for r in rows:
def _f(v, fmt):
return format(v, fmt) if isinstance(v, (int, float)) else "-"
psnr_str = "inf" if r.get("mean_psnr") == math.inf else _f(r.get("mean_psnr"), ".2f")
print(
f" {r['quant']:<28}{_f(r.get('file_size_mib'), '>9'):>9}"
@ -343,9 +383,11 @@ def _print_table(rows: list[dict]) -> None:
def _recommend(args: argparse.Namespace, rows: list[dict]) -> None:
# The smallest-on-disk non-reference quant that stays within the quality budget.
passing = [
r for r in rows
r
for r in rows
if not r["is_reference"]
and r.get("mean_ssim") is not None and r["mean_ssim"] >= args.ssim_threshold
and r.get("mean_ssim") is not None
and r["mean_ssim"] >= args.ssim_threshold
and (r.get("mean_psnr") is None or r["mean_psnr"] >= args.psnr_threshold)
and r.get("file_size_mib") is not None
]
@ -355,9 +397,12 @@ def _recommend(args: argparse.Namespace, rows: list[dict]) -> None:
print(" no candidate quant met the quality budget; keep the reference quant.", flush = True)
return
best = min(passing, key = lambda r: r["file_size_mib"])
print(f" smallest quant within budget: {best['quant']} "
f"({best['file_size_mib']} MB, SSIM {best['mean_ssim']}, PSNR "
f"{'inf' if best['mean_psnr'] == math.inf else best['mean_psnr']})", flush = True)
print(
f" smallest quant within budget: {best['quant']} "
f"({best['file_size_mib']} MB, SSIM {best['mean_ssim']}, PSNR "
f"{'inf' if best['mean_psnr'] == math.inf else best['mean_psnr']})",
flush = True,
)
# ── self-test (CPU, no GPU/model) ─────────────────────────────────────────────
@ -371,17 +416,23 @@ def _selftest() -> int:
base = rng.integers(0, 256, (128, 128, 3), dtype = np.uint8)
a = Image.fromarray(base)
b = Image.fromarray(base) # identical
noisy = Image.fromarray(np.clip(base.astype(int) + rng.integers(-40, 40, base.shape), 0, 255).astype(np.uint8))
noisy = Image.fromarray(
np.clip(base.astype(int) + rng.integers(-40, 40, base.shape), 0, 255).astype(np.uint8)
)
checks = []
checks.append(("identical PSNR is inf", psnr(a, b) == math.inf))
checks.append(("identical SSIM ~ 1.0", abs(ssim(a, b) - 1.0) < 1e-9))
checks.append(("noisy PSNR is finite + lower", math.isfinite(psnr(a, noisy)) and psnr(a, noisy) < 60))
checks.append(
("noisy PSNR is finite + lower", math.isfinite(psnr(a, noisy)) and psnr(a, noisy) < 60)
)
checks.append(("noisy SSIM < identical", ssim(a, noisy) < ssim(a, b)))
checks.append(("shape mismatch -> 0", psnr(a, Image.fromarray(base[:64])) == 0.0))
# box mean of a constant field equals the constant
const = np.full((32, 32), 7.0)
checks.append(("box mean of constant is constant", abs(_box_mean(const, 7).mean() - 7.0) < 1e-9))
checks.append(
("box mean of constant is constant", abs(_box_mean(const, 7).mean() - 7.0) < 1e-9)
)
ok = True
for name, passed in checks:
@ -399,12 +450,24 @@ def _build_parser() -> argparse.ArgumentParser:
description = "Image quality-vs-quant harness 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("--reference-quant", default = "z-image-turbo-BF16.gguf",
help = "high-fidelity reference GGUF filename")
p.add_argument("--quants", nargs = "*", default = [
"z-image-turbo-Q8_0.gguf", "z-image-turbo-Q4_K_M.gguf", "z-image-turbo-Q2_K.gguf",
], help = "candidate GGUF filenames to score against the reference")
p.add_argument(
"--model", default = "unsloth/Z-Image-Turbo-GGUF", help = "GGUF repo id or local path"
)
p.add_argument(
"--reference-quant",
default = "z-image-turbo-BF16.gguf",
help = "high-fidelity reference GGUF filename",
)
p.add_argument(
"--quants",
nargs = "*",
default = [
"z-image-turbo-Q8_0.gguf",
"z-image-turbo-Q4_K_M.gguf",
"z-image-turbo-Q2_K.gguf",
],
help = "candidate GGUF filenames to score against the reference",
)
p.add_argument("--base-repo", default = None)
p.add_argument("--family-override", default = None)
p.add_argument("--prompts", nargs = "*", default = DEFAULT_PROMPTS)
@ -415,10 +478,18 @@ def _build_parser() -> argparse.ArgumentParser:
p.add_argument("--guidance", type = float, default = 0.0)
p.add_argument("--memory-mode", default = None, choices = ["auto", "fast", "balanced", "low_vram"])
p.add_argument("--clip", action = "store_true", help = "also compute CLIP text + image scores")
p.add_argument("--psnr-threshold", type = float, default = 30.0,
help = "min mean PSNR (dB) vs reference for the recommendation")
p.add_argument("--ssim-threshold", type = float, default = 0.92,
help = "min mean SSIM vs reference for the recommendation")
p.add_argument(
"--psnr-threshold",
type = float,
default = 30.0,
help = "min mean PSNR (dB) vs reference for the recommendation",
)
p.add_argument(
"--ssim-threshold",
type = float,
default = 0.92,
help = "min mean SSIM vs reference for the recommendation",
)
p.add_argument("--out-dir", default = "outputs/diffusion_quality")
p.add_argument("--selftest", action = "store_true", help = "CPU metric sanity check; no GPU/model")
return p

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

@ -465,13 +465,20 @@ class DiffusionBackend:
# Opt-in speed optims run BEFORE placement (channels_last / compile
# must precede CPU offload). Off by default -> bit-identical output.
speed_applied = apply_speed_optims(
pipe, target, is_gguf = bool(gguf_filename), family = fam,
speed_mode = speed_mode or SPEED_OFF, logger = logger,
pipe,
target,
is_gguf = bool(gguf_filename),
family = fam,
speed_mode = speed_mode or SPEED_OFF,
logger = logger,
)
# Cast the dense companion text encoder(s) to fp8 storage (opt-in),
# also before placement so the offload hooks move the smaller weights.
fp8_cast = apply_fp8_text_encoder(
pipe, target, enable = text_encoder_fp8, logger = logger,
pipe,
target,
enable = text_encoder_fp8,
logger = logger,
)
# Decide placement from MEASURED free device memory vs the model's
@ -508,7 +515,12 @@ class DiffusionBackend:
logger.info(
"diffusion.loaded: repo=%s base=%s device=%s offload=%s tiling=%s reasons=%s",
repo_id, base, device, effective_policy, effective_tiling, "; ".join(plan.reasons),
repo_id,
base,
device,
effective_policy,
effective_tiling,
"; ".join(plan.reasons),
)
return self.status()

15
studio/backend/core/inference/diffusion_memory.py
View file

@ -177,7 +177,6 @@ def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]:
def _xpu_memory() -> tuple[Optional[int], Optional[int]]:
try:
import torch
mem_get_info = getattr(getattr(torch, "xpu", None), "mem_get_info", None)
if callable(mem_get_info):
free, total = mem_get_info()
@ -191,7 +190,6 @@ def _system_memory_mib() -> tuple[Optional[int], Optional[int]]:
"""(total, available) host RAM in MiB, via psutil then POSIX sysconf."""
try:
import psutil
vm = psutil.virtual_memory()
return int(vm.total // (1024 * 1024)), int(vm.available // (1024 * 1024))
except Exception:
@ -212,7 +210,6 @@ def file_size_mib(path: Any) -> Optional[int]:
"""On-disk size of ``path`` in MiB, or None if it can't be stat'd."""
try:
from pathlib import Path
return max(1, int(Path(path).expanduser().stat().st_size // (1024 * 1024)))
except Exception:
return None
@ -421,7 +418,11 @@ def plan_diffusion_memory(
def apply_memory_plan(
pipe: Any, plan: MemoryPlan, *, device: str, logger: Any = None,
pipe: Any,
plan: MemoryPlan,
*,
device: str,
logger: Any = None,
) -> tuple[str, bool]:
"""Apply ``plan`` to a freshly built diffusers pipeline: enable the VAE memory
savers then place / offload the weights. Exactly one placement call runs, so the
@ -452,7 +453,8 @@ def apply_memory_plan(
if logger is not None:
logger.warning(
"diffusion.memory: sequential offload failed (%s); "
"falling back to whole-module offload", exc,
"falling back to whole-module offload",
exc,
)
pipe.enable_model_cpu_offload()
policy = OFFLOAD_MODEL
@ -511,6 +513,7 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool:
if logger is not None:
logger.warning(
"diffusion.memory: group offload failed (%s); falling back to "
"whole-module offload", exc,
"whole-module offload",
exc,
)
return False

7
studio/backend/core/inference/diffusion_precision.py
View file

@ -37,14 +37,17 @@ def fp8_text_encoder_supported(target: Any) -> bool:
return False
try:
import torch
return getattr(target, "dtype", None) is torch.bfloat16 and hasattr(torch, "float8_e4m3fn")
except Exception:
return False
def apply_fp8_text_encoder(
pipe: Any, target: Any, *, enable: bool, logger: Any = None,
pipe: Any,
target: Any,
*,
enable: bool,
logger: Any = None,
) -> list[str]:
"""Cast each text encoder's linear weights to fp8 storage (compute stays the
target dtype). Returns the names of the encoders actually cast (empty when

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

@ -62,7 +62,6 @@ def compile_eligible(target: Any, *, is_gguf: bool, family: Any) -> bool:
def _is_bfloat16(dtype: Any) -> bool:
try:
import torch
return dtype is torch.bfloat16
except Exception:
return str(dtype).endswith("bfloat16")
@ -107,7 +106,6 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool:
return False
try:
import torch
vae.to(memory_format = torch.channels_last)
return True
except Exception as exc: # noqa: BLE001 — optimisation only

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

@ -1704,21 +1704,21 @@ class DiffusionLoadRequest(BaseModel):
memory_mode: Optional[Literal["auto", "fast", "balanced", "low_vram"]] = Field(
None,
description = "Memory policy: auto (measured), fast (resident), balanced "
"(stream the transformer, near-resident speed, moderate VRAM "
"cut), low_vram (offload every component, lowest VRAM, slower). "
"Overrides cpu_offload when set.",
"(stream the transformer, near-resident speed, moderate VRAM "
"cut), low_vram (offload every component, lowest VRAM, slower). "
"Overrides cpu_offload when set.",
)
speed_mode: Optional[Literal["off", "default", "max"]] = Field(
None,
description = "Opt-in speed optims (default off -> bit-identical output): "
"default (channels_last + regional torch.compile where eligible), "
"max (also TF32 + fused QKV).",
"default (channels_last + regional torch.compile where eligible), "
"max (also TF32 + fused QKV).",
)
text_encoder_fp8: bool = Field(
False,
description = "Cast the companion text encoder(s) to fp8 storage (~2x smaller, "
"CUDA + bf16 only). A memory-vs-quality tradeoff (shifts fine "
"detail), not free; pairs well with balanced mode.",
"CUDA + bf16 only). A memory-vs-quality tradeoff (shifts fine "
"detail), not free; pairs well with balanced mode.",
)

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

@ -828,7 +828,9 @@ def test_load_memory_mode_low_vram_engages_model_offload(fake_runtime, tmp_path,
assert pipe.offloaded is True and pipe.moved_to is None # offload owns placement
def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(fake_runtime, tmp_path, monkeypatch):
def test_load_explicit_cpu_offload_engages_model_offload_on_cuda(
fake_runtime, tmp_path, monkeypatch
):
# cpu_offload=True with no mode: auto would stay resident (budget unknown under
# the stub), but the explicit flag forces whole-module offload.
(tmp_path / "m.gguf").write_bytes(b"x")

63
studio/backend/tests/test_diffusion_memory.py
View file

@ -36,7 +36,12 @@ from core.inference.diffusion_memory import (
)
def _target(*, device = "cuda", backend = "cuda", supports_offload = True):
def _target(
*,
device = "cuda",
backend = "cuda",
supports_offload = True,
):
"""A duck-typed stand-in for DiffusionDeviceTarget (only the fields the
planner / snapshot read)."""
return types.SimpleNamespace(
@ -221,18 +226,36 @@ def test_auto_stays_resident_when_budget_unknown():
def test_explicit_modes_force_policy_regardless_of_budget():
roomy = _discrete(80000)
assert plan_diffusion_memory(
target = _target(), device_memory = roomy, model_dense_mib = 1000,
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_FAST,
).offload_policy == OFFLOAD_NONE
assert plan_diffusion_memory(
target = _target(), device_memory = roomy, model_dense_mib = 1000,
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_BALANCED,
).offload_policy == OFFLOAD_GROUP
assert plan_diffusion_memory(
target = _target(), device_memory = roomy, model_dense_mib = 1000,
runtime_headroom_mib = 1000, requested_mode = MEMORY_MODE_LOW_VRAM,
).offload_policy == OFFLOAD_MODEL
assert (
plan_diffusion_memory(
target = _target(),
device_memory = roomy,
model_dense_mib = 1000,
runtime_headroom_mib = 1000,
requested_mode = MEMORY_MODE_FAST,
).offload_policy
== OFFLOAD_NONE
)
assert (
plan_diffusion_memory(
target = _target(),
device_memory = roomy,
model_dense_mib = 1000,
runtime_headroom_mib = 1000,
requested_mode = MEMORY_MODE_BALANCED,
).offload_policy
== OFFLOAD_GROUP
)
assert (
plan_diffusion_memory(
target = _target(),
device_memory = roomy,
model_dense_mib = 1000,
runtime_headroom_mib = 1000,
requested_mode = MEMORY_MODE_LOW_VRAM,
).offload_policy
== OFFLOAD_MODEL
)
def test_fast_falls_back_to_model_offload_when_it_does_not_fit():
@ -381,8 +404,10 @@ def test_apply_vae_tiling_falls_back_to_vae_submodule():
class _VaeOnly:
def __init__(self):
self.vae = types.SimpleNamespace(
tiled = False, sliced = False,
enable_tiling = self._tile, enable_slicing = self._slice,
tiled = False,
sliced = False,
enable_tiling = self._tile,
enable_slicing = self._slice,
)
def _tile(self):
@ -409,7 +434,9 @@ def test_apply_group_falls_back_to_model_without_transformer():
def test_apply_sequential_offload():
pipe = _RecordingPipe()
effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
effective, _ = apply_memory_plan(
pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda"
)
assert "sequential_offload" in pipe.calls and "to:cuda" not in pipe.calls
assert effective == OFFLOAD_SEQUENTIAL
@ -422,7 +449,9 @@ def test_apply_sequential_falls_back_to_model_offload_when_unsupported():
raise RuntimeError("sequential offload not supported for this transformer")
pipe = _NoSeqPipe()
effective, _ = apply_memory_plan(pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda")
effective, _ = apply_memory_plan(
pipe, _manual_plan(OFFLOAD_SEQUENTIAL, tiling = True), device = "cuda"
)
assert effective == OFFLOAD_MODEL
assert "model_offload" in pipe.calls

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

@ -24,9 +24,16 @@ from core.inference.diffusion_speed import (
)
def _target(*, device = "cuda", dtype = "bfloat16", compile_ok = True):
def _target(
*,
device = "cuda",
dtype = "bfloat16",
compile_ok = True,
):
return types.SimpleNamespace(
device = device, dtype = dtype, supports_default_torch_compile = compile_ok,
device = device,
dtype = dtype,
supports_default_torch_compile = compile_ok,
)
@ -78,7 +85,12 @@ def test_compile_eligible_requires_non_gguf_bf16_cuda_friendly(monkeypatch):
class _Pipe:
def __init__(self, *, with_compile = False, with_fuse = False) -> None:
def __init__(
self,
*,
with_compile = False,
with_fuse = False,
) -> None:
self.vae = types.SimpleNamespace(mem_format = None, to = self._vae_to)
self.transformer = types.SimpleNamespace()
if with_compile:
@ -144,8 +156,11 @@ def test_speed_max_tf32_only_on_cuda(monkeypatch):
_stub_torch(monkeypatch)
pipe = _Pipe()
applied = apply_speed_optims(
pipe, _target(device = "mps", compile_ok = False), is_gguf = True,
family = _family(), speed_mode = SPEED_MAX,
pipe,
_target(device = "mps", compile_ok = False),
is_gguf = True,
family = _family(),
speed_mode = SPEED_MAX,
)
assert applied["tf32"] is False # not CUDA -> no TF32