chore: drop dead experiment-scaffold references

- Remove the pytest collect_ignore_glob that skipped WIP scaffolds
  (test_experiments / test_unet_conversion / standalone_test). Those
  files no longer exist and only imported the experimental
  coreml_suite.experiments / convert_apple modules.
- Drop .gitignore entries for the removed experiment and alternate-env
  directories (experiments/, experiment_results/, apple_env/, comfy_env/,
  coremlsuite-venv/).

.github/workflows/tier0.yml

@@ -5,10 +5,10 @@ on:
     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.
+# Minimal-deps run: Tier 0 must work without ComfyUI, coremltools, or
+# python_coreml_stable_diffusion (Linux CI image won't have them). The
+# in-tree purity gate (tests/unit/test_tier0_purity.py) double-checks
+# that the suite hasn't started leaking framework imports.
 jobs:
   unit:
     runs-on: ubuntu-latest
@@ -16,12 +16,18 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
-      - uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v5
         with:
-          enable-cache: true
+          python-version: "3.11"
 
-      - name: uv sync
-        run: uv sync --no-install-project
+      - name: Install Tier 0 deps
+        run: |
+          python -m pip install --upgrade pip
+          # Tier 0 only needs torch + numpy + pytest; everything else
+          # is Mac-only.
+          python -m pip install \
+            "torch==2.0.1" "numpy<1.25" \
+            "pytest>=8" "pytest-xdist"
 
       - name: Run Tier 0
-        run: uv run pytest -m unit tests/ -v
+        run: pytest -m unit tests/ -v

.github/workflows/tier1.yml

@@ -0,0 +1,31 @@
+name: Tier 1 — Smoke (macOS-ARM)
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    # Gate behind the run-tier1 label too, so external PRs that touch
+    # only docs don't burn a minute of macOS-ARM time. Maintainers can
+    # always re-run via the run-tier1 label.
+    types: [opened, synchronize, reopened, labeled]
+
+jobs:
+  smoke:
+    if: |
+      github.event_name == 'push' ||
+      github.event.action != 'labeled' ||
+      contains(github.event.pull_request.labels.*.name, 'run-tier1')
+    runs-on: macos-14  # M1, Apple Silicon hosted runner
+    timeout-minutes: 20
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v3
+        with:
+          enable-cache: true
+
+      - name: uv sync
+        run: uv sync --no-install-project
+
+      - name: Run Tier 1 (synthetic micro-UNet smoke)
+        run: uv run pytest -m smoke tests/ -v

.github/workflows/tier2.yml

@@ -30,7 +30,7 @@ jobs:
       # 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
+      #   - PR label / dispatch -> the pinned requires-comfyui SHA + the frozen
       #     `comfy` uv group. Reproducible merge gate, immune to overnight drift.
       - name: Resolve ComfyUI ref + mode
         run: |
@@ -38,12 +38,10 @@ jobs:
             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"
+            PIN="$(sed -nE 's/^requires-comfyui *= *"==?([0-9a-f]+)".*/\1/p' pyproject.toml)"
+            if [ -z "$PIN" ]; then echo "could not parse requires-comfyui from pyproject.toml"; exit 1; fi
+            echo "COMFY_MODE=pinned" >> "$GITHUB_ENV"
+            echo "COMFY_REF=$PIN"    >> "$GITHUB_ENV"
           fi
 
       - name: Set up ComfyUI checkout
@@ -115,18 +113,10 @@ jobs:
           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).
+        # The golden-image workflow drives the Core ML Converter node, so the
+        # UNet is converted on demand on the first run and reused from the
+        # runner-local cache afterwards.
         run: uv run --no-sync pytest -m m2 tests/ -v
 
       - name: Stop ComfyUI server

.gitignore

@@ -3,6 +3,4 @@ __pycache__/
 models/
 .venv/
 test_results/
-*.log
-.DS_Store
-.claude/
+tests/m2/_latest_generated.png

.python-version

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

LICENSE

@@ -1,21 +1,674 @@
-MIT License
+                    GNU GENERAL PUBLIC LICENSE
+                       Version 3, 29 June 2007
 
