aszc-dev
0bbd8d8e0d
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.
69 lines
2.5 KiB
Python
69 lines
2.5 KiB
Python
"""Pure out_name composition for the Core ML UNet artifact.
|
|
|
|
Extracted from CoreMLConverter.convert in Phase 3 so the filename contract
|
|
can be tested + reused without instantiating the node. The string is the
|
|
cache key: every workflow that references a converted .mlpackage depends
|
|
on it staying byte-for-byte identical.
|
|
"""
|
|
from typing import Iterable, Tuple
|
|
|
|
ATTN_SUFFIX = {
|
|
"SPLIT_EINSUM": "se",
|
|
"SPLIT_EINSUM_V2": "se2",
|
|
"ORIGINAL": "orig",
|
|
}
|
|
|
|
# Phase 6: palettization bits. "none" = no quantization (default; keeps the
|
|
# pre-Phase-6 filename intact so existing workflows still resolve their
|
|
# cached .mlpackage). Numeric values append a `_q<bits>` suffix.
|
|
QUANT_NBITS_VALUES = ("none", "8", "6", "4")
|
|
|
|
|
|
def compose_out_name(
|
|
*,
|
|
ckpt_name: str,
|
|
batch_size: int,
|
|
width: int,
|
|
height: int,
|
|
controlnet_support: bool,
|
|
attention_implementation: str,
|
|
lora_names: Iterable[str] = (),
|
|
quantize_nbits: str = "none",
|
|
) -> str:
|
|
"""Build the .mlpackage stem from convert() parameters.
|
|
|
|
Locked behaviour (Phase 2 characterization tests):
|
|
- first '.' in ckpt_name wins (`a.b.c.safetensors` -> `a`)
|
|
- spaces collapse to underscores
|
|
- LoRA names are taken stem-only, sorted, joined with '_' and
|
|
prefixed with '_' when present (caller is expected to pass a
|
|
sorted list; we sort defensively)
|
|
- controlnet adds `_cn`
|
|
- attn suffix is `_se` | `_se2` | `_orig`
|
|
|
|
Phase 6 addition:
|
|
- quantize_nbits "none" (default) appends nothing — existing
|
|
unquantized .mlpackages keep the old filename
|
|
- "4" / "6" / "8" appends `_q<bits>` after the attn suffix
|
|
"""
|
|
if quantize_nbits not in QUANT_NBITS_VALUES:
|
|
raise ValueError(
|
|
f"quantize_nbits={quantize_nbits!r} not in {QUANT_NBITS_VALUES}"
|
|
)
|
|
stem = ckpt_name.split(".")[0]
|
|
sorted_names = sorted(lora_names)
|
|
lora_str = "_" + "_".join(name.split(".")[0] for name in sorted_names) if sorted_names else ""
|
|
cn_suffix = "_cn" if controlnet_support else ""
|
|
attn_suffix = "_" + ATTN_SUFFIX[attention_implementation]
|
|
quant_suffix = f"_q{quantize_nbits}" if quantize_nbits != "none" else ""
|
|
out_name = (
|
|
f"{stem}{lora_str}_{batch_size}x{width}x{height}"
|
|
f"{cn_suffix}{attn_suffix}{quant_suffix}"
|
|
)
|
|
return out_name.replace(" ", "_")
|
|
|
|
|
|
def lora_names_from_params(lora_params: Iterable[Tuple[str, float]]) -> list[str]:
|
|
"""Mirror the sort applied inside CoreMLConverter.convert."""
|
|
return [name for name, _ in sorted(lora_params, key=lambda pair: pair[0])]
|