Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 aszc-dev
|
45be061edb |
@@ -0,0 +1,23 @@
|
||||
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
|
||||
@@ -3,4 +3,3 @@ __pycache__/
|
||||
models/
|
||||
.venv/
|
||||
test_results/
|
||||
.claude/
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
3.12
|
||||
@@ -1,618 +0,0 @@
|
||||
# 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,21 +1,674 @@
|
||||
MIT License
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 3, 29 June 2007
|
||||
|
||||
Copyright (c) 2023-2026 Adrian Szczepański
|
||||
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
|
||||
Everyone is permitted to copy and distribute verbatim copies
|
||||
of this license document, but changing it is not allowed.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
Preamble
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
The GNU General Public License is a free, copyleft license for
|
||||
software and other kinds of works.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
The licenses for most software and other practical works are designed
|
||||
to take away your freedom to share and change the works. By contrast,
|
||||
the GNU General Public License is intended to guarantee your freedom to
|
||||
share and change all versions of a program--to make sure it remains free
|
||||
software for all its users. We, the Free Software Foundation, use the
|
||||
GNU General Public License for most of our software; it applies also to
|
||||
any other work released this way by its authors. You can apply it to
|
||||
your programs, too.
|
||||
|
||||
When we speak of free software, we are referring to freedom, not
|
||||
price. Our General Public Licenses are designed to make sure that you
|
||||
have the freedom to distribute copies of free software (and charge for
|
||||
them if you wish), that you receive source code or can get it if you
|
||||
want it, that you can change the software or use pieces of it in new
|
||||
free programs, and that you know you can do these things.
|
||||
|
||||
To protect your rights, we need to prevent others from denying you
|
||||
these rights or asking you to surrender the rights. Therefore, you have
|
||||
certain responsibilities if you distribute copies of the software, or if
|
||||
you modify it: responsibilities to respect the freedom of others.
|
||||
|
||||
For example, if you distribute copies of such a program, whether
|
||||
gratis or for a fee, you must pass on to the recipients the same
|
||||
freedoms that you received. You must make sure that they, too, receive
|
||||
or can get the source code. And you must show them these terms so they
|
||||
know their rights.
|
||||
|
||||
Developers that use the GNU GPL protect your rights with two steps:
|
||||
(1) assert copyright on the software, and (2) offer you this License
|
||||
giving you legal permission to copy, distribute and/or modify it.
|
||||
|
||||
For the developers' and authors' protection, the GPL clearly explains
|
||||
that there is no warranty for this free software. For both users' and
|
||||
authors' sake, the GPL requires that modified versions be marked as
|
||||
changed, so that their problems will not be attributed erroneously to
|
||||
authors of previous versions.
|
||||
|
||||
Some devices are designed to deny users access to install or run
|
||||
modified versions of the software inside them, although the manufacturer
|
||||
can do so. This is fundamentally incompatible with the aim of
|
||||
protecting users' freedom to change the software. The systematic
|
||||
pattern of such abuse occurs in the area of products for individuals to
|
||||
use, which is precisely where it is most unacceptable. Therefore, we
|
||||
have designed this version of the GPL to prohibit the practice for those
|
||||
products. If such problems arise substantially in other domains, we
|
||||
stand ready to extend this provision to those domains in future versions
|
||||
of the GPL, as needed to protect the freedom of users.
|
||||
|
||||
Finally, every program is threatened constantly by software patents.
|
||||
States should not allow patents to restrict development and use of
|
||||
software on general-purpose computers, but in those that do, we wish to
|
||||
avoid the special danger that patents applied to a free program could
|
||||
make it effectively proprietary. To prevent this, the GPL assures that
|
||||
patents cannot be used to render the program non-free.
|
||||
|
||||
The precise terms and conditions for copying, distribution and
|
||||
modification follow.
|
||||
|
||||
TERMS AND CONDITIONS
|
||||
|
||||
0. Definitions.
|
||||
|
||||
"This License" refers to version 3 of the GNU General Public License.
|
||||
|
||||
"Copyright" also means copyright-like laws that apply to other kinds of
|
||||
works, such as semiconductor masks.
|
||||
|
||||
"The Program" refers to any copyrightable work licensed under this
|
||||
License. Each licensee is addressed as "you". "Licensees" and
|
||||
"recipients" may be individuals or organizations.
|
||||
|
||||
To "modify" a work means to copy from or adapt all or part of the work
|
||||
in a fashion requiring copyright permission, other than the making of an
|
||||
exact copy. The resulting work is called a "modified version" of the
|
||||
earlier work or a work "based on" the earlier work.
|
||||
|
||||
A "covered work" means either the unmodified Program or a work based
|
||||
on the Program.
|
||||
|
||||
To "propagate" a work means to do anything with it that, without
|
||||
permission, would make you directly or secondarily liable for
|
||||
infringement under applicable copyright law, except executing it on a
|
||||
computer or modifying a private copy. Propagation includes copying,
|
||||
distribution (with or without modification), making available to the
|
||||
public, and in some countries other activities as well.
|
||||
|
||||
To "convey" a work means any kind of propagation that enables other
|
||||
parties to make or receive copies. Mere interaction with a user through
|
||||
a computer network, with no transfer of a copy, is not conveying.
|
||||
|
||||
An interactive user interface displays "Appropriate Legal Notices"
|
||||
to the extent that it includes a convenient and prominently visible
|
||||
feature that (1) displays an appropriate copyright notice, and (2)
|
||||
tells the user that there is no warranty for the work (except to the
|
||||
extent that warranties are provided), that licensees may convey the
|
||||
work under this License, and how to view a copy of this License. If
|
||||
the interface presents a list of user commands or options, such as a
|
||||
menu, a prominent item in the list meets this criterion.
|
||||
|
||||
1. Source Code.
|
||||
|
||||
The "source code" for a work means the preferred form of the work
|
||||
for making modifications to it. "Object code" means any non-source
|
||||
form of a work.
|
||||
|
||||
A "Standard Interface" means an interface that either is an official
|
||||
standard defined by a recognized standards body, or, in the case of
|
||||
interfaces specified for a particular programming language, one that
|
||||
is widely used among developers working in that language.
|
||||
|
||||
The "System Libraries" of an executable work include anything, other
|
||||
than the work as a whole, that (a) is included in the normal form of
|
||||
packaging a Major Component, but which is not part of that Major
|
||||
Component, and (b) serves only to enable use of the work with that
|
||||
Major Component, or to implement a Standard Interface for which an
|
||||
implementation is available to the public in source code form. A
|
||||
"Major Component", in this context, means a major essential component
|
||||
(kernel, window system, and so on) of the specific operating system
|
||||
(if any) on which the executable work runs, or a compiler used to
|
||||
produce the work, or an object code interpreter used to run it.
|
||||
|
||||
The "Corresponding Source" for a work in object code form means all
|
||||
the source code needed to generate, install, and (for an executable
|
||||
work) run the object code and to modify the work, including scripts to
|
||||
control those activities. However, it does not include the work's
|
||||
System Libraries, or general-purpose tools or generally available free
|
||||
programs which are used unmodified in performing those activities but
|
||||
which are not part of the work. For example, Corresponding Source
|
||||
includes interface definition files associated with source files for
|
||||
the work, and the source code for shared libraries and dynamically
|
||||
linked subprograms that the work is specifically designed to require,
|
||||
such as by intimate data communication or control flow between those
|
||||
subprograms and other parts of the work.
|
||||
|
||||
The Corresponding Source need not include anything that users
|
||||
can regenerate automatically from other parts of the Corresponding
|
||||
Source.
|
||||
|
||||
The Corresponding Source for a work in source code form is that
|
||||
same work.
|
||||
|
||||
2. Basic Permissions.
|
||||
|
||||
All rights granted under this License are granted for the term of
|
||||
copyright on the Program, and are irrevocable provided the stated
|
||||
conditions are met. This License explicitly affirms your unlimited
|
||||
permission to run the unmodified Program. The output from running a
|
||||
covered work is covered by this License only if the output, given its
|
||||
content, constitutes a covered work. This License acknowledges your
|
||||
rights of fair use or other equivalent, as provided by copyright law.
|
||||
|
||||
You may make, run and propagate covered works that you do not
|
||||
convey, without conditions so long as your license otherwise remains
|
||||
in force. You may convey covered works to others for the sole purpose
|
||||
of having them make modifications exclusively for you, or provide you
|
||||
with facilities for running those works, provided that you comply with
|
||||
the terms of this License in conveying all material for which you do
|
||||
not control copyright. Those thus making or running the covered works
|
||||
for you must do so exclusively on your behalf, under your direction
|
||||
and control, on terms that prohibit them from making any copies of
|
||||
your copyrighted material outside their relationship with you.
|
||||
|
||||
Conveying under any other circumstances is permitted solely under
|
||||
the conditions stated below. Sublicensing is not allowed; section 10
|
||||
makes it unnecessary.
|
||||
|
||||
3. Protecting Users' Legal Rights From Anti-Circumvention Law.
|
||||
|
||||
No covered work shall be deemed part of an effective technological
|
||||
measure under any applicable law fulfilling obligations under article
|
||||
11 of the WIPO copyright treaty adopted on 20 December 1996, or
|
||||
similar laws prohibiting or restricting circumvention of such
|
||||
measures.
|
||||
|
||||
When you convey a covered work, you waive any legal power to forbid
|
||||
circumvention of technological measures to the extent such circumvention
|
||||
is effected by exercising rights under this License with respect to
|
||||
the covered work, and you disclaim any intention to limit operation or
|
||||
modification of the work as a means of enforcing, against the work's
|
||||
users, your or third parties' legal rights to forbid circumvention of
|
||||
technological measures.
|
||||
|
||||
4. Conveying Verbatim Copies.
|
||||
|
||||
You may convey verbatim copies of the Program's source code as you
|
||||
receive it, in any medium, provided that you conspicuously and
|
||||
appropriately publish on each copy an appropriate copyright notice;
|
||||
keep intact all notices stating that this License and any
|
||||
non-permissive terms added in accord with section 7 apply to the code;
|
||||
keep intact all notices of the absence of any warranty; and give all
|
||||
recipients a copy of this License along with the Program.
|
||||
|
||||
You may charge any price or no price for each copy that you convey,
|
||||
and you may offer support or warranty protection for a fee.
|
||||
|
||||
5. Conveying Modified Source Versions.
|
||||
|
||||
You may convey a work based on the Program, or the modifications to
|
||||
produce it from the Program, in the form of source code under the
|
||||
terms of section 4, provided that you also meet all of these conditions:
|
||||
|
||||
a) The work must carry prominent notices stating that you modified
|
||||
it, and giving a relevant date.
|
||||
|
||||
b) The work must carry prominent notices stating that it is
|
||||
released under this License and any conditions added under section
|
||||
7. This requirement modifies the requirement in section 4 to
|
||||
"keep intact all notices".
|
||||
|
||||
c) You must license the entire work, as a whole, under this
|
||||
License to anyone who comes into possession of a copy. This
|
||||
License will therefore apply, along with any applicable section 7
|
||||
additional terms, to the whole of the work, and all its parts,
|
||||
regardless of how they are packaged. This License gives no
|
||||
permission to license the work in any other way, but it does not
|
||||
invalidate such permission if you have separately received it.
|
||||
|
||||
d) If the work has interactive user interfaces, each must display
|
||||
Appropriate Legal Notices; however, if the Program has interactive
|
||||
interfaces that do not display Appropriate Legal Notices, your
|
||||
work need not make them do so.
|
||||
|
||||
A compilation of a covered work with other separate and independent
|
||||
works, which are not by their nature extensions of the covered work,
|
||||
and which are not combined with it such as to form a larger program,
|
||||
in or on a volume of a storage or distribution medium, is called an
|
||||
"aggregate" if the compilation and its resulting copyright are not
|
||||
used to limit the access or legal rights of the compilation's users
|
||||
beyond what the individual works permit. Inclusion of a covered work
|
||||
in an aggregate does not cause this License to apply to the other
|
||||
parts of the aggregate.
|
||||
|
||||
6. Conveying Non-Source Forms.
|
||||
|
||||
You may convey a covered work in object code form under the terms
|
||||
of sections 4 and 5, provided that you also convey the
|
||||
machine-readable Corresponding Source under the terms of this License,
|
||||
in one of these ways:
|
||||
|
||||
a) Convey the object code in, or embodied in, a physical product
|
||||
(including a physical distribution medium), accompanied by the
|
||||
Corresponding Source fixed on a durable physical medium
|
||||
customarily used for software interchange.
|
||||
|
||||
b) Convey the object code in, or embodied in, a physical product
|
||||
(including a physical distribution medium), accompanied by a
|
||||
written offer, valid for at least three years and valid for as
|
||||
long as you offer spare parts or customer support for that product
|
||||
model, to give anyone who possesses the object code either (1) a
|
||||
copy of the Corresponding Source for all the software in the
|
||||
product that is covered by this License, on a durable physical
|
||||
medium customarily used for software interchange, for a price no
|
||||
more than your reasonable cost of physically performing this
|
||||
conveying of source, or (2) access to copy the
|
||||
Corresponding Source from a network server at no charge.
|
||||
|
||||
c) Convey individual copies of the object code with a copy of the
|
||||
written offer to provide the Corresponding Source. This
|
||||
alternative is allowed only occasionally and noncommercially, and
|
||||
only if you received the object code with such an offer, in accord
|
||||
with subsection 6b.
|
||||
|
||||
d) Convey the object code by offering access from a designated
|
||||
place (gratis or for a charge), and offer equivalent access to the
|
||||
Corresponding Source in the same way through the same place at no
|
||||
further charge. You need not require recipients to copy the
|
||||
Corresponding Source along with the object code. If the place to
|
||||
copy the object code is a network server, the Corresponding Source
|
||||
may be on a different server (operated by you or a third party)
|
||||
that supports equivalent copying facilities, provided you maintain
|
||||
clear directions next to the object code saying where to find the
|
||||
Corresponding Source. Regardless of what server hosts the
|
||||
Corresponding Source, you remain obligated to ensure that it is
|
||||
available for as long as needed to satisfy these requirements.
|
||||
|
||||
e) Convey the object code using peer-to-peer transmission, provided
|
||||
you inform other peers where the object code and Corresponding
|
||||
Source of the work are being offered to the general public at no
|
||||
charge under subsection 6d.
|
||||
|
||||
A separable portion of the object code, whose source code is excluded
|
||||
from the Corresponding Source as a System Library, need not be
|
||||
included in conveying the object code work.
|
||||
|
||||
A "User Product" is either (1) a "consumer product", which means any
|
||||
tangible personal property which is normally used for personal, family,
|
||||
or household purposes, or (2) anything designed or sold for incorporation
|
||||
into a dwelling. In determining whether a product is a consumer product,
|
||||
doubtful cases shall be resolved in favor of coverage. For a particular
|
||||
product received by a particular user, "normally used" refers to a
|
||||
typical or common use of that class of product, regardless of the status
|
||||
of the particular user or of the way in which the particular user
|
||||
actually uses, or expects or is expected to use, the product. A product
|
||||
is a consumer product regardless of whether the product has substantial
|
||||
commercial, industrial or non-consumer uses, unless such uses represent
|
||||
the only significant mode of use of the product.
|
||||
|
||||
"Installation Information" for a User Product means any methods,
|
||||
procedures, authorization keys, or other information required to install
|
||||
and execute modified versions of a covered work in that User Product from
|
||||
a modified version of its Corresponding Source. The information must
|
||||
suffice to ensure that the continued functioning of the modified object
|
||||
code is in no case prevented or interfered with solely because
|
||||
modification has been made.
|
||||
|
||||
If you convey an object code work under this section in, or with, or
|
||||
specifically for use in, a User Product, and the conveying occurs as
|
||||
part of a transaction in which the right of possession and use of the
|
||||
User Product is transferred to the recipient in perpetuity or for a
|
||||
fixed term (regardless of how the transaction is characterized), the
|
||||
Corresponding Source conveyed under this section must be accompanied
|
||||
by the Installation Information. But this requirement does not apply
|
||||
if neither you nor any third party retains the ability to install
|
||||
modified object code on the User Product (for example, the work has
|
||||
been installed in ROM).
|
||||
|
||||
The requirement to provide Installation Information does not include a
|
||||
requirement to continue to provide support service, warranty, or updates
|
||||
for a work that has been modified or installed by the recipient, or for
|
||||
the User Product in which it has been modified or installed. Access to a
|
||||
network may be denied when the modification itself materially and
|
||||
adversely affects the operation of the network or violates the rules and
|
||||
protocols for communication across the network.
|
||||
|
||||
Corresponding Source conveyed, and Installation Information provided,
|
||||
in accord with this section must be in a format that is publicly
|
||||
documented (and with an implementation available to the public in
|
||||
source code form), and must require no special password or key for
|
||||
unpacking, reading or copying.
|
||||
|
||||
7. Additional Terms.
|
||||
|
||||
"Additional permissions" are terms that supplement the terms of this
|
||||
License by making exceptions from one or more of its conditions.
|
||||
Additional permissions that are applicable to the entire Program shall
|
||||
be treated as though they were included in this License, to the extent
|
||||
that they are valid under applicable law. If additional permissions
|
||||
apply only to part of the Program, that part may be used separately
|
||||
under those permissions, but the entire Program remains governed by
|
||||
this License without regard to the additional permissions.
|
||||
|
||||
When you convey a copy of a covered work, you may at your option
|
||||
remove any additional permissions from that copy, or from any part of
|
||||
it. (Additional permissions may be written to require their own
|
||||
removal in certain cases when you modify the work.) You may place
|
||||
additional permissions on material, added by you to a covered work,
|
||||
for which you have or can give appropriate copyright permission.
|
||||
|
||||
Notwithstanding any other provision of this License, for material you
|
||||
add to a covered work, you may (if authorized by the copyright holders of
|
||||
that material) supplement the terms of this License with terms:
|
||||
|
||||
a) Disclaiming warranty or limiting liability differently from the
|
||||
terms of sections 15 and 16 of this License; or
|
||||
|
||||
b) Requiring preservation of specified reasonable legal notices or
|
||||
author attributions in that material or in the Appropriate Legal
|
||||
Notices displayed by works containing it; or
|
||||
|
||||
c) Prohibiting misrepresentation of the origin of that material, or
|
||||
requiring that modified versions of such material be marked in
|
||||
reasonable ways as different from the original version; or
|
||||
|
||||
d) Limiting the use for publicity purposes of names of licensors or
|
||||
authors of the material; or
|
||||
|
||||
e) Declining to grant rights under trademark law for use of some
|
||||
trade names, trademarks, or service marks; or
|
||||
|
||||
f) Requiring indemnification of licensors and authors of that
|
||||
material by anyone who conveys the material (or modified versions of
|
||||
it) with contractual assumptions of liability to the recipient, for
|
||||
any liability that these contractual assumptions directly impose on
|
||||
those licensors and authors.
|
||||
|
||||
All other non-permissive additional terms are considered "further
|
||||
restrictions" within the meaning of section 10. If the Program as you
|
||||
received it, or any part of it, contains a notice stating that it is
|
||||
governed by this License along with a term that is a further
|
||||
restriction, you may remove that term. If a license document contains
|
||||
a further restriction but permits relicensing or conveying under this
|
||||
License, you may add to a covered work material governed by the terms
|
||||
of that license document, provided that the further restriction does
|
||||
not survive such relicensing or conveying.
|
||||
|
||||
If you add terms to a covered work in accord with this section, you
|
||||
must place, in the relevant source files, a statement of the
|
||||
additional terms that apply to those files, or a notice indicating
|
||||
where to find the applicable terms.
|
||||
|
||||
Additional terms, permissive or non-permissive, may be stated in the
|
||||
form of a separately written license, or stated as exceptions;
|
||||
the above requirements apply either way.
|
||||
|
||||
8. Termination.
|
||||
|
||||
You may not propagate or modify a covered work except as expressly
|
||||
provided under this License. Any attempt otherwise to propagate or
|
||||
modify it is void, and will automatically terminate your rights under
|
||||
this License (including any patent licenses granted under the third
|
||||
paragraph of section 11).
|
||||
|
||||
However, if you cease all violation of this License, then your
|
||||
license from a particular copyright holder is reinstated (a)
|
||||
provisionally, unless and until the copyright holder explicitly and
|
||||
finally terminates your license, and (b) permanently, if the copyright
|
||||
holder fails to notify you of the violation by some reasonable means
|
||||
prior to 60 days after the cessation.
|
||||
|
||||
Moreover, your license from a particular copyright holder is
|
||||
reinstated permanently if the copyright holder notifies you of the
|
||||
violation by some reasonable means, this is the first time you have
|
||||
received notice of violation of this License (for any work) from that
|
||||
copyright holder, and you cure the violation prior to 30 days after
|
||||
your receipt of the notice.
|
||||
|
||||
Termination of your rights under this section does not terminate the
|
||||
licenses of parties who have received copies or rights from you under
|
||||
this License. If your rights have been terminated and not permanently
|
||||
reinstated, you do not qualify to receive new licenses for the same
|
||||
material under section 10.
|
||||
|
||||
9. Acceptance Not Required for Having Copies.
|
||||
|
||||
You are not required to accept this License in order to receive or
|
||||
run a copy of the Program. Ancillary propagation of a covered work
|
||||
occurring solely as a consequence of using peer-to-peer transmission
|
||||
to receive a copy likewise does not require acceptance. However,
|
||||
nothing other than this License grants you permission to propagate or
|
||||
modify any covered work. These actions infringe copyright if you do
|
||||
not accept this License. Therefore, by modifying or propagating a
|
||||
covered work, you indicate your acceptance of this License to do so.
|
||||
|
||||
10. Automatic Licensing of Downstream Recipients.
|
||||
|
||||
Each time you convey a covered work, the recipient automatically
|
||||
receives a license from the original licensors, to run, modify and
|
||||
propagate that work, subject to this License. You are not responsible
|
||||
for enforcing compliance by third parties with this License.
|
||||
|
||||
An "entity transaction" is a transaction transferring control of an
|
||||
organization, or substantially all assets of one, or subdividing an
|
||||
organization, or merging organizations. If propagation of a covered
|
||||
work results from an entity transaction, each party to that
|
||||
transaction who receives a copy of the work also receives whatever
|
||||
licenses to the work the party's predecessor in interest had or could
|
||||
give under the previous paragraph, plus a right to possession of the
|
||||
Corresponding Source of the work from the predecessor in interest, if
|
||||
the predecessor has it or can get it with reasonable efforts.
|
||||
|
||||
You may not impose any further restrictions on the exercise of the
|
||||
rights granted or affirmed under this License. For example, you may
|
||||
not impose a license fee, royalty, or other charge for exercise of
|
||||
rights granted under this License, and you may not initiate litigation
|
||||
(including a cross-claim or counterclaim in a lawsuit) alleging that
|
||||
any patent claim is infringed by making, using, selling, offering for
|
||||
sale, or importing the Program or any portion of it.
|
||||
|
||||
11. Patents.
|
||||
|
||||
A "contributor" is a copyright holder who authorizes use under this
|
||||
License of the Program or a work on which the Program is based. The
|
||||
work thus licensed is called the contributor's "contributor version".
|
||||
|
||||
A contributor's "essential patent claims" are all patent claims
|
||||
owned or controlled by the contributor, whether already acquired or
|
||||
hereafter acquired, that would be infringed by some manner, permitted
|
||||
by this License, of making, using, or selling its contributor version,
|
||||
but do not include claims that would be infringed only as a
|
||||
consequence of further modification of the contributor version. For
|
||||
purposes of this definition, "control" includes the right to grant
|
||||
patent sublicenses in a manner consistent with the requirements of
|
||||
this License.
|
||||
|
||||
Each contributor grants you a non-exclusive, worldwide, royalty-free
|
||||
patent license under the contributor's essential patent claims, to
|
||||
make, use, sell, offer for sale, import and otherwise run, modify and
|
||||
propagate the contents of its contributor version.
|
||||
|
||||
In the following three paragraphs, a "patent license" is any express
|
||||
agreement or commitment, however denominated, not to enforce a patent
|
||||
(such as an express permission to practice a patent or covenant not to
|
||||
sue for patent infringement). To "grant" such a patent license to a
|
||||
party means to make such an agreement or commitment not to enforce a
|
||||
patent against the party.
|
||||
|
||||
If you convey a covered work, knowingly relying on a patent license,
|
||||
and the Corresponding Source of the work is not available for anyone
|
||||
to copy, free of charge and under the terms of this License, through a
|
||||
publicly available network server or other readily accessible means,
|
||||
then you must either (1) cause the Corresponding Source to be so
|
||||
available, or (2) arrange to deprive yourself of the benefit of the
|
||||
patent license for this particular work, or (3) arrange, in a manner
|
||||
consistent with the requirements of this License, to extend the patent
|
||||
license to downstream recipients. "Knowingly relying" means you have
|
||||
actual knowledge that, but for the patent license, your conveying the
|
||||
covered work in a country, or your recipient's use of the covered work
|
||||
in a country, would infringe one or more identifiable patents in that
|
||||
country that you have reason to believe are valid.
|
||||
|
||||
If, pursuant to or in connection with a single transaction or
|
||||
arrangement, you convey, or propagate by procuring conveyance of, a
|
||||
covered work, and grant a patent license to some of the parties
|
||||
receiving the covered work authorizing them to use, propagate, modify
|
||||
or convey a specific copy of the covered work, then the patent license
|
||||
you grant is automatically extended to all recipients of the covered
|
||||
work and works based on it.
|
||||
|
||||
A patent license is "discriminatory" if it does not include within
|
||||
the scope of its coverage, prohibits the exercise of, or is
|
||||
conditioned on the non-exercise of one or more of the rights that are
|
||||
specifically granted under this License. You may not convey a covered
|
||||
work if you are a party to an arrangement with a third party that is
|
||||
in the business of distributing software, under which you make payment
|
||||
to the third party based on the extent of your activity of conveying
|
||||
the work, and under which the third party grants, to any of the
|
||||
parties who would receive the covered work from you, a discriminatory
|
||||
patent license (a) in connection with copies of the covered work
|
||||
conveyed by you (or copies made from those copies), or (b) primarily
|
||||
for and in connection with specific products or compilations that
|
||||
contain the covered work, unless you entered into that arrangement,
|
||||
or that patent license was granted, prior to 28 March 2007.
|
||||
|
||||
Nothing in this License shall be construed as excluding or limiting
|
||||
any implied license or other defenses to infringement that may
|
||||
otherwise be available to you under applicable patent law.
|
||||
|
||||
12. No Surrender of Others' Freedom.
|
||||
|
||||
If conditions are imposed on you (whether by court order, agreement or
|
||||
otherwise) that contradict the conditions of this License, they do not
|
||||
excuse you from the conditions of this License. If you cannot convey a
|
||||
covered work so as to satisfy simultaneously your obligations under this
|
||||
License and any other pertinent obligations, then as a consequence you may
|
||||
not convey it at all. For example, if you agree to terms that obligate you
|
||||
to collect a royalty for further conveying from those to whom you convey
|
||||
the Program, the only way you could satisfy both those terms and this
|
||||
License would be to refrain entirely from conveying the Program.
|
||||
|
||||
13. Use with the GNU Affero General Public License.
|
||||
|
||||
Notwithstanding any other provision of this License, you have
|
||||
permission to link or combine any covered work with a work licensed
|
||||
under version 3 of the GNU Affero General Public License into a single
|
||||
combined work, and to convey the resulting work. The terms of this
|
||||
License will continue to apply to the part which is the covered work,
|
||||
but the special requirements of the GNU Affero General Public License,
|
||||
section 13, concerning interaction through a network will apply to the
|
||||
combination as such.
|
||||
|
||||
14. Revised Versions of this License.
|
||||
|
||||
The Free Software Foundation may publish revised and/or new versions of
|
||||
the GNU General Public License from time to time. Such new versions will
|
||||
be similar in spirit to the present version, but may differ in detail to
|
||||
address new problems or concerns.
|
||||
|
||||
Each version is given a distinguishing version number. If the
|
||||
Program specifies that a certain numbered version of the GNU General
|
||||
Public License "or any later version" applies to it, you have the
|
||||
option of following the terms and conditions either of that numbered
|
||||
version or of any later version published by the Free Software
|
||||
Foundation. If the Program does not specify a version number of the
|
||||
GNU General Public License, you may choose any version ever published
|
||||
by the Free Software Foundation.
|
||||
|
||||
If the Program specifies that a proxy can decide which future
|
||||
versions of the GNU General Public License can be used, that proxy's
|
||||
public statement of acceptance of a version permanently authorizes you
|
||||
to choose that version for the Program.
|
||||
|
||||
Later license versions may give you additional or different
|
||||
permissions. However, no additional obligations are imposed on any
|
||||
author or copyright holder as a result of your choosing to follow a
|
||||
later version.
|
||||
|
||||
15. Disclaimer of Warranty.
|
||||
|
||||
THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
|
||||
APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
|
||||
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
|
||||
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
|
||||
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
|
||||
PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
|
||||
IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
|
||||
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
|
||||
|
||||
16. Limitation of Liability.
|
||||
|
||||
IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
|
||||
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
|
||||
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
|
||||
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
|
||||
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
|
||||
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
|
||||
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
|
||||
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
|
||||
SUCH DAMAGES.
|
||||
|
||||
17. Interpretation of Sections 15 and 16.
|
||||
|
||||
If the disclaimer of warranty and limitation of liability provided
|
||||
above cannot be given local legal effect according to their terms,
|
||||
reviewing courts shall apply local law that most closely approximates
|
||||
an absolute waiver of all civil liability in connection with the
|
||||
Program, unless a warranty or assumption of liability accompanies a
|
||||
copy of the Program in return for a fee.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
How to Apply These Terms to Your New Programs
|
||||
|
||||
If you develop a new program, and you want it to be of the greatest
|
||||
possible use to the public, the best way to achieve this is to make it
|
||||
free software which everyone can redistribute and change under these terms.
|
||||
|
||||
To do so, attach the following notices to the program. It is safest
|
||||
to attach them to the start of each source file to most effectively
|
||||
state the exclusion of warranty; and each file should have at least
|
||||
the "copyright" line and a pointer to where the full notice is found.
|
||||
|
||||
<one line to give the program's name and a brief idea of what it does.>
|
||||
Copyright (C) <year> <name of author>
|
||||
|
||||
This program is free software: you can redistribute it and/or modify
|
||||
it under the terms of the GNU General Public License as published by
|
||||
the Free Software Foundation, either version 3 of the License, or
|
||||
(at your option) any later version.
|
||||
|
||||
This program is distributed in the hope that it will be useful,
|
||||
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
GNU General Public License for more details.
|
||||
|
||||
You should have received a copy of the GNU General Public License
|
||||
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
Also add information on how to contact you by electronic and paper mail.
|
||||
|
||||
If the program does terminal interaction, make it output a short
|
||||
notice like this when it starts in an interactive mode:
|
||||
|
||||
<program> Copyright (C) <year> <name of author>
|
||||
This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
|
||||
This is free software, and you are welcome to redistribute it
|
||||
under certain conditions; type `show c' for details.
|
||||
|
||||
The hypothetical commands `show w' and `show c' should show the appropriate
|
||||
parts of the General Public License. Of course, your program's commands
|
||||
might be different; for a GUI interface, you would use an "about box".
|
||||
|
||||
You should also get your employer (if you work as a programmer) or school,
|
||||
if any, to sign a "copyright disclaimer" for the program, if necessary.
|
||||
For more information on this, and how to apply and follow the GNU GPL, see
|
||||
<https://www.gnu.org/licenses/>.
|
||||
|
||||
The GNU General Public License does not permit incorporating your program
|
||||
into proprietary programs. If your program is a subroutine library, you
|
||||
may consider it more useful to permit linking proprietary applications with
|
||||
the library. If this is what you want to do, use the GNU Lesser General
|
||||
Public License instead of this License. But first, please read
|
||||
<https://www.gnu.org/licenses/why-not-lgpl.html>.
|
||||
|
||||
@@ -11,6 +11,9 @@ from coreml_suite.nodes import (
|
||||
CoreMLConverter,
|
||||
COREML_LOAD_LORA,
|
||||
)
|
||||
from coreml_suite.lcm import (
|
||||
COREML_CONVERT_LCM,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CoreMLUNetLoader": CoreMLLoaderUNet,
|
||||
@@ -19,6 +22,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"CoreMLModelAdapter": CoreMLModelAdapter,
|
||||
"Core ML LoRA Loader": COREML_LOAD_LORA,
|
||||
"Core ML Converter": CoreMLConverter,
|
||||
"Core ML LCM Converter": COREML_CONVERT_LCM,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CoreMLUNetLoader": "Load Core ML UNet",
|
||||
@@ -27,4 +31,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
|
||||
"Core ML LoRA Loader": "Load LoRA to use with Core ML",
|
||||
"Core ML Converter": "Convert Checkpoint to Core ML",
|
||||
"Core ML LCM Converter": "Convert LCM to Core ML",
|
||||
}
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
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_diffusion import ModelVersion
|
||||
from coreml_suite.model_version import ModelVersion
|
||||
|
||||
|
||||
config_map = {
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
"""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.
|
||||
"""
|
||||
@@ -0,0 +1,239 @@
|
||||
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
|
||||
@@ -0,0 +1,20 @@
|
||||
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
|
||||
@@ -0,0 +1,61 @@
|
||||
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
|
||||
@@ -0,0 +1,54 @@
|
||||
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"}
|
||||
@@ -0,0 +1,322 @@
|
||||
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,
|
||||
)
|
||||
@@ -0,0 +1,68 @@
|
||||
"""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])]
|
||||
@@ -1,8 +1,3 @@
|
||||
"""LCM runtime support (sampler-side).
|
||||
from .nodes import COREML_CONVERT_LCM
|
||||
|
||||
The dedicated LCM converter node was removed once the standard ``CoreMLConverter``
|
||||
gained model-version auto-detection (full-distill LCM is detected from the
|
||||
checkpoint). What remains here is runtime sampling support — ``utils`` patches the
|
||||
model sampling and supplies the guidance embedding when a converted UNet exposes
|
||||
``timestep_cond``.
|
||||
"""
|
||||
__all__ = ["COREML_CONVERT_LCM"]
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
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
|
||||
|
||||
logging.basicConfig()
|
||||
logger = logging.getLogger(__name__)
|
||||
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():
|
||||
ref_unet = UNet2DConditionModel.from_pretrained(
|
||||
MODEL_VERSION,
|
||||
subfolder="unet",
|
||||
device_map=None,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
|
||||
cml_unet = CoreMLUNetWrapper(
|
||||
apply_attention_implementation(ref_unet.eval(), "SPLIT_EINSUM"),
|
||||
ModelVersion.LCM,
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
scheduler = LCMScheduler.from_pretrained(MODEL_VERSION, subfolder="scheduler")
|
||||
scheduler.set_timesteps(50, get_torch_device(), 50)
|
||||
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
|
||||
|
||||
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, 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,
|
||||
sample_size: tuple[int, int] = (64, 64),
|
||||
controlnet_support: bool = False,
|
||||
lora_paths: list[str] = None,
|
||||
):
|
||||
lora_paths = lora_paths or []
|
||||
coreml_unet, ref_unet = get_unets()
|
||||
|
||||
for lora_path in lora_paths:
|
||||
lora_sd, network_alphas = LoraLoaderMixin.lora_state_dict(lora_path)
|
||||
LoraLoaderMixin.load_lora_into_unet(lora_sd, network_alphas, ref_unet)
|
||||
ref_unet.fuse_lora()
|
||||
|
||||
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.config, batch_size
|
||||
)
|
||||
|
||||
scheduler = get_scheduler()
|
||||
|
||||
sample_inputs = get_sample_input(
|
||||
batch_size, encoder_hidden_states_shape, sample_shape, scheduler
|
||||
)
|
||||
|
||||
if controlnet_support:
|
||||
sample_inputs |= add_cnet_support(sample_shape, ref_unet)
|
||||
|
||||
sample_inputs_spec = get_unet_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"], out_path
|
||||
)
|
||||
|
||||
del traced_unet
|
||||
gc.collect()
|
||||
|
||||
coreml_unet.save(out_path)
|
||||
logger.info(f"Saved unet into {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
h = 512
|
||||
w = 512
|
||||
sample_size = (h // 8, w // 8)
|
||||
batch_size = 4
|
||||
|
||||
cn_support_str = "_cn" if True else ""
|
||||
|
||||
out_name = f"{MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
|
||||
|
||||
out_path = get_out_path("unet", f"{out_name}")
|
||||
if not os.path.exists(out_path):
|
||||
convert(out_path=out_path, sample_size=sample_size, batch_size=batch_size)
|
||||
@@ -0,0 +1,70 @@
|
||||
import os
|
||||
|
||||
from coremltools import ComputeUnit
|
||||
|
||||
from coreml_suite import COREML_NODE
|
||||
from coreml_suite.coreml_model import CoreMLModel
|
||||
|
||||
|
||||
class COREML_CONVERT_LCM(COREML_NODE):
|
||||
"""Converts a LCM model to Core ML."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
||||
"width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
||||
"compute_unit": (
|
||||
[
|
||||
ComputeUnit.CPU_AND_NE.name,
|
||||
ComputeUnit.CPU_AND_GPU.name,
|
||||
ComputeUnit.ALL.name,
|
||||
ComputeUnit.CPU_ONLY.name,
|
||||
],
|
||||
),
|
||||
"controlnet_support": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("COREML_UNET",)
|
||||
RETURN_NAMES = ("coreml_model",)
|
||||
FUNCTION = "convert"
|
||||
|
||||
def convert(self, height, width, batch_size, compute_unit, controlnet_support):
|
||||
"""Converts a LCM model to Core ML.
|
||||
|
||||
Args:
|
||||
height (int): Height of the target image.
|
||||
width (int): Width of the target image.
|
||||
batch_size (int): Batch size.
|
||||
compute_unit (str): Compute unit to use when loading the model.
|
||||
|
||||
Returns:
|
||||
coreml_model: The converted Core ML model.
|
||||
|
||||
The converted model is also saved to "models/unet" directory and
|
||||
can be loaded with the "LCMCoreMLLoaderUNet" node.
|
||||
"""
|
||||
from coreml_suite.lcm import converter as lcm_converter
|
||||
|
||||
h = height
|
||||
w = width
|
||||
sample_size = (h // 8, w // 8)
|
||||
batch_size = batch_size
|
||||
cn_support_str = "_cn" if controlnet_support else ""
|
||||
|
||||
out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
|
||||
|
||||
out_path = lcm_converter.get_out_path("unet", f"{out_name}")
|
||||
|
||||
if not os.path.exists(out_path):
|
||||
lcm_converter.convert(
|
||||
out_path=out_path,
|
||||
sample_size=sample_size,
|
||||
batch_size=batch_size,
|
||||
controlnet_support=controlnet_support,
|
||||
)
|
||||
|
||||
return (CoreMLModel(out_path, compute_unit),)
|
||||
@@ -0,0 +1,98 @@
|
||||
from diffusers import UNet2DConditionModel
|
||||
from diffusers.models.embeddings import TimestepEmbedding
|
||||
|
||||
|
||||
class UNet2DConditionModelLCM(UNet2DConditionModel):
|
||||
def __init__(
|
||||
self,
|
||||
time_cond_proj_dim=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
timestep_input_dim = self.config.block_out_channels[0]
|
||||
time_embed_dim = self.config.block_out_channels[0] * 4
|
||||
|
||||
time_embedding = TimestepEmbedding(
|
||||
timestep_input_dim, time_embed_dim, cond_proj_dim=time_cond_proj_dim
|
||||
)
|
||||
self.time_embedding = time_embedding
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample,
|
||||
timestep,
|
||||
encoder_hidden_states,
|
||||
timestep_cond,
|
||||
*additional_residuals,
|
||||
):
|
||||
# 0. Project (or look-up) time embeddings
|
||||
t_emb = self.time_proj(timestep)
|
||||
emb = self.time_embedding(t_emb, timestep_cond)
|
||||
|
||||
# 1. center input if necessary
|
||||
if self.config.center_input_sample:
|
||||
sample = 2 * sample - 1.0
|
||||
|
||||
# 2. pre-process
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
# 3. down
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if (
|
||||
hasattr(downsample_block, "attentions")
|
||||
and downsample_block.attentions is not None
|
||||
):
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if additional_residuals:
|
||||
new_down_block_res_samples = ()
|
||||
for i, down_block_res_sample in enumerate(down_block_res_samples):
|
||||
down_block_res_sample = down_block_res_sample + additional_residuals[i]
|
||||
new_down_block_res_samples += (down_block_res_sample,)
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# 4. mid
|
||||
sample = self.mid_block(
|
||||
sample, emb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
|
||||
if additional_residuals:
|
||||
sample = sample + additional_residuals[-1]
|
||||
|
||||
# 5. up
|
||||
for upsample_block in self.up_blocks:
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[
|
||||
: -len(upsample_block.resnets)
|
||||
]
|
||||
|
||||
if (
|
||||
hasattr(upsample_block, "attentions")
|
||||
and upsample_block.attentions is not None
|
||||
):
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples
|
||||
)
|
||||
|
||||
# 6. post-process
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
return (sample,)
|
||||
@@ -0,0 +1,8 @@
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class ModelVersion(Enum):
|
||||
SD15 = "sd15"
|
||||
SDXL = "sdxl"
|
||||
SDXL_REFINER = "sdxl_refiner"
|
||||
LCM = "lcm"
|
||||
+31
-51
@@ -4,9 +4,16 @@ 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 nodes import KSampler, LoraLoader, KSamplerAdvanced
|
||||
|
||||
from coreml_suite.models import (
|
||||
@@ -17,26 +24,6 @@ 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):
|
||||
@@ -225,26 +212,24 @@ class CoreMLModelAdapter(COREML_NODE):
|
||||
|
||||
|
||||
class CoreMLConverter(COREML_NODE):
|
||||
"""Converts a Stable Diffusion checkpoint (UNet) to Core ML.
|
||||
|
||||
The model version (SD15 / SDXL / SDXL refiner / LCM) is auto-detected from
|
||||
the checkpoint's architecture, so there is no version dropdown — one node
|
||||
converts every supported family, including full-distill LCM.
|
||||
"""
|
||||
"""Converts a LCM model to Core ML."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
||||
"model_version": (
|
||||
[
|
||||
ModelVersion.SD15.name,
|
||||
ModelVersion.SDXL.name,
|
||||
],
|
||||
),
|
||||
"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": (
|
||||
_discover(
|
||||
"list_attention_impls",
|
||||
["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"],
|
||||
),
|
||||
list(ATTENTION_IMPLEMENTATIONS),
|
||||
),
|
||||
"compute_unit": (
|
||||
[
|
||||
@@ -262,10 +247,7 @@ 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": (
|
||||
_discover("list_quant_modes", ["none", "8", "6", "4"]),
|
||||
{"default": "none"},
|
||||
),
|
||||
"quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}),
|
||||
"lora_params": ("LORA_PARAMS",),
|
||||
},
|
||||
}
|
||||
@@ -277,6 +259,7 @@ class CoreMLConverter(COREML_NODE):
|
||||
def convert(
|
||||
self,
|
||||
ckpt_name,
|
||||
model_version,
|
||||
height,
|
||||
width,
|
||||
batch_size,
|
||||
@@ -286,11 +269,9 @@ class CoreMLConverter(COREML_NODE):
|
||||
quantize_nbits="none",
|
||||
lora_params=None,
|
||||
):
|
||||
"""Converts a checkpoint's UNet to Core ML.
|
||||
"""Converts a LCM model to Core ML.
|
||||
|
||||
Args:
|
||||
ckpt_name (str): Checkpoint to convert; its model version is
|
||||
auto-detected from the weights.
|
||||
height (int): Height of the target image.
|
||||
width (int): Width of the target image.
|
||||
batch_size (int): Batch size.
|
||||
@@ -300,8 +281,10 @@ class CoreMLConverter(COREML_NODE):
|
||||
coreml_model: The converted Core ML model.
|
||||
|
||||
The converted model is also saved to "models/unet" directory and
|
||||
can be loaded with the "Load Core ML UNet" node.
|
||||
can be loaded with the "LCMCoreMLLoaderUNet" node.
|
||||
"""
|
||||
model_version = ModelVersion[model_version]
|
||||
|
||||
lora_params = lora_params or {}
|
||||
lora_params = [(k, v[0]) for k, v in lora_params.items()]
|
||||
lora_params = sorted(lora_params, key=lambda lora: lora[0])
|
||||
@@ -310,16 +293,14 @@ class CoreMLConverter(COREML_NODE):
|
||||
h = height
|
||||
w = width
|
||||
sample_size = (h // 8, w // 8)
|
||||
import coreml_diffusion
|
||||
|
||||
out_name = coreml_diffusion.compose_out_name(
|
||||
out_name = 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=coreml_diffusion.lora_names_from_params(lora_params),
|
||||
lora_names=lora_names_from_params(lora_params),
|
||||
quantize_nbits=quantize_nbits,
|
||||
)
|
||||
|
||||
@@ -330,14 +311,13 @@ class CoreMLConverter(COREML_NODE):
|
||||
logger.info(f"Attention implementation: {attention_implementation}")
|
||||
|
||||
if lora_params:
|
||||
logger.info("LoRAs used:")
|
||||
logger.info(f"LoRAs used:")
|
||||
for lora_param in lora_params:
|
||||
logger.info(f" {lora_param[0]} - strength: {lora_param[1]}")
|
||||
|
||||
# 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")
|
||||
from coreml_suite import converter
|
||||
|
||||
unet_out_path = converter.get_out_path("unet", f"{out_name}")
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
|
||||
config_filename = ckpt_name.split(".")[0] + ".yaml"
|
||||
@@ -345,10 +325,10 @@ class CoreMLConverter(COREML_NODE):
|
||||
if config_path:
|
||||
logger.info(f"Using config file {config_path}")
|
||||
|
||||
coreml_diffusion.convert(
|
||||
ckpt_path,
|
||||
None, # model_version auto-detected from the checkpoint
|
||||
unet_out_path,
|
||||
converter.convert(
|
||||
ckpt_path=ckpt_path,
|
||||
model_version=model_version,
|
||||
unet_out_path=unet_out_path,
|
||||
sample_size=sample_size,
|
||||
batch_size=batch_size,
|
||||
controlnet_support=controlnet_support,
|
||||
|
||||
+9
-16
@@ -1,34 +1,27 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[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.1.2"
|
||||
version = "2.0.2"
|
||||
license = "MIT"
|
||||
requires-python = ">=3.12"
|
||||
requires-python = ">=3.12,<3.13"
|
||||
packages = [{ include = "coreml_suite" }]
|
||||
dependencies = [
|
||||
# torch is provided by the host (ComfyUI) and intentionally left unpinned
|
||||
# here: a hard torch cap would downgrade the host's torch and break its
|
||||
# torchvision/torchaudio ABI. coreml-diffusion pulls torch>=2.7 transitively.
|
||||
# >=0.1.6: model-version auto-detection (convert(model_version=None)) and the
|
||||
# dropped <3.13 Python cap (kept in sync with this package's requires-python).
|
||||
"coreml-diffusion>=0.1.6,<0.2",
|
||||
"torch>=2.7,<2.8",
|
||||
"coremltools>=9,<10",
|
||||
"numpy>=2,<3",
|
||||
"diffusers>=0.30",
|
||||
"peft>=0.13",
|
||||
"omegaconf>=2.3",
|
||||
"transformers>=4.44",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["coreml_suite"]
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "aszc-dev"
|
||||
DisplayName = "ComfyUI-CoreMLSuite"
|
||||
Icon = "https://raw.githubusercontent.com/aszc-dev/ComfyUI-CoreMLSuite/main/assets/snake.png"
|
||||
Icon = ""
|
||||
requires-comfyui = ">=0.3.27"
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
+4
-1
@@ -1,4 +1,7 @@
|
||||
coreml-diffusion>=0.1.4,<0.2
|
||||
torch>=2.7,<2.8
|
||||
coremltools>=9,<10
|
||||
numpy>=2,<3
|
||||
diffusers>=0.30
|
||||
peft>=0.13
|
||||
omegaconf>=2.3
|
||||
transformers>=4.44
|
||||
|
||||
@@ -1,208 +0,0 @@
|
||||
# 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.
|
||||
@@ -107,6 +107,7 @@
|
||||
"10": {
|
||||
"inputs": {
|
||||
"ckpt_name": "dreamshaper_8.safetensors",
|
||||
"model_version": "SD15",
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"batch_size": 1,
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
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)
|
||||
@@ -0,0 +1,138 @@
|
||||
"""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}"
|
||||
)
|
||||
@@ -0,0 +1,197 @@
|
||||
"""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}
|
||||
)
|
||||
@@ -0,0 +1,183 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user