feat(extraction): split conversion into the coreml-diffusion package (E0–E5) (#63)
* 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.
This commit is contained in:
aszc
@@ -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
|
||||
@@ -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: <path>; symbol counts (move / stay / already-framework-free)
|
||||
- python_coreml_stable_diffusion usage (verified by grep): conversion=<list> runtime=<list>
|
||||
- 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): <path>
|
||||
- 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: <path>
|
||||
- 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): <map>
|
||||
- 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 <x> / 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: <value> (> 25?)
|
||||
- Dead converter shims removed / reduced: <list>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 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: <none / list>
|
||||
```
|
||||
**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==<version>` from PyPI, or `git+...@<tag>`
|
||||
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: <path/branch>; final package name: <name>
|
||||
- Suite depends on package by pinned version: <spec>
|
||||
- Suite e2e [M2-ANE] PSNR after split: <value> (> 25?)
|
||||
- Conversion code fully absent from suite: confirmed (grep output)
|
||||
- Compatibility matrix documented: <link>
|
||||
- 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<bits>` 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: <link>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 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: <path/hash>
|
||||
- LCM node now uses unified from_single_file path; HF hardcode removed: confirmed
|
||||
- [M2-ANE] LCM golden after change: identical / within tol <x> / DIVERGED (STOP)
|
||||
- Status flipped lcm→VERIFIED in package (minor bump <ver>): 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): <link>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 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.**
|
||||
@@ -1,5 +0,0 @@
|
||||
ATTENTION_IMPLEMENTATIONS = (
|
||||
"SPLIT_EINSUM",
|
||||
"SPLIT_EINSUM_V2",
|
||||
"ORIGINAL",
|
||||
)
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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"}
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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<bits>` suffix.
|
||||
QUANT_NBITS_VALUES = ("none", "8", "6", "4")
|
||||
|
||||
|
||||
def compose_out_name(
|
||||
*,
|
||||
ckpt_name: str,
|
||||
batch_size: int,
|
||||
width: int,
|
||||
height: int,
|
||||
controlnet_support: bool,
|
||||
attention_implementation: str,
|
||||
lora_names: Iterable[str] = (),
|
||||
quantize_nbits: str = "none",
|
||||
) -> str:
|
||||
"""Build the .mlpackage stem from convert() parameters.
|
||||
|
||||
Locked behaviour (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<bits>` after the attn suffix
|
||||
"""
|
||||
if quantize_nbits not in QUANT_NBITS_VALUES:
|
||||
raise ValueError(
|
||||
f"quantize_nbits={quantize_nbits!r} not in {QUANT_NBITS_VALUES}"
|
||||
)
|
||||
stem = ckpt_name.split(".")[0]
|
||||
sorted_names = sorted(lora_names)
|
||||
lora_str = "_" + "_".join(name.split(".")[0] for name in sorted_names) if sorted_names else ""
|
||||
cn_suffix = "_cn" if controlnet_support else ""
|
||||
attn_suffix = "_" + ATTN_SUFFIX[attention_implementation]
|
||||
quant_suffix = f"_q{quantize_nbits}" if quantize_nbits != "none" else ""
|
||||
out_name = (
|
||||
f"{stem}{lora_str}_{batch_size}x{width}x{height}"
|
||||
f"{cn_suffix}{attn_suffix}{quant_suffix}"
|
||||
)
|
||||
return out_name.replace(" ", "_")
|
||||
|
||||
|
||||
def lora_names_from_params(lora_params: Iterable[Tuple[str, float]]) -> list[str]:
|
||||
"""Mirror the sort applied inside CoreMLConverter.convert."""
|
||||
return [name for name, _ in sorted(lora_params, key=lambda pair: pair[0])]
|
||||
+31
-146
@@ -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..")
|
||||
|
||||
@@ -1,8 +0,0 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ModelVersion(Enum):
|
||||
SD15 = "sd15"
|
||||
SDXL = "sdxl"
|
||||
SDXL_REFINER = "sdxl_refiner"
|
||||
LCM = "lcm"
|
||||
+42
-24
@@ -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,
|
||||
|
||||
+9
-4
@@ -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"
|
||||
|
||||
+1
-3
@@ -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
|
||||
|
||||
@@ -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@<sha>`.** 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.
|
||||
@@ -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)
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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}
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user