feat(phase6): opt-in k-means weight palettization (quantize_nbits)
Phase 6 of the modernization plan: add weight palettization to the
Core ML converter as an opt-in knob, so the SD1.5 / SDXL UNet can
ship at 1/2, 1/2.7 or 1/4 of its current size with ANE-friendly
inference.
CoreMLConverter (and the LCM converter) gains a `quantize_nbits`
dropdown: `none` (default — identical to pre-Phase-6 behavior and
filenames, so existing cached .mlpackages still resolve) / `8` / `6` /
`4`. The value is encoded as `_q<bits>` after the attn suffix, so the
unquantized model and the three palettized variants coexist on disk
under distinct cache keys.
Implementation
- core/naming.compose_out_name: accepts `quantize_nbits`, validates
against {none, 8, 6, 4}, appends `_q<bits>` (none = empty).
- converter.convert_unet: after ct.convert + before .save, runs
coremltools.optimize.coreml.palettize_weights with
OpPalettizerConfig(mode="kmeans", nbits=...) when the value is not
"none". Adds a `Palettization took Xs` log line.
- converter.convert / nodes.CoreMLConverter.convert: pipe the new arg
through; the ComfyUI node exposes it as a dropdown with default
"none" so existing workflows are unchanged at load time.
- bench/scripts/convert_sd15.py: QUANT_NBITS env knob; uses the
pure compose_out_name (replaces the inline string formatter).
Test infra
- tests/unit/test_characterization_out_name.py: 6 new tests pinning
the `_q<bits>` suffix contract, the "none" passthrough (backward
compat), the cn + lora + quant combination, and the invalid-value
ValueError. Total Tier 0 now at 94.
- Makefile gains `bench-quant` (runs the matrix script) and
`convert-quant` (converts q8, q6, q4 sequentially).
- bench/scripts/quant_matrix.py (new): loads each variant, runs
REPEATS forward passes with a fixed seed, then computes the
noise_pred PSNR of each quantized variant against the unquantized
baseline. Writes bench/results/quant_matrix_<sha>.{json,md}.
README
- New "Quantization (Phase 6, opt-in)" section: tradeoff table
measured on M2 Pro SD1.5 1x512x512 SPLIT_EINSUM (sizes 1641/822/
617/412 MB; fwd 197/187/183/180 ms; PSNR 53.5 / 40.2 / 27.5 dB),
plus per-chip/RAM recommendations.
Default-path safety
- "none" produces the same out_name as Phase 5 -> existing
v1-5-pruned-emaonly_1x512x512_se_unet.mlmodelc is still picked up
unchanged; the m2 golden image test continues to anchor.
This commit is contained in:
aszc-dev
@@ -49,21 +49,32 @@ def main() -> int:
|
||||
width = int(os.environ.get("WIDTH", "512"))
|
||||
batch_size = int(os.environ.get("BATCH_SIZE", "1"))
|
||||
controlnet = os.environ.get("CONTROLNET", "0") not in ("0", "false", "False", "")
|
||||
quantize_nbits = os.environ.get("QUANT_NBITS", "none")
|
||||
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
if not ckpt_path:
|
||||
log.error("checkpoint not found: %s under %s", ckpt_name, COMFY_DIR / "models" / "checkpoints")
|
||||
return 2
|
||||
|
||||
attn_suffix = {"SPLIT_EINSUM": "se", "SPLIT_EINSUM_V2": "se2", "ORIGINAL": "orig"}[attn]
|
||||
cn_suffix = "_cn" if controlnet else ""
|
||||
stem = ckpt_name.split(".")[0]
|
||||
out_name = f"{stem}_{batch_size}x{width}x{height}{cn_suffix}_{attn_suffix}"
|
||||
from coreml_suite.core.naming import compose_out_name
|
||||
|
||||
out_name = compose_out_name(
|
||||
ckpt_name=ckpt_name,
|
||||
batch_size=batch_size,
|
||||
width=width,
|
||||
height=height,
|
||||
controlnet_support=controlnet,
|
||||
attention_implementation=attn,
|
||||
quantize_nbits=quantize_nbits,
|
||||
)
|
||||
unet_out_path = converter.get_out_path("unet", out_name)
|
||||
|
||||
log.info("repo_root=%s comfy_dir=%s", REPO_ROOT, COMFY_DIR)
|
||||
log.info("ckpt=%s out_name=%s", ckpt_path, out_name)
|
||||
log.info("attn=%s size=%dx%d batch=%d controlnet=%s", attn, width, height, batch_size, controlnet)
|
||||
log.info(
|
||||
"attn=%s size=%dx%d batch=%d controlnet=%s quant=%s",
|
||||
attn, width, height, batch_size, controlnet, quantize_nbits,
|
||||
)
|
||||
|
||||
converter.convert(
|
||||
ckpt_path=ckpt_path,
|
||||
@@ -75,6 +86,7 @@ def main() -> int:
|
||||
lora_weights=[],
|
||||
attn_impl=attn,
|
||||
config_path=None,
|
||||
quantize_nbits=quantize_nbits,
|
||||
)
|
||||
|
||||
target_path = converter.compile_model(out_path=unet_out_path, out_name=out_name, submodule_name="unet")
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
#!/usr/bin/env python3
|
||||
"""bench/scripts/quant_matrix.py — Phase 6 quantization tradeoff matrix.
|
||||
|
||||
For each variant in {none, 8, 6, 4}, load the corresponding .mlmodelc,
|
||||
run a synthetic forward pass with a fixed seed, and record:
|
||||
- on-disk size
|
||||
- load time
|
||||
- forward-pass median latency
|
||||
- PSNR of noise_pred vs the `none` baseline (proxy for output drift)
|
||||
|
||||
Writes a JSON + Markdown summary into bench/results/quant_matrix_<sha>.{json,md}.
|
||||
|
||||
Assumes converted .mlmodelc files exist under $COMFY_DIR/models/unet/
|
||||
following the convention from coreml_suite.core.naming.compose_out_name
|
||||
(e.g. v1-5-pruned-emaonly_1x512x512_se_unet.mlmodelc for the baseline,
|
||||
v1-5-pruned-emaonly_1x512x512_se_q4_unet.mlmodelc for 4-bit, etc.).
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_DIR = Path(os.environ.get("COMFY_DIR", REPO_ROOT.parents[1])).resolve()
|
||||
MODELS_DIR = COMFY_DIR / "models" / "unet"
|
||||
|
||||
VARIANTS = ["none", "8", "6", "4"]
|
||||
CKPT_STEM = os.environ.get("CKPT_STEM", "v1-5-pruned-emaonly")
|
||||
BATCH = int(os.environ.get("BATCH", "1"))
|
||||
WIDTH = int(os.environ.get("WIDTH", "512"))
|
||||
HEIGHT = int(os.environ.get("HEIGHT", "512"))
|
||||
ATTN_SUFFIX = os.environ.get("ATTN_SUFFIX", "se")
|
||||
COMPUTE_UNIT = os.environ.get("COMPUTE_UNIT", "CPU_AND_NE")
|
||||
REPEATS = int(os.environ.get("REPEATS", "20"))
|
||||
INPUT_SEED = int(os.environ.get("INPUT_SEED", "0"))
|
||||
|
||||
|
||||
def variant_path(nbits: str) -> Path:
|
||||
quant_suffix = f"_q{nbits}" if nbits != "none" else ""
|
||||
name = f"{CKPT_STEM}_{BATCH}x{WIDTH}x{HEIGHT}_{ATTN_SUFFIX}{quant_suffix}_unet.mlmodelc"
|
||||
return MODELS_DIR / name
|
||||
|
||||
|
||||
def dir_size_bytes(path: Path) -> int:
|
||||
total = 0
|
||||
for root, _dirs, files in os.walk(path):
|
||||
for f in files:
|
||||
try:
|
||||
total += (Path(root) / f).stat().st_size
|
||||
except OSError:
|
||||
pass
|
||||
return total
|
||||
|
||||
|
||||
def psnr(a: np.ndarray, b: np.ndarray) -> float:
|
||||
"""Float PSNR for noise_pred tensors normalized to a [-1, 1]-ish range."""
|
||||
a = a.astype(np.float64)
|
||||
b = b.astype(np.float64)
|
||||
peak = max(float(np.max(np.abs(a))), float(np.max(np.abs(b))), 1.0)
|
||||
mse = float(np.mean((a - b) ** 2))
|
||||
if mse == 0:
|
||||
return 100.0
|
||||
return 20.0 * float(np.log10(peak / np.sqrt(mse)))
|
||||
|
||||
|
||||
def measure_variant(
|
||||
nbits: str, ref_output: "np.ndarray | None"
|
||||
) -> "tuple[dict[str, Any], np.ndarray | None]":
|
||||
from python_coreml_stable_diffusion.coreml_model import CoreMLModel
|
||||
|
||||
path = variant_path(nbits)
|
||||
result: dict[str, Any] = {"nbits": nbits, "path": str(path)}
|
||||
if not path.exists():
|
||||
result["error"] = f"model not found at {path}"
|
||||
return result, None
|
||||
result["size_bytes"] = dir_size_bytes(path)
|
||||
|
||||
rng = np.random.default_rng(INPUT_SEED)
|
||||
t0 = time.perf_counter()
|
||||
model = CoreMLModel(str(path), COMPUTE_UNIT, "compiled")
|
||||
result["load_time_s"] = round(time.perf_counter() - t0, 4)
|
||||
|
||||
expected = dict(model.expected_inputs)
|
||||
fixed_inputs = {
|
||||
name: rng.standard_normal(tuple(int(d) for d in spec["shape"])).astype(np.float16)
|
||||
for name, spec in expected.items()
|
||||
}
|
||||
|
||||
# Warmup
|
||||
out = model(**fixed_inputs)
|
||||
times = []
|
||||
for _ in range(REPEATS):
|
||||
# Fresh randomness per repeat for steady-state timing, but keep
|
||||
# output capture deterministic (we run fixed_inputs at end).
|
||||
live = {
|
||||
name: rng.standard_normal(tuple(int(d) for d in spec["shape"])).astype(np.float16)
|
||||
for name, spec in expected.items()
|
||||
}
|
||||
t0 = time.perf_counter()
|
||||
model(**live)
|
||||
times.append((time.perf_counter() - t0) * 1000.0)
|
||||
result["fwd_ms_median"] = round(statistics.median(times), 3)
|
||||
result["fwd_ms_min"] = round(min(times), 3)
|
||||
|
||||
# Deterministic forward for PSNR comparison across variants.
|
||||
deterministic_out = model(**fixed_inputs)["noise_pred"]
|
||||
if ref_output is None:
|
||||
result["psnr_db_vs_none"] = None
|
||||
else:
|
||||
result["psnr_db_vs_none"] = round(psnr(ref_output, deterministic_out), 2)
|
||||
return result, deterministic_out
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
sha = subprocess.check_output(
|
||||
["git", "-C", str(REPO_ROOT), "rev-parse", "--short", "HEAD"],
|
||||
stderr=subprocess.DEVNULL,
|
||||
).decode().strip()
|
||||
except Exception:
|
||||
sha = "nogit"
|
||||
|
||||
out_dir = REPO_ROOT / "bench" / "results"
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"[quant-matrix] git={sha} compute_unit={COMPUTE_UNIT} repeats={REPEATS}", file=sys.stderr)
|
||||
|
||||
rows: list[dict[str, Any]] = []
|
||||
ref_output: np.ndarray | None = None
|
||||
for v in VARIANTS:
|
||||
print(f"[quant-matrix] measuring nbits={v} ...", file=sys.stderr)
|
||||
try:
|
||||
row, out_tensor = measure_variant(v, ref_output)
|
||||
if v == "none":
|
||||
ref_output = out_tensor
|
||||
rows.append(row)
|
||||
except Exception as exc:
|
||||
rows.append({"nbits": v, "error": f"{type(exc).__name__}: {exc}"})
|
||||
|
||||
baseline_size = next((r["size_bytes"] for r in rows if r.get("nbits") == "none" and "size_bytes" in r), None)
|
||||
for r in rows:
|
||||
if "size_bytes" in r and baseline_size:
|
||||
r["size_ratio_vs_none"] = round(r["size_bytes"] / baseline_size, 3)
|
||||
|
||||
report = {
|
||||
"git_sha": sha,
|
||||
"compute_unit": COMPUTE_UNIT,
|
||||
"repeats": REPEATS,
|
||||
"input_seed": INPUT_SEED,
|
||||
"ckpt_stem": CKPT_STEM,
|
||||
"rows": rows,
|
||||
}
|
||||
(out_dir / f"quant_matrix_{sha}.json").write_text(json.dumps(report, indent=2))
|
||||
|
||||
lines = [
|
||||
f"# Quantization tradeoff matrix — `{sha}`",
|
||||
"",
|
||||
f"- compute unit: {COMPUTE_UNIT}",
|
||||
f"- repeats per variant: {REPEATS}",
|
||||
f"- ckpt: {CKPT_STEM}, {BATCH}x{WIDTH}x{HEIGHT}, attn={ATTN_SUFFIX}",
|
||||
"",
|
||||
"| nbits | size (MB) | size vs none | load (s) | fwd median (ms) | fwd min (ms) | PSNR vs none (dB) | error |",
|
||||
"|---|---|---|---|---|---|---|---|",
|
||||
]
|
||||
for r in rows:
|
||||
if "error" in r and "size_bytes" not in r:
|
||||
lines.append(f"| {r.get('nbits','?')} | - | - | - | - | - | - | {r['error']} |")
|
||||
continue
|
||||
size_mb = round(r.get("size_bytes", 0) / (1024 * 1024), 1)
|
||||
lines.append(
|
||||
f"| {r['nbits']} | {size_mb} | {r.get('size_ratio_vs_none','-')} | "
|
||||
f"{r.get('load_time_s','-')} | {r.get('fwd_ms_median','-')} | "
|
||||
f"{r.get('fwd_ms_min','-')} | {r.get('psnr_db_vs_none') if r.get('psnr_db_vs_none') is not None else '-'} | "
|
||||
f"{r.get('error','')} |"
|
||||
)
|
||||
(out_dir / f"quant_matrix_{sha}.md").write_text("\n".join(lines) + "\n")
|
||||
print(f"[quant-matrix] wrote {out_dir / f'quant_matrix_{sha}.md'}", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user