-Copyright (c) 2023-2026 Adrian Szczepański
+ Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
+ Everyone is permitted to copy and distribute verbatim copies
+ of this license document, but changing it is not allowed.
 
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
+                            Preamble
 
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
+  The GNU General Public License is a free, copyleft license for
+software and other kinds of works.
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+  The licenses for most software and other practical works are designed
+to take away your freedom to share and change the works.  By contrast,
+the GNU General Public License is intended to guarantee your freedom to
+share and change all versions of a program--to make sure it remains free
+software for all its users.  We, the Free Software Foundation, use the
+GNU General Public License for most of our software; it applies also to
+any other work released this way by its authors.  You can apply it to
+your programs, too.
+
+  When we speak of free software, we are referring to freedom, not
+price.  Our General Public Licenses are designed to make sure that you
+have the freedom to distribute copies of free software (and charge for
+them if you wish), that you receive source code or can get it if you
+want it, that you can change the software or use pieces of it in new
+free programs, and that you know you can do these things.
+
+  To protect your rights, we need to prevent others from denying you
+these rights or asking you to surrender the rights.  Therefore, you have
+certain responsibilities if you distribute copies of the software, or if
+you modify it: responsibilities to respect the freedom of others.
+
+  For example, if you distribute copies of such a program, whether
+gratis or for a fee, you must pass on to the recipients the same
+freedoms that you received.  You must make sure that they, too, receive
+or can get the source code.  And you must show them these terms so they
+know their rights.
+
+  Developers that use the GNU GPL protect your rights with two steps:
+(1) assert copyright on the software, and (2) offer you this License
+giving you legal permission to copy, distribute and/or modify it.
+
+  For the developers' and authors' protection, the GPL clearly explains
+that there is no warranty for this free software.  For both users' and
+authors' sake, the GPL requires that modified versions be marked as
+changed, so that their problems will not be attributed erroneously to
+authors of previous versions.
+
+  Some devices are designed to deny users access to install or run
+modified versions of the software inside them, although the manufacturer
+can do so.  This is fundamentally incompatible with the aim of
+protecting users' freedom to change the software.  The systematic
+pattern of such abuse occurs in the area of products for individuals to
+use, which is precisely where it is most unacceptable.  Therefore, we
+have designed this version of the GPL to prohibit the practice for those
+products.  If such problems arise substantially in other domains, we
+stand ready to extend this provision to those domains in future versions
+of the GPL, as needed to protect the freedom of users.
+
+  Finally, every program is threatened constantly by software patents.
+States should not allow patents to restrict development and use of
+software on general-purpose computers, but in those that do, we wish to
+avoid the special danger that patents applied to a free program could
+make it effectively proprietary.  To prevent this, the GPL assures that
+patents cannot be used to render the program non-free.
+
+  The precise terms and conditions for copying, distribution and
+modification follow.
+
+                       TERMS AND CONDITIONS
+
+  0. Definitions.
+
+  "This License" refers to version 3 of the GNU General Public License.
+
+  "Copyright" also means copyright-like laws that apply to other kinds of
+works, such as semiconductor masks.
+
+  "The Program" refers to any copyrightable work licensed under this
+License.  Each licensee is addressed as "you".  "Licensees" and
+"recipients" may be individuals or organizations.
+
+  To "modify" a work means to copy from or adapt all or part of the work
+in a fashion requiring copyright permission, other than the making of an
+exact copy.  The resulting work is called a "modified version" of the
+earlier work or a work "based on" the earlier work.
+
+  A "covered work" means either the unmodified Program or a work based
+on the Program.
+
+  To "propagate" a work means to do anything with it that, without
+permission, would make you directly or secondarily liable for
+infringement under applicable copyright law, except executing it on a
+computer or modifying a private copy.  Propagation includes copying,
+distribution (with or without modification), making available to the
+public, and in some countries other activities as well.
+
+  To "convey" a work means any kind of propagation that enables other
+parties to make or receive copies.  Mere interaction with a user through
+a computer network, with no transfer of a copy, is not conveying.
+
+  An interactive user interface displays "Appropriate Legal Notices"
+to the extent that it includes a convenient and prominently visible
+feature that (1) displays an appropriate copyright notice, and (2)
+tells the user that there is no warranty for the work (except to the
+extent that warranties are provided), that licensees may convey the
+work under this License, and how to view a copy of this License.  If
+the interface presents a list of user commands or options, such as a
+menu, a prominent item in the list meets this criterion.
+
+  1. Source Code.
+
+  The "source code" for a work means the preferred form of the work
+for making modifications to it.  "Object code" means any non-source
+form of a work.
+
+  A "Standard Interface" means an interface that either is an official
+standard defined by a recognized standards body, or, in the case of
+interfaces specified for a particular programming language, one that
+is widely used among developers working in that language.
+
+  The "System Libraries" of an executable work include anything, other
+than the work as a whole, that (a) is included in the normal form of
+packaging a Major Component, but which is not part of that Major
+Component, and (b) serves only to enable use of the work with that
+Major Component, or to implement a Standard Interface for which an
+implementation is available to the public in source code form.  A
+"Major Component", in this context, means a major essential component
+(kernel, window system, and so on) of the specific operating system
+(if any) on which the executable work runs, or a compiler used to
+produce the work, or an object code interpreter used to run it.
+
+  The "Corresponding Source" for a work in object code form means all
+the source code needed to generate, install, and (for an executable
+work) run the object code and to modify the work, including scripts to
+control those activities.  However, it does not include the work's
+System Libraries, or general-purpose tools or generally available free
+programs which are used unmodified in performing those activities but
+which are not part of the work.  For example, Corresponding Source
+includes interface definition files associated with source files for
+the work, and the source code for shared libraries and dynamically
+linked subprograms that the work is specifically designed to require,
+such as by intimate data communication or control flow between those
+subprograms and other parts of the work.
+
+  The Corresponding Source need not include anything that users
+can regenerate automatically from other parts of the Corresponding
+Source.
+
+  The Corresponding Source for a work in source code form is that
+same work.
+
+  2. Basic Permissions.
+
+  All rights granted under this License are granted for the term of
+copyright on the Program, and are irrevocable provided the stated
+conditions are met.  This License explicitly affirms your unlimited
+permission to run the unmodified Program.  The output from running a
+covered work is covered by this License only if the output, given its
+content, constitutes a covered work.  This License acknowledges your
+rights of fair use or other equivalent, as provided by copyright law.
+
+  You may make, run and propagate covered works that you do not
+convey, without conditions so long as your license otherwise remains
+in force.  You may convey covered works to others for the sole purpose
+of having them make modifications exclusively for you, or provide you
+with facilities for running those works, provided that you comply with
+the terms of this License in conveying all material for which you do
+not control copyright.  Those thus making or running the covered works
+for you must do so exclusively on your behalf, under your direction
+and control, on terms that prohibit them from making any copies of
+your copyrighted material outside their relationship with you.
+
+  Conveying under any other circumstances is permitted solely under
+the conditions stated below.  Sublicensing is not allowed; section 10
+makes it unnecessary.
+
+  3. Protecting Users' Legal Rights From Anti-Circumvention Law.
+
+  No covered work shall be deemed part of an effective technological
+measure under any applicable law fulfilling obligations under article
+11 of the WIPO copyright treaty adopted on 20 December 1996, or
+similar laws prohibiting or restricting circumvention of such
+measures.
+
+  When you convey a covered work, you waive any legal power to forbid
+circumvention of technological measures to the extent such circumvention
+is effected by exercising rights under this License with respect to
+the covered work, and you disclaim any intention to limit operation or
+modification of the work as a means of enforcing, against the work's
+users, your or third parties' legal rights to forbid circumvention of
+technological measures.
+
+  4. Conveying Verbatim Copies.
+
+  You may convey verbatim copies of the Program's source code as you
+receive it, in any medium, provided that you conspicuously and
+appropriately publish on each copy an appropriate copyright notice;
+keep intact all notices stating that this License and any
+non-permissive terms added in accord with section 7 apply to the code;
+keep intact all notices of the absence of any warranty; and give all
+recipients a copy of this License along with the Program.
+
+  You may charge any price or no price for each copy that you convey,
+and you may offer support or warranty protection for a fee.
+
+  5. Conveying Modified Source Versions.
+
+  You may convey a work based on the Program, or the modifications to
+produce it from the Program, in the form of source code under the
+terms of section 4, provided that you also meet all of these conditions:
+
+    a) The work must carry prominent notices stating that you modified
+    it, and giving a relevant date.
+
+    b) The work must carry prominent notices stating that it is
+    released under this License and any conditions added under section
+    7.  This requirement modifies the requirement in section 4 to
+    "keep intact all notices".
+
+    c) You must license the entire work, as a whole, under this
+    License to anyone who comes into possession of a copy.  This
+    License will therefore apply, along with any applicable section 7
+    additional terms, to the whole of the work, and all its parts,
+    regardless of how they are packaged.  This License gives no
+    permission to license the work in any other way, but it does not
+    invalidate such permission if you have separately received it.
+
+    d) If the work has interactive user interfaces, each must display
+    Appropriate Legal Notices; however, if the Program has interactive
+    interfaces that do not display Appropriate Legal Notices, your
+    work need not make them do so.
+
+  A compilation of a covered work with other separate and independent
+works, which are not by their nature extensions of the covered work,
+and which are not combined with it such as to form a larger program,
+in or on a volume of a storage or distribution medium, is called an
+"aggregate" if the compilation and its resulting copyright are not
+used to limit the access or legal rights of the compilation's users
+beyond what the individual works permit.  Inclusion of a covered work
+in an aggregate does not cause this License to apply to the other
+parts of the aggregate.
+
+  6. Conveying Non-Source Forms.
+
+  You may convey a covered work in object code form under the terms
+of sections 4 and 5, provided that you also convey the
+machine-readable Corresponding Source under the terms of this License,
+in one of these ways:
+
+    a) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by the
+    Corresponding Source fixed on a durable physical medium
+    customarily used for software interchange.
+
+    b) Convey the object code in, or embodied in, a physical product
+    (including a physical distribution medium), accompanied by a
+    written offer, valid for at least three years and valid for as
+    long as you offer spare parts or customer support for that product
+    model, to give anyone who possesses the object code either (1) a
+    copy of the Corresponding Source for all the software in the
+    product that is covered by this License, on a durable physical
+    medium customarily used for software interchange, for a price no
+    more than your reasonable cost of physically performing this
+    conveying of source, or (2) access to copy the
+    Corresponding Source from a network server at no charge.
+
+    c) Convey individual copies of the object code with a copy of the
+    written offer to provide the Corresponding Source.  This
+    alternative is allowed only occasionally and noncommercially, and
+    only if you received the object code with such an offer, in accord
+    with subsection 6b.
+
+    d) Convey the object code by offering access from a designated
+    place (gratis or for a charge), and offer equivalent access to the
+    Corresponding Source in the same way through the same place at no
+    further charge.  You need not require recipients to copy the
+    Corresponding Source along with the object code.  If the place to
+    copy the object code is a network server, the Corresponding Source
+    may be on a different server (operated by you or a third party)
+    that supports equivalent copying facilities, provided you maintain
+    clear directions next to the object code saying where to find the
+    Corresponding Source.  Regardless of what server hosts the
+    Corresponding Source, you remain obligated to ensure that it is
+    available for as long as needed to satisfy these requirements.
+
+    e) Convey the object code using peer-to-peer transmission, provided
+    you inform other peers where the object code and Corresponding
+    Source of the work are being offered to the general public at no
+    charge under subsection 6d.
+
+  A separable portion of the object code, whose source code is excluded
+from the Corresponding Source as a System Library, need not be
+included in conveying the object code work.
+
+  A "User Product" is either (1) a "consumer product", which means any
+tangible personal property which is normally used for personal, family,
+or household purposes, or (2) anything designed or sold for incorporation
+into a dwelling.  In determining whether a product is a consumer product,
+doubtful cases shall be resolved in favor of coverage.  For a particular
+product received by a particular user, "normally used" refers to a
+typical or common use of that class of product, regardless of the status
+of the particular user or of the way in which the particular user
+actually uses, or expects or is expected to use, the product.  A product
+is a consumer product regardless of whether the product has substantial
+commercial, industrial or non-consumer uses, unless such uses represent
+the only significant mode of use of the product.
+
+  "Installation Information" for a User Product means any methods,
+procedures, authorization keys, or other information required to install
+and execute modified versions of a covered work in that User Product from
+a modified version of its Corresponding Source.  The information must
+suffice to ensure that the continued functioning of the modified object
+code is in no case prevented or interfered with solely because
+modification has been made.
+
+  If you convey an object code work under this section in, or with, or
+specifically for use in, a User Product, and the conveying occurs as
+part of a transaction in which the right of possession and use of the
+User Product is transferred to the recipient in perpetuity or for a
+fixed term (regardless of how the transaction is characterized), the
+Corresponding Source conveyed under this section must be accompanied
+by the Installation Information.  But this requirement does not apply
+if neither you nor any third party retains the ability to install
+modified object code on the User Product (for example, the work has
+been installed in ROM).
+
+  The requirement to provide Installation Information does not include a
+requirement to continue to provide support service, warranty, or updates
+for a work that has been modified or installed by the recipient, or for
+the User Product in which it has been modified or installed.  Access to a
+network may be denied when the modification itself materially and
+adversely affects the operation of the network or violates the rules and
+protocols for communication across the network.
+
+  Corresponding Source conveyed, and Installation Information provided,
+in accord with this section must be in a format that is publicly
+documented (and with an implementation available to the public in
+source code form), and must require no special password or key for
+unpacking, reading or copying.
+
+  7. Additional Terms.
+
+  "Additional permissions" are terms that supplement the terms of this
+License by making exceptions from one or more of its conditions.
+Additional permissions that are applicable to the entire Program shall
+be treated as though they were included in this License, to the extent
+that they are valid under applicable law.  If additional permissions
+apply only to part of the Program, that part may be used separately
+under those permissions, but the entire Program remains governed by
+this License without regard to the additional permissions.
+
+  When you convey a copy of a covered work, you may at your option
+remove any additional permissions from that copy, or from any part of
+it.  (Additional permissions may be written to require their own
+removal in certain cases when you modify the work.)  You may place
+additional permissions on material, added by you to a covered work,
+for which you have or can give appropriate copyright permission.
+
+  Notwithstanding any other provision of this License, for material you
+add to a covered work, you may (if authorized by the copyright holders of
+that material) supplement the terms of this License with terms:
+
+    a) Disclaiming warranty or limiting liability differently from the
+    terms of sections 15 and 16 of this License; or
+
+    b) Requiring preservation of specified reasonable legal notices or
+    author attributions in that material or in the Appropriate Legal
+    Notices displayed by works containing it; or
+
+    c) Prohibiting misrepresentation of the origin of that material, or
+    requiring that modified versions of such material be marked in
+    reasonable ways as different from the original version; or
+
+    d) Limiting the use for publicity purposes of names of licensors or
+    authors of the material; or
+
+    e) Declining to grant rights under trademark law for use of some
+    trade names, trademarks, or service marks; or
+
+    f) Requiring indemnification of licensors and authors of that
+    material by anyone who conveys the material (or modified versions of
+    it) with contractual assumptions of liability to the recipient, for
+    any liability that these contractual assumptions directly impose on
+    those licensors and authors.
+
+  All other non-permissive additional terms are considered "further
+restrictions" within the meaning of section 10.  If the Program as you
+received it, or any part of it, contains a notice stating that it is
+governed by this License along with a term that is a further
+restriction, you may remove that term.  If a license document contains
+a further restriction but permits relicensing or conveying under this
+License, you may add to a covered work material governed by the terms
+of that license document, provided that the further restriction does
+not survive such relicensing or conveying.
+
+  If you add terms to a covered work in accord with this section, you
+must place, in the relevant source files, a statement of the
+additional terms that apply to those files, or a notice indicating
+where to find the applicable terms.
+
+  Additional terms, permissive or non-permissive, may be stated in the
+form of a separately written license, or stated as exceptions;
+the above requirements apply either way.
+
+  8. Termination.
+
+  You may not propagate or modify a covered work except as expressly
+provided under this License.  Any attempt otherwise to propagate or
+modify it is void, and will automatically terminate your rights under
+this License (including any patent licenses granted under the third
+paragraph of section 11).
+
+  However, if you cease all violation of this License, then your
+license from a particular copyright holder is reinstated (a)
+provisionally, unless and until the copyright holder explicitly and
+finally terminates your license, and (b) permanently, if the copyright
+holder fails to notify you of the violation by some reasonable means
+prior to 60 days after the cessation.
+
+  Moreover, your license from a particular copyright holder is
+reinstated permanently if the copyright holder notifies you of the
+violation by some reasonable means, this is the first time you have
+received notice of violation of this License (for any work) from that
+copyright holder, and you cure the violation prior to 30 days after
+your receipt of the notice.
+
+  Termination of your rights under this section does not terminate the
+licenses of parties who have received copies or rights from you under
+this License.  If your rights have been terminated and not permanently
+reinstated, you do not qualify to receive new licenses for the same
+material under section 10.
+
+  9. Acceptance Not Required for Having Copies.
+
+  You are not required to accept this License in order to receive or
+run a copy of the Program.  Ancillary propagation of a covered work
+occurring solely as a consequence of using peer-to-peer transmission
+to receive a copy likewise does not require acceptance.  However,
+nothing other than this License grants you permission to propagate or
+modify any covered work.  These actions infringe copyright if you do
+not accept this License.  Therefore, by modifying or propagating a
+covered work, you indicate your acceptance of this License to do so.
+
+  10. Automatic Licensing of Downstream Recipients.
+
+  Each time you convey a covered work, the recipient automatically
+receives a license from the original licensors, to run, modify and
+propagate that work, subject to this License.  You are not responsible
+for enforcing compliance by third parties with this License.
+
+  An "entity transaction" is a transaction transferring control of an
+organization, or substantially all assets of one, or subdividing an
+organization, or merging organizations.  If propagation of a covered
+work results from an entity transaction, each party to that
+transaction who receives a copy of the work also receives whatever
+licenses to the work the party's predecessor in interest had or could
+give under the previous paragraph, plus a right to possession of the
+Corresponding Source of the work from the predecessor in interest, if
+the predecessor has it or can get it with reasonable efforts.
+
+  You may not impose any further restrictions on the exercise of the
+rights granted or affirmed under this License.  For example, you may
+not impose a license fee, royalty, or other charge for exercise of
+rights granted under this License, and you may not initiate litigation
+(including a cross-claim or counterclaim in a lawsuit) alleging that
+any patent claim is infringed by making, using, selling, offering for
+sale, or importing the Program or any portion of it.
+
+  11. Patents.
+
+  A "contributor" is a copyright holder who authorizes use under this
+License of the Program or a work on which the Program is based.  The
+work thus licensed is called the contributor's "contributor version".
+
+  A contributor's "essential patent claims" are all patent claims
+owned or controlled by the contributor, whether already acquired or
+hereafter acquired, that would be infringed by some manner, permitted
+by this License, of making, using, or selling its contributor version,
+but do not include claims that would be infringed only as a
+consequence of further modification of the contributor version.  For
+purposes of this definition, "control" includes the right to grant
+patent sublicenses in a manner consistent with the requirements of
+this License.
+
+  Each contributor grants you a non-exclusive, worldwide, royalty-free
+patent license under the contributor's essential patent claims, to
+make, use, sell, offer for sale, import and otherwise run, modify and
+propagate the contents of its contributor version.
+
+  In the following three paragraphs, a "patent license" is any express
+agreement or commitment, however denominated, not to enforce a patent
+(such as an express permission to practice a patent or covenant not to
+sue for patent infringement).  To "grant" such a patent license to a
+party means to make such an agreement or commitment not to enforce a
+patent against the party.
+
+  If you convey a covered work, knowingly relying on a patent license,
+and the Corresponding Source of the work is not available for anyone
+to copy, free of charge and under the terms of this License, through a
+publicly available network server or other readily accessible means,
+then you must either (1) cause the Corresponding Source to be so
+available, or (2) arrange to deprive yourself of the benefit of the
+patent license for this particular work, or (3) arrange, in a manner
+consistent with the requirements of this License, to extend the patent
+license to downstream recipients.  "Knowingly relying" means you have
+actual knowledge that, but for the patent license, your conveying the
+covered work in a country, or your recipient's use of the covered work
+in a country, would infringe one or more identifiable patents in that
+country that you have reason to believe are valid.
+
+  If, pursuant to or in connection with a single transaction or
+arrangement, you convey, or propagate by procuring conveyance of, a
+covered work, and grant a patent license to some of the parties
+receiving the covered work authorizing them to use, propagate, modify
+or convey a specific copy of the covered work, then the patent license
+you grant is automatically extended to all recipients of the covered
+work and works based on it.
+
+  A patent license is "discriminatory" if it does not include within
+the scope of its coverage, prohibits the exercise of, or is
+conditioned on the non-exercise of one or more of the rights that are
+specifically granted under this License.  You may not convey a covered
+work if you are a party to an arrangement with a third party that is
+in the business of distributing software, under which you make payment
+to the third party based on the extent of your activity of conveying
+the work, and under which the third party grants, to any of the
+parties who would receive the covered work from you, a discriminatory
+patent license (a) in connection with copies of the covered work
+conveyed by you (or copies made from those copies), or (b) primarily
+for and in connection with specific products or compilations that
+contain the covered work, unless you entered into that arrangement,
+or that patent license was granted, prior to 28 March 2007.
+
+  Nothing in this License shall be construed as excluding or limiting
+any implied license or other defenses to infringement that may
+otherwise be available to you under applicable patent law.
+
+  12. No Surrender of Others' Freedom.
+
+  If conditions are imposed on you (whether by court order, agreement or
+otherwise) that contradict the conditions of this License, they do not
+excuse you from the conditions of this License.  If you cannot convey a
+covered work so as to satisfy simultaneously your obligations under this
+License and any other pertinent obligations, then as a consequence you may
+not convey it at all.  For example, if you agree to terms that obligate you
+to collect a royalty for further conveying from those to whom you convey
+the Program, the only way you could satisfy both those terms and this
+License would be to refrain entirely from conveying the Program.
+
+  13. Use with the GNU Affero General Public License.
+
+  Notwithstanding any other provision of this License, you have
+permission to link or combine any covered work with a work licensed
+under version 3 of the GNU Affero General Public License into a single
+combined work, and to convey the resulting work.  The terms of this
+License will continue to apply to the part which is the covered work,
+but the special requirements of the GNU Affero General Public License,
+section 13, concerning interaction through a network will apply to the
+combination as such.
+
+  14. Revised Versions of this License.
+
+  The Free Software Foundation may publish revised and/or new versions of
+the GNU General Public License from time to time.  Such new versions will
+be similar in spirit to the present version, but may differ in detail to
+address new problems or concerns.
+
+  Each version is given a distinguishing version number.  If the
+Program specifies that a certain numbered version of the GNU General
+Public License "or any later version" applies to it, you have the
+option of following the terms and conditions either of that numbered
+version or of any later version published by the Free Software
+Foundation.  If the Program does not specify a version number of the
+GNU General Public License, you may choose any version ever published
+by the Free Software Foundation.
+
+  If the Program specifies that a proxy can decide which future
+versions of the GNU General Public License can be used, that proxy's
+public statement of acceptance of a version permanently authorizes you
+to choose that version for the Program.
+
+  Later license versions may give you additional or different
+permissions.  However, no additional obligations are imposed on any
+author or copyright holder as a result of your choosing to follow a
+later version.
+
+  15. Disclaimer of Warranty.
+
+  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
+APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
+HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
+OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
+THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
+IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
+ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
+
+  16. Limitation of Liability.
+
+  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
+WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
+THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
+GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
+USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
+DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
+PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
+EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
+SUCH DAMAGES.
+
+  17. Interpretation of Sections 15 and 16.
+
+  If the disclaimer of warranty and limitation of liability provided
+above cannot be given local legal effect according to their terms,
+reviewing courts shall apply local law that most closely approximates
+an absolute waiver of all civil liability in connection with the
+Program, unless a warranty or assumption of liability accompanies a
+copy of the Program in return for a fee.
+
+                     END OF TERMS AND CONDITIONS
+
+            How to Apply These Terms to Your New Programs
+
+  If you develop a new program, and you want it to be of the greatest
+possible use to the public, the best way to achieve this is to make it
+free software which everyone can redistribute and change under these terms.
+
+  To do so, attach the following notices to the program.  It is safest
+to attach them to the start of each source file to most effectively
+state the exclusion of warranty; and each file should have at least
+the "copyright" line and a pointer to where the full notice is found.
+
+    <one line to give the program's name and a brief idea of what it does.>
+    Copyright (C) <year>  <name of author>
+
+    This program is free software: you can redistribute it and/or modify
+    it under the terms of the GNU General Public License as published by
+    the Free Software Foundation, either version 3 of the License, or
+    (at your option) any later version.
+
+    This program is distributed in the hope that it will be useful,
+    but WITHOUT ANY WARRANTY; without even the implied warranty of
+    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+    GNU General Public License for more details.
+
+    You should have received a copy of the GNU General Public License
+    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+Also add information on how to contact you by electronic and paper mail.
+
+  If the program does terminal interaction, make it output a short
+notice like this when it starts in an interactive mode:
+
+    <program>  Copyright (C) <year>  <name of author>
+    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
+    This is free software, and you are welcome to redistribute it
+    under certain conditions; type `show c' for details.
+
+The hypothetical commands `show w' and `show c' should show the appropriate
+parts of the General Public License.  Of course, your program's commands
+might be different; for a GUI interface, you would use an "about box".
+
+  You should also get your employer (if you work as a programmer) or school,
+if any, to sign a "copyright disclaimer" for the program, if necessary.
+For more information on this, and how to apply and follow the GNU GPL, see
+<https://www.gnu.org/licenses/>.
+
+  The GNU General Public License does not permit incorporating your program
+into proprietary programs.  If your program is a subroutine library, you
+may consider it more useful to permit linking proprietary applications with
+the library.  If this is what you want to do, use the GNU Lesser General
+Public License instead of this License.  But first, please read
+<https://www.gnu.org/licenses/why-not-lgpl.html>.

