Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5e57024336 | ||
|
|
8d28964831 | ||
|
|
a199e749bc | ||
|
|
a3d555cdc5 | ||
|
|
d90546b6bb | ||
|
|
f054bbe991 | ||
|
|
20ec450f9b | ||
|
|
d0cca3c3f4 | ||
|
|
65a2de2fab | ||
|
|
02b6e8ece3 | ||
|
|
7678a07ed5 | ||
|
|
43b77e8471 | ||
|
|
c96059ff0b | ||
|
|
3224d62342 | ||
|
|
2fb135df03 | ||
|
|
66e83c2f2f | ||
|
|
fb7188e5a2 | ||
|
|
4096466f8c | ||
|
|
b8c263b763 | ||
|
|
56cff2bd91 | ||
|
|
7b3f8fc29e | ||
|
|
adaecd3f66 | ||
|
|
aa60cda09b | ||
|
|
5f7fcd6df3 | ||
|
|
e89cff6d01 | ||
|
|
9f90083126 | ||
|
|
5c774ddc5e | ||
|
|
b8197c21ef | ||
|
|
0c78803b25 | ||
|
|
763ca3961b | ||
|
|
ae9a9874c5 | ||
|
|
67c902f761 | ||
|
|
bb44b4a35f | ||
|
|
46d1124573 | ||
|
|
ead01c08dd | ||
|
|
b10effc7c2 | ||
|
|
b1d2e82677 | ||
|
|
9f650acb79 | ||
|
|
c6d6917827 | ||
|
|
63377ebd73 | ||
|
|
42ff10cd43 | ||
|
|
da3a8e13d3 | ||
|
|
8092a19173 | ||
|
|
5477e3d71a | ||
|
|
a8d2d6ec46 | ||
|
|
44cffbb8b8 | ||
|
|
6907d4910f | ||
|
|
1930be5c98 | ||
|
|
45be6761d1 | ||
|
|
fc1132a5d5 | ||
|
|
e440f725a4 | ||
|
|
f9f25fbeb7 | ||
|
|
4a1359b6b5 | ||
|
|
971e60aa09 | ||
|
|
8bcdeab234 | ||
|
|
c9e403b1d8 | ||
|
|
7492f0b486 | ||
|
|
c09221945d | ||
|
|
6864c233e3 | ||
|
|
6ccf41e5c9 | ||
|
|
73aa2d11d3 | ||
|
|
4c438e1ee6 | ||
|
|
fa0735746c | ||
|
|
c26099b334 | ||
|
|
27f1a19131 | ||
|
|
701443f59e | ||
|
|
6d095a67a2 | ||
|
|
bb73e686a0 | ||
|
|
967ab7f269 | ||
|
|
c51d9041a4 | ||
|
|
eeae4bd6e3 | ||
|
|
1ebd9e72ae | ||
|
|
c01c60e3c1 | ||
|
|
44a380ffdf | ||
|
|
e22d8187cd | ||
|
|
1937f39cca | ||
|
|
b90591dfd4 | ||
|
|
1aa5a19b2a | ||
|
|
8a814b7a56 | ||
|
|
213088241d | ||
|
|
9d509ad8f4 | ||
|
|
0092ad5e75 | ||
|
|
99a0a9996d | ||
|
|
db0aea3d9c | ||
|
|
dfdc1bf520 | ||
|
|
d63df5b62f | ||
|
|
901ea6da16 | ||
|
|
dd438f66cc | ||
|
|
41797203d7 | ||
|
|
4d83603c98 | ||
|
|
8a3e9332e1 | ||
|
|
6319d2aedb | ||
|
|
d0629b4efc | ||
|
|
c043e1f9aa | ||
|
|
133f943472 |
@@ -0,0 +1,25 @@
|
||||
name: Publish to Comfy registry
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "pyproject.toml"
|
||||
|
||||
permissions:
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
publish-node:
|
||||
name: Publish Custom Node to registry
|
||||
runs-on: ubuntu-latest
|
||||
if: ${{ github.repository_owner == 'aszc-dev' }}
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
with:
|
||||
## Add your own personal access token to your Github Repository secrets and reference it here.
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
@@ -0,0 +1,27 @@
|
||||
name: Tier 0 — Unit (Linux)
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [main]
|
||||
pull_request:
|
||||
|
||||
# Deps are resolved from pyproject.toml via uv, so the toolchain pins live in
|
||||
# one place. Tier 0 must run without ComfyUI; the in-tree purity gate
|
||||
# (tests/unit/test_tier0_purity.py) enforces that the suite hasn't started
|
||||
# leaking framework imports.
|
||||
jobs:
|
||||
unit:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: astral-sh/setup-uv@v7
|
||||
with:
|
||||
enable-cache: true
|
||||
|
||||
- name: uv sync
|
||||
run: uv sync --no-install-project
|
||||
|
||||
- name: Run Tier 0
|
||||
run: uv run pytest -m unit tests/ -v
|
||||
@@ -0,0 +1,134 @@
|
||||
name: Tier 2 — M2 / ANE (self-hosted)
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
# `labeled` fires when run-m2 is first added; `synchronize`/`reopened`
|
||||
# re-run on every subsequent push while the label is present, so the
|
||||
# result tracks the PR head instead of going stale. The `if` below keeps
|
||||
# the run gated on the run-m2 label for all pull_request events.
|
||||
types: [labeled, synchronize, reopened]
|
||||
schedule:
|
||||
# Nightly at 04:00 UTC (~05/06 in PL). Keeps the M2 path honest
|
||||
# without burning the runner on every PR.
|
||||
- cron: "0 4 * * *"
|
||||
workflow_dispatch:
|
||||
|
||||
jobs:
|
||||
m2:
|
||||
if: |
|
||||
github.event_name == 'schedule' ||
|
||||
github.event_name == 'workflow_dispatch' ||
|
||||
(github.event_name == 'pull_request' &&
|
||||
contains(github.event.pull_request.labels.*.name, 'run-m2'))
|
||||
# Self-hosted Apple Silicon runner. Prerequisites: COMFY_DIR pointing at
|
||||
# a runner-owned ComfyUI clone, plus a cached SD1.5 checkpoint.
|
||||
runs-on: [self-hosted, macOS, ARM64, coreml]
|
||||
timeout-minutes: 90
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
# Hybrid ComfyUI strategy:
|
||||
# - schedule (nightly) -> latest origin/master + ComfyUI's own
|
||||
# requirements.txt (constrained). Canary for upstream API breakage.
|
||||
# - PR label / dispatch -> the requires-comfyui version tag + the frozen
|
||||
# `comfy` uv group. Reproducible merge gate, immune to overnight drift.
|
||||
- name: Resolve ComfyUI ref + mode
|
||||
run: |
|
||||
if [ "$GITHUB_EVENT_NAME" = "schedule" ]; then
|
||||
echo "COMFY_MODE=latest" >> "$GITHUB_ENV"
|
||||
echo "COMFY_REF=master" >> "$GITHUB_ENV"
|
||||
else
|
||||
# requires-comfyui is a semver constraint (e.g. ">=0.3.27"); pin the
|
||||
# gate to the matching ComfyUI release tag (vX.Y.Z).
|
||||
VERSION="$(sed -nE 's/^requires-comfyui *= *"[^0-9]*([0-9]+\.[0-9]+\.[0-9]+).*/\1/p' pyproject.toml)"
|
||||
if [ -z "$VERSION" ]; then echo "could not parse requires-comfyui from pyproject.toml"; exit 1; fi
|
||||
echo "COMFY_MODE=pinned" >> "$GITHUB_ENV"
|
||||
echo "COMFY_REF=v$VERSION" >> "$GITHUB_ENV"
|
||||
fi
|
||||
|
||||
- name: Set up ComfyUI checkout
|
||||
# COMFY_DIR is exported by the self-hosted runner's .env and MUST be a
|
||||
# runner-owned ComfyUI clone (never your dev checkout — this step does
|
||||
# git reset --hard and rewrites custom_nodes). Cloned on first run.
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ -z "${COMFY_DIR:-}" ]; then echo "COMFY_DIR unset"; exit 1; fi
|
||||
# Init-in-place rather than `git clone`: COMFY_DIR may already hold the
|
||||
# cached checkpoint (models/checkpoints) or converted .mlmodelc, and
|
||||
# `git clone` refuses a non-empty target. init + fetch + `checkout -f`
|
||||
# populates the ComfyUI tree while leaving untracked files (the
|
||||
# checkpoint, the cached models) untouched — so setup order is free.
|
||||
if [ ! -d "$COMFY_DIR/.git" ]; then
|
||||
echo "initialising ComfyUI repo in $COMFY_DIR"
|
||||
mkdir -p "$COMFY_DIR"
|
||||
git -C "$COMFY_DIR" init -q
|
||||
fi
|
||||
git -C "$COMFY_DIR" remote get-url origin >/dev/null 2>&1 \
|
||||
|| git -C "$COMFY_DIR" remote add origin https://github.com/comfyanonymous/ComfyUI.git
|
||||
git -C "$COMFY_DIR" fetch --quiet origin
|
||||
if [ "$COMFY_MODE" = "latest" ]; then
|
||||
git -C "$COMFY_DIR" checkout -f -B master origin/master
|
||||
else
|
||||
git -C "$COMFY_DIR" checkout -f "$COMFY_REF"
|
||||
fi
|
||||
COMFY_SHA="$(git -C "$COMFY_DIR" rev-parse HEAD)"
|
||||
echo "COMFY_SHA=$COMFY_SHA" >> "$GITHUB_ENV"
|
||||
echo "Tier 2 mode=$COMFY_MODE, ComfyUI \`$COMFY_SHA\`" >> "$GITHUB_STEP_SUMMARY"
|
||||
|
||||
# Point ComfyUI's custom-node loader at this checkout. Refresh the
|
||||
# symlink only; refuse to clobber a real directory (guards against a
|
||||
# COMFY_DIR that is accidentally a dev checkout).
|
||||
NODE_LINK="$COMFY_DIR/custom_nodes/ComfyUI-CoreMLSuite"
|
||||
if [ -e "$NODE_LINK" ] && [ ! -L "$NODE_LINK" ]; then
|
||||
echo "ERROR: $NODE_LINK is a real directory, not a symlink."
|
||||
echo "COMFY_DIR must be a runner-owned ComfyUI, not your dev checkout."
|
||||
exit 1
|
||||
fi
|
||||
mkdir -p "$COMFY_DIR/custom_nodes"
|
||||
ln -sfn "$GITHUB_WORKSPACE" "$NODE_LINK"
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if [ "$COMFY_MODE" = "latest" ]; then
|
||||
# Node deps (our coremltools-9 toolchain), then ComfyUI's own
|
||||
# requirements for the pulled SHA, capped by the toolchain ceiling.
|
||||
uv sync
|
||||
uv pip install -r "$COMFY_DIR/requirements.txt" \
|
||||
-c constraints/comfy-ceiling.txt
|
||||
else
|
||||
# Pinned gate: the frozen group mirrors the known-good pinned SHA.
|
||||
uv sync --group comfy
|
||||
fi
|
||||
|
||||
- name: Start ComfyUI server (background)
|
||||
run: |
|
||||
cd "$COMFY_DIR"
|
||||
nohup "$GITHUB_WORKSPACE/.venv/bin/python" main.py --port 8188 --cpu-vae > /tmp/comfyui-ci.log 2>&1 &
|
||||
# Poll the HTTP endpoint for readiness — robust to startup-banner
|
||||
# wording / colored-log changes in a floating-latest ComfyUI.
|
||||
for _ in $(seq 1 90); do
|
||||
if curl -sf -o /dev/null http://127.0.0.1:8188/system_stats; then
|
||||
echo "comfy ready (ComfyUI ${COMFY_SHA:-unknown})"; exit 0
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
echo "comfy failed to start"; tail -100 /tmp/comfyui-ci.log; exit 1
|
||||
|
||||
- name: Purge cached Core ML UNets (force fresh conversion)
|
||||
# The converter skips when a model of the same name already exists. That
|
||||
# cache key is conversion *parameters* only, not the conversion code or
|
||||
# toolchain — so a stale model would let a conversion regression pass.
|
||||
# Clear it so every Tier 2 run exercises the full convert -> compile ->
|
||||
# sample path end to end.
|
||||
run: |
|
||||
rm -rf "$COMFY_DIR"/models/unet/*.mlpackage "$COMFY_DIR"/models/unet/*.mlmodelc || true
|
||||
|
||||
- name: Run Tier 2 (m2 marker)
|
||||
# Drives the Core ML Converter node, which converts the UNet from the
|
||||
# checkpoint on every run (cache purged above).
|
||||
run: uv run --no-sync pytest -m m2 tests/ -v
|
||||
|
||||
- name: Stop ComfyUI server
|
||||
if: always()
|
||||
run: pkill -f "main.py.*8188" || true
|
||||
@@ -1,3 +1,6 @@
|
||||
playground/
|
||||
experiments/
|
||||
__pycache__/
|
||||
models/
|
||||
.venv/
|
||||
test_results/
|
||||
.claude/
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
3.12
|
||||
@@ -0,0 +1,618 @@
|
||||
# ComfyUI-CoreMLSuite — Converter Extraction Spec for Claude Code
|
||||
|
||||
> **Companion to `MODERNIZATION_SPEC.md`.** That spec hardens the repo and (Phase 3)
|
||||
> splits the *inference* math from the framework. **This** spec splits the *conversion*
|
||||
> path (`safetensors → CoreML`) out into a standalone, `comfy`-free, pip-installable
|
||||
> package that CoreMLSuite then depends on — and that other projects (incl. on-device
|
||||
> iOS tooling) can reuse.
|
||||
>
|
||||
> **Same discipline as the modernization spec:** safety-net first, behavior-preserving
|
||||
> until told otherwise, one phase = one branch = one PR, `STOP — VALIDATE` gate between
|
||||
> every phase, golden-latent as the regression anchor. `[M2]` = needs macOS/Apple Silicon;
|
||||
> `[M2-ANE]` = needs the Neural Engine. Everything else must run on plain Linux/CI.
|
||||
|
||||
---
|
||||
|
||||
## 0. How to work (read first — non-negotiable)
|
||||
|
||||
1. **Behavior-preserving until Phase E6.** Phases E1–E5 must not change image output, node
|
||||
names, `INPUT_TYPES` field names, or `NODE_CLASS_MAPPINGS` keys. The node graph is the
|
||||
public contract; saved user-workflow JSON breaks if these change.
|
||||
2. **The conversion package produces an artifact and stops there.** Its job ends at a written
|
||||
`.mlpackage` / `.mlmodelc` on disk. It must NOT import `comfy`, `folder_paths`, or
|
||||
`comfy_extras`, and must NOT know ComfyUI's `models/unet` layout. Paths are *inputs*.
|
||||
3. **The runtime loader stays in the suite.** The loader is the **local** `coreml_suite.coreml_model.CoreMLModel`
|
||||
— a thin wrapper over `coremltools.models.MLModel` (NOT Apple's
|
||||
`python_coreml_stable_diffusion.coreml_model.CoreMLModel`, which is no longer used; see #58).
|
||||
It *runs* a compiled model in Python — a desktop/Python inference concern, not a conversion
|
||||
concern. It is NOT moved into the package. (On iOS the `.mlmodelc` is loaded natively; the
|
||||
package's output is the deliverable, not a Python runner.)
|
||||
4. **Decouple in-repo before splitting repos.** Phases E1–E4 create the package *inside this
|
||||
repo* and prove equivalence. The physical second-repo split is Phase E5, only after the
|
||||
golden latent is proven identical. Do not create a second repository before Gate E4 passes.
|
||||
5. **Reuse the existing regression anchor.** The golden latent / PSNR anchor from
|
||||
`MODERNIZATION_SPEC.md` Phase 2 is the cross-cutting proof for every gate here. If it is not
|
||||
yet captured, capture it first (it is a prerequisite for E2 onward).
|
||||
6. **No new runtime dependencies** without flagging in the gate report (name, why, license, size).
|
||||
7. **A failing gate means stop and report**, not work around into the next phase.
|
||||
8. **Tooling is `uv`, not bare `pip`/`venv`.** Every environment/install/lock step uses the
|
||||
project's `uv` toolchain: `uv venv`, `uv pip install`, `uv pip install -e .`, `uv lock`,
|
||||
`uv run pytest`, `uv export`/`uv pip freeze` for baselines. Where this spec says "fresh venv",
|
||||
read "`uv venv` + `uv pip install`". Reserve `uv pip` (not `pip`) inside that venv too.
|
||||
9. **The package is the single source of truth for *what is possible*; the node is a thin,
|
||||
discovery-driven frontend.** See the "Interface contract" pillar below — this is the
|
||||
maintainer's hard requirement and it overrides the earlier (now-rescinded) "freeze the
|
||||
dropdown list" instruction.
|
||||
|
||||
---
|
||||
|
||||
## Interface contract (the maintainer's hard requirement) — read before any phase
|
||||
|
||||
Two coupled guarantees must hold once the package is split out:
|
||||
|
||||
**(A) Updating the converter must NOT require updating CoreMLSuite.**
|
||||
This is satisfied by treating the package's public surface as a versioned contract:
|
||||
- `convert(...)` and `compile_model(...)` are **keyword-only with defaults** for everything
|
||||
past the genuinely-required positionals (`ckpt_path`, `model_version`, `out_path`). New
|
||||
capabilities are added as new keyword args with defaults, so an old Suite's call still
|
||||
validates against a newer package. **Never** reorder or rename existing parameters.
|
||||
- `compose_out_name` (the `.mlpackage` filename = the cache key) **moves into the package** and
|
||||
is versioned with it. The Suite must not carry its own copy; if the package changes the naming
|
||||
scheme that is a **major** bump (old cached artifacts stop resolving).
|
||||
|
||||
**(B) CoreMLSuite must be able to list *new* conversion types WITHOUT a Suite code change or
|
||||
version bump.** Today the node hardcodes its dropdowns:
|
||||
```python
|
||||
"model_version": ([ModelVersion.SD15.name, ModelVersion.SDXL.name],), # hand-typed, also INCOMPLETE (no LCM / SDXL_REFINER)
|
||||
"attention_implementation": (list(ATTENTION_IMPLEMENTATIONS),), # from coreml_suite.attention
|
||||
"quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}), # from coreml_suite.core.naming
|
||||
```
|
||||
These are replaced by **runtime discovery calls into the package**, evaluated inside
|
||||
`INPUT_TYPES` (ComfyUI re-evaluates `INPUT_TYPES` on every plugin load):
|
||||
```python
|
||||
import coreml_diffusion
|
||||
"model_version": (coreml_diffusion.list_model_versions(),),
|
||||
"attention_implementation": (coreml_diffusion.list_attention_impls(),),
|
||||
"quantize_nbits": (coreml_diffusion.list_quant_modes(), {"default": "none"}),
|
||||
```
|
||||
Effect: `uv pip install -U coreml_diffusion` + ComfyUI restart surfaces any newly-added type in the old
|
||||
plugin's dropdown — **no Suite edit, no Suite version bump.** This is the requirement.
|
||||
|
||||
**The cost, stated honestly (accept this trade-off explicitly at Gate E0):**
|
||||
- The Suite becomes a "dumb" frontend; the package is the sole authority on what conversions
|
||||
exist. The Suite can no longer guarantee its saved workflows are valid against *arbitrary*
|
||||
future package versions.
|
||||
- Therefore the package's discovery identifiers (`ModelVersion` values, attn-impl strings, quant
|
||||
modes) are an **ADDITIVE-ONLY contract**: the package may *add* identifiers freely (minor bump,
|
||||
no Suite change); **removing or renaming an identifier is a breaking change requiring a MAJOR
|
||||
bump and a migration note**, because a saved workflow JSON references these strings verbatim.
|
||||
Without this rule, "no version bump" silently becomes "randomly broken workflows."
|
||||
- `INPUT_TYPES` must **fail soft** when the package is missing/old: wrap the discovery calls so a
|
||||
missing `coreml_diffusion` (or an old one lacking a `list_*` function) yields a sane fallback list and a
|
||||
logged warning, instead of the node failing to register and disappearing from the menu.
|
||||
|
||||
**Discovery API the package must expose (stable names):**
|
||||
```python
|
||||
coreml_diffusion.list_model_versions() -> list[str] # VERIFIED ones only, e.g. ["SD15","SDXL"] today (.name — see seam.md)
|
||||
coreml_diffusion.list_attention_impls() -> list[str] # ["SPLIT_EINSUM","SPLIT_EINSUM_V2","ORIGINAL"]
|
||||
coreml_diffusion.list_quant_modes() -> list[str] # ["none","8","6","4"]
|
||||
coreml_diffusion.CONTRACT_VERSION: str # bump rules above; Suite may log/compare it
|
||||
```
|
||||
These return the *display strings already used today*, so existing workflows keep validating.
|
||||
|
||||
**Verification status is a PACKAGE property, not a node hardcode (maintainer's intent).**
|
||||
The Suite wants to expose *every model the converter can verifiably convert*. Today `lcm` and
|
||||
`sdxl_refiner` are absent from the converter node not because the Suite chooses to hide them, but
|
||||
because they lack a full golden/PSNR verification. So the gating lives in the package as a status:
|
||||
```python
|
||||
from enum import Enum
|
||||
class Status(Enum):
|
||||
VERIFIED = "verified" # has a golden anchor + passing [M2-ANE] check
|
||||
EXPERIMENTAL = "experimental" # convertible but not yet anchored/verified
|
||||
|
||||
# internal registry, single source of truth.
|
||||
# KEY by ModelVersion enum MEMBER so list_* can emit .name. Keying by the lowercase
|
||||
# .value string returns ["sd15",...], which the node reverses via ModelVersion[...] -> KeyError.
|
||||
_MODEL_STATUS = {ModelVersion.SD15: Status.VERIFIED, ModelVersion.SDXL: Status.VERIFIED,
|
||||
ModelVersion.SDXL_REFINER: Status.EXPERIMENTAL, ModelVersion.LCM: Status.EXPERIMENTAL}
|
||||
|
||||
def list_model_versions(include_experimental: bool = False) -> list[str]:
|
||||
return [v.name for v, s in _MODEL_STATUS.items() # .name -> "SD15","SDXL"; node reverses with ModelVersion[...]
|
||||
if s is Status.VERIFIED or (include_experimental and s is Status.EXPERIMENTAL)]
|
||||
```
|
||||
Consequence: **promoting a model to VERIFIED in the package expands the Suite's dropdown with no
|
||||
Suite change and no Suite bump** — exactly the requirement. The act of verification (E-LCM
|
||||
produces an LCM golden anchor; same later for refiner) is what flips the status. The Suite's
|
||||
converter node calls `list_model_versions()` (verified-only); a power-user/CLI path may pass
|
||||
`include_experimental=True`. Promotion VERIFIED-from-EXPERIMENTAL is additive (minor bump);
|
||||
demotion or removal is breaking (major bump + note).
|
||||
|
||||
---
|
||||
|
||||
## Naming & layout (chosen — frozen at Gate E0)
|
||||
|
||||
**Distribution name (PyPI):** `coreml-diffusion`. **Import name (Python):** `coreml_diffusion`.
|
||||
(PyPI normalizes `-`/`_`; the distribution uses the hyphen, the importable module the underscore.)
|
||||
Availability checked: both `coreml-diffusion` and the near variants were free on PyPI at E0.
|
||||
|
||||
**Why this name (the positioning it encodes):** the project's niche is *diffusion models on Apple
|
||||
Neural Engine via CoreML, inside ComfyUI and on-device* — **not** Stable Diffusion specifically.
|
||||
`sd*` was rejected because it falsely narrows scope to SD; `coreml-diffusion` keeps `coreml` on the
|
||||
front for discoverability while `diffusion` honestly states the scope (SD/SDXL/LCM today, Flux and
|
||||
other diffusion architectures later) **without** promising arbitrary non-diffusion torch models,
|
||||
whose tracing/shape/sample-input pipeline differs. The name must not be re-narrowed to SD in
|
||||
future docs. ANE is the *differentiator* (documented in the README), but `coreml` was chosen over
|
||||
`ane` in the name for search discoverability per maintainer decision.
|
||||
|
||||
Target package layout (framework-free — zero `comfy` imports):
|
||||
|
||||
```
|
||||
coreml_diffusion/
|
||||
__init__.py # public API surface (see "Public API" below)
|
||||
model_version.py # ModelVersion enum — the SINGLE source of truth, no comfy
|
||||
attention.py # ATTENTION_IMPLEMENTATIONS tuple (from coreml_suite/attention.py) + apply_attention_implementation
|
||||
pipeline.py # get_pipeline (from_single_file), get_unet (cml UNet from ref unet)
|
||||
unet.py # UNet2DConditionModelLCM (moved from coreml_suite/lcm/unet.py)
|
||||
inputs.py # get_sample_input, lcm_inputs, sdxl_inputs,
|
||||
# get_encoder_hidden_states_shape, get_coreml_inputs, get_inputs_spec
|
||||
controlnet.py # add_cnet_support (conversion-side residual SHAPE calc only)
|
||||
convert.py # convert_unet, convert (orchestration), convert_to_coreml, load_coreml_model
|
||||
compile.py # compile_coreml_model
|
||||
quantize.py # (Phase E6 / MODERNIZATION Phase 6 lands here) palettization 4/6/8-bit
|
||||
cli.py # console entry point: `coreml-diffusion convert ...`
|
||||
pyproject.toml # standalone packaging (at E5)
|
||||
```
|
||||
|
||||
What stays in `coreml_suite/` (the ComfyUI side, thinned):
|
||||
- `nodes.py` — still owns **name-encoding** (`out_name` construction), path resolution via
|
||||
`folder_paths`, the node `INPUT_TYPES`/mappings, and wrapping the result in `CoreMLModel`.
|
||||
- `models.py`, `latents.py`, `controlnet.py` (inference parts), `lcm/utils.py`, `config.py`
|
||||
(inference config build) — untouched by this spec except the import-source of `ModelVersion`.
|
||||
|
||||
### Public API (the contract `coreml_diffusion` exposes)
|
||||
```python
|
||||
from coreml_diffusion import ModelVersion, convert, compile_model, compose_out_name
|
||||
from coreml_diffusion import list_model_versions, list_attention_impls, list_quant_modes, CONTRACT_VERSION
|
||||
|
||||
# Mirror the CURRENT converter.py signature, made keyword-only past the required positionals
|
||||
# and with paths/device injected (no folder_paths, no comfy.model_management):
|
||||
# convert(ckpt_path, model_version, out_path, *,
|
||||
# batch_size=1, sample_size=(64, 64), controlnet_support=False,
|
||||
# lora_weights=None, attn_impl=list_attention_impls()[0], config_path=None,
|
||||
# quantize_nbits="none", device=None) -> None # side effect: writes out_path
|
||||
# (current convert() returns None and writes via convert_unet → coreml_unet.save; keep that,
|
||||
# or change to `return out_path` as a deliberate, documented improvement — pick one at E0.)
|
||||
# compile_model(src_path, out_dir, final_name) -> str # returns compiled .mlmodelc path
|
||||
```
|
||||
Note: `convert` takes an **explicit `out_path`** — no `folder_paths`. `device` is injected
|
||||
(defaults to torch's default device). `compose_out_name` lives here (cache-key contract) and the
|
||||
node imports it from the package. The `list_*` discovery functions back the node's dropdowns.
|
||||
|
||||
---
|
||||
|
||||
## The import chains to cut (root cause inventory) — REVISED against current code
|
||||
|
||||
> **State note (verified):** the code moved on since the original draft. Several chains are
|
||||
> already cut. Re-verify each line by `grep` before acting; do not assume the original draft.
|
||||
|
||||
**Already done (verify, then skip):**
|
||||
- ✅ `converter.py` already imports `from coreml_suite.model_version import ModelVersion`, and
|
||||
`model_version.py` is **clean** (`from enum import Enum` only — zero comfy). The old
|
||||
"converter → config → comfy" chain is **already broken**. `config.py` still imports comfy, but
|
||||
it is **inference-side** (`get_model_config` via `supported_models_base`/`latent_formats`) —
|
||||
*not* on the conversion path. Do **not** treat `config.py` as a converter dependency.
|
||||
- ✅ `converter.py` now uses `diffusers.UNet2DConditionModel.from_single_file` and a local
|
||||
`CoreMLUNetWrapper` (in `coreml_suite/conversion/unet.py`) — it is **no longer** importing the
|
||||
Apple `python_coreml_stable_diffusion.unet.UNet2DConditionModel*` internals on the main path.
|
||||
A `coreml_suite/conversion/` subpackage already exists (`attention`, `shapes`, `trace`, `unet`).
|
||||
- ✅ Name-encoding already extracted to `coreml_suite/core/naming.py` (`compose_out_name`,
|
||||
`lora_names_from_params`, `ATTN_SUFFIX`, `QUANT_NBITS_VALUES`) **with characterization tests**
|
||||
(`tests/unit/test_characterization_out_name.py`). The pure-naming split is done.
|
||||
- ✅ Quantization is **already implemented** in `converter.py` (`quantize_nbits`, k-means
|
||||
`palettize_weights`) and surfaced as an optional node input. Phase E6 is therefore *move*, not
|
||||
*build* (see revised E6).
|
||||
|
||||
**Still to cut (the real remaining work):**
|
||||
1. `coreml_suite/converter.py::get_out_path` → `from folder_paths import get_folder_paths`.
|
||||
Main converter still reaches into ComfyUI's model dir. **Cut: `out_path` is an injected arg;
|
||||
`folder_paths` resolution moves up into the node** (the node already computes `out_name`).
|
||||
2. `coreml_suite/lcm/converter.py` → still has its **own** `from folder_paths import
|
||||
get_folder_paths` (`get_out_path`) and (per original draft) `comfy.model_management`. Verify
|
||||
the current LCM file and cut both: inject `out_path` and `device`.
|
||||
3. Global mutation of the attention impl: confirm where it now lives. Main path appears to route
|
||||
through `coreml_suite/conversion/attention.apply_attention_implementation` (cleaner than the
|
||||
old global), but `lcm/converter.py` may still set a module global at import. **Ensure the
|
||||
package sets attention per-call, never at import time.**
|
||||
4. **Duplication LCM vs main:** `lcm/converter.py` still carries its own copies of
|
||||
`convert_to_coreml`, `load_coreml_model`, `get_out_path`, `get_sample_input` (the LCM variant
|
||||
takes a `scheduler` arg), and hardcodes `SimianLuo/LCM_Dreamshaper_v7`. **Dedupe into the
|
||||
single `coreml_diffusion` implementation;** the HF-hardcode consolidation is the *behavior-changing*
|
||||
part → deferred to optional **E-LCM**, not E1–E5.
|
||||
5. **`compose_out_name` ownership:** currently in `coreml_suite/core/naming.py` and called by the
|
||||
node. Per the Interface-contract pillar it must **move into the package** (it is the cache-key
|
||||
contract) and the node must import it from `coreml_diffusion`, not keep a copy.
|
||||
|
||||
---
|
||||
|
||||
## Phase E0 — Seam decision & inventory (no code change)
|
||||
|
||||
**Objective:** lock the cut line, the interface contract, and naming so later phases don't drift.
|
||||
|
||||
### Tasks
|
||||
1. Produce `docs/extraction/seam.md`: a table of every symbol in `converter.py`,
|
||||
`lcm/converter.py`, `lcm/unet.py`, **plus the already-extracted `conversion/` subpackage
|
||||
(`attention`, `shapes`, `trace`, `unet`) and `core/naming.py`**, classified
|
||||
**CONVERSION → coreml_diffusion** vs **STAYS (comfy/node)**. Note which are already framework-free.
|
||||
2. ~~Confirm the current `python_coreml_stable_diffusion` footprint.~~ **DONE (seam.md §6):
|
||||
footprint is ZERO** — no runtime imports anywhere; only a docstring mention in
|
||||
`core/__init__.py:4`. Main path uses `diffusers` + local `CoreMLUNetWrapper`; the runtime
|
||||
`CoreMLModel` (STAYS in suite) is a local coremltools wrapper, not Apple's. No shape/attn helper
|
||||
comes from Apple (local `conversion/shapes.py`, `conversion/attention.py`).
|
||||
3. **Decide the interface contract concretely (the maintainer's hard requirement):**
|
||||
- Discovery functions `list_model_versions / list_attention_impls / list_quant_modes` live in
|
||||
the package and return today's display strings verbatim. Node `INPUT_TYPES` calls them.
|
||||
- `ModelVersion` values, attn-impl strings, quant modes are **ADDITIVE-ONLY** across package
|
||||
versions; removal/rename = MAJOR bump + migration note. Write this into the package's
|
||||
versioning policy doc now.
|
||||
- `compose_out_name` moves to the package; node imports it (no copy). Confirm the
|
||||
characterization tests in `test_characterization_out_name.py` will be re-pointed, not
|
||||
duplicated.
|
||||
- **Resolve the `model_version` dropdown question (maintainer decided):** the Suite exposes
|
||||
*every model the converter can verifiably convert*. `lcm` and `sdxl_refiner` are absent today
|
||||
only because they lack a golden/PSNR verification — **not** because the node hardcodes a
|
||||
short list. Encode this as a **status registry in the package** (`VERIFIED` vs
|
||||
`EXPERIMENTAL`); `list_model_versions()` returns VERIFIED-only by default. The converter node
|
||||
calls it plainly. Promoting LCM/refiner to VERIFIED (after E-LCM / a refiner anchor) expands
|
||||
the dropdown with **no Suite change**. Do NOT add permanent per-node filtering — the gate is
|
||||
verification status, owned by the package.
|
||||
4. ~~Confirm the `ml-stable-diffusion` git dep is pinned.~~ **N/A — already removed (#58).** Verified:
|
||||
zero `python_coreml_stable_diffusion` imports in the repo; `CoreMLModel` is now a local
|
||||
coremltools wrapper; the dep is absent from `pyproject.toml`/`requirements.txt`. No SHA to pin.
|
||||
|
||||
### STOP — VALIDATE (Gate E0)
|
||||
```
|
||||
## Gate E0 report
|
||||
- seam.md committed: <path>; symbol counts (move / stay / already-framework-free)
|
||||
- python_coreml_stable_diffusion usage (verified by grep): conversion=<list> runtime=<list>
|
||||
- Discovery API signatures frozen: list_model_versions (verified-only) / list_attention_impls / list_quant_modes
|
||||
- Status registry decided: sd15+sdxl=VERIFIED, lcm+sdxl_refiner=EXPERIMENTAL (gated, not hidden)
|
||||
- Additive-only contract policy doc written (incl. promotion=minor, demotion/removal=major): <path>
|
||||
- model_version dropdown: expose all (incl. LCM/REFINER) / filtered per node — DECISION: <...>
|
||||
- compose_out_name move-not-copy confirmed; tests re-point plan: <...>
|
||||
- LCM consolidation deferred to optional E-LCM: YES/NO
|
||||
- ml-stable-diffusion: N/A — already removed (#58), not a dependency (was: pin-or-BLOCKER)
|
||||
- Package name in-repo: coreml_diffusion (final PyPI name deferred to E5)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase E1 — Establish `coreml_diffusion` package + discovery API (mostly verification)
|
||||
|
||||
**Objective:** stand up the package namespace and the discovery surface. Much of the comfy-chain
|
||||
cut is **already done** — this phase mostly *verifies* that and adds the discovery functions.
|
||||
|
||||
### Tasks
|
||||
1. **Verify (don't redo):** `coreml_suite/model_version.py` is already clean (`Enum` only). Confirm
|
||||
`import coreml_suite.model_version` works with **no comfy** (`uv run python -c "..."` in a
|
||||
comfy-free `uv venv`). If true, E1's original "extract ModelVersion" task is already satisfied.
|
||||
2. Create the `coreml_diffusion/` package skeleton with `__init__.py` exporting the **discovery API**
|
||||
backed by the *existing* sources of truth for now (re-export `ModelVersion`, the
|
||||
`ATTENTION_IMPLEMENTATIONS` tuple, and `QUANT_NBITS_VALUES`) so values are byte-identical:
|
||||
```python
|
||||
def list_model_versions(): return [v.name for v in ModelVersion] # .name -> "SD15" (node reverses via ModelVersion[...]; .value KeyErrors)
|
||||
def list_attention_impls(): return list(ATTENTION_IMPLEMENTATIONS)
|
||||
def list_quant_modes(): return list(QUANT_NBITS_VALUES)
|
||||
CONTRACT_VERSION = "1.0"
|
||||
```
|
||||
(At this stage `coreml_diffusion` may live inside the repo and import from `coreml_suite.*`; the
|
||||
physical move of implementation happens in E2. The point of E1 is to freeze the *contract*.)
|
||||
3. **Decided (`.name`):** the node renders `ModelVersion.SD15.name` (`"SD15"`) and reverses the
|
||||
dropdown string via `ModelVersion[model_version]` (name lookup, `nodes.py:286`). Discovery API
|
||||
therefore returns `.name`; `.value` (`"sd15"`) would `KeyError`. Recorded in `seam.md` §5.
|
||||
|
||||
### Acceptance criteria
|
||||
- `uv run python -c "import coreml_diffusion; print(coreml_diffusion.list_model_versions(), coreml_diffusion.list_quant_modes())"`
|
||||
works in a **comfy-free** `uv venv` and prints today's exact strings.
|
||||
- Existing characterization tests pass unchanged.
|
||||
- No node behavior change yet (node still uses its current hardcoded lists in E1).
|
||||
|
||||
### STOP — VALIDATE (Gate E1)
|
||||
```
|
||||
## Gate E1 report
|
||||
- model_version.py confirmed comfy-free (uv, no comfy): PASS/FAIL
|
||||
- coreml_diffusion.list_* returns byte-identical strings to current dropdowns: YES/NO (show values)
|
||||
- .name vs .value decision for model_version discovery: <...>
|
||||
- CONTRACT_VERSION set; additive-only policy linked: <path>
|
||||
- Characterization tests unchanged & green (uv run pytest): YES/NO
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase E2 — Move conversion code into `coreml_diffusion` (in-repo, dedup, behavior-preserving)
|
||||
|
||||
**Objective:** physically relocate the conversion mechanics into the framework-free package,
|
||||
collapsing the two duplicate converters into one, with paths/device injected.
|
||||
|
||||
### Tasks
|
||||
1. Move into `coreml_diffusion/`: `pipeline.py` (`get_pipeline`, `get_unet`), `unet.py`
|
||||
(`UNet2DConditionModelLCM`), `inputs.py` (sample/lcm/sdxl input builders +
|
||||
`get_encoder_hidden_states_shape` + `get_coreml_inputs` + `get_inputs_spec`),
|
||||
`controlnet.py` (`add_cnet_support`), `convert.py` (`convert_unet`, `convert`,
|
||||
`convert_to_coreml`, `load_coreml_model`), `compile.py` (`compile_coreml_model`).
|
||||
2. **Dedupe LCM vs main** (the real remaining duplication): delete `lcm/converter.py`'s copies of
|
||||
`convert_to_coreml` / `load_coreml_model` / `get_out_path` / `get_sample_input` (LCM variant
|
||||
carries a `scheduler` arg — fold that into the shared `get_sample_input` as an optional param)
|
||||
in favor of the single `coreml_diffusion` implementation. The main path's helpers
|
||||
(`get_unet`/`get_encoder_hidden_states_shape`/`get_coreml_inputs`/`convert_unet`/`convert`) and
|
||||
the `conversion/` subpackage (`attention`, `shapes`, `trace`, `unet`) move as-is.
|
||||
3. **Inject paths**: replace `get_out_path`'s `folder_paths` reach-in with an injected `out_path`
|
||||
argument on `convert(...)`; `folder_paths` resolution moves up into the node (which already
|
||||
computes `out_name`). No `folder_paths` import anywhere in `coreml_diffusion`.
|
||||
4. **Inject device** where the LCM path used `comfy.model_management` (verify it still does):
|
||||
`convert(..., device=None)`, default to torch's default device.
|
||||
5. **Attention per-call, never at import:** main path already routes through
|
||||
`conversion/attention.apply_attention_implementation` — keep that. If `lcm/converter.py` still
|
||||
sets any module global at import, remove it; the package sets attention from the `attn_impl`
|
||||
arg inside `convert`.
|
||||
6. **Move `compose_out_name` into the package** (`coreml_diffusion/naming.py`); re-point
|
||||
`test_characterization_out_name.py` imports to `coreml_diffusion.naming` — assertions and values
|
||||
unchanged. The node will import it from the package in E3.
|
||||
7. Leave **thin shims** in `coreml_suite/converter.py` and `coreml_suite/lcm/converter.py` that
|
||||
re-export from `coreml_diffusion`, preserving the old call signatures the nodes use (nodes untouched
|
||||
this phase). Shims map comfy `folder_paths`/device into package args.
|
||||
|
||||
### Acceptance criteria
|
||||
- `uv run pytest -m unit` (Tier 0) imports `coreml_diffusion.*` with **no comfy / no MPS** and is green on Linux.
|
||||
- The dedup leaves exactly one implementation of each previously-duplicated function.
|
||||
- Characterization tests pass unchanged after the `compose_out_name` re-point.
|
||||
- `[M2]` A real SD1.5 conversion via the shim still produces a loadable model.
|
||||
- `[M2-ANE]` **Golden latent identical / within tolerance** to the MODERNIZATION Phase 2 anchor
|
||||
(same seed/prompt) — proves the move + dedup changed nothing.
|
||||
|
||||
### STOP — VALIDATE (Gate E2 — first regression gate)
|
||||
```
|
||||
## Gate E2 report
|
||||
- Tier 0 import of coreml_diffusion without comfy/MPS (uv run): PASS/FAIL
|
||||
- LCM/main duplicated funcs collapsed to one (list old→new): <map>
|
||||
- compose_out_name moved to package; char-tests re-pointed & green: YES/NO
|
||||
- Paths injected (no folder_paths in package): confirmed
|
||||
- Device injected (no comfy.model_management in package): confirmed
|
||||
- Attention set per-call, not at import (both main & lcm): confirmed
|
||||
- [M2-ANE] Golden latent vs Phase-2 anchor: identical / within tol <x> / DIVERGED (STOP)
|
||||
- Node INPUT_TYPES / mappings untouched: confirmed (diff)
|
||||
```
|
||||
**If the golden latent diverged at all, STOP and report — do not continue.**
|
||||
|
||||
---
|
||||
|
||||
## Phase E3 — Thin the nodes onto the package (behavior-preserving)
|
||||
|
||||
**Objective:** remove the shims; have the ComfyUI nodes call `coreml_diffusion` directly, keeping the
|
||||
node contract byte-identical.
|
||||
|
||||
### Tasks
|
||||
1. `CoreMLConverter.convert` (in `coreml_suite/nodes.py`): keep the `folder_paths`-based path
|
||||
resolution **in the node**; import `compose_out_name` from `coreml_diffusion` (not `coreml_suite.core`);
|
||||
call `coreml_diffusion.convert(...)` and `coreml_diffusion.compile_model(...)` directly; wrap the compiled path
|
||||
in `CoreMLModel`.
|
||||
2. **Wire the dropdowns to discovery (the maintainer's hard requirement).** Replace the hardcoded
|
||||
`INPUT_TYPES` lists with fail-soft discovery calls:
|
||||
```python
|
||||
def _discover(fn, fallback):
|
||||
try:
|
||||
import coreml_diffusion
|
||||
return getattr(coreml_diffusion, fn)()
|
||||
except Exception as e: # missing/old package, or import error
|
||||
logger.warning(f"coreml_diffusion.{fn} unavailable ({e}); using fallback {fallback}")
|
||||
return fallback
|
||||
...
|
||||
"model_version": (_discover("list_model_versions", ["SD15", "SDXL"]),),
|
||||
"attention_implementation": (_discover("list_attention_impls", ["SPLIT_EINSUM","SPLIT_EINSUM_V2","ORIGINAL"]),),
|
||||
"quantize_nbits": (_discover("list_quant_modes", ["none","8","6","4"]), {"default": "none"}),
|
||||
```
|
||||
This is what makes "update the package → new types appear in the old node, no Suite bump" true.
|
||||
3. `COREML_CONVERT_LCM` (in `coreml_suite/lcm/nodes.py`): route through `coreml_diffusion` for the shared
|
||||
mechanics. **Keep the existing LCM behavior/HF-hardcode for now** — consolidation is optional E-LCM.
|
||||
4. Delete the now-dead `coreml_suite/converter.py` / `coreml_suite/lcm/converter.py` shims (or
|
||||
reduce to a one-line re-export if anything external imports them — grep first).
|
||||
|
||||
### Acceptance criteria
|
||||
- `NODE_CLASS_MAPPINGS` / `NODE_DISPLAY_NAME_MAPPINGS` keys: **unchanged** (diff `__init__.py`).
|
||||
- Every `INPUT_TYPES` **field name** unchanged. Dropdown **values**: the discovery calls must
|
||||
return **a superset of** today's values, with every previously-present value still present and
|
||||
spelled identically (additive-only). *(This deliberately replaces the original spec's
|
||||
"values must be byte-identical/frozen" criterion — the maintainer requires the list be
|
||||
extensible at runtime. Frozen-field-names + additive-only-values is the new contract.)*
|
||||
- With `coreml_diffusion` **absent**, the node still registers and shows the fallback lists (fail-soft).
|
||||
- `[M2-ANE]` Golden latent still identical to the Phase-2 anchor.
|
||||
- `[M2-ANE]` The committed e2e workflow `tests/integration/...` still passes (PSNR > 25).
|
||||
|
||||
### STOP — VALIDATE (Gate E3)
|
||||
```
|
||||
## Gate E3 report
|
||||
- Node mappings diff: empty (confirmed)
|
||||
- INPUT_TYPES field-names diff: empty (confirmed)
|
||||
- Dropdown values: superset of prior, all prior values still present & identical: YES/NO (show)
|
||||
- Fail-soft with coreml_diffusion absent (node still registers): PASS/FAIL
|
||||
- compose_out_name now imported from coreml_diffusion (no node-side copy): confirmed
|
||||
- [M2-ANE] Golden latent vs anchor: identical / within tol / DIVERGED (STOP)
|
||||
- [M2-ANE] e2e workflow PSNR: <value> (> 25?)
|
||||
- Dead converter shims removed / reduced: <list>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase E4 — Standalone packaging & CLI (still in-repo)
|
||||
|
||||
**Objective:** make `coreml_diffusion` independently installable and usable without ComfyUI, with a CLI
|
||||
suitable for the planned article and for on-device/iOS conversion workflows.
|
||||
|
||||
### Tasks
|
||||
1. Add `coreml_diffusion/pyproject.toml`: name (working `coreml_diffusion`), `requires-python`, dependencies
|
||||
= `coremltools` (pinned to the MODERNIZATION-validated version), `diffusers`, `transformers`,
|
||||
`peft`, `omegaconf`, `numpy`, `torch`. **No `ml-stable-diffusion`** (already removed in #58, see
|
||||
§0.3) and **no comfy**. Suite pins `transformers>=4.44`/`peft>=0.13`/`omegaconf>=2.3` today;
|
||||
grep-confirm each is on the conversion path before listing it. A `[project.scripts]` entry:
|
||||
`coreml-diffusion = "coreml_diffusion.cli:main"`.
|
||||
2. `coreml_diffusion/cli.py`: `coreml-diffusion convert --ckpt PATH --model-version sd15 --out PATH
|
||||
[--height --width --batch-size --attn-impl --controlnet --lora NAME:STRENGTH ... --config PATH]`
|
||||
and `coreml-diffusion compile --src PATH --out-dir DIR --name NAME`. Mirrors `convert()`/`compile_model()`.
|
||||
3. Tier-0 Linux tests for the CLI **arg→call mapping** (mock the heavy `convert`); the real
|
||||
convert remains `[M2]`. Add a `[M2]` smoke test: convert a tiny synthetic UNet end-to-end.
|
||||
4. README for the package: install, CLI usage, "produce a `.mlpackage`/`.mlmodelc` for use in a
|
||||
Swift/iOS app", and the ANE positioning note (low-power, GPU-free, embeddable; SD1.5/SDXL on
|
||||
ANE, **not** a Flux-speed claim).
|
||||
|
||||
### Acceptance criteria
|
||||
- Fresh `python -m venv` + `uv pip install ./coreml-diffusion` (no ComfyUI present) imports and runs
|
||||
`coreml-diffusion --help` and the arg-mapping tests on Linux.
|
||||
- `[M2]` `coreml-diffusion convert` produces a model file identical (golden) to the node path.
|
||||
|
||||
### STOP — VALIDATE (Gate E4)
|
||||
```
|
||||
## Gate E4 report
|
||||
- uv pip install ./coreml-diffusion in comfy-free venv: PASS/FAIL (log)
|
||||
- CLI arg→call tests (Tier 0, Linux): green
|
||||
- [M2] CLI-produced model golden vs node-produced model: identical / DIVERGED
|
||||
- Package deps list (with pinned SHAs/versions + licenses):
|
||||
- New runtime deps vs suite before: <none / list>
|
||||
```
|
||||
**This is the gate that proves the package stands alone. Do not split repos before it passes.**
|
||||
|
||||
---
|
||||
|
||||
## Phase E5 — Physical split into a second repository
|
||||
|
||||
**Objective:** move `coreml_diffusion/` to its own repo; CoreMLSuite depends on it by pinned version.
|
||||
|
||||
### Tasks
|
||||
1. Create the new repo (maintainer action — agent prepares the tree, not the GitHub repo).
|
||||
Choose final distributable name; rename imports if changed (single sweep, recorded).
|
||||
2. CoreMLSuite `pyproject.toml` / `requirements.txt`: replace the conversion-only deps with a
|
||||
pinned dependency on the new package (`coreml_diffusion==<version>` from PyPI, or `git+...@<tag>`
|
||||
until first PyPI release). (There is no `git+...ml-stable-diffusion` line to remove — already
|
||||
gone since #58.)
|
||||
3. ~~Keep `python_coreml_stable_diffusion` for the loader.~~ **Void.** The loader is the local
|
||||
`coreml_suite/coreml_model.py` over `coremltools`; the suite keeps `coremltools` as a direct dep
|
||||
for it. No Apple lib involved.
|
||||
4. Set up the new repo's CI: Tier 0 on Linux (import + arg-mapping + input-shape math),
|
||||
`[M2]`/`[M2-ANE]` on a self-hosted/macOS-ARM runner reusing the golden-latent anchor.
|
||||
5. Versioning: SemVer; first release `0.1.0`. Document the compatibility matrix
|
||||
(coreml_diffusion ↔ coremltools version ↔ diffusers version). No ml-stable-diffusion axis.
|
||||
|
||||
### Acceptance criteria
|
||||
- CoreMLSuite installs in a fresh venv pulling the new package; e2e workflow still passes `[M2-ANE]`.
|
||||
- New repo CI green on Linux (Tier 0) and `[M2-ANE]` golden latent matches the anchor.
|
||||
- No conversion code remains in CoreMLSuite (grep: no `ct.convert`, no `from_single_file`,
|
||||
no `torch.jit.trace`).
|
||||
|
||||
### STOP — VALIDATE (Gate E5)
|
||||
```
|
||||
## Gate E5 report
|
||||
- New repo tree prepared: <path/branch>; final package name: <name>
|
||||
- Suite depends on package by pinned version: <spec>
|
||||
- Suite e2e [M2-ANE] PSNR after split: <value> (> 25?)
|
||||
- Conversion code fully absent from suite: confirmed (grep output)
|
||||
- Compatibility matrix documented: <link>
|
||||
- First release tag: 0.1.0
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase E6 — Quantization travels WITH the conversion code (already implemented → move)
|
||||
|
||||
**Objective:** quantization is **already implemented** (k-means `palettize_weights` in
|
||||
`converter.py`, `quantize_nbits` node input, `_q<bits>` filename suffix, README tradeoff table).
|
||||
There is nothing to *build*. It simply **moves with the conversion code in E2** as part of
|
||||
`convert_unet`. This phase is a checkpoint that it survived the extraction intact, plus exposing
|
||||
it through the CLI.
|
||||
|
||||
### Tasks
|
||||
1. Confirm the palettization block moved cleanly into `coreml_diffusion` (lives in `convert.py` or a
|
||||
`quantize.py` helper called from `convert_unet`). Default `"none"` stays byte-identical.
|
||||
2. Expose via CLI flag `--quantize {none,8,6,4}` (E4 already lists this) and via
|
||||
`list_quant_modes()` discovery (E1/E3).
|
||||
3. The existing README tradeoff table (SD1.5 1×512×512 SPLIT_EINSUM: none/8/6/4 → size/ms/PSNR)
|
||||
moves to the package README. Re-confirm one row `[M2-ANE]` so the article can cite a live number.
|
||||
|
||||
### Acceptance criteria
|
||||
- Default (`none`) output byte-identical to pre-extraction (covered by the E2/E3 golden latent).
|
||||
- `coreml-diffusion convert --quantize 4` produces a `_q4` artifact matching the node's `_q4` artifact `[M2]`.
|
||||
- `list_quant_modes()` drives the node dropdown (no hardcoded copy remains).
|
||||
|
||||
### STOP — VALIDATE (Gate E6)
|
||||
```
|
||||
## Gate E6 report
|
||||
- Palettization relocated into coreml_diffusion, called from convert_unet: confirmed
|
||||
- Default none output identical (golden): YES/NO
|
||||
- [M2] CLI --quantize {8,6,4} artifacts match node artifacts: YES/NO
|
||||
- Tradeoff table in package README with at least one re-confirmed [M2-ANE] row: <link>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Phase E-LCM — FIRST task after the split: clean up LCM + verify → promote (behavior-changing, gated)
|
||||
|
||||
> Promoted from "optional, someday" to **the first thing after E5**, per maintainer intent: the
|
||||
> Suite should expose every verifiably-convertible model, and LCM is the obvious first cleanup.
|
||||
|
||||
Two coupled goals:
|
||||
1. **Consolidate the LCM path.** Make the LCM node use the unified `from_single_file` path in
|
||||
`coreml_diffusion.convert(model_version=LCM, ...)` instead of the hardcoded `SimianLuo/LCM_Dreamshaper_v7`
|
||||
HF download; drop the duplicated LCM helpers (already deduped in E2). **Behavior change** ⇒
|
||||
capture an LCM golden anchor *before* the change, then prove within-tolerance after.
|
||||
2. **Verify → promote.** Once the LCM conversion has a passing `[M2-ANE]` golden anchor, flip
|
||||
`_MODEL_STATUS["lcm"] = Status.VERIFIED` **in the package** (minor bump). The Suite's dropdown
|
||||
gains `lcm` automatically — no Suite change, no Suite bump. This is the end-to-end proof that
|
||||
the discovery contract works as designed.
|
||||
|
||||
Repeat the same recipe for `sdxl_refiner` when it gets an anchor (separate small gate). Do NOT
|
||||
bundle E-LCM into E1–E5; it changes behavior and must stand on its own golden.
|
||||
|
||||
### STOP — VALIDATE (Gate E-LCM)
|
||||
```
|
||||
## Gate E-LCM report
|
||||
- LCM golden anchor captured BEFORE change: <path/hash>
|
||||
- LCM node now uses unified from_single_file path; HF hardcode removed: confirmed
|
||||
- [M2-ANE] LCM golden after change: identical / within tol <x> / DIVERGED (STOP)
|
||||
- Status flipped lcm→VERIFIED in package (minor bump <ver>): confirmed
|
||||
- Suite dropdown now lists lcm with NO Suite code change / NO Suite bump: confirmed (diff empty)
|
||||
- LCM node accepts a checkpoint arg now (documented breaking-ish UI note): <link>
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Article deliverable (after E4)
|
||||
|
||||
Once the CLI exists and stands alone, the "convert a Comfy/A1111 workflow into an on-device iOS
|
||||
app" write-up becomes a clean tutorial: `coreml-diffusion convert` → `.mlmodelc` → load in Swift/CoreML.
|
||||
Frame the niche honestly per the README note above (ANE feasibility & power, not raw Flux speed).
|
||||
|
||||
---
|
||||
|
||||
## Quick reference: extraction gate discipline
|
||||
|
||||
```
|
||||
E0 Seam decision, interface contract, discovery API frozen → Gate E0 (cut line + additive-only policy?)
|
||||
E1 Stand up coreml_diffusion + discovery API (mostly verify) → Gate E1 (list_* byte-identical, comfy-free?)
|
||||
E2 Move conversion code, dedup LCM/main, inject paths/device→ Gate E2 (golden identical? duplicates gone?) ← first regression gate
|
||||
E3 Thin nodes onto package + wire discovery dropdowns → Gate E3 (field-names frozen, values additive, fail-soft, golden identical?)
|
||||
E4 Standalone packaging + CLI (uv) → Gate E4 (uv pip install w/o comfy? CLI golden?) ← proves it stands alone
|
||||
E5 Physical second-repo split → Gate E5 (suite depends on pkg? conversion absent?)
|
||||
E-LCM FIRST post-split: clean up LCM, verify → promote → Gate E-LCM (LCM golden? dropdown gains lcm w/ no Suite bump?)
|
||||
E6 Quantization checkpoint (already built → moved in E2) → Gate E6 (default identical? CLI quant matches?)
|
||||
(refiner) same recipe as E-LCM when an anchor exists → own small gate (promote sdxl_refiner→VERIFIED)
|
||||
```
|
||||
|
||||
**Interface-contract invariants (the maintainer's hard requirement), restated:**
|
||||
- Package API is keyword-only-with-defaults past the required positionals → converter updates
|
||||
don't force Suite updates.
|
||||
- Node dropdowns are discovery-driven (`coreml_diffusion.list_*`) + fail-soft → new conversion types
|
||||
appear in the old plugin with `uv pip install -U coreml_diffusion`, **no Suite code change, no bump**.
|
||||
- Discovery identifiers are **additive-only**; removal/rename = MAJOR bump + migration note.
|
||||
- `compose_out_name` (cache key) lives in the package, single copy.
|
||||
|
||||
|
||||
**Golden rule (inherited): never cross a gate with a failing acceptance criterion.
|
||||
Stop, report, wait. The golden latent is the single source of truth that the extraction
|
||||
changed nothing.**
|
||||
@@ -1,674 +1,21 @@
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 3, 29 June 2007
|
||||
|
||||
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.
|
||||
|
||||
Preamble
|
||||
|
||||
The GNU General Public License is a free, copyleft license for
|
||||
software and other kinds of works.
|
||||
|
||||
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>.
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2023-2026 Adrian Szczepański
|
||||
|
||||
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:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
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.
|
||||
|
||||
@@ -2,14 +2,14 @@
|
||||
|
||||
## Overview
|
||||
|
||||
Welcome! I've developed a set of custom nodes for ComfyUI that allows you to use Core ML models in your ComfyUI
|
||||
workflows.
|
||||
Welcome! In this repository you'll find a set of custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
||||
that allows you to use Core ML models in your ComfyUI workflows.
|
||||
These models are designed to leverage the Apple Neural Engine (ANE) on Apple Silicon (M1/M2) machines,
|
||||
thereby enhancing your workflows and improving performance.
|
||||
|
||||
If you're not sure how to obtain these models, you can download them
|
||||
[here](https://huggingface.co/coreml-community) or convert your own models using
|
||||
[coremltools](https://github.com/apple/ml-stable-diffusion).
|
||||
[here](https://huggingface.co/coreml-community) or convert your own checkpoints
|
||||
directly with the conversion nodes in this suite (see [How to use](#how-to-use)).
|
||||
|
||||
In simple terms, think of Core ML models as a tool that can help your ComfyUI work faster and more efficiently.
|
||||
For instance, during my tests on an M2 Pro 32GB machine,
|
||||
@@ -48,6 +48,8 @@ That's it! You're now ready to start enhancing your ComfyUI workflows with Core
|
||||
- **VAE**: Variational Autoencoder. A model that learns a latent representation of images. It's used as a prior in
|
||||
Stable Diffusion.
|
||||
- **Checkpoint**: A file that contains the weights of a model. It's used to load models in Stable Diffusion.
|
||||
- **LCM**: [Latent Consistency Model](https://latent-consistency-models.github.io/). A type of model designed to
|
||||
generate images with as few steps as possible.
|
||||
|
||||
> [!NOTE]
|
||||
> Note on Compute Units:
|
||||
@@ -64,6 +66,12 @@ These custom nodes come with a host of features, including:
|
||||
- Support for ANE (Apple Neural Engine)
|
||||
- Support for CPU and GPU
|
||||
- Support for `mlmodelc` and `mlpackage` files
|
||||
- Support for SDXL models
|
||||
- Support for LCM models
|
||||
- Support for LoRAs
|
||||
- SD1.5 -> Core ML conversion
|
||||
- SDXL -> Core ML conversion
|
||||
- LCM -> Core ML conversion
|
||||
|
||||
> [!NOTE]
|
||||
> Please note that using Core ML models can take a bit longer to load initially.
|
||||
@@ -73,9 +81,43 @@ These custom nodes come with a host of features, including:
|
||||
> [!NOTE]
|
||||
> This repository will continue to be updated with more nodes and features over time.
|
||||
|
||||
## Conversion & Acknowledgements
|
||||
|
||||
The Core ML conversion pipeline in this repository began as an adaptation of
|
||||
Apple's [ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion),
|
||||
which pioneered running Stable Diffusion on the Apple Neural Engine. The
|
||||
implementation has since diverged and no longer depends on that package:
|
||||
|
||||
- UNet conversion runs natively on `diffusers`' `UNet2DConditionModel`.
|
||||
- The ANE-friendly attention path (`SPLIT_EINSUM`, `SPLIT_EINSUM_V2`) is
|
||||
reimplemented as standalone `diffusers` attention processors.
|
||||
- The toolchain tracks current ComfyUI (NumPy 2, Torch 2.7, coremltools 9,
|
||||
Python 3.12).
|
||||
|
||||
The goal is to keep iterating on these methods independently and to explore
|
||||
support beyond SD1.5.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Breaking change in 2.0.0.** The converted Core ML UNet now takes
|
||||
> `encoder_hidden_states` in the native `diffusers` layout
|
||||
> `(batch, tokens, hidden)` instead of the previous
|
||||
> `(batch, hidden, 1, tokens)`. Core ML models converted with earlier versions
|
||||
> are not compatible with 2.0.0 and must be re-converted.
|
||||
|
||||
## Installation
|
||||
|
||||
The installation process is simple!
|
||||
### Using ComfyUI-Manager
|
||||
|
||||
The easiest way to install the custom nodes is to use the ComfyUI-Manager. You can find the installation instructions
|
||||
[here](https://github.com/ltdrdata/ComfyUI-Manager#installation). Once you've installed the ComfyUI-Manager, you can
|
||||
install the custom nodes by following these steps:
|
||||
|
||||
- Open the ComfyUI-Manager by clicking the `Manager` button in the ComfyUI toolbar.
|
||||
- Click the `Install Custom Nodes` button.
|
||||
- Search for `Core ML` and click the `Install` button.
|
||||
- Restart ComfyUI.
|
||||
|
||||
### Manual Installation
|
||||
|
||||
1. Clone this repository into the custom_nodes directory of your ComfyUI. If you're not sure how to do this, you can
|
||||
download the repository as a zip file and extract it into the same directory.
|
||||
@@ -121,10 +163,6 @@ node is a `coreml_model` object that can be used with the Core ML Sampler.
|
||||
- **Outputs**:
|
||||
- **coreml_model**: A Core ML model that can be used with the Core ML Sampler.
|
||||
|
||||
> [!NOTE]
|
||||
> Some models are designed to support ControlNet. If you're using such a model,
|
||||
> make sure to provide a ControlNet input; otherwise, the model will use random noise as ControlNet input.
|
||||
|
||||
#### Core ML Sampler (`CoreMLSampler`)
|
||||
|
||||

|
||||
@@ -143,6 +181,118 @@ resulting latent as you normally would in your workflow.
|
||||
- **LATENT**: The latent image output by the Core ML model. This can be decoded using a VAE Decoder or used as input
|
||||
to the next node in your workflow.
|
||||
|
||||
#### Checkpoint Converter
|
||||
|
||||

|
||||
|
||||
You can use this node to convert any **SD1.5** based checkpoint to a Core ML model. The converted model is stored in the
|
||||
`models/unet` directory and can be used with the `Core ML UNet Loader`. The conversion parameters are encoded in
|
||||
the node name, so if the model already exists, the node will not convert it again.
|
||||
|
||||
- **Inputs**:
|
||||
- **ckpt_name**: The name of the checkpoint to convert. This should be the name of the checkpoint file stored in the
|
||||
`models/checkpoints` directory.
|
||||
- **model_version**: Whether the model is based on SD1.5 or SDXL.
|
||||
- **height**: The desired height of the image generated by the model. The default is 512. Any positive multiple of 8 is accepted.
|
||||
- **width**: The desired width of the image generated by the model. The default is 512. Any positive multiple of 8 is accepted.
|
||||
- **batch_size**: The batch size of generated images. If you're planning to generate batches of images, you can try
|
||||
increasing this value to speed up the generation process. The default is 1.
|
||||
- **attention_implementation**: The attention implementation used when converting the model. Choose SPLIT_EINSUM or
|
||||
SPLIT_EINSUM_V2 for better ANE support. Choose ORIGINAL for better GPU support.
|
||||
- **compute_unit**: The hardware on which the model should run. This is used only when loading the model and doesn't
|
||||
affect the conversion process.
|
||||
- **controlnet_support**: For the model to support ControlNet, it must be converted with this option set to True.
|
||||
The
|
||||
default is False.
|
||||
- **lora_params** [optional]: Optional LoRA names and weights. If provided, the model will be converted with LoRA(s)
|
||||
baked in. More on loading LoRAs below.
|
||||
- **Outputs**:
|
||||
- **coreml_model**: The converted Core ML model that can be used with Core ML Sampler.
|
||||
|
||||
> [!NOTE]
|
||||
> Some models use a custom config .yaml file. If you're using such a model, you'll need to place the config file in the
|
||||
> `models/configs` directory. The config file should be named the same as the checkpoint file. For example, if the
|
||||
> checkpoint file is named `juggernaut_aftermath.safetensors`, the config file should be
|
||||
> named `juggernaut_aftermath.yaml`.
|
||||
> The config file will be automatically loaded during conversion.
|
||||
|
||||
> [!NOTE]
|
||||
> For now, the converter relies heavilty on the model name to determine the conversion parameters. This means that if
|
||||
> you change the model name, the node will convert the model again. Other than that, if you find the name too long or
|
||||
> confusing, you can change it to anything you want.
|
||||
|
||||
#### LoRA Loader
|
||||
|
||||

|
||||
|
||||
This node allows you to load LoRAs and bake them into a model. Since this is a workaround (as model weights can't be
|
||||
modified
|
||||
after conversion), there are a few caveats to keep in mind:
|
||||
|
||||
- The LoRA weights and _strength_model_ parameter are baked into the model. This means that you can't change them
|
||||
after conversion. This also means that you need to convert the model again if you want to change the LoRA weights.
|
||||
- Loading LoRA affects CLIP, which is not a part of Core ML workflow, so you'll need to load CLIP separately,
|
||||
either using `CLIPLoader` or `CheckpointLoaderSimple`. (See [example workflows](#example-workflows) for more details.)
|
||||
- After conversion, if you want to load the model using `CoreMLUnetLoader`, you'll need to apply the same LoRAs to
|
||||
CLIP manually. (See [example workflows](#example-workflows) for more details.)
|
||||
- The LoRA names are encoded in the model name. This means that if you change the name of the LoRA file,
|
||||
you'll need to change the model name as well, or the node will convert the model again. (Model strength is not
|
||||
encoded, so if you want to change it, you'll need to delete the converted model manually)
|
||||
- _strength_clip_ parameter only affects the CLIP model and is not baked into the converted model. This means that
|
||||
you can change it after conversion.
|
||||
|
||||
- **Inputs**:
|
||||
- **lora_name**: The name of the LoRA to load.
|
||||
- **strength_model**: The strength of the LoRA model.
|
||||
- **strength_clip**: The strength of the LoRA CLIP.
|
||||
- **lora_params** [optional]: Optional output from other LoRA Loaders.
|
||||
- **clip**: The CLIP model to use with the LoRA. This can be either output of the
|
||||
`CLIPLoader`/`CheckpointLoaderSimple` or other LoRA Loaders.
|
||||
- **Outputs**:
|
||||
- **lora_params**: The LoRA parameters that can be passed to the Core ML Converter or other LoRA Loaders.
|
||||
- **CLIP**: The CLIP model with LoRA applied.
|
||||
|
||||
#### LCM Converter
|
||||
|
||||

|
||||
|
||||
This node converts [SimianLuo/LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7) model to Core
|
||||
ML. The converted model is stored in the `models/unet` directory and can be used with the Core ML UNet Loader. The
|
||||
conversion parameteres are encoded in the node name, so if the model already exists, the node will not convert it again.
|
||||
|
||||
- **Inputs**:
|
||||
- **height**: The desired height of the image generated by the model. The default is 512. Must be a multiple of 8.
|
||||
- **width**: The desired width of the image generated by the model. The default is 512. Must be a multiple of 8.
|
||||
- **batch_size**: The batch size of generated images. If you're planning to generate batches of images, you can try
|
||||
increasing this value to speed up the generation process. The default is 1.
|
||||
- **compute_unit**: The hardware on which the model should run. This is used only when loading the model and
|
||||
doesn't affect the conversion process.
|
||||
- **controlnet_support**: For the model to support ControlNet, it must be converted with this option set to True.
|
||||
The default is False.
|
||||
|
||||
> [!NOTE]
|
||||
> The conversion process can take a while, so please be patient.
|
||||
|
||||
> [!NOTE]
|
||||
> When using the LCM model with Core ML Sampler, please set _sampler_name_ to `lcm` and _scheduler_ to `sgm_uniform`.
|
||||
|
||||
#### Core ML Adapter (Experimental) (`CoreMLModelAdapter`)
|
||||
|
||||

|
||||
|
||||
This node allows you to use a Core ML as a standard ComfyUI model. This is an experimental node and may not work with
|
||||
all models and nodes. Please use with caution and pay attention to the expected inputs of the model.
|
||||
|
||||
- **Input**:
|
||||
- **coreml_model**: The Core ML model to use as a ComfyUI model.
|
||||
- **Output**:
|
||||
- **MODEL**: The Core ML model wrapped in a ComfyUI model.
|
||||
|
||||
> [!NOTE]
|
||||
> While this approach allows you to use Core ML models with many ComfyUI nodes (both standard and custom), the
|
||||
> expected inputs of the model will not be checked, which may cause errors. Please make sure to use a model compatible
|
||||
> with the expected parameters.
|
||||
|
||||
### Example Workflows
|
||||
|
||||
> [!NOTE]
|
||||
@@ -157,8 +307,8 @@ can use any CLIP or VAE model as long as it's compatible with Stable Diffusion v
|
||||
|
||||
1. **Loading text encoder (CLIP) and VAE models separately**
|
||||
- This workflow uses CLIP and VAE models available
|
||||
[here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/text_encoder/model.safetensors) and
|
||||
[here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/vae/diffusion_pytorch_model.safetensors).
|
||||
[here](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors) and
|
||||
[here](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/vae/diffusion_pytorch_model.safetensors).
|
||||
Once downloaded, place the models in the`models/clip` and `models/vae` directories respectively.
|
||||
- The Core ML UNet model is available
|
||||
[here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip).
|
||||
@@ -166,7 +316,7 @@ can use any CLIP or VAE model as long as it's compatible with Stable Diffusion v
|
||||

|
||||
2. **Loading text encoder (CLIP) and VAE models from checkpoint file**
|
||||
- This workflow loads the CLIP and VAE models from the checkpoint file available
|
||||
[here](https://huggingface.co/runwayml/stable-diffusion-v1-5/blob/main/v1-5-pruned-emaonly.safetensors).
|
||||
[here](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors).
|
||||
Once downloaded, place the model in the`models/checkpoints` directory.
|
||||
- The Core ML UNet model is available
|
||||
[here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip).
|
||||
@@ -180,19 +330,117 @@ being loaded using the standard ComfyUI nodes. Please refer to
|
||||
the [basic txt2img workflow](#basic-txt2img-with-core-ml-unet-loader) for more details on how to load the CLIP and VAE
|
||||
models.
|
||||
The ControlNet model used in this workflow is available
|
||||
[here](https://huggingface.co/lllyasviel/ControlNet-v1-1/blob/main/control_v11p_sd15_lineart.pth).
|
||||
[here](https://huggingface.co/lllyasviel/control_v11p_sd15_scribble/blob/main/diffusion_pytorch_model.fp16.safetensors).
|
||||
Once downloaded, place the model in the `models/controlnet` directory.
|
||||

|
||||
|
||||
#### Checkpoint conversion
|
||||
|
||||
This workflow uses the Checkpoint Converter to convert the checkpoint file. See
|
||||
[Checkpoint Converter](#checkpoint-converter) description for more details.
|
||||
|
||||

|
||||
|
||||
#### Checkpoint conversion with LoRA
|
||||
|
||||
This workflow uses the Checkpoint Converter to convert the checkpoint file with LoRA. See
|
||||
[LoRA Loader](#lora-loader) description to read more about the caveats of using LoRA.
|
||||
|
||||

|
||||
|
||||
#### LCM LoRA conversion
|
||||
|
||||
Please note that you can use multiple LoRAs with the same model. To do this, you'll need to use multiple LoRA Loaders.
|
||||
> [!IMPORTANT]
|
||||
> In this example, the model is passed through the adapter and `ModelSamplingDiscrete` nodes to a standard ComfyUI's
|
||||
> KSampler (not Core ML Sampler). ModelSamplingDiscrete needs to be used to sample models with LCM LoRAs properly.
|
||||
|
||||

|
||||
|
||||
#### Loader with LoRAs
|
||||
|
||||
This workflow uses the Core ML UNet Loader to load a model with LoRAs. The CLIP must be loaded separately and passed
|
||||
through the same LoRA nodes as during conversion. See [LoRA Loader](#lora-loader) description to read more about the
|
||||
caveats of using LoRA. Since _lora_name_ and _strength_model_ are baked into the model, it is not necessary to pass
|
||||
them as inputs to the loader.
|
||||
> [!IMPORTANT]
|
||||
> In this example, the model is passed through the adapter and `ModelSamplingDiscrete` nodes to a standard ComfyUI's
|
||||
> KSampler (not Core ML Sampler). ModelSamplingDiscrete needs to be used to sample models with LCM LoRAs properly.
|
||||
|
||||

|
||||
|
||||
#### LCM conversion with ControlNet
|
||||
|
||||
This workflow uses LCM converter to
|
||||
convert [SimianLuo/LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7)
|
||||
model to Core ML. The converted model can then be used with or without ControlNet to generate images.
|
||||

|
||||
|
||||
#### SDXL Base + Refiner conversion
|
||||
|
||||
This is a basic workflow for SDXL. You add LoRAs and ControlNets the same way as in the previous examples.
|
||||
You can also skip the refiner step.
|
||||
|
||||
The models used in this workflow are available at the following links:
|
||||
|
||||
- [Base model + text_encoder (clip) + text_encoder_2 (clip2)](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
|
||||
- [Refiner model](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0)
|
||||
- [VAE](https://huggingface.co/stabilityai/sdxl-vae)
|
||||
|
||||
> [!IMPORTANT]
|
||||
> SDXL on ANE is not supported. If loading of the model gets stuck, please try using CPU_AND_GPU or CPU_ONLY.
|
||||
> For best results, use ORIGINAL attention implementation.
|
||||
|
||||

|
||||
|
||||
## Quantization (opt-in)
|
||||
|
||||
The `Core ML Converter` and `Core ML LCM Converter` nodes accept an
|
||||
optional `quantize_nbits` dropdown that runs k-means weight palettization
|
||||
(`coremltools.optimize.coreml.palettize_weights`) on the UNet before save.
|
||||
|
||||
Values: `none` (default — no quantization, identical to unquantized
|
||||
behavior and filenames), `8`, `6`, `4`. The number is appended to the
|
||||
.mlpackage stem as `_q<bits>` so quantized and unquantized variants
|
||||
coexist on disk and in cache.
|
||||
|
||||
### SD1.5 1×512×512 SPLIT_EINSUM tradeoffs (M2 Pro, ANE)
|
||||
|
||||
Measured with 20 UNet forward passes at a fixed seed for the PSNR
|
||||
comparison:
|
||||
|
||||
| nbits | size (MB) | size vs none | fwd median (ms) | PSNR vs `none` (dB) |
|
||||
|---|---:|---:|---:|---:|
|
||||
| none | 1641 | 1.000 | 197.1 | — |
|
||||
| 8 | 822 | 0.501 | 186.6 | 53.5 |
|
||||
| 6 | 617 | 0.376 | 183.0 | 40.2 |
|
||||
| 4 | 412 | 0.251 | 179.8 | 27.5 |
|
||||
|
||||
PSNR here is computed on the raw `noise_pred` output of a single UNet
|
||||
forward at a fixed seed, not on the final decoded image — it isolates
|
||||
the quantization-induced drift from sampler / VAE noise. Final-image
|
||||
PSNR is comfortably higher (the sampler averages over 20 steps).
|
||||
|
||||
### Recommended settings per chip / RAM
|
||||
|
||||
- **8 GB RAM (M1 base, M2 base):** `nbits=4`. ~4× smaller model, still
|
||||
loads, PSNR 27 dB is visually identical at SD1.5 sizes.
|
||||
- **16 GB RAM (M1/M2/M3 Pro):** `nbits=6` is the sweet spot — ~2.7×
|
||||
smaller, PSNR 40 dB, no perceptible quality drop.
|
||||
- **32 GB+ RAM (Max / Ultra):** `nbits=8` if you want the safety
|
||||
margin, `none` if you want bit-identical output for golden testing.
|
||||
|
||||
The default stays `none` so existing workflows produce byte-for-byte
|
||||
identical output.
|
||||
|
||||
## Limitations
|
||||
|
||||
- Core ML models are fixed in terms of their inputs and outputs.
|
||||
This means you'll need to use latent images of the same size as the input of the model (512x512 is the default for
|
||||
SD1.5).
|
||||
However, you can convert the model to a different input size using tools available
|
||||
in the [apple/ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion) repository.
|
||||
- For now, only Stable Diffusion v1.5 is supported.
|
||||
- LoRA is not supported yet.
|
||||
However, you can re-convert the model to a different input size using the
|
||||
conversion nodes in this suite (set the desired width and height).
|
||||
- SD2.1 models are not supported.
|
||||
|
||||
[^1]:
|
||||
Unless [EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes)
|
||||
|
||||
@@ -3,13 +3,28 @@ import sys
|
||||
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
|
||||
from coreml_suite import CoreMLLoaderUNet, CoreMLSampler
|
||||
from coreml_suite.nodes import (
|
||||
CoreMLLoaderUNet,
|
||||
CoreMLSampler,
|
||||
CoreMLSamplerAdvanced,
|
||||
CoreMLModelAdapter,
|
||||
CoreMLConverter,
|
||||
COREML_LOAD_LORA,
|
||||
)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"CoreMLUNetLoader": CoreMLLoaderUNet,
|
||||
"CoreMLSampler": CoreMLSampler,
|
||||
"CoreMLSamplerAdvanced": CoreMLSamplerAdvanced,
|
||||
"CoreMLModelAdapter": CoreMLModelAdapter,
|
||||
"Core ML LoRA Loader": COREML_LOAD_LORA,
|
||||
"Core ML Converter": CoreMLConverter,
|
||||
}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"CoreMLUNetLoader": "Load Core ML UNet",
|
||||
"CoreMLSampler": "Core ML Sampler",
|
||||
"CoreMLSamplerAdvanced": "Core ML Sampler (Advanced)",
|
||||
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
|
||||
"Core ML LoRA Loader": "Load LoRA to use with Core ML",
|
||||
"Core ML Converter": "Convert Checkpoint to Core ML",
|
||||
}
|
||||
|
||||
|
After Width: | Height: | Size: 34 KiB |
|
After Width: | Height: | Size: 387 KiB |
|
After Width: | Height: | Size: 94 KiB |
|
After Width: | Height: | Size: 416 KiB |
|
After Width: | Height: | Size: 462 KiB |
|
After Width: | Height: | Size: 476 KiB |
|
After Width: | Height: | Size: 51 KiB |
|
After Width: | Height: | Size: 474 KiB |
|
After Width: | Height: | Size: 54 KiB |
|
After Width: | Height: | Size: 1.5 MiB |
|
Before Width: | Height: | Size: 469 KiB After Width: | Height: | Size: 508 KiB |
@@ -0,0 +1,4 @@
|
||||
"""Top-level conftest: prevent pytest from importing the repo-root
|
||||
__init__.py (the ComfyUI custom-node entry point pulls in comfy + nodes,
|
||||
which breaks the Tier-0 'no-framework' promise)."""
|
||||
collect_ignore = ["__init__.py"]
|
||||
@@ -0,0 +1,18 @@
|
||||
# Toolchain ceiling for installing a floating-latest ComfyUI's requirements.txt
|
||||
# in the Tier 2 nightly canary (.github/workflows/tier2.yml, latest mode).
|
||||
#
|
||||
# ComfyUI's requirements.txt requests bare `torch`/`torchvision`/`torchaudio`
|
||||
# and `numpy>=1.25.0`, which would float past the versions coremltools 9 /
|
||||
# apple-ml-stable-diffusion have been validated against.
|
||||
# These constraints cap the resolution so the canary keeps testing the same
|
||||
# toolchain the suite actually ships.
|
||||
#
|
||||
# If upstream ComfyUI ever hard-requires something beyond these bounds, the
|
||||
# install FAILS — and that failure is the signal we want: it means the host
|
||||
# outgrew the pinned toolchain and coremltools / ml-stable-diffusion need a
|
||||
# deliberate bump, not a silent float.
|
||||
torch>=2.7,<2.8
|
||||
torchvision>=0.22,<0.23
|
||||
torchaudio>=2.7,<2.8
|
||||
numpy>=1.25,<2
|
||||
coremltools>=9,<10
|
||||
@@ -1,4 +1,2 @@
|
||||
from coreml_suite.loaders import CoreMLLoaderUNet
|
||||
from coreml_suite.samplers import CoreMLSampler
|
||||
|
||||
__all__ = ["CoreMLLoaderUNet", "CoreMLSampler"]
|
||||
class COREML_NODE:
|
||||
CATEGORY = "Core ML Suite"
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
import torch
|
||||
|
||||
from comfy import supported_models_base
|
||||
from comfy import latent_formats
|
||||
from comfy.model_detection import convert_config
|
||||
|
||||
from coreml_diffusion import ModelVersion
|
||||
|
||||
|
||||
config_map = {
|
||||
ModelVersion.SD15: {
|
||||
"use_checkpoint": False,
|
||||
"image_size": 32,
|
||||
"out_channels": 4,
|
||||
"use_spatial_transformer": True,
|
||||
"legacy": False,
|
||||
"adm_in_channels": None,
|
||||
"dtype": torch.float16,
|
||||
"in_channels": 4,
|
||||
"model_channels": 320,
|
||||
"num_res_blocks": 2,
|
||||
"attention_resolutions": [1, 2, 4],
|
||||
"transformer_depth": [1, 1, 1, 0],
|
||||
"channel_mult": [1, 2, 4, 4],
|
||||
"transformer_depth_middle": 1,
|
||||
"use_linear_in_transformer": False,
|
||||
"context_dim": 768,
|
||||
"num_heads": 8,
|
||||
"disable_unet_model_creation": True,
|
||||
},
|
||||
ModelVersion.SDXL: {
|
||||
"use_checkpoint": False,
|
||||
"image_size": 32,
|
||||
"out_channels": 4,
|
||||
"use_spatial_transformer": True,
|
||||
"legacy": False,
|
||||
"num_classes": "sequential",
|
||||
"adm_in_channels": 2816,
|
||||
"dtype": torch.float16,
|
||||
"in_channels": 4,
|
||||
"model_channels": 320,
|
||||
"num_res_blocks": 2,
|
||||
"attention_resolutions": [2, 4],
|
||||
"transformer_depth": [0, 2, 10],
|
||||
"channel_mult": [1, 2, 4],
|
||||
"transformer_depth_middle": 10,
|
||||
"use_linear_in_transformer": True,
|
||||
"context_dim": 2048,
|
||||
"num_head_channels": 64,
|
||||
"disable_unet_model_creation": True,
|
||||
},
|
||||
ModelVersion.SDXL_REFINER: {
|
||||
"use_checkpoint": False,
|
||||
"image_size": 32,
|
||||
"out_channels": 4,
|
||||
"use_spatial_transformer": True,
|
||||
"legacy": False,
|
||||
"num_classes": "sequential",
|
||||
"adm_in_channels": 2560,
|
||||
"dtype": torch.float16,
|
||||
"in_channels": 4,
|
||||
"model_channels": 384,
|
||||
"num_res_blocks": 2,
|
||||
"attention_resolutions": [2, 4],
|
||||
"transformer_depth": [0, 4, 4, 0],
|
||||
"channel_mult": [1, 2, 4, 4],
|
||||
"transformer_depth_middle": 4,
|
||||
"use_linear_in_transformer": True,
|
||||
"context_dim": 1280,
|
||||
"num_head_channels": 64,
|
||||
"disable_unet_model_creation": True,
|
||||
},
|
||||
}
|
||||
|
||||
latent_format_map = {
|
||||
ModelVersion.SD15: latent_formats.SD15,
|
||||
ModelVersion.SDXL: latent_formats.SDXL,
|
||||
ModelVersion.SDXL_REFINER: latent_formats.SDXL,
|
||||
}
|
||||
|
||||
|
||||
def get_model_config(model_version: ModelVersion):
|
||||
unet_config = convert_config(config_map[model_version])
|
||||
config = supported_models_base.BASE(unet_config)
|
||||
config.latent_format = latent_format_map[model_version]()
|
||||
return config
|
||||
|
||||
|
||||
def unet_config_from_diffusers_unet(state_dict):
|
||||
match = {}
|
||||
attention_resolutions = []
|
||||
|
||||
attn_res = 1
|
||||
for i in range(5):
|
||||
k = "down_blocks.{}.attentions.1.transformer_blocks.0.attn2.to_k.weight".format(
|
||||
i
|
||||
)
|
||||
if k in state_dict:
|
||||
match["context_dim"] = state_dict[k].shape[1]
|
||||
attention_resolutions.append(attn_res)
|
||||
attn_res *= 2
|
||||
|
||||
match["attention_resolutions"] = attention_resolutions
|
||||
|
||||
match["model_channels"] = state_dict["conv_in.weight"].shape[0]
|
||||
match["in_channels"] = state_dict["conv_in.weight"].shape[1]
|
||||
match["adm_in_channels"] = None
|
||||
if "class_embedding.linear_1.weight" in state_dict:
|
||||
match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[
|
||||
1
|
||||
]
|
||||
elif "add_embedding.linear_1.weight" in state_dict:
|
||||
match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]
|
||||
|
||||
print(match)
|
||||
@@ -0,0 +1,14 @@
|
||||
"""Compatibility shim — re-exports from coreml_suite.core.controlnet."""
|
||||
from coreml_suite.core.controlnet import (
|
||||
chunk_control,
|
||||
expand_inputs,
|
||||
extract_residual_kwargs,
|
||||
no_control,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"chunk_control",
|
||||
"expand_inputs",
|
||||
"extract_residual_kwargs",
|
||||
"no_control",
|
||||
]
|
||||
@@ -0,0 +1,10 @@
|
||||
"""Framework-free pure-logic core of ComfyUI-CoreMLSuite.
|
||||
|
||||
Modules under this package must NOT import `comfy`, `coremltools`,
|
||||
`python_coreml_stable_diffusion`, `folder_paths`, `nodes`, or any other
|
||||
ComfyUI / Apple runtime. Only `numpy` and `torch` are allowed.
|
||||
|
||||
The thin adapters in `coreml_suite.{latents,controlnet,models}` keep the
|
||||
old public import paths working so `coreml_suite/nodes.py` and downstream
|
||||
ComfyUI workflows are unchanged.
|
||||
"""
|
||||
@@ -0,0 +1,67 @@
|
||||
"""Pure helpers around the ControlNet residual inputs of the Core ML UNet.
|
||||
|
||||
Re-exported by coreml_suite.controlnet. Characterization tests cover
|
||||
shapes, dtype (fp16), and zero-fill fallback.
|
||||
"""
|
||||
from itertools import chain
|
||||
from math import ceil
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.latents import chunk_batch
|
||||
|
||||
|
||||
def expand_inputs(inputs):
|
||||
expanded = inputs.copy()
|
||||
for k, v in inputs.items():
|
||||
if isinstance(v, np.ndarray):
|
||||
expanded[k] = np.concatenate([v] * 2) if v.shape[0] == 1 else v
|
||||
elif isinstance(v, torch.Tensor):
|
||||
expanded[k] = torch.cat([v] * 2) if v.shape[0] == 1 else v
|
||||
elif isinstance(v, list):
|
||||
expanded[k] = v * 2 if len(v) == 1 else v
|
||||
elif isinstance(v, dict):
|
||||
expand_inputs(v)
|
||||
return expanded
|
||||
|
||||
|
||||
def extract_residual_kwargs(expected_inputs, control):
|
||||
if "additional_residual_0" not in expected_inputs.keys():
|
||||
return {}
|
||||
if control is None:
|
||||
return no_control(expected_inputs)
|
||||
|
||||
residual_kwargs = {
|
||||
"additional_residual_{}".format(i): r.cpu().numpy().astype(np.float16)
|
||||
for i, r in enumerate(chain(control["output"], control["middle"]))
|
||||
}
|
||||
return residual_kwargs
|
||||
|
||||
|
||||
def no_control(expected_inputs):
|
||||
shapes_dict = {
|
||||
k: v["shape"] for k, v in expected_inputs.items() if k.startswith("additional")
|
||||
}
|
||||
residual_kwargs = {
|
||||
k: torch.zeros(*shape).cpu().numpy().astype(dtype=np.float16)
|
||||
for k, shape in shapes_dict.items()
|
||||
}
|
||||
return residual_kwargs
|
||||
|
||||
|
||||
def chunk_control(cn, target_size):
|
||||
if cn is None:
|
||||
return [None] * target_size
|
||||
|
||||
num_chunks = ceil(cn["output"][0].shape[0] / target_size)
|
||||
|
||||
out = [{"output": [], "middle": []} for _ in range(num_chunks)]
|
||||
|
||||
for k, v in cn.items():
|
||||
for i, x in enumerate(v):
|
||||
chunks = chunk_batch(x, (target_size, *x.shape[1:]))
|
||||
for j, chunk in enumerate(chunks):
|
||||
out[j][k].append(chunk)
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Pure transform from torch sampler inputs to Core ML UNet kwargs.
|
||||
|
||||
Characterization tests cover SD1.5 / SDXL base / SDXL refiner / LCM
|
||||
variants and the chunked-batch fan-out.
|
||||
"""
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.controlnet import extract_residual_kwargs, chunk_control
|
||||
from coreml_suite.core.latents import chunk_batch
|
||||
|
||||
|
||||
class CoreMLInputs:
|
||||
def __init__(self, x, t, context, control, **kwargs):
|
||||
self.x = x
|
||||
self.t = t
|
||||
self.context = context
|
||||
self.control = control
|
||||
self.time_ids = kwargs.get("time_ids")
|
||||
self.text_embeds = kwargs.get("text_embeds")
|
||||
self.ts_cond = kwargs.get("timestep_cond")
|
||||
|
||||
def coreml_kwargs(self, expected_inputs):
|
||||
sample = self.x.cpu().numpy().astype(np.float16)
|
||||
|
||||
context = self.context.cpu().numpy().astype(np.float16)
|
||||
|
||||
t = self.t.cpu().numpy().astype(np.float16)
|
||||
|
||||
model_input_kwargs = {
|
||||
"sample": sample,
|
||||
"encoder_hidden_states": context,
|
||||
"timestep": t,
|
||||
}
|
||||
residual_kwargs = extract_residual_kwargs(expected_inputs, self.control)
|
||||
model_input_kwargs |= residual_kwargs
|
||||
|
||||
# LCM
|
||||
if self.ts_cond is not None:
|
||||
model_input_kwargs["timestep_cond"] = (
|
||||
self.ts_cond.cpu().numpy().astype(np.float16)
|
||||
)
|
||||
|
||||
# SDXL
|
||||
if "text_embeds" in expected_inputs:
|
||||
model_input_kwargs["text_embeds"] = (
|
||||
self.text_embeds.cpu().numpy().astype(np.float16)
|
||||
)
|
||||
if "time_ids" in expected_inputs:
|
||||
model_input_kwargs["time_ids"] = (
|
||||
self.time_ids.cpu().numpy().astype(np.float16)
|
||||
)
|
||||
|
||||
return model_input_kwargs
|
||||
|
||||
def chunks(self, expected_inputs):
|
||||
sample_shape = expected_inputs["sample"]["shape"]
|
||||
timestep_shape = expected_inputs["timestep"]["shape"]
|
||||
context_shape = expected_inputs["encoder_hidden_states"]["shape"]
|
||||
|
||||
chunked_x = chunk_batch(self.x, sample_shape)
|
||||
ts = list(torch.full((len(chunked_x), timestep_shape[0]), self.t[0]))
|
||||
chunked_context = chunk_batch(self.context, context_shape)
|
||||
|
||||
chunked_control = [None] * len(chunked_x)
|
||||
if self.control is not None:
|
||||
chunked_control = chunk_control(self.control, sample_shape[0])
|
||||
|
||||
chunked_ts_cond = [None] * len(chunked_x)
|
||||
if self.ts_cond is not None:
|
||||
ts_cond_shape = expected_inputs["timestep_cond"]["shape"]
|
||||
chunked_ts_cond = chunk_batch(self.ts_cond, ts_cond_shape)
|
||||
|
||||
chunked_time_ids = [None] * len(chunked_x)
|
||||
if expected_inputs.get("time_ids") is not None:
|
||||
time_ids_shape = expected_inputs["time_ids"]["shape"]
|
||||
if self.time_ids is None:
|
||||
self.time_ids = torch.zeros(len(chunked_x), *time_ids_shape[1:]).to(
|
||||
self.x.device
|
||||
)
|
||||
chunked_time_ids = chunk_batch(self.time_ids, time_ids_shape)
|
||||
|
||||
chunked_text_embeds = [None] * len(chunked_x)
|
||||
if expected_inputs.get("text_embeds") is not None:
|
||||
text_embeds_shape = expected_inputs["text_embeds"]["shape"]
|
||||
if self.text_embeds is None:
|
||||
self.text_embeds = torch.zeros(
|
||||
len(chunked_x), *text_embeds_shape[1:]
|
||||
).to(self.x.device)
|
||||
chunked_text_embeds = chunk_batch(self.text_embeds, text_embeds_shape)
|
||||
|
||||
return [
|
||||
CoreMLInputs(
|
||||
x,
|
||||
t,
|
||||
context,
|
||||
control,
|
||||
timestep_cond=ts_cond,
|
||||
time_ids=time_ids,
|
||||
text_embeds=text_embeds,
|
||||
)
|
||||
for x, t, context, control, ts_cond, time_ids, text_embeds in zip(
|
||||
chunked_x,
|
||||
ts,
|
||||
chunked_context,
|
||||
chunked_control,
|
||||
chunked_ts_cond,
|
||||
chunked_time_ids,
|
||||
chunked_text_embeds,
|
||||
)
|
||||
]
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Pure batch-chunking helpers for Core ML's fixed-shape UNet inputs.
|
||||
|
||||
Re-exported by coreml_suite.latents. Characterization tests cover the
|
||||
contract (padding-zero regions, truncation in merge_chunks,
|
||||
identity-passthrough when shape already matches).
|
||||
"""
|
||||
import torch
|
||||
|
||||
|
||||
def chunk_batch(input_tensor, target_shape):
|
||||
if input_tensor.shape == target_shape:
|
||||
return [input_tensor]
|
||||
|
||||
batch_size = input_tensor.shape[0]
|
||||
target_batch_size = target_shape[0]
|
||||
|
||||
num_chunks = batch_size // target_batch_size
|
||||
if num_chunks == 0:
|
||||
padding = torch.zeros(target_batch_size - batch_size, *target_shape[1:]).to(
|
||||
input_tensor.device
|
||||
)
|
||||
return [torch.cat((input_tensor, padding), dim=0)]
|
||||
|
||||
mod = batch_size % target_batch_size
|
||||
if mod != 0:
|
||||
chunks = list(torch.chunk(input_tensor[:-mod], num_chunks))
|
||||
padding = torch.zeros(target_batch_size - mod, *target_shape[1:]).to(
|
||||
input_tensor.device
|
||||
)
|
||||
padded = torch.cat((input_tensor[-mod:], padding), dim=0)
|
||||
chunks.append(padded)
|
||||
return chunks
|
||||
|
||||
chunks = list(torch.chunk(input_tensor, num_chunks))
|
||||
return chunks
|
||||
|
||||
|
||||
def merge_chunks(chunks, orig_shape):
|
||||
merged = torch.cat(chunks, dim=0)
|
||||
if merged.shape == orig_shape:
|
||||
return merged
|
||||
return merged[: orig_shape[0]]
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Pure SDXL detection + time_ids/text_embeds assembly.
|
||||
|
||||
The framework-coupled adapter `add_sdxl_model_options` lives in models.py
|
||||
and delegates the math here. Characterization tests cover base (len 6) vs
|
||||
refiner (len 5) and the closure free-vars produced by
|
||||
`sdxl_model_function_wrapper`.
|
||||
"""
|
||||
import torch
|
||||
|
||||
|
||||
def is_sdxl(coreml_model):
|
||||
return (
|
||||
"time_ids" in coreml_model.expected_inputs
|
||||
and "text_embeds" in coreml_model.expected_inputs
|
||||
)
|
||||
|
||||
|
||||
def is_sdxl_base(coreml_model):
|
||||
return (
|
||||
is_sdxl(coreml_model)
|
||||
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 6
|
||||
)
|
||||
|
||||
|
||||
def is_sdxl_refiner(coreml_model):
|
||||
return (
|
||||
is_sdxl(coreml_model)
|
||||
and coreml_model.expected_inputs["time_ids"]["shape"][1] == 5
|
||||
)
|
||||
|
||||
|
||||
def build_sdxl_time_ids(pos_dict, neg_dict, *, is_base: bool, is_refiner: bool):
|
||||
"""Compose the (2, N) time_ids tensor for the SDXL Core ML UNet.
|
||||
|
||||
- base: N=6 -> [h, w, crop_h, crop_w, target_h, target_w]
|
||||
- refiner: N=5 -> [h, w, crop_h, crop_w, aesthetic_score]
|
||||
- neither: N=4 -> [h, w, crop_h, crop_w] (edge case kept for parity)
|
||||
"""
|
||||
pos_time_ids = [
|
||||
pos_dict.get("height", 768),
|
||||
pos_dict.get("width", 768),
|
||||
pos_dict.get("crop_h", 0),
|
||||
pos_dict.get("crop_w", 0),
|
||||
]
|
||||
neg_time_ids = [
|
||||
neg_dict.get("height", 768),
|
||||
neg_dict.get("width", 768),
|
||||
neg_dict.get("crop_h", 0),
|
||||
neg_dict.get("crop_w", 0),
|
||||
]
|
||||
|
||||
if is_base:
|
||||
pos_time_ids += [
|
||||
pos_dict.get("target_height", 768),
|
||||
pos_dict.get("target_width", 768),
|
||||
]
|
||||
neg_time_ids += [
|
||||
neg_dict.get("target_height", 768),
|
||||
neg_dict.get("target_width", 768),
|
||||
]
|
||||
|
||||
if is_refiner:
|
||||
pos_time_ids += [pos_dict.get("aesthetic_score", 6)]
|
||||
neg_time_ids += [neg_dict.get("aesthetic_score", 2.5)]
|
||||
|
||||
return torch.tensor([pos_time_ids, neg_time_ids])
|
||||
|
||||
|
||||
def build_sdxl_text_embeds(pos_pooled, neg_pooled):
|
||||
"""Concat pos then neg along the batch dim. Locked contract."""
|
||||
return torch.cat((pos_pooled, neg_pooled))
|
||||
|
||||
|
||||
def sdxl_model_function_wrapper(time_ids, text_embeds, refiner=False):
|
||||
def wrapper(model_function, params):
|
||||
x = params["input"]
|
||||
t = params["timestep"]
|
||||
c = params["c"]
|
||||
|
||||
context = c.get("c_crossattn")
|
||||
|
||||
if context is None:
|
||||
return torch.zeros_like(x)
|
||||
|
||||
if refiner and context is not None:
|
||||
# converted refiner accepts only g clip
|
||||
c["c_crossattn"] = context[:, :, 768:]
|
||||
|
||||
return model_function(x, t, **c, time_ids=time_ids, text_embeds=text_embeds)
|
||||
|
||||
return wrapper
|
||||
@@ -0,0 +1,42 @@
|
||||
import time
|
||||
|
||||
import coremltools as ct
|
||||
|
||||
from coreml_suite.logger import logger
|
||||
|
||||
|
||||
class CoreMLModel:
|
||||
"""Small runtime wrapper around coremltools.models.MLModel.
|
||||
|
||||
This keeps the inference path independent from apple/ml-stable-diffusion's
|
||||
CoreMLModel wrapper while preserving the contract used by the sampler code:
|
||||
``expected_inputs`` and callable prediction.
|
||||
"""
|
||||
|
||||
def __init__(self, model_path, compute_unit):
|
||||
self.model_path = model_path
|
||||
self.compute_unit = self._compute_unit(compute_unit)
|
||||
|
||||
logger.info(f"Loading {model_path} to {self.compute_unit.name}")
|
||||
start = time.time()
|
||||
self.model = ct.models.MLModel(model_path, compute_units=self.compute_unit)
|
||||
logger.info(f"Loading {model_path} took {time.time() - start:.1f} seconds")
|
||||
|
||||
self.expected_inputs = self._expected_inputs()
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
return self.model.predict(kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _compute_unit(compute_unit):
|
||||
if isinstance(compute_unit, ct.ComputeUnit):
|
||||
return compute_unit
|
||||
return ct.ComputeUnit[compute_unit]
|
||||
|
||||
def _expected_inputs(self):
|
||||
return {
|
||||
feature.name: {
|
||||
"shape": tuple(feature.type.multiArrayType.shape),
|
||||
}
|
||||
for feature in self.model.get_spec().description.input
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
"""Compatibility shim — re-exports from coreml_suite.core.latents."""
|
||||
from coreml_suite.core.latents import chunk_batch, merge_chunks
|
||||
|
||||
__all__ = ["chunk_batch", "merge_chunks"]
|
||||
@@ -0,0 +1,8 @@
|
||||
"""LCM runtime support (sampler-side).
|
||||
|
||||
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``.
|
||||
"""
|
||||
@@ -0,0 +1,73 @@
|
||||
import torch
|
||||
|
||||
from comfy.model_management import get_torch_device
|
||||
from comfy_extras.nodes_model_advanced import ModelSamplingDiscreteDistilled, LCM
|
||||
|
||||
|
||||
def is_lcm(coreml_model):
|
||||
return "timestep_cond" in coreml_model.expected_inputs
|
||||
|
||||
|
||||
def get_w_embedding(w, embedding_dim=512, dtype=torch.float32):
|
||||
assert len(w.shape) == 1
|
||||
w = w * 1000.0
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
|
||||
emb = w.to(dtype)[:, None] * emb[None, :]
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
||||
if embedding_dim % 2 == 1: # zero pad
|
||||
emb = torch.nn.functional.pad(emb, (0, 1))
|
||||
assert emb.shape == (w.shape[0], embedding_dim)
|
||||
return emb
|
||||
|
||||
|
||||
def model_function_wrapper(w_embedding):
|
||||
def wrapper(model_function, params):
|
||||
x = params["input"]
|
||||
t = params["timestep"]
|
||||
c = params["c"]
|
||||
|
||||
context = c.get("c_crossattn")
|
||||
|
||||
if context is None:
|
||||
return torch.zeros_like(x)
|
||||
|
||||
return model_function(x, t, **c, timestep_cond=w_embedding)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def lcm_patch(model):
|
||||
m = model.clone()
|
||||
sampling_type = LCM
|
||||
sampling_base = ModelSamplingDiscreteDistilled
|
||||
|
||||
class ModelSamplingAdvanced(sampling_base, sampling_type):
|
||||
pass
|
||||
|
||||
model_sampling = ModelSamplingAdvanced()
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
|
||||
return m
|
||||
|
||||
|
||||
def add_lcm_model_options(model_patcher, cfg, latent_image):
|
||||
mp = model_patcher.clone()
|
||||
|
||||
latent = latent_image["samples"].to(get_torch_device())
|
||||
batch_size = latent.shape[0]
|
||||
dtype = latent.dtype
|
||||
device = get_torch_device()
|
||||
|
||||
w = torch.tensor(cfg).repeat(batch_size)
|
||||
w_embedding = get_w_embedding(w, embedding_dim=256).to(device=device, dtype=dtype)
|
||||
|
||||
model_options = {
|
||||
"model_function_wrapper": model_function_wrapper(w_embedding),
|
||||
"sampler_cfg_function": lambda x: x["cond"].to(device),
|
||||
}
|
||||
mp.model_options |= model_options
|
||||
|
||||
return mp
|
||||
@@ -1,84 +0,0 @@
|
||||
import os.path
|
||||
|
||||
from coremltools import ComputeUnit
|
||||
from python_coreml_stable_diffusion.coreml_model import CoreMLModel
|
||||
|
||||
import folder_paths
|
||||
|
||||
from coreml_suite.logger import logger
|
||||
|
||||
|
||||
class CoreMLLoader:
|
||||
PACKAGE_DIRNAME = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"coreml_name": (list(s.coreml_filenames().keys()),),
|
||||
"compute_unit": (
|
||||
[
|
||||
ComputeUnit.CPU_AND_NE.name,
|
||||
ComputeUnit.CPU_AND_GPU.name,
|
||||
ComputeUnit.ALL.name,
|
||||
ComputeUnit.CPU_ONLY.name,
|
||||
],
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "Core ML Suite"
|
||||
|
||||
@classmethod
|
||||
def coreml_filenames(cls):
|
||||
extensions = (".mlmodelc", ".mlpackage")
|
||||
all_paths = folder_paths.get_filename_list_(cls.PACKAGE_DIRNAME)[1]
|
||||
coreml_paths = folder_paths.filter_files_extensions(all_paths, extensions)
|
||||
|
||||
return {os.path.split(p)[-1]: p for p in coreml_paths}
|
||||
|
||||
def load(self, coreml_name, compute_unit):
|
||||
logger.info(f"Loading {coreml_name} to {compute_unit}")
|
||||
|
||||
coreml_path = self.coreml_filenames()[coreml_name]
|
||||
|
||||
sources = "compiled" if coreml_name.endswith(".mlmodelc") else "packages"
|
||||
|
||||
return self._load(coreml_path, compute_unit, sources)
|
||||
|
||||
def _load(self, coreml_path, compute_unit, sources):
|
||||
return (CoreMLModel(coreml_path, compute_unit, sources),)
|
||||
|
||||
|
||||
class CoreMLLoaderCkpt(CoreMLLoader):
|
||||
PACKAGE_DIRNAME = "checkpoints"
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
|
||||
def load(self, coreml_name, compute_unit):
|
||||
# TODO: Implement this
|
||||
pass
|
||||
|
||||
|
||||
class CoreMLLoaderTextEncoder(CoreMLLoader):
|
||||
PACKAGE_DIRNAME = "clip"
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
|
||||
def load(self, coreml_name, compute_unit):
|
||||
# TODO: Implement this
|
||||
pass
|
||||
|
||||
|
||||
class CoreMLLoaderUNet(CoreMLLoader):
|
||||
PACKAGE_DIRNAME = "unet"
|
||||
RETURN_TYPES = ("COREML_UNET",)
|
||||
RETURN_NAMES = ("coreml_model",)
|
||||
|
||||
|
||||
class CoreMLLoaderVAE(CoreMLLoader):
|
||||
PACKAGE_DIRNAME = "vae"
|
||||
RETURN_TYPES = ("VAE",)
|
||||
|
||||
def load(self, coreml_name, compute_unit):
|
||||
# TODO: Implement this
|
||||
pass
|
||||
@@ -1,63 +1,147 @@
|
||||
import numpy as np
|
||||
"""Framework-coupled glue between Core ML UNets and ComfyUI's sampler stack.
|
||||
|
||||
Pure math (CoreMLInputs, SDXL detection, time_ids/text_embeds assembly,
|
||||
sdxl_model_function_wrapper) lives in coreml_suite.core.*.
|
||||
This module is what touches comfy.*: model_base, ModelPatcher, the
|
||||
diffusion_model wrapper, and the maintainer-facing add_sdxl_model_options
|
||||
adapter.
|
||||
"""
|
||||
import torch
|
||||
|
||||
from comfy import model_base
|
||||
from comfy.model_management import get_torch_device
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
|
||||
from comfy import supported_models_base
|
||||
from comfy.latent_formats import SD15
|
||||
from comfy.model_base import BaseModel
|
||||
from coreml_suite.config import get_model_config, ModelVersion
|
||||
from coreml_suite.core.inputs import CoreMLInputs
|
||||
from coreml_suite.core.latents import merge_chunks
|
||||
from coreml_suite.core.sdxl import (
|
||||
build_sdxl_text_embeds,
|
||||
build_sdxl_time_ids,
|
||||
is_sdxl,
|
||||
is_sdxl_base,
|
||||
is_sdxl_refiner,
|
||||
sdxl_model_function_wrapper,
|
||||
)
|
||||
from coreml_suite.lcm.utils import is_lcm
|
||||
from coreml_suite.logger import logger
|
||||
|
||||
from coreml_suite.utils import expand_inputs, extract_residual_kwargs
|
||||
__all__ = [
|
||||
"CoreMLInputs",
|
||||
"CoreMLModelWrapper",
|
||||
"CoreMLModelWrapperLCM",
|
||||
"add_sdxl_model_options",
|
||||
"get_latent_image",
|
||||
"get_model_patcher",
|
||||
"is_sdxl",
|
||||
"is_sdxl_base",
|
||||
"is_sdxl_refiner",
|
||||
"sdxl_model_function_wrapper",
|
||||
]
|
||||
|
||||
|
||||
def get_model_config():
|
||||
# TODO: This is a dummy model config, but it should be enough to
|
||||
# get the model to load - implement a proper model config
|
||||
model_config = supported_models_base.BASE({})
|
||||
model_config.latent_format = SD15()
|
||||
model_config.unet_config = {
|
||||
"disable_unet_model_creation": True,
|
||||
"num_res_blocks": 2,
|
||||
"attention_resolutions": [1, 2, 4],
|
||||
"channel_mult": [1, 2, 4, 4],
|
||||
"transformer_depth": [1, 1, 1, 0],
|
||||
class CoreMLModelWrapper:
|
||||
def __init__(self, coreml_model):
|
||||
self.coreml_model = coreml_model
|
||||
self.dtype = torch.float16
|
||||
|
||||
def __call__(self, x, t, context, control, transformer_options=None, **kwargs):
|
||||
inputs = CoreMLInputs(x, t, context, control, **kwargs)
|
||||
input_list = inputs.chunks(self.expected_inputs)
|
||||
|
||||
chunked_out = [
|
||||
self.get_torch_outputs(
|
||||
self.coreml_model(**input_kwargs.coreml_kwargs(self.expected_inputs)),
|
||||
x.device,
|
||||
)
|
||||
for input_kwargs in input_list
|
||||
]
|
||||
merged_out = merge_chunks(chunked_out, x.shape)
|
||||
|
||||
return merged_out
|
||||
|
||||
@staticmethod
|
||||
def get_torch_outputs(model_output, device):
|
||||
return torch.from_numpy(model_output["noise_pred"]).to(device)
|
||||
|
||||
@property
|
||||
def expected_inputs(self):
|
||||
return self.coreml_model.expected_inputs
|
||||
|
||||
@property
|
||||
def is_lcm(self):
|
||||
return is_lcm(self.coreml_model)
|
||||
|
||||
@property
|
||||
def is_sdxl_base(self):
|
||||
return is_sdxl_base(self.coreml_model)
|
||||
|
||||
@property
|
||||
def is_sdxl_refiner(self):
|
||||
return is_sdxl_refiner(self.coreml_model)
|
||||
|
||||
@property
|
||||
def config(self):
|
||||
if self.is_sdxl_base:
|
||||
return get_model_config(ModelVersion.SDXL)
|
||||
|
||||
if self.is_sdxl_refiner:
|
||||
return get_model_config(ModelVersion.SDXL_REFINER)
|
||||
|
||||
return get_model_config(ModelVersion.SD15)
|
||||
|
||||
|
||||
class CoreMLModelWrapperLCM(CoreMLModelWrapper):
|
||||
def __init__(self, coreml_model):
|
||||
super().__init__(coreml_model)
|
||||
self.config = None
|
||||
|
||||
|
||||
def add_sdxl_model_options(model_patcher, positive, negative):
|
||||
mp = model_patcher.clone()
|
||||
|
||||
pos_dict = positive[0][1]
|
||||
neg_dict = negative[0][1]
|
||||
|
||||
is_base = model_patcher.model.diffusion_model.is_sdxl_base
|
||||
is_refiner = model_patcher.model.diffusion_model.is_sdxl_refiner
|
||||
|
||||
time_ids = build_sdxl_time_ids(
|
||||
pos_dict, neg_dict, is_base=is_base, is_refiner=is_refiner
|
||||
)
|
||||
text_embeds = build_sdxl_text_embeds(
|
||||
pos_dict["pooled_output"], neg_dict["pooled_output"]
|
||||
)
|
||||
|
||||
mp.model_options |= {
|
||||
"model_function_wrapper": sdxl_model_function_wrapper(
|
||||
time_ids, text_embeds, is_refiner
|
||||
),
|
||||
}
|
||||
return model_config
|
||||
return mp
|
||||
|
||||
|
||||
class CoreMLModelWrapper(BaseModel):
|
||||
def __init__(self, model_config, coreml_model):
|
||||
super().__init__(model_config)
|
||||
self.diffusion_model = coreml_model
|
||||
def get_latent_image(coreml_model, latent_image):
|
||||
if latent_image is not None:
|
||||
return latent_image
|
||||
|
||||
def apply_model(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
c_concat=None,
|
||||
c_crossattn=None,
|
||||
c_adm=None,
|
||||
control=None,
|
||||
transformer_options={},
|
||||
):
|
||||
sample = x.cpu().numpy().astype(np.float16)
|
||||
logger.warning("No latent image provided, using empty tensor.")
|
||||
expected = coreml_model.expected_inputs["sample"]["shape"]
|
||||
batch_size = max(expected[0] // 2, 1)
|
||||
latent_image = {"samples": torch.zeros(batch_size, *expected[1:])}
|
||||
return latent_image
|
||||
|
||||
context = c_crossattn.cpu().numpy().astype(np.float16)
|
||||
context = context.transpose(0, 2, 1)[:, :, None, :]
|
||||
|
||||
t = t.cpu().numpy().astype(np.float16)
|
||||
def get_model_patcher(coreml_model):
|
||||
wrapped_model = CoreMLModelWrapper(coreml_model)
|
||||
|
||||
model_input_kwargs = {
|
||||
"sample": sample,
|
||||
"encoder_hidden_states": context,
|
||||
"timestep": t,
|
||||
}
|
||||
residual_kwargs = extract_residual_kwargs(self.diffusion_model, control)
|
||||
model_input_kwargs |= residual_kwargs
|
||||
model_input_kwargs = expand_inputs(model_input_kwargs)
|
||||
if wrapped_model.is_sdxl_base:
|
||||
model = model_base.SDXL(wrapped_model.config, device=get_torch_device())
|
||||
elif wrapped_model.is_sdxl_refiner:
|
||||
model = model_base.SDXLRefiner(wrapped_model.config, device=get_torch_device())
|
||||
else:
|
||||
model = model_base.BaseModel(wrapped_model.config, device=get_torch_device())
|
||||
|
||||
np_out = self.diffusion_model(**model_input_kwargs)["noise_pred"]
|
||||
return torch.from_numpy(np_out).to(x.device)
|
||||
|
||||
def get_dtype(self):
|
||||
# Hardcoding torch-compatible dtype (used for memory allocation)
|
||||
return torch.float16
|
||||
model.diffusion_model = wrapped_model
|
||||
model_patcher = ModelPatcher(model, get_torch_device(), None)
|
||||
return model_patcher
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
import os
|
||||
|
||||
from coremltools import ComputeUnit
|
||||
|
||||
import folder_paths
|
||||
from coreml_suite import COREML_NODE
|
||||
from coreml_suite.coreml_model import CoreMLModel
|
||||
from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
|
||||
from coreml_suite.logger import logger
|
||||
from nodes import KSampler, LoraLoader, KSamplerAdvanced
|
||||
|
||||
from coreml_suite.models import (
|
||||
add_sdxl_model_options,
|
||||
is_sdxl,
|
||||
get_model_patcher,
|
||||
get_latent_image,
|
||||
)
|
||||
|
||||
|
||||
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):
|
||||
old_required = KSampler.INPUT_TYPES()["required"].copy()
|
||||
old_required.pop("model")
|
||||
old_required.pop("negative")
|
||||
old_required.pop("latent_image")
|
||||
new_required = {"coreml_model": ("COREML_UNET",)}
|
||||
return {
|
||||
"required": new_required | old_required,
|
||||
"optional": {"negative": ("CONDITIONING",), "latent_image": ("LATENT",)},
|
||||
}
|
||||
|
||||
def sample(
|
||||
self,
|
||||
coreml_model,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
negative=None,
|
||||
latent_image=None,
|
||||
denoise=1.0,
|
||||
):
|
||||
model_patcher = get_model_patcher(coreml_model)
|
||||
latent_image = get_latent_image(coreml_model, latent_image)
|
||||
|
||||
if is_lcm(coreml_model):
|
||||
negative = [[None, {}]]
|
||||
positive[0][1]["control_apply_to_uncond"] = False
|
||||
model_patcher = add_lcm_model_options(model_patcher, cfg, latent_image)
|
||||
model_patcher = lcm_patch(model_patcher)
|
||||
else:
|
||||
assert (
|
||||
negative is not None
|
||||
), "Negative conditioning is optional only for LCM models."
|
||||
|
||||
if is_sdxl(coreml_model):
|
||||
model_patcher = add_sdxl_model_options(model_patcher, positive, negative)
|
||||
|
||||
return super().sample(
|
||||
model_patcher,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
negative,
|
||||
latent_image,
|
||||
denoise,
|
||||
)
|
||||
|
||||
|
||||
class CoreMLSamplerAdvanced(COREML_NODE, KSamplerAdvanced):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
old_required = KSamplerAdvanced.INPUT_TYPES()["required"].copy()
|
||||
old_required.pop("model")
|
||||
old_required.pop("negative")
|
||||
old_required.pop("latent_image")
|
||||
new_required = {"coreml_model": ("COREML_UNET",)}
|
||||
return {
|
||||
"required": new_required | old_required,
|
||||
"optional": {"negative": ("CONDITIONING",), "latent_image": ("LATENT",)},
|
||||
}
|
||||
|
||||
def sample(
|
||||
self,
|
||||
coreml_model,
|
||||
add_noise,
|
||||
noise_seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
start_at_step,
|
||||
end_at_step,
|
||||
return_with_leftover_noise,
|
||||
negative=None,
|
||||
latent_image=None,
|
||||
denoise=1.0,
|
||||
):
|
||||
model_patcher = get_model_patcher(coreml_model)
|
||||
latent_image = get_latent_image(coreml_model, latent_image)
|
||||
|
||||
if is_lcm(coreml_model):
|
||||
negative = [[None, {}]]
|
||||
positive[0][1]["control_apply_to_uncond"] = False
|
||||
model_patcher = add_lcm_model_options(model_patcher, cfg, latent_image)
|
||||
model_patcher = lcm_patch(model_patcher)
|
||||
else:
|
||||
assert (
|
||||
negative is not None
|
||||
), "Negative conditioning is optional only for LCM models."
|
||||
|
||||
if is_sdxl(coreml_model):
|
||||
model_patcher = add_sdxl_model_options(model_patcher, positive, negative)
|
||||
|
||||
return super().sample(
|
||||
model_patcher,
|
||||
add_noise,
|
||||
noise_seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
negative,
|
||||
latent_image,
|
||||
start_at_step,
|
||||
end_at_step,
|
||||
return_with_leftover_noise,
|
||||
denoise,
|
||||
)
|
||||
|
||||
|
||||
class CoreMLLoader(COREML_NODE):
|
||||
PACKAGE_DIRNAME = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"coreml_name": (list(s.coreml_filenames().keys()),),
|
||||
"compute_unit": (
|
||||
[
|
||||
ComputeUnit.CPU_AND_NE.name,
|
||||
ComputeUnit.CPU_AND_GPU.name,
|
||||
ComputeUnit.ALL.name,
|
||||
ComputeUnit.CPU_ONLY.name,
|
||||
],
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "load"
|
||||
|
||||
@classmethod
|
||||
def coreml_filenames(cls):
|
||||
extensions = (".mlpackage",)
|
||||
all_paths = folder_paths.get_filename_list_(cls.PACKAGE_DIRNAME)[1]
|
||||
coreml_paths = folder_paths.filter_files_extensions(all_paths, extensions)
|
||||
|
||||
return {os.path.split(p)[-1]: p for p in coreml_paths}
|
||||
|
||||
def load(self, coreml_name, compute_unit):
|
||||
logger.info(f"Loading {coreml_name} to {compute_unit}")
|
||||
|
||||
coreml_path = self.coreml_filenames()[coreml_name]
|
||||
|
||||
return (CoreMLModel(coreml_path, compute_unit),)
|
||||
|
||||
|
||||
class CoreMLLoaderUNet(CoreMLLoader):
|
||||
PACKAGE_DIRNAME = "unet"
|
||||
RETURN_TYPES = ("COREML_UNET",)
|
||||
RETURN_NAMES = ("coreml_model",)
|
||||
|
||||
|
||||
class CoreMLModelAdapter(COREML_NODE):
|
||||
"""
|
||||
Adapter Node to use CoreML models as Comfy models. This is an experimental
|
||||
feature and may not work as expected.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"coreml_model": ("COREML_UNET",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
|
||||
FUNCTION = "wrap"
|
||||
CATEGORY = "Core ML Suite"
|
||||
|
||||
def wrap(self, coreml_model):
|
||||
model_patcher = get_model_patcher(coreml_model)
|
||||
return (model_patcher,)
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
||||
"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"],
|
||||
),
|
||||
),
|
||||
"compute_unit": (
|
||||
[
|
||||
ComputeUnit.CPU_AND_NE.name,
|
||||
ComputeUnit.CPU_AND_GPU.name,
|
||||
ComputeUnit.ALL.name,
|
||||
ComputeUnit.CPU_ONLY.name,
|
||||
],
|
||||
),
|
||||
"controlnet_support": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
# k-means weight palettization. Kept optional so workflows
|
||||
# that omit it still validate — ComfyUI rejects a prompt that
|
||||
# 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"},
|
||||
),
|
||||
"lora_params": ("LORA_PARAMS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("COREML_UNET",)
|
||||
RETURN_NAMES = ("coreml_model",)
|
||||
FUNCTION = "convert"
|
||||
|
||||
def convert(
|
||||
self,
|
||||
ckpt_name,
|
||||
height,
|
||||
width,
|
||||
batch_size,
|
||||
attention_implementation,
|
||||
compute_unit,
|
||||
controlnet_support,
|
||||
quantize_nbits="none",
|
||||
lora_params=None,
|
||||
):
|
||||
"""Converts a checkpoint's UNet 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.
|
||||
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 "Load Core ML UNet" node.
|
||||
"""
|
||||
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])
|
||||
lora_weights = [(self.lora_path(lora[0]), lora[1]) for lora in lora_params]
|
||||
|
||||
h = height
|
||||
w = width
|
||||
sample_size = (h // 8, w // 8)
|
||||
import coreml_diffusion
|
||||
|
||||
out_name = coreml_diffusion.compose_out_name(
|
||||
ckpt_name=ckpt_name,
|
||||
batch_size=batch_size,
|
||||
width=w,
|
||||
height=h,
|
||||
controlnet_support=controlnet_support,
|
||||
attention_implementation=attention_implementation,
|
||||
lora_names=coreml_diffusion.lora_names_from_params(lora_params),
|
||||
quantize_nbits=quantize_nbits,
|
||||
)
|
||||
|
||||
logger.info(f"Converting {ckpt_name} to {out_name}")
|
||||
logger.info(f"Batch size: {batch_size}")
|
||||
logger.info(f"Width: {w}, Height: {h}")
|
||||
logger.info(f"ControlNet support: {controlnet_support}")
|
||||
logger.info(f"Attention implementation: {attention_implementation}")
|
||||
|
||||
if lora_params:
|
||||
logger.info("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")
|
||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||
|
||||
config_filename = ckpt_name.split(".")[0] + ".yaml"
|
||||
config_path = folder_paths.get_full_path("configs", config_filename)
|
||||
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,
|
||||
sample_size=sample_size,
|
||||
batch_size=batch_size,
|
||||
controlnet_support=controlnet_support,
|
||||
lora_weights=lora_weights,
|
||||
attn_impl=attention_implementation,
|
||||
config_path=config_path,
|
||||
quantize_nbits=quantize_nbits,
|
||||
)
|
||||
return (CoreMLModel(unet_out_path, compute_unit),)
|
||||
|
||||
@staticmethod
|
||||
def lora_path(lora_name):
|
||||
return folder_paths.get_full_path("loras", lora_name)
|
||||
|
||||
|
||||
class COREML_LOAD_LORA(COREML_NODE, LoraLoader):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
required = LoraLoader.INPUT_TYPES()["required"].copy()
|
||||
required.pop("model")
|
||||
return {
|
||||
"required": required,
|
||||
"optional": {"lora_params": ("LORA_PARAMS",)},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("CLIP", "LORA_PARAMS")
|
||||
RETURN_NAMES = ("CLIP", "lora_params")
|
||||
|
||||
def load_lora(
|
||||
self, clip, lora_name, strength_model, strength_clip, lora_params=None
|
||||
):
|
||||
_, lora_clip = super().load_lora(
|
||||
None, clip, lora_name, strength_model, strength_clip
|
||||
)
|
||||
|
||||
lora_params = lora_params or {}
|
||||
lora_params[lora_name] = (strength_model, strength_clip)
|
||||
|
||||
return lora_clip, lora_params
|
||||
@@ -1,74 +0,0 @@
|
||||
import torch
|
||||
from torchvision.transforms.functional import resize
|
||||
|
||||
from comfy.model_management import get_torch_device
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from coreml_suite.logger import logger
|
||||
from nodes import KSampler
|
||||
|
||||
from coreml_suite.models import CoreMLModelWrapper, get_model_config
|
||||
|
||||
|
||||
def reshape_latent_image(latent_image, target_shape):
|
||||
if latent_image is None:
|
||||
logger.warning("No latent image provided, using zeros.")
|
||||
return {"samples": torch.zeros(target_shape)}
|
||||
|
||||
if latent_image["samples"].shape == target_shape:
|
||||
return latent_image
|
||||
|
||||
logger.warning(
|
||||
"Latent image shape does not match model input shape,"
|
||||
" resizing to match models expected input shape."
|
||||
)
|
||||
resized = resize(latent_image["samples"], target_shape[-2:])
|
||||
return {"samples": resized}
|
||||
|
||||
|
||||
class CoreMLSampler(KSampler):
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
old_required = KSampler.INPUT_TYPES()["required"].copy()
|
||||
old_required.pop("model")
|
||||
old_required.pop("latent_image")
|
||||
new_required = {"coreml_model": ("COREML_UNET",)}
|
||||
return {
|
||||
"required": new_required | old_required,
|
||||
"optional": {"latent_image": ("LATENT",)},
|
||||
}
|
||||
|
||||
CATEGORY = "Core ML Suite"
|
||||
|
||||
def sample(
|
||||
self,
|
||||
coreml_model,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
negative,
|
||||
latent_image=None,
|
||||
denoise=1.0,
|
||||
):
|
||||
sample_shape = coreml_model.expected_inputs["sample"]["shape"]
|
||||
latent_image = reshape_latent_image(latent_image, sample_shape)
|
||||
latent_image["samples"] = latent_image["samples"][0:1]
|
||||
|
||||
model_config = get_model_config()
|
||||
wrapped_model = CoreMLModelWrapper(model_config, coreml_model)
|
||||
model = ModelPatcher(wrapped_model, get_torch_device(), None)
|
||||
|
||||
return super().sample(
|
||||
model,
|
||||
seed,
|
||||
steps,
|
||||
cfg,
|
||||
sampler_name,
|
||||
scheduler,
|
||||
positive,
|
||||
negative,
|
||||
latent_image,
|
||||
denoise,
|
||||
)
|
||||
@@ -1,62 +0,0 @@
|
||||
from itertools import chain
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from coreml_suite.logger import logger
|
||||
|
||||
|
||||
def expand_inputs(inputs):
|
||||
expanded = inputs.copy()
|
||||
for k, v in inputs.items():
|
||||
if isinstance(v, np.ndarray):
|
||||
expanded[k] = np.concatenate([v] * 2) if v.shape[0] == 1 else v
|
||||
elif isinstance(v, torch.Tensor):
|
||||
expanded[k] = torch.cat([v] * 2) if v.shape[0] == 1 else v
|
||||
elif isinstance(v, list):
|
||||
expanded[k] = v * 2 if len(v) == 1 else v
|
||||
elif isinstance(v, dict):
|
||||
expand_inputs(v)
|
||||
return expanded
|
||||
|
||||
|
||||
def extract_residual_kwargs(model, control):
|
||||
if "additional_residual_0" not in model.expected_inputs.keys():
|
||||
return {}
|
||||
if control is None:
|
||||
return no_control(model)
|
||||
|
||||
residual_kwargs = {
|
||||
"additional_residual_{}".format(i): r.cpu().numpy().astype(np.float16)
|
||||
for i, r in enumerate(chain(control["output"], control["middle"]))
|
||||
}
|
||||
return residual_kwargs
|
||||
|
||||
|
||||
def no_control(model):
|
||||
# Dirty hack to get the expected input shape when doing partial ControlNet
|
||||
# 0.18215 is the latent scale factor (IDK, it kinda works)
|
||||
# TODO: Find a better way to do this or tweak the values
|
||||
|
||||
logger.warning(
|
||||
"No ControlNet input, despite the model supports it. "
|
||||
"Using random noise as ControlNet residuals. "
|
||||
"For better results, please use a ControlNet or a model "
|
||||
"that does not support ControlNet."
|
||||
)
|
||||
residuals_names = [
|
||||
name
|
||||
for name in model.expected_inputs.keys()
|
||||
if name.startswith("additional_residual")
|
||||
]
|
||||
residual_kwargs = {
|
||||
"additional_residual_{}".format(i): 0.18215
|
||||
* torch.randn(
|
||||
*model.expected_inputs["additional_residual_{}".format(i)]["shape"]
|
||||
)
|
||||
.cpu()
|
||||
.numpy()
|
||||
.astype(dtype=np.float16)
|
||||
for i in range(len(residuals_names))
|
||||
}
|
||||
return residual_kwargs
|
||||
@@ -0,0 +1,63 @@
|
||||
[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"
|
||||
license = "MIT"
|
||||
requires-python = ">=3.12"
|
||||
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",
|
||||
"coremltools>=9,<10",
|
||||
"numpy>=2,<3",
|
||||
]
|
||||
|
||||
[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"
|
||||
requires-comfyui = ">=0.3.27"
|
||||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"pillow>=12.2.0",
|
||||
"psutil>=7.2.2",
|
||||
"pytest>=9.0.3",
|
||||
]
|
||||
comfy = [
|
||||
"comfyui-frontend-package==1.14.6",
|
||||
"torchvision",
|
||||
"torchaudio",
|
||||
"torchsde",
|
||||
"einops",
|
||||
"tokenizers>=0.13.3",
|
||||
"safetensors>=0.4.2",
|
||||
"aiohttp>=3.11.8",
|
||||
"yarl>=1.18.0",
|
||||
"kornia>=0.7.1",
|
||||
"spandrel",
|
||||
"soundfile",
|
||||
"sentencepiece",
|
||||
]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
markers = [
|
||||
"unit: framework-free unit test (Tier 0)",
|
||||
"smoke: macOS-ARM smoke test on a synthetic micro-model (Tier 1)",
|
||||
"m2: requires Apple Silicon + Neural Engine (Tier 2)",
|
||||
]
|
||||
testpaths = ["tests"]
|
||||
addopts = ["--import-mode=importlib", "--confcutdir=tests"]
|
||||
@@ -1,2 +1,4 @@
|
||||
git+https://github.com/apple/ml-stable-diffusion.git
|
||||
coremltools
|
||||
coreml-diffusion>=0.1.4,<0.2
|
||||
coremltools>=9,<10
|
||||
numpy>=2,<3
|
||||
diffusers>=0.30
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
# Conversion Extraction — Seam Inventory (`docs/extraction/seam.md`)
|
||||
|
||||
> **Gate E0 deliverable.** Symbol-by-symbol cut line between the future `coreml_diffusion`
|
||||
> package (CONVERSION) and what stays in `coreml_suite` (the ComfyUI side).
|
||||
>
|
||||
> **Confidence legend:**
|
||||
> - ✅ **verified** — read directly from the current source in this repo.
|
||||
> - 🔍 **confirm** — inferred / partially seen; Claude Code must `grep`-verify before acting.
|
||||
>
|
||||
> **Cut rule:** a symbol goes to `coreml_diffusion` iff it participates in producing the `.mlpackage`
|
||||
> artifact AND can be made free of `comfy` / `folder_paths` / `comfy_extras`. The runtime
|
||||
> *loader* that **runs** a compiled model stays in the suite.
|
||||
|
||||
---
|
||||
|
||||
## 1. File-level map
|
||||
|
||||
| File | Side | Status | Note |
|
||||
|---|---|---|---|
|
||||
| `coreml_suite/model_version.py` | **coreml_diffusion** | ✅ | Already `Enum`-only, zero comfy. Becomes pkg source of truth. |
|
||||
| `coreml_suite/attention.py` | **coreml_diffusion** | ✅ | `ATTENTION_IMPLEMENTATIONS` tuple; pure constant. |
|
||||
| `coreml_suite/core/naming.py` | **coreml_diffusion** | ✅ | `compose_out_name` = cache-key contract. Move (not copy). |
|
||||
| `coreml_suite/converter.py` | **coreml_diffusion** (mostly) | ✅ | Main conversion. One symbol stays-adjacent: `get_out_path` (folder_paths) is replaced by injected `out_path`. |
|
||||
| `coreml_suite/conversion/attention.py` | **coreml_diffusion** | ✅ | `apply_attention_implementation`. Imports `logging`,`torch` only — no comfy. |
|
||||
| `coreml_suite/conversion/shapes.py` | **coreml_diffusion** | ✅ | `conv2d_output_shape`. Pure math, no imports. |
|
||||
| `coreml_suite/conversion/trace.py` | **coreml_diffusion** | ✅ | Imports `types.MethodType`, `diffusers...Transformer2DModel` only — torch/diffusers. |
|
||||
| `coreml_suite/conversion/unet.py` | **coreml_diffusion** | ✅ | `CoreMLUNetWrapper`. Imports `torch` only — no comfy. |
|
||||
| `coreml_suite/lcm/converter.py` | **coreml_diffusion** (after dedup) | ✅ | Dup helpers deleted; `MODEL_VERSION` HF-hardcode (L22) → E-LCM. `folder_paths` (L111) + `comfy.model_management` (L54) confirmed present → CUT. |
|
||||
| `coreml_suite/lcm/unet.py` | **coreml_diffusion** | ✅ | `UNet2DConditionModelLCM(UNet2DConditionModel)`. diffusers-only, no comfy. |
|
||||
| `coreml_suite/config.py` | **STAYS** | ✅ | Imports `comfy.supported_models_base`/`latent_formats`/`model_detection`. **Inference-side** (`get_model_config`), NOT conversion. |
|
||||
| `coreml_suite/coreml_model.py` | **STAYS** | ✅ | `CoreMLModel` = runtime loader (runs `.mlpackage`). Desktop/Python inference; not used on iOS. |
|
||||
| `coreml_suite/nodes.py` | **STAYS** | ✅ | Nodes; will call `coreml_diffusion` + own `folder_paths` path resolution + discovery dropdowns. |
|
||||
| `coreml_suite/lcm/nodes.py` | **STAYS** | ✅ | `COREML_CONVERT_LCM` node. |
|
||||
| `coreml_suite/models.py` | **STAYS** | ✅ | Inference: `add_sdxl_model_options`, `is_sdxl`, `get_model_patcher`, `get_latent_image`. |
|
||||
| `coreml_suite/latents.py` | **STAYS** | ✅ | Inference chunking (MODERNIZATION Phase 3 target, not this spec). |
|
||||
| `coreml_suite/controlnet.py` | **STAYS** | ✅ | Inference-side controlnet. Distinct from converter `add_cnet_support`. |
|
||||
| `coreml_suite/lcm/utils.py` | **STAYS** | ✅ | `add_lcm_model_options`, `lcm_patch`, `is_lcm`; imports `comfy_extras`. Inference. |
|
||||
| `coreml_suite/logger.py` | **both / copy** | ✅ | Trivial. Package gets its own logger; suite keeps its. |
|
||||
|
||||
---
|
||||
|
||||
## 2. Symbol-level: `coreml_suite/converter.py` (main conversion)
|
||||
|
||||
| Symbol | Side | Status | Cut action |
|
||||
|---|---|---|---|
|
||||
| `DEFAULT_TRACE_TIMESTEP`, `TEXT_TOKEN_SEQUENCE_LENGTH` | coreml_diffusion | ✅ | Move as-is (module constants). |
|
||||
| `get_unet(model_version, ref_unet, attention_implementation)` | coreml_diffusion | ✅ | Move. Uses `conversion.{trace,attention,unet}`. No comfy. |
|
||||
| `get_encoder_hidden_states_shape(ref_unet, batch_size)` | coreml_diffusion | ✅ | Move. Reads `ref_unet.config.cross_attention_dim`. Pure. |
|
||||
| `get_coreml_inputs(sample_inputs)` | coreml_diffusion | ✅ | Move. `ct.TensorType` build. |
|
||||
| `load_coreml_model(out_path)` | coreml_diffusion | ✅ | Move. `ct.models.MLModel(out_path)`. (Dedup target vs LCM copy.) |
|
||||
| `convert_to_coreml(submodule, ts_module, inputs, names, out_path)` | coreml_diffusion | ✅ | Move. `ct.convert(...)`. (Dedup target vs LCM copy.) |
|
||||
| `get_sample_input(batch, ehs_shape, sample_shape)` | coreml_diffusion | ✅ | Move. **Merge** with LCM variant (LCM passes extra `scheduler` → optional param). |
|
||||
| `lcm_inputs(sample_unet_inputs)` | coreml_diffusion | ✅ | Move. Adds `timestep_cond`. |
|
||||
| `sdxl_inputs(sample_unet_inputs, ref_unet, model_version)` | coreml_diffusion | ✅ | Move. `time_ids`/`text_embeds`/`add_embeds`. |
|
||||
| `add_cnet_support(sample_shape, ref_unet)` | coreml_diffusion | ✅ | Move. Builds `additional_residual_*` inputs from unet block channels. |
|
||||
| `convert_unet(ref_unet, model_version, unet_out_path, ...)` | coreml_diffusion | ✅ | Move. Orchestrates trace→convert→**quant (palettize)**→save. Quant travels here (E6). |
|
||||
| `convert(ckpt_path, model_version, unet_out_path, ...)` | coreml_diffusion | ✅ | Move. **Make kw-only past `ckpt_path,model_version,out_path`** (contract). Validates `attn_impl`. |
|
||||
| `load_unet(ckpt_path, config_path)` | coreml_diffusion | ✅ | Move. `UNet2DConditionModel.from_single_file`. |
|
||||
| `get_out_path(submodule_name, model_name)` | **STAYS (node)** | ✅ | Uses `folder_paths.get_folder_paths`. **Delete from converter; node resolves path and passes `out_path` in.** |
|
||||
|
||||
**Apple `python_coreml_stable_diffusion` footprint on this path:** ✅ **none.** Verified by grep:
|
||||
zero imports in `converter.py` / `conversion/*`. Main path uses `diffusers` +
|
||||
local `CoreMLUNetWrapper`. (And the runtime `CoreMLModel` is now a local coremltools wrapper too —
|
||||
see §6 stale-spec note.)
|
||||
|
||||
---
|
||||
|
||||
## 3. Symbol-level: `coreml_suite/lcm/converter.py` (LCM — dedup + defer)
|
||||
|
||||
| Symbol | Side | Status | Cut action |
|
||||
|---|---|---|---|
|
||||
| `load_coreml_model` (LCM copy) | DELETE | ✅ | Duplicate of main. Remove; use `coreml_diffusion.load_coreml_model`. |
|
||||
| `convert_to_coreml` (LCM copy) | DELETE | ✅ | Duplicate of main. Remove. |
|
||||
| `get_out_path` (LCM copy, folder_paths) | DELETE | ✅ | Duplicate + comfy. Remove; node injects `out_path`. |
|
||||
| `get_sample_input(..., scheduler)` (LCM copy) | MERGE → coreml_diffusion | ✅ | Fold `scheduler` into shared `get_sample_input` as optional param. |
|
||||
| `MODEL_NAME` (= LCM_Dreamshaper) | **E-LCM** | ✅ | HF hardcode. Removing it is the behavior change → E-LCM, not E2. |
|
||||
| `convert(out_path, sample_size, batch_size, controlnet_support)` (LCM, L190) | coreml_diffusion (via unified) | ✅ | Route through `coreml_diffusion.convert(model_version=LCM, ...)` in E-LCM. |
|
||||
| `from comfy.model_management import get_torch_device` (L54, in `get_scheduler`) | **CUT** | ✅ | Confirmed present. Inject `device`. |
|
||||
| module-global attention set at import | n/a | ✅ | **No module global.** Attention already per-call: `get_unets` (L36) calls `apply_attention_implementation(ref_unet, "SPLIT_EINSUM")`. No `ATTENTION_IMPLEMENTATION_IN_EFFECT` anywhere in repo. (Note: LCM hardcodes `"SPLIT_EINSUM"` — pass `attn_impl` through in dedup.) |
|
||||
|
||||
---
|
||||
|
||||
## 4. Symbol-level: `coreml_suite/core/naming.py` → `coreml_diffusion/naming.py`
|
||||
|
||||
| Symbol | Side | Status | Cut action |
|
||||
|---|---|---|---|
|
||||
| `compose_out_name(...)` | coreml_diffusion | ✅ | **Move** (cache-key contract). Node imports from pkg. |
|
||||
| `lora_names_from_params(...)` | coreml_diffusion | ✅ | Move. |
|
||||
| `ATTN_SUFFIX` dict | coreml_diffusion | ✅ | Move. |
|
||||
| `QUANT_NBITS_VALUES` | coreml_diffusion | ✅ | Move; backs `list_quant_modes()`. |
|
||||
| `tests/unit/test_characterization_out_name.py` | re-point | ✅ | Change import to `coreml_diffusion.naming`. Assertions/values **unchanged**. |
|
||||
|
||||
---
|
||||
|
||||
## 5. Discovery API + status registry (new in `coreml_diffusion/__init__.py`)
|
||||
|
||||
```python
|
||||
from enum import Enum
|
||||
|
||||
class Status(Enum):
|
||||
VERIFIED = "verified" # has a golden anchor + passing [M2-ANE] check
|
||||
EXPERIMENTAL = "experimental" # convertible, not yet anchored/verified
|
||||
|
||||
# Single source of truth. Suite gates on this, NOT on a hardcoded node list.
|
||||
# KEY by ModelVersion enum MEMBER (not a bare string) so list_model_versions can
|
||||
# emit .name — see the .name decision below. Keying by the lowercase .value string
|
||||
# (as an earlier draft of this block did) returns ["sd15",...], which the node then
|
||||
# reverses via ModelVersion[...] → KeyError. Do NOT key by .value.
|
||||
_MODEL_STATUS = {
|
||||
ModelVersion.SD15: Status.VERIFIED,
|
||||
ModelVersion.SDXL: Status.VERIFIED,
|
||||
ModelVersion.SDXL_REFINER: Status.EXPERIMENTAL, # → VERIFIED after a refiner golden anchor
|
||||
ModelVersion.LCM: Status.EXPERIMENTAL, # → VERIFIED after E-LCM golden anchor
|
||||
}
|
||||
|
||||
def list_model_versions(include_experimental: bool = False) -> list[str]:
|
||||
return [v.name for v, s in _MODEL_STATUS.items() # .name → "SD15","SDXL" (see decision)
|
||||
if s is Status.VERIFIED or (include_experimental and s is Status.EXPERIMENTAL)]
|
||||
|
||||
def list_attention_impls() -> list[str]: # from attention.ATTENTION_IMPLEMENTATIONS
|
||||
...
|
||||
def list_quant_modes() -> list[str]: # from naming.QUANT_NBITS_VALUES
|
||||
...
|
||||
|
||||
CONTRACT_VERSION = "1.0"
|
||||
# Additive-only: adding an id or promoting EXPERIMENTAL→VERIFIED = minor bump (Suite unaffected).
|
||||
# Removing/renaming an id, or demoting VERIFIED→EXPERIMENTAL = MAJOR bump + migration note.
|
||||
```
|
||||
|
||||
**Decision check (`.name` vs `.value`): RESOLVED → `.name`.** ✅
|
||||
Verified in current source:
|
||||
- Node renders `ModelVersion.SD15.name` / `ModelVersion.SDXL.name` → `"SD15"`, `"SDXL"`
|
||||
(`nodes.py:224-225`).
|
||||
- Node reverses the dropdown string with `model_version = ModelVersion[model_version]`
|
||||
(`nodes.py:286`) — i.e. **lookup by NAME**. Feeding it a `.value` (`"sd15"`) raises `KeyError`.
|
||||
- Enum values are lowercase (`model_version.py`: `SD15="sd15"`, `SDXL="sdxl"`,
|
||||
`SDXL_REFINER="sdxl_refiner"`, `LCM="lcm"`).
|
||||
- `compose_out_name` does NOT consume the model_version string (grep of `core/naming.py` empty) —
|
||||
no coupling there, so no constraint from that side.
|
||||
|
||||
**Decision:** `list_model_versions()` returns `.name` (uppercase). Saved workflows store `"SD15"`,
|
||||
node already validates them via `ModelVersion[...]`. The `_MODEL_STATUS` block above was corrected
|
||||
to key by enum member and emit `.name`. **The earlier `v.value` form was a latent bug.**
|
||||
|
||||
---
|
||||
|
||||
## 6. `python_coreml_stable_diffusion` split (Gate E0 line to fill by grep)
|
||||
|
||||
| Use | Side | Status |
|
||||
|---|---|---|
|
||||
| `coreml_model.CoreMLModel` (runs compiled model) | **STAYS** (suite runtime) | ✅ — **local class**, not Apple's |
|
||||
| `unet.UNet2DConditionModel*` internals | **gone** — `converter.py:319` uses `diffusers.UNet2DConditionModel.from_single_file` | ✅ |
|
||||
| `AttentionImplementations` enum | gone — local `apply_attention_implementation` + `attention.py` tuple | ✅ |
|
||||
| `calculate_conv2d_output_shape` | gone — replaced by `conversion/shapes.conv2d_output_shape` | ✅ |
|
||||
|
||||
> ### ⚠️ SPEC IS STALE: `ml-stable-diffusion` is already fully removed
|
||||
> Commit #58 ("replace apple/ml-stable-diffusion with native diffusers conversion") already did
|
||||
> the de-Apple work. Verified now:
|
||||
> - **Zero** `python_coreml_stable_diffusion` runtime imports anywhere in `coreml_suite` (only a
|
||||
> docstring mention at `core/__init__.py:4`).
|
||||
> - `coreml_suite/coreml_model.py:8` `CoreMLModel` is a **local** wrapper over
|
||||
> `coremltools.models.MLModel` (`coreml_model.py:22`) — it does **not** import Apple's class.
|
||||
> - `ml-stable-diffusion` / `python_coreml_stable_diffusion` appears in **neither** `pyproject.toml`
|
||||
> **nor** `requirements.txt`. It is not a dependency at all.
|
||||
>
|
||||
> **Consequences for the spec (correct these in CONVERTER_EXTRACTION_SPEC.md):**
|
||||
> - §0.3 premise ("runtime loader = `python_coreml_stable_diffusion.coreml_model.CoreMLModel`,
|
||||
> stays in suite") is **wrong**: the loader is already the local `coreml_model.CoreMLModel`. The
|
||||
> "stays in suite" conclusion still holds; the identity does not.
|
||||
> - **Gate E0 item "ml-stable-diffusion pinned SHA — BLOCKER if unpinned" is MOOT** — there is no
|
||||
> such dep to pin. Mark it N/A, not BLOCKER.
|
||||
> - **E4/E5 dependency lists must drop `git+...ml-stable-diffusion@<sha>`.** Package runtime deps
|
||||
> are: `coremltools`, `diffusers`, `peft` (LoRA), `omegaconf` (config), `numpy`, `torch`. Confirm
|
||||
> `peft`/`omegaconf` actually used before listing (grep at E4).
|
||||
> - The "keep `python_coreml_stable_diffusion` as a suite dep for the loader" instruction in E5 is
|
||||
> **void** — coremltools backs the loader.
|
||||
|
||||
---
|
||||
|
||||
## 7. Pre-flight checklist before E1 (run these greps)
|
||||
|
||||
```
|
||||
grep -rn "import comfy" coreml_suite/conversion coreml_suite/converter.py coreml_suite/lcm/converter.py coreml_suite/lcm/unet.py
|
||||
grep -rn "folder_paths" coreml_suite/converter.py coreml_suite/lcm/converter.py
|
||||
grep -rn "model_management" coreml_suite/lcm
|
||||
grep -rn "python_coreml_stable_diffusion" coreml_suite
|
||||
grep -rn "ATTENTION_IMPLEMENTATION_IN_EFFECT" coreml_suite
|
||||
grep -rn "SimianLuo\|LCM_Dreamshaper" coreml_suite/lcm
|
||||
```
|
||||
Every 🔍 above resolves to ✅ or a correction once these run. Do not start moving code (E2)
|
||||
with any 🔍 unresolved on the CONVERSION side.
|
||||
|
||||
**STATUS (run 2026-05-26): all 🔍 resolved.** Summary of what the greps found:
|
||||
- `conversion/*`, `lcm/unet.py`: comfy-free (torch/diffusers only). ✅
|
||||
- `converter.py`: only comfy reach-in is `folder_paths` in `get_out_path` (L91-94) → inject `out_path`.
|
||||
- `lcm/converter.py`: `folder_paths` (L111-114) + `comfy.model_management.get_torch_device` (L54)
|
||||
→ cut both. Dup helpers (`load_coreml_model`,`convert_to_coreml`,`get_out_path`,`get_sample_input`)
|
||||
confirmed → dedup E2. `MODEL_VERSION="SimianLuo/LCM_Dreamshaper_v7"` (L22) → E-LCM.
|
||||
- No attention module-global anywhere (`ATTENTION_IMPLEMENTATION_IN_EFFECT` absent); already per-call.
|
||||
LCM hardcodes `"SPLIT_EINSUM"` in `get_unets` — thread `attn_impl` through during dedup.
|
||||
- `.name` vs `.value`: **decided `.name`** (node reverses via `ModelVersion[...]`). §5 corrected.
|
||||
- `ml-stable-diffusion`: **already gone** (#58). §6 stale-spec note added — fix the spec's E0/E4/E5
|
||||
dep + pinning items.
|
||||
|
||||
Two grep blind-spots to note (the checklist above doesn't cover them, but cheap to add): the
|
||||
`folder_paths` grep only scans the two converter files — also grep `coreml_suite/lcm/utils.py`
|
||||
(it imports `comfy.model_management` at L3, but it's inference/STAYS, so fine) and confirm no other
|
||||
`conversion/` file grew a comfy import since.
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Pytest bootstrap for ComfyUI-CoreMLSuite tests.
|
||||
|
||||
- Adds the ComfyUI checkout to sys.path so the framework-coupled modules
|
||||
that transitively import `comfy.*` resolve when pytest is invoked from
|
||||
this package's root.
|
||||
- Auto-applies tier markers based on the directory a test lives in, so
|
||||
individual files don't have to repeat @pytest.mark.unit / .smoke.
|
||||
"""
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
COMFY_DIR = REPO_ROOT.parents[1]
|
||||
|
||||
for p in (str(COMFY_DIR), str(REPO_ROOT)):
|
||||
if p not in sys.path:
|
||||
sys.path.insert(0, p)
|
||||
|
||||
|
||||
_TIER_BY_DIR = {
|
||||
"tests/unit": "unit",
|
||||
"tests/m2": "m2",
|
||||
"tests/integration": "m2",
|
||||
"tests/smoke": "smoke",
|
||||
}
|
||||
|
||||
# When the user asks for a single tier (-m unit / -m smoke), skip the other
|
||||
# directories at collection time. Tier-0 cannot afford to import tests/smoke
|
||||
# files because they pull in coremltools which Linux CI won't have.
|
||||
_TIER_DIRS = {
|
||||
"unit": ("/tests/unit/",),
|
||||
"m2": ("/tests/m2/", "/tests/integration/"),
|
||||
"smoke": ("/tests/smoke/",),
|
||||
}
|
||||
|
||||
|
||||
def pytest_ignore_collect(collection_path, config):
|
||||
expr = config.option.markexpr
|
||||
if expr not in _TIER_DIRS:
|
||||
return None
|
||||
allowed = _TIER_DIRS[expr]
|
||||
rel = str(collection_path).replace("\\", "/")
|
||||
if "/tests/" not in rel:
|
||||
return None
|
||||
# Always allow tests/ root + the tier's own dirs.
|
||||
if rel.endswith("/tests"):
|
||||
return None
|
||||
if any(frag in rel + "/" for frag in allowed):
|
||||
return None
|
||||
return True
|
||||
|
||||
|
||||
def pytest_collection_modifyitems(config, items):
|
||||
for item in items:
|
||||
path = str(item.fspath).replace("\\", "/")
|
||||
for fragment, marker in _TIER_BY_DIR.items():
|
||||
if f"/{fragment}/" in path:
|
||||
item.add_marker(getattr(pytest.mark, marker))
|
||||
break
|
||||
@@ -0,0 +1,181 @@
|
||||
{
|
||||
"3": {
|
||||
"inputs": {
|
||||
"seed": 0,
|
||||
"steps": 20,
|
||||
"cfg": 8,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
"denoise": 1,
|
||||
"model": [
|
||||
"4",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"6",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"7",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"5",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "KSampler",
|
||||
"_meta": {
|
||||
"title": "KSampler"
|
||||
}
|
||||
},
|
||||
"4": {
|
||||
"inputs": {
|
||||
"ckpt_name": "dreamshaper_8.safetensors"
|
||||
},
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"_meta": {
|
||||
"title": "Load Checkpoint"
|
||||
}
|
||||
},
|
||||
"5": {
|
||||
"inputs": {
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"batch_size": 1
|
||||
},
|
||||
"class_type": "EmptyLatentImage",
|
||||
"_meta": {
|
||||
"title": "Empty Latent Image"
|
||||
}
|
||||
},
|
||||
"6": {
|
||||
"inputs": {
|
||||
"text": "beautiful scenery nature glass bottle landscape, purple galaxy bottle",
|
||||
"clip": [
|
||||
"4",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"7": {
|
||||
"inputs": {
|
||||
"text": "text, watermark",
|
||||
"clip": [
|
||||
"4",
|
||||
1
|
||||
]
|
||||
},
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "CLIP Text Encode (Prompt)"
|
||||
}
|
||||
},
|
||||
"8": {
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"3",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"4",
|
||||
2
|
||||
]
|
||||
},
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {
|
||||
"title": "VAE Decode"
|
||||
}
|
||||
},
|
||||
"9": {
|
||||
"inputs": {
|
||||
"filename_prefix": "E2E-1.5-MPS",
|
||||
"images": [
|
||||
"8",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SaveImage",
|
||||
"_meta": {
|
||||
"title": "Save Image"
|
||||
}
|
||||
},
|
||||
"10": {
|
||||
"inputs": {
|
||||
"ckpt_name": "dreamshaper_8.safetensors",
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"batch_size": 1,
|
||||
"attention_implementation": "SPLIT_EINSUM",
|
||||
"compute_unit": "CPU_AND_NE",
|
||||
"controlnet_support": false
|
||||
},
|
||||
"class_type": "Core ML Converter",
|
||||
"_meta": {
|
||||
"title": "Convert Checkpoint to Core ML"
|
||||
}
|
||||
},
|
||||
"11": {
|
||||
"inputs": {
|
||||
"seed": 0,
|
||||
"steps": 20,
|
||||
"cfg": 8,
|
||||
"sampler_name": "dpmpp_2m",
|
||||
"scheduler": "karras",
|
||||
"denoise": 1,
|
||||
"coreml_model": [
|
||||
"10",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"6",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"7",
|
||||
0
|
||||
],
|
||||
"latent_image": [
|
||||
"5",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "CoreMLSampler",
|
||||
"_meta": {
|
||||
"title": "Core ML Sampler"
|
||||
}
|
||||
},
|
||||
"13": {
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"11",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"4",
|
||||
2
|
||||
]
|
||||
},
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {
|
||||
"title": "VAE Decode"
|
||||
}
|
||||
},
|
||||
"14": {
|
||||
"inputs": {
|
||||
"filename_prefix": "E2E-1.5-CoreML",
|
||||
"images": [
|
||||
"13",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "SaveImage",
|
||||
"_meta": {
|
||||
"title": "Save Image"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
After Width: | Height: | Size: 448 KiB |
@@ -0,0 +1 @@
|
||||
e89344e544d4edfbd3ebe9a1c78dadb2729f53549666052b74ac7308f326f4fc
|
||||
@@ -0,0 +1,170 @@
|
||||
"""[M2-ANE] golden-image anchor.
|
||||
|
||||
Runs the e2e SD1.5 + CoreML workflow against a local ComfyUI server, fetches
|
||||
the generated PNG, and asserts both:
|
||||
- byte-identical SHA256 against the stored golden, OR
|
||||
- PSNR >= GOLDEN_PSNR_MIN_DB against the stored golden PNG.
|
||||
|
||||
The hash is the strict gate (a refactor that doesn't touch the math
|
||||
should hit it). PSNR is the soft gate that tolerates the drift a
|
||||
toolchain bump injects through different MIL graphs / kernel selection
|
||||
/ fp accumulation order — anything below the threshold is treated as a
|
||||
regression.
|
||||
|
||||
The 20 dB default absorbs Apple Neural Engine run-to-run nondeterminism:
|
||||
the same model and seed can drift several dB between runs as the 20
|
||||
sampling steps amplify tiny per-step UNet differences (kernel selection /
|
||||
fp accumulation order). Same-scene ANE outputs have been observed at
|
||||
~23 dB, so 20 leaves margin while still catching gross regressions — a
|
||||
broken image lands far lower. Bump it up for pure-refactor PRs that must
|
||||
not change math; down for toolchain bumps.
|
||||
|
||||
Skips entirely on non-Apple-Silicon hosts or when the server / converted
|
||||
model is missing, so the unit lane on Linux still passes.
|
||||
|
||||
The first run with no golden writes one and fails so it's reviewed before
|
||||
being committed.
|
||||
"""
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import shutil
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
COMFY_DIR = Path(os.environ.get("COMFY_DIR", REPO_ROOT.parents[1])).resolve()
|
||||
COMFY_HOST = os.environ.get("COMFY_HOST", "localhost")
|
||||
COMFY_PORT = int(os.environ.get("COMFY_PORT", "8188"))
|
||||
COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
|
||||
|
||||
CKPT_NAME = os.environ.get("CKPT_NAME", "v1-5-pruned-emaonly.safetensors")
|
||||
WORKFLOW_PATH = (
|
||||
REPO_ROOT / "tests" / "integration" / "workflows" / "e2e-1.5-basic-conversion.json"
|
||||
)
|
||||
GOLDEN_DIR = Path(__file__).parent / "goldens"
|
||||
GOLDEN_HASH_PATH = GOLDEN_DIR / "sd15_seed42.sha256"
|
||||
GOLDEN_PNG_PATH = GOLDEN_DIR / "sd15_seed42.png"
|
||||
GOLDEN_PSNR_MIN_DB = float(os.environ.get("GOLDEN_PSNR_MIN_DB", "20"))
|
||||
SEED = 42
|
||||
|
||||
|
||||
def _server_reachable() -> bool:
|
||||
try:
|
||||
with urllib.request.urlopen(f"{COMFY_URL}/prompt", timeout=3) as r:
|
||||
return r.status == 200
|
||||
except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError):
|
||||
return False
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def comfy_server():
|
||||
if platform.machine() != "arm64":
|
||||
pytest.skip("requires Apple Silicon")
|
||||
if not _server_reachable():
|
||||
pytest.skip(f"ComfyUI server not reachable at {COMFY_URL}")
|
||||
return COMFY_URL
|
||||
|
||||
|
||||
def _http_post_json(path: str, payload: dict) -> dict:
|
||||
data = json.dumps(payload).encode("utf-8")
|
||||
req = urllib.request.Request(
|
||||
f"{COMFY_URL}{path}", data=data,
|
||||
headers={"Content-Type": "application/json"}, method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=300) as r:
|
||||
return json.loads(r.read().decode())
|
||||
|
||||
|
||||
def _http_get_json(path: str, timeout: int = 300) -> dict:
|
||||
"""ComfyUI runs UNet inference on its single asyncio loop, so GET /prompt
|
||||
blocks while the queued prompt is executing. Use a generous timeout."""
|
||||
with urllib.request.urlopen(f"{COMFY_URL}{path}", timeout=timeout) as r:
|
||||
return json.loads(r.read().decode())
|
||||
|
||||
|
||||
def _drain_queue(timeout_s: int = 600) -> None:
|
||||
deadline = time.time() + timeout_s
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
q = _http_get_json("/prompt")
|
||||
except (urllib.error.URLError, TimeoutError):
|
||||
# Transient block while server executes; retry until our overall
|
||||
# deadline expires.
|
||||
continue
|
||||
if q.get("exec_info", {}).get("queue_remaining", -1) == 0:
|
||||
return
|
||||
time.sleep(2)
|
||||
raise TimeoutError(f"queue did not drain within {timeout_s}s")
|
||||
|
||||
|
||||
def _post_workflow_and_collect_png() -> bytes:
|
||||
workflow = json.loads(WORKFLOW_PATH.read_text())
|
||||
for nid in ("4", "10"):
|
||||
if nid in workflow:
|
||||
workflow[nid]["inputs"]["ckpt_name"] = CKPT_NAME
|
||||
for nid in ("3", "11"):
|
||||
if nid in workflow and "seed" in workflow[nid].get("inputs", {}):
|
||||
workflow[nid]["inputs"]["seed"] = SEED
|
||||
# Drop the MPS reference branch — only the Core ML pipeline is needed here.
|
||||
for nid in ("3", "8", "9"):
|
||||
workflow.pop(nid, None)
|
||||
|
||||
_http_post_json("/prompt", {"prompt": workflow})
|
||||
_drain_queue()
|
||||
|
||||
comfy_out = COMFY_DIR / "output"
|
||||
matches = sorted(comfy_out.glob("E2E-1.5-CoreML_*.png"), reverse=True)
|
||||
if not matches:
|
||||
raise FileNotFoundError(f"no Core ML image under {comfy_out}")
|
||||
return matches[0].read_bytes()
|
||||
|
||||
|
||||
def _psnr(a: np.ndarray, b: np.ndarray) -> float:
|
||||
mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
|
||||
if mse == 0:
|
||||
return 100.0
|
||||
return 20.0 * float(np.log10(255.0 / np.sqrt(mse)))
|
||||
|
||||
|
||||
def test_sd15_seed42_image_matches_golden(comfy_server):
|
||||
GOLDEN_DIR.mkdir(parents=True, exist_ok=True)
|
||||
png_bytes = _post_workflow_and_collect_png()
|
||||
sha = hashlib.sha256(png_bytes).hexdigest()
|
||||
|
||||
if not GOLDEN_HASH_PATH.exists() or not GOLDEN_PNG_PATH.exists():
|
||||
GOLDEN_HASH_PATH.write_text(sha + "\n")
|
||||
# Persist the PNG too for visual diffing + PSNR.
|
||||
tmp_path = Path(__file__).parent / "_latest_generated.png"
|
||||
tmp_path.write_bytes(png_bytes)
|
||||
shutil.copy2(tmp_path, GOLDEN_PNG_PATH)
|
||||
pytest.fail(
|
||||
f"No golden present; wrote {GOLDEN_HASH_PATH.name} and "
|
||||
f"{GOLDEN_PNG_PATH.name}. Review the image and re-run."
|
||||
)
|
||||
|
||||
expected_hash = GOLDEN_HASH_PATH.read_text().strip()
|
||||
if sha == expected_hash:
|
||||
return
|
||||
|
||||
# Hash drift: fall back to PSNR to distinguish a refactor-safe rounding
|
||||
# change from a real regression.
|
||||
a = np.array(Image.open(GOLDEN_PNG_PATH).convert("RGB"))
|
||||
b_path = Path(__file__).parent / "_latest_generated.png"
|
||||
b_path.write_bytes(png_bytes)
|
||||
b = np.array(Image.open(b_path).convert("RGB"))
|
||||
if a.shape != b.shape:
|
||||
pytest.fail(f"shape mismatch: golden={a.shape} actual={b.shape}")
|
||||
psnr_db = _psnr(a, b)
|
||||
assert psnr_db >= GOLDEN_PSNR_MIN_DB, (
|
||||
f"hash drifted (got {sha[:12]}.., expected {expected_hash[:12]}..) and "
|
||||
f"PSNR {psnr_db:.2f} dB < {GOLDEN_PSNR_MIN_DB} dB threshold; "
|
||||
f"diff PNG at {b_path}"
|
||||
)
|
||||
@@ -1,19 +0,0 @@
|
||||
import pytest
|
||||
|
||||
import torch
|
||||
|
||||
from coreml_suite.samplers import reshape_latent_image
|
||||
|
||||
|
||||
def test_fix_latents_no_latent_image():
|
||||
reshaped = reshape_latent_image(None, (2, 4, 64, 64))
|
||||
assert reshaped["samples"].shape == (2, 4, 64, 64)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"latent_shape", [(2, 4, 64, 64), (2, 4, 128, 128), (2, 4, 32, 32), (2, 4, 128, 64)]
|
||||
)
|
||||
def test_reshape_latents(latent_shape):
|
||||
latent_image = {"samples": torch.zeros(latent_shape)}
|
||||
reshaped = reshape_latent_image(latent_image, (2, 4, 64, 64))
|
||||
assert reshaped["samples"].shape == (2, 4, 64, 64)
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Characterization tests for coreml_suite.controlnet.
|
||||
|
||||
Locks shapes + dtypes + zero-fill behavior of expand_inputs / no_control /
|
||||
extract_residual_kwargs / chunk_control. These pure helpers feed the Core ML
|
||||
UNet's additional_residual_N inputs; any drift here silently breaks
|
||||
ControlNet-based workflows.
|
||||
"""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.controlnet import (
|
||||
chunk_control,
|
||||
expand_inputs,
|
||||
extract_residual_kwargs,
|
||||
no_control,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _deterministic_seed():
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
SD15_RESIDUAL_SPEC = {
|
||||
"additional_residual_0": {"shape": (2, 320, 64, 64)},
|
||||
"additional_residual_1": {"shape": (2, 640, 32, 32)},
|
||||
"additional_residual_2": {"shape": (2, 1280, 8, 8)},
|
||||
}
|
||||
NON_RESIDUAL_SPEC = {
|
||||
"sample": {"shape": (2, 4, 64, 64)},
|
||||
"encoder_hidden_states": {"shape": (2, 77, 768)},
|
||||
}
|
||||
|
||||
|
||||
# ---------- expand_inputs ----------------------------------------------------
|
||||
|
||||
|
||||
def test_expand_inputs_doubles_singleton_numpy():
|
||||
inputs = {"a": np.ones((1, 4), dtype=np.float32)}
|
||||
out = expand_inputs(inputs)
|
||||
assert out["a"].shape == (2, 4)
|
||||
assert np.array_equal(out["a"], np.ones((2, 4)))
|
||||
|
||||
|
||||
def test_expand_inputs_doubles_singleton_torch():
|
||||
inputs = {"a": torch.ones(1, 4)}
|
||||
out = expand_inputs(inputs)
|
||||
assert out["a"].shape == (2, 4)
|
||||
assert torch.equal(out["a"], torch.ones(2, 4))
|
||||
|
||||
|
||||
def test_expand_inputs_doubles_singleton_list():
|
||||
inputs = {"a": [42]}
|
||||
out = expand_inputs(inputs)
|
||||
assert out["a"] == [42, 42]
|
||||
|
||||
|
||||
def test_expand_inputs_skips_already_batched():
|
||||
"""batch > 1 inputs are returned unchanged (same object identity)."""
|
||||
arr = np.ones((2, 4), dtype=np.float32)
|
||||
tensor = torch.ones(3, 4)
|
||||
lst = [1, 2]
|
||||
out = expand_inputs({"a": arr, "b": tensor, "c": lst})
|
||||
assert out["a"] is arr
|
||||
assert out["b"] is tensor
|
||||
assert out["c"] is lst
|
||||
|
||||
|
||||
def test_expand_inputs_preserves_unknown_value_types():
|
||||
# Strings/None pass through untouched — locks current permissive contract.
|
||||
inputs = {"s": "hello", "none": None, "int": 7}
|
||||
out = expand_inputs(inputs)
|
||||
assert out == {"s": "hello", "none": None, "int": 7}
|
||||
|
||||
|
||||
# ---------- no_control -------------------------------------------------------
|
||||
|
||||
|
||||
def test_no_control_returns_zero_fp16_for_residuals():
|
||||
out = no_control({**SD15_RESIDUAL_SPEC, **NON_RESIDUAL_SPEC})
|
||||
# Only additional_residual_* keys are produced.
|
||||
assert set(out.keys()) == set(SD15_RESIDUAL_SPEC.keys())
|
||||
for key, spec in SD15_RESIDUAL_SPEC.items():
|
||||
arr = out[key]
|
||||
assert arr.shape == spec["shape"]
|
||||
assert arr.dtype == np.float16
|
||||
assert np.all(arr == 0)
|
||||
|
||||
|
||||
def test_no_control_returns_empty_when_no_residuals():
|
||||
out = no_control(NON_RESIDUAL_SPEC)
|
||||
assert out == {}
|
||||
|
||||
|
||||
# ---------- extract_residual_kwargs -----------------------------------------
|
||||
|
||||
|
||||
def test_extract_residual_kwargs_empty_when_model_has_no_residual_inputs():
|
||||
out = extract_residual_kwargs(NON_RESIDUAL_SPEC, control={"output": [], "middle": []})
|
||||
assert out == {}
|
||||
|
||||
|
||||
def test_extract_residual_kwargs_none_control_returns_no_control_shapes():
|
||||
out = extract_residual_kwargs(SD15_RESIDUAL_SPEC, control=None)
|
||||
assert set(out.keys()) == set(SD15_RESIDUAL_SPEC.keys())
|
||||
for key, spec in SD15_RESIDUAL_SPEC.items():
|
||||
assert out[key].shape == spec["shape"]
|
||||
assert out[key].dtype == np.float16
|
||||
assert np.all(out[key] == 0)
|
||||
|
||||
|
||||
def test_extract_residual_kwargs_flattens_output_then_middle_and_casts_fp16():
|
||||
"""output residuals come first (indexed 0..N-1), then middle residuals
|
||||
(indexed N..M-1). Values come out of CPU as fp16 numpy arrays."""
|
||||
control = {
|
||||
"output": [torch.ones(2, 320, 64, 64) * 0.5, torch.ones(2, 640, 32, 32) * 2.0],
|
||||
"middle": [torch.ones(2, 1280, 8, 8) * -1.0],
|
||||
}
|
||||
out = extract_residual_kwargs(SD15_RESIDUAL_SPEC, control)
|
||||
assert set(out.keys()) == {"additional_residual_0", "additional_residual_1", "additional_residual_2"}
|
||||
assert out["additional_residual_0"].shape == (2, 320, 64, 64)
|
||||
assert out["additional_residual_1"].shape == (2, 640, 32, 32)
|
||||
assert out["additional_residual_2"].shape == (2, 1280, 8, 8)
|
||||
for arr in out.values():
|
||||
assert arr.dtype == np.float16
|
||||
# Locked order: index 0 == first output residual (0.5), index 2 == middle (-1.0).
|
||||
assert np.allclose(out["additional_residual_0"], 0.5)
|
||||
assert np.allclose(out["additional_residual_1"], 2.0)
|
||||
assert np.allclose(out["additional_residual_2"], -1.0)
|
||||
|
||||
|
||||
# ---------- chunk_control ----------------------------------------------------
|
||||
|
||||
|
||||
def test_chunk_control_none_returns_list_of_nones_with_length_target():
|
||||
"""`no_control` path: when there's no control, you get [None] * target_size
|
||||
(NOT [None, None] regardless of target — this is the contract today)."""
|
||||
assert chunk_control(None, 1) == [None]
|
||||
assert chunk_control(None, 2) == [None, None]
|
||||
assert chunk_control(None, 4) == [None, None, None, None]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch,target,expected_chunks",
|
||||
[(1, 2, 1), (2, 2, 1), (3, 2, 2), (4, 2, 2), (5, 3, 2), (9, 4, 3)],
|
||||
)
|
||||
def test_chunk_control_shapes_after_chunking(batch, target, expected_chunks):
|
||||
cn = {
|
||||
"output": [
|
||||
torch.randn(batch, 320, 64, 64),
|
||||
torch.randn(batch, 640, 32, 32),
|
||||
],
|
||||
"middle": [torch.randn(batch, 1280, 8, 8)],
|
||||
}
|
||||
chunks = chunk_control(cn, target)
|
||||
assert len(chunks) == expected_chunks
|
||||
for c in chunks:
|
||||
assert c["output"][0].shape == (target, 320, 64, 64)
|
||||
assert c["output"][1].shape == (target, 640, 32, 32)
|
||||
assert c["middle"][0].shape == (target, 1280, 8, 8)
|
||||
|
||||
|
||||
def test_chunk_control_preserves_keys_order():
|
||||
"""Output dicts contain exactly {"output", "middle"} in that order."""
|
||||
cn = {
|
||||
"output": [torch.zeros(2, 4, 4, 4)],
|
||||
"middle": [torch.zeros(2, 4, 4, 4)],
|
||||
}
|
||||
chunks = chunk_control(cn, 2)
|
||||
assert list(chunks[0].keys()) == ["output", "middle"]
|
||||
|
||||
|
||||
def test_chunk_control_zero_pads_remainder():
|
||||
"""A batch=3, target=2 split puts the third row alongside a zero row."""
|
||||
cn = {
|
||||
"output": [torch.arange(3 * 4).reshape(3, 1, 2, 2).float()],
|
||||
"middle": [torch.arange(3 * 4).reshape(3, 1, 2, 2).float()],
|
||||
}
|
||||
chunks = chunk_control(cn, 2)
|
||||
assert len(chunks) == 2
|
||||
last_out = chunks[-1]["output"][0]
|
||||
# First row is the original third row; second row is padding zeros.
|
||||
assert torch.equal(last_out[0], cn["output"][0][2])
|
||||
assert torch.equal(last_out[1], torch.zeros(1, 2, 2))
|
||||
@@ -0,0 +1,228 @@
|
||||
"""Characterization tests for coreml_suite.models.CoreMLInputs.
|
||||
|
||||
Locks the shape transforms applied by chunks() and coreml_kwargs() for the
|
||||
four model variants the suite supports: SD1.5, LCM (SD1.5 + timestep_cond),
|
||||
SDXL base (time_ids len 6), and SDXL refiner (time_ids len 5).
|
||||
|
||||
These contracts feed the Core ML UNet at runtime; if a refactor silently
|
||||
re-shapes them, generation breaks.
|
||||
"""
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.inputs import CoreMLInputs
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _deterministic_seed():
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
|
||||
|
||||
# ---------- expected_inputs fixtures (mirror real model expectations) -------
|
||||
|
||||
SD15_EXPECTED = {
|
||||
"sample": {"shape": (2, 4, 64, 64)},
|
||||
"timestep": {"shape": (2,)},
|
||||
"encoder_hidden_states": {"shape": (2, 77, 768)},
|
||||
}
|
||||
|
||||
SD15_WITH_CN = {
|
||||
**SD15_EXPECTED,
|
||||
"additional_residual_0": {"shape": (2, 320, 64, 64)},
|
||||
"additional_residual_1": {"shape": (2, 640, 32, 32)},
|
||||
}
|
||||
|
||||
LCM_EXPECTED = {
|
||||
**SD15_EXPECTED,
|
||||
"timestep_cond": {"shape": (2, 256)},
|
||||
}
|
||||
|
||||
SDXL_BASE_EXPECTED = {
|
||||
"sample": {"shape": (2, 4, 128, 128)},
|
||||
"timestep": {"shape": (2,)},
|
||||
"encoder_hidden_states": {"shape": (2, 77, 2048)},
|
||||
"time_ids": {"shape": (2, 6)},
|
||||
"text_embeds": {"shape": (2, 1280)},
|
||||
}
|
||||
|
||||
SDXL_REFINER_EXPECTED = {
|
||||
"sample": {"shape": (2, 4, 128, 128)},
|
||||
"timestep": {"shape": (2,)},
|
||||
"encoder_hidden_states": {"shape": (2, 77, 1280)},
|
||||
"time_ids": {"shape": (2, 5)},
|
||||
"text_embeds": {"shape": (2, 1280)},
|
||||
}
|
||||
|
||||
|
||||
def _sd15_inputs(batch=1, with_control=False, with_ts_cond=False):
|
||||
x = torch.randn(batch, 4, 64, 64)
|
||||
t = torch.full((batch,), 999.0)
|
||||
context = torch.randn(batch, 77, 768)
|
||||
control = None
|
||||
if with_control:
|
||||
control = {
|
||||
"output": [torch.randn(batch, 320, 64, 64), torch.randn(batch, 640, 32, 32)],
|
||||
"middle": [],
|
||||
}
|
||||
kwargs = {}
|
||||
if with_ts_cond:
|
||||
kwargs["timestep_cond"] = torch.randn(batch, 256)
|
||||
return CoreMLInputs(x, t, context, control, **kwargs)
|
||||
|
||||
|
||||
def _sdxl_inputs(batch=1, refiner=False):
|
||||
x = torch.randn(batch, 4, 128, 128)
|
||||
t = torch.full((batch,), 999.0)
|
||||
ctx_dim = 1280 if refiner else 2048
|
||||
context = torch.randn(batch, 77, ctx_dim)
|
||||
time_ids_dim = 5 if refiner else 6
|
||||
time_ids = torch.randn(batch, time_ids_dim)
|
||||
text_embeds = torch.randn(batch, 1280)
|
||||
return CoreMLInputs(
|
||||
x, t, context, control=None, time_ids=time_ids, text_embeds=text_embeds
|
||||
)
|
||||
|
||||
|
||||
# ---------- coreml_kwargs ---------------------------------------------------
|
||||
|
||||
|
||||
def test_coreml_kwargs_sd15_shapes_and_fp16():
|
||||
out = _sd15_inputs(batch=1).coreml_kwargs(SD15_EXPECTED)
|
||||
assert set(out.keys()) == {"sample", "encoder_hidden_states", "timestep"}
|
||||
assert out["sample"].shape == (1, 4, 64, 64)
|
||||
assert out["sample"].dtype == np.float16
|
||||
# encoder_hidden_states keeps Comfy's native (b, seq, dim) layout.
|
||||
assert out["encoder_hidden_states"].shape == (1, 77, 768)
|
||||
assert out["encoder_hidden_states"].dtype == np.float16
|
||||
assert out["timestep"].shape == (1,)
|
||||
assert out["timestep"].dtype == np.float16
|
||||
|
||||
|
||||
def test_coreml_kwargs_sd15_with_controlnet_emits_residuals():
|
||||
inputs = _sd15_inputs(batch=1, with_control=True)
|
||||
out = inputs.coreml_kwargs(SD15_WITH_CN)
|
||||
assert "additional_residual_0" in out
|
||||
assert "additional_residual_1" in out
|
||||
assert out["additional_residual_0"].shape == (1, 320, 64, 64)
|
||||
assert out["additional_residual_1"].shape == (1, 640, 32, 32)
|
||||
|
||||
|
||||
def test_coreml_kwargs_sd15_without_controlnet_zero_fills_residuals():
|
||||
inputs = _sd15_inputs(batch=1, with_control=False)
|
||||
out = inputs.coreml_kwargs(SD15_WITH_CN)
|
||||
assert np.all(out["additional_residual_0"] == 0)
|
||||
assert np.all(out["additional_residual_1"] == 0)
|
||||
|
||||
|
||||
def test_coreml_kwargs_lcm_adds_timestep_cond():
|
||||
inputs = _sd15_inputs(batch=1, with_ts_cond=True)
|
||||
out = inputs.coreml_kwargs(LCM_EXPECTED)
|
||||
assert "timestep_cond" in out
|
||||
assert out["timestep_cond"].shape == (1, 256)
|
||||
assert out["timestep_cond"].dtype == np.float16
|
||||
|
||||
|
||||
def test_coreml_kwargs_lcm_skips_timestep_cond_when_not_provided():
|
||||
"""timestep_cond is only forwarded when the input supplied one — even if
|
||||
the model's expected_inputs lists it."""
|
||||
inputs = _sd15_inputs(batch=1, with_ts_cond=False)
|
||||
out = inputs.coreml_kwargs(LCM_EXPECTED)
|
||||
assert "timestep_cond" not in out
|
||||
|
||||
|
||||
def test_coreml_kwargs_sdxl_base_emits_time_ids_and_text_embeds():
|
||||
out = _sdxl_inputs(batch=1, refiner=False).coreml_kwargs(SDXL_BASE_EXPECTED)
|
||||
assert out["time_ids"].shape == (1, 6)
|
||||
assert out["text_embeds"].shape == (1, 1280)
|
||||
assert out["time_ids"].dtype == np.float16
|
||||
assert out["text_embeds"].dtype == np.float16
|
||||
|
||||
|
||||
def test_coreml_kwargs_sdxl_refiner_uses_len5_time_ids():
|
||||
out = _sdxl_inputs(batch=1, refiner=True).coreml_kwargs(SDXL_REFINER_EXPECTED)
|
||||
assert out["time_ids"].shape == (1, 5)
|
||||
|
||||
|
||||
# ---------- chunks ----------------------------------------------------------
|
||||
|
||||
|
||||
def test_chunks_sd15_pad_to_batch2_returns_one_chunk():
|
||||
chunked = _sd15_inputs(batch=1).chunks(SD15_EXPECTED)
|
||||
assert len(chunked) == 1
|
||||
c = chunked[0]
|
||||
assert c.x.shape == (2, 4, 64, 64)
|
||||
assert c.t.shape == (2,)
|
||||
# context shape: (b, seq, dim) padded along batch dim.
|
||||
assert c.context.shape == (2, 77, 768)
|
||||
assert c.control is None
|
||||
assert c.ts_cond is None
|
||||
assert c.time_ids is None
|
||||
assert c.text_embeds is None
|
||||
|
||||
|
||||
def test_chunks_sd15_with_controlnet_chunks_residuals_too():
|
||||
chunked = _sd15_inputs(batch=1, with_control=True).chunks(SD15_EXPECTED)
|
||||
assert len(chunked) == 1
|
||||
cn = chunked[0].control
|
||||
assert cn is not None
|
||||
assert cn["output"][0].shape == (2, 320, 64, 64)
|
||||
assert cn["output"][1].shape == (2, 640, 32, 32)
|
||||
|
||||
|
||||
def test_chunks_lcm_carries_timestep_cond_per_chunk():
|
||||
chunked = _sd15_inputs(batch=1, with_ts_cond=True).chunks(LCM_EXPECTED)
|
||||
assert len(chunked) == 1
|
||||
assert chunked[0].ts_cond is not None
|
||||
assert chunked[0].ts_cond.shape == (2, 256)
|
||||
|
||||
|
||||
def test_chunks_sdxl_base_propagates_time_ids_and_text_embeds():
|
||||
chunked = _sdxl_inputs(batch=1, refiner=False).chunks(SDXL_BASE_EXPECTED)
|
||||
assert len(chunked) == 1
|
||||
c = chunked[0]
|
||||
assert c.time_ids is not None and c.time_ids.shape == (2, 6)
|
||||
assert c.text_embeds is not None and c.text_embeds.shape == (2, 1280)
|
||||
|
||||
|
||||
def test_chunks_sdxl_refiner_uses_len5_time_ids():
|
||||
chunked = _sdxl_inputs(batch=1, refiner=True).chunks(SDXL_REFINER_EXPECTED)
|
||||
assert chunked[0].time_ids.shape == (2, 5)
|
||||
|
||||
|
||||
def test_chunks_sdxl_synthesizes_zero_time_ids_when_caller_omits():
|
||||
"""If the model expects time_ids but caller passed nothing, the suite
|
||||
fabricates a zero-filled tensor. Lock that fallback."""
|
||||
x = torch.randn(1, 4, 128, 128)
|
||||
t = torch.full((1,), 999.0)
|
||||
context = torch.randn(1, 77, 2048)
|
||||
inputs = CoreMLInputs(x, t, context, control=None)
|
||||
chunked = inputs.chunks(SDXL_BASE_EXPECTED)
|
||||
assert chunked[0].time_ids.shape == (2, 6)
|
||||
assert torch.equal(chunked[0].time_ids, torch.zeros(2, 6))
|
||||
assert chunked[0].text_embeds.shape == (2, 1280)
|
||||
assert torch.equal(chunked[0].text_embeds, torch.zeros(2, 1280))
|
||||
|
||||
|
||||
def test_chunks_splits_batch_into_multiple_target2_chunks():
|
||||
"""batch=5 with target_batch=2 -> 3 chunks (last padded)."""
|
||||
chunked = _sd15_inputs(batch=5).chunks(SD15_EXPECTED)
|
||||
assert len(chunked) == 3
|
||||
for c in chunked:
|
||||
assert c.x.shape == (2, 4, 64, 64)
|
||||
assert c.context.shape == (2, 77, 768)
|
||||
# Last chunk's second batch row is the zero-pad.
|
||||
assert torch.equal(chunked[-1].x[1], torch.zeros(4, 64, 64))
|
||||
|
||||
|
||||
def test_chunks_timestep_is_broadcast_from_first_value():
|
||||
"""t is rebuilt from t[0] across all chunks: locks current behavior that
|
||||
discards any per-row timestep variation."""
|
||||
x = torch.randn(2, 4, 64, 64)
|
||||
t = torch.tensor([42.0, 99.0]) # the second value will be lost
|
||||
context = torch.randn(2, 77, 768)
|
||||
inputs = CoreMLInputs(x, t, context, control=None)
|
||||
chunked = inputs.chunks(SD15_EXPECTED)
|
||||
assert chunked[0].t.shape == (2,)
|
||||
assert torch.equal(chunked[0].t, torch.full((2,), 42.0))
|
||||
@@ -0,0 +1,118 @@
|
||||
"""Characterization tests for coreml_suite.latents.
|
||||
|
||||
Locks the *current* behavior of chunk_batch / merge_chunks — including the
|
||||
zero-pad regions and the truncation in merge — so a refactor
|
||||
cannot silently shift either contract.
|
||||
"""
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.latents import chunk_batch, merge_chunks
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _deterministic_seed():
|
||||
torch.manual_seed(0)
|
||||
|
||||
|
||||
def _const_tensor(batch, *rest):
|
||||
return torch.arange(batch * 4 * 8 * 8, dtype=torch.float32).reshape(batch, 4, 8, 8)
|
||||
|
||||
|
||||
# ---------- chunk_batch ------------------------------------------------------
|
||||
|
||||
|
||||
def test_chunk_batch_passthrough_when_shape_matches():
|
||||
x = _const_tensor(2)
|
||||
out = chunk_batch(x, (2, 4, 8, 8))
|
||||
assert len(out) == 1
|
||||
# passthrough: the same object identity is returned (no copy).
|
||||
assert out[0] is x
|
||||
|
||||
|
||||
def test_chunk_batch_pads_single_chunk_when_input_smaller():
|
||||
"""batch=1, target=2 -> one padded chunk; the second row is exact zero."""
|
||||
x = _const_tensor(1)
|
||||
out = chunk_batch(x, (2, 4, 8, 8))
|
||||
assert len(out) == 1
|
||||
assert out[0].shape == (2, 4, 8, 8)
|
||||
assert torch.equal(out[0][0], x[0])
|
||||
assert torch.equal(out[0][1], torch.zeros(4, 8, 8))
|
||||
|
||||
|
||||
def test_chunk_batch_splits_exact_multiple():
|
||||
"""batch=4, target=2 -> two chunks, no padding."""
|
||||
x = _const_tensor(4)
|
||||
out = chunk_batch(x, (2, 4, 8, 8))
|
||||
assert len(out) == 2
|
||||
assert out[0].shape == (2, 4, 8, 8)
|
||||
assert out[1].shape == (2, 4, 8, 8)
|
||||
assert torch.equal(out[0], x[:2])
|
||||
assert torch.equal(out[1], x[2:])
|
||||
|
||||
|
||||
def test_chunk_batch_pads_remainder_chunk():
|
||||
"""batch=5, target=2 -> chunks=[x[0:2], x[2:4]] then [x[4], 0]."""
|
||||
x = _const_tensor(5)
|
||||
out = chunk_batch(x, (2, 4, 8, 8))
|
||||
assert len(out) == 3
|
||||
assert torch.equal(out[0], x[0:2])
|
||||
assert torch.equal(out[1], x[2:4])
|
||||
last = out[-1]
|
||||
assert last.shape == (2, 4, 8, 8)
|
||||
assert torch.equal(last[0], x[4])
|
||||
# The remainder row is zero-padded; lock that exact contract.
|
||||
assert torch.equal(last[1], torch.zeros(4, 8, 8))
|
||||
assert last[1].sum() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,target,expected_chunks",
|
||||
[
|
||||
(1, 4, 1),
|
||||
(3, 2, 2),
|
||||
(5, 3, 2),
|
||||
(9, 4, 3),
|
||||
],
|
||||
)
|
||||
def test_chunk_batch_pad_region_is_zero(batch_size, target, expected_chunks):
|
||||
x = _const_tensor(batch_size)
|
||||
out = chunk_batch(x, (target, 4, 8, 8))
|
||||
assert len(out) == expected_chunks
|
||||
mod = batch_size % target
|
||||
if mod == 0 and batch_size >= target:
|
||||
return
|
||||
last = out[-1]
|
||||
pad_rows = target - (mod if (mod != 0 and batch_size >= target) else batch_size)
|
||||
pad_region = last[-pad_rows:]
|
||||
assert torch.equal(pad_region, torch.zeros_like(pad_region))
|
||||
|
||||
|
||||
# ---------- merge_chunks -----------------------------------------------------
|
||||
|
||||
|
||||
def test_merge_chunks_exact_concat():
|
||||
x = _const_tensor(4)
|
||||
chunks = chunk_batch(x, (2, 4, 8, 8))
|
||||
merged = merge_chunks(chunks, x.shape)
|
||||
assert merged.shape == x.shape
|
||||
assert torch.equal(merged, x)
|
||||
|
||||
|
||||
def test_merge_chunks_truncates_padding():
|
||||
"""Round-trip with a padded last chunk drops the pad rows."""
|
||||
x = _const_tensor(5)
|
||||
chunks = chunk_batch(x, (2, 4, 8, 8))
|
||||
merged = merge_chunks(chunks, x.shape)
|
||||
assert merged.shape == x.shape
|
||||
assert torch.equal(merged, x)
|
||||
|
||||
|
||||
def test_merge_chunks_singleton_returns_equal_copy_when_shape_matches():
|
||||
"""A singleton chunk list still goes through torch.cat, so we get a new
|
||||
tensor equal to the input — locked here because a refactor might be tempted
|
||||
to short-circuit and accidentally return the same object."""
|
||||
x = _const_tensor(2)
|
||||
out = merge_chunks([x], x.shape)
|
||||
assert torch.equal(out, x)
|
||||
assert out is not x
|
||||
@@ -0,0 +1,127 @@
|
||||
"""Characterization tests for the SDXL options math.
|
||||
|
||||
The SDXL time_ids / text_embeds math lives in
|
||||
coreml_suite.core.sdxl as pure builders. The framework adapter
|
||||
add_sdxl_model_options lives in models.py; here we just lock the pure
|
||||
math.
|
||||
"""
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.sdxl import (
|
||||
build_sdxl_text_embeds,
|
||||
build_sdxl_time_ids,
|
||||
sdxl_model_function_wrapper,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _deterministic_seed():
|
||||
torch.manual_seed(0)
|
||||
|
||||
|
||||
# ---------- build_sdxl_time_ids: base (len 6) -------------------------------
|
||||
|
||||
|
||||
def test_build_time_ids_base_defaults():
|
||||
out = build_sdxl_time_ids({}, {}, is_base=True, is_refiner=False)
|
||||
expected = torch.tensor([[768, 768, 0, 0, 768, 768], [768, 768, 0, 0, 768, 768]])
|
||||
assert out.shape == (2, 6)
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
|
||||
def test_build_time_ids_base_respects_overrides():
|
||||
pos = {"height": 1024, "width": 512, "crop_h": 8, "crop_w": 4,
|
||||
"target_height": 1024, "target_width": 1024}
|
||||
neg = {"height": 256, "width": 256, "crop_h": 0, "crop_w": 0,
|
||||
"target_height": 256, "target_width": 256}
|
||||
out = build_sdxl_time_ids(pos, neg, is_base=True, is_refiner=False)
|
||||
expected = torch.tensor([[1024, 512, 8, 4, 1024, 1024], [256, 256, 0, 0, 256, 256]])
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
|
||||
# ---------- build_sdxl_time_ids: refiner (len 5) ----------------------------
|
||||
|
||||
|
||||
def test_build_time_ids_refiner_defaults():
|
||||
out = build_sdxl_time_ids({}, {}, is_base=False, is_refiner=True)
|
||||
expected = torch.tensor([[768, 768, 0, 0, 6.0], [768, 768, 0, 0, 2.5]])
|
||||
assert out.shape == (2, 5)
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
|
||||
def test_build_time_ids_refiner_respects_aesthetic_score():
|
||||
pos = {"aesthetic_score": 8.5}
|
||||
neg = {"aesthetic_score": 1.5}
|
||||
out = build_sdxl_time_ids(pos, neg, is_base=False, is_refiner=True)
|
||||
expected = torch.tensor([[768, 768, 0, 0, 8.5], [768, 768, 0, 0, 1.5]])
|
||||
assert torch.equal(out, expected)
|
||||
|
||||
|
||||
# ---------- build_sdxl_time_ids: edge case ----------------------------------
|
||||
|
||||
|
||||
def test_build_time_ids_neither_base_nor_refiner_returns_len4():
|
||||
out = build_sdxl_time_ids({}, {}, is_base=False, is_refiner=False)
|
||||
assert out.shape == (2, 4)
|
||||
|
||||
|
||||
# ---------- build_sdxl_text_embeds ------------------------------------------
|
||||
|
||||
|
||||
def test_text_embeds_concat_pos_then_neg():
|
||||
pos = torch.full((1, 1280), 1.0)
|
||||
neg = torch.full((1, 1280), -1.0)
|
||||
out = build_sdxl_text_embeds(pos, neg)
|
||||
assert out.shape == (2, 1280)
|
||||
assert torch.equal(out[0], pos[0])
|
||||
assert torch.equal(out[1], neg[0])
|
||||
|
||||
|
||||
# ---------- sdxl_model_function_wrapper closure -----------------------------
|
||||
|
||||
|
||||
def test_wrapper_captures_time_ids_text_embeds_refiner_via_closure():
|
||||
time_ids = torch.zeros(2, 6)
|
||||
text_embeds = torch.zeros(2, 1280)
|
||||
wrapper = sdxl_model_function_wrapper(time_ids, text_embeds, refiner=False)
|
||||
closure = inspect.getclosurevars(wrapper).nonlocals
|
||||
assert closure["time_ids"] is time_ids
|
||||
assert closure["text_embeds"] is text_embeds
|
||||
assert closure["refiner"] is False
|
||||
|
||||
|
||||
def test_wrapper_returns_zero_when_context_missing():
|
||||
"""When c_crossattn is None the wrapper short-circuits to zeros_like(x).
|
||||
Locked here because the refactor mustn't change this default."""
|
||||
wrapper = sdxl_model_function_wrapper(torch.zeros(2, 6), torch.zeros(2, 1280))
|
||||
x = torch.randn(2, 4, 16, 16)
|
||||
out = wrapper(
|
||||
model_function=lambda *a, **kw: pytest.fail("model_function must not run"),
|
||||
params={"input": x, "timestep": torch.zeros(2), "c": {}},
|
||||
)
|
||||
assert torch.equal(out, torch.zeros_like(x))
|
||||
|
||||
|
||||
def test_wrapper_refiner_truncates_context_to_g_clip():
|
||||
"""refiner=True slices c_crossattn[:, :, 768:] before forwarding."""
|
||||
captured = {}
|
||||
|
||||
def fake_model(x, t, **c):
|
||||
captured["context_shape"] = c["c_crossattn"].shape
|
||||
captured["time_ids_shape"] = c["time_ids"].shape
|
||||
return x
|
||||
|
||||
wrapper = sdxl_model_function_wrapper(
|
||||
torch.zeros(2, 5), torch.zeros(2, 1280), refiner=True
|
||||
)
|
||||
x = torch.randn(2, 4, 16, 16)
|
||||
context = torch.randn(2, 77, 2048) # 768 + 1280 dims
|
||||
wrapper(
|
||||
model_function=fake_model,
|
||||
params={"input": x, "timestep": torch.zeros(2), "c": {"c_crossattn": context}},
|
||||
)
|
||||
assert captured["context_shape"] == (2, 77, 1280)
|
||||
assert captured["time_ids_shape"] == (2, 5)
|
||||
@@ -0,0 +1,122 @@
|
||||
"""Smoke tests for the pure batch-chunking helpers in coreml_suite.core.
|
||||
|
||||
Uses torch.device('cpu') instead of comfy.model_management.get_torch_device
|
||||
so Tier 0 runs without ComfyUI.
|
||||
"""
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from coreml_suite.core.controlnet import chunk_control
|
||||
from coreml_suite.core.inputs import CoreMLInputs
|
||||
from coreml_suite.core.latents import chunk_batch, merge_chunks
|
||||
|
||||
|
||||
CPU = torch.device("cpu")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def expected_inputs():
|
||||
return {
|
||||
"sample": {"shape": (2, 4, 64, 64)},
|
||||
"timestep": {"shape": (2,)},
|
||||
"timestep_cond": {"shape": (2, 256)},
|
||||
"encoder_hidden_states": {"shape": (2, 77, 768)},
|
||||
"additional_residual_0": {"shape": (2, 320, 64, 64)},
|
||||
"additional_residual_1": {"shape": (2, 640, 32, 32)},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 5, 9])
|
||||
def test_batch_chunking(batch_size):
|
||||
latent_image = torch.randn(batch_size, 4, 64, 64).to(CPU)
|
||||
target_shape = (4, 4, 64, 64)
|
||||
|
||||
chunked = chunk_batch(latent_image, target_shape)
|
||||
|
||||
for chunk in chunked:
|
||||
assert chunk.shape == target_shape
|
||||
|
||||
if batch_size % target_shape[0] != 0:
|
||||
assert chunked[-1][batch_size % target_shape[0] :].sum() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2, 4, 5, 9])
|
||||
def test_merge_chunks(batch_size):
|
||||
input_tensor = torch.randn(batch_size, 4, 64, 64).to(CPU)
|
||||
target_shape = (4, 4, 64, 64)
|
||||
chunked = chunk_batch(input_tensor, target_shape)
|
||||
|
||||
merged = merge_chunks(chunked, input_tensor.shape)
|
||||
|
||||
assert merged.shape == input_tensor.shape
|
||||
assert torch.equal(input_tensor, merged)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def inputs():
|
||||
x = torch.randn(1, 4, 64, 64).to(CPU)
|
||||
t = torch.randn([1]).to(CPU)
|
||||
c_crossattn = torch.randn(1, 77, 768).to(CPU)
|
||||
control = {
|
||||
"output": [
|
||||
torch.randn(1, 320, 64, 64).to(CPU),
|
||||
torch.randn(1, 640, 32, 32).to(CPU),
|
||||
],
|
||||
}
|
||||
timestep_cond = torch.randn(1, 256).to(CPU)
|
||||
|
||||
return CoreMLInputs(x, t, c_crossattn, control, timestep_cond=timestep_cond)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"b, target_size, num_chunks",
|
||||
[
|
||||
(1, 2, 1),
|
||||
(1, 1, 1),
|
||||
(2, 2, 1),
|
||||
(3, 2, 2),
|
||||
(4, 2, 2),
|
||||
(5, 3, 2),
|
||||
(9, 4, 3),
|
||||
],
|
||||
)
|
||||
def test_chunking_controlnet(b, target_size, num_chunks):
|
||||
cn = {
|
||||
"output": [
|
||||
torch.randn(b, 320, 64, 64).to(CPU),
|
||||
torch.randn(b, 640, 32, 32).to(CPU),
|
||||
],
|
||||
"middle": [
|
||||
torch.randn(b, 1280, 8, 8).to(CPU),
|
||||
],
|
||||
}
|
||||
|
||||
chunked = chunk_control(cn, target_size)
|
||||
|
||||
assert len(chunked) == num_chunks
|
||||
for chunk in chunked:
|
||||
assert chunk["output"][0].shape == (target_size, 320, 64, 64)
|
||||
assert chunk["output"][1].shape == (target_size, 640, 32, 32)
|
||||
assert chunk["middle"][0].shape == (target_size, 1280, 8, 8)
|
||||
|
||||
|
||||
def test_chunking_no_control():
|
||||
cn = None
|
||||
target_size = 2
|
||||
|
||||
chunked = chunk_control(cn, target_size)
|
||||
|
||||
assert chunked == [None, None]
|
||||
|
||||
|
||||
def test_chunking_inputs(expected_inputs, inputs):
|
||||
chunked = inputs.chunks(expected_inputs)
|
||||
|
||||
assert len(chunked) == 1
|
||||
|
||||
assert chunked[0].x.shape == (2, 4, 64, 64)
|
||||
assert chunked[0].t.shape == (2,)
|
||||
assert chunked[0].context.shape == (2, 77, 768)
|
||||
assert chunked[0].control["output"][0].shape == (2, 320, 64, 64)
|
||||
assert chunked[0].control["output"][1].shape == (2, 640, 32, 32)
|
||||
assert chunked[0].ts_cond.shape == (2, 256)
|
||||
@@ -0,0 +1,16 @@
|
||||
from coreml_suite.controlnet import no_control
|
||||
|
||||
|
||||
def test_no_control():
|
||||
expected_inputs = {
|
||||
"additional_residual_0": {"shape": (2, 2, 2)},
|
||||
"additional_residual_1": {"shape": (2, 4, 4)},
|
||||
"additional_residual_2": {"shape": (2, 8, 8)},
|
||||
}
|
||||
|
||||
residual_kwargs = no_control(expected_inputs)
|
||||
|
||||
assert len(residual_kwargs) == 3
|
||||
assert residual_kwargs["additional_residual_0"].shape == (2, 2, 2)
|
||||
assert residual_kwargs["additional_residual_1"].shape == (2, 4, 4)
|
||||
assert residual_kwargs["additional_residual_2"].shape == (2, 8, 8)
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Gate: prove the Tier-0 lane is framework-free.
|
||||
|
||||
In a pure `pytest -m unit` run, none of the banned runtime modules
|
||||
(comfy, coremltools, python_coreml_stable_diffusion, folder_paths,
|
||||
nodes, comfy_extras, diffusers, diffusionkit) may be in sys.modules
|
||||
after collection. If they are, a tests/unit/ file is transitively
|
||||
pulling them in and the Tier-0 promise — "runs on Linux with no Mac
|
||||
stack" — is broken.
|
||||
|
||||
When other tiers are also collected, framework modules may be imported
|
||||
deliberately (e.g. smoke pulls in coremltools), so the check is skipped
|
||||
unless the run is purely `-m unit` — Tier-0 purity is only meaningful
|
||||
when nothing else is loaded.
|
||||
"""
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
BANNED_ROOTS = {
|
||||
"comfy",
|
||||
"comfy_extras",
|
||||
"coremltools",
|
||||
"python_coreml_stable_diffusion",
|
||||
"folder_paths",
|
||||
"nodes",
|
||||
"diffusers",
|
||||
"diffusionkit",
|
||||
}
|
||||
|
||||
|
||||
def test_no_framework_modules_loaded_by_unit_tier(request):
|
||||
markexpr = request.config.option.markexpr
|
||||
if markexpr != "unit":
|
||||
pytest.skip(
|
||||
"purity gate only meaningful in a pure `-m unit` run "
|
||||
f"(got markexpr={markexpr!r}); other tiers are expected to "
|
||||
"import comfy/coremltools."
|
||||
)
|
||||
loaded = {name for name in sys.modules if name.split(".")[0] in BANNED_ROOTS}
|
||||
assert not loaded, (
|
||||
f"Tier-0 leakage: these framework modules are in sys.modules after "
|
||||
f"collecting tests/unit/: {sorted(loaded)}. Pure-core promise broken."
|
||||
)
|
||||