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.
Phase 4 of the modernization plan: institutionalize the 3-tier strategy
so future changes are guarded automatically, and pin down the
self-hosted M2 path the maintainer's hardware needs.
Tier dispatch
- Makefile targets test-unit / test-smoke / test-m2 / bench (plus
ci-tier0 / ci-tier1 wrappers that echo env first). check-macos-arm
fails fast on non-Apple-Silicon hosts.
Tier 1 smoke
- tests/smoke/test_synthetic_unet.py: builds a TinyUNet (conv-in,
time/text projections, conv-out), traces it, ct.convert to
mlprogram + fp16 CPU_ONLY, loads back via CoreMLModel and asserts
expected_inputs + named output. Runs in ~2s; auto-skips on
non-Apple-Silicon. Catches coremltools / ml-stable-diffusion API
drift without needing a real SD checkpoint or the ANE.
GitHub Actions
- .github/workflows/tier0.yml: ubuntu-latest on every push/PR, ~10
min budget, minimal-deps install (torch==2.0.1, numpy<1.25, pytest)
-> pytest -m unit.
- .github/workflows/tier1.yml: macos-14 (M1) on push/PR; opt-in via
run-tier1 label on labeled PRs to spare external-doc PRs.
- .github/workflows/tier2.yml: self-hosted [macOS, ARM64, coreml] on
PR label run-m2 / nightly cron / workflow_dispatch. Starts ComfyUI
with --cpu-vae, runs pytest -m m2 + bench/run.py, uploads bench
results.
Integration coverage moved
- Removed tests/integration/test_basic_conversion_1_5.py: it required
an MPS reference image (broken on macOS 26 + torch 2.0.1, see
Phase 1 Gate) and a checkpoint the maintainer doesn't have on disk
(dreamshaper_8). The same coverage now lives in
tests/m2/test_golden_image.py: deterministic numerical pass/fail
(SHA256 + PSNR fallback) against a stored golden, Core ML pipeline
only. No more human eyeballing.
Docs
- docs/ci-m2.md: one-time runner registration steps, COMFY_DIR
persistence, baseline model pre-conversion, trigger semantics, what
to do when the runner is offline, and the migration note from
integration -> m2 golden.
Sanity check
- Temporarily set convert_to="BREAKAGE_CANARY_NOT_A_REAL_FORMAT" in
the smoke test; Tier 1 surfaced
NotImplementedError: Backend converter BREAKAGE_CANARY_NOT_A_REAL_FORMAT not implemented
immediately. Reverted.
Local verification
- make test-unit -> 88/88 passed in 2.09s
- make test-smoke -> 1/1 passed in 1.99s
aszc-dev