README.md

@@ -1,113 +1,430 @@
 # Core ML Suite for ComfyUI
 
-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.
+## Overview
 
-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.
+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.
 
-> [!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.
+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).
 
-## Installation
+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.
 
-### ComfyUI-Manager (recommended)
+## Getting Started
 
-Open **Manager → Install Custom Nodes**, search for `Core ML`, click
-**Install**, and restart ComfyUI.
+To start using custom nodes in your ComfyUI, follow these simple steps:
 
-### Manual
+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.
 
-```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
-```
+That's it! You're now ready to start enhancing your ComfyUI workflows with Core ML models.
 
-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.
-
-## Quickstart
-
-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.
-
-See [docs/workflows.md](docs/workflows.md) for complete example graphs (txt2img,
-ControlNet, LoRA, LCM, SDXL).
-
-## Which compute unit should I pick?
-
-The **compute unit** selects the hardware Core ML runs on. Pair it with the
-attention implementation chosen at conversion time:
-
-| 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) |
-
-`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).
-
-## Documentation
-
-- [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.
+- 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** — 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.
+- **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.
+
+## Installation
+
+### Using ComfyUI-Manager
+
+The easiest way to install the custom nodes is to use the ComfyUI-Manager. You can find the installation instructions
+[here](https://github.com/ltdrdata/ComfyUI-Manager#installation). Once you've installed the ComfyUI-Manager, you can
+install the custom nodes by following these steps:
+
+- Open the ComfyUI-Manager by clicking the `Manager` button in the ComfyUI toolbar.
+- Click the `Install Custom Nodes` button.
+- Search for `Core ML` and click the `Install` button.
+- Restart ComfyUI.
+
+### Manual Installation
+
+1. Clone this repository into the custom_nodes directory of your ComfyUI. If you're not sure how to do this, you can
+   download the repository as a zip file and extract it into the same directory.
+    ```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:
+
+    ```bash
+    cd /path/to/comfyui/custom_nodes/ComfyUI-CoreMLSuite
+    pip install -r requirements.txt
+    ```
+
+## How to use
+
+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.
+
+### Available Nodes
+
+#### Core ML UNet Loader (`CoreMLUnetLoader`)
+
+![CoreMLUnetLoader](./assets/unet_loader.png?raw=true)
+
+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.
+
+- **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.
+
+#### Core ML Sampler (`CoreMLSampler`)
+
+![CoreMLSampler](./assets/sampler.png?raw=true)
+
+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)
 
 > [!IMPORTANT]
