docs: sync with converter consolidation

Docs described the pre-#67 suite: a removed standalone LCM converter
node, a removed model_version input, a nonexistent
CoreMLDetailerHookProvider node, and a deleted lcm/converter.py file.

- drop LCM converter node docs; LCM checkpoints are auto-detected by
  the consolidated CoreMLConverter
- remove model_version input and invented 512-768 resolution range
- replace phantom detailer-hook fix with real workarounds
- fix Python support claim (3.12+ per requires-python)
- update LCM support-matrix row, drop stale line reference

Refs #67, #68

.github/workflows/tier0.yml

@@ -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

.github/workflows/tier2.yml

@@ -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

.gitignore

@@ -1,3 +1,8 @@
 playground/
-experiments/
 __pycache__/
+models/
+.venv/
+test_results/
+*.log
+.DS_Store
+.claude/

.python-version

@@ -0,0 +1 @@
+3.12

LICENSE

@@ -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.

README.md

@@ -1,389 +1,113 @@
 # Core ML Suite for ComfyUI
 
-## Overview
+Custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that run
+Stable Diffusion UNets as [Core ML](https://developer.apple.com/documentation/coreml)
+models on Apple Silicon (M1/M2/M3). Core ML can use the Apple Neural Engine
+(ANE), which is unavailable to PyTorch — on an M2 Pro 32 GB, SD1.5 at 512×512
+generates roughly **1.5–2× faster** than the standard PyTorch/MPS path.
 
-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.
+You convert a Stable Diffusion checkpoint to a Core ML model with the nodes in
+this suite, then sample from it like any other ComfyUI workflow.
 
-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).
-
-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,
-the use of Core ML models sped up the generation of 512x512 images by a factor
-of approximately 1.5 to 2 times.
-
-## Getting Started
-
-To start using custom nodes in your ComfyUI, follow these simple steps:
-
-1. Clone or download this repository: You can do this directly into the custom_nodes directory of your ComfyUI.
-2. Install the dependencies: You'll need to use a package manager like pip to do this.
-
-That's it! You're now ready to start enhancing your ComfyUI workflows with Core ML models.
-
-- Check [Installation](#installation) for more details on installation.
-- Check [How to use](#how-to-use) for more details on how to use the custom nodes.
-- Check [Example Workflows](#example-workflows) for some example workflows.
-
-## Glossary
-
-- **Core ML**: A machine learning framework developed by Apple. It's used to run machine learning models on Apple
-  devices.
-- **Core ML Model**: A machine learning model that can be run on Apple devices using Core ML.
-- **mlmodelc**: A compiled Core ML model. This is the recommended format for Core ML models.
-- **mlpackage**: A Core ML model packaged in a directory. This is the default format for Core ML models.
-- **ANE**: Apple Neural Engine. A hardware accelerator for machine learning tasks on Apple devices.
-- **Compute Unit**: A Core ML option that allows you to specify the hardware on which the model should run.
-    - **CPU_AND_ANE**: A Core ML compute unit option that allows the model to run on both the CPU and ANE. This is the
-      default option.
-    - **CPU_AND_GPU**: A Core ML compute unit option that allows the model to run on both the CPU and GPU.
-    - **CPU_ONLY**: A Core ML compute unit option that allows the model to run on the CPU only.
-    - **ALL**: A Core ML compute unit option that allows the model to run on all available hardware.
-- **CLIP**: Contrastive Language-Image Pre-training. A model that learns visual concepts from natural language
-  supervision. It's used as a text encoder in Stable Diffusion.
-- **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:
-> For the model to run on the ANE, the model must be converted with the `--attention-implementation SPLIT_EINSUM`
-> option.
-> Models converted with `--attention-implementation ORIGINAL` will run on GPU instead of ANE.
-
-## Features
-
-These custom nodes come with a host of features, including:
-
-- Loading Core ML Unet models
-- Support for ControlNet
-- 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.
-> For the best experience, I recommend using the compiled models
-> (.mlmodelc files) instead of the .mlpackage files.
-
-> [!NOTE]  
-> This repository will continue to be updated with more nodes and features over time.
+> [!IMPORTANT]
+> **Convert your own checkpoints — that is the only supported path.** This
+> suite uses its own input dimensions, naming convention, and metadata
+> (produced by the [coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion)
+> package). Pre-converted Core ML models from elsewhere (e.g. the
+> coreml-community Hugging Face org) are **not** supported. Conversion is cheap
+> and runs on your machine, so there is no need to download Core ML models.
 
 ## Installation
 
-### Using ComfyUI-Manager
+### ComfyUI-Manager (recommended)
 
-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 **Manager → Install Custom Nodes**, search for `Core ML`, click
+**Install**, and restart ComfyUI.
 
-- 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
 
-### Manual Installation
+```bash
+cd /path/to/comfyui/custom_nodes
+git clone https://github.com/aszc-dev/ComfyUI-CoreMLSuite.git
+cd ComfyUI-CoreMLSuite
+pip install -r requirements.txt
+```
 
-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.
-    ```bash
-    cd /path/to/comfyui/custom_nodes
-    git clone https://github.com/aszc-dev/ComfyUI-CoreMLSuite.git
-    ```
-2. Next, install the required dependencies using pip or another package manager:
+Dependencies (`coreml-diffusion`, `coremltools`, `numpy`, `diffusers`) install
+from PyPI. PyTorch is intentionally **not** pinned — it is provided by your
+ComfyUI host, and a hard cap here would downgrade it and break ComfyUI.
 
-    ```bash
-    cd /path/to/comfyui/custom_nodes/ComfyUI-CoreMLSuite
-    pip install -r requirements.txt
-    ```
+## Quickstart
 
-## How to use
+1. Put a SD1.5 checkpoint in `models/checkpoints`.
+2. Add the **Convert Checkpoint to Core ML** node, select the checkpoint, and
+   queue once. It writes a `.mlpackage` to `models/unet` (cached by name — it
+   won't reconvert next time).
+3. Sample with the **Core ML Sampler** node, decoding the latent with a normal
+   VAE Decode. CLIP and VAE come from standard ComfyUI nodes.
 
-Once you've installed the custom nodes, you can start using them in your ComfyUI workflows.
-To do this, you need to add the nodes to your workflow. You can do this by right-clicking on the workflow canvas and
-selecting the nodes from the list of available nodes (the nodes are in the `Core ML Suite` category).
-You can also double-click the canvas and use the search bar to find the nodes. The list of available nodes is given
-below.
+See [docs/workflows.md](docs/workflows.md) for complete example graphs (txt2img,
+ControlNet, LoRA, LCM, SDXL).
 
-### Available Nodes
+## Which compute unit should I pick?
 
-#### Core ML UNet Loader (`CoreMLUnetLoader`)
+The **compute unit** selects the hardware Core ML runs on. Pair it with the
+attention implementation chosen at conversion time:
 
-![CoreMLUnetLoader](./assets/unet_loader.png?raw=true)
+| Model | Convert with | Load with | Runs on |
+|---|---|---|---|
+| SD1.5 @ 512×512 | `SPLIT_EINSUM` | `CPU_AND_NE` | Neural Engine (fastest) |
+| SD1.5 @ larger sizes | `ORIGINAL` | `CPU_AND_GPU` | GPU |
+| SDXL | `ORIGINAL` | `CPU_AND_GPU` | GPU (ANE unsupported) |
 
-This node allows you to load a Core ML UNet model and use it in your ComfyUI workflow. Place the converted
-.mlpackage or .mlmodelc file in ComfyUI's `models/unet` directory and use the node to load the model. The output of the
-node is a `coreml_model` object that can be used with the Core ML Sampler.
+`CPU_AND_NE` is usually the fastest option for SD1.5 — often faster than `ALL`.
+This suite uses Core ML compute units only; it never touches PyTorch MPS, so
+`PYTORCH_ENABLE_MPS_FALLBACK` is irrelevant to these nodes. Full reasoning and
+benchmarks: [docs/hardware.md](docs/hardware.md).
 
-- **Inputs**:
-    - **model_name**: The name of the model to load. This should be the name of the .mlpackage or .mlmodelc file.
-    - **compute_unit**: The hardware on which the model should run. This can be one of the following:
-        - `CPU_AND_ANE`: The model will run on both the CPU and ANE. This is the default option. It works best with
-          models
-          converted with `--attention-implementation SPLIT_EINSUM` or `--attention-implementation SPLIT_EINSUM_V2`.
-        - `CPU_AND_GPU`: The model will run on both the CPU and GPU. It works best with models converted with
-          `--attention-implementation ORIGINAL`.
-        - `CPU_ONLY`: The model will run on the CPU only.
-        - `ALL`: The model will run on all available hardware.
-- **Outputs**:
-    - **coreml_model**: A Core ML model that can be used with the Core ML Sampler.
+## Documentation
 
-#### Core ML Sampler (`CoreMLSampler`)
+- [Hardware & compute units](docs/hardware.md) — ANE vs GPU vs MPS, attention
+  implementations, which to choose.
+- [Nodes](docs/nodes.md) — full reference for every node.
+- [Conversion](docs/conversion.md) — how conversion works, caching,
+  quantization.
+- [Example workflows](docs/workflows.md) — annotated example graphs.
+- [FAQ](docs/faq.md) — answers to common questions.
+- [Troubleshooting](docs/troubleshooting.md) — common errors and fixes.
+- [Limitations & support matrix](docs/limitations.md) — what is and isn't
+  supported.
 
-![CoreMLSampler](./assets/sampler.png?raw=true)
+## Glossary
 
-This node allows you to generate images using a Core ML model. The node takes a Core ML model as input and outputs a
-latent image similar to the latent image output by the KSampler. This means that you can use the
-resulting latent as you normally would in your workflow.
-
-- **Inputs**:
-    - **coreml_model**: The Core ML model to use for sampling. This should be the output of the Core ML UNet Loader.
-    - **latent_image** [optional]: The latent image to use for sampling. If provided, should be of the same size as the
-      input of the Core ML model. If not provided, the node will create a latent suitable for the Core ML model used.
-      Useful in img2img workflows.
-    - ... _(the rest of the inputs are the same as the KSampler)_
-- **Outputs**:
-    - **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
-
-![CoreMLConverter](./assets/checkpoint_converter.png?raw=true)
-
-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. 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.
-    - **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
-
-![LoRALoader](./assets/lora_loader.png?raw=true)
-
-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
-
-![LCMConverter](./assets/lcm_converter.png?raw=true)
-
-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`)
-
-![CoreMLModelAdapter](./assets/adapter.png?raw=true)
-
-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]
-> The models used are just an example. Feel free to experiment with different models and see what works best for you.
-
-#### Basic txt2img with Core ML UNet loader
-
-This is a basic txt2img workflow that uses the Core ML UNet loader to load a model. The CLIP and VAE models
-are loaded using the standard ComfyUI nodes. In the first example, the text encoder (CLIP) and VAE models are loaded
-separately. In the second example, the text encoder and VAE models are loaded from the checkpoint file. Note that you
-can use any CLIP or VAE model as long as it's compatible with Stable Diffusion v1.5.
-
-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).
-      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).
-      Once downloaded, place the model in the `models/unet` directory.  
-      ![coreml-unet+clip+vae](./assets/unet+sampler+clip+vae.png?raw=true)
-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).
-      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).
-      Once downloaded, place the model in the `models/unet` directory.  
-      ![coreml-unet+checkpoint](./assets/unet+sampler+checkpoint.png?raw=true)
-
-#### ControlNet with Core ML UNet loader
-
-This workflow uses the Core ML UNet loader to load a Core ML UNet model that supports ControlNet. The ControlNet is
-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/control_v11p_sd15_scribble/blob/main/diffusion_pytorch_model.fp16.safetensors).
-Once downloaded, place the model in the `models/controlnet` directory.
-![coreml-unet+controlnet](./assets/unet+sampler+controlnet.png?raw=true)
-
-#### Checkpoint conversion
-
-This workflow uses the Checkpoint Converter to convert the checkpoint file. See
-[Checkpoint Converter](#checkpoint-converter) description for more details.
-
-![checkpoint-converter](./assets/basic_conversion.png?raw=true)
-
-#### 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.
-
-![checkpoint-converter+lora](./assets/conversion+lora.png?raw=true)
-
-#### 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.
-
-![multiple-loras](./assets/conversion+lcm_lora.png?raw=true)
-
-#### 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.
-
-![loader+lora](./assets/loader+lcm_lora.png?raw=true)
-
-#### 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.
-![lcm+controlnet](./assets/lcm+controlnet.png?raw=true)
-
-#### 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)
+- **Core ML** — Apple's on-device machine-learning framework.
+- **`.mlpackage`** — the Core ML model format this suite produces and loads.
+- **ANE** — Apple Neural Engine, a hardware accelerator for ML.
+- **Compute unit** — which hardware Core ML uses (`CPU_AND_NE`, `CPU_AND_GPU`,
+  `CPU_ONLY`, `ALL`).
+- **Attention implementation** — `SPLIT_EINSUM` / `SPLIT_EINSUM_V2` (ANE-friendly)
+  or `ORIGINAL` (GPU-friendly), chosen at conversion.
 
 > [!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.
+> **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)`.
+> Models converted with earlier versions are not compatible and must be
+> re-converted.
 
-![sdxl](./assets/sdxl_conversion.png?raw=true)
+## Acknowledgements
 
-## 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.
-- SD2.1 models are not supported.
-
-[^1]:
-Unless [EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes)
-is used during conversion. Needs more testing.
+The conversion pipeline began as an adaptation of Apple's
+[ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion), which
+pioneered running Stable Diffusion on the Neural Engine. It has since diverged
+and no longer depends on that package: UNet conversion runs natively on
+`diffusers`' `UNet2DConditionModel`, the ANE attention path (`SPLIT_EINSUM`,
+`SPLIT_EINSUM_V2`) is reimplemented as standalone `diffusers` attention
+processors, and the toolchain tracks current ComfyUI (NumPy 2, Torch 2.7+,
+coremltools 9, Python 3.12+). Conversion now lives in the separate
+[coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) package.
 
 ## Support
 
-I'm here to help! If you have any questions or suggestions, don't hesitate to open an issue and I'll do my best
-to assist you.
+Questions or suggestions? Open an
+[issue](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/issues).

__init__.py

@@ -11,9 +11,6 @@ from coreml_suite.nodes import (
     CoreMLConverter,
     COREML_LOAD_LORA,
 )
-from coreml_suite.lcm import (
-    COREML_CONVERT_LCM,
-)
 
 NODE_CLASS_MAPPINGS = {
     "CoreMLUNetLoader": CoreMLLoaderUNet,
@@ -22,7 +19,6 @@ NODE_CLASS_MAPPINGS = {
     "CoreMLModelAdapter": CoreMLModelAdapter,
     "Core ML LoRA Loader": COREML_LOAD_LORA,
     "Core ML Converter": CoreMLConverter,
-    "Core ML LCM Converter": COREML_CONVERT_LCM,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "CoreMLUNetLoader": "Load Core ML UNet",
@@ -31,5 +27,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CoreMLModelAdapter": "Core ML Adapter (Experimental)",
     "Core ML LoRA Loader": "Load LoRA to use with Core ML",
     "Core ML Converter": "Convert Checkpoint to Core ML",
-    "Core ML LCM Converter": "Convert LCM to Core ML",
 }

conftest.py

@@ -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"]

constraints/comfy-ceiling.txt

@@ -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

coreml_suite/config.py

@@ -1,17 +1,10 @@
-from enum import Enum
-
 import torch
 
 from comfy import supported_models_base
 from comfy import latent_formats
 from comfy.model_detection import convert_config
 
-
-class ModelVersion(Enum):
-    SD15 = "sd15"
-    SDXL = "sdxl"
-    SDXL_REFINER = "sdxl_refiner"
-    LCM = "lcm"
+from coreml_diffusion import ModelVersion
 
 
 config_map = {

coreml_suite/controlnet.py

@@ -1,62 +1,14 @@
-from itertools import chain
-from math import ceil
+"""Compatibility shim — re-exports from coreml_suite.core.controlnet."""
+from coreml_suite.core.controlnet import (
+    chunk_control,
+    expand_inputs,
+    extract_residual_kwargs,
+    no_control,
+)
 
-import numpy as np
-import torch
-
-from coreml_suite.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
+__all__ = [
+    "chunk_control",
+    "expand_inputs",
+    "extract_residual_kwargs",
+    "no_control",
+]

coreml_suite/converter.py

@@ -1,362 +0,0 @@
-import gc
-import os
-import shutil
-import time
-from typing import Union
-
-import coremltools as ct
-import numpy as np
-import python_coreml_stable_diffusion.unet
-import torch
-from diffusers import (
-    StableDiffusionPipeline,
-    LatentConsistencyModelPipeline,
-    StableDiffusionXLPipeline,
-)
-from python_coreml_stable_diffusion.unet import (
-    UNet2DConditionModel,
-    UNet2DConditionModelXL,
-    AttentionImplementations,
-)
-
-from coreml_suite.config import ModelVersion
-from coreml_suite.lcm.unet import UNet2DConditionModelLCM
-from coreml_suite.logger import logger
-from folder_paths import get_folder_paths
-
-
-class StableDiffusionLCMPipeline(LatentConsistencyModelPipeline):
-    pass
-
-
-MODEL_TYPE_TO_UNET_CLS = {
-    ModelVersion.SD15: UNet2DConditionModel,
-    ModelVersion.SDXL: UNet2DConditionModelXL,
-    ModelVersion.LCM: UNet2DConditionModelLCM,
-}
-
-MODEL_TYPE_TO_PIPE_CLS = {
-    ModelVersion.SD15: StableDiffusionPipeline,
-    ModelVersion.SDXL: StableDiffusionXLPipeline,
-    ModelVersion.LCM: StableDiffusionLCMPipeline,
-}
-
-
-def get_unet(model_type: ModelVersion, ref_pipe):
-    ref_unet = ref_pipe.unet
-
-    unet_cls = MODEL_TYPE_TO_UNET_CLS[model_type]
-    cml_unet = unet_cls.from_config(ref_unet.config).eval()
-    cml_unet.load_state_dict(ref_unet.state_dict(), strict=False)
-
-    return cml_unet
-
-
-def get_encoder_hidden_states_shape(ref_pipe, batch_size):
-    text_encoder = (
-        ref_pipe.text_encoder_2
-        if hasattr(ref_pipe, "text_encoder_2")
-        else ref_pipe.text_encoder
-    )
-
-    text_token_sequence_length = text_encoder.config.max_position_embeddings
-    hidden_size = (text_encoder.config.hidden_size,)
-
-    encoder_hidden_states_shape = (
-        batch_size,
-        ref_pipe.unet.config.cross_attention_dim or hidden_size,
-        1,
-        text_token_sequence_length,
-    )
-
-    return encoder_hidden_states_shape
-
-
-def get_coreml_inputs(sample_inputs):
-    coreml_sample_unet_inputs = {
-        k: v.numpy().astype(np.float16) for k, v in sample_inputs.items()
-    }
-    return [
-        ct.TensorType(
-            name=k,
-            shape=v.shape,
-            dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
-        )
-        for k, v in coreml_sample_unet_inputs.items()
-    ]
-
-
-def load_coreml_model(out_path):
-    logger.info(f"Loading model from {out_path}")
-
-    start = time.time()
-    coreml_model = ct.models.MLModel(out_path)
-    logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds")
-
-    return coreml_model
-
-
-def convert_to_coreml(
-    submodule_name, torchscript_module, sample_inputs, output_names, out_path
-):
-    if os.path.exists(out_path):
-        logger.info(f"Skipping export because {out_path} already exists")
-        coreml_model = load_coreml_model(out_path)
-    else:
-        logger.info(f"Converting {submodule_name} to CoreML..")
-        coreml_model = ct.convert(
-            torchscript_module,
-            convert_to="mlprogram",
-            minimum_deployment_target=ct.target.macOS13,
-            inputs=sample_inputs,
-            outputs=[
-                ct.TensorType(name=name, dtype=np.float32) for name in output_names
-            ],
-            skip_model_load=True,
-        )
-
-        del torchscript_module
-        gc.collect()
-
-    return coreml_model
-
-
-def get_out_path(submodule_name, model_name):
-    fname = f"{model_name}_{submodule_name}.mlpackage"
-    unet_path = get_folder_paths(submodule_name)[0]
-    out_path = os.path.join(unet_path, fname)
-    return out_path
-
-
-def compile_coreml_model(source_model_path, output_dir, final_name):
-    """Compiles Core ML models using the coremlcompiler utility from Xcode toolchain"""
-    target_path = os.path.join(output_dir, f"{final_name}.mlmodelc")
-    if os.path.exists(target_path):
-        logger.warning(f"Found existing compiled model at {target_path}! Skipping..")
-        return target_path
-
-    logger.info(f"Compiling {source_model_path}")
-    source_model_name = os.path.basename(os.path.splitext(source_model_path)[0])
-
-    os.system(f"xcrun coremlcompiler compile {source_model_path} {output_dir}")
-    compiled_output = os.path.join(output_dir, f"{source_model_name}.mlmodelc")
-    shutil.move(compiled_output, target_path)
-
-    return target_path
-
-
-def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape, scheduler):
-    sample_unet_inputs = dict(
-        [
-            ("sample", torch.rand(*sample_shape)),
-            (
-                "timestep",
-                torch.tensor([scheduler.timesteps[0].item()] * batch_size).to(
-                    torch.float32
-                ),
-            ),
-            ("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)),
-        ]
-    )
-    return sample_unet_inputs
-
-
-def lcm_inputs(sample_unet_inputs):
-    batch_size = sample_unet_inputs["sample"].shape[0]
-    return {"timestep_cond": torch.randn(batch_size, 256).to(torch.float32)}
-
-
-def sdxl_inputs(sample_unet_inputs, ref_pipe):
-    sample_shape = sample_unet_inputs["sample"].shape
-    batch_size = sample_shape[0]
-    h = sample_shape[2] * 8
-    w = sample_shape[3] * 8
-    original_size = (h, w)
-    crops_coords_top_left = (0, 0)
-
-    is_refiner = (
-        hasattr(ref_pipe.config, "requires_aesthetics_score")
-        and ref_pipe.config.requires_aesthetics_score
-    )
-
-    if is_refiner:
-        aesthetic_score = (6.0,)
-        time_ids_list = list(original_size + crops_coords_top_left + aesthetic_score)
-    else:
-        target_size = (h, w)
-        time_ids_list = list(original_size + crops_coords_top_left + target_size)
-
-    time_ids = torch.tensor(time_ids_list).repeat(batch_size, 1).to(torch.int64)
-    text_embeds_shape = (batch_size, ref_pipe.text_encoder_2.config.hidden_size)
-
-    return {
-        "time_ids": time_ids,
-        "text_embeds": torch.randn(*text_embeds_shape).to(torch.float32),
-    }
-
-
-def get_inputs_spec(inputs):
-    inputs_spec = {k: (v.shape, v.dtype) for k, v in inputs.items()}
-    return inputs_spec
-
-
-def add_cnet_support(sample_shape, reference_unet):
-    from python_coreml_stable_diffusion.unet import calculate_conv2d_output_shape
-
-    additional_residuals_shapes = []
-
-    batch_size = sample_shape[0]
-    h, w = sample_shape[2:]
-
-    # conv_in
-    out_h, out_w = calculate_conv2d_output_shape(
-        h,
-        w,
-        reference_unet.conv_in,
-    )
-    additional_residuals_shapes.append(
-        (batch_size, reference_unet.conv_in.out_channels, out_h, out_w)
-    )
-
-    # down_blocks
-    for down_block in reference_unet.down_blocks:
-        additional_residuals_shapes += [
-            (batch_size, resnet.out_channels, out_h, out_w)
-            for resnet in down_block.resnets
-        ]
-        if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
-            for downsampler in down_block.downsamplers:
-                out_h, out_w = calculate_conv2d_output_shape(
-                    out_h, out_w, downsampler.conv
-                )
-            additional_residuals_shapes.append(
-                (
-                    batch_size,
-                    down_block.downsamplers[-1].conv.out_channels,
-                    out_h,
-                    out_w,
-                )
-            )
-
-    # mid_block
-    additional_residuals_shapes.append(
-        (batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w)
-    )
-
-    additional_inputs = {}
-    for i, shape in enumerate(additional_residuals_shapes):
-        sample_residual_input = torch.rand(*shape)
-        additional_inputs[f"additional_residual_{i}"] = sample_residual_input
-
-    return additional_inputs
-
-
-def convert_unet(
-    ref_pipe,
-    model_version: ModelVersion,
-    unet_out_path: str,
-    batch_size: int = 1,
-    sample_size: tuple[int, int] = (64, 64),
-    controlnet_support: bool = False,
-):
-    coreml_unet = get_unet(model_version, ref_pipe)
-    ref_unet = ref_pipe.unet
-
-    sample_shape = (
-        batch_size,  # B
-        ref_unet.config.in_channels,  # C
-        sample_size[0],  # H
-        sample_size[1],  # W
-    )
-
-    encoder_hidden_states_shape = get_encoder_hidden_states_shape(ref_pipe, batch_size)
-
-    scheduler = ref_pipe.scheduler
-    scheduler.set_timesteps(50)
-
-    sample_inputs = get_sample_input(
-        batch_size, encoder_hidden_states_shape, sample_shape, scheduler
-    )
-
-    if model_version == ModelVersion.LCM:
-        sample_inputs |= lcm_inputs(sample_inputs)
-
-    if model_version == ModelVersion.SDXL:
-        sample_inputs |= sdxl_inputs(sample_inputs, ref_pipe)
-
-    if controlnet_support:
-        sample_inputs |= add_cnet_support(sample_shape, ref_unet)
-
-    sample_inputs_spec = get_inputs_spec(sample_inputs)
-
-    logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
-    logger.info("JIT tracing..")
-    traced_unet = torch.jit.trace(
-        coreml_unet, example_inputs=list(sample_inputs.values())
-    )
-    logger.info("Done.")
-
-    coreml_sample_inputs = get_coreml_inputs(sample_inputs)
-
-    coreml_unet = convert_to_coreml(
-        "unet", traced_unet, coreml_sample_inputs, ["noise_pred"], unet_out_path
-    )
-
-    del traced_unet
-    gc.collect()
-
-    coreml_unet.save(unet_out_path)
-    logger.info(f"Saved unet into {unet_out_path}")
-
-
-def convert(
-    ckpt_path: str,
-    model_version: ModelVersion,
-    unet_out_path: str,
-    batch_size: int = 1,
-    sample_size: tuple[int, int] = (64, 64),
-    controlnet_support: bool = False,
-    lora_weights: list[tuple[Union[str, os.PathLike], float]] = None,
-    attn_impl: str = AttentionImplementations.SPLIT_EINSUM.name,
-    config_path: str = None,
-):
-    if os.path.exists(unet_out_path):
-        logger.info(f"Found existing model at {unet_out_path}! Skipping..")
-        return
-
-    python_coreml_stable_diffusion.unet.ATTENTION_IMPLEMENTATION_IN_EFFECT = (
-        AttentionImplementations(attn_impl)
-    )
-
-    ref_pipe = get_pipeline(ckpt_path, config_path, model_version)
-
-    for i, lora_weight in enumerate(lora_weights or []):
-        lora_path, strength = lora_weight
-        adapter_name = f"lora_{i}"
-        ref_pipe.load_lora_weights(lora_path, adapter_name=adapter_name)
-        ref_pipe.set_adapters([adapter_name], adapter_weights=[strength])
-        ref_pipe.fuse_lora()
-
-    convert_unet(
-        ref_pipe,
-        model_version,
-        unet_out_path,
-        batch_size,
-        sample_size,
-        controlnet_support,
-    )
-
-
-def get_pipeline(ckpt_path, config_path, model_version):
-    pipe_cls = MODEL_TYPE_TO_PIPE_CLS[model_version]
-    ref_pipe = pipe_cls.from_single_file(ckpt_path, original_config_file=config_path)
-    return ref_pipe
-
-
-def compile_model(out_path, out_name, submodule_name):
-    # Compile the model
-    target_path = compile_coreml_model(
-        out_path, get_folder_paths(submodule_name)[0], f"{out_name}_{submodule_name}"
-    )
-    logger.info(f"Compiled {out_path} to {target_path}")
-    return target_path

coreml_suite/core/__init__.py

@@ -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.
+"""

coreml_suite/core/controlnet.py

@@ -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

coreml_suite/core/inputs.py

@@ -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,
+            )
+        ]

coreml_suite/core/latents.py

@@ -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]]

coreml_suite/core/sdxl.py

@@ -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

coreml_suite/coreml_model.py

@@ -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
+        }

coreml_suite/latents.py

@@ -1,36 +1,4 @@
-import torch
+"""Compatibility shim — re-exports from coreml_suite.core.latents."""
+from coreml_suite.core.latents import chunk_batch, merge_chunks
 
-
-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]]
+__all__ = ["chunk_batch", "merge_chunks"]

coreml_suite/lcm/__init__.py

@@ -1,3 +1,8 @@
-from .nodes import COREML_CONVERT_LCM
+"""LCM runtime support (sampler-side).
 
-__all__ = ["COREML_CONVERT_LCM"]
+The dedicated LCM converter node was removed once the standard ``CoreMLConverter``
+gained model-version auto-detection (full-distill LCM is detected from the
+checkpoint). What remains here is runtime sampling support — ``utils`` patches the
+model sampling and supplies the guidance embedding when a converted UNet exposes
+``timestep_cond``.
+"""

coreml_suite/lcm/converter.py

@@ -1,297 +0,0 @@
-import os
-import shutil
-import logging
-import time
-import gc
-
-import numpy as np
-import torch
-from diffusers import UNet2DConditionModel, LCMScheduler
-from diffusers.loaders import LoraLoaderMixin
-
-from comfy.model_management import get_torch_device
-from coreml_suite.lcm.unet import UNet2DConditionModelLCM
-
-from transformers import CLIPTextModel
-import coremltools as ct
-
-from folder_paths import get_folder_paths
-
-logging.basicConfig()
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-
-MODEL_VERSION = "SimianLuo/LCM_Dreamshaper_v7"
-MODEL_NAME = MODEL_VERSION.split("/")[-1] + "_4k"
-
-import python_coreml_stable_diffusion.unet as unet
-
-unet.ATTENTION_IMPLEMENTATION_IN_EFFECT = unet.AttentionImplementations.SPLIT_EINSUM
-
-
-def get_unets():
-    ref_unet = UNet2DConditionModel.from_pretrained(
-        MODEL_VERSION,
-        subfolder="unet",
-        device_map=None,
-        low_cpu_mem_usage=False,
-    )
-
-    cml_unet = UNet2DConditionModelLCM.from_config(ref_unet.config).eval()
-    cml_unet.load_state_dict(ref_unet.state_dict(), strict=False)
-
-    return cml_unet, ref_unet
-
-
-def get_encoder_hidden_states_shape(unet_config, batch_size):
-    text_encoder = CLIPTextModel.from_pretrained(
-        MODEL_VERSION, subfolder="text_encoder"
-    )
-
-    text_token_sequence_length = text_encoder.config.max_position_embeddings
-    hidden_size = (text_encoder.config.hidden_size,)
-
-    encoder_hidden_states_shape = (
-        batch_size,
-        unet_config.cross_attention_dim or hidden_size,
-        1,
-        text_token_sequence_length,
-    )
-
-    return encoder_hidden_states_shape
-
-
-def get_scheduler():
-    scheduler = LCMScheduler.from_pretrained(MODEL_VERSION, subfolder="scheduler")
-    scheduler.set_timesteps(50, get_torch_device(), 50)
-    return scheduler
-
-
-def get_coreml_inputs(sample_inputs):
-    coreml_sample_unet_inputs = {
-        k: v.numpy().astype(np.float16) for k, v in sample_inputs.items()
-    }
-    return [
-        ct.TensorType(
-            name=k,
-            shape=v.shape,
-            dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
-        )
-        for k, v in coreml_sample_unet_inputs.items()
-    ]
-
-
-def load_coreml_model(out_path):
-    logger.info(f"Loading model from {out_path}")
-
-    start = time.time()
-    coreml_model = ct.models.MLModel(out_path)
-    logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds")
-
-    return coreml_model
-
-
-def convert_to_coreml(
-    submodule_name, torchscript_module, sample_inputs, output_names, out_path
-):
-    if os.path.exists(out_path):
-        logger.info(f"Skipping export because {out_path} already exists")
-        coreml_model = load_coreml_model(out_path)
-    else:
-        logger.info(f"Converting {submodule_name} to CoreML..")
-        coreml_model = ct.convert(
-            torchscript_module,
-            convert_to="mlprogram",
-            minimum_deployment_target=ct.target.macOS13,
-            inputs=sample_inputs,
-            outputs=[
-                ct.TensorType(name=name, dtype=np.float32) for name in output_names
-            ],
-            skip_model_load=True,
-        )
-
-        del torchscript_module
-        gc.collect()
-
-    return coreml_model
-
-
-def get_out_path(submodule_name, model_name):
-    fname = f"{model_name}_{submodule_name}.mlpackage"
-    unet_path = get_folder_paths(submodule_name)[0]
-    out_path = os.path.join(unet_path, fname)
-    return out_path
-
-
-def compile_coreml_model(source_model_path, output_dir, final_name):
-    """Compiles Core ML models using the coremlcompiler utility from Xcode toolchain"""
-    target_path = os.path.join(output_dir, f"{final_name}.mlmodelc")
-    if os.path.exists(target_path):
-        logger.warning(f"Found existing compiled model at {target_path}! Skipping..")
-        return target_path
-
-    logger.info(f"Compiling {source_model_path}")
-    source_model_name = os.path.basename(os.path.splitext(source_model_path)[0])
-
-    os.system(f"xcrun coremlcompiler compile {source_model_path} {output_dir}")
-    compiled_output = os.path.join(output_dir, f"{source_model_name}.mlmodelc")
-    shutil.move(compiled_output, target_path)
-
-    return target_path
-
-
-def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape, scheduler):
-    sample_unet_inputs = dict(
-        [
-            ("sample", torch.rand(*sample_shape)),
-            (
-                "timestep",
-                torch.tensor([scheduler.timesteps[0].item()] * batch_size).to(
-                    torch.float32
-                ),
-            ),
-            ("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)),
-            ("timestep_cond", torch.randn(batch_size, 256).to(torch.float32)),
-        ]
-    )
-    return sample_unet_inputs
-
-
-def get_unet_inputs_spec(sample_unet_inputs):
-    sample_unet_inputs_spec = {
-        k: (v.shape, v.dtype) for k, v in sample_unet_inputs.items()
-    }
-    return sample_unet_inputs_spec
-
-
-def add_cnet_support(sample_shape, reference_unet):
-    from python_coreml_stable_diffusion.unet import calculate_conv2d_output_shape
-
-    additional_residuals_shapes = []
-
-    batch_size = sample_shape[0]
-    h, w = sample_shape[2:]
-
-    # conv_in
-    out_h, out_w = calculate_conv2d_output_shape(
-        h,
-        w,
-        reference_unet.conv_in,
-    )
-    additional_residuals_shapes.append(
-        (batch_size, reference_unet.conv_in.out_channels, out_h, out_w)
-    )
-
-    # down_blocks
-    for down_block in reference_unet.down_blocks:
-        additional_residuals_shapes += [
-            (batch_size, resnet.out_channels, out_h, out_w)
-            for resnet in down_block.resnets
-        ]
-        if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
-            for downsampler in down_block.downsamplers:
-                out_h, out_w = calculate_conv2d_output_shape(
-                    out_h, out_w, downsampler.conv
-                )
-            additional_residuals_shapes.append(
-                (
-                    batch_size,
-                    down_block.downsamplers[-1].conv.out_channels,
-                    out_h,
-                    out_w,
-                )
-            )
-
-    # mid_block
-    additional_residuals_shapes.append(
-        (batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w)
-    )
-
-    additional_inputs = {}
-    for i, shape in enumerate(additional_residuals_shapes):
-        sample_residual_input = torch.rand(*shape)
-        additional_inputs[f"additional_residual_{i}"] = sample_residual_input
-
-    return additional_inputs
-
-
-def convert(
-    out_path: str,
-    batch_size: int = 1,
-    sample_size: tuple[int, int] = (64, 64),
-    controlnet_support: bool = False,
-    lora_paths: list[str] = None,
-):
-    lora_paths = lora_paths or []
-    coreml_unet, ref_unet = get_unets()
-
-    for lora_path in lora_paths:
-        lora_sd, network_alphas = LoraLoaderMixin.lora_state_dict(lora_path)
-        LoraLoaderMixin.load_lora_into_unet(lora_sd, network_alphas, ref_unet)
-        ref_unet.fuse_lora()
-
-    sample_shape = (
-        batch_size,  # B
-        ref_unet.config.in_channels,  # C
-        sample_size[0],  # H
-        sample_size[1],  # W
-    )
-
-    encoder_hidden_states_shape = get_encoder_hidden_states_shape(
-        ref_unet.config, batch_size
-    )
-
-    scheduler = get_scheduler()
-
-    sample_inputs = get_sample_input(
-        batch_size, encoder_hidden_states_shape, sample_shape, scheduler
-    )
-
-    if controlnet_support:
-        sample_inputs |= add_cnet_support(sample_shape, ref_unet)
-
-    sample_inputs_spec = get_unet_inputs_spec(sample_inputs)
-
-    logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
-    logger.info("JIT tracing..")
-    traced_unet = torch.jit.trace(
-        coreml_unet, example_inputs=list(sample_inputs.values())
-    )
-    logger.info("Done.")
-
-    coreml_sample_inputs = get_coreml_inputs(sample_inputs)
-
-    coreml_unet = convert_to_coreml(
-        "unet", traced_unet, coreml_sample_inputs, ["noise_pred"], out_path
-    )
-
-    del traced_unet
-    gc.collect()
-
-    coreml_unet.save(out_path)
-    logger.info(f"Saved unet into {out_path}")
-
-
-def compile_model(out_path, out_name):
-    # Compile the model
-    target_path = compile_coreml_model(
-        out_path, get_folder_paths("unet")[0], f"{out_name}_unet"
-    )
-    logger.info(f"Compiled {out_path} to {target_path}")
-    return target_path
-
-
-if __name__ == "__main__":
-    h = 512
-    w = 512
-    sample_size = (h // 8, w // 8)
-    batch_size = 4
-
-    cn_support_str = "_cn" if True else ""
-
-    out_name = f"{MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
-
-    out_path = get_out_path("unet", f"{out_name}")
-    if not os.path.exists(out_path):
-        convert(out_path=out_path, sample_size=sample_size, batch_size=batch_size)
-    compile_model(out_path=out_path, out_name=out_name)

coreml_suite/lcm/nodes.py

@@ -1,70 +0,0 @@
-import os
-
-from coremltools import ComputeUnit
-from python_coreml_stable_diffusion.coreml_model import CoreMLModel
-
-from coreml_suite import COREML_NODE
-from coreml_suite.lcm import converter as lcm_converter
-
-
-class COREML_CONVERT_LCM(COREML_NODE):
-    """Converts a LCM model to Core ML."""
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
-                "width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
-                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
-                "compute_unit": (
-                    [
-                        ComputeUnit.CPU_AND_NE.name,
-                        ComputeUnit.CPU_AND_GPU.name,
-                        ComputeUnit.ALL.name,
-                        ComputeUnit.CPU_ONLY.name,
-                    ],
-                ),
-                "controlnet_support": ("BOOLEAN", {"default": False}),
-            }
-        }
-
-    RETURN_TYPES = ("COREML_UNET",)
-    RETURN_NAMES = ("coreml_model",)
-    FUNCTION = "convert"
-
-    def convert(self, height, width, batch_size, compute_unit, controlnet_support):
-        """Converts a LCM model to Core ML.
-
-        Args:
-            height (int): Height of the target image.
-            width (int): Width of the target image.
-            batch_size (int): Batch size.
-            compute_unit (str): Compute unit to use when loading the model.
-
-        Returns:
-            coreml_model: The converted Core ML model.
-
-        The converted model is also saved to "models/unet" directory and
-        can be loaded with the "LCMCoreMLLoaderUNet" node.
-        """
-        h = height
-        w = width
-        sample_size = (h // 8, w // 8)
-        batch_size = batch_size
-        cn_support_str = "_cn" if controlnet_support else ""
-
-        out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
-
-        out_path = lcm_converter.get_out_path("unet", f"{out_name}")
-
-        if not os.path.exists(out_path):
-            lcm_converter.convert(
-                out_path=out_path,
-                sample_size=sample_size,
-                batch_size=batch_size,
-                controlnet_support=controlnet_support,
-            )
-        target_path = lcm_converter.compile_model(out_path=out_path, out_name=out_name)
-
-        return (CoreMLModel(target_path, compute_unit, "compiled"),)

coreml_suite/lcm/unet.py

@@ -1,99 +0,0 @@
-from overrides import overrides
-from python_coreml_stable_diffusion.unet import UNet2DConditionModel, TimestepEmbedding
-
-
-class UNet2DConditionModelLCM(UNet2DConditionModel):
-    def __init__(
-        self,
-        time_cond_proj_dim=None,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-        timestep_input_dim = self.config.block_out_channels[0]
-        time_embed_dim = self.config.block_out_channels[0] * 4
-
-        time_embedding = TimestepEmbedding(
-            timestep_input_dim, time_embed_dim, cond_proj_dim=time_cond_proj_dim
-        )
-        self.time_embedding = time_embedding
-
-    @overrides(check_signature=False)
-    def forward(
-        self,
-        sample,
-        timestep,
-        encoder_hidden_states,
-        timestep_cond,
-        *additional_residuals,
-    ):
-        # 0. Project (or look-up) time embeddings
-        t_emb = self.time_proj(timestep)
-        emb = self.time_embedding(t_emb, timestep_cond)
-
-        # 1. center input if necessary
-        if self.config.center_input_sample:
-            sample = 2 * sample - 1.0
-
-        # 2. pre-process
-        sample = self.conv_in(sample)
-
-        # 3. down
-        down_block_res_samples = (sample,)
-        for downsample_block in self.down_blocks:
-            if (
-                hasattr(downsample_block, "attentions")
-                and downsample_block.attentions is not None
-            ):
-                sample, res_samples = downsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    encoder_hidden_states=encoder_hidden_states,
-                )
-            else:
-                sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
-
-            down_block_res_samples += res_samples
-
-        if additional_residuals:
-            new_down_block_res_samples = ()
-            for i, down_block_res_sample in enumerate(down_block_res_samples):
-                down_block_res_sample = down_block_res_sample + additional_residuals[i]
-                new_down_block_res_samples += (down_block_res_sample,)
-            down_block_res_samples = new_down_block_res_samples
-
-        # 4. mid
-        sample = self.mid_block(
-            sample, emb, encoder_hidden_states=encoder_hidden_states
-        )
-
-        if additional_residuals:
-            sample = sample + additional_residuals[-1]
-
-        # 5. up
-        for upsample_block in self.up_blocks:
-            res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
-            down_block_res_samples = down_block_res_samples[
-                : -len(upsample_block.resnets)
-            ]
-
-            if (
-                hasattr(upsample_block, "attentions")
-                and upsample_block.attentions is not None
-            ):
-                sample = upsample_block(
-                    hidden_states=sample,
-                    temb=emb,
-                    res_hidden_states_tuple=res_samples,
-                    encoder_hidden_states=encoder_hidden_states,
-                )
-            else:
-                sample = upsample_block(
-                    hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples
-                )
-
-        # 6. post-process
-        sample = self.conv_norm_out(sample)
-        sample = self.conv_act(sample)
-        sample = self.conv_out(sample)
-
-        return (sample,)

coreml_suite/models.py

@@ -1,15 +1,44 @@
-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 coreml_suite.config import get_model_config, ModelVersion
-from coreml_suite.controlnet import extract_residual_kwargs, chunk_control
-from coreml_suite.latents import chunk_batch, merge_chunks
+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
 
+__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",
+]
+
 
 class CoreMLModelWrapper:
     def __init__(self, coreml_model):
@@ -68,204 +97,27 @@ class CoreMLModelWrapperLCM(CoreMLModelWrapper):
         self.config = None
 
 
-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)
-        context = context.transpose(0, 2, 1)[:, :, None, :]
-
-        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"]
-        hidden_shape = expected_inputs["encoder_hidden_states"]["shape"]
-        context_shape = (hidden_shape[0], hidden_shape[3], hidden_shape[1])
-
-        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,
-            )
-        ]
-
-
-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 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
-
-
 def add_sdxl_model_options(model_patcher, positive, negative):
     mp = model_patcher.clone()
 
     pos_dict = positive[0][1]
     neg_dict = negative[0][1]
 
-    pos_pooled = pos_dict["pooled_output"]
-    neg_pooled = neg_dict["pooled_output"]
-
-    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 model_patcher.model.diffusion_model.is_sdxl_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),
-        ]
-
+    is_base = model_patcher.model.diffusion_model.is_sdxl_base
     is_refiner = model_patcher.model.diffusion_model.is_sdxl_refiner
-    if is_refiner:
-        pos_time_ids += [
-            pos_dict.get("aesthetic_score", 6),
-        ]
 
-        neg_time_ids += [
-            neg_dict.get("aesthetic_score", 2.5),
-        ]
+    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"]
+    )
 
-    time_ids = torch.tensor([pos_time_ids, neg_time_ids])
-    text_embeds = torch.cat((pos_pooled, neg_pooled))
-
-    model_options = {
+    mp.model_options |= {
         "model_function_wrapper": sdxl_model_function_wrapper(
             time_ids, text_embeds, is_refiner
         ),
     }
-    mp.model_options |= model_options
-
     return mp
 
 

coreml_suite/nodes.py

@@ -1,13 +1,10 @@
 import os
 
 from coremltools import ComputeUnit
-from python_coreml_stable_diffusion.coreml_model import CoreMLModel
-from python_coreml_stable_diffusion.unet import AttentionImplementations
 
 import folder_paths
 from coreml_suite import COREML_NODE
-from coreml_suite import converter
-from coreml_suite.config import ModelVersion
+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
@@ -20,6 +17,26 @@ from coreml_suite.models import (
 )
 
 
+def _discover(fn_name, fallback):
+    """Populate a converter dropdown from coreml_diffusion's discovery API.
+
+    Fails soft: if the package is missing, too old to expose ``fn_name``, or
+    errors, the node still registers with the fallback list instead of vanishing
+    from the menu. Evaluated on every INPUT_TYPES call, so installing a newer
+    coreml_diffusion surfaces new conversion types with no Suite change.
+    """
+    try:
+        import coreml_diffusion
+
+        return getattr(coreml_diffusion, fn_name)()
+    except Exception as exc:  # missing/old package, import error, etc.
+        logger.warning(
+            f"coreml_diffusion.{fn_name} unavailable ({exc}); "
+            f"using fallback {fallback}"
+        )
+        return fallback
+
+
 class CoreMLSampler(COREML_NODE, KSampler):
     @classmethod
     def INPUT_TYPES(s):
@@ -163,7 +180,7 @@ class CoreMLLoader(COREML_NODE):
 
     @classmethod
     def coreml_filenames(cls):
-        extensions = (".mlmodelc", ".mlpackage")
+        extensions = (".mlpackage",)
         all_paths = folder_paths.get_filename_list_(cls.PACKAGE_DIRNAME)[1]
         coreml_paths = folder_paths.filter_files_extensions(all_paths, extensions)
 
@@ -174,9 +191,7 @@ class CoreMLLoader(COREML_NODE):
 
         coreml_path = self.coreml_filenames()[coreml_name]
 
-        sources = "compiled" if coreml_name.endswith(".mlmodelc") else "packages"
-
-        return (CoreMLModel(coreml_path, compute_unit, sources),)
+        return (CoreMLModel(coreml_path, compute_unit),)
 
 
 class CoreMLLoaderUNet(CoreMLLoader):
@@ -210,28 +225,26 @@ class CoreMLModelAdapter(COREML_NODE):
 
 
 class CoreMLConverter(COREML_NODE):
-    """Converts a LCM model to Core ML."""
+    """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"),),
-                "model_version": (
-                    [
-                        ModelVersion.SD15.name,
-                        ModelVersion.SDXL.name,
-                    ],
-                ),
-                "height": ("INT", {"default": 512, "min": 256, "max": 2048, "step": 8}),
-                "width": ("INT", {"default": 512, "min": 256, "max": 2048, "step": 8}),
+                "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": (
-                    [
-                        AttentionImplementations.SPLIT_EINSUM.name,
-                        AttentionImplementations.SPLIT_EINSUM_V2.name,
-                        AttentionImplementations.ORIGINAL.name,
-                    ],
+                    _discover(
+                        "list_attention_impls",
+                        ["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"],
+                    ),
                 ),
                 "compute_unit": (
                     [
@@ -244,6 +257,15 @@ class CoreMLConverter(COREML_NODE):
                 "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",),
             },
         }
@@ -255,18 +277,20 @@ class CoreMLConverter(COREML_NODE):
     def convert(
         self,
         ckpt_name,
-        model_version,
         height,
         width,
         batch_size,
         attention_implementation,
         compute_unit,
         controlnet_support,
+        quantize_nbits="none",
         lora_params=None,
     ):
-        """Converts a LCM model to Core ML.
+        """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.
@@ -276,10 +300,8 @@ class CoreMLConverter(COREML_NODE):
             coreml_model: The converted Core ML model.
 
         The converted model is also saved to "models/unet" directory and
-        can be loaded with the "LCMCoreMLLoaderUNet" node.
+        can be loaded with the "Load Core ML UNet" node.
         """
-        model_version = ModelVersion[model_version]
-
         lora_params = lora_params or {}
         lora_params = [(k, v[0]) for k, v in lora_params.items()]
         lora_params = sorted(lora_params, key=lambda lora: lora[0])
@@ -288,24 +310,19 @@ class CoreMLConverter(COREML_NODE):
         h = height
         w = width
         sample_size = (h // 8, w // 8)
-        batch_size = batch_size
-        cn_support_str = "_cn" if controlnet_support else ""
-        lora_str = (
-            "_" + "_".join(lora_param[0].split(".")[0] for lora_param in lora_params)
-            if lora_params
-            else ""
-        )
+        import coreml_diffusion
 
-        attn_str = (
-            "_"
-            + {"SPLIT_EINSUM": "se", "SPLIT_EINSUM_V2": "se2", "ORIGINAL": "orig"}[
-                attention_implementation
-            ]
+        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,
         )
 
-        out_name = f"{ckpt_name.split('.')[0]}{lora_str}_{batch_size}x{w}x{h}{cn_support_str}{attn_str}"
-        out_name = out_name.replace(" ", "_")
-
         logger.info(f"Converting {ckpt_name} to {out_name}")
         logger.info(f"Batch size: {batch_size}")
         logger.info(f"Width: {w}, Height: {h}")
@@ -313,11 +330,14 @@ class CoreMLConverter(COREML_NODE):
         logger.info(f"Attention implementation: {attention_implementation}")
 
         if lora_params:
-            logger.info(f"LoRAs used:")
+            logger.info("LoRAs used:")
             for lora_param in lora_params:
                 logger.info(f"  {lora_param[0]} - strength: {lora_param[1]}")
 
-        unet_out_path = converter.get_out_path("unet", f"{out_name}")
+        # Resolve the ComfyUI models/unet path here (a node concern); the package
+        # takes the output path as an injected argument.
+        unet_path = folder_paths.get_folder_paths("unet")[0]
+        unet_out_path = os.path.join(unet_path, f"{out_name}_unet.mlpackage")
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
 
         config_filename = ckpt_name.split(".")[0] + ".yaml"
@@ -325,22 +345,19 @@ class CoreMLConverter(COREML_NODE):
         if config_path:
             logger.info(f"Using config file {config_path}")
 
-        converter.convert(
-            ckpt_path=ckpt_path,
-            model_version=model_version,
-            unet_out_path=unet_out_path,
+        coreml_diffusion.convert(
+            ckpt_path,
+            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,
         )
-        unet_target_path = converter.compile_model(
-            out_path=unet_out_path, out_name=out_name, submodule_name="unet"
-        )
-
-        return (CoreMLModel(unet_target_path, compute_unit, "compiled"),)
+        return (CoreMLModel(unet_out_path, compute_unit),)
 
     @staticmethod
     def lora_path(lora_name):

docs/conversion.md

@@ -0,0 +1,93 @@
+# Conversion
+
+## Conversion is the only supported path
+
+You always start from a Stable Diffusion checkpoint (`.safetensors` / `.ckpt`)
+and convert it with the **Convert Checkpoint to Core ML** node. The model
+version (SD1.5, SDXL, SDXL refiner, full-distill LCM) is auto-detected from the
+checkpoint. Pre-converted Core ML models from elsewhere are not supported,
+because:
+
+- The suite uses its own input **dimensions**, **naming convention**, and
+  **metadata**, all produced by the
+  [coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) package.
+- Apple's `ml-stable-diffusion` (which most community Core ML models target) is
+  effectively obsolete, and the layouts differ (see the 2.0.0
+  `encoder_hidden_states` change in the README).
+- Conversion is cheap and one-time, so there is no value in maintaining
+  backwards compatibility with foreign formats.
+
+The output is always a **`.mlpackage`**. The suite no longer compiles to
+`.mlmodelc` (it didn't work with the inference backend), so there is **no Xcode
+or `coremlcompiler` dependency**.
+
+## One-time conversion and name-based caching
+
+Conversion runs **once**, not on every queue. The converter encodes all
+conversion parameters into the output filename (via `coreml_diffusion.compose_out_name`,
+called in `coreml_suite/nodes.py`):
+
+- checkpoint name, `batch_size`, `width`, `height`
+- `controlnet_support`, `attention_implementation`
+- baked LoRA names, `quantize_nbits`
+
+The result is written as `<encoded-name>_unet.mlpackage` in `models/unet`. If a
+file with that name already exists, it is reused and conversion is skipped. Change
+any parameter → new name → new conversion; keep them the same → the cached model
+is loaded instantly.
+
+This is why the recommended workflow is **convert once, then load**: run the
+converter a single time, then in day-to-day use load the `.mlpackage` with the
+**Load Core ML UNet** node. (You can also leave the converter node in the graph;
+it short-circuits to the cached file.)
+
+> [!NOTE]
+> The converter relies on the filename to decide whether to reconvert. If you
+> rename the `.mlpackage`, it will be converted again. You can otherwise rename
+> it freely if the auto-generated name is too long.
+
+## Quantization
+
+The converter node accepts an optional `quantize_nbits` dropdown that runs
+k-means weight palettization (`coremltools.optimize.coreml.palettize_weights`) on
+the UNet before saving.
+
+Values: `none` (default — no quantization, identical output and filename to
+before), `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:
+
+| 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 quantization drift from
+sampler/VAE noise. Final-image PSNR is comfortably higher (the sampler averages
+over many steps).
+
+### Recommended settings per chip / RAM
+
+- **8 GB (M1/M2 base):** `nbits=4`. ~4× smaller, still loads, 27 dB is visually
+  identical at SD1.5 sizes.
+- **16 GB (M1/M2/M3 Pro):** `nbits=6` — the sweet spot, ~2.7× smaller, 40 dB, no
+  perceptible quality drop.
+- **32 GB+ (Max / Ultra):** `nbits=8` for a safety margin, or `none` for
+  bit-identical output (golden testing).
+
+The default stays `none`, so existing workflows produce byte-for-byte identical
+output.
+
+## Where conversion lives
+
+The conversion engine was extracted into the standalone
+[coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) PyPI package.
+The nodes in this suite resolve ComfyUI paths and call into it; node names,
+inputs, and outputs are unchanged, so the split has effectively no user-facing
+impact beyond `pip install` pulling one more dependency.

docs/faq.md

@@ -0,0 +1,77 @@
+# FAQ
+
+## What's the difference between ANE, GPU, and MPS, and which do I pick?
+
+ANE is the Neural Engine (Core ML only), GPU is the Metal GPU (Core ML or
+PyTorch), MPS is PyTorch's GPU backend. This suite uses **Core ML compute units
+only** and never touches MPS. Short answer: SD1.5 at 512×512 → convert
+`SPLIT_EINSUM`, load `CPU_AND_NE`; larger sizes or SDXL → convert `ORIGINAL`,
+load `CPU_AND_GPU`. Full reasoning: [hardware](hardware.md).
+
+## Do I still need `PYTORCH_ENABLE_MPS_FALLBACK=1`?
+
+Not for these nodes — Core ML inference doesn't use PyTorch MPS. It may still
+matter for other parts of your ComfyUI graph, but it has no effect on Core ML
+sampling.
+
+## Why is my Core ML SDXL workflow no faster than the default nodes?
+
+Because **SDXL can't run on the ANE** — the speedup comes from the Neural Engine,
+and SDXL falls back to the GPU, running at roughly MPS-equivalent speed. This is a
+known limitation, not a misconfiguration. The ANE benefit is real for SD1.5. See
+[limitations](limitations.md).
+
+## Where do I get Core ML models?
+
+You convert them yourself — that's the only supported path. See
+[conversion](conversion.md). Downloaded Core ML models (e.g. coreml-community) use
+different dimensions/metadata and are not supported.
+
+## Is conversion run every time I queue, or once?
+
+Once. Parameters are encoded in the output filename, so an already-converted model
+is reused and conversion is skipped. Convert once, then load the `.mlpackage`. See
+[conversion → caching](conversion.md#one-time-conversion-and-name-based-caching).
+
+## Does a converted model produce the same output as the original?
+
+With the default `quantize_nbits = none`, the converted UNet output matches the
+source within numerical rounding (the golden test in `tests/m2/test_golden_image.py`
+gates on PSNR ≥ 20 dB on the decoded image). Quantization (`8`/`6`/`4`) introduces
+measured, bounded drift — see the [PSNR table](conversion.md#quantization). For
+bit-identical output, keep `none`.
+
+## Are `.mlpackage` models safe to use?
+
+`.mlpackage` is a declarative Core ML model format — it carries weights and a
+compute graph, not arbitrary executable code or Python pickle, so its safety
+profile is comparable to `safetensors`. In practice this matters little here,
+since the only supported models are ones you convert locally from your own
+checkpoints.
+
+## Are LoRAs reliable?
+
+Partially. Some LoRAs convert cleanly; others produce poor or broken output —
+there's no firm rule, so test per-LoRA. LoRA weights and `strength_model` are
+baked in at conversion and can't be changed afterward; for some LCM-LoRA cases the
+[Core ML Adapter](nodes.md#core-ml-adapter-experimental-coremlmodeladapter) path
+is more reliable. Treat LoRA support as experimental. See
+[troubleshooting](troubleshooting.md).
+
+## Does the experimental Adapter cost performance vs the Core ML Sampler?
+
+Yes, a little. The Adapter wraps the model in a ComfyUI `ModelPatcher` so standard
+samplers work, which adds per-step interface overhead the native Core ML Sampler
+avoids. Use the native sampler unless you specifically need a `MODEL` (e.g.
+`ModelSamplingDiscrete` for LCM LoRAs).
+
+## Which Python versions work?
+
+Python 3.12 or newer (`requires-python >=3.12`). Older 3.12 install failures
+came from the now-removed `ml-stable-diffusion` build, not from this suite.
+
+## Long prompts crash my workflow
+
+Core ML has a hard **77-token** prompt limit and doesn't auto-chunk long prompts.
+Split the prompt across multiple CLIP Text Encode nodes and merge with Conditioning
+(Combine). See [troubleshooting](troubleshooting.md).

docs/hardware.md

@@ -0,0 +1,95 @@
+# Hardware & Compute Units
+
+This page explains how the suite maps to Apple Silicon hardware, the difference
+between ANE, GPU, and MPS, and how to choose a compute unit and attention
+implementation.
+
+## ANE vs GPU vs MPS
+
+Three terms get conflated:
+
+- **ANE (Apple Neural Engine)** — a dedicated ML accelerator on Apple Silicon.
+  Only Core ML can target it; PyTorch cannot. This is the whole reason the suite
+  exists.
+- **GPU** — the Metal GPU. Reachable both by Core ML (as a compute unit) and by
+  PyTorch (via MPS).
+- **MPS (Metal Performance Shaders)** — PyTorch's GPU backend on macOS. This is
+  the path standard ComfyUI nodes use.
+
+**This suite uses Core ML compute units only — it never runs the UNet through
+PyTorch/MPS.** Consequently `PYTORCH_ENABLE_MPS_FALLBACK` has no effect on these
+nodes. It may still matter for the rest of your ComfyUI graph (CLIP, VAE,
+samplers on non-Core ML models), but not for Core ML inference itself.
+
+Rough performance picture (SD1.5, maintainer- and user-reported):
+
+- ANE is meaningfully faster than MPS — on the order of **50–100%** for SD1.5.
+- Core ML on the GPU is only marginally faster than PyTorch/MPS.
+
+So the speedup comes from the Neural Engine, which means it depends on being able
+to actually run on the ANE (see [attention implementations](#attention-implementations)
+and the [SDXL caveat](#sdxl-and-the-ane)).
+
+## Compute units
+
+The **compute unit** is set on the loader/converter node and tells Core ML which
+hardware to use. It is applied when the model is loaded
+(`coreml_suite/coreml_model.py:22`), not during conversion.
+
+| Value | Hardware | Best paired with |
+|---|---|---|
+| `CPU_AND_NE` (default) | CPU + Neural Engine | `SPLIT_EINSUM` / `SPLIT_EINSUM_V2` |
+| `CPU_AND_GPU` | CPU + Metal GPU | `ORIGINAL` |
+| `CPU_ONLY` | CPU only | fallback / debugging |
+| `ALL` | all available hardware | rarely optimal — see below |
+
+Notes:
+
+- Every option includes the CPU; there is no GPU-and-ANE-without-CPU combination.
+- `NE` in `CPU_AND_NE` is the Neural Engine (Apple's enum spells it `NE`, not
+  `ANE`).
+- **`CPU_AND_NE` is often faster than `ALL`.** Letting Core ML use everything can
+  be *slower* on non-Max chips, where memory bandwidth is the bottleneck. Try
+  `CPU_AND_NE` first for SD1.5.
+
+## Attention implementations
+
+Chosen at conversion time on the **Convert Checkpoint to Core ML** node. It
+decides whether the model can run on the ANE:
+
+- **`SPLIT_EINSUM`** — ANE-friendly attention. Use for the Neural Engine.
+- **`SPLIT_EINSUM_V2`** — a variant; in practice ≈ `SPLIT_EINSUM` for most users.
+- **`ORIGINAL`** — standard attention. Runs on the GPU, not the ANE.
+
+The implementation and the compute unit must agree: a `SPLIT_EINSUM` model wants
+`CPU_AND_NE`; an `ORIGINAL` model wants `CPU_AND_GPU`.
+
+## Which should I pick?
+
+| Scenario | Attention | Compute unit |
+|---|---|---|
+| SD1.5 at 512×512 | `SPLIT_EINSUM` | `CPU_AND_NE` |
+| SD1.5 at larger sizes (e.g. 768) | `ORIGINAL` | `CPU_AND_GPU` |
+| SDXL / SDXL Turbo | `ORIGINAL` | `CPU_AND_GPU` |
+
+### Resolution crossover
+
+ANE shines at small latents; the GPU scales better as resolution grows. In user
+benchmarks:
+
+- At **512×512**, ANE + `SPLIT_EINSUM` wins by roughly **10%** over the GPU path.
+- At **768×512**, GPU + `ORIGINAL` pulls ahead by roughly **10%**, and the larger
+  image is about 2× slower overall.
+
+If you mostly work at 512×512, convert with `SPLIT_EINSUM` and load on
+`CPU_AND_NE`. If you routinely go larger, an `ORIGINAL` + GPU model may be
+faster.
+
+### SDXL and the ANE
+
+SDXL (and SDXL Turbo) **cannot run on the ANE** — the dual-text-encoder UNet
+exceeds what the Neural Engine path supports. SDXL therefore runs at roughly
+MPS-equivalent speed with no ANE speedup. If a Core ML SDXL workflow feels no
+faster than the standard nodes, this is why. Convert SDXL with `ORIGINAL` and
+load with `CPU_AND_GPU` or `CPU_ONLY`. See
+[limitations](limitations.md) for the full picture.

docs/limitations.md

@@ -0,0 +1,53 @@
+# Limitations & Support Matrix
+
+## Support matrix
+
+| Feature | Status | Notes |
+|---|---|---|
+| SD1.5 | ✅ Full | ANE via `SPLIT_EINSUM`; the primary, fastest path |
+| SDXL / SDXL Turbo | ⚠️ Partial | GPU only (no ANE), no speedup; possible quality loss vs source. Don't run Turbo at 1024² |
+| SD2.1 | ❌ Unsupported | |
+| Inpainting checkpoints (9-channel) | ❌ Unsupported | |
+| ControlNet | ✅ Supported | Convert the checkpoint with `controlnet_support = True` |
+| LoRA | ⚠️ Experimental | Inconsistent per-LoRA; baked at conversion, immutable afterward |
+| LCM | ⚠️ Experimental | Full-distill LCM checkpoints auto-detected by the converter |
+| SVD | ❌ Not supported | |
+| AnimateDiff | ❌ Not supported | Motion modules need pre-conversion injection; not feasible today |
+| IPAdapter | ❌ Not supported | Needs a real `MODEL` the Core ML wrapper can't provide |
+| Core ML Adapter | ⚠️ Experimental | Works for many nodes; fails for merges/IPAdapter/etc. |
+
+## Fixed input/output shapes
+
+A Core ML model is converted for one specific resolution and batch size. To work
+at a different size, re-convert with the new width/height (conversion is cheap and
+cached by name). This is also why detailers and latent-upscale workflows that
+rescale mid-graph break — see [troubleshooting](troubleshooting.md).
+
+There is experimental support for flexible shapes via
+[EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes),
+but it is **much slower** — user benchmarks show roughly **5×** the per-iteration
+time on every run, not just the first. Fixed-shape models per resolution are the
+practical choice.
+
+## SDXL on the Neural Engine
+
+SDXL and SDXL Turbo cannot run on the ANE — the dual-text-encoder UNet exceeds the
+supported Neural Engine path. They run on the GPU at roughly MPS-equivalent speed,
+so Core ML offers no speed advantage for SDXL, and converted output may look
+degraded versus the safetensors original (an upstream conversion artifact). Use
+`ORIGINAL` + `CPU_AND_GPU`. See [hardware](hardware.md).
+
+## Experimental Core ML Adapter
+
+The Adapter wraps a Core ML model to look like a standard ComfyUI `MODEL`, which
+covers many standard and custom nodes. But it can't fully emulate a real model:
+operations that need genuine `MODEL` internals — model merges, IPAdapter, some
+LoRA flows, detailers — generally won't work, and the model's
+fixed input shapes aren't validated, so mismatches error at runtime. Prefer the
+native Core ML Sampler when you don't need the `MODEL` type.
+
+## Prompt length
+
+Core ML enforces a hard 77-token prompt limit with no auto-chunking. Split long
+prompts across multiple CLIP Text Encode nodes and merge with Conditioning
+(Combine).

docs/nodes.md

@@ -0,0 +1,157 @@
+# Node Reference
+
+All nodes live in the **Core ML Suite** category. Right-click the canvas →
+**Add Node → Core ML Suite**, or double-click and search.
+
+| Display name | Class | Purpose |
+|---|---|---|
+| Load Core ML UNet | `CoreMLUNetLoader` | Load a converted `.mlpackage` |
+| Core ML Sampler | `CoreMLSampler` | Sample (KSampler-style) |
+| Core ML Sampler (Advanced) | `CoreMLSamplerAdvanced` | Sample (KSamplerAdvanced-style) |
+| Core ML Adapter (Experimental) | `CoreMLModelAdapter` | Wrap as a standard `MODEL` |
+| Load LoRA to use with Core ML | `Core ML LoRA Loader` | Bake LoRA(s) at conversion |
+| Convert Checkpoint to Core ML | `Core ML Converter` | Convert a checkpoint |
+
+---
+
+## Load Core ML UNet (`CoreMLUNetLoader`)
+
+![Load Core ML UNet](../assets/unet_loader.png?raw=true)
+
+Loads a converted `.mlpackage` from `models/unet` and outputs a `coreml_model`
+for the samplers. Only `.mlpackage` files are listed — this suite no longer uses
+`.mlmodelc`.
+
+- **Inputs**
+  - `coreml_name` — the `.mlpackage` to load from `models/unet`.
+  - `compute_unit` — hardware to run on: `CPU_AND_NE` (default), `CPU_AND_GPU`,
+    `CPU_ONLY`, `ALL`. See [hardware](hardware.md).
+- **Output**
+  - `coreml_model` — for the Core ML Sampler or Adapter.
+
+---
+
+## Core ML Sampler (`CoreMLSampler`)
+
+![Core ML Sampler](../assets/sampler.png?raw=true)
+
+Generates a latent from a Core ML model. Behaves like the standard KSampler and
+outputs a `LATENT` you can decode or feed downstream.
+
+- **Inputs**
+  - `coreml_model` — output of the loader or a converter.
+  - `latent_image` *(optional)* — must match the model's input size. If omitted,
+    a suitable empty latent is created. Provide one for img2img.
+  - `negative` *(optional)* — required for normal models; optional for LCM.
+  - Remaining inputs (`seed`, `steps`, `cfg`, `sampler_name`, `scheduler`,
+    `positive`, `denoise`) match the KSampler.
+- **Output**
+  - `LATENT` — decode with a VAE Decode, or use downstream.
+
+---
+
+## Core ML Sampler (Advanced) (`CoreMLSamplerAdvanced`)
+
+The KSamplerAdvanced counterpart of the Core ML Sampler — same Core ML input,
+plus the advanced sampling controls. Use it for partial denoising, fixed noise,
+and multi-stage (e.g. SDXL base → refiner) workflows.
+
+- **Inputs**
+  - `coreml_model` — output of the loader or a converter.
+  - `add_noise`, `noise_seed`, `start_at_step`, `end_at_step`,
+    `return_with_leftover_noise` — as in KSamplerAdvanced.
+  - `steps`, `cfg`, `sampler_name`, `scheduler`, `positive` — as usual.
+  - `latent_image` *(optional)*, `negative` *(optional, required for non-LCM)*.
+- **Output**
+  - `LATENT`.
+
+---
+
+## Core ML Adapter (Experimental) (`CoreMLModelAdapter`)
+
+![Core ML Adapter](../assets/adapter.png?raw=true)
+
+Wraps a Core ML model so it presents as a standard ComfyUI `MODEL`, letting you
+feed it to the normal KSampler and many other nodes (e.g. `ModelSamplingDiscrete`
+for LCM LoRAs).
+
+- **Input**
+  - `coreml_model`.
+- **Output**
+  - `MODEL` — a Core ML model wrapped as a ComfyUI model.
+
+> [!NOTE]
+> Experimental. The wrapper presents a `MODEL` interface but cannot fully
+> emulate one — model merges, IPAdapter, and similar advanced uses generally
+> won't work, and the model's fixed input shapes are not validated, so mismatched
+> inputs error at runtime. The native Core ML Sampler is faster when you don't
+> need the `MODEL` type. See the [FAQ](faq.md) and [limitations](limitations.md).
+
+---
+
+## Load LoRA to use with Core ML (`Core ML LoRA Loader`)
+
+![LoRA Loader](../assets/lora_loader.png?raw=true)
+
+Collects LoRA name + `strength_model` to bake into the model at conversion, and
+applies the LoRA to CLIP (which is not part of the Core ML path). Chain multiple
+loaders for multiple LoRAs.
+
+Because a converted model is immutable, the baked weights and `strength_model`
+**cannot** be changed afterward — changing them means re-converting. `strength_clip`
+only affects CLIP and can be changed freely. After conversion, when loading with
+`CoreMLUNetLoader`, apply the same LoRAs to CLIP manually (see
+[workflows](workflows.md)).
+
+- **Inputs**
+  - `lora_name`, `strength_model`, `strength_clip`.
+  - `clip` — from `CLIPLoader` / `CheckpointLoaderSimple` or another LoRA loader.
+  - `lora_params` *(optional)* — chain from another LoRA loader.
+- **Outputs**
+  - `CLIP` — with the LoRA applied.
+  - `lora_params` — pass to the converter or the next LoRA loader.
+
+> [!NOTE]
+> LoRA support is experimental and inconsistent — some LoRAs convert cleanly,
+> others produce poor results. Test per-LoRA. See [troubleshooting](troubleshooting.md).
+
+---
+
+## Convert Checkpoint to Core ML (`Core ML Converter`)
+
+![Checkpoint Converter](../assets/checkpoint_converter.png?raw=true)
+
+Converts a checkpoint from `models/checkpoints` to a Core ML `.mlpackage` in
+`models/unet`. The model version (SD1.5, SDXL, SDXL refiner, or full-distill
+LCM) is auto-detected from the checkpoint's architecture — there is no version
+dropdown. The conversion parameters are encoded in the output name, so an
+already-converted model is reused instead of re-converted. See
+[conversion](conversion.md) for details.
+
+- **Inputs**
+  - `ckpt_name` — checkpoint in `models/checkpoints`.
+  - `height`, `width` — target image size; any positive multiple of 8 (default
+    512). The model's input size is fixed at these values.
+  - `batch_size` — default 1; raise to convert a batch-capable model.
+  - `attention_implementation` — `SPLIT_EINSUM` / `SPLIT_EINSUM_V2` (ANE) or
+    `ORIGINAL` (GPU). See [hardware](hardware.md).
+  - `compute_unit` — used only when loading the result; does not affect
+    conversion.
+  - `controlnet_support` — set `True` to make the model usable with ControlNet
+    (default `False`).
+  - `quantize_nbits` *(optional)* — `none` (default), `8`, `6`, `4`. See
+    [conversion → quantization](conversion.md#quantization).
+  - `lora_params` *(optional)* — from the LoRA loader, to bake LoRAs in.
+- **Output**
+  - `coreml_model`.
+
+> [!NOTE]
+> Some checkpoints need a custom config `.yaml`. Place it in `models/configs`
+> named like the checkpoint (e.g. `juggernaut.safetensors` →
+> `juggernaut.yaml`); it is loaded automatically during conversion.
+
+> [!NOTE]
+> Full-distill LCM checkpoints (e.g.
+> [LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7)) are
+> detected and converted like any other checkpoint. When sampling an LCM model,
+> set `sampler_name` to `lcm` and `scheduler` to `sgm_uniform`.

docs/troubleshooting.md

@@ -0,0 +1,86 @@
+# Troubleshooting
+
+## `Expected shape … got …` / latent size mismatch
+
+The most common error. A Core ML model has **fixed** input dimensions — a model
+converted for 512×512 expects a 64×64 latent and rejects any other size (batch
+size is handled and doesn't matter; only width/height are fixed).
+
+**Fix:** set your Empty Latent (or upstream latent) to exactly the resolution the
+model was converted for, or re-convert at the size you want.
+
+## Old `.mlmodelc` model, or `metadata.json` not found
+
+This suite no longer produces or loads `.mlmodelc`; the loader lists `.mlpackage`
+only. Models from an older version (or downloaded community models) with a
+`.mlmodelc` structure won't load.
+
+**Fix:** re-convert the checkpoint with **Convert Checkpoint to Core ML**. No
+Xcode or `coremlcompiler` is required — that dependency was removed.
+
+## Prompt too long (`Expected size 154 but got 77`, or a crash)
+
+Core ML enforces a hard **77-token** prompt limit and does not auto-chunk like
+A1111/ComfyUI.
+
+**Fix:** split the prompt across multiple CLIP Text Encode nodes and merge them
+with **Conditioning (Combine)**.
+
+## `cannot import name 'ModelSamplingDiscreteLCM'`
+
+A ComfyUI refactor renamed this symbol.
+
+**Fix:** update the suite (fixed in PR #29) and re-run
+`pip install -r requirements.txt`.
+
+## LoRA loader `ImportError`
+
+`peft` became a required dependency.
+
+**Fix:** `pip install -r requirements.txt`. This recurs after ComfyUI-Manager
+updates if requirements aren't reinstalled.
+
+## ControlNet has no effect
+
+ControlNet support is baked at conversion. If the checkpoint was converted with
+`controlnet_support = False`, ControlNet does nothing.
+
+**Fix:** re-convert with `controlnet_support = True`. The ControlNet model itself
+needs no conversion, and `.fp16.safetensors` vs `.safetensors` makes no
+difference.
+
+## LoRAs produce garbage
+
+LoRA support is inconsistent — some work, some don't, with no firm rule. Test
+per-LoRA. For some LCM-LoRA setups, routing through the
+[Core ML Adapter](nodes.md#core-ml-adapter-experimental-coremlmodeladapter) is
+more reliable than the basic loader path. Remember weights are baked at conversion
+and can't be changed afterward.
+
+## FaceDetailer / detailers error on size
+
+Detailers rescale latents internally (e.g. 512 → 1024), which breaks the model's
+fixed input shape. There is no workaround node — a Core ML model only accepts
+the resolution it was converted for.
+
+**Fix:** convert a second model at the detailer's internal resolution and use it
+for the detailing pass, or run the detailer with a standard (non–Core ML) model.
+
+## Inpainting checkpoint errors (`tensor size 9 vs 4`)
+
+SD1.5 inpainting checkpoints use a 9-channel input and are **not supported**. This
+error is expected, not a bug.
+
+## Errors mentioning `python_coreml_stable_diffusion` or `ml-stable-diffusion`
+
+You're on a stale install. That dependency was removed; old install scripts tried
+`pip install git+…/ml-stable-diffusion.git`, which fails on modern Python.
+
+**Fix:** reinstall the current suite (`pip install -r requirements.txt`, which
+pulls `coreml-diffusion` from PyPI).
+
+## `all input tensors must be on the same device (mps:0 and cpu)` / ControlNet residual shape `(2,…) vs (1,…)`
+
+Old bugs that have been fixed.
+
+**Fix:** update to the latest version.

docs/workflows.md

@@ -0,0 +1,103 @@
+# Example Workflows
+
+> [!NOTE]
+> The models referenced are examples — substitute your own. Every workflow
+> starts from a checkpoint you convert yourself (see [conversion](conversion.md));
+> there is no Core ML model to download.
+
+## Basic txt2img
+
+Convert a SD1.5 checkpoint, then sample from it. CLIP and VAE come from standard
+ComfyUI nodes — either loaded separately or pulled from the checkpoint.
+
+1. Place a SD1.5 checkpoint in `models/checkpoints` (e.g.
+   [v1-5-pruned-emaonly](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors)).
+2. **Convert Checkpoint to Core ML** → queue once → a `.mlpackage` lands in
+   `models/unet`.
+3. **Load Core ML UNet** (or wire the converter output straight in) →
+   **Core ML Sampler** → **VAE Decode**.
+
+**CLIP and VAE from the checkpoint:**
+
+![Core ML UNet + checkpoint](../assets/unet+sampler+checkpoint.png?raw=true)
+
+**CLIP and VAE loaded separately** — use any SD1.5-compatible
+[CLIP](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors)
+and [VAE](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/vae/diffusion_pytorch_model.safetensors),
+placed in `models/clip` and `models/vae`:
+
+![Core ML UNet + CLIP + VAE](../assets/unet+sampler+clip+vae.png?raw=true)
+
+## ControlNet
+
+Convert the checkpoint with `controlnet_support = True`, then wire a standard
+ComfyUI ControlNet. The ControlNet model itself needs no conversion. Place it in
+`models/controlnet` (e.g.
+[control_v11p_sd15_scribble](https://huggingface.co/lllyasviel/control_v11p_sd15_scribble/blob/main/diffusion_pytorch_model.fp16.safetensors)).
+
+![Core ML UNet + ControlNet](../assets/unet+sampler+controlnet.png?raw=true)
+
+## Checkpoint conversion
+
+The minimal conversion graph. See
+[Convert Checkpoint to Core ML](nodes.md#convert-checkpoint-to-core-ml-core-ml-converter).
+
+![Checkpoint converter](../assets/basic_conversion.png?raw=true)
+
+## Conversion with LoRA
+
+Bake LoRA(s) into the model at conversion. Read the
+[LoRA caveats](nodes.md#load-lora-to-use-with-core-ml-core-ml-lora-loader) first
+— baked weights are immutable, and support is inconsistent per-LoRA.
+
+![Checkpoint converter + LoRA](../assets/conversion+lora.png?raw=true)
+
+## LCM LoRA conversion
+
+Chain multiple LoRA loaders to use several LoRAs with one model.
+
+> [!IMPORTANT]
+> Here the model goes through the **Core ML Adapter** and `ModelSamplingDiscrete`
+> into the standard ComfyUI KSampler (not the Core ML Sampler).
+> `ModelSamplingDiscrete` is required to sample LCM LoRAs correctly.
+
+![Multiple LoRAs](../assets/conversion+lcm_lora.png?raw=true)
+
+## Loading a model with baked LoRAs
+
+Load a model that already has LoRAs baked in. CLIP must be loaded separately and
+passed through the same LoRA nodes used at conversion. Since `lora_name` and
+`strength_model` are baked in, they need not be passed to the loader.
+
+> [!IMPORTANT]
+> As above, the model goes through the Core ML Adapter + `ModelSamplingDiscrete`
+> into the standard KSampler.
+
+![Loader + LoRA](../assets/loader+lcm_lora.png?raw=true)
+
+## LCM conversion with ControlNet
+
+Convert a full-distill LCM checkpoint (e.g.
+[LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7)) with
+the standard **Convert Checkpoint to Core ML** node — the LCM architecture is
+auto-detected. Use it with or without ControlNet. When sampling, set
+`sampler_name` to `lcm` and `scheduler` to `sgm_uniform`.
+
+![LCM + ControlNet](../assets/lcm+controlnet.png?raw=true)
+
+## SDXL Base + Refiner
+
+A basic SDXL graph. Add LoRAs and ControlNets as in the SD1.5 examples; the
+refiner step is optional.
+
+Models:
+[base + text encoders](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0),
+[refiner](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0),
+[VAE](https://huggingface.co/stabilityai/sdxl-vae).
+
+> [!IMPORTANT]
+> SDXL does not run on the ANE. Convert with `ORIGINAL` and load with
+> `CPU_AND_GPU` (or `CPU_ONLY`). If loading hangs on `CPU_AND_NE`, that is the
+> cause. See [limitations](limitations.md).
+
+![SDXL](../assets/sdxl_conversion.png?raw=true)

pyproject.toml

@@ -1,15 +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 = "1.0.1"
-license = { file = "LICENSE" }
-dependencies = ["git+https://github.com/apple/ml-stable-diffusion.git", "coremltools>=7.1", "overrides", "diffusers>=0.22", "peft>=0.6.2", "omegaconf>=2.3"]
+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"
-#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.hatch.build.targets.wheel]
+packages = ["coreml_suite"]
 
 [tool.comfy]
 PublisherId = "aszc-dev"
 DisplayName = "ComfyUI-CoreMLSuite"
-Icon = ""
+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"]

requirements.txt

@@ -1,6 +1,4 @@
-git+https://github.com/apple/ml-stable-diffusion.git
-coremltools>=7.1
-overrides
-diffusers>=0.22
-peft>=0.6.2
-omegaconf>=2.3
+coreml-diffusion>=0.1.4,<0.2
+coremltools>=9,<10
+numpy>=2,<3
+diffusers>=0.30

tests/conftest.py

@@ -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

tests/integration/test_basic_conversion_1_5.py

@@ -1,85 +0,0 @@
-import json
-import os
-
-import pytest
-import requests
-import torch
-from PIL import Image
-import numpy as np
-
-from folder_paths import get_save_image_path, get_output_directory
-
-IMAGE_PREFIX = "E2E-1.5"
-IMAGE_PREFIX_CML = f"{IMAGE_PREFIX}-CoreML"
-IMAGE_PREFIX_MPS = f"{IMAGE_PREFIX}-MPS"
-
-
-class OutputImageRepository:
-    def __init__(self, name_prefix):
-        self.name_prefix = name_prefix
-
-    def list_images(self):
-        full_output_folder, _, _, _, _ = get_save_image_path(
-            self.name_prefix, get_output_directory(), 512, 512
-        )
-        return full_output_folder, os.listdir(full_output_folder)
-
-    def delete_images(self):
-        full_output_folder, images = self.list_images()
-        for image in images:
-            os.remove(os.path.join(full_output_folder, image))
-
-    def get_latest_image(self, prefix):
-        full_output_folder, images = self.list_images()
-        for image in sorted(images, reverse=True):
-            if image.startswith(prefix):
-                return os.path.join(full_output_folder, image)
-        return None
-
-
-@pytest.fixture(scope="function")
-def output_image_repository():
-    repo = OutputImageRepository(IMAGE_PREFIX)
-    yield repo
-    repo.delete_images()
-
-
-def test_basic_conversion_1_5(output_image_repository):
-    with open("tests/integration/workflows/e2e-1.5-basic-conversion.json") as f:
-        prompt = json.load(f)
-    prompt = randomize_seed_in_prompt(prompt)
-    queue_prompt(prompt)
-
-    coreml_img_path = output_image_repository.get_latest_image(IMAGE_PREFIX_CML)
-    mps_img_path = output_image_repository.get_latest_image(IMAGE_PREFIX_MPS)
-
-    coreml_image = Image.open(coreml_img_path)
-    mps_image = Image.open(mps_img_path)
-
-    assert psnr(np.array(coreml_image), np.array(mps_image)) > 25
-
-
-def psnr(img1, img2):
-    mse = np.mean((img1 - img2) ** 2)
-    if mse == 0:
-        return 100
-    PIXEL_MAX = 255.0
-    return 20 * np.log10(PIXEL_MAX / np.sqrt(mse))
-
-
-def queue_prompt(prompt: dict):
-    p = {"prompt": prompt}
-    data = json.dumps(p).encode("utf-8")
-    req = requests.post("http://localhost:8188/prompt", data=data)
-    assert req.status_code == 200
-    while True:
-        req = requests.get("http://localhost:8188/prompt")
-        if req.json()["exec_info"]["queue_remaining"] == 0:
-            break
-
-
-def randomize_seed_in_prompt(prompt):
-    seed = torch.random.seed()
-    prompt["3"]["inputs"]["seed"] = seed
-    prompt["11"]["inputs"]["seed"] = seed
-    return prompt

tests/integration/workflows/e2e-1.5-basic-conversion.json

@@ -107,7 +107,6 @@
   "10": {
     "inputs": {
       "ckpt_name": "dreamshaper_8.safetensors",
-      "model_version": "SD15",
       "height": 512,
       "width": 512,
       "batch_size": 1,

tests/m2/goldens/sd15_seed42.sha256

@@ -0,0 +1 @@
+e89344e544d4edfbd3ebe9a1c78dadb2729f53549666052b74ac7308f326f4fc

tests/m2/test_golden_image.py

@@ -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}"
+    )

tests/unit/test_characterization_controlnet.py

@@ -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))

tests/unit/test_characterization_inputs.py

@@ -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))

tests/unit/test_characterization_latents.py

@@ -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

tests/unit/test_characterization_sdxl_options.py

@@ -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)

tests/unit/test_chunks.py

@@ -1,37 +1,34 @@
-import pytest
+"""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 comfy.model_management import get_torch_device
-from coreml_suite.latents import chunk_batch, merge_chunks
-from coreml_suite.controlnet import chunk_control
-from coreml_suite.models import (
-    CoreMLInputs,
-)
-from coreml_suite.config import get_model_config
+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():
-    expected = {
+    return {
         "sample": {"shape": (2, 4, 64, 64)},
         "timestep": {"shape": (2,)},
         "timestep_cond": {"shape": (2, 256)},
-        "encoder_hidden_states": {"shape": (2, 768, 1, 77)},
+        "encoder_hidden_states": {"shape": (2, 77, 768)},
         "additional_residual_0": {"shape": (2, 320, 64, 64)},
         "additional_residual_1": {"shape": (2, 640, 32, 32)},
     }
-    return expected
-
-
-@pytest.fixture
-def model_config():
-    return get_model_config()
 
 
 @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(get_torch_device())
+    latent_image = torch.randn(batch_size, 4, 64, 64).to(CPU)
     target_shape = (4, 4, 64, 64)
 
     chunked = chunk_batch(latent_image, target_shape)
@@ -45,7 +42,7 @@ def test_batch_chunking(batch_size):
 
 @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(get_torch_device())
+    input_tensor = torch.randn(batch_size, 4, 64, 64).to(CPU)
     target_shape = (4, 4, 64, 64)
     chunked = chunk_batch(input_tensor, target_shape)
 
@@ -57,16 +54,16 @@ def test_merge_chunks(batch_size):
 
 @pytest.fixture
 def inputs():
-    x = torch.randn(1, 4, 64, 64).to(get_torch_device())
-    t = torch.randn([1]).to(get_torch_device())
-    c_crossattn = torch.randn(1, 77, 768).to(get_torch_device())
+    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(get_torch_device()),
-            torch.randn(1, 640, 32, 32).to(get_torch_device()),
+            torch.randn(1, 320, 64, 64).to(CPU),
+            torch.randn(1, 640, 32, 32).to(CPU),
         ],
     }
-    timestep_cond = torch.randn(1, 256).to(get_torch_device())
+    timestep_cond = torch.randn(1, 256).to(CPU)
 
     return CoreMLInputs(x, t, c_crossattn, control, timestep_cond=timestep_cond)
 
@@ -86,11 +83,11 @@ def inputs():
 def test_chunking_controlnet(b, target_size, num_chunks):
     cn = {
         "output": [
-            torch.randn(b, 320, 64, 64).to(get_torch_device()),
-            torch.randn(b, 640, 32, 32).to(get_torch_device()),
+            torch.randn(b, 320, 64, 64).to(CPU),
+            torch.randn(b, 640, 32, 32).to(CPU),
         ],
         "middle": [
-            torch.randn(b, 1280, 8, 8).to(get_torch_device()),
+            torch.randn(b, 1280, 8, 8).to(CPU),
         ],
     }
 

tests/unit/test_tier0_purity.py

@@ -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."
+    )