Commit Graph
10 Commits
Author SHA1 Message Date
aszc-dev 0bbd8d8e0d 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.
2026-05-25 01:30:29 +02:00
aszc-dev 66e83c2f2f Change syntax to support older Python versions 2024-07-04 08:44:37 +02:00
aszc-dev b8197c21ef Converting refiner works 2023-11-24 12:14:15 +01:00
aszc-dev 0c78803b25 Base SDXL conversion works 2023-11-24 12:14:15 +01:00
aszc-dev 9f650acb79 Load .yaml config if present 2023-11-17 22:55:13 +01:00
aszc-dev c6d6917827 Setting LoRA model weights works 2023-11-17 22:55:13 +01:00
aszc-dev 8092a19173 Enable choosing attention implementation during conversion 2023-11-17 22:55:13 +01:00
aszc-dev a8d2d6ec46 Move lora related code around, remove clip stuff 2023-11-17 22:55:13 +01:00
aszc-dev 6907d4910f Remove ckpt loading when loading lora clip 2023-11-17 22:55:13 +01:00
aszc-dev 45be6761d1 Basic conversion + LoRA support works 2023-11-17 22:55:13 +01:00