-> **Breaking change in 2.0.0.** The converted Core ML UNet now takes
-> `encoder_hidden_states` in the native `diffusers` layout
-> `(batch, tokens, hidden)` instead of the previous `(batch, hidden, 1, tokens)`.
-> Models converted with earlier versions are not compatible and must be
-> re-converted.
+> 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.
 
-## Acknowledgements
+![sdxl](./assets/sdxl_conversion.png?raw=true)
 
-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.
+## Quantization (opt-in)
+
+The `Core ML Converter` and `Core ML LCM Converter` nodes accept an
+optional `quantize_nbits` dropdown that runs k-means weight palettization
+(`coremltools.optimize.coreml.palettize_weights`) on the UNet before save.
+
+Values: `none` (default — no quantization, identical to unquantized
+behavior and filenames), `8`, `6`, `4`. The number is appended to the
+.mlpackage stem as `_q<bits>` so quantized and unquantized variants
+coexist on disk and in cache.
+
+### SD1.5 1×512×512 SPLIT_EINSUM tradeoffs (M2 Pro, ANE)
+
+Measured with 20 UNet forward passes at a fixed seed for the PSNR
+comparison:
+
+| nbits | size (MB) | size vs none | fwd median (ms) | PSNR vs `none` (dB) |
+|---|---:|---:|---:|---:|
+| none | 1641 | 1.000 | 197.1 | — |
+| 8    |  822 | 0.501 | 186.6 | 53.5 |
+| 6    |  617 | 0.376 | 183.0 | 40.2 |
+| 4    |  412 | 0.251 | 179.8 | 27.5 |
+
+PSNR here is computed on the raw `noise_pred` output of a single UNet
+forward at a fixed seed, not on the final decoded image — it isolates
+the quantization-induced drift from sampler / VAE noise. Final-image
+PSNR is comfortably higher (the sampler averages over 20 steps).
+
+### Recommended settings per chip / RAM
+
+- **8 GB RAM (M1 base, M2 base):** `nbits=4`. ~4× smaller model, still
+  loads, PSNR 27 dB is visually identical at SD1.5 sizes.
+- **16 GB RAM (M1/M2/M3 Pro):** `nbits=6` is the sweet spot — ~2.7×
+  smaller, PSNR 40 dB, no perceptible quality drop.
+- **32 GB+ RAM (Max / Ultra):** `nbits=8` if you want the safety
+  margin, `none` if you want bit-identical output for golden testing.
+
+The default stays `none` so existing workflows produce byte-for-byte
+identical output — the golden-image anchor (`tests/m2/test_golden_image.py`)
+verifies this on every Tier 2 run.
+
+## 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.
 
 ## Support
 
