From d90546b6bb945f229437df77e9b7aa02e1bccf0e Mon Sep 17 00:00:00 2001 From: aszc Date: Tue, 26 May 2026 22:12:56 +0200 Subject: [PATCH] =?UTF-8?q?feat(extraction):=20split=20conversion=20into?= =?UTF-8?q?=20the=20coreml-diffusion=20package=20(E0=E2=80=93E5)=20(#63)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs(extraction): E0 seam inventory + correct stale spec assumptions Resolve all pre-flight greps for the converter-extraction seam: - conversion/* and lcm/unet.py confirmed comfy-free - converter.py: only folder_paths reach-in is get_out_path - lcm/converter.py: folder_paths + comfy.model_management to cut; dup helpers and SimianLuo HF hardcode confirmed - no attention module-global; already per-call Correct two stale assumptions verified against current source: - ml-stable-diffusion is already fully removed (#58); CoreMLModel is a local coremltools wrapper, not Apple's. Drop the dep-pinning blocker and the package/suite dep lines that assumed it. - model_version discovery must emit .name (node reverses via ModelVersion[...]); the .value form in the draft would KeyError on every saved workflow. * feat(extraction): E1 coreml_diffusion package + discovery API Stand up the framework-free coreml_diffusion namespace and freeze its versioned discovery contract. The package re-exports from already comfy-free coreml_suite sources (model_version, attention, core.naming); the conversion implementation moves in E2. - list_model_versions/list_attention_impls/list_quant_modes return today's exact dropdown strings, so wiring the node onto them (E3) changes no value and breaks no saved workflow. - Status/_MODEL_STATUS registry gates VERIFIED vs EXPERIMENTAL in the package, so promoting a model expands the node dropdown with no Suite change (additive-only contract; CONTRACT_VERSION=1.0). - Tier-0 test pins the contract and proves comfy/diffusers/coremltools are not pulled on import. Node untouched; zero behavior change. * refactor(extraction): E2 move conversion mechanics into coreml_diffusion Physically relocate the framework-free conversion code into the package and collapse the duplicated LCM/main helpers, behavior-preserving. - coreml_suite/conversion/ -> coreml_diffusion/conversion/ (attention, shapes, trace, unet) - coreml_suite/core/naming.py -> coreml_diffusion/naming.py (the cache-key contract now lives with the package; tests re-pointed) - coreml_suite/converter.py logic -> coreml_diffusion/convert.py, with convert() made keyword-only past (ckpt_path, model_version, out_path) per the interface contract; out_path is injected (no folder_paths) - dedup: load_coreml_model / convert_to_coreml / get_coreml_inputs / add_cnet_support / get_encoder_hidden_states_shape / inputs-spec now defined once in the package; get_sample_input gains an optional scheduler arg so the LCM path shares it (same keys/order/dtypes) - coreml_suite.{converter,lcm.converter} reduced to comfy-side shims: folder_paths path resolution and the LCM scheduler's comfy.model_management stay here; the package imports neither - __init__ keeps discovery + compose_out_name eager; convert is lazy via __getattr__ so 'import coreml_diffusion' stays Tier-0 pure Nodes untouched (E3 thins them onto the package). Tier-0 (109) and smoke (3, real coremltools conversion) green; [M2-ANE] golden pending a server. * refactor(extraction): E3 thin nodes onto coreml_diffusion + discovery dropdowns The CoreMLConverter node now calls coreml_diffusion directly instead of the coreml_suite.converter shim, and its dropdowns are populated at runtime from the package's discovery API. - INPUT_TYPES dropdowns (model_version / attention_implementation / quantize_nbits) now come from a fail-soft _discover() that calls coreml_diffusion.list_*; a missing/old package falls back to a literal list and logs a warning instead of de-registering the node. Installing a newer coreml_diffusion surfaces new conversion types with no Suite change. - folder_paths path resolution moved inline into the node; the package's convert() takes the output path as an injected positional. - compose_out_name / lora_names_from_params now imported from coreml_diffusion (lazily, inside convert) — no node-side copy. - deleted the dead coreml_suite/converter.py and coreml_suite/core/naming.py shims (no remaining importers). Field names, RETURN_TYPES/NAMES and NODE_*_MAPPINGS unchanged; dropdown values are a superset of the prior literals (additive-only). Tier-0 (109) and smoke (3) green; [M2-ANE] golden re-runs on push. * refactor(extraction): E5 depend on external coreml-diffusion package Conversion code now lives in the standalone coreml-diffusion repo. The Suite deletes its in-tree copy and depends on the package instead. - removed coreml_diffusion/ (whole package), coreml_suite/model_version.py and coreml_suite/attention.py (moved to the package as its source of truth), and the tests that moved with them (discovery, conversion_helpers, out_name; smoke synthetic_unet + split_einsum) - re-pointed ModelVersion imports (config.py, nodes.py, lcm/converter.py) to coreml_diffusion - pyproject: drop the coreml_diffusion package include and the conversion-only deps (peft/omegaconf/transformers, now transitive via coreml-diffusion); add coreml-diffusion as a dependency with a local path source until it is published (switch to git tag/PyPI once the repo exists, so CI can resolve it) Suite Tier-0 green (75); conversion code fully absent from the Suite. The comfy node still imports coreml_diffusion (installed package) for ModelVersion + the discovery dropdowns + convert. * build(extraction): pin coreml-diffusion to git tag v0.1.0 Switch the coreml-diffusion source from a local path to the published git tag so CI can resolve it. Suite Tier-0 green resolving from the tag. * ci(extraction): drop Suite smoke tier (moved to coreml-diffusion) The conversion smoke tests moved to the coreml-diffusion repo, which runs its own Tier 1. The Suite's smoke lane had no tests left (pytest exit 5). The Suite keeps Tier 0 (inference units) and the m2 golden e2e. * chore(release): v2.1.0; wire coreml-diffusion into requirements.txt Minor bump: the conversion path moved to the external coreml-diffusion package (node graph + artifact cache keys unchanged, golden-verified). requirements.txt (used by ComfyUI Manager) now installs coreml-diffusion from the v0.1.0 tag and drops the conversion-only deps now provided transitively. --- .github/workflows/tier1.yml | 23 - CONVERTER_EXTRACTION_SPEC.md | 618 +++++++++++++++++++ coreml_suite/attention.py | 5 - coreml_suite/config.py | 2 +- coreml_suite/conversion/__init__.py | 9 - coreml_suite/conversion/attention.py | 239 ------- coreml_suite/conversion/shapes.py | 20 - coreml_suite/conversion/trace.py | 61 -- coreml_suite/conversion/unet.py | 54 -- coreml_suite/converter.py | 322 ---------- coreml_suite/core/naming.py | 68 -- coreml_suite/lcm/converter.py | 177 +----- coreml_suite/model_version.py | 8 - coreml_suite/nodes.py | 66 +- pyproject.toml | 13 +- requirements.txt | 4 +- seam.md | 208 +++++++ tests/smoke/__init__.py | 0 tests/smoke/test_split_einsum_attention.py | 41 -- tests/smoke/test_synthetic_unet.py | 138 ----- tests/unit/test_characterization_out_name.py | 197 ------ tests/unit/test_conversion_helpers.py | 183 ------ uv.lock | 22 +- 23 files changed, 926 insertions(+), 1552 deletions(-) delete mode 100644 .github/workflows/tier1.yml create mode 100644 CONVERTER_EXTRACTION_SPEC.md delete mode 100644 coreml_suite/attention.py delete mode 100644 coreml_suite/conversion/__init__.py delete mode 100644 coreml_suite/conversion/attention.py delete mode 100644 coreml_suite/conversion/shapes.py delete mode 100644 coreml_suite/conversion/trace.py delete mode 100644 coreml_suite/conversion/unet.py delete mode 100644 coreml_suite/converter.py delete mode 100644 coreml_suite/core/naming.py delete mode 100644 coreml_suite/model_version.py create mode 100644 seam.md delete mode 100644 tests/smoke/__init__.py delete mode 100644 tests/smoke/test_split_einsum_attention.py delete mode 100644 tests/smoke/test_synthetic_unet.py delete mode 100644 tests/unit/test_characterization_out_name.py delete mode 100644 tests/unit/test_conversion_helpers.py diff --git a/.github/workflows/tier1.yml b/.github/workflows/tier1.yml deleted file mode 100644 index f805981..0000000 --- a/.github/workflows/tier1.yml +++ /dev/null @@ -1,23 +0,0 @@ -name: Tier 1 — Smoke (macOS self-hosted) - -# macOS smoke tests run on the self-hosted Apple Silicon runner instead of -# GitHub-hosted macOS (10x minute multiplier), which exhausts the included -# Actions minutes too quickly. -on: - push: - branches: [main] - pull_request: - -jobs: - smoke: - runs-on: [self-hosted, macOS, ARM64, coreml] - timeout-minutes: 20 - steps: - - uses: actions/checkout@v4 - - # The self-hosted runner provides uv; no setup-uv action needed. - - name: uv sync - run: uv sync --no-install-project - - - name: Run Tier 1 (synthetic micro-UNet smoke) - run: uv run pytest -m smoke tests/ -v diff --git a/CONVERTER_EXTRACTION_SPEC.md b/CONVERTER_EXTRACTION_SPEC.md new file mode 100644 index 0000000..26bf3ff --- /dev/null +++ b/CONVERTER_EXTRACTION_SPEC.md @@ -0,0 +1,618 @@ +# ComfyUI-CoreMLSuite — Converter Extraction Spec for Claude Code + +> **Companion to `MODERNIZATION_SPEC.md`.** That spec hardens the repo and (Phase 3) +> splits the *inference* math from the framework. **This** spec splits the *conversion* +> path (`safetensors → CoreML`) out into a standalone, `comfy`-free, pip-installable +> package that CoreMLSuite then depends on — and that other projects (incl. on-device +> iOS tooling) can reuse. +> +> **Same discipline as the modernization spec:** safety-net first, behavior-preserving +> until told otherwise, one phase = one branch = one PR, `STOP — VALIDATE` gate between +> every phase, golden-latent as the regression anchor. `[M2]` = needs macOS/Apple Silicon; +> `[M2-ANE]` = needs the Neural Engine. Everything else must run on plain Linux/CI. + +--- + +## 0. How to work (read first — non-negotiable) + +1. **Behavior-preserving until Phase E6.** Phases E1–E5 must not change image output, node + names, `INPUT_TYPES` field names, or `NODE_CLASS_MAPPINGS` keys. The node graph is the + public contract; saved user-workflow JSON breaks if these change. +2. **The conversion package produces an artifact and stops there.** Its job ends at a written + `.mlpackage` / `.mlmodelc` on disk. It must NOT import `comfy`, `folder_paths`, or + `comfy_extras`, and must NOT know ComfyUI's `models/unet` layout. Paths are *inputs*. +3. **The runtime loader stays in the suite.** The loader is the **local** `coreml_suite.coreml_model.CoreMLModel` + — a thin wrapper over `coremltools.models.MLModel` (NOT Apple's + `python_coreml_stable_diffusion.coreml_model.CoreMLModel`, which is no longer used; see #58). + It *runs* a compiled model in Python — a desktop/Python inference concern, not a conversion + concern. It is NOT moved into the package. (On iOS the `.mlmodelc` is loaded natively; the + package's output is the deliverable, not a Python runner.) +4. **Decouple in-repo before splitting repos.** Phases E1–E4 create the package *inside this + repo* and prove equivalence. The physical second-repo split is Phase E5, only after the + golden latent is proven identical. Do not create a second repository before Gate E4 passes. +5. **Reuse the existing regression anchor.** The golden latent / PSNR anchor from + `MODERNIZATION_SPEC.md` Phase 2 is the cross-cutting proof for every gate here. If it is not + yet captured, capture it first (it is a prerequisite for E2 onward). +6. **No new runtime dependencies** without flagging in the gate report (name, why, license, size). +7. **A failing gate means stop and report**, not work around into the next phase. +8. **Tooling is `uv`, not bare `pip`/`venv`.** Every environment/install/lock step uses the + project's `uv` toolchain: `uv venv`, `uv pip install`, `uv pip install -e .`, `uv lock`, + `uv run pytest`, `uv export`/`uv pip freeze` for baselines. Where this spec says "fresh venv", + read "`uv venv` + `uv pip install`". Reserve `uv pip` (not `pip`) inside that venv too. +9. **The package is the single source of truth for *what is possible*; the node is a thin, + discovery-driven frontend.** See the "Interface contract" pillar below — this is the + maintainer's hard requirement and it overrides the earlier (now-rescinded) "freeze the + dropdown list" instruction. + +--- + +## Interface contract (the maintainer's hard requirement) — read before any phase + +Two coupled guarantees must hold once the package is split out: + +**(A) Updating the converter must NOT require updating CoreMLSuite.** +This is satisfied by treating the package's public surface as a versioned contract: +- `convert(...)` and `compile_model(...)` are **keyword-only with defaults** for everything + past the genuinely-required positionals (`ckpt_path`, `model_version`, `out_path`). New + capabilities are added as new keyword args with defaults, so an old Suite's call still + validates against a newer package. **Never** reorder or rename existing parameters. +- `compose_out_name` (the `.mlpackage` filename = the cache key) **moves into the package** and + is versioned with it. The Suite must not carry its own copy; if the package changes the naming + scheme that is a **major** bump (old cached artifacts stop resolving). + +**(B) CoreMLSuite must be able to list *new* conversion types WITHOUT a Suite code change or +version bump.** Today the node hardcodes its dropdowns: +```python +"model_version": ([ModelVersion.SD15.name, ModelVersion.SDXL.name],), # hand-typed, also INCOMPLETE (no LCM / SDXL_REFINER) +"attention_implementation": (list(ATTENTION_IMPLEMENTATIONS),), # from coreml_suite.attention +"quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}), # from coreml_suite.core.naming +``` +These are replaced by **runtime discovery calls into the package**, evaluated inside +`INPUT_TYPES` (ComfyUI re-evaluates `INPUT_TYPES` on every plugin load): +```python +import coreml_diffusion +"model_version": (coreml_diffusion.list_model_versions(),), +"attention_implementation": (coreml_diffusion.list_attention_impls(),), +"quantize_nbits": (coreml_diffusion.list_quant_modes(), {"default": "none"}), +``` +Effect: `uv pip install -U coreml_diffusion` + ComfyUI restart surfaces any newly-added type in the old +plugin's dropdown — **no Suite edit, no Suite version bump.** This is the requirement. + +**The cost, stated honestly (accept this trade-off explicitly at Gate E0):** +- The Suite becomes a "dumb" frontend; the package is the sole authority on what conversions + exist. The Suite can no longer guarantee its saved workflows are valid against *arbitrary* + future package versions. +- Therefore the package's discovery identifiers (`ModelVersion` values, attn-impl strings, quant + modes) are an **ADDITIVE-ONLY contract**: the package may *add* identifiers freely (minor bump, + no Suite change); **removing or renaming an identifier is a breaking change requiring a MAJOR + bump and a migration note**, because a saved workflow JSON references these strings verbatim. + Without this rule, "no version bump" silently becomes "randomly broken workflows." +- `INPUT_TYPES` must **fail soft** when the package is missing/old: wrap the discovery calls so a + missing `coreml_diffusion` (or an old one lacking a `list_*` function) yields a sane fallback list and a + logged warning, instead of the node failing to register and disappearing from the menu. + +**Discovery API the package must expose (stable names):** +```python +coreml_diffusion.list_model_versions() -> list[str] # VERIFIED ones only, e.g. ["SD15","SDXL"] today (.name — see seam.md) +coreml_diffusion.list_attention_impls() -> list[str] # ["SPLIT_EINSUM","SPLIT_EINSUM_V2","ORIGINAL"] +coreml_diffusion.list_quant_modes() -> list[str] # ["none","8","6","4"] +coreml_diffusion.CONTRACT_VERSION: str # bump rules above; Suite may log/compare it +``` +These return the *display strings already used today*, so existing workflows keep validating. + +**Verification status is a PACKAGE property, not a node hardcode (maintainer's intent).** +The Suite wants to expose *every model the converter can verifiably convert*. Today `lcm` and +`sdxl_refiner` are absent from the converter node not because the Suite chooses to hide them, but +because they lack a full golden/PSNR verification. So the gating lives in the package as a status: +```python +from enum import Enum +class Status(Enum): + VERIFIED = "verified" # has a golden anchor + passing [M2-ANE] check + EXPERIMENTAL = "experimental" # convertible but not yet anchored/verified + +# internal registry, single source of truth. +# KEY by ModelVersion enum MEMBER so list_* can emit .name. Keying by the lowercase +# .value string returns ["sd15",...], which the node reverses via ModelVersion[...] -> KeyError. +_MODEL_STATUS = {ModelVersion.SD15: Status.VERIFIED, ModelVersion.SDXL: Status.VERIFIED, + ModelVersion.SDXL_REFINER: Status.EXPERIMENTAL, ModelVersion.LCM: Status.EXPERIMENTAL} + +def list_model_versions(include_experimental: bool = False) -> list[str]: + return [v.name for v, s in _MODEL_STATUS.items() # .name -> "SD15","SDXL"; node reverses with ModelVersion[...] + if s is Status.VERIFIED or (include_experimental and s is Status.EXPERIMENTAL)] +``` +Consequence: **promoting a model to VERIFIED in the package expands the Suite's dropdown with no +Suite change and no Suite bump** — exactly the requirement. The act of verification (E-LCM +produces an LCM golden anchor; same later for refiner) is what flips the status. The Suite's +converter node calls `list_model_versions()` (verified-only); a power-user/CLI path may pass +`include_experimental=True`. Promotion VERIFIED-from-EXPERIMENTAL is additive (minor bump); +demotion or removal is breaking (major bump + note). + +--- + +## Naming & layout (chosen — frozen at Gate E0) + +**Distribution name (PyPI):** `coreml-diffusion`. **Import name (Python):** `coreml_diffusion`. +(PyPI normalizes `-`/`_`; the distribution uses the hyphen, the importable module the underscore.) +Availability checked: both `coreml-diffusion` and the near variants were free on PyPI at E0. + +**Why this name (the positioning it encodes):** the project's niche is *diffusion models on Apple +Neural Engine via CoreML, inside ComfyUI and on-device* — **not** Stable Diffusion specifically. +`sd*` was rejected because it falsely narrows scope to SD; `coreml-diffusion` keeps `coreml` on the +front for discoverability while `diffusion` honestly states the scope (SD/SDXL/LCM today, Flux and +other diffusion architectures later) **without** promising arbitrary non-diffusion torch models, +whose tracing/shape/sample-input pipeline differs. The name must not be re-narrowed to SD in +future docs. ANE is the *differentiator* (documented in the README), but `coreml` was chosen over +`ane` in the name for search discoverability per maintainer decision. + +Target package layout (framework-free — zero `comfy` imports): + +``` +coreml_diffusion/ + __init__.py # public API surface (see "Public API" below) + model_version.py # ModelVersion enum — the SINGLE source of truth, no comfy + attention.py # ATTENTION_IMPLEMENTATIONS tuple (from coreml_suite/attention.py) + apply_attention_implementation + pipeline.py # get_pipeline (from_single_file), get_unet (cml UNet from ref unet) + unet.py # UNet2DConditionModelLCM (moved from coreml_suite/lcm/unet.py) + inputs.py # get_sample_input, lcm_inputs, sdxl_inputs, + # get_encoder_hidden_states_shape, get_coreml_inputs, get_inputs_spec + controlnet.py # add_cnet_support (conversion-side residual SHAPE calc only) + convert.py # convert_unet, convert (orchestration), convert_to_coreml, load_coreml_model + compile.py # compile_coreml_model + quantize.py # (Phase E6 / MODERNIZATION Phase 6 lands here) palettization 4/6/8-bit + cli.py # console entry point: `coreml-diffusion convert ...` + pyproject.toml # standalone packaging (at E5) +``` + +What stays in `coreml_suite/` (the ComfyUI side, thinned): +- `nodes.py` — still owns **name-encoding** (`out_name` construction), path resolution via + `folder_paths`, the node `INPUT_TYPES`/mappings, and wrapping the result in `CoreMLModel`. +- `models.py`, `latents.py`, `controlnet.py` (inference parts), `lcm/utils.py`, `config.py` + (inference config build) — untouched by this spec except the import-source of `ModelVersion`. + +### Public API (the contract `coreml_diffusion` exposes) +```python +from coreml_diffusion import ModelVersion, convert, compile_model, compose_out_name +from coreml_diffusion import list_model_versions, list_attention_impls, list_quant_modes, CONTRACT_VERSION + +# Mirror the CURRENT converter.py signature, made keyword-only past the required positionals +# and with paths/device injected (no folder_paths, no comfy.model_management): +# convert(ckpt_path, model_version, out_path, *, +# batch_size=1, sample_size=(64, 64), controlnet_support=False, +# lora_weights=None, attn_impl=list_attention_impls()[0], config_path=None, +# quantize_nbits="none", device=None) -> None # side effect: writes out_path +# (current convert() returns None and writes via convert_unet → coreml_unet.save; keep that, +# or change to `return out_path` as a deliberate, documented improvement — pick one at E0.) +# compile_model(src_path, out_dir, final_name) -> str # returns compiled .mlmodelc path +``` +Note: `convert` takes an **explicit `out_path`** — no `folder_paths`. `device` is injected +(defaults to torch's default device). `compose_out_name` lives here (cache-key contract) and the +node imports it from the package. The `list_*` discovery functions back the node's dropdowns. + +--- + +## The import chains to cut (root cause inventory) — REVISED against current code + +> **State note (verified):** the code moved on since the original draft. Several chains are +> already cut. Re-verify each line by `grep` before acting; do not assume the original draft. + +**Already done (verify, then skip):** +- ✅ `converter.py` already imports `from coreml_suite.model_version import ModelVersion`, and + `model_version.py` is **clean** (`from enum import Enum` only — zero comfy). The old + "converter → config → comfy" chain is **already broken**. `config.py` still imports comfy, but + it is **inference-side** (`get_model_config` via `supported_models_base`/`latent_formats`) — + *not* on the conversion path. Do **not** treat `config.py` as a converter dependency. +- ✅ `converter.py` now uses `diffusers.UNet2DConditionModel.from_single_file` and a local + `CoreMLUNetWrapper` (in `coreml_suite/conversion/unet.py`) — it is **no longer** importing the + Apple `python_coreml_stable_diffusion.unet.UNet2DConditionModel*` internals on the main path. + A `coreml_suite/conversion/` subpackage already exists (`attention`, `shapes`, `trace`, `unet`). +- ✅ Name-encoding already extracted to `coreml_suite/core/naming.py` (`compose_out_name`, + `lora_names_from_params`, `ATTN_SUFFIX`, `QUANT_NBITS_VALUES`) **with characterization tests** + (`tests/unit/test_characterization_out_name.py`). The pure-naming split is done. +- ✅ Quantization is **already implemented** in `converter.py` (`quantize_nbits`, k-means + `palettize_weights`) and surfaced as an optional node input. Phase E6 is therefore *move*, not + *build* (see revised E6). + +**Still to cut (the real remaining work):** +1. `coreml_suite/converter.py::get_out_path` → `from folder_paths import get_folder_paths`. + Main converter still reaches into ComfyUI's model dir. **Cut: `out_path` is an injected arg; + `folder_paths` resolution moves up into the node** (the node already computes `out_name`). +2. `coreml_suite/lcm/converter.py` → still has its **own** `from folder_paths import + get_folder_paths` (`get_out_path`) and (per original draft) `comfy.model_management`. Verify + the current LCM file and cut both: inject `out_path` and `device`. +3. Global mutation of the attention impl: confirm where it now lives. Main path appears to route + through `coreml_suite/conversion/attention.apply_attention_implementation` (cleaner than the + old global), but `lcm/converter.py` may still set a module global at import. **Ensure the + package sets attention per-call, never at import time.** +4. **Duplication LCM vs main:** `lcm/converter.py` still carries its own copies of + `convert_to_coreml`, `load_coreml_model`, `get_out_path`, `get_sample_input` (the LCM variant + takes a `scheduler` arg), and hardcodes `SimianLuo/LCM_Dreamshaper_v7`. **Dedupe into the + single `coreml_diffusion` implementation;** the HF-hardcode consolidation is the *behavior-changing* + part → deferred to optional **E-LCM**, not E1–E5. +5. **`compose_out_name` ownership:** currently in `coreml_suite/core/naming.py` and called by the + node. Per the Interface-contract pillar it must **move into the package** (it is the cache-key + contract) and the node must import it from `coreml_diffusion`, not keep a copy. + +--- + +## Phase E0 — Seam decision & inventory (no code change) + +**Objective:** lock the cut line, the interface contract, and naming so later phases don't drift. + +### Tasks +1. Produce `docs/extraction/seam.md`: a table of every symbol in `converter.py`, + `lcm/converter.py`, `lcm/unet.py`, **plus the already-extracted `conversion/` subpackage + (`attention`, `shapes`, `trace`, `unet`) and `core/naming.py`**, classified + **CONVERSION → coreml_diffusion** vs **STAYS (comfy/node)**. Note which are already framework-free. +2. ~~Confirm the current `python_coreml_stable_diffusion` footprint.~~ **DONE (seam.md §6): + footprint is ZERO** — no runtime imports anywhere; only a docstring mention in + `core/__init__.py:4`. Main path uses `diffusers` + local `CoreMLUNetWrapper`; the runtime + `CoreMLModel` (STAYS in suite) is a local coremltools wrapper, not Apple's. No shape/attn helper + comes from Apple (local `conversion/shapes.py`, `conversion/attention.py`). +3. **Decide the interface contract concretely (the maintainer's hard requirement):** + - Discovery functions `list_model_versions / list_attention_impls / list_quant_modes` live in + the package and return today's display strings verbatim. Node `INPUT_TYPES` calls them. + - `ModelVersion` values, attn-impl strings, quant modes are **ADDITIVE-ONLY** across package + versions; removal/rename = MAJOR bump + migration note. Write this into the package's + versioning policy doc now. + - `compose_out_name` moves to the package; node imports it (no copy). Confirm the + characterization tests in `test_characterization_out_name.py` will be re-pointed, not + duplicated. + - **Resolve the `model_version` dropdown question (maintainer decided):** the Suite exposes + *every model the converter can verifiably convert*. `lcm` and `sdxl_refiner` are absent today + only because they lack a golden/PSNR verification — **not** because the node hardcodes a + short list. Encode this as a **status registry in the package** (`VERIFIED` vs + `EXPERIMENTAL`); `list_model_versions()` returns VERIFIED-only by default. The converter node + calls it plainly. Promoting LCM/refiner to VERIFIED (after E-LCM / a refiner anchor) expands + the dropdown with **no Suite change**. Do NOT add permanent per-node filtering — the gate is + verification status, owned by the package. +4. ~~Confirm the `ml-stable-diffusion` git dep is pinned.~~ **N/A — already removed (#58).** Verified: + zero `python_coreml_stable_diffusion` imports in the repo; `CoreMLModel` is now a local + coremltools wrapper; the dep is absent from `pyproject.toml`/`requirements.txt`. No SHA to pin. + +### STOP — VALIDATE (Gate E0) +``` +## Gate E0 report +- seam.md committed: ; symbol counts (move / stay / already-framework-free) +- python_coreml_stable_diffusion usage (verified by grep): conversion= runtime= +- Discovery API signatures frozen: list_model_versions (verified-only) / list_attention_impls / list_quant_modes +- Status registry decided: sd15+sdxl=VERIFIED, lcm+sdxl_refiner=EXPERIMENTAL (gated, not hidden) +- Additive-only contract policy doc written (incl. promotion=minor, demotion/removal=major): +- model_version dropdown: expose all (incl. LCM/REFINER) / filtered per node — DECISION: <...> +- compose_out_name move-not-copy confirmed; tests re-point plan: <...> +- LCM consolidation deferred to optional E-LCM: YES/NO +- ml-stable-diffusion: N/A — already removed (#58), not a dependency (was: pin-or-BLOCKER) +- Package name in-repo: coreml_diffusion (final PyPI name deferred to E5) +``` + +--- + +## Phase E1 — Establish `coreml_diffusion` package + discovery API (mostly verification) + +**Objective:** stand up the package namespace and the discovery surface. Much of the comfy-chain +cut is **already done** — this phase mostly *verifies* that and adds the discovery functions. + +### Tasks +1. **Verify (don't redo):** `coreml_suite/model_version.py` is already clean (`Enum` only). Confirm + `import coreml_suite.model_version` works with **no comfy** (`uv run python -c "..."` in a + comfy-free `uv venv`). If true, E1's original "extract ModelVersion" task is already satisfied. +2. Create the `coreml_diffusion/` package skeleton with `__init__.py` exporting the **discovery API** + backed by the *existing* sources of truth for now (re-export `ModelVersion`, the + `ATTENTION_IMPLEMENTATIONS` tuple, and `QUANT_NBITS_VALUES`) so values are byte-identical: + ```python + def list_model_versions(): return [v.name for v in ModelVersion] # .name -> "SD15" (node reverses via ModelVersion[...]; .value KeyErrors) + def list_attention_impls(): return list(ATTENTION_IMPLEMENTATIONS) + def list_quant_modes(): return list(QUANT_NBITS_VALUES) + CONTRACT_VERSION = "1.0" + ``` + (At this stage `coreml_diffusion` may live inside the repo and import from `coreml_suite.*`; the + physical move of implementation happens in E2. The point of E1 is to freeze the *contract*.) +3. **Decided (`.name`):** the node renders `ModelVersion.SD15.name` (`"SD15"`) and reverses the + dropdown string via `ModelVersion[model_version]` (name lookup, `nodes.py:286`). Discovery API + therefore returns `.name`; `.value` (`"sd15"`) would `KeyError`. Recorded in `seam.md` §5. + +### Acceptance criteria +- `uv run python -c "import coreml_diffusion; print(coreml_diffusion.list_model_versions(), coreml_diffusion.list_quant_modes())"` + works in a **comfy-free** `uv venv` and prints today's exact strings. +- Existing characterization tests pass unchanged. +- No node behavior change yet (node still uses its current hardcoded lists in E1). + +### STOP — VALIDATE (Gate E1) +``` +## Gate E1 report +- model_version.py confirmed comfy-free (uv, no comfy): PASS/FAIL +- coreml_diffusion.list_* returns byte-identical strings to current dropdowns: YES/NO (show values) +- .name vs .value decision for model_version discovery: <...> +- CONTRACT_VERSION set; additive-only policy linked: +- Characterization tests unchanged & green (uv run pytest): YES/NO +``` + +--- + +## Phase E2 — Move conversion code into `coreml_diffusion` (in-repo, dedup, behavior-preserving) + +**Objective:** physically relocate the conversion mechanics into the framework-free package, +collapsing the two duplicate converters into one, with paths/device injected. + +### Tasks +1. Move into `coreml_diffusion/`: `pipeline.py` (`get_pipeline`, `get_unet`), `unet.py` + (`UNet2DConditionModelLCM`), `inputs.py` (sample/lcm/sdxl input builders + + `get_encoder_hidden_states_shape` + `get_coreml_inputs` + `get_inputs_spec`), + `controlnet.py` (`add_cnet_support`), `convert.py` (`convert_unet`, `convert`, + `convert_to_coreml`, `load_coreml_model`), `compile.py` (`compile_coreml_model`). +2. **Dedupe LCM vs main** (the real remaining duplication): delete `lcm/converter.py`'s copies of + `convert_to_coreml` / `load_coreml_model` / `get_out_path` / `get_sample_input` (LCM variant + carries a `scheduler` arg — fold that into the shared `get_sample_input` as an optional param) + in favor of the single `coreml_diffusion` implementation. The main path's helpers + (`get_unet`/`get_encoder_hidden_states_shape`/`get_coreml_inputs`/`convert_unet`/`convert`) and + the `conversion/` subpackage (`attention`, `shapes`, `trace`, `unet`) move as-is. +3. **Inject paths**: replace `get_out_path`'s `folder_paths` reach-in with an injected `out_path` + argument on `convert(...)`; `folder_paths` resolution moves up into the node (which already + computes `out_name`). No `folder_paths` import anywhere in `coreml_diffusion`. +4. **Inject device** where the LCM path used `comfy.model_management` (verify it still does): + `convert(..., device=None)`, default to torch's default device. +5. **Attention per-call, never at import:** main path already routes through + `conversion/attention.apply_attention_implementation` — keep that. If `lcm/converter.py` still + sets any module global at import, remove it; the package sets attention from the `attn_impl` + arg inside `convert`. +6. **Move `compose_out_name` into the package** (`coreml_diffusion/naming.py`); re-point + `test_characterization_out_name.py` imports to `coreml_diffusion.naming` — assertions and values + unchanged. The node will import it from the package in E3. +7. Leave **thin shims** in `coreml_suite/converter.py` and `coreml_suite/lcm/converter.py` that + re-export from `coreml_diffusion`, preserving the old call signatures the nodes use (nodes untouched + this phase). Shims map comfy `folder_paths`/device into package args. + +### Acceptance criteria +- `uv run pytest -m unit` (Tier 0) imports `coreml_diffusion.*` with **no comfy / no MPS** and is green on Linux. +- The dedup leaves exactly one implementation of each previously-duplicated function. +- Characterization tests pass unchanged after the `compose_out_name` re-point. +- `[M2]` A real SD1.5 conversion via the shim still produces a loadable model. +- `[M2-ANE]` **Golden latent identical / within tolerance** to the MODERNIZATION Phase 2 anchor + (same seed/prompt) — proves the move + dedup changed nothing. + +### STOP — VALIDATE (Gate E2 — first regression gate) +``` +## Gate E2 report +- Tier 0 import of coreml_diffusion without comfy/MPS (uv run): PASS/FAIL +- LCM/main duplicated funcs collapsed to one (list old→new): +- compose_out_name moved to package; char-tests re-pointed & green: YES/NO +- Paths injected (no folder_paths in package): confirmed +- Device injected (no comfy.model_management in package): confirmed +- Attention set per-call, not at import (both main & lcm): confirmed +- [M2-ANE] Golden latent vs Phase-2 anchor: identical / within tol / DIVERGED (STOP) +- Node INPUT_TYPES / mappings untouched: confirmed (diff) +``` +**If the golden latent diverged at all, STOP and report — do not continue.** + +--- + +## Phase E3 — Thin the nodes onto the package (behavior-preserving) + +**Objective:** remove the shims; have the ComfyUI nodes call `coreml_diffusion` directly, keeping the +node contract byte-identical. + +### Tasks +1. `CoreMLConverter.convert` (in `coreml_suite/nodes.py`): keep the `folder_paths`-based path + resolution **in the node**; import `compose_out_name` from `coreml_diffusion` (not `coreml_suite.core`); + call `coreml_diffusion.convert(...)` and `coreml_diffusion.compile_model(...)` directly; wrap the compiled path + in `CoreMLModel`. +2. **Wire the dropdowns to discovery (the maintainer's hard requirement).** Replace the hardcoded + `INPUT_TYPES` lists with fail-soft discovery calls: + ```python + def _discover(fn, fallback): + try: + import coreml_diffusion + return getattr(coreml_diffusion, fn)() + except Exception as e: # missing/old package, or import error + logger.warning(f"coreml_diffusion.{fn} unavailable ({e}); using fallback {fallback}") + return fallback + ... + "model_version": (_discover("list_model_versions", ["SD15", "SDXL"]),), + "attention_implementation": (_discover("list_attention_impls", ["SPLIT_EINSUM","SPLIT_EINSUM_V2","ORIGINAL"]),), + "quantize_nbits": (_discover("list_quant_modes", ["none","8","6","4"]), {"default": "none"}), + ``` + This is what makes "update the package → new types appear in the old node, no Suite bump" true. +3. `COREML_CONVERT_LCM` (in `coreml_suite/lcm/nodes.py`): route through `coreml_diffusion` for the shared + mechanics. **Keep the existing LCM behavior/HF-hardcode for now** — consolidation is optional E-LCM. +4. Delete the now-dead `coreml_suite/converter.py` / `coreml_suite/lcm/converter.py` shims (or + reduce to a one-line re-export if anything external imports them — grep first). + +### Acceptance criteria +- `NODE_CLASS_MAPPINGS` / `NODE_DISPLAY_NAME_MAPPINGS` keys: **unchanged** (diff `__init__.py`). +- Every `INPUT_TYPES` **field name** unchanged. Dropdown **values**: the discovery calls must + return **a superset of** today's values, with every previously-present value still present and + spelled identically (additive-only). *(This deliberately replaces the original spec's + "values must be byte-identical/frozen" criterion — the maintainer requires the list be + extensible at runtime. Frozen-field-names + additive-only-values is the new contract.)* +- With `coreml_diffusion` **absent**, the node still registers and shows the fallback lists (fail-soft). +- `[M2-ANE]` Golden latent still identical to the Phase-2 anchor. +- `[M2-ANE]` The committed e2e workflow `tests/integration/...` still passes (PSNR > 25). + +### STOP — VALIDATE (Gate E3) +``` +## Gate E3 report +- Node mappings diff: empty (confirmed) +- INPUT_TYPES field-names diff: empty (confirmed) +- Dropdown values: superset of prior, all prior values still present & identical: YES/NO (show) +- Fail-soft with coreml_diffusion absent (node still registers): PASS/FAIL +- compose_out_name now imported from coreml_diffusion (no node-side copy): confirmed +- [M2-ANE] Golden latent vs anchor: identical / within tol / DIVERGED (STOP) +- [M2-ANE] e2e workflow PSNR: (> 25?) +- Dead converter shims removed / reduced: +``` + +--- + +## Phase E4 — Standalone packaging & CLI (still in-repo) + +**Objective:** make `coreml_diffusion` independently installable and usable without ComfyUI, with a CLI +suitable for the planned article and for on-device/iOS conversion workflows. + +### Tasks +1. Add `coreml_diffusion/pyproject.toml`: name (working `coreml_diffusion`), `requires-python`, dependencies + = `coremltools` (pinned to the MODERNIZATION-validated version), `diffusers`, `transformers`, + `peft`, `omegaconf`, `numpy`, `torch`. **No `ml-stable-diffusion`** (already removed in #58, see + §0.3) and **no comfy**. Suite pins `transformers>=4.44`/`peft>=0.13`/`omegaconf>=2.3` today; + grep-confirm each is on the conversion path before listing it. A `[project.scripts]` entry: + `coreml-diffusion = "coreml_diffusion.cli:main"`. +2. `coreml_diffusion/cli.py`: `coreml-diffusion convert --ckpt PATH --model-version sd15 --out PATH + [--height --width --batch-size --attn-impl --controlnet --lora NAME:STRENGTH ... --config PATH]` + and `coreml-diffusion compile --src PATH --out-dir DIR --name NAME`. Mirrors `convert()`/`compile_model()`. +3. Tier-0 Linux tests for the CLI **arg→call mapping** (mock the heavy `convert`); the real + convert remains `[M2]`. Add a `[M2]` smoke test: convert a tiny synthetic UNet end-to-end. +4. README for the package: install, CLI usage, "produce a `.mlpackage`/`.mlmodelc` for use in a + Swift/iOS app", and the ANE positioning note (low-power, GPU-free, embeddable; SD1.5/SDXL on + ANE, **not** a Flux-speed claim). + +### Acceptance criteria +- Fresh `python -m venv` + `uv pip install ./coreml-diffusion` (no ComfyUI present) imports and runs + `coreml-diffusion --help` and the arg-mapping tests on Linux. +- `[M2]` `coreml-diffusion convert` produces a model file identical (golden) to the node path. + +### STOP — VALIDATE (Gate E4) +``` +## Gate E4 report +- uv pip install ./coreml-diffusion in comfy-free venv: PASS/FAIL (log) +- CLI arg→call tests (Tier 0, Linux): green +- [M2] CLI-produced model golden vs node-produced model: identical / DIVERGED +- Package deps list (with pinned SHAs/versions + licenses): +- New runtime deps vs suite before: +``` +**This is the gate that proves the package stands alone. Do not split repos before it passes.** + +--- + +## Phase E5 — Physical split into a second repository + +**Objective:** move `coreml_diffusion/` to its own repo; CoreMLSuite depends on it by pinned version. + +### Tasks +1. Create the new repo (maintainer action — agent prepares the tree, not the GitHub repo). + Choose final distributable name; rename imports if changed (single sweep, recorded). +2. CoreMLSuite `pyproject.toml` / `requirements.txt`: replace the conversion-only deps with a + pinned dependency on the new package (`coreml_diffusion==` from PyPI, or `git+...@` + until first PyPI release). (There is no `git+...ml-stable-diffusion` line to remove — already + gone since #58.) +3. ~~Keep `python_coreml_stable_diffusion` for the loader.~~ **Void.** The loader is the local + `coreml_suite/coreml_model.py` over `coremltools`; the suite keeps `coremltools` as a direct dep + for it. No Apple lib involved. +4. Set up the new repo's CI: Tier 0 on Linux (import + arg-mapping + input-shape math), + `[M2]`/`[M2-ANE]` on a self-hosted/macOS-ARM runner reusing the golden-latent anchor. +5. Versioning: SemVer; first release `0.1.0`. Document the compatibility matrix + (coreml_diffusion ↔ coremltools version ↔ diffusers version). No ml-stable-diffusion axis. + +### Acceptance criteria +- CoreMLSuite installs in a fresh venv pulling the new package; e2e workflow still passes `[M2-ANE]`. +- New repo CI green on Linux (Tier 0) and `[M2-ANE]` golden latent matches the anchor. +- No conversion code remains in CoreMLSuite (grep: no `ct.convert`, no `from_single_file`, + no `torch.jit.trace`). + +### STOP — VALIDATE (Gate E5) +``` +## Gate E5 report +- New repo tree prepared: ; final package name: +- Suite depends on package by pinned version: +- Suite e2e [M2-ANE] PSNR after split: (> 25?) +- Conversion code fully absent from suite: confirmed (grep output) +- Compatibility matrix documented: +- First release tag: 0.1.0 +``` + +--- + +## Phase E6 — Quantization travels WITH the conversion code (already implemented → move) + +**Objective:** quantization is **already implemented** (k-means `palettize_weights` in +`converter.py`, `quantize_nbits` node input, `_q` filename suffix, README tradeoff table). +There is nothing to *build*. It simply **moves with the conversion code in E2** as part of +`convert_unet`. This phase is a checkpoint that it survived the extraction intact, plus exposing +it through the CLI. + +### Tasks +1. Confirm the palettization block moved cleanly into `coreml_diffusion` (lives in `convert.py` or a + `quantize.py` helper called from `convert_unet`). Default `"none"` stays byte-identical. +2. Expose via CLI flag `--quantize {none,8,6,4}` (E4 already lists this) and via + `list_quant_modes()` discovery (E1/E3). +3. The existing README tradeoff table (SD1.5 1×512×512 SPLIT_EINSUM: none/8/6/4 → size/ms/PSNR) + moves to the package README. Re-confirm one row `[M2-ANE]` so the article can cite a live number. + +### Acceptance criteria +- Default (`none`) output byte-identical to pre-extraction (covered by the E2/E3 golden latent). +- `coreml-diffusion convert --quantize 4` produces a `_q4` artifact matching the node's `_q4` artifact `[M2]`. +- `list_quant_modes()` drives the node dropdown (no hardcoded copy remains). + +### STOP — VALIDATE (Gate E6) +``` +## Gate E6 report +- Palettization relocated into coreml_diffusion, called from convert_unet: confirmed +- Default none output identical (golden): YES/NO +- [M2] CLI --quantize {8,6,4} artifacts match node artifacts: YES/NO +- Tradeoff table in package README with at least one re-confirmed [M2-ANE] row: +``` + +--- + +## Phase E-LCM — FIRST task after the split: clean up LCM + verify → promote (behavior-changing, gated) + +> Promoted from "optional, someday" to **the first thing after E5**, per maintainer intent: the +> Suite should expose every verifiably-convertible model, and LCM is the obvious first cleanup. + +Two coupled goals: +1. **Consolidate the LCM path.** Make the LCM node use the unified `from_single_file` path in + `coreml_diffusion.convert(model_version=LCM, ...)` instead of the hardcoded `SimianLuo/LCM_Dreamshaper_v7` + HF download; drop the duplicated LCM helpers (already deduped in E2). **Behavior change** ⇒ + capture an LCM golden anchor *before* the change, then prove within-tolerance after. +2. **Verify → promote.** Once the LCM conversion has a passing `[M2-ANE]` golden anchor, flip + `_MODEL_STATUS["lcm"] = Status.VERIFIED` **in the package** (minor bump). The Suite's dropdown + gains `lcm` automatically — no Suite change, no Suite bump. This is the end-to-end proof that + the discovery contract works as designed. + +Repeat the same recipe for `sdxl_refiner` when it gets an anchor (separate small gate). Do NOT +bundle E-LCM into E1–E5; it changes behavior and must stand on its own golden. + +### STOP — VALIDATE (Gate E-LCM) +``` +## Gate E-LCM report +- LCM golden anchor captured BEFORE change: +- LCM node now uses unified from_single_file path; HF hardcode removed: confirmed +- [M2-ANE] LCM golden after change: identical / within tol / DIVERGED (STOP) +- Status flipped lcm→VERIFIED in package (minor bump ): confirmed +- Suite dropdown now lists lcm with NO Suite code change / NO Suite bump: confirmed (diff empty) +- LCM node accepts a checkpoint arg now (documented breaking-ish UI note): +``` + +--- + +## Article deliverable (after E4) + +Once the CLI exists and stands alone, the "convert a Comfy/A1111 workflow into an on-device iOS +app" write-up becomes a clean tutorial: `coreml-diffusion convert` → `.mlmodelc` → load in Swift/CoreML. +Frame the niche honestly per the README note above (ANE feasibility & power, not raw Flux speed). + +--- + +## Quick reference: extraction gate discipline + +``` +E0 Seam decision, interface contract, discovery API frozen → Gate E0 (cut line + additive-only policy?) +E1 Stand up coreml_diffusion + discovery API (mostly verify) → Gate E1 (list_* byte-identical, comfy-free?) +E2 Move conversion code, dedup LCM/main, inject paths/device→ Gate E2 (golden identical? duplicates gone?) ← first regression gate +E3 Thin nodes onto package + wire discovery dropdowns → Gate E3 (field-names frozen, values additive, fail-soft, golden identical?) +E4 Standalone packaging + CLI (uv) → Gate E4 (uv pip install w/o comfy? CLI golden?) ← proves it stands alone +E5 Physical second-repo split → Gate E5 (suite depends on pkg? conversion absent?) +E-LCM FIRST post-split: clean up LCM, verify → promote → Gate E-LCM (LCM golden? dropdown gains lcm w/ no Suite bump?) +E6 Quantization checkpoint (already built → moved in E2) → Gate E6 (default identical? CLI quant matches?) +(refiner) same recipe as E-LCM when an anchor exists → own small gate (promote sdxl_refiner→VERIFIED) +``` + +**Interface-contract invariants (the maintainer's hard requirement), restated:** +- Package API is keyword-only-with-defaults past the required positionals → converter updates + don't force Suite updates. +- Node dropdowns are discovery-driven (`coreml_diffusion.list_*`) + fail-soft → new conversion types + appear in the old plugin with `uv pip install -U coreml_diffusion`, **no Suite code change, no bump**. +- Discovery identifiers are **additive-only**; removal/rename = MAJOR bump + migration note. +- `compose_out_name` (cache key) lives in the package, single copy. + + +**Golden rule (inherited): never cross a gate with a failing acceptance criterion. +Stop, report, wait. The golden latent is the single source of truth that the extraction +changed nothing.** diff --git a/coreml_suite/attention.py b/coreml_suite/attention.py deleted file mode 100644 index ffeb979..0000000 --- a/coreml_suite/attention.py +++ /dev/null @@ -1,5 +0,0 @@ -ATTENTION_IMPLEMENTATIONS = ( - "SPLIT_EINSUM", - "SPLIT_EINSUM_V2", - "ORIGINAL", -) diff --git a/coreml_suite/config.py b/coreml_suite/config.py index 54a6562..ebf582d 100644 --- a/coreml_suite/config.py +++ b/coreml_suite/config.py @@ -4,7 +4,7 @@ from comfy import supported_models_base from comfy import latent_formats from comfy.model_detection import convert_config -from coreml_suite.model_version import ModelVersion +from coreml_diffusion import ModelVersion config_map = { diff --git a/coreml_suite/conversion/__init__.py b/coreml_suite/conversion/__init__.py deleted file mode 100644 index 6eee9cb..0000000 --- a/coreml_suite/conversion/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Core ML conversion helpers. - -The conversion approach originates from Apple's ml-stable-diffusion -(https://github.com/apple/ml-stable-diffusion). This implementation has since -diverged: it runs natively on diffusers' UNet2DConditionModel with its own -SPLIT_EINSUM / SPLIT_EINSUM_V2 attention processors and no longer depends on -that package. The intent is to keep iterating on these methods independently -while tracking current tooling. -""" diff --git a/coreml_suite/conversion/attention.py b/coreml_suite/conversion/attention.py deleted file mode 100644 index de6f966..0000000 --- a/coreml_suite/conversion/attention.py +++ /dev/null @@ -1,239 +0,0 @@ -import logging - -import torch - -logger = logging.getLogger(__name__) - -CHUNK_SIZE = 512 - - -def apply_attention_implementation(unet, attention_implementation): - if attention_implementation == "ORIGINAL": - return unet - - if attention_implementation == "SPLIT_EINSUM": - unet.set_attn_processor(SplitEinsumAttnProcessor()) - return unet - - if attention_implementation == "SPLIT_EINSUM_V2": - unet.set_attn_processor(SplitEinsumV2AttnProcessor()) - return unet - - raise ValueError(f"Unsupported attention implementation: {attention_implementation}") - - -class SplitEinsumAttnProcessor: - def __call__( - self, - attn, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - temb=None, - *args, - **kwargs, - ): - return _attention_forward( - attn, - hidden_states, - encoder_hidden_states, - attention_mask, - temb, - split_einsum, - ) - - -class SplitEinsumV2AttnProcessor: - def __call__( - self, - attn, - hidden_states, - encoder_hidden_states=None, - attention_mask=None, - temb=None, - *args, - **kwargs, - ): - return _attention_forward( - attn, - hidden_states, - encoder_hidden_states, - attention_mask, - temb, - split_einsum_v2, - ) - - -def _attention_forward( - attn, - hidden_states, - encoder_hidden_states, - attention_mask, - temb, - attention_fn, -): - residual = hidden_states - - if attn.spatial_norm is not None: - hidden_states = attn.spatial_norm(hidden_states, temb) - - input_ndim = hidden_states.ndim - if input_ndim == 4: - batch_size, channel, height, width = hidden_states.shape - hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) - else: - batch_size, _, channel = hidden_states.shape - height = None - width = None - - batch_size, key_sequence_length, _ = ( - hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape - ) - - if attention_mask is not None: - attention_mask = attn.prepare_attention_mask( - attention_mask, - key_sequence_length, - batch_size, - ) - attention_mask = _prepare_split_einsum_mask( - attention_mask, - batch_size, - attn.heads, - key_sequence_length, - ) - - if attn.group_norm is not None: - hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) - - query = attn.to_q(hidden_states) - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - elif attn.norm_cross: - encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) - - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - batch_size = query.shape[0] - dim_head = attn.inner_kv_dim // attn.heads - - query = _linear_projection_to_bchw(query) - key = _linear_projection_to_bchw(key) - value = _linear_projection_to_bchw(value) - - hidden_states = attention_fn( - query, - key, - value, - attention_mask, - attn.heads, - dim_head, - ) - hidden_states = hidden_states.squeeze(2).transpose(1, 2) - hidden_states = hidden_states.reshape(batch_size, -1, attn.inner_dim) - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - - if input_ndim == 4: - hidden_states = hidden_states.transpose(-1, -2).reshape( - batch_size, - channel, - height, - width, - ) - - if attn.residual_connection: - hidden_states = hidden_states + residual - - hidden_states = hidden_states / attn.rescale_output_factor - return hidden_states - - -def split_einsum(q, k, v, mask, heads, dim_head): - q_heads = _split_heads(q, heads, dim_head) - k = k.transpose(1, 3) - k_heads = [ - k[:, :, :, head_idx * dim_head : (head_idx + 1) * dim_head] - for head_idx in range(heads) - ] - v_heads = _split_heads(v, heads, dim_head) - - weights = [ - torch.einsum("bchq,bkhc->bkhq", query, key) * (dim_head**-0.5) - for query, key in zip(q_heads, k_heads) - ] - if mask is not None: - weights = [weight + mask for weight in weights] - - weights = [weight.softmax(dim=1) for weight in weights] - outputs = [ - torch.einsum("bkhq,bchk->bchq", weight, value) - for weight, value in zip(weights, v_heads) - ] - return torch.cat(outputs, dim=1) - - -def split_einsum_v2(q, k, v, mask, heads, dim_head): - query_length = q.size(3) - num_chunks = query_length // CHUNK_SIZE - if num_chunks == 0: - logger.info( - "SPLIT_EINSUM_V2 query sequence is shorter than %s; using SPLIT_EINSUM.", - CHUNK_SIZE, - ) - return split_einsum(q, k, v, mask, heads, dim_head) - - q_heads = _split_heads(q, heads, dim_head) - q_chunks = [ - [ - head[..., chunk_idx * CHUNK_SIZE : (chunk_idx + 1) * CHUNK_SIZE] - for chunk_idx in range(num_chunks) - ] - for head in q_heads - ] - - k = k.transpose(1, 3) - k_heads = [ - k[:, :, :, head_idx * dim_head : (head_idx + 1) * dim_head] - for head_idx in range(heads) - ] - v_heads = _split_heads(v, heads, dim_head) - - head_outputs = [] - for query_chunks, key, value in zip(q_chunks, k_heads, v_heads): - chunk_outputs = [] - for query_chunk in query_chunks: - weights = torch.einsum("bchq,bkhc->bkhq", query_chunk, key) - weights = weights * (dim_head**-0.5) - if mask is not None: - weights = weights + mask - weights = weights.softmax(dim=1) - chunk_outputs.append(torch.einsum("bkhq,bchk->bchq", weights, value)) - head_outputs.append(torch.cat(chunk_outputs, dim=3)) - - return torch.cat(head_outputs, dim=1) - - -def _split_heads(x, heads, dim_head): - return [ - x[:, head_idx * dim_head : (head_idx + 1) * dim_head, :, :] - for head_idx in range(heads) - ] - - -def _linear_projection_to_bchw(x): - return x.transpose(1, 2).unsqueeze(2) - - -def _prepare_split_einsum_mask(mask, batch_size, heads, key_sequence_length): - if mask.ndim == 2: - mask = mask[:, None, :] - if mask.shape[0] == batch_size * heads: - mask = mask.reshape(batch_size, heads, -1, key_sequence_length) - mask = mask[:, 0] - if mask.ndim == 3: - mask = mask[:, :, None, None] - return mask diff --git a/coreml_suite/conversion/shapes.py b/coreml_suite/conversion/shapes.py deleted file mode 100644 index 1991e73..0000000 --- a/coreml_suite/conversion/shapes.py +++ /dev/null @@ -1,20 +0,0 @@ -def conv2d_output_shape(height, width, conv): - """Return the spatial output shape for a torch.nn.Conv2d-like module.""" - kernel_h, kernel_w = _pair(conv.kernel_size) - stride_h, stride_w = _pair(conv.stride) - pad_h, pad_w = _pair(conv.padding) - dilation_h, dilation_w = _pair(conv.dilation) - - out_h = _conv_output_dim(height, kernel_h, stride_h, pad_h, dilation_h) - out_w = _conv_output_dim(width, kernel_w, stride_w, pad_w, dilation_w) - return out_h, out_w - - -def _conv_output_dim(size, kernel, stride, padding, dilation): - return ((size + (2 * padding) - (dilation * (kernel - 1)) - 1) // stride) + 1 - - -def _pair(value): - if isinstance(value, tuple): - return value - return value, value diff --git a/coreml_suite/conversion/trace.py b/coreml_suite/conversion/trace.py deleted file mode 100644 index ab443bf..0000000 --- a/coreml_suite/conversion/trace.py +++ /dev/null @@ -1,61 +0,0 @@ -from types import MethodType - -from diffusers.models.transformers.transformer_2d import Transformer2DModel - - -def prepare_unet_for_coreml_trace(unet): - for module in unet.modules(): - if isinstance(module, Transformer2DModel): - module._operate_on_continuous_inputs = MethodType( - _operate_on_continuous_inputs, - module, - ) - module._get_output_for_continuous_inputs = MethodType( - _get_output_for_continuous_inputs, - module, - ) - return unet - - -def _operate_on_continuous_inputs(self, hidden_states): - hidden_states = self.norm(hidden_states) - - if not self.use_linear_projection: - hidden_states = self.proj_in(hidden_states) - inner_dim = self.inner_dim - hidden_states = hidden_states.flatten(2).transpose(1, 2) - else: - inner_dim = hidden_states.shape[1] - hidden_states = hidden_states.flatten(2).transpose(1, 2) - hidden_states = self.proj_in(hidden_states) - - return hidden_states, inner_dim - - -def _get_output_for_continuous_inputs( - self, - hidden_states, - residual, - batch_size, - height, - width, - inner_dim, -): - if not self.use_linear_projection: - hidden_states = hidden_states.transpose(1, 2).reshape( - batch_size, - inner_dim, - height, - width, - ) - hidden_states = self.proj_out(hidden_states) - else: - hidden_states = self.proj_out(hidden_states) - hidden_states = hidden_states.transpose(1, 2).reshape( - batch_size, - inner_dim, - height, - width, - ) - - return hidden_states + residual diff --git a/coreml_suite/conversion/unet.py b/coreml_suite/conversion/unet.py deleted file mode 100644 index 0dd01cd..0000000 --- a/coreml_suite/conversion/unet.py +++ /dev/null @@ -1,54 +0,0 @@ -import torch - - -class CoreMLUNetWrapper(torch.nn.Module): - """Adapt diffusers UNet inputs to CoreMLSuite's stable Core ML contract.""" - - def __init__(self, unet, model_version): - super().__init__() - self.unet = unet - self.model_version = model_version - - def forward(self, sample, timestep, encoder_hidden_states, *extra_inputs): - input_index = 0 - timestep_cond = None - if self._is_lcm: - timestep_cond = extra_inputs[input_index] - input_index += 1 - - added_cond_kwargs = None - if self._is_sdxl: - time_ids = extra_inputs[input_index] - text_embeds = extra_inputs[input_index + 1] - input_index += 2 - added_cond_kwargs = { - "time_ids": time_ids, - "text_embeds": text_embeds, - } - - additional_residuals = extra_inputs[input_index:] - down_residuals = None - mid_residual = None - if additional_residuals: - down_residuals = tuple(additional_residuals[:-1]) - mid_residual = additional_residuals[-1] - - outputs = self.unet( - sample, - timestep, - encoder_hidden_states=encoder_hidden_states, - timestep_cond=timestep_cond, - added_cond_kwargs=added_cond_kwargs, - down_block_additional_residuals=down_residuals, - mid_block_additional_residual=mid_residual, - return_dict=False, - ) - return outputs[0] - - @property - def _is_lcm(self): - return self.model_version.name == "LCM" - - @property - def _is_sdxl(self): - return self.model_version.name in {"SDXL", "SDXL_REFINER"} diff --git a/coreml_suite/converter.py b/coreml_suite/converter.py deleted file mode 100644 index 8acc149..0000000 --- a/coreml_suite/converter.py +++ /dev/null @@ -1,322 +0,0 @@ -import gc -import os -import time - -import coremltools as ct -import numpy as np -import torch -from diffusers import UNet2DConditionModel - -from coreml_suite.attention import ATTENTION_IMPLEMENTATIONS -from coreml_suite.conversion.attention import apply_attention_implementation -from coreml_suite.conversion.shapes import conv2d_output_shape -from coreml_suite.conversion.trace import prepare_unet_for_coreml_trace -from coreml_suite.conversion.unet import CoreMLUNetWrapper -from coreml_suite.logger import logger -from coreml_suite.model_version import ModelVersion - -DEFAULT_TRACE_TIMESTEP = 999.0 -TEXT_TOKEN_SEQUENCE_LENGTH = 77 - - -def get_unet(model_version: ModelVersion, ref_unet, attention_implementation): - ref_unet = prepare_unet_for_coreml_trace(ref_unet) - unet = apply_attention_implementation( - ref_unet.eval(), - attention_implementation, - ) - return CoreMLUNetWrapper(unet, model_version) - - -def get_encoder_hidden_states_shape(ref_unet, batch_size): - encoder_hidden_states_shape = ( - batch_size, - TEXT_TOKEN_SEQUENCE_LENGTH, - ref_unet.config.cross_attention_dim, - ) - - return encoder_hidden_states_shape - - -def get_coreml_inputs(sample_inputs): - coreml_sample_unet_inputs = { - k: v.numpy().astype(np.float16) for k, v in sample_inputs.items() - } - return [ - ct.TensorType( - name=k, - shape=v.shape, - dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype, - ) - for k, v in coreml_sample_unet_inputs.items() - ] - - -def load_coreml_model(out_path): - logger.info(f"Loading model from {out_path}") - - start = time.time() - coreml_model = ct.models.MLModel(out_path) - logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds") - - return coreml_model - - -def convert_to_coreml( - submodule_name, torchscript_module, sample_inputs, output_names, out_path -): - if os.path.exists(out_path): - logger.info(f"Skipping export because {out_path} already exists") - coreml_model = load_coreml_model(out_path) - else: - logger.info(f"Converting {submodule_name} to CoreML..") - coreml_model = ct.convert( - torchscript_module, - convert_to="mlprogram", - minimum_deployment_target=ct.target.macOS13, - inputs=sample_inputs, - outputs=[ - ct.TensorType(name=name, dtype=np.float32) for name in output_names - ], - skip_model_load=True, - ) - - del torchscript_module - gc.collect() - - return coreml_model - - -def get_out_path(submodule_name, model_name): - from folder_paths import get_folder_paths - - fname = f"{model_name}_{submodule_name}.mlpackage" - unet_path = get_folder_paths(submodule_name)[0] - out_path = os.path.join(unet_path, fname) - return out_path - - -def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape): - sample_unet_inputs = dict( - [ - ("sample", torch.rand(*sample_shape)), - ( - "timestep", - torch.tensor([DEFAULT_TRACE_TIMESTEP] * batch_size).to(torch.float32), - ), - ("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)), - ] - ) - return sample_unet_inputs - - -def lcm_inputs(sample_unet_inputs): - batch_size = sample_unet_inputs["sample"].shape[0] - return {"timestep_cond": torch.randn(batch_size, 256).to(torch.float32)} - - -def sdxl_inputs(sample_unet_inputs, ref_unet, model_version): - sample_shape = sample_unet_inputs["sample"].shape - batch_size = sample_shape[0] - h = sample_shape[2] * 8 - w = sample_shape[3] * 8 - original_size = (h, w) - crops_coords_top_left = (0, 0) - - is_refiner = model_version == ModelVersion.SDXL_REFINER - - if is_refiner: - aesthetic_score = (6.0,) - time_ids_list = list(original_size + crops_coords_top_left + aesthetic_score) - else: - target_size = (h, w) - time_ids_list = list(original_size + crops_coords_top_left + target_size) - - time_ids = torch.tensor(time_ids_list).repeat(batch_size, 1).to(torch.int64) - text_embeds_shape = (batch_size, get_sdxl_text_embeds_dim(ref_unet, len(time_ids_list))) - - return { - "time_ids": time_ids, - "text_embeds": torch.randn(*text_embeds_shape).to(torch.float32), - } - - -def get_sdxl_text_embeds_dim(ref_unet, time_ids_dim): - projection_dim = ref_unet.config.projection_class_embeddings_input_dim - time_embed_dim = ref_unet.config.addition_time_embed_dim - return projection_dim - (time_ids_dim * time_embed_dim) - - -def get_inputs_spec(inputs): - inputs_spec = {k: (v.shape, v.dtype) for k, v in inputs.items()} - return inputs_spec - - -def add_cnet_support(sample_shape, reference_unet): - additional_residuals_shapes = [] - - batch_size = sample_shape[0] - h, w = sample_shape[2:] - - # conv_in - out_h, out_w = conv2d_output_shape( - h, - w, - reference_unet.conv_in, - ) - additional_residuals_shapes.append( - (batch_size, reference_unet.conv_in.out_channels, out_h, out_w) - ) - - # down_blocks - for down_block in reference_unet.down_blocks: - additional_residuals_shapes += [ - (batch_size, resnet.out_channels, out_h, out_w) - for resnet in down_block.resnets - ] - if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None: - for downsampler in down_block.downsamplers: - out_h, out_w = conv2d_output_shape(out_h, out_w, downsampler.conv) - additional_residuals_shapes.append( - ( - batch_size, - down_block.downsamplers[-1].conv.out_channels, - out_h, - out_w, - ) - ) - - # mid_block - additional_residuals_shapes.append( - (batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w) - ) - - additional_inputs = {} - for i, shape in enumerate(additional_residuals_shapes): - sample_residual_input = torch.rand(*shape) - additional_inputs[f"additional_residual_{i}"] = sample_residual_input - - return additional_inputs - - -def convert_unet( - ref_unet, - model_version: ModelVersion, - unet_out_path: str, - batch_size: int = 1, - sample_size: tuple[int, int] = (64, 64), - controlnet_support: bool = False, - attention_implementation: str = ATTENTION_IMPLEMENTATIONS[0], - quantize_nbits: str = "none", -): - coreml_unet = get_unet(model_version, ref_unet, attention_implementation) - - sample_shape = ( - batch_size, # B - ref_unet.config.in_channels, # C - sample_size[0], # H - sample_size[1], # W - ) - - encoder_hidden_states_shape = get_encoder_hidden_states_shape(ref_unet, batch_size) - - sample_inputs = get_sample_input( - batch_size, encoder_hidden_states_shape, sample_shape - ) - - if model_version == ModelVersion.LCM: - sample_inputs |= lcm_inputs(sample_inputs) - - if model_version in {ModelVersion.SDXL, ModelVersion.SDXL_REFINER}: - sample_inputs |= sdxl_inputs(sample_inputs, ref_unet, model_version) - - if controlnet_support: - sample_inputs |= add_cnet_support(sample_shape, ref_unet) - - sample_inputs_spec = get_inputs_spec(sample_inputs) - - logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}") - logger.info("JIT tracing..") - traced_unet = torch.jit.trace( - coreml_unet, example_inputs=list(sample_inputs.values()) - ) - logger.info("Done.") - - coreml_sample_inputs = get_coreml_inputs(sample_inputs) - - coreml_unet = convert_to_coreml( - "unet", traced_unet, coreml_sample_inputs, ["noise_pred"], unet_out_path - ) - - del traced_unet - gc.collect() - - if quantize_nbits != "none": - # Opt-in k-means weight palettization. The default path - # (quantize_nbits="none") leaves the traced UNet untouched. - from coremltools.optimize.coreml import ( - OpPalettizerConfig, - OptimizationConfig, - palettize_weights, - ) - - nbits = int(quantize_nbits) - logger.info(f"Palettizing UNet weights to {nbits}-bit (kmeans)..") - t0 = time.time() - cfg = OptimizationConfig( - global_config=OpPalettizerConfig(mode="kmeans", nbits=nbits) - ) - coreml_unet = palettize_weights(coreml_unet, config=cfg) - logger.info(f"Palettization took {time.time() - t0:.1f}s") - - coreml_unet.save(unet_out_path) - logger.info(f"Saved unet into {unet_out_path}") - - -def convert( - ckpt_path: str, - model_version: ModelVersion, - unet_out_path: str, - batch_size: int = 1, - sample_size: tuple[int, int] = (64, 64), - controlnet_support: bool = False, - lora_weights: list[tuple[str | os.PathLike, float]] = None, - attn_impl: str = ATTENTION_IMPLEMENTATIONS[0], - config_path: str = None, - quantize_nbits: str = "none", -): - if os.path.exists(unet_out_path): - logger.info(f"Found existing model at {unet_out_path}! Skipping..") - return - - if attn_impl not in ATTENTION_IMPLEMENTATIONS: - raise ValueError( - f"Unsupported attention implementation {attn_impl!r}. " - f"Expected one of {ATTENTION_IMPLEMENTATIONS}." - ) - ref_unet = load_unet(ckpt_path, config_path) - - for i, lora_weight in enumerate(lora_weights or []): - lora_path, strength = lora_weight - adapter_name = f"lora_{i}" - ref_unet.load_lora_adapter(lora_path, adapter_name=adapter_name) - ref_unet.set_adapters([adapter_name], weights=[strength]) - ref_unet.fuse_lora() - - convert_unet( - ref_unet, - model_version, - unet_out_path, - batch_size, - sample_size, - controlnet_support, - attention_implementation=attn_impl, - quantize_nbits=quantize_nbits, - ) - - -def load_unet(ckpt_path, config_path): - return UNet2DConditionModel.from_single_file( - ckpt_path, - original_config=config_path, - ) diff --git a/coreml_suite/core/naming.py b/coreml_suite/core/naming.py deleted file mode 100644 index 71d186d..0000000 --- a/coreml_suite/core/naming.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Pure out_name composition for the Core ML UNet artifact. - -Extracted from CoreMLConverter.convert 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", -} - -# Palettization bits. "none" = no quantization (default; keeps the -# unquantized filename intact so existing workflows still resolve their -# cached .mlpackage). Numeric values append a `_q` 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 (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` - - Quantization: - - quantize_nbits "none" (default) appends nothing — existing - unquantized .mlpackages keep the old filename - - "4" / "6" / "8" appends `_q` 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])] diff --git a/coreml_suite/lcm/converter.py b/coreml_suite/lcm/converter.py index c2508ca..3214375 100644 --- a/coreml_suite/lcm/converter.py +++ b/coreml_suite/lcm/converter.py @@ -1,19 +1,36 @@ +"""LCM-specific conversion orchestration (comfy-side). + +E2 deduped the generic helpers (input building, Core ML export, residual-shape +calc) into ``coreml_diffusion.convert`` — this file now imports them instead of +carrying near-identical copies. What stays here is the genuinely LCM-specific +path: the hardcoded ``SimianLuo/LCM_Dreamshaper_v7`` download and the scheduler +that supplies the trace timestep. Consolidating that into the unified +``coreml_diffusion.convert(model_version=LCM, ...)`` path is a behavior change +deferred to E-LCM (it needs its own golden anchor). + +``get_scheduler`` keeps using ``comfy.model_management`` because it runs on the +comfy side; the conversion package itself stays comfy-free. +""" +import gc +import logging import os -import logging -import time -import gc -import numpy as np import torch from diffusers import UNet2DConditionModel, LCMScheduler from diffusers.loaders import LoraLoaderMixin -from coreml_suite.conversion.attention import apply_attention_implementation -from coreml_suite.conversion.shapes import conv2d_output_shape -from coreml_suite.conversion.unet import CoreMLUNetWrapper -from coreml_suite.model_version import ModelVersion - -import coremltools as ct +from coreml_diffusion.conversion.attention import apply_attention_implementation +from coreml_diffusion.conversion.unet import CoreMLUNetWrapper +from coreml_diffusion.convert import ( + add_cnet_support, + convert_to_coreml, + get_coreml_inputs, + get_encoder_hidden_states_shape, + get_inputs_spec, + get_sample_input, + lcm_inputs, +) +from coreml_diffusion import ModelVersion logging.basicConfig() logger = logging.getLogger(__name__) @@ -21,7 +38,6 @@ logger.setLevel(logging.DEBUG) MODEL_VERSION = "SimianLuo/LCM_Dreamshaper_v7" MODEL_NAME = MODEL_VERSION.split("/")[-1] + "_4k" -TEXT_TOKEN_SEQUENCE_LENGTH = 77 def get_unets(): @@ -40,16 +56,6 @@ def get_unets(): return cml_unet, ref_unet -def get_encoder_hidden_states_shape(unet_config, batch_size): - encoder_hidden_states_shape = ( - batch_size, - TEXT_TOKEN_SEQUENCE_LENGTH, - unet_config.cross_attention_dim, - ) - - return encoder_hidden_states_shape - - def get_scheduler(): from comfy.model_management import get_torch_device @@ -58,55 +64,6 @@ def get_scheduler(): return scheduler -def get_coreml_inputs(sample_inputs): - coreml_sample_unet_inputs = { - k: v.numpy().astype(np.float16) for k, v in sample_inputs.items() - } - return [ - ct.TensorType( - name=k, - shape=v.shape, - dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype, - ) - for k, v in coreml_sample_unet_inputs.items() - ] - - -def load_coreml_model(out_path): - logger.info(f"Loading model from {out_path}") - - start = time.time() - coreml_model = ct.models.MLModel(out_path) - logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds") - - return coreml_model - - -def convert_to_coreml( - submodule_name, torchscript_module, sample_inputs, output_names, out_path -): - if os.path.exists(out_path): - logger.info(f"Skipping export because {out_path} already exists") - coreml_model = load_coreml_model(out_path) - else: - logger.info(f"Converting {submodule_name} to CoreML..") - coreml_model = ct.convert( - torchscript_module, - convert_to="mlprogram", - minimum_deployment_target=ct.target.macOS13, - inputs=sample_inputs, - outputs=[ - ct.TensorType(name=name, dtype=np.float32) for name in output_names - ], - skip_model_load=True, - ) - - del torchscript_module - gc.collect() - - return coreml_model - - def get_out_path(submodule_name, model_name): from folder_paths import get_folder_paths @@ -116,77 +73,6 @@ def get_out_path(submodule_name, model_name): return out_path -def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape, scheduler): - sample_unet_inputs = dict( - [ - ("sample", torch.rand(*sample_shape)), - ( - "timestep", - torch.tensor([scheduler.timesteps[0].item()] * batch_size).to( - torch.float32 - ), - ), - ("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)), - ("timestep_cond", torch.randn(batch_size, 256).to(torch.float32)), - ] - ) - return sample_unet_inputs - - -def get_unet_inputs_spec(sample_unet_inputs): - sample_unet_inputs_spec = { - k: (v.shape, v.dtype) for k, v in sample_unet_inputs.items() - } - return sample_unet_inputs_spec - - -def add_cnet_support(sample_shape, reference_unet): - additional_residuals_shapes = [] - - batch_size = sample_shape[0] - h, w = sample_shape[2:] - - # conv_in - out_h, out_w = conv2d_output_shape( - h, - w, - reference_unet.conv_in, - ) - additional_residuals_shapes.append( - (batch_size, reference_unet.conv_in.out_channels, out_h, out_w) - ) - - # down_blocks - for down_block in reference_unet.down_blocks: - additional_residuals_shapes += [ - (batch_size, resnet.out_channels, out_h, out_w) - for resnet in down_block.resnets - ] - if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None: - for downsampler in down_block.downsamplers: - out_h, out_w = conv2d_output_shape(out_h, out_w, downsampler.conv) - additional_residuals_shapes.append( - ( - batch_size, - down_block.downsamplers[-1].conv.out_channels, - out_h, - out_w, - ) - ) - - # mid_block - additional_residuals_shapes.append( - (batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w) - ) - - additional_inputs = {} - for i, shape in enumerate(additional_residuals_shapes): - sample_residual_input = torch.rand(*shape) - additional_inputs[f"additional_residual_{i}"] = sample_residual_input - - return additional_inputs - - def convert( out_path: str, batch_size: int = 1, @@ -209,20 +95,19 @@ def convert( sample_size[1], # W ) - encoder_hidden_states_shape = get_encoder_hidden_states_shape( - ref_unet.config, batch_size - ) + encoder_hidden_states_shape = get_encoder_hidden_states_shape(ref_unet, batch_size) scheduler = get_scheduler() sample_inputs = get_sample_input( - batch_size, encoder_hidden_states_shape, sample_shape, scheduler + batch_size, encoder_hidden_states_shape, sample_shape, scheduler=scheduler ) + sample_inputs |= lcm_inputs(sample_inputs) if controlnet_support: sample_inputs |= add_cnet_support(sample_shape, ref_unet) - sample_inputs_spec = get_unet_inputs_spec(sample_inputs) + sample_inputs_spec = get_inputs_spec(sample_inputs) logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}") logger.info("JIT tracing..") diff --git a/coreml_suite/model_version.py b/coreml_suite/model_version.py deleted file mode 100644 index d4d2fa1..0000000 --- a/coreml_suite/model_version.py +++ /dev/null @@ -1,8 +0,0 @@ -from enum import Enum - - -class ModelVersion(Enum): - SD15 = "sd15" - SDXL = "sdxl" - SDXL_REFINER = "sdxl_refiner" - LCM = "lcm" diff --git a/coreml_suite/nodes.py b/coreml_suite/nodes.py index f6d9301..75009d7 100644 --- a/coreml_suite/nodes.py +++ b/coreml_suite/nodes.py @@ -4,16 +4,10 @@ from coremltools import ComputeUnit import folder_paths from coreml_suite import COREML_NODE -from coreml_suite.attention import ATTENTION_IMPLEMENTATIONS from coreml_suite.coreml_model import CoreMLModel -from coreml_suite.core.naming import ( - QUANT_NBITS_VALUES, - compose_out_name, - lora_names_from_params, -) from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm from coreml_suite.logger import logger -from coreml_suite.model_version import ModelVersion +from coreml_diffusion import ModelVersion from nodes import KSampler, LoraLoader, KSamplerAdvanced from coreml_suite.models import ( @@ -24,6 +18,26 @@ from coreml_suite.models import ( ) +def _discover(fn_name, fallback): + """Populate a converter dropdown from coreml_diffusion's discovery API. + + Fails soft: if the package is missing, too old to expose ``fn_name``, or + errors, the node still registers with the fallback list instead of vanishing + from the menu. Evaluated on every INPUT_TYPES call, so installing a newer + coreml_diffusion surfaces new conversion types with no Suite change. + """ + try: + import coreml_diffusion + + return getattr(coreml_diffusion, fn_name)() + except Exception as exc: # missing/old package, import error, etc. + logger.warning( + f"coreml_diffusion.{fn_name} unavailable ({exc}); " + f"using fallback {fallback}" + ) + return fallback + + class CoreMLSampler(COREML_NODE, KSampler): @classmethod def INPUT_TYPES(s): @@ -219,17 +233,15 @@ class CoreMLConverter(COREML_NODE): return { "required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"),), - "model_version": ( - [ - ModelVersion.SD15.name, - ModelVersion.SDXL.name, - ], - ), + "model_version": (_discover("list_model_versions", ["SD15", "SDXL"]),), "height": ("INT", {"default": 512, "min": 8, "step": 8}), "width": ("INT", {"default": 512, "min": 8, "step": 8}), "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}), "attention_implementation": ( - list(ATTENTION_IMPLEMENTATIONS), + _discover( + "list_attention_impls", + ["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"], + ), ), "compute_unit": ( [ @@ -247,7 +259,10 @@ class CoreMLConverter(COREML_NODE): # omits any `required` input. When omitted it defaults to # "none", identical to unquantized behavior and filename, so # existing cached .mlpackages still resolve. - "quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}), + "quantize_nbits": ( + _discover("list_quant_modes", ["none", "8", "6", "4"]), + {"default": "none"}, + ), "lora_params": ("LORA_PARAMS",), }, } @@ -293,14 +308,16 @@ class CoreMLConverter(COREML_NODE): h = height w = width sample_size = (h // 8, w // 8) - out_name = compose_out_name( + import coreml_diffusion + + out_name = coreml_diffusion.compose_out_name( ckpt_name=ckpt_name, batch_size=batch_size, width=w, height=h, controlnet_support=controlnet_support, attention_implementation=attention_implementation, - lora_names=lora_names_from_params(lora_params), + lora_names=coreml_diffusion.lora_names_from_params(lora_params), quantize_nbits=quantize_nbits, ) @@ -315,9 +332,10 @@ class CoreMLConverter(COREML_NODE): for lora_param in lora_params: logger.info(f" {lora_param[0]} - strength: {lora_param[1]}") - from coreml_suite import converter - - unet_out_path = converter.get_out_path("unet", f"{out_name}") + # Resolve the ComfyUI models/unet path here (a node concern); the package + # takes the output path as an injected argument. + unet_path = folder_paths.get_folder_paths("unet")[0] + unet_out_path = os.path.join(unet_path, f"{out_name}_unet.mlpackage") ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) config_filename = ckpt_name.split(".")[0] + ".yaml" @@ -325,10 +343,10 @@ class CoreMLConverter(COREML_NODE): if config_path: logger.info(f"Using config file {config_path}") - converter.convert( - ckpt_path=ckpt_path, - model_version=model_version, - unet_out_path=unet_out_path, + coreml_diffusion.convert( + ckpt_path, + model_version, + unet_out_path, sample_size=sample_size, batch_size=batch_size, controlnet_support=controlnet_support, diff --git a/pyproject.toml b/pyproject.toml index 94b39e6..e58f62a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,23 +1,28 @@ [project] name = "comfyui-coremlsuite" description = "This extension contains a set of custom nodes for ComfyUI that allow you to use Core ML models in your ComfyUI workflows." -version = "2.0.2" +version = "2.1.0" license = "MIT" requires-python = ">=3.12,<3.13" packages = [{ include = "coreml_suite" }] dependencies = [ + "coreml-diffusion", "torch>=2.7,<2.8", "coremltools>=9,<10", "numpy>=2,<3", + # diffusers is still imported directly by the comfy-side LCM converter + # (coreml_suite/lcm/converter.py) until E-LCM folds it into the package. "diffusers>=0.30", - "peft>=0.13", - "omegaconf>=2.3", - "transformers>=4.44", ] [project.urls] Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite" +[tool.uv.sources] +# Pinned to a published tag so CI can resolve it. Bump the tag when depending on +# a newer coreml-diffusion release. +coreml-diffusion = { git = "https://github.com/aszc-dev/coreml-diffusion", tag = "v0.1.0" } + [tool.comfy] PublisherId = "aszc-dev" DisplayName = "ComfyUI-CoreMLSuite" diff --git a/requirements.txt b/requirements.txt index 2824535..1a126bc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,5 @@ +coreml-diffusion @ git+https://github.com/aszc-dev/coreml-diffusion@v0.1.0 torch>=2.7,<2.8 coremltools>=9,<10 numpy>=2,<3 diffusers>=0.30 -peft>=0.13 -omegaconf>=2.3 -transformers>=4.44 diff --git a/seam.md b/seam.md new file mode 100644 index 0000000..5d2d9e8 --- /dev/null +++ b/seam.md @@ -0,0 +1,208 @@ +# Conversion Extraction — Seam Inventory (`docs/extraction/seam.md`) + +> **Gate E0 deliverable.** Symbol-by-symbol cut line between the future `coreml_diffusion` +> package (CONVERSION) and what stays in `coreml_suite` (the ComfyUI side). +> +> **Confidence legend:** +> - ✅ **verified** — read directly from the current source in this repo. +> - 🔍 **confirm** — inferred / partially seen; Claude Code must `grep`-verify before acting. +> +> **Cut rule:** a symbol goes to `coreml_diffusion` iff it participates in producing the `.mlpackage` +> artifact AND can be made free of `comfy` / `folder_paths` / `comfy_extras`. The runtime +> *loader* that **runs** a compiled model stays in the suite. + +--- + +## 1. File-level map + +| File | Side | Status | Note | +|---|---|---|---| +| `coreml_suite/model_version.py` | **coreml_diffusion** | ✅ | Already `Enum`-only, zero comfy. Becomes pkg source of truth. | +| `coreml_suite/attention.py` | **coreml_diffusion** | ✅ | `ATTENTION_IMPLEMENTATIONS` tuple; pure constant. | +| `coreml_suite/core/naming.py` | **coreml_diffusion** | ✅ | `compose_out_name` = cache-key contract. Move (not copy). | +| `coreml_suite/converter.py` | **coreml_diffusion** (mostly) | ✅ | Main conversion. One symbol stays-adjacent: `get_out_path` (folder_paths) is replaced by injected `out_path`. | +| `coreml_suite/conversion/attention.py` | **coreml_diffusion** | ✅ | `apply_attention_implementation`. Imports `logging`,`torch` only — no comfy. | +| `coreml_suite/conversion/shapes.py` | **coreml_diffusion** | ✅ | `conv2d_output_shape`. Pure math, no imports. | +| `coreml_suite/conversion/trace.py` | **coreml_diffusion** | ✅ | Imports `types.MethodType`, `diffusers...Transformer2DModel` only — torch/diffusers. | +| `coreml_suite/conversion/unet.py` | **coreml_diffusion** | ✅ | `CoreMLUNetWrapper`. Imports `torch` only — no comfy. | +| `coreml_suite/lcm/converter.py` | **coreml_diffusion** (after dedup) | ✅ | Dup helpers deleted; `MODEL_VERSION` HF-hardcode (L22) → E-LCM. `folder_paths` (L111) + `comfy.model_management` (L54) confirmed present → CUT. | +| `coreml_suite/lcm/unet.py` | **coreml_diffusion** | ✅ | `UNet2DConditionModelLCM(UNet2DConditionModel)`. diffusers-only, no comfy. | +| `coreml_suite/config.py` | **STAYS** | ✅ | Imports `comfy.supported_models_base`/`latent_formats`/`model_detection`. **Inference-side** (`get_model_config`), NOT conversion. | +| `coreml_suite/coreml_model.py` | **STAYS** | ✅ | `CoreMLModel` = runtime loader (runs `.mlpackage`). Desktop/Python inference; not used on iOS. | +| `coreml_suite/nodes.py` | **STAYS** | ✅ | Nodes; will call `coreml_diffusion` + own `folder_paths` path resolution + discovery dropdowns. | +| `coreml_suite/lcm/nodes.py` | **STAYS** | ✅ | `COREML_CONVERT_LCM` node. | +| `coreml_suite/models.py` | **STAYS** | ✅ | Inference: `add_sdxl_model_options`, `is_sdxl`, `get_model_patcher`, `get_latent_image`. | +| `coreml_suite/latents.py` | **STAYS** | ✅ | Inference chunking (MODERNIZATION Phase 3 target, not this spec). | +| `coreml_suite/controlnet.py` | **STAYS** | ✅ | Inference-side controlnet. Distinct from converter `add_cnet_support`. | +| `coreml_suite/lcm/utils.py` | **STAYS** | ✅ | `add_lcm_model_options`, `lcm_patch`, `is_lcm`; imports `comfy_extras`. Inference. | +| `coreml_suite/logger.py` | **both / copy** | ✅ | Trivial. Package gets its own logger; suite keeps its. | + +--- + +## 2. Symbol-level: `coreml_suite/converter.py` (main conversion) + +| Symbol | Side | Status | Cut action | +|---|---|---|---| +| `DEFAULT_TRACE_TIMESTEP`, `TEXT_TOKEN_SEQUENCE_LENGTH` | coreml_diffusion | ✅ | Move as-is (module constants). | +| `get_unet(model_version, ref_unet, attention_implementation)` | coreml_diffusion | ✅ | Move. Uses `conversion.{trace,attention,unet}`. No comfy. | +| `get_encoder_hidden_states_shape(ref_unet, batch_size)` | coreml_diffusion | ✅ | Move. Reads `ref_unet.config.cross_attention_dim`. Pure. | +| `get_coreml_inputs(sample_inputs)` | coreml_diffusion | ✅ | Move. `ct.TensorType` build. | +| `load_coreml_model(out_path)` | coreml_diffusion | ✅ | Move. `ct.models.MLModel(out_path)`. (Dedup target vs LCM copy.) | +| `convert_to_coreml(submodule, ts_module, inputs, names, out_path)` | coreml_diffusion | ✅ | Move. `ct.convert(...)`. (Dedup target vs LCM copy.) | +| `get_sample_input(batch, ehs_shape, sample_shape)` | coreml_diffusion | ✅ | Move. **Merge** with LCM variant (LCM passes extra `scheduler` → optional param). | +| `lcm_inputs(sample_unet_inputs)` | coreml_diffusion | ✅ | Move. Adds `timestep_cond`. | +| `sdxl_inputs(sample_unet_inputs, ref_unet, model_version)` | coreml_diffusion | ✅ | Move. `time_ids`/`text_embeds`/`add_embeds`. | +| `add_cnet_support(sample_shape, ref_unet)` | coreml_diffusion | ✅ | Move. Builds `additional_residual_*` inputs from unet block channels. | +| `convert_unet(ref_unet, model_version, unet_out_path, ...)` | coreml_diffusion | ✅ | Move. Orchestrates trace→convert→**quant (palettize)**→save. Quant travels here (E6). | +| `convert(ckpt_path, model_version, unet_out_path, ...)` | coreml_diffusion | ✅ | Move. **Make kw-only past `ckpt_path,model_version,out_path`** (contract). Validates `attn_impl`. | +| `load_unet(ckpt_path, config_path)` | coreml_diffusion | ✅ | Move. `UNet2DConditionModel.from_single_file`. | +| `get_out_path(submodule_name, model_name)` | **STAYS (node)** | ✅ | Uses `folder_paths.get_folder_paths`. **Delete from converter; node resolves path and passes `out_path` in.** | + +**Apple `python_coreml_stable_diffusion` footprint on this path:** ✅ **none.** Verified by grep: +zero imports in `converter.py` / `conversion/*`. Main path uses `diffusers` + +local `CoreMLUNetWrapper`. (And the runtime `CoreMLModel` is now a local coremltools wrapper too — +see §6 stale-spec note.) + +--- + +## 3. Symbol-level: `coreml_suite/lcm/converter.py` (LCM — dedup + defer) + +| Symbol | Side | Status | Cut action | +|---|---|---|---| +| `load_coreml_model` (LCM copy) | DELETE | ✅ | Duplicate of main. Remove; use `coreml_diffusion.load_coreml_model`. | +| `convert_to_coreml` (LCM copy) | DELETE | ✅ | Duplicate of main. Remove. | +| `get_out_path` (LCM copy, folder_paths) | DELETE | ✅ | Duplicate + comfy. Remove; node injects `out_path`. | +| `get_sample_input(..., scheduler)` (LCM copy) | MERGE → coreml_diffusion | ✅ | Fold `scheduler` into shared `get_sample_input` as optional param. | +| `MODEL_NAME` (= LCM_Dreamshaper) | **E-LCM** | ✅ | HF hardcode. Removing it is the behavior change → E-LCM, not E2. | +| `convert(out_path, sample_size, batch_size, controlnet_support)` (LCM, L190) | coreml_diffusion (via unified) | ✅ | Route through `coreml_diffusion.convert(model_version=LCM, ...)` in E-LCM. | +| `from comfy.model_management import get_torch_device` (L54, in `get_scheduler`) | **CUT** | ✅ | Confirmed present. Inject `device`. | +| module-global attention set at import | n/a | ✅ | **No module global.** Attention already per-call: `get_unets` (L36) calls `apply_attention_implementation(ref_unet, "SPLIT_EINSUM")`. No `ATTENTION_IMPLEMENTATION_IN_EFFECT` anywhere in repo. (Note: LCM hardcodes `"SPLIT_EINSUM"` — pass `attn_impl` through in dedup.) | + +--- + +## 4. Symbol-level: `coreml_suite/core/naming.py` → `coreml_diffusion/naming.py` + +| Symbol | Side | Status | Cut action | +|---|---|---|---| +| `compose_out_name(...)` | coreml_diffusion | ✅ | **Move** (cache-key contract). Node imports from pkg. | +| `lora_names_from_params(...)` | coreml_diffusion | ✅ | Move. | +| `ATTN_SUFFIX` dict | coreml_diffusion | ✅ | Move. | +| `QUANT_NBITS_VALUES` | coreml_diffusion | ✅ | Move; backs `list_quant_modes()`. | +| `tests/unit/test_characterization_out_name.py` | re-point | ✅ | Change import to `coreml_diffusion.naming`. Assertions/values **unchanged**. | + +--- + +## 5. Discovery API + status registry (new in `coreml_diffusion/__init__.py`) + +```python +from enum import Enum + +class Status(Enum): + VERIFIED = "verified" # has a golden anchor + passing [M2-ANE] check + EXPERIMENTAL = "experimental" # convertible, not yet anchored/verified + +# Single source of truth. Suite gates on this, NOT on a hardcoded node list. +# KEY by ModelVersion enum MEMBER (not a bare string) so list_model_versions can +# emit .name — see the .name decision below. Keying by the lowercase .value string +# (as an earlier draft of this block did) returns ["sd15",...], which the node then +# reverses via ModelVersion[...] → KeyError. Do NOT key by .value. +_MODEL_STATUS = { + ModelVersion.SD15: Status.VERIFIED, + ModelVersion.SDXL: Status.VERIFIED, + ModelVersion.SDXL_REFINER: Status.EXPERIMENTAL, # → VERIFIED after a refiner golden anchor + ModelVersion.LCM: Status.EXPERIMENTAL, # → VERIFIED after E-LCM golden anchor +} + +def list_model_versions(include_experimental: bool = False) -> list[str]: + return [v.name for v, s in _MODEL_STATUS.items() # .name → "SD15","SDXL" (see decision) + if s is Status.VERIFIED or (include_experimental and s is Status.EXPERIMENTAL)] + +def list_attention_impls() -> list[str]: # from attention.ATTENTION_IMPLEMENTATIONS + ... +def list_quant_modes() -> list[str]: # from naming.QUANT_NBITS_VALUES + ... + +CONTRACT_VERSION = "1.0" +# Additive-only: adding an id or promoting EXPERIMENTAL→VERIFIED = minor bump (Suite unaffected). +# Removing/renaming an id, or demoting VERIFIED→EXPERIMENTAL = MAJOR bump + migration note. +``` + +**Decision check (`.name` vs `.value`): RESOLVED → `.name`.** ✅ +Verified in current source: +- Node renders `ModelVersion.SD15.name` / `ModelVersion.SDXL.name` → `"SD15"`, `"SDXL"` + (`nodes.py:224-225`). +- Node reverses the dropdown string with `model_version = ModelVersion[model_version]` + (`nodes.py:286`) — i.e. **lookup by NAME**. Feeding it a `.value` (`"sd15"`) raises `KeyError`. +- Enum values are lowercase (`model_version.py`: `SD15="sd15"`, `SDXL="sdxl"`, + `SDXL_REFINER="sdxl_refiner"`, `LCM="lcm"`). +- `compose_out_name` does NOT consume the model_version string (grep of `core/naming.py` empty) — + no coupling there, so no constraint from that side. + +**Decision:** `list_model_versions()` returns `.name` (uppercase). Saved workflows store `"SD15"`, +node already validates them via `ModelVersion[...]`. The `_MODEL_STATUS` block above was corrected +to key by enum member and emit `.name`. **The earlier `v.value` form was a latent bug.** + +--- + +## 6. `python_coreml_stable_diffusion` split (Gate E0 line to fill by grep) + +| Use | Side | Status | +|---|---|---| +| `coreml_model.CoreMLModel` (runs compiled model) | **STAYS** (suite runtime) | ✅ — **local class**, not Apple's | +| `unet.UNet2DConditionModel*` internals | **gone** — `converter.py:319` uses `diffusers.UNet2DConditionModel.from_single_file` | ✅ | +| `AttentionImplementations` enum | gone — local `apply_attention_implementation` + `attention.py` tuple | ✅ | +| `calculate_conv2d_output_shape` | gone — replaced by `conversion/shapes.conv2d_output_shape` | ✅ | + +> ### ⚠️ SPEC IS STALE: `ml-stable-diffusion` is already fully removed +> Commit #58 ("replace apple/ml-stable-diffusion with native diffusers conversion") already did +> the de-Apple work. Verified now: +> - **Zero** `python_coreml_stable_diffusion` runtime imports anywhere in `coreml_suite` (only a +> docstring mention at `core/__init__.py:4`). +> - `coreml_suite/coreml_model.py:8` `CoreMLModel` is a **local** wrapper over +> `coremltools.models.MLModel` (`coreml_model.py:22`) — it does **not** import Apple's class. +> - `ml-stable-diffusion` / `python_coreml_stable_diffusion` appears in **neither** `pyproject.toml` +> **nor** `requirements.txt`. It is not a dependency at all. +> +> **Consequences for the spec (correct these in CONVERTER_EXTRACTION_SPEC.md):** +> - §0.3 premise ("runtime loader = `python_coreml_stable_diffusion.coreml_model.CoreMLModel`, +> stays in suite") is **wrong**: the loader is already the local `coreml_model.CoreMLModel`. The +> "stays in suite" conclusion still holds; the identity does not. +> - **Gate E0 item "ml-stable-diffusion pinned SHA — BLOCKER if unpinned" is MOOT** — there is no +> such dep to pin. Mark it N/A, not BLOCKER. +> - **E4/E5 dependency lists must drop `git+...ml-stable-diffusion@`.** Package runtime deps +> are: `coremltools`, `diffusers`, `peft` (LoRA), `omegaconf` (config), `numpy`, `torch`. Confirm +> `peft`/`omegaconf` actually used before listing (grep at E4). +> - The "keep `python_coreml_stable_diffusion` as a suite dep for the loader" instruction in E5 is +> **void** — coremltools backs the loader. + +--- + +## 7. Pre-flight checklist before E1 (run these greps) + +``` +grep -rn "import comfy" coreml_suite/conversion coreml_suite/converter.py coreml_suite/lcm/converter.py coreml_suite/lcm/unet.py +grep -rn "folder_paths" coreml_suite/converter.py coreml_suite/lcm/converter.py +grep -rn "model_management" coreml_suite/lcm +grep -rn "python_coreml_stable_diffusion" coreml_suite +grep -rn "ATTENTION_IMPLEMENTATION_IN_EFFECT" coreml_suite +grep -rn "SimianLuo\|LCM_Dreamshaper" coreml_suite/lcm +``` +Every 🔍 above resolves to ✅ or a correction once these run. Do not start moving code (E2) +with any 🔍 unresolved on the CONVERSION side. + +**STATUS (run 2026-05-26): all 🔍 resolved.** Summary of what the greps found: +- `conversion/*`, `lcm/unet.py`: comfy-free (torch/diffusers only). ✅ +- `converter.py`: only comfy reach-in is `folder_paths` in `get_out_path` (L91-94) → inject `out_path`. +- `lcm/converter.py`: `folder_paths` (L111-114) + `comfy.model_management.get_torch_device` (L54) + → cut both. Dup helpers (`load_coreml_model`,`convert_to_coreml`,`get_out_path`,`get_sample_input`) + confirmed → dedup E2. `MODEL_VERSION="SimianLuo/LCM_Dreamshaper_v7"` (L22) → E-LCM. +- No attention module-global anywhere (`ATTENTION_IMPLEMENTATION_IN_EFFECT` absent); already per-call. + LCM hardcodes `"SPLIT_EINSUM"` in `get_unets` — thread `attn_impl` through during dedup. +- `.name` vs `.value`: **decided `.name`** (node reverses via `ModelVersion[...]`). §5 corrected. +- `ml-stable-diffusion`: **already gone** (#58). §6 stale-spec note added — fix the spec's E0/E4/E5 + dep + pinning items. + +Two grep blind-spots to note (the checklist above doesn't cover them, but cheap to add): the +`folder_paths` grep only scans the two converter files — also grep `coreml_suite/lcm/utils.py` +(it imports `comfy.model_management` at L3, but it's inference/STAYS, so fine) and confirm no other +`conversion/` file grew a comfy import since. diff --git a/tests/smoke/__init__.py b/tests/smoke/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/smoke/test_split_einsum_attention.py b/tests/smoke/test_split_einsum_attention.py deleted file mode 100644 index 44c89d6..0000000 --- a/tests/smoke/test_split_einsum_attention.py +++ /dev/null @@ -1,41 +0,0 @@ -import platform - -import pytest -import torch -from diffusers.models.attention_processor import Attention, AttnProcessor - -from coreml_suite.conversion.attention import ( - SplitEinsumAttnProcessor, - SplitEinsumV2AttnProcessor, -) - - -pytestmark = pytest.mark.skipif( - platform.system() != "Darwin" or platform.machine() != "arm64", - reason="Tier 1 requires macOS on Apple Silicon", -) - - -@pytest.mark.parametrize( - "processor", - [ - SplitEinsumAttnProcessor(), - SplitEinsumV2AttnProcessor(), - ], -) -def test_split_einsum_processor_matches_diffusers_attention(processor): - torch.manual_seed(0) - reference = Attention(query_dim=32, heads=4, dim_head=8, dropout=0.0) - reference.set_processor(AttnProcessor()) - - candidate = Attention(query_dim=32, heads=4, dim_head=8, dropout=0.0) - candidate.load_state_dict(reference.state_dict()) - candidate.set_processor(processor) - - hidden_states = torch.randn(2, 17, 32) - encoder_hidden_states = torch.randn(2, 11, 32) - - expected = reference(hidden_states, encoder_hidden_states=encoder_hidden_states) - actual = candidate(hidden_states, encoder_hidden_states=encoder_hidden_states) - - assert torch.allclose(actual, expected, atol=1e-5) diff --git a/tests/smoke/test_synthetic_unet.py b/tests/smoke/test_synthetic_unet.py deleted file mode 100644 index 1f14f84..0000000 --- a/tests/smoke/test_synthetic_unet.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Tier 1 smoke: convert a synthetic micro-UNet through coremltools and load -it back with CoreMLSuite's runtime CoreMLModel wrapper. - -Purpose: catch API breakage in coremltools *without* needing a real SD -checkpoint, the ANE, or a converted .mlmodelc on disk. -Runs in minutes on a hosted macOS-ARM runner (no Apple internal stuff). - -What it asserts: - - coremltools.convert still accepts the call shape we use today - - the resulting .mlpackage round-trips through CoreMLSuite's CoreMLModel - - expected_inputs exposes the input names/shapes we declared - - calling the model returns the named output (`noise_pred`) - -Auto-skips on non-Apple-Silicon hosts so Tier 0 CI on Linux ignores it. -""" -import platform -import shutil -from types import SimpleNamespace - -import numpy as np -import pytest -import torch -import torch.nn as nn - -from coreml_suite.conversion.unet import CoreMLUNetWrapper - - -pytestmark = pytest.mark.skipif( - platform.system() != "Darwin" or platform.machine() != "arm64", - reason="Tier 1 requires macOS on Apple Silicon", -) - - -# Tiny shapes — large enough to exercise conv2d + linear + addition kernels in -# coremltools, small enough that conversion finishes in seconds on CPU. -SAMPLE_SHAPE = (1, 4, 8, 8) -TIMESTEP_SHAPE = (1,) -ENCODER_SHAPE = (1, 4, 64) # native diffusers encoder_hidden_states (batch, tokens, hidden) -OUT_NAME = "noise_pred" - - -class TinyUNet(nn.Module): - """Minimal UNet-shaped graph: conv -> add(time+context) -> conv. - - Not a real diffusion model. Just enough op variety to exercise the - PyTorch -> MIL frontend in coremltools and confirm we can still wire - the inputs/outputs the way CoreMLSuite's runtime expects. - """ - - def __init__(self): - super().__init__() - self.conv_in = nn.Conv2d(4, 8, kernel_size=3, padding=1) - self.conv_out = nn.Conv2d(8, 4, kernel_size=3, padding=1) - self.time_proj = nn.Linear(1, 8) - self.text_proj = nn.Linear(64, 8) - - def forward( - self, - sample, - timestep, - encoder_hidden_states, - timestep_cond=None, - added_cond_kwargs=None, - down_block_additional_residuals=None, - mid_block_additional_residual=None, - return_dict=True, - ): - h = self.conv_in(sample) - t_emb = self.time_proj(timestep.unsqueeze(-1)).view(1, 8, 1, 1) - c_emb = self.text_proj(encoder_hidden_states.mean(1)).view(1, 8, 1, 1) - h = h + t_emb + c_emb - return (self.conv_out(h),) - - -@pytest.fixture(scope="module") -def tiny_mlpackage(tmp_path_factory): - """Convert TinyUNet once per test session and reuse the .mlpackage.""" - import coremltools as ct - - torch.manual_seed(0) - model = CoreMLUNetWrapper( - TinyUNet().eval(), - SimpleNamespace(name="SD15"), - ) - example = ( - torch.randn(*SAMPLE_SHAPE), - torch.randn(*TIMESTEP_SHAPE), - torch.randn(*ENCODER_SHAPE), - ) - traced = torch.jit.trace(model, example) - - mlmodel = ct.convert( - traced, - inputs=[ - ct.TensorType(name="sample", shape=SAMPLE_SHAPE, dtype=np.float16), - ct.TensorType(name="timestep", shape=TIMESTEP_SHAPE, dtype=np.float16), - ct.TensorType(name="encoder_hidden_states", shape=ENCODER_SHAPE, dtype=np.float16), - ], - outputs=[ct.TensorType(name=OUT_NAME, dtype=np.float16)], - compute_units=ct.ComputeUnit.CPU_ONLY, - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=ct.target.macOS13, - ) - - out_dir = tmp_path_factory.mktemp("tiny_unet") - pkg_path = out_dir / "tiny.mlpackage" - mlmodel.save(str(pkg_path)) - yield pkg_path - shutil.rmtree(out_dir, ignore_errors=True) - - -def test_coremltools_convert_round_trips_via_coreml_model(tiny_mlpackage): - from coreml_suite.coreml_model import CoreMLModel - - model = CoreMLModel(str(tiny_mlpackage), "CPU_ONLY") - - # expected_inputs is the contract our wrappers depend on. Lock the shape - # of the dict + a sample entry. - expected = dict(model.expected_inputs) - assert set(expected.keys()) == {"sample", "timestep", "encoder_hidden_states"} - assert tuple(expected["sample"]["shape"]) == SAMPLE_SHAPE - assert tuple(expected["timestep"]["shape"]) == TIMESTEP_SHAPE - assert tuple(expected["encoder_hidden_states"]["shape"]) == ENCODER_SHAPE - - # Forward pass: drive the model the way CoreMLModelWrapper does. - rng = np.random.default_rng(0) - inputs = { - "sample": rng.standard_normal(SAMPLE_SHAPE).astype(np.float16), - "timestep": rng.standard_normal(TIMESTEP_SHAPE).astype(np.float16), - "encoder_hidden_states": rng.standard_normal(ENCODER_SHAPE).astype(np.float16), - } - out = model(**inputs) - assert isinstance(out, dict), f"unexpected output type: {type(out)}" - assert OUT_NAME in out, f"missing output {OUT_NAME!r}; got {sorted(out)}" - assert out[OUT_NAME].shape == SAMPLE_SHAPE, ( - f"output shape drift: got {out[OUT_NAME].shape}, expected {SAMPLE_SHAPE}" - ) diff --git a/tests/unit/test_characterization_out_name.py b/tests/unit/test_characterization_out_name.py deleted file mode 100644 index a246a50..0000000 --- a/tests/unit/test_characterization_out_name.py +++ /dev/null @@ -1,197 +0,0 @@ -"""Characterization tests for the .mlpackage filename composition. - -The filename composition is the pure -coreml_suite.core.naming.compose_out_name function. CoreMLConverter.convert -calls it; testing the pure function avoids monkey-patching heavy converter -internals just to capture the string. -""" -import pytest - -from coreml_suite.core.naming import compose_out_name, lora_names_from_params - - -# ---------- attention suffixes ---------------------------------------------- - - -@pytest.mark.parametrize( - "attn_name,suffix", - [ - ("SPLIT_EINSUM", "se"), - ("SPLIT_EINSUM_V2", "se2"), - ("ORIGINAL", "orig"), - ], -) -def test_attention_suffix(attn_name, suffix): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation=attn_name, - ) - assert out == f"dreamshaper_8_1x512x512_{suffix}" - - -# ---------- batch / size ---------------------------------------------------- - - -def test_includes_batch_and_size(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=4, width=768, height=1024, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - ) - assert out == "dreamshaper_8_4x768x1024_se" - - -# ---------- ControlNet ------------------------------------------------------ - - -def test_appends_cn_suffix_when_controlnet_support_true(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=True, - attention_implementation="SPLIT_EINSUM", - ) - assert out == "dreamshaper_8_1x512x512_cn_se" - - -# ---------- ckpt name massage ----------------------------------------------- - - -def test_drops_extension_at_first_period(): - out = compose_out_name( - ckpt_name="my.checkpoint.v2.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - ) - assert out == "my_1x512x512_se" - - -def test_replaces_spaces_with_underscores(): - out = compose_out_name( - ckpt_name="dream shaper 8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - ) - assert out == "dream_shaper_8_1x512x512_se" - - -# ---------- LoRA suffixes --------------------------------------------------- - - -def test_single_lora(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - lora_names=["epi_noiseoffset.safetensors"], - ) - assert out == "dreamshaper_8_epi_noiseoffset_1x512x512_se" - - -def test_multiple_loras_sorted(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - lora_names=["zoom.safetensors", "alpha.safetensors", "moody.safetensors"], - ) - assert out == "dreamshaper_8_alpha_moody_zoom_1x512x512_se" - - -def test_lora_plus_controlnet(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=True, - attention_implementation="SPLIT_EINSUM", - lora_names=["a.safetensors"], - ) - assert out == "dreamshaper_8_a_1x512x512_cn_se" - - -# ---------- sdxl combinations ----------------------------------------------- - - -def test_sdxl_1024_original_gpu(): - out = compose_out_name( - ckpt_name="sd_xl_base_1.0.safetensors", - batch_size=1, width=1024, height=1024, - controlnet_support=False, - attention_implementation="ORIGINAL", - ) - assert out == "sd_xl_base_1_1x1024x1024_orig" - - -# ---------- lora_names_from_params helper ---------------------------------- - - -def test_lora_names_from_params_sorts_by_name(): - names = lora_names_from_params([ - ("zebra.safetensors", 1.0), - ("apple.safetensors", 0.5), - ("mango.safetensors", 0.7), - ]) - assert names == ["apple.safetensors", "mango.safetensors", "zebra.safetensors"] - - -def test_lora_names_from_params_empty_list(): - assert lora_names_from_params([]) == [] - - -# ---------- quantize_nbits suffix ------------------------------------------ - - -def test_quantize_nbits_none_appends_nothing(): - """'none' is the default and must keep the unquantized filename so - existing cached .mlpackages still resolve.""" - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - quantize_nbits="none", - ) - assert out == "dreamshaper_8_1x512x512_se" - - -@pytest.mark.parametrize("nbits,suffix", [("4", "_q4"), ("6", "_q6"), ("8", "_q8")]) -def test_quantize_nbits_appends_q_suffix(nbits, suffix): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - quantize_nbits=nbits, - ) - assert out == f"dreamshaper_8_1x512x512_se{suffix}" - - -def test_quantize_nbits_with_controlnet_and_lora(): - out = compose_out_name( - ckpt_name="dreamshaper_8.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=True, - attention_implementation="SPLIT_EINSUM", - lora_names=["a.safetensors"], - quantize_nbits="6", - ) - assert out == "dreamshaper_8_a_1x512x512_cn_se_q6" - - -def test_quantize_nbits_invalid_raises(): - import pytest as _pytest - with _pytest.raises(ValueError, match="quantize_nbits"): - compose_out_name( - ckpt_name="x.safetensors", - batch_size=1, width=512, height=512, - controlnet_support=False, - attention_implementation="SPLIT_EINSUM", - quantize_nbits="16", # not in {none, 8, 6, 4} - ) diff --git a/tests/unit/test_conversion_helpers.py b/tests/unit/test_conversion_helpers.py deleted file mode 100644 index 0b3c53d..0000000 --- a/tests/unit/test_conversion_helpers.py +++ /dev/null @@ -1,183 +0,0 @@ -from types import SimpleNamespace - -import torch - -from coreml_suite.conversion.attention import ( - SplitEinsumAttnProcessor, - SplitEinsumV2AttnProcessor, - apply_attention_implementation, - split_einsum, - split_einsum_v2, -) -from coreml_suite.conversion.shapes import conv2d_output_shape -from coreml_suite.conversion.unet import CoreMLUNetWrapper - - -class RecordingUNet(torch.nn.Module): - def __init__(self): - super().__init__() - self.call = None - - def forward( - self, - sample, - timestep, - encoder_hidden_states, - timestep_cond=None, - added_cond_kwargs=None, - down_block_additional_residuals=None, - mid_block_additional_residual=None, - return_dict=True, - **kwargs, - ): - self.call = { - "sample": sample, - "timestep": timestep, - "encoder_hidden_states": encoder_hidden_states, - "timestep_cond": timestep_cond, - "added_cond_kwargs": added_cond_kwargs, - "down_block_additional_residuals": down_block_additional_residuals, - "mid_block_additional_residual": mid_block_additional_residual, - "return_dict": return_dict, - } - return (sample + 1,) - - -def test_conv2d_output_shape_matches_torch_conv2d_contract(): - conv = torch.nn.Conv2d( - 4, - 8, - kernel_size=(3, 5), - stride=(2, 3), - padding=(1, 2), - dilation=(1, 2), - ) - - assert conv2d_output_shape(17, 19, conv) == (9, 5) - - -def test_unet_wrapper_passes_context_through_for_sd15(): - unet = RecordingUNet() - wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="SD15")) - - sample = torch.randn(2, 4, 8, 8) - timestep = torch.randn(2) - context = torch.randn(2, 77, 768) - - out = wrapper(sample, timestep, context) - - assert torch.equal(out, sample + 1) - assert unet.call["encoder_hidden_states"] is context - assert unet.call["return_dict"] is False - - -def test_unet_wrapper_routes_lcm_sdxl_and_controlnet_inputs(): - unet = RecordingUNet() - wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="LCM")) - - sample = torch.randn(1, 4, 8, 8) - timestep = torch.randn(1) - context = torch.randn(1, 77, 768) - timestep_cond = torch.randn(1, 256) - down_residual = torch.randn(1, 320, 8, 8) - mid_residual = torch.randn(1, 1280, 1, 1) - - wrapper(sample, timestep, context, timestep_cond, down_residual, mid_residual) - - assert unet.call["timestep_cond"] is timestep_cond - assert len(unet.call["down_block_additional_residuals"]) == 1 - assert unet.call["down_block_additional_residuals"][0] is down_residual - assert unet.call["mid_block_additional_residual"] is mid_residual - - -def test_unet_wrapper_routes_sdxl_added_conditioning(): - unet = RecordingUNet() - wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="SDXL")) - - sample = torch.randn(1, 4, 8, 8) - timestep = torch.randn(1) - context = torch.randn(1, 77, 2048) - time_ids = torch.randn(1, 6) - text_embeds = torch.randn(1, 1280) - - wrapper(sample, timestep, context, time_ids, text_embeds) - - assert unet.call["added_cond_kwargs"]["time_ids"] is time_ids - assert unet.call["added_cond_kwargs"]["text_embeds"] is text_embeds - - -def test_split_einsum_matches_original_attention_math(): - torch.manual_seed(0) - batch = 2 - heads = 3 - dim_head = 4 - sequence = 16 - channels = heads * dim_head - q = torch.randn(batch, channels, 1, sequence) - k = torch.randn(batch, channels, 1, sequence) - v = torch.randn(batch, channels, 1, sequence) - - expected = _original_attention(q, k, v, None, heads, dim_head) - - # split-einsum reorders the float32 reductions vs the reference, so equality - # only holds up to rounding; the drift exceeds allclose's default atol on - # some BLAS backends (e.g. Linux x86 CI). - assert torch.allclose(split_einsum(q, k, v, None, heads, dim_head), expected, atol=1e-6) - assert torch.allclose(split_einsum_v2(q, k, v, None, heads, dim_head), expected, atol=1e-6) - - -def test_split_einsum_v2_chunked_path_matches_original_attention_math(): - torch.manual_seed(0) - batch = 1 - heads = 2 - dim_head = 2 - sequence = 512 - channels = heads * dim_head - q = torch.randn(batch, channels, 1, sequence) - k = torch.randn(batch, channels, 1, sequence) - v = torch.randn(batch, channels, 1, sequence) - - expected = _original_attention(q, k, v, None, heads, dim_head) - - assert torch.allclose( - split_einsum_v2(q, k, v, None, heads, dim_head), - expected, - atol=1e-6, - ) - - -def test_apply_attention_implementation_sets_split_processors(): - unet = RecordingProcessorUNet() - - assert apply_attention_implementation(unet, "ORIGINAL") is unet - assert unet.processor is None - - apply_attention_implementation(unet, "SPLIT_EINSUM") - assert isinstance(unet.processor, SplitEinsumAttnProcessor) - - apply_attention_implementation(unet, "SPLIT_EINSUM_V2") - assert isinstance(unet.processor, SplitEinsumV2AttnProcessor) - - -class RecordingProcessorUNet: - def __init__(self): - self.processor = None - - def set_attn_processor(self, processor): - self.processor = processor - - -def _original_attention(q, k, v, mask, heads, dim_head): - batch = q.size(0) - mh_q = q.view(batch, heads, dim_head, -1) - mh_k = k.view(batch, heads, dim_head, -1) - mh_v = v.view(batch, heads, dim_head, -1) - - weights = torch.einsum("bhcq,bhck->bhqk", mh_q, mh_k) - weights = weights * (dim_head**-0.5) - if mask is not None: - weights = weights + mask - weights = weights.softmax(dim=3) - - attn = torch.einsum("bhqk,bhck->bhcq", weights, mh_v) - return attn.contiguous().view(batch, heads * dim_head, 1, -1) diff --git a/uv.lock b/uv.lock index 35d7415..c2dbec9 100644 --- a/uv.lock +++ b/uv.lock @@ -209,13 +209,11 @@ name = "comfyui-coremlsuite" version = "2.0.2" source = { virtual = "." } dependencies = [ + { name = "coreml-diffusion" }, { name = "coremltools" }, { name = "diffusers" }, { name = "numpy" }, - { name = "omegaconf" }, - { name = "peft" }, { name = "torch" }, - { name = "transformers" }, ] [package.dev-dependencies] @@ -242,13 +240,11 @@ dev = [ [package.metadata] requires-dist = [ + { name = "coreml-diffusion", git = "https://github.com/aszc-dev/coreml-diffusion?tag=v0.1.0" }, { name = "coremltools", specifier = ">=9,<10" }, { name = "diffusers", specifier = ">=0.30" }, { name = "numpy", specifier = ">=2,<3" }, - { name = "omegaconf", specifier = ">=2.3" }, - { name = "peft", specifier = ">=0.13" }, { name = "torch", specifier = ">=2.7,<2.8" }, - { name = "transformers", specifier = ">=4.44" }, ] [package.metadata.requires-dev] @@ -282,6 +278,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/00/25/26c0d75006470d32bfc9eb39ac61a22600bfa509d672405ff3fe1561e2a8/comfyui_frontend_package-1.14.6-py3-none-any.whl", hash = "sha256:1044e30ff3c025dfb63f4c68ecd808a77050c3a5cc3c1f3e6421ea39800bbf40", size = 34873501, upload-time = "2025-03-27T15:50:30.726Z" }, ] +[[package]] +name = "coreml-diffusion" +version = "0.1.0" +source = { git = "https://github.com/aszc-dev/coreml-diffusion?tag=v0.1.0#0cc9382526ba3d1f2e6c8454836638d34b0977b5" } +dependencies = [ + { name = "coremltools" }, + { name = "diffusers" }, + { name = "numpy" }, + { name = "omegaconf" }, + { name = "peft" }, + { name = "torch" }, + { name = "transformers" }, +] + [[package]] name = "coremltools" version = "9.0"