-Questions or suggestions? Open an
-[issue](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/issues).
+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.

__init__.py

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

coreml_suite/config.py

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

coreml_suite/converter.py

@@ -0,0 +1,383 @@
+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,
+    quantize_nbits: str = "none",
+):
+    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()
+
+    if quantize_nbits != "none":
+        # Opt-in k-means weight palettization. The default path
+        # (quantize_nbits="none") leaves the traced UNet untouched.
+        from coremltools.optimize.coreml import (
+            OpPalettizerConfig,
+            OptimizationConfig,
+            palettize_weights,
+        )
+
+        nbits = int(quantize_nbits)
+        logger.info(f"Palettizing UNet weights to {nbits}-bit (kmeans)..")
+        t0 = time.time()
+        cfg = OptimizationConfig(
+            global_config=OpPalettizerConfig(mode="kmeans", nbits=nbits)
+        )
+        coreml_unet = palettize_weights(coreml_unet, config=cfg)
+        logger.info(f"Palettization took {time.time() - t0:.1f}s")
+
+    coreml_unet.save(unet_out_path)
+    logger.info(f"Saved unet into {unet_out_path}")
+
+
+def convert(
+    ckpt_path: str,
+    model_version: ModelVersion,
+    unet_out_path: str,
+    batch_size: int = 1,
+    sample_size: tuple[int, int] = (64, 64),
+    controlnet_support: bool = False,
+    lora_weights: list[tuple[Union[str, os.PathLike], float]] = None,
+    attn_impl: str = AttentionImplementations.SPLIT_EINSUM.name,
+    config_path: str = None,
+    quantize_nbits: str = "none",
+):
+    if os.path.exists(unet_out_path):
+        logger.info(f"Found existing model at {unet_out_path}! Skipping..")
+        return
+
+    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,
+        quantize_nbits=quantize_nbits,
+    )
+
+
+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/inputs.py

@@ -24,6 +24,7 @@ class CoreMLInputs:
         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)
 
@@ -56,7 +57,8 @@ class CoreMLInputs:
     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"]
+        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]))

coreml_suite/core/naming.py

@@ -0,0 +1,68 @@
+"""Pure out_name composition for the Core ML UNet artifact.
+
+Extracted from CoreMLConverter.convert so the filename contract
+can be tested + reused without instantiating the node. The string is the
+cache key: every workflow that references a converted .mlpackage depends
+on it staying byte-for-byte identical.
+"""
+from typing import Iterable, Tuple
+
+ATTN_SUFFIX = {
+    "SPLIT_EINSUM": "se",
+    "SPLIT_EINSUM_V2": "se2",
+    "ORIGINAL": "orig",
+}
+
+# Palettization bits. "none" = no quantization (default; keeps the
+# unquantized filename intact so existing workflows still resolve their
+# cached .mlpackage). Numeric values append a `_q<bits>` suffix.
+QUANT_NBITS_VALUES = ("none", "8", "6", "4")
+
+
+def compose_out_name(
+    *,
+    ckpt_name: str,
+    batch_size: int,
+    width: int,
+    height: int,
+    controlnet_support: bool,
+    attention_implementation: str,
+    lora_names: Iterable[str] = (),
+    quantize_nbits: str = "none",
+) -> str:
+    """Build the .mlpackage stem from convert() parameters.
+
+    Locked behaviour (characterization tests):
+      - first '.' in ckpt_name wins (`a.b.c.safetensors` -> `a`)
+      - spaces collapse to underscores
+      - LoRA names are taken stem-only, sorted, joined with '_' and
+        prefixed with '_' when present (caller is expected to pass a
+        sorted list; we sort defensively)
+      - controlnet adds `_cn`
+      - attn suffix is `_se` | `_se2` | `_orig`
+
+    Quantization:
+      - quantize_nbits "none" (default) appends nothing — existing
+        unquantized .mlpackages keep the old filename
+      - "4" / "6" / "8" appends `_q<bits>` after the attn suffix
+    """
+    if quantize_nbits not in QUANT_NBITS_VALUES:
+        raise ValueError(
+            f"quantize_nbits={quantize_nbits!r} not in {QUANT_NBITS_VALUES}"
+        )
+    stem = ckpt_name.split(".")[0]
+    sorted_names = sorted(lora_names)
+    lora_str = "_" + "_".join(name.split(".")[0] for name in sorted_names) if sorted_names else ""
+    cn_suffix = "_cn" if controlnet_support else ""
+    attn_suffix = "_" + ATTN_SUFFIX[attention_implementation]
+    quant_suffix = f"_q{quantize_nbits}" if quantize_nbits != "none" else ""
+    out_name = (
+        f"{stem}{lora_str}_{batch_size}x{width}x{height}"
+        f"{cn_suffix}{attn_suffix}{quant_suffix}"
+    )
+    return out_name.replace(" ", "_")
+
+
+def lora_names_from_params(lora_params: Iterable[Tuple[str, float]]) -> list[str]:
+    """Mirror the sort applied inside CoreMLConverter.convert."""
+    return [name for name, _ in sorted(lora_params, key=lambda pair: pair[0])]

coreml_suite/coreml_model.py

@@ -1,42 +0,0 @@
-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/lcm/__init__.py

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

coreml_suite/lcm/converter.py

@@ -0,0 +1,297 @@
+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

@@ -0,0 +1,70 @@
+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

@@ -0,0 +1,99 @@
+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/nodes.py

@@ -1,10 +1,18 @@
 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.coreml_model import CoreMLModel
+from coreml_suite import converter
+from coreml_suite.config import ModelVersion
+from coreml_suite.core.naming import (
+    QUANT_NBITS_VALUES,
+    compose_out_name,
+    lora_names_from_params,
+)
 from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
 from coreml_suite.logger import logger
 from nodes import KSampler, LoraLoader, KSamplerAdvanced
@@ -17,26 +25,6 @@ from coreml_suite.models import (
 )
 
 
-def _discover(fn_name, fallback):
-    """Populate a converter dropdown from coreml_diffusion's discovery API.
-
-    Fails soft: if the package is missing, too old to expose ``fn_name``, or
-    errors, the node still registers with the fallback list instead of vanishing
-    from the menu. Evaluated on every INPUT_TYPES call, so installing a newer
-    coreml_diffusion surfaces new conversion types with no Suite change.
-    """
-    try:
-        import coreml_diffusion
-
-        return getattr(coreml_diffusion, fn_name)()
-    except Exception as exc:  # missing/old package, import error, etc.
-        logger.warning(
-            f"coreml_diffusion.{fn_name} unavailable ({exc}); "
-            f"using fallback {fallback}"
-        )
-        return fallback
-
-
 class CoreMLSampler(COREML_NODE, KSampler):
     @classmethod
     def INPUT_TYPES(s):
@@ -180,7 +168,7 @@ class CoreMLLoader(COREML_NODE):
 
     @classmethod
     def coreml_filenames(cls):
-        extensions = (".mlpackage",)
+        extensions = (".mlmodelc", ".mlpackage")
         all_paths = folder_paths.get_filename_list_(cls.PACKAGE_DIRNAME)[1]
         coreml_paths = folder_paths.filter_files_extensions(all_paths, extensions)
 
@@ -191,7 +179,9 @@ class CoreMLLoader(COREML_NODE):
 
         coreml_path = self.coreml_filenames()[coreml_name]
 
-        return (CoreMLModel(coreml_path, compute_unit),)
+        sources = "compiled" if coreml_name.endswith(".mlmodelc") else "packages"
+
+        return (CoreMLModel(coreml_path, compute_unit, sources),)
 
 
 class CoreMLLoaderUNet(CoreMLLoader):
@@ -225,26 +215,28 @@ class CoreMLModelAdapter(COREML_NODE):
 
 
 class CoreMLConverter(COREML_NODE):
-    """Converts a Stable Diffusion checkpoint (UNet) to Core ML.
-
-    The model version (SD15 / SDXL / SDXL refiner / LCM) is auto-detected from
-    the checkpoint's architecture, so there is no version dropdown — one node
-    converts every supported family, including full-distill LCM.
-    """
+    """Converts a LCM model to Core ML."""
 
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
                 "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
-                "height": ("INT", {"default": 512, "min": 8, "step": 8}),
-                "width": ("INT", {"default": 512, "min": 8, "step": 8}),
+                "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}),
                 "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
                 "attention_implementation": (
-                    _discover(
-                        "list_attention_impls",
-                        ["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"],
-                    ),
+                    [
+                        AttentionImplementations.SPLIT_EINSUM.name,
+                        AttentionImplementations.SPLIT_EINSUM_V2.name,
+                        AttentionImplementations.ORIGINAL.name,
+                    ],
                 ),
                 "compute_unit": (
                     [
@@ -262,10 +254,7 @@ class CoreMLConverter(COREML_NODE):
                 # omits any `required` input. When omitted it defaults to
                 # "none", identical to unquantized behavior and filename, so
                 # existing cached .mlpackages still resolve.
-                "quantize_nbits": (
-                    _discover("list_quant_modes", ["none", "8", "6", "4"]),
-                    {"default": "none"},
-                ),
+                "quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}),
                 "lora_params": ("LORA_PARAMS",),
             },
         }
@@ -277,6 +266,7 @@ class CoreMLConverter(COREML_NODE):
     def convert(
         self,
         ckpt_name,
+        model_version,
         height,
         width,
         batch_size,
@@ -286,11 +276,9 @@ class CoreMLConverter(COREML_NODE):
         quantize_nbits="none",
         lora_params=None,
     ):
-        """Converts a checkpoint's UNet to Core ML.
+        """Converts a LCM model to Core ML.
 
         Args:
-            ckpt_name (str): Checkpoint to convert; its model version is
-                auto-detected from the weights.
             height (int): Height of the target image.
             width (int): Width of the target image.
             batch_size (int): Batch size.
@@ -300,8 +288,10 @@ class CoreMLConverter(COREML_NODE):
             coreml_model: The converted Core ML model.
 
         The converted model is also saved to "models/unet" directory and
-        can be loaded with the "Load Core ML UNet" node.
+        can be loaded with the "LCMCoreMLLoaderUNet" node.
         """
+        model_version = ModelVersion[model_version]
+
         lora_params = lora_params or {}
         lora_params = [(k, v[0]) for k, v in lora_params.items()]
         lora_params = sorted(lora_params, key=lambda lora: lora[0])
@@ -310,16 +300,14 @@ class CoreMLConverter(COREML_NODE):
         h = height
         w = width
         sample_size = (h // 8, w // 8)
-        import coreml_diffusion
-
-        out_name = coreml_diffusion.compose_out_name(
+        out_name = compose_out_name(
             ckpt_name=ckpt_name,
             batch_size=batch_size,
             width=w,
             height=h,
             controlnet_support=controlnet_support,
             attention_implementation=attention_implementation,
-            lora_names=coreml_diffusion.lora_names_from_params(lora_params),
+            lora_names=lora_names_from_params(lora_params),
             quantize_nbits=quantize_nbits,
         )
 
@@ -330,14 +318,11 @@ class CoreMLConverter(COREML_NODE):
         logger.info(f"Attention implementation: {attention_implementation}")
 
         if lora_params:
-            logger.info("LoRAs used:")
+            logger.info(f"LoRAs used:")
             for lora_param in lora_params:
                 logger.info(f"  {lora_param[0]} - strength: {lora_param[1]}")
 
-        # Resolve the ComfyUI models/unet path here (a node concern); the package
-        # takes the output path as an injected argument.
-        unet_path = folder_paths.get_folder_paths("unet")[0]
-        unet_out_path = os.path.join(unet_path, f"{out_name}_unet.mlpackage")
+        unet_out_path = converter.get_out_path("unet", f"{out_name}")
         ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
 
         config_filename = ckpt_name.split(".")[0] + ".yaml"
@@ -345,10 +330,10 @@ class CoreMLConverter(COREML_NODE):
         if config_path:
             logger.info(f"Using config file {config_path}")
 
-        coreml_diffusion.convert(
-            ckpt_path,
-            None,  # model_version auto-detected from the checkpoint
-            unet_out_path,
+        converter.convert(
+            ckpt_path=ckpt_path,
+            model_version=model_version,
+            unet_out_path=unet_out_path,
             sample_size=sample_size,
             batch_size=batch_size,
             controlnet_support=controlnet_support,
@@ -357,7 +342,11 @@ class CoreMLConverter(COREML_NODE):
             config_path=config_path,
             quantize_nbits=quantize_nbits,
         )
-        return (CoreMLModel(unet_out_path, compute_unit),)
+        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"),)
 
     @staticmethod
     def lora_path(lora_name):

docs/conversion.md

@@ -1,93 +0,0 @@
-# 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

@@ -1,77 +0,0 @@
-# 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

@@ -1,95 +0,0 @@
-# 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

@@ -1,53 +0,0 @@
-# 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

@@ -1,157 +0,0 @@
-# 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

@@ -1,86 +0,0 @@
-# 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

@@ -1,103 +0,0 @@
-# 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,42 +1,50 @@
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
 [project]
 name = "comfyui-coremlsuite"
 description = "This extension contains a set of custom nodes for ComfyUI that allow you to use Core ML models in your ComfyUI workflows."
-version = "2.1.2"
+version = "1.0.1"
 license = "MIT"
-requires-python = ">=3.12"
+requires-python = ">=3.12,<3.13"
+packages = [{ include = "coreml_suite" }]
 dependencies = [
-    # torch is provided by the host (ComfyUI) and intentionally left unpinned
-    # here: a hard torch cap would downgrade the host's torch and break its
-    # torchvision/torchaudio ABI. coreml-diffusion pulls torch>=2.7 transitively.
-    # >=0.1.6: model-version auto-detection (convert(model_version=None)) and the
-    # dropped <3.13 Python cap (kept in sync with this package's requires-python).
-    "coreml-diffusion>=0.1.6,<0.2",
+    # Python 3.12, coremltools 9, torch 2.7.
+    # numpy stays in the 1.24..1.x range — none of our modules need
+    # numpy 2, and coremltools + ml-stable-diffusion's SD UNet trace
+    # hit hard bugs under numpy 2 (`_cast` int(ndarray) strictness and
+    # `view` mixed-Var shape lists).
+    # torch 2.7 is the latest version coremltools 9's PyTorch frontend
+    # has been tested against.
+    "python-coreml-stable-diffusion @ git+https://github.com/apple/ml-stable-diffusion.git@e5d960c41a6a4ab200b8db379194127607b1c590",
+    "torch>=2.7,<2.8",
     "coremltools>=9,<10",
-    "numpy>=2,<3",
+    "numpy>=1.24,<2",
+    "overrides",
+    "diffusers>=0.22",
+    "peft>=0.6.2",
+    "omegaconf>=2.3",
 ]
 
 [project.urls]
 Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite"
-
-[tool.hatch.build.targets.wheel]
-packages = ["coreml_suite"]
+#  Used by Comfy Registry https://comfyregistry.org
 
 [tool.comfy]
 PublisherId = "aszc-dev"
 DisplayName = "ComfyUI-CoreMLSuite"
-Icon = "https://raw.githubusercontent.com/aszc-dev/ComfyUI-CoreMLSuite/main/assets/snake.png"
-requires-comfyui = ">=0.3.27"
+Icon = ""
+# Pinned to the ComfyUI commit this toolchain was validated against.
+requires-comfyui = "==ab5413351eee61f3d7f10c74e75286df0058bb18"
 
 [dependency-groups]
 dev = [
     "pillow>=12.2.0",
     "psutil>=7.2.2",
-    "pytest>=9.0.3",
 ]
+# ComfyUI runtime deps that aren't part of our package's runtime contract
+# but are needed to spin up the ComfyUI server for Tier 2 integration tests.
+# Kept in a uv group so `uv sync --group comfy` brings them in without
+# polluting the published metadata (and without re-bumping our torch pin
+# via `uv pip install -r ComfyUI/requirements.txt`, which would float to
+# the latest torch and break the coremltools 9 compatibility ceiling).
 comfy = [
     "comfyui-frontend-package==1.14.6",
     "torchvision",
@@ -53,11 +61,34 @@ comfy = [
     "sentencepiece",
 ]
 
+[tool.uv]
+# ml-stable-diffusion's setup.py hard-pins numpy<1.24, diffusers==0.30.2
+# and transformers==4.44.2, which blocks the modern torch / coremltools
+# combo on Python 3.12. Override the four blocking pins; the .unet /
+# .coreml_model symbols we actually import are stable across the bumped
+# versions.
+override-dependencies = [
+    "numpy>=1.24,<2",
+    "diffusers>=0.30",
+    "transformers>=4.44",
+    "huggingface-hub>=0.24",
+]
+
 [tool.pytest.ini_options]
+# Tier markers gate which environment a test needs.
+# - unit: framework-free pure-logic tests (Tier 0; run without ComfyUI on
+#   Linux).
+# - m2:   needs an Apple Silicon Mac with the Neural Engine (Tier 2),
+#   typically a self-hosted runner or local M-series box.
+# - smoke: lightweight checks that need Apple Silicon + coremltools but no
+#   ANE/real model (Tier 1).
 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)",
+    "smoke: macOS-ARM smoke test on a synthetic micro-model (Tier 1)",
 ]
 testpaths = ["tests"]
+# importlib mode keeps pytest from importing the repo-root __init__.py
+# (which is the ComfyUI custom-node entry and pulls in comfy + nodes).
+# Without this Tier-0 leaks the entire ComfyUI runtime on collection.
 addopts = ["--import-mode=importlib", "--confcutdir=tests"]

requirements.txt

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

tests/conftest.py

@@ -4,7 +4,7 @@
   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.
+  individual files don't have to repeat @pytest.mark.unit / .m2.
 """
 import sys
 from pathlib import Path
@@ -26,9 +26,9 @@ _TIER_BY_DIR = {
     "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.
+# When the user asks for a single tier (-m unit / -m m2), skip the other
+# directories at collection time. Tier-0 cannot afford to import tests/m2
+# files because they pull in PIL + ComfyUI runtime which Linux CI won't have.
 _TIER_DIRS = {
     "unit": ("/tests/unit/",),
     "m2": ("/tests/m2/", "/tests/integration/"),

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

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

tests/smoke/test_synthetic_unet.py

@@ -0,0 +1,122 @@
+"""Tier 1 smoke: convert a synthetic micro-UNet through coremltools and load
+it back with python_coreml_stable_diffusion's CoreMLModel.
+
+Purpose: catch API breakage in coremltools / ml-stable-diffusion *without*
+needing a real SD checkpoint, the ANE, or a converted .mlmodelc on disk.
+Runs in minutes on a hosted macOS-ARM runner (no Apple internal stuff).
+
+What it asserts:
+  - coremltools.convert still accepts the call shape we use today
+  - the resulting .mlpackage round-trips through CoreMLModel
+  - expected_inputs exposes the input names/shapes we declared
+  - calling the model returns the named output (`noise_pred`)
+
+Auto-skips on non-Apple-Silicon hosts so Tier 0 CI on Linux ignores it.
+"""
+import platform
+import shutil
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+
+pytestmark = pytest.mark.skipif(
+    platform.system() != "Darwin" or platform.machine() != "arm64",
+    reason="Tier 1 requires macOS on Apple Silicon",
+)
+
+
+# Tiny shapes — large enough to exercise conv2d + linear + addition kernels in
+# coremltools, small enough that conversion finishes in seconds on CPU.
+SAMPLE_SHAPE = (1, 4, 8, 8)
+TIMESTEP_SHAPE = (1,)
+ENCODER_SHAPE = (1, 64, 1, 4)  # matches SD's transposed encoder_hidden_states layout
+OUT_NAME = "noise_pred"
+
+
+class TinyUNet(nn.Module):
+    """Minimal UNet-shaped graph: conv -> add(time+context) -> conv.
+
+    Not a real diffusion model. Just enough op variety to exercise the
+    PyTorch -> MIL frontend in coremltools and confirm we can still wire
+    the inputs/outputs the way ml-stable-diffusion expects.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.conv_in = nn.Conv2d(4, 8, kernel_size=3, padding=1)
+        self.conv_out = nn.Conv2d(8, 4, kernel_size=3, padding=1)
+        self.time_proj = nn.Linear(1, 8)
+        self.text_proj = nn.Linear(64, 8)
+
+    def forward(self, sample, timestep, encoder_hidden_states):
+        h = self.conv_in(sample)
+        t_emb = self.time_proj(timestep.unsqueeze(-1)).view(1, 8, 1, 1)
+        c_emb = self.text_proj(encoder_hidden_states.squeeze(2).mean(-1)).view(1, 8, 1, 1)
+        h = h + t_emb + c_emb
+        return self.conv_out(h)
+
+
+@pytest.fixture(scope="module")
+def tiny_mlpackage(tmp_path_factory):
+    """Convert TinyUNet once per test session and reuse the .mlpackage."""
+    import coremltools as ct
+
+    torch.manual_seed(0)
+    model = TinyUNet().eval()
+    example = (
+        torch.randn(*SAMPLE_SHAPE),
+        torch.randn(*TIMESTEP_SHAPE),
+        torch.randn(*ENCODER_SHAPE),
+    )
+    traced = torch.jit.trace(model, example)
+
+    mlmodel = ct.convert(
+        traced,
+        inputs=[
+            ct.TensorType(name="sample", shape=SAMPLE_SHAPE, dtype=np.float16),
+            ct.TensorType(name="timestep", shape=TIMESTEP_SHAPE, dtype=np.float16),
+            ct.TensorType(name="encoder_hidden_states", shape=ENCODER_SHAPE, dtype=np.float16),
+        ],
+        outputs=[ct.TensorType(name=OUT_NAME, dtype=np.float16)],
+        compute_units=ct.ComputeUnit.CPU_ONLY,
+        compute_precision=ct.precision.FLOAT16,
+        convert_to="mlprogram",
+        minimum_deployment_target=ct.target.macOS13,
+    )
+
+    out_dir = tmp_path_factory.mktemp("tiny_unet")
+    pkg_path = out_dir / "tiny.mlpackage"
+    mlmodel.save(str(pkg_path))
+    yield pkg_path
+    shutil.rmtree(out_dir, ignore_errors=True)
+
+
+def test_coremltools_convert_round_trips_via_coreml_model(tiny_mlpackage):
+    from python_coreml_stable_diffusion.coreml_model import CoreMLModel
+
+    model = CoreMLModel(str(tiny_mlpackage), "CPU_ONLY", "packages")
+
+    # expected_inputs is the contract our wrappers depend on. Lock the shape
+    # of the dict + a sample entry.
+    expected = dict(model.expected_inputs)
+    assert set(expected.keys()) == {"sample", "timestep", "encoder_hidden_states"}
+    assert tuple(expected["sample"]["shape"]) == SAMPLE_SHAPE
+    assert tuple(expected["timestep"]["shape"]) == TIMESTEP_SHAPE
+    assert tuple(expected["encoder_hidden_states"]["shape"]) == ENCODER_SHAPE
+
+    # Forward pass: drive the model the way CoreMLModelWrapper does.
+    rng = np.random.default_rng(0)
+    inputs = {
+        "sample": rng.standard_normal(SAMPLE_SHAPE).astype(np.float16),
+        "timestep": rng.standard_normal(TIMESTEP_SHAPE).astype(np.float16),
+        "encoder_hidden_states": rng.standard_normal(ENCODER_SHAPE).astype(np.float16),
+    }
+    out = model(**inputs)
+    assert isinstance(out, dict), f"unexpected output type: {type(out)}"
+    assert OUT_NAME in out, f"missing output {OUT_NAME!r}; got {sorted(out)}"
+    assert out[OUT_NAME].shape == SAMPLE_SHAPE, (
+        f"output shape drift: got {out[OUT_NAME].shape}, expected {SAMPLE_SHAPE}"
+    )

tests/unit/test_characterization_controlnet.py

@@ -30,7 +30,7 @@ SD15_RESIDUAL_SPEC = {
 }
 NON_RESIDUAL_SPEC = {
     "sample": {"shape": (2, 4, 64, 64)},
-    "encoder_hidden_states": {"shape": (2, 77, 768)},
+    "encoder_hidden_states": {"shape": (2, 768, 1, 77)},
 }
 
 

tests/unit/test_characterization_inputs.py

@@ -25,7 +25,7 @@ def _deterministic_seed():
 SD15_EXPECTED = {
     "sample": {"shape": (2, 4, 64, 64)},
     "timestep": {"shape": (2,)},
-    "encoder_hidden_states": {"shape": (2, 77, 768)},
+    "encoder_hidden_states": {"shape": (2, 768, 1, 77)},
 }
 
 SD15_WITH_CN = {
@@ -42,7 +42,7 @@ LCM_EXPECTED = {
 SDXL_BASE_EXPECTED = {
     "sample": {"shape": (2, 4, 128, 128)},
     "timestep": {"shape": (2,)},
-    "encoder_hidden_states": {"shape": (2, 77, 2048)},
+    "encoder_hidden_states": {"shape": (2, 2048, 1, 77)},
     "time_ids": {"shape": (2, 6)},
     "text_embeds": {"shape": (2, 1280)},
 }
@@ -50,7 +50,7 @@ SDXL_BASE_EXPECTED = {
 SDXL_REFINER_EXPECTED = {
     "sample": {"shape": (2, 4, 128, 128)},
     "timestep": {"shape": (2,)},
-    "encoder_hidden_states": {"shape": (2, 77, 1280)},
+    "encoder_hidden_states": {"shape": (2, 1280, 1, 77)},
     "time_ids": {"shape": (2, 5)},
     "text_embeds": {"shape": (2, 1280)},
 }
@@ -93,8 +93,8 @@ def test_coreml_kwargs_sd15_shapes_and_fp16():
     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)
+    # encoder_hidden_states is transposed (b, seq, dim) -> (b, dim, 1, seq).
+    assert out["encoder_hidden_states"].shape == (1, 768, 1, 77)
     assert out["encoder_hidden_states"].dtype == np.float16
     assert out["timestep"].shape == (1,)
     assert out["timestep"].dtype == np.float16

tests/unit/test_characterization_out_name.py

@@ -0,0 +1,197 @@
+"""Characterization tests for the .mlpackage filename composition.
+
+The filename composition is the pure
+coreml_suite.core.naming.compose_out_name function. CoreMLConverter.convert
+calls it; testing the pure function avoids monkey-patching heavy converter
+internals just to capture the string.
+"""
+import pytest
+
+from coreml_suite.core.naming import compose_out_name, lora_names_from_params
+
+
+# ---------- attention suffixes ----------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "attn_name,suffix",
+    [
+        ("SPLIT_EINSUM", "se"),
+        ("SPLIT_EINSUM_V2", "se2"),
+        ("ORIGINAL", "orig"),
+    ],
+)
+def test_attention_suffix(attn_name, suffix):
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation=attn_name,
+    )
+    assert out == f"dreamshaper_8_1x512x512_{suffix}"
+
+
+# ---------- batch / size ----------------------------------------------------
+
+
+def test_includes_batch_and_size():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=4, width=768, height=1024,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+    )
+    assert out == "dreamshaper_8_4x768x1024_se"
+
+
+# ---------- ControlNet ------------------------------------------------------
+
+
+def test_appends_cn_suffix_when_controlnet_support_true():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=True,
+        attention_implementation="SPLIT_EINSUM",
+    )
+    assert out == "dreamshaper_8_1x512x512_cn_se"
+
+
+# ---------- ckpt name massage -----------------------------------------------
+
+
+def test_drops_extension_at_first_period():
+    out = compose_out_name(
+        ckpt_name="my.checkpoint.v2.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+    )
+    assert out == "my_1x512x512_se"
+
+
+def test_replaces_spaces_with_underscores():
+    out = compose_out_name(
+        ckpt_name="dream shaper 8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+    )
+    assert out == "dream_shaper_8_1x512x512_se"
+
+
+# ---------- LoRA suffixes ---------------------------------------------------
+
+
+def test_single_lora():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+        lora_names=["epi_noiseoffset.safetensors"],
+    )
+    assert out == "dreamshaper_8_epi_noiseoffset_1x512x512_se"
+
+
+def test_multiple_loras_sorted():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+        lora_names=["zoom.safetensors", "alpha.safetensors", "moody.safetensors"],
+    )
+    assert out == "dreamshaper_8_alpha_moody_zoom_1x512x512_se"
+
+
+def test_lora_plus_controlnet():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=True,
+        attention_implementation="SPLIT_EINSUM",
+        lora_names=["a.safetensors"],
+    )
+    assert out == "dreamshaper_8_a_1x512x512_cn_se"
+
+
+# ---------- sdxl combinations -----------------------------------------------
+
+
+def test_sdxl_1024_original_gpu():
+    out = compose_out_name(
+        ckpt_name="sd_xl_base_1.0.safetensors",
+        batch_size=1, width=1024, height=1024,
+        controlnet_support=False,
+        attention_implementation="ORIGINAL",
+    )
+    assert out == "sd_xl_base_1_1x1024x1024_orig"
+
+
+# ---------- lora_names_from_params helper ----------------------------------
+
+
+def test_lora_names_from_params_sorts_by_name():
+    names = lora_names_from_params([
+        ("zebra.safetensors", 1.0),
+        ("apple.safetensors", 0.5),
+        ("mango.safetensors", 0.7),
+    ])
+    assert names == ["apple.safetensors", "mango.safetensors", "zebra.safetensors"]
+
+
+def test_lora_names_from_params_empty_list():
+    assert lora_names_from_params([]) == []
+
+
+# ---------- quantize_nbits suffix ------------------------------------------
+
+
+def test_quantize_nbits_none_appends_nothing():
+    """'none' is the default and must keep the unquantized filename so
+    existing cached .mlpackages still resolve."""
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+        quantize_nbits="none",
+    )
+    assert out == "dreamshaper_8_1x512x512_se"
+
+
+@pytest.mark.parametrize("nbits,suffix", [("4", "_q4"), ("6", "_q6"), ("8", "_q8")])
+def test_quantize_nbits_appends_q_suffix(nbits, suffix):
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=False,
+        attention_implementation="SPLIT_EINSUM",
+        quantize_nbits=nbits,
+    )
+    assert out == f"dreamshaper_8_1x512x512_se{suffix}"
+
+
+def test_quantize_nbits_with_controlnet_and_lora():
+    out = compose_out_name(
+        ckpt_name="dreamshaper_8.safetensors",
+        batch_size=1, width=512, height=512,
+        controlnet_support=True,
+        attention_implementation="SPLIT_EINSUM",
+        lora_names=["a.safetensors"],
+        quantize_nbits="6",
+    )
+    assert out == "dreamshaper_8_a_1x512x512_cn_se_q6"
+
+
+def test_quantize_nbits_invalid_raises():
+    import pytest as _pytest
+    with _pytest.raises(ValueError, match="quantize_nbits"):
+        compose_out_name(
+            ckpt_name="x.safetensors",
+            batch_size=1, width=512, height=512,
+            controlnet_support=False,
+            attention_implementation="SPLIT_EINSUM",
+            quantize_nbits="16",  # not in {none, 8, 6, 4}
+        )

tests/unit/test_characterization_sdxl_options.py

@@ -2,8 +2,8 @@
 
 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.
+add_sdxl_model_options (in models.py) is exercised separately by the m2
+golden image test; here we just lock the pure math.
 """
 import inspect
 

tests/unit/test_chunks.py

@@ -20,7 +20,7 @@ def expected_inputs():
         "sample": {"shape": (2, 4, 64, 64)},
         "timestep": {"shape": (2,)},
         "timestep_cond": {"shape": (2, 256)},
-        "encoder_hidden_states": {"shape": (2, 77, 768)},
+        "encoder_hidden_states": {"shape": (2, 768, 1, 77)},
         "additional_residual_0": {"shape": (2, 320, 64, 64)},
         "additional_residual_1": {"shape": (2, 640, 32, 32)},
     }

tests/unit/test_tier0_purity.py

@@ -7,10 +7,10 @@ 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.
+When other tiers (m2 / integration) are also collected, comfy is
+expected in sys.modules (integration imports it deliberately), so the
+check is skipped in mixed runs — Tier-0 purity is only meaningful when
+nothing else is loaded.
 """
 import sys