Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1008f144ad | ||
|
|
8f94f0eea5 | ||
|
|
5e57024336 | ||
|
|
8d28964831 | ||
|
|
a199e749bc | ||
|
|
a3d555cdc5 | ||
|
|
d90546b6bb | ||
|
|
f054bbe991 |
@@ -1,23 +0,0 @@
|
|||||||
name: Tier 1 — Smoke (macOS self-hosted)
|
|
||||||
|
|
||||||
# macOS smoke tests run on the self-hosted Apple Silicon runner instead of
|
|
||||||
# GitHub-hosted macOS (10x minute multiplier), which exhausts the included
|
|
||||||
# Actions minutes too quickly.
|
|
||||||
on:
|
|
||||||
push:
|
|
||||||
branches: [main]
|
|
||||||
pull_request:
|
|
||||||
|
|
||||||
jobs:
|
|
||||||
smoke:
|
|
||||||
runs-on: [self-hosted, macOS, ARM64, coreml]
|
|
||||||
timeout-minutes: 20
|
|
||||||
steps:
|
|
||||||
- uses: actions/checkout@v4
|
|
||||||
|
|
||||||
# The self-hosted runner provides uv; no setup-uv action needed.
|
|
||||||
- name: uv sync
|
|
||||||
run: uv sync --no-install-project
|
|
||||||
|
|
||||||
- name: Run Tier 1 (synthetic micro-UNet smoke)
|
|
||||||
run: uv run pytest -m smoke tests/ -v
|
|
||||||
@@ -3,3 +3,6 @@ __pycache__/
|
|||||||
models/
|
models/
|
||||||
.venv/
|
.venv/
|
||||||
test_results/
|
test_results/
|
||||||
|
*.log
|
||||||
|
.DS_Store
|
||||||
|
.claude/
|
||||||
|
|||||||
@@ -0,0 +1 @@
|
|||||||
|
3.12
|
||||||
@@ -1,674 +1,21 @@
|
|||||||
GNU GENERAL PUBLIC LICENSE
|
MIT 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
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
of this license document, but changing it is not allowed.
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
Preamble
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
The GNU General Public License is a free, copyleft license for
|
furnished to do so, subject to the following conditions:
|
||||||
software and other kinds of works.
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
The licenses for most software and other practical works are designed
|
copies or substantial portions of the Software.
|
||||||
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
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
share and change all versions of a program--to make sure it remains free
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
software for all its users. We, the Free Software Foundation, use the
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
GNU General Public License for most of our software; it applies also to
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
any other work released this way by its authors. You can apply it to
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
your programs, too.
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
||||||
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>.
|
|
||||||
|
|||||||
@@ -1,452 +1,113 @@
|
|||||||
# Core ML Suite for ComfyUI
|
# Core ML Suite for ComfyUI
|
||||||
|
|
||||||
## Overview
|
Custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI) that run
|
||||||
|
Stable Diffusion UNets as [Core ML](https://developer.apple.com/documentation/coreml)
|
||||||
|
models on Apple Silicon (M1/M2/M3). Core ML can use the Apple Neural Engine
|
||||||
|
(ANE), which is unavailable to PyTorch — on an M2 Pro 32 GB, SD1.5 at 512×512
|
||||||
|
generates roughly **1.5–2× faster** than the standard PyTorch/MPS path.
|
||||||
|
|
||||||
Welcome! In this repository you'll find a set of custom nodes for [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
You convert a Stable Diffusion checkpoint to a Core ML model with the nodes in
|
||||||
that allows you to use Core ML models in your ComfyUI workflows.
|
this suite, then sample from it like any other ComfyUI workflow.
|
||||||
These models are designed to leverage the Apple Neural Engine (ANE) on Apple Silicon (M1/M2) machines,
|
|
||||||
thereby enhancing your workflows and improving performance.
|
|
||||||
|
|
||||||
If you're not sure how to obtain these models, you can download them
|
> [!IMPORTANT]
|
||||||
[here](https://huggingface.co/coreml-community) or convert your own checkpoints
|
> **Convert your own checkpoints — that is the only supported path.** This
|
||||||
directly with the conversion nodes in this suite (see [How to use](#how-to-use)).
|
> 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.
|
||||||
|
|
||||||
In simple terms, think of Core ML models as a tool that can help your ComfyUI work faster and more efficiently.
|
## Installation
|
||||||
For instance, during my tests on an M2 Pro 32GB machine,
|
|
||||||
the use of Core ML models sped up the generation of 512x512 images by a factor
|
|
||||||
of approximately 1.5 to 2 times.
|
|
||||||
|
|
||||||
## Getting Started
|
### ComfyUI-Manager (recommended)
|
||||||
|
|
||||||
To start using custom nodes in your ComfyUI, follow these simple steps:
|
Open **Manager → Install Custom Nodes**, search for `Core ML`, click
|
||||||
|
**Install**, and restart ComfyUI.
|
||||||
|
|
||||||
1. Clone or download this repository: You can do this directly into the custom_nodes directory of your ComfyUI.
|
### Manual
|
||||||
2. Install the dependencies: You'll need to use a package manager like pip to do this.
|
|
||||||
|
|
||||||
That's it! You're now ready to start enhancing your ComfyUI workflows with Core ML models.
|
```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
|
||||||
|
```
|
||||||
|
|
||||||
- Check [Installation](#installation) for more details on installation.
|
Dependencies (`coreml-diffusion`, `coremltools`, `numpy`, `diffusers`) install
|
||||||
- Check [How to use](#how-to-use) for more details on how to use the custom nodes.
|
from PyPI. PyTorch is intentionally **not** pinned — it is provided by your
|
||||||
- Check [Example Workflows](#example-workflows) for some example workflows.
|
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.
|
||||||
|
|
||||||
## Glossary
|
## Glossary
|
||||||
|
|
||||||
- **Core ML**: A machine learning framework developed by Apple. It's used to run machine learning models on Apple
|
- **Core ML** — Apple's on-device machine-learning framework.
|
||||||
devices.
|
- **`.mlpackage`** — the Core ML model format this suite produces and loads.
|
||||||
- **Core ML Model**: A machine learning model that can be run on Apple devices using Core ML.
|
- **ANE** — Apple Neural Engine, a hardware accelerator for ML.
|
||||||
- **mlmodelc**: A compiled Core ML model. This is the recommended format for Core ML models.
|
- **Compute unit** — which hardware Core ML uses (`CPU_AND_NE`, `CPU_AND_GPU`,
|
||||||
- **mlpackage**: A Core ML model packaged in a directory. This is the default format for Core ML models.
|
`CPU_ONLY`, `ALL`).
|
||||||
- **ANE**: Apple Neural Engine. A hardware accelerator for machine learning tasks on Apple devices.
|
- **Attention implementation** — `SPLIT_EINSUM` / `SPLIT_EINSUM_V2` (ANE-friendly)
|
||||||
- **Compute Unit**: A Core ML option that allows you to specify the hardware on which the model should run.
|
or `ORIGINAL` (GPU-friendly), chosen at conversion.
|
||||||
- **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.
|
|
||||||
|
|
||||||
## Conversion & Acknowledgements
|
|
||||||
|
|
||||||
The Core ML conversion pipeline in this repository began as an adaptation of
|
|
||||||
Apple's [ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion),
|
|
||||||
which pioneered running Stable Diffusion on the Apple Neural Engine. The
|
|
||||||
implementation has since diverged and no longer depends on that package:
|
|
||||||
|
|
||||||
- UNet conversion runs natively on `diffusers`' `UNet2DConditionModel`.
|
|
||||||
- The ANE-friendly attention path (`SPLIT_EINSUM`, `SPLIT_EINSUM_V2`) is
|
|
||||||
reimplemented as standalone `diffusers` attention processors.
|
|
||||||
- The toolchain tracks current ComfyUI (NumPy 2, Torch 2.7, coremltools 9,
|
|
||||||
Python 3.12).
|
|
||||||
|
|
||||||
The goal is to keep iterating on these methods independently and to explore
|
|
||||||
support beyond SD1.5.
|
|
||||||
|
|
||||||
> [!IMPORTANT]
|
> [!IMPORTANT]
|
||||||
> **Breaking change in 2.0.0.** The converted Core ML UNet now takes
|
> **Breaking change in 2.0.0.** The converted Core ML UNet now takes
|
||||||
> `encoder_hidden_states` in the native `diffusers` layout
|
> `encoder_hidden_states` in the native `diffusers` layout
|
||||||
> `(batch, tokens, hidden)` instead of the previous
|
> `(batch, tokens, hidden)` instead of the previous `(batch, hidden, 1, tokens)`.
|
||||||
> `(batch, hidden, 1, tokens)`. Core ML models converted with earlier versions
|
> Models converted with earlier versions are not compatible and must be
|
||||||
> are not compatible with 2.0.0 and must be re-converted.
|
> re-converted.
|
||||||
|
|
||||||
## Installation
|
## Acknowledgements
|
||||||
|
|
||||||
### Using ComfyUI-Manager
|
The conversion pipeline began as an adaptation of Apple's
|
||||||
|
[ml-stable-diffusion](https://github.com/apple/ml-stable-diffusion), which
|
||||||
The easiest way to install the custom nodes is to use the ComfyUI-Manager. You can find the installation instructions
|
pioneered running Stable Diffusion on the Neural Engine. It has since diverged
|
||||||
[here](https://github.com/ltdrdata/ComfyUI-Manager#installation). Once you've installed the ComfyUI-Manager, you can
|
and no longer depends on that package: UNet conversion runs natively on
|
||||||
install the custom nodes by following these steps:
|
`diffusers`' `UNet2DConditionModel`, the ANE attention path (`SPLIT_EINSUM`,
|
||||||
|
`SPLIT_EINSUM_V2`) is reimplemented as standalone `diffusers` attention
|
||||||
- Open the ComfyUI-Manager by clicking the `Manager` button in the ComfyUI toolbar.
|
processors, and the toolchain tracks current ComfyUI (NumPy 2, Torch 2.7+,
|
||||||
- Click the `Install Custom Nodes` button.
|
coremltools 9, Python 3.12+). Conversion now lives in the separate
|
||||||
- Search for `Core ML` and click the `Install` button.
|
[coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) package.
|
||||||
- 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`)
|
|
||||||
|
|
||||||

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

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

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

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

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

|
|
||||||
|
|
||||||
This node allows you to use a Core ML as a standard ComfyUI model. This is an experimental node and may not work with
|
|
||||||
all models and nodes. Please use with caution and pay attention to the expected inputs of the model.
|
|
||||||
|
|
||||||
- **Input**:
|
|
||||||
- **coreml_model**: The Core ML model to use as a ComfyUI model.
|
|
||||||
- **Output**:
|
|
||||||
- **MODEL**: The Core ML model wrapped in a ComfyUI model.
|
|
||||||
|
|
||||||
> [!NOTE]
|
|
||||||
> While this approach allows you to use Core ML models with many ComfyUI nodes (both standard and custom), the
|
|
||||||
> expected inputs of the model will not be checked, which may cause errors. Please make sure to use a model compatible
|
|
||||||
> with the expected parameters.
|
|
||||||
|
|
||||||
### Example Workflows
|
|
||||||
|
|
||||||
> [!NOTE]
|
|
||||||
> 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/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors) and
|
|
||||||
[here](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/vae/diffusion_pytorch_model.safetensors).
|
|
||||||
Once downloaded, place the models in the`models/clip` and `models/vae` directories respectively.
|
|
||||||
- The Core ML UNet model is available
|
|
||||||
[here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip).
|
|
||||||
Once downloaded, place the model in the `models/unet` directory.
|
|
||||||

|
|
||||||
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/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors).
|
|
||||||
Once downloaded, place the model in the`models/checkpoints` directory.
|
|
||||||
- The Core ML UNet model is available
|
|
||||||
[here](https://huggingface.co/coreml-community/coreml-stable-diffusion-v1-5_cn/blob/main/split_einsum/stable-diffusion-_v1-5_split-einsum_cn.zip).
|
|
||||||
Once downloaded, place the model in the `models/unet` directory.
|
|
||||||

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

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

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

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

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

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

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

|
|
||||||
|
|
||||||
## Quantization (opt-in)
|
|
||||||
|
|
||||||
The `Core ML Converter` and `Core ML LCM Converter` nodes accept an
|
|
||||||
optional `quantize_nbits` dropdown that runs k-means weight palettization
|
|
||||||
(`coremltools.optimize.coreml.palettize_weights`) on the UNet before save.
|
|
||||||
|
|
||||||
Values: `none` (default — no quantization, identical to unquantized
|
|
||||||
behavior and filenames), `8`, `6`, `4`. The number is appended to the
|
|
||||||
.mlpackage stem as `_q<bits>` so quantized and unquantized variants
|
|
||||||
coexist on disk and in cache.
|
|
||||||
|
|
||||||
### SD1.5 1×512×512 SPLIT_EINSUM tradeoffs (M2 Pro, ANE)
|
|
||||||
|
|
||||||
Measured with 20 UNet forward passes at a fixed seed for the PSNR
|
|
||||||
comparison:
|
|
||||||
|
|
||||||
| nbits | size (MB) | size vs none | fwd median (ms) | PSNR vs `none` (dB) |
|
|
||||||
|---|---:|---:|---:|---:|
|
|
||||||
| none | 1641 | 1.000 | 197.1 | — |
|
|
||||||
| 8 | 822 | 0.501 | 186.6 | 53.5 |
|
|
||||||
| 6 | 617 | 0.376 | 183.0 | 40.2 |
|
|
||||||
| 4 | 412 | 0.251 | 179.8 | 27.5 |
|
|
||||||
|
|
||||||
PSNR here is computed on the raw `noise_pred` output of a single UNet
|
|
||||||
forward at a fixed seed, not on the final decoded image — it isolates
|
|
||||||
the quantization-induced drift from sampler / VAE noise. Final-image
|
|
||||||
PSNR is comfortably higher (the sampler averages over 20 steps).
|
|
||||||
|
|
||||||
### Recommended settings per chip / RAM
|
|
||||||
|
|
||||||
- **8 GB RAM (M1 base, M2 base):** `nbits=4`. ~4× smaller model, still
|
|
||||||
loads, PSNR 27 dB is visually identical at SD1.5 sizes.
|
|
||||||
- **16 GB RAM (M1/M2/M3 Pro):** `nbits=6` is the sweet spot — ~2.7×
|
|
||||||
smaller, PSNR 40 dB, no perceptible quality drop.
|
|
||||||
- **32 GB+ RAM (Max / Ultra):** `nbits=8` if you want the safety
|
|
||||||
margin, `none` if you want bit-identical output for golden testing.
|
|
||||||
|
|
||||||
The default stays `none` so existing workflows produce byte-for-byte
|
|
||||||
identical output.
|
|
||||||
|
|
||||||
## Limitations
|
|
||||||
|
|
||||||
- Core ML models are fixed in terms of their inputs and outputs.
|
|
||||||
This means you'll need to use latent images of the same size as the input of the model (512x512 is the default for
|
|
||||||
SD1.5).
|
|
||||||
However, you can re-convert the model to a different input size using the
|
|
||||||
conversion nodes in this suite (set the desired width and height).
|
|
||||||
- SD2.1 models are not supported.
|
|
||||||
|
|
||||||
[^1]:
|
|
||||||
Unless [EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes)
|
|
||||||
is used during conversion. Needs more testing.
|
|
||||||
|
|
||||||
## Support
|
## Support
|
||||||
|
|
||||||
I'm here to help! If you have any questions or suggestions, don't hesitate to open an issue and I'll do my best
|
Questions or suggestions? Open an
|
||||||
to assist you.
|
[issue](https://github.com/aszc-dev/ComfyUI-CoreMLSuite/issues).
|
||||||
|
|||||||
@@ -11,9 +11,6 @@ from coreml_suite.nodes import (
|
|||||||
CoreMLConverter,
|
CoreMLConverter,
|
||||||
COREML_LOAD_LORA,
|
COREML_LOAD_LORA,
|
||||||
)
|
)
|
||||||
from coreml_suite.lcm import (
|
|
||||||
COREML_CONVERT_LCM,
|
|
||||||
)
|
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
NODE_CLASS_MAPPINGS = {
|
||||||
"CoreMLUNetLoader": CoreMLLoaderUNet,
|
"CoreMLUNetLoader": CoreMLLoaderUNet,
|
||||||
@@ -22,7 +19,6 @@ NODE_CLASS_MAPPINGS = {
|
|||||||
"CoreMLModelAdapter": CoreMLModelAdapter,
|
"CoreMLModelAdapter": CoreMLModelAdapter,
|
||||||
"Core ML LoRA Loader": COREML_LOAD_LORA,
|
"Core ML LoRA Loader": COREML_LOAD_LORA,
|
||||||
"Core ML Converter": CoreMLConverter,
|
"Core ML Converter": CoreMLConverter,
|
||||||
"Core ML LCM Converter": COREML_CONVERT_LCM,
|
|
||||||
}
|
}
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
"CoreMLUNetLoader": "Load Core ML UNet",
|
"CoreMLUNetLoader": "Load Core ML UNet",
|
||||||
@@ -31,5 +27,4 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
|||||||
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
|
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
|
||||||
"Core ML LoRA Loader": "Load LoRA to use with Core ML",
|
"Core ML LoRA Loader": "Load LoRA to use with Core ML",
|
||||||
"Core ML Converter": "Convert Checkpoint to Core ML",
|
"Core ML Converter": "Convert Checkpoint to Core ML",
|
||||||
"Core ML LCM Converter": "Convert LCM to Core ML",
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,5 +0,0 @@
|
|||||||
ATTENTION_IMPLEMENTATIONS = (
|
|
||||||
"SPLIT_EINSUM",
|
|
||||||
"SPLIT_EINSUM_V2",
|
|
||||||
"ORIGINAL",
|
|
||||||
)
|
|
||||||
@@ -4,7 +4,7 @@ from comfy import supported_models_base
|
|||||||
from comfy import latent_formats
|
from comfy import latent_formats
|
||||||
from comfy.model_detection import convert_config
|
from comfy.model_detection import convert_config
|
||||||
|
|
||||||
from coreml_suite.model_version import ModelVersion
|
from coreml_diffusion import ModelVersion
|
||||||
|
|
||||||
|
|
||||||
config_map = {
|
config_map = {
|
||||||
|
|||||||
@@ -1,9 +0,0 @@
|
|||||||
"""Core ML conversion helpers.
|
|
||||||
|
|
||||||
The conversion approach originates from Apple's ml-stable-diffusion
|
|
||||||
(https://github.com/apple/ml-stable-diffusion). This implementation has since
|
|
||||||
diverged: it runs natively on diffusers' UNet2DConditionModel with its own
|
|
||||||
SPLIT_EINSUM / SPLIT_EINSUM_V2 attention processors and no longer depends on
|
|
||||||
that package. The intent is to keep iterating on these methods independently
|
|
||||||
while tracking current tooling.
|
|
||||||
"""
|
|
||||||
@@ -1,239 +0,0 @@
|
|||||||
import logging
|
|
||||||
|
|
||||||
import torch
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
CHUNK_SIZE = 512
|
|
||||||
|
|
||||||
|
|
||||||
def apply_attention_implementation(unet, attention_implementation):
|
|
||||||
if attention_implementation == "ORIGINAL":
|
|
||||||
return unet
|
|
||||||
|
|
||||||
if attention_implementation == "SPLIT_EINSUM":
|
|
||||||
unet.set_attn_processor(SplitEinsumAttnProcessor())
|
|
||||||
return unet
|
|
||||||
|
|
||||||
if attention_implementation == "SPLIT_EINSUM_V2":
|
|
||||||
unet.set_attn_processor(SplitEinsumV2AttnProcessor())
|
|
||||||
return unet
|
|
||||||
|
|
||||||
raise ValueError(f"Unsupported attention implementation: {attention_implementation}")
|
|
||||||
|
|
||||||
|
|
||||||
class SplitEinsumAttnProcessor:
|
|
||||||
def __call__(
|
|
||||||
self,
|
|
||||||
attn,
|
|
||||||
hidden_states,
|
|
||||||
encoder_hidden_states=None,
|
|
||||||
attention_mask=None,
|
|
||||||
temb=None,
|
|
||||||
*args,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
return _attention_forward(
|
|
||||||
attn,
|
|
||||||
hidden_states,
|
|
||||||
encoder_hidden_states,
|
|
||||||
attention_mask,
|
|
||||||
temb,
|
|
||||||
split_einsum,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
class SplitEinsumV2AttnProcessor:
|
|
||||||
def __call__(
|
|
||||||
self,
|
|
||||||
attn,
|
|
||||||
hidden_states,
|
|
||||||
encoder_hidden_states=None,
|
|
||||||
attention_mask=None,
|
|
||||||
temb=None,
|
|
||||||
*args,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
return _attention_forward(
|
|
||||||
attn,
|
|
||||||
hidden_states,
|
|
||||||
encoder_hidden_states,
|
|
||||||
attention_mask,
|
|
||||||
temb,
|
|
||||||
split_einsum_v2,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _attention_forward(
|
|
||||||
attn,
|
|
||||||
hidden_states,
|
|
||||||
encoder_hidden_states,
|
|
||||||
attention_mask,
|
|
||||||
temb,
|
|
||||||
attention_fn,
|
|
||||||
):
|
|
||||||
residual = hidden_states
|
|
||||||
|
|
||||||
if attn.spatial_norm is not None:
|
|
||||||
hidden_states = attn.spatial_norm(hidden_states, temb)
|
|
||||||
|
|
||||||
input_ndim = hidden_states.ndim
|
|
||||||
if input_ndim == 4:
|
|
||||||
batch_size, channel, height, width = hidden_states.shape
|
|
||||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
|
||||||
else:
|
|
||||||
batch_size, _, channel = hidden_states.shape
|
|
||||||
height = None
|
|
||||||
width = None
|
|
||||||
|
|
||||||
batch_size, key_sequence_length, _ = (
|
|
||||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
|
||||||
)
|
|
||||||
|
|
||||||
if attention_mask is not None:
|
|
||||||
attention_mask = attn.prepare_attention_mask(
|
|
||||||
attention_mask,
|
|
||||||
key_sequence_length,
|
|
||||||
batch_size,
|
|
||||||
)
|
|
||||||
attention_mask = _prepare_split_einsum_mask(
|
|
||||||
attention_mask,
|
|
||||||
batch_size,
|
|
||||||
attn.heads,
|
|
||||||
key_sequence_length,
|
|
||||||
)
|
|
||||||
|
|
||||||
if attn.group_norm is not None:
|
|
||||||
hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
|
|
||||||
|
|
||||||
query = attn.to_q(hidden_states)
|
|
||||||
|
|
||||||
if encoder_hidden_states is None:
|
|
||||||
encoder_hidden_states = hidden_states
|
|
||||||
elif attn.norm_cross:
|
|
||||||
encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
|
|
||||||
|
|
||||||
key = attn.to_k(encoder_hidden_states)
|
|
||||||
value = attn.to_v(encoder_hidden_states)
|
|
||||||
|
|
||||||
batch_size = query.shape[0]
|
|
||||||
dim_head = attn.inner_kv_dim // attn.heads
|
|
||||||
|
|
||||||
query = _linear_projection_to_bchw(query)
|
|
||||||
key = _linear_projection_to_bchw(key)
|
|
||||||
value = _linear_projection_to_bchw(value)
|
|
||||||
|
|
||||||
hidden_states = attention_fn(
|
|
||||||
query,
|
|
||||||
key,
|
|
||||||
value,
|
|
||||||
attention_mask,
|
|
||||||
attn.heads,
|
|
||||||
dim_head,
|
|
||||||
)
|
|
||||||
hidden_states = hidden_states.squeeze(2).transpose(1, 2)
|
|
||||||
hidden_states = hidden_states.reshape(batch_size, -1, attn.inner_dim)
|
|
||||||
|
|
||||||
hidden_states = attn.to_out[0](hidden_states)
|
|
||||||
hidden_states = attn.to_out[1](hidden_states)
|
|
||||||
|
|
||||||
if input_ndim == 4:
|
|
||||||
hidden_states = hidden_states.transpose(-1, -2).reshape(
|
|
||||||
batch_size,
|
|
||||||
channel,
|
|
||||||
height,
|
|
||||||
width,
|
|
||||||
)
|
|
||||||
|
|
||||||
if attn.residual_connection:
|
|
||||||
hidden_states = hidden_states + residual
|
|
||||||
|
|
||||||
hidden_states = hidden_states / attn.rescale_output_factor
|
|
||||||
return hidden_states
|
|
||||||
|
|
||||||
|
|
||||||
def split_einsum(q, k, v, mask, heads, dim_head):
|
|
||||||
q_heads = _split_heads(q, heads, dim_head)
|
|
||||||
k = k.transpose(1, 3)
|
|
||||||
k_heads = [
|
|
||||||
k[:, :, :, head_idx * dim_head : (head_idx + 1) * dim_head]
|
|
||||||
for head_idx in range(heads)
|
|
||||||
]
|
|
||||||
v_heads = _split_heads(v, heads, dim_head)
|
|
||||||
|
|
||||||
weights = [
|
|
||||||
torch.einsum("bchq,bkhc->bkhq", query, key) * (dim_head**-0.5)
|
|
||||||
for query, key in zip(q_heads, k_heads)
|
|
||||||
]
|
|
||||||
if mask is not None:
|
|
||||||
weights = [weight + mask for weight in weights]
|
|
||||||
|
|
||||||
weights = [weight.softmax(dim=1) for weight in weights]
|
|
||||||
outputs = [
|
|
||||||
torch.einsum("bkhq,bchk->bchq", weight, value)
|
|
||||||
for weight, value in zip(weights, v_heads)
|
|
||||||
]
|
|
||||||
return torch.cat(outputs, dim=1)
|
|
||||||
|
|
||||||
|
|
||||||
def split_einsum_v2(q, k, v, mask, heads, dim_head):
|
|
||||||
query_length = q.size(3)
|
|
||||||
num_chunks = query_length // CHUNK_SIZE
|
|
||||||
if num_chunks == 0:
|
|
||||||
logger.info(
|
|
||||||
"SPLIT_EINSUM_V2 query sequence is shorter than %s; using SPLIT_EINSUM.",
|
|
||||||
CHUNK_SIZE,
|
|
||||||
)
|
|
||||||
return split_einsum(q, k, v, mask, heads, dim_head)
|
|
||||||
|
|
||||||
q_heads = _split_heads(q, heads, dim_head)
|
|
||||||
q_chunks = [
|
|
||||||
[
|
|
||||||
head[..., chunk_idx * CHUNK_SIZE : (chunk_idx + 1) * CHUNK_SIZE]
|
|
||||||
for chunk_idx in range(num_chunks)
|
|
||||||
]
|
|
||||||
for head in q_heads
|
|
||||||
]
|
|
||||||
|
|
||||||
k = k.transpose(1, 3)
|
|
||||||
k_heads = [
|
|
||||||
k[:, :, :, head_idx * dim_head : (head_idx + 1) * dim_head]
|
|
||||||
for head_idx in range(heads)
|
|
||||||
]
|
|
||||||
v_heads = _split_heads(v, heads, dim_head)
|
|
||||||
|
|
||||||
head_outputs = []
|
|
||||||
for query_chunks, key, value in zip(q_chunks, k_heads, v_heads):
|
|
||||||
chunk_outputs = []
|
|
||||||
for query_chunk in query_chunks:
|
|
||||||
weights = torch.einsum("bchq,bkhc->bkhq", query_chunk, key)
|
|
||||||
weights = weights * (dim_head**-0.5)
|
|
||||||
if mask is not None:
|
|
||||||
weights = weights + mask
|
|
||||||
weights = weights.softmax(dim=1)
|
|
||||||
chunk_outputs.append(torch.einsum("bkhq,bchk->bchq", weights, value))
|
|
||||||
head_outputs.append(torch.cat(chunk_outputs, dim=3))
|
|
||||||
|
|
||||||
return torch.cat(head_outputs, dim=1)
|
|
||||||
|
|
||||||
|
|
||||||
def _split_heads(x, heads, dim_head):
|
|
||||||
return [
|
|
||||||
x[:, head_idx * dim_head : (head_idx + 1) * dim_head, :, :]
|
|
||||||
for head_idx in range(heads)
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def _linear_projection_to_bchw(x):
|
|
||||||
return x.transpose(1, 2).unsqueeze(2)
|
|
||||||
|
|
||||||
|
|
||||||
def _prepare_split_einsum_mask(mask, batch_size, heads, key_sequence_length):
|
|
||||||
if mask.ndim == 2:
|
|
||||||
mask = mask[:, None, :]
|
|
||||||
if mask.shape[0] == batch_size * heads:
|
|
||||||
mask = mask.reshape(batch_size, heads, -1, key_sequence_length)
|
|
||||||
mask = mask[:, 0]
|
|
||||||
if mask.ndim == 3:
|
|
||||||
mask = mask[:, :, None, None]
|
|
||||||
return mask
|
|
||||||
@@ -1,20 +0,0 @@
|
|||||||
def conv2d_output_shape(height, width, conv):
|
|
||||||
"""Return the spatial output shape for a torch.nn.Conv2d-like module."""
|
|
||||||
kernel_h, kernel_w = _pair(conv.kernel_size)
|
|
||||||
stride_h, stride_w = _pair(conv.stride)
|
|
||||||
pad_h, pad_w = _pair(conv.padding)
|
|
||||||
dilation_h, dilation_w = _pair(conv.dilation)
|
|
||||||
|
|
||||||
out_h = _conv_output_dim(height, kernel_h, stride_h, pad_h, dilation_h)
|
|
||||||
out_w = _conv_output_dim(width, kernel_w, stride_w, pad_w, dilation_w)
|
|
||||||
return out_h, out_w
|
|
||||||
|
|
||||||
|
|
||||||
def _conv_output_dim(size, kernel, stride, padding, dilation):
|
|
||||||
return ((size + (2 * padding) - (dilation * (kernel - 1)) - 1) // stride) + 1
|
|
||||||
|
|
||||||
|
|
||||||
def _pair(value):
|
|
||||||
if isinstance(value, tuple):
|
|
||||||
return value
|
|
||||||
return value, value
|
|
||||||
@@ -1,61 +0,0 @@
|
|||||||
from types import MethodType
|
|
||||||
|
|
||||||
from diffusers.models.transformers.transformer_2d import Transformer2DModel
|
|
||||||
|
|
||||||
|
|
||||||
def prepare_unet_for_coreml_trace(unet):
|
|
||||||
for module in unet.modules():
|
|
||||||
if isinstance(module, Transformer2DModel):
|
|
||||||
module._operate_on_continuous_inputs = MethodType(
|
|
||||||
_operate_on_continuous_inputs,
|
|
||||||
module,
|
|
||||||
)
|
|
||||||
module._get_output_for_continuous_inputs = MethodType(
|
|
||||||
_get_output_for_continuous_inputs,
|
|
||||||
module,
|
|
||||||
)
|
|
||||||
return unet
|
|
||||||
|
|
||||||
|
|
||||||
def _operate_on_continuous_inputs(self, hidden_states):
|
|
||||||
hidden_states = self.norm(hidden_states)
|
|
||||||
|
|
||||||
if not self.use_linear_projection:
|
|
||||||
hidden_states = self.proj_in(hidden_states)
|
|
||||||
inner_dim = self.inner_dim
|
|
||||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
|
||||||
else:
|
|
||||||
inner_dim = hidden_states.shape[1]
|
|
||||||
hidden_states = hidden_states.flatten(2).transpose(1, 2)
|
|
||||||
hidden_states = self.proj_in(hidden_states)
|
|
||||||
|
|
||||||
return hidden_states, inner_dim
|
|
||||||
|
|
||||||
|
|
||||||
def _get_output_for_continuous_inputs(
|
|
||||||
self,
|
|
||||||
hidden_states,
|
|
||||||
residual,
|
|
||||||
batch_size,
|
|
||||||
height,
|
|
||||||
width,
|
|
||||||
inner_dim,
|
|
||||||
):
|
|
||||||
if not self.use_linear_projection:
|
|
||||||
hidden_states = hidden_states.transpose(1, 2).reshape(
|
|
||||||
batch_size,
|
|
||||||
inner_dim,
|
|
||||||
height,
|
|
||||||
width,
|
|
||||||
)
|
|
||||||
hidden_states = self.proj_out(hidden_states)
|
|
||||||
else:
|
|
||||||
hidden_states = self.proj_out(hidden_states)
|
|
||||||
hidden_states = hidden_states.transpose(1, 2).reshape(
|
|
||||||
batch_size,
|
|
||||||
inner_dim,
|
|
||||||
height,
|
|
||||||
width,
|
|
||||||
)
|
|
||||||
|
|
||||||
return hidden_states + residual
|
|
||||||
@@ -1,54 +0,0 @@
|
|||||||
import torch
|
|
||||||
|
|
||||||
|
|
||||||
class CoreMLUNetWrapper(torch.nn.Module):
|
|
||||||
"""Adapt diffusers UNet inputs to CoreMLSuite's stable Core ML contract."""
|
|
||||||
|
|
||||||
def __init__(self, unet, model_version):
|
|
||||||
super().__init__()
|
|
||||||
self.unet = unet
|
|
||||||
self.model_version = model_version
|
|
||||||
|
|
||||||
def forward(self, sample, timestep, encoder_hidden_states, *extra_inputs):
|
|
||||||
input_index = 0
|
|
||||||
timestep_cond = None
|
|
||||||
if self._is_lcm:
|
|
||||||
timestep_cond = extra_inputs[input_index]
|
|
||||||
input_index += 1
|
|
||||||
|
|
||||||
added_cond_kwargs = None
|
|
||||||
if self._is_sdxl:
|
|
||||||
time_ids = extra_inputs[input_index]
|
|
||||||
text_embeds = extra_inputs[input_index + 1]
|
|
||||||
input_index += 2
|
|
||||||
added_cond_kwargs = {
|
|
||||||
"time_ids": time_ids,
|
|
||||||
"text_embeds": text_embeds,
|
|
||||||
}
|
|
||||||
|
|
||||||
additional_residuals = extra_inputs[input_index:]
|
|
||||||
down_residuals = None
|
|
||||||
mid_residual = None
|
|
||||||
if additional_residuals:
|
|
||||||
down_residuals = tuple(additional_residuals[:-1])
|
|
||||||
mid_residual = additional_residuals[-1]
|
|
||||||
|
|
||||||
outputs = self.unet(
|
|
||||||
sample,
|
|
||||||
timestep,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
timestep_cond=timestep_cond,
|
|
||||||
added_cond_kwargs=added_cond_kwargs,
|
|
||||||
down_block_additional_residuals=down_residuals,
|
|
||||||
mid_block_additional_residual=mid_residual,
|
|
||||||
return_dict=False,
|
|
||||||
)
|
|
||||||
return outputs[0]
|
|
||||||
|
|
||||||
@property
|
|
||||||
def _is_lcm(self):
|
|
||||||
return self.model_version.name == "LCM"
|
|
||||||
|
|
||||||
@property
|
|
||||||
def _is_sdxl(self):
|
|
||||||
return self.model_version.name in {"SDXL", "SDXL_REFINER"}
|
|
||||||
@@ -1,322 +0,0 @@
|
|||||||
import gc
|
|
||||||
import os
|
|
||||||
import time
|
|
||||||
|
|
||||||
import coremltools as ct
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from diffusers import UNet2DConditionModel
|
|
||||||
|
|
||||||
from coreml_suite.attention import ATTENTION_IMPLEMENTATIONS
|
|
||||||
from coreml_suite.conversion.attention import apply_attention_implementation
|
|
||||||
from coreml_suite.conversion.shapes import conv2d_output_shape
|
|
||||||
from coreml_suite.conversion.trace import prepare_unet_for_coreml_trace
|
|
||||||
from coreml_suite.conversion.unet import CoreMLUNetWrapper
|
|
||||||
from coreml_suite.logger import logger
|
|
||||||
from coreml_suite.model_version import ModelVersion
|
|
||||||
|
|
||||||
DEFAULT_TRACE_TIMESTEP = 999.0
|
|
||||||
TEXT_TOKEN_SEQUENCE_LENGTH = 77
|
|
||||||
|
|
||||||
|
|
||||||
def get_unet(model_version: ModelVersion, ref_unet, attention_implementation):
|
|
||||||
ref_unet = prepare_unet_for_coreml_trace(ref_unet)
|
|
||||||
unet = apply_attention_implementation(
|
|
||||||
ref_unet.eval(),
|
|
||||||
attention_implementation,
|
|
||||||
)
|
|
||||||
return CoreMLUNetWrapper(unet, model_version)
|
|
||||||
|
|
||||||
|
|
||||||
def get_encoder_hidden_states_shape(ref_unet, batch_size):
|
|
||||||
encoder_hidden_states_shape = (
|
|
||||||
batch_size,
|
|
||||||
TEXT_TOKEN_SEQUENCE_LENGTH,
|
|
||||||
ref_unet.config.cross_attention_dim,
|
|
||||||
)
|
|
||||||
|
|
||||||
return encoder_hidden_states_shape
|
|
||||||
|
|
||||||
|
|
||||||
def get_coreml_inputs(sample_inputs):
|
|
||||||
coreml_sample_unet_inputs = {
|
|
||||||
k: v.numpy().astype(np.float16) for k, v in sample_inputs.items()
|
|
||||||
}
|
|
||||||
return [
|
|
||||||
ct.TensorType(
|
|
||||||
name=k,
|
|
||||||
shape=v.shape,
|
|
||||||
dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
|
|
||||||
)
|
|
||||||
for k, v in coreml_sample_unet_inputs.items()
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def load_coreml_model(out_path):
|
|
||||||
logger.info(f"Loading model from {out_path}")
|
|
||||||
|
|
||||||
start = time.time()
|
|
||||||
coreml_model = ct.models.MLModel(out_path)
|
|
||||||
logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds")
|
|
||||||
|
|
||||||
return coreml_model
|
|
||||||
|
|
||||||
|
|
||||||
def convert_to_coreml(
|
|
||||||
submodule_name, torchscript_module, sample_inputs, output_names, out_path
|
|
||||||
):
|
|
||||||
if os.path.exists(out_path):
|
|
||||||
logger.info(f"Skipping export because {out_path} already exists")
|
|
||||||
coreml_model = load_coreml_model(out_path)
|
|
||||||
else:
|
|
||||||
logger.info(f"Converting {submodule_name} to CoreML..")
|
|
||||||
coreml_model = ct.convert(
|
|
||||||
torchscript_module,
|
|
||||||
convert_to="mlprogram",
|
|
||||||
minimum_deployment_target=ct.target.macOS13,
|
|
||||||
inputs=sample_inputs,
|
|
||||||
outputs=[
|
|
||||||
ct.TensorType(name=name, dtype=np.float32) for name in output_names
|
|
||||||
],
|
|
||||||
skip_model_load=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
del torchscript_module
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
return coreml_model
|
|
||||||
|
|
||||||
|
|
||||||
def get_out_path(submodule_name, model_name):
|
|
||||||
from folder_paths import get_folder_paths
|
|
||||||
|
|
||||||
fname = f"{model_name}_{submodule_name}.mlpackage"
|
|
||||||
unet_path = get_folder_paths(submodule_name)[0]
|
|
||||||
out_path = os.path.join(unet_path, fname)
|
|
||||||
return out_path
|
|
||||||
|
|
||||||
|
|
||||||
def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape):
|
|
||||||
sample_unet_inputs = dict(
|
|
||||||
[
|
|
||||||
("sample", torch.rand(*sample_shape)),
|
|
||||||
(
|
|
||||||
"timestep",
|
|
||||||
torch.tensor([DEFAULT_TRACE_TIMESTEP] * batch_size).to(torch.float32),
|
|
||||||
),
|
|
||||||
("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)),
|
|
||||||
]
|
|
||||||
)
|
|
||||||
return sample_unet_inputs
|
|
||||||
|
|
||||||
|
|
||||||
def lcm_inputs(sample_unet_inputs):
|
|
||||||
batch_size = sample_unet_inputs["sample"].shape[0]
|
|
||||||
return {"timestep_cond": torch.randn(batch_size, 256).to(torch.float32)}
|
|
||||||
|
|
||||||
|
|
||||||
def sdxl_inputs(sample_unet_inputs, ref_unet, model_version):
|
|
||||||
sample_shape = sample_unet_inputs["sample"].shape
|
|
||||||
batch_size = sample_shape[0]
|
|
||||||
h = sample_shape[2] * 8
|
|
||||||
w = sample_shape[3] * 8
|
|
||||||
original_size = (h, w)
|
|
||||||
crops_coords_top_left = (0, 0)
|
|
||||||
|
|
||||||
is_refiner = model_version == ModelVersion.SDXL_REFINER
|
|
||||||
|
|
||||||
if is_refiner:
|
|
||||||
aesthetic_score = (6.0,)
|
|
||||||
time_ids_list = list(original_size + crops_coords_top_left + aesthetic_score)
|
|
||||||
else:
|
|
||||||
target_size = (h, w)
|
|
||||||
time_ids_list = list(original_size + crops_coords_top_left + target_size)
|
|
||||||
|
|
||||||
time_ids = torch.tensor(time_ids_list).repeat(batch_size, 1).to(torch.int64)
|
|
||||||
text_embeds_shape = (batch_size, get_sdxl_text_embeds_dim(ref_unet, len(time_ids_list)))
|
|
||||||
|
|
||||||
return {
|
|
||||||
"time_ids": time_ids,
|
|
||||||
"text_embeds": torch.randn(*text_embeds_shape).to(torch.float32),
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def get_sdxl_text_embeds_dim(ref_unet, time_ids_dim):
|
|
||||||
projection_dim = ref_unet.config.projection_class_embeddings_input_dim
|
|
||||||
time_embed_dim = ref_unet.config.addition_time_embed_dim
|
|
||||||
return projection_dim - (time_ids_dim * time_embed_dim)
|
|
||||||
|
|
||||||
|
|
||||||
def get_inputs_spec(inputs):
|
|
||||||
inputs_spec = {k: (v.shape, v.dtype) for k, v in inputs.items()}
|
|
||||||
return inputs_spec
|
|
||||||
|
|
||||||
|
|
||||||
def add_cnet_support(sample_shape, reference_unet):
|
|
||||||
additional_residuals_shapes = []
|
|
||||||
|
|
||||||
batch_size = sample_shape[0]
|
|
||||||
h, w = sample_shape[2:]
|
|
||||||
|
|
||||||
# conv_in
|
|
||||||
out_h, out_w = conv2d_output_shape(
|
|
||||||
h,
|
|
||||||
w,
|
|
||||||
reference_unet.conv_in,
|
|
||||||
)
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(batch_size, reference_unet.conv_in.out_channels, out_h, out_w)
|
|
||||||
)
|
|
||||||
|
|
||||||
# down_blocks
|
|
||||||
for down_block in reference_unet.down_blocks:
|
|
||||||
additional_residuals_shapes += [
|
|
||||||
(batch_size, resnet.out_channels, out_h, out_w)
|
|
||||||
for resnet in down_block.resnets
|
|
||||||
]
|
|
||||||
if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
|
|
||||||
for downsampler in down_block.downsamplers:
|
|
||||||
out_h, out_w = conv2d_output_shape(out_h, out_w, downsampler.conv)
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(
|
|
||||||
batch_size,
|
|
||||||
down_block.downsamplers[-1].conv.out_channels,
|
|
||||||
out_h,
|
|
||||||
out_w,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
# mid_block
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w)
|
|
||||||
)
|
|
||||||
|
|
||||||
additional_inputs = {}
|
|
||||||
for i, shape in enumerate(additional_residuals_shapes):
|
|
||||||
sample_residual_input = torch.rand(*shape)
|
|
||||||
additional_inputs[f"additional_residual_{i}"] = sample_residual_input
|
|
||||||
|
|
||||||
return additional_inputs
|
|
||||||
|
|
||||||
|
|
||||||
def convert_unet(
|
|
||||||
ref_unet,
|
|
||||||
model_version: ModelVersion,
|
|
||||||
unet_out_path: str,
|
|
||||||
batch_size: int = 1,
|
|
||||||
sample_size: tuple[int, int] = (64, 64),
|
|
||||||
controlnet_support: bool = False,
|
|
||||||
attention_implementation: str = ATTENTION_IMPLEMENTATIONS[0],
|
|
||||||
quantize_nbits: str = "none",
|
|
||||||
):
|
|
||||||
coreml_unet = get_unet(model_version, ref_unet, attention_implementation)
|
|
||||||
|
|
||||||
sample_shape = (
|
|
||||||
batch_size, # B
|
|
||||||
ref_unet.config.in_channels, # C
|
|
||||||
sample_size[0], # H
|
|
||||||
sample_size[1], # W
|
|
||||||
)
|
|
||||||
|
|
||||||
encoder_hidden_states_shape = get_encoder_hidden_states_shape(ref_unet, batch_size)
|
|
||||||
|
|
||||||
sample_inputs = get_sample_input(
|
|
||||||
batch_size, encoder_hidden_states_shape, sample_shape
|
|
||||||
)
|
|
||||||
|
|
||||||
if model_version == ModelVersion.LCM:
|
|
||||||
sample_inputs |= lcm_inputs(sample_inputs)
|
|
||||||
|
|
||||||
if model_version in {ModelVersion.SDXL, ModelVersion.SDXL_REFINER}:
|
|
||||||
sample_inputs |= sdxl_inputs(sample_inputs, ref_unet, model_version)
|
|
||||||
|
|
||||||
if controlnet_support:
|
|
||||||
sample_inputs |= add_cnet_support(sample_shape, ref_unet)
|
|
||||||
|
|
||||||
sample_inputs_spec = get_inputs_spec(sample_inputs)
|
|
||||||
|
|
||||||
logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
|
|
||||||
logger.info("JIT tracing..")
|
|
||||||
traced_unet = torch.jit.trace(
|
|
||||||
coreml_unet, example_inputs=list(sample_inputs.values())
|
|
||||||
)
|
|
||||||
logger.info("Done.")
|
|
||||||
|
|
||||||
coreml_sample_inputs = get_coreml_inputs(sample_inputs)
|
|
||||||
|
|
||||||
coreml_unet = convert_to_coreml(
|
|
||||||
"unet", traced_unet, coreml_sample_inputs, ["noise_pred"], unet_out_path
|
|
||||||
)
|
|
||||||
|
|
||||||
del traced_unet
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
if quantize_nbits != "none":
|
|
||||||
# Opt-in k-means weight palettization. The default path
|
|
||||||
# (quantize_nbits="none") leaves the traced UNet untouched.
|
|
||||||
from coremltools.optimize.coreml import (
|
|
||||||
OpPalettizerConfig,
|
|
||||||
OptimizationConfig,
|
|
||||||
palettize_weights,
|
|
||||||
)
|
|
||||||
|
|
||||||
nbits = int(quantize_nbits)
|
|
||||||
logger.info(f"Palettizing UNet weights to {nbits}-bit (kmeans)..")
|
|
||||||
t0 = time.time()
|
|
||||||
cfg = OptimizationConfig(
|
|
||||||
global_config=OpPalettizerConfig(mode="kmeans", nbits=nbits)
|
|
||||||
)
|
|
||||||
coreml_unet = palettize_weights(coreml_unet, config=cfg)
|
|
||||||
logger.info(f"Palettization took {time.time() - t0:.1f}s")
|
|
||||||
|
|
||||||
coreml_unet.save(unet_out_path)
|
|
||||||
logger.info(f"Saved unet into {unet_out_path}")
|
|
||||||
|
|
||||||
|
|
||||||
def convert(
|
|
||||||
ckpt_path: str,
|
|
||||||
model_version: ModelVersion,
|
|
||||||
unet_out_path: str,
|
|
||||||
batch_size: int = 1,
|
|
||||||
sample_size: tuple[int, int] = (64, 64),
|
|
||||||
controlnet_support: bool = False,
|
|
||||||
lora_weights: list[tuple[str | os.PathLike, float]] = None,
|
|
||||||
attn_impl: str = ATTENTION_IMPLEMENTATIONS[0],
|
|
||||||
config_path: str = None,
|
|
||||||
quantize_nbits: str = "none",
|
|
||||||
):
|
|
||||||
if os.path.exists(unet_out_path):
|
|
||||||
logger.info(f"Found existing model at {unet_out_path}! Skipping..")
|
|
||||||
return
|
|
||||||
|
|
||||||
if attn_impl not in ATTENTION_IMPLEMENTATIONS:
|
|
||||||
raise ValueError(
|
|
||||||
f"Unsupported attention implementation {attn_impl!r}. "
|
|
||||||
f"Expected one of {ATTENTION_IMPLEMENTATIONS}."
|
|
||||||
)
|
|
||||||
ref_unet = load_unet(ckpt_path, config_path)
|
|
||||||
|
|
||||||
for i, lora_weight in enumerate(lora_weights or []):
|
|
||||||
lora_path, strength = lora_weight
|
|
||||||
adapter_name = f"lora_{i}"
|
|
||||||
ref_unet.load_lora_adapter(lora_path, adapter_name=adapter_name)
|
|
||||||
ref_unet.set_adapters([adapter_name], weights=[strength])
|
|
||||||
ref_unet.fuse_lora()
|
|
||||||
|
|
||||||
convert_unet(
|
|
||||||
ref_unet,
|
|
||||||
model_version,
|
|
||||||
unet_out_path,
|
|
||||||
batch_size,
|
|
||||||
sample_size,
|
|
||||||
controlnet_support,
|
|
||||||
attention_implementation=attn_impl,
|
|
||||||
quantize_nbits=quantize_nbits,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def load_unet(ckpt_path, config_path):
|
|
||||||
return UNet2DConditionModel.from_single_file(
|
|
||||||
ckpt_path,
|
|
||||||
original_config=config_path,
|
|
||||||
)
|
|
||||||
@@ -1,68 +0,0 @@
|
|||||||
"""Pure out_name composition for the Core ML UNet artifact.
|
|
||||||
|
|
||||||
Extracted from CoreMLConverter.convert so the filename contract
|
|
||||||
can be tested + reused without instantiating the node. The string is the
|
|
||||||
cache key: every workflow that references a converted .mlpackage depends
|
|
||||||
on it staying byte-for-byte identical.
|
|
||||||
"""
|
|
||||||
from typing import Iterable, Tuple
|
|
||||||
|
|
||||||
ATTN_SUFFIX = {
|
|
||||||
"SPLIT_EINSUM": "se",
|
|
||||||
"SPLIT_EINSUM_V2": "se2",
|
|
||||||
"ORIGINAL": "orig",
|
|
||||||
}
|
|
||||||
|
|
||||||
# Palettization bits. "none" = no quantization (default; keeps the
|
|
||||||
# unquantized filename intact so existing workflows still resolve their
|
|
||||||
# cached .mlpackage). Numeric values append a `_q<bits>` suffix.
|
|
||||||
QUANT_NBITS_VALUES = ("none", "8", "6", "4")
|
|
||||||
|
|
||||||
|
|
||||||
def compose_out_name(
|
|
||||||
*,
|
|
||||||
ckpt_name: str,
|
|
||||||
batch_size: int,
|
|
||||||
width: int,
|
|
||||||
height: int,
|
|
||||||
controlnet_support: bool,
|
|
||||||
attention_implementation: str,
|
|
||||||
lora_names: Iterable[str] = (),
|
|
||||||
quantize_nbits: str = "none",
|
|
||||||
) -> str:
|
|
||||||
"""Build the .mlpackage stem from convert() parameters.
|
|
||||||
|
|
||||||
Locked behaviour (characterization tests):
|
|
||||||
- first '.' in ckpt_name wins (`a.b.c.safetensors` -> `a`)
|
|
||||||
- spaces collapse to underscores
|
|
||||||
- LoRA names are taken stem-only, sorted, joined with '_' and
|
|
||||||
prefixed with '_' when present (caller is expected to pass a
|
|
||||||
sorted list; we sort defensively)
|
|
||||||
- controlnet adds `_cn`
|
|
||||||
- attn suffix is `_se` | `_se2` | `_orig`
|
|
||||||
|
|
||||||
Quantization:
|
|
||||||
- quantize_nbits "none" (default) appends nothing — existing
|
|
||||||
unquantized .mlpackages keep the old filename
|
|
||||||
- "4" / "6" / "8" appends `_q<bits>` after the attn suffix
|
|
||||||
"""
|
|
||||||
if quantize_nbits not in QUANT_NBITS_VALUES:
|
|
||||||
raise ValueError(
|
|
||||||
f"quantize_nbits={quantize_nbits!r} not in {QUANT_NBITS_VALUES}"
|
|
||||||
)
|
|
||||||
stem = ckpt_name.split(".")[0]
|
|
||||||
sorted_names = sorted(lora_names)
|
|
||||||
lora_str = "_" + "_".join(name.split(".")[0] for name in sorted_names) if sorted_names else ""
|
|
||||||
cn_suffix = "_cn" if controlnet_support else ""
|
|
||||||
attn_suffix = "_" + ATTN_SUFFIX[attention_implementation]
|
|
||||||
quant_suffix = f"_q{quantize_nbits}" if quantize_nbits != "none" else ""
|
|
||||||
out_name = (
|
|
||||||
f"{stem}{lora_str}_{batch_size}x{width}x{height}"
|
|
||||||
f"{cn_suffix}{attn_suffix}{quant_suffix}"
|
|
||||||
)
|
|
||||||
return out_name.replace(" ", "_")
|
|
||||||
|
|
||||||
|
|
||||||
def lora_names_from_params(lora_params: Iterable[Tuple[str, float]]) -> list[str]:
|
|
||||||
"""Mirror the sort applied inside CoreMLConverter.convert."""
|
|
||||||
return [name for name, _ in sorted(lora_params, key=lambda pair: pair[0])]
|
|
||||||
@@ -1,3 +1,8 @@
|
|||||||
from .nodes import COREML_CONVERT_LCM
|
"""LCM runtime support (sampler-side).
|
||||||
|
|
||||||
__all__ = ["COREML_CONVERT_LCM"]
|
The dedicated LCM converter node was removed once the standard ``CoreMLConverter``
|
||||||
|
gained model-version auto-detection (full-distill LCM is detected from the
|
||||||
|
checkpoint). What remains here is runtime sampling support — ``utils`` patches the
|
||||||
|
model sampling and supplies the guidance embedding when a converted UNet exposes
|
||||||
|
``timestep_cond``.
|
||||||
|
"""
|
||||||
|
|||||||
@@ -1,259 +0,0 @@
|
|||||||
import os
|
|
||||||
import logging
|
|
||||||
import time
|
|
||||||
import gc
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from diffusers import UNet2DConditionModel, LCMScheduler
|
|
||||||
from diffusers.loaders import LoraLoaderMixin
|
|
||||||
|
|
||||||
from coreml_suite.conversion.attention import apply_attention_implementation
|
|
||||||
from coreml_suite.conversion.shapes import conv2d_output_shape
|
|
||||||
from coreml_suite.conversion.unet import CoreMLUNetWrapper
|
|
||||||
from coreml_suite.model_version import ModelVersion
|
|
||||||
|
|
||||||
import coremltools as ct
|
|
||||||
|
|
||||||
logging.basicConfig()
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
logger.setLevel(logging.DEBUG)
|
|
||||||
|
|
||||||
MODEL_VERSION = "SimianLuo/LCM_Dreamshaper_v7"
|
|
||||||
MODEL_NAME = MODEL_VERSION.split("/")[-1] + "_4k"
|
|
||||||
TEXT_TOKEN_SEQUENCE_LENGTH = 77
|
|
||||||
|
|
||||||
|
|
||||||
def get_unets():
|
|
||||||
ref_unet = UNet2DConditionModel.from_pretrained(
|
|
||||||
MODEL_VERSION,
|
|
||||||
subfolder="unet",
|
|
||||||
device_map=None,
|
|
||||||
low_cpu_mem_usage=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
cml_unet = CoreMLUNetWrapper(
|
|
||||||
apply_attention_implementation(ref_unet.eval(), "SPLIT_EINSUM"),
|
|
||||||
ModelVersion.LCM,
|
|
||||||
)
|
|
||||||
|
|
||||||
return cml_unet, ref_unet
|
|
||||||
|
|
||||||
|
|
||||||
def get_encoder_hidden_states_shape(unet_config, batch_size):
|
|
||||||
encoder_hidden_states_shape = (
|
|
||||||
batch_size,
|
|
||||||
TEXT_TOKEN_SEQUENCE_LENGTH,
|
|
||||||
unet_config.cross_attention_dim,
|
|
||||||
)
|
|
||||||
|
|
||||||
return encoder_hidden_states_shape
|
|
||||||
|
|
||||||
|
|
||||||
def get_scheduler():
|
|
||||||
from comfy.model_management import get_torch_device
|
|
||||||
|
|
||||||
scheduler = LCMScheduler.from_pretrained(MODEL_VERSION, subfolder="scheduler")
|
|
||||||
scheduler.set_timesteps(50, get_torch_device(), 50)
|
|
||||||
return scheduler
|
|
||||||
|
|
||||||
|
|
||||||
def get_coreml_inputs(sample_inputs):
|
|
||||||
coreml_sample_unet_inputs = {
|
|
||||||
k: v.numpy().astype(np.float16) for k, v in sample_inputs.items()
|
|
||||||
}
|
|
||||||
return [
|
|
||||||
ct.TensorType(
|
|
||||||
name=k,
|
|
||||||
shape=v.shape,
|
|
||||||
dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
|
|
||||||
)
|
|
||||||
for k, v in coreml_sample_unet_inputs.items()
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def load_coreml_model(out_path):
|
|
||||||
logger.info(f"Loading model from {out_path}")
|
|
||||||
|
|
||||||
start = time.time()
|
|
||||||
coreml_model = ct.models.MLModel(out_path)
|
|
||||||
logger.info(f"Loading {out_path} took {time.time() - start:.1f} seconds")
|
|
||||||
|
|
||||||
return coreml_model
|
|
||||||
|
|
||||||
|
|
||||||
def convert_to_coreml(
|
|
||||||
submodule_name, torchscript_module, sample_inputs, output_names, out_path
|
|
||||||
):
|
|
||||||
if os.path.exists(out_path):
|
|
||||||
logger.info(f"Skipping export because {out_path} already exists")
|
|
||||||
coreml_model = load_coreml_model(out_path)
|
|
||||||
else:
|
|
||||||
logger.info(f"Converting {submodule_name} to CoreML..")
|
|
||||||
coreml_model = ct.convert(
|
|
||||||
torchscript_module,
|
|
||||||
convert_to="mlprogram",
|
|
||||||
minimum_deployment_target=ct.target.macOS13,
|
|
||||||
inputs=sample_inputs,
|
|
||||||
outputs=[
|
|
||||||
ct.TensorType(name=name, dtype=np.float32) for name in output_names
|
|
||||||
],
|
|
||||||
skip_model_load=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
del torchscript_module
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
return coreml_model
|
|
||||||
|
|
||||||
|
|
||||||
def get_out_path(submodule_name, model_name):
|
|
||||||
from folder_paths import get_folder_paths
|
|
||||||
|
|
||||||
fname = f"{model_name}_{submodule_name}.mlpackage"
|
|
||||||
unet_path = get_folder_paths(submodule_name)[0]
|
|
||||||
out_path = os.path.join(unet_path, fname)
|
|
||||||
return out_path
|
|
||||||
|
|
||||||
|
|
||||||
def get_sample_input(batch_size, encoder_hidden_states_shape, sample_shape, scheduler):
|
|
||||||
sample_unet_inputs = dict(
|
|
||||||
[
|
|
||||||
("sample", torch.rand(*sample_shape)),
|
|
||||||
(
|
|
||||||
"timestep",
|
|
||||||
torch.tensor([scheduler.timesteps[0].item()] * batch_size).to(
|
|
||||||
torch.float32
|
|
||||||
),
|
|
||||||
),
|
|
||||||
("encoder_hidden_states", torch.rand(*encoder_hidden_states_shape)),
|
|
||||||
("timestep_cond", torch.randn(batch_size, 256).to(torch.float32)),
|
|
||||||
]
|
|
||||||
)
|
|
||||||
return sample_unet_inputs
|
|
||||||
|
|
||||||
|
|
||||||
def get_unet_inputs_spec(sample_unet_inputs):
|
|
||||||
sample_unet_inputs_spec = {
|
|
||||||
k: (v.shape, v.dtype) for k, v in sample_unet_inputs.items()
|
|
||||||
}
|
|
||||||
return sample_unet_inputs_spec
|
|
||||||
|
|
||||||
|
|
||||||
def add_cnet_support(sample_shape, reference_unet):
|
|
||||||
additional_residuals_shapes = []
|
|
||||||
|
|
||||||
batch_size = sample_shape[0]
|
|
||||||
h, w = sample_shape[2:]
|
|
||||||
|
|
||||||
# conv_in
|
|
||||||
out_h, out_w = conv2d_output_shape(
|
|
||||||
h,
|
|
||||||
w,
|
|
||||||
reference_unet.conv_in,
|
|
||||||
)
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(batch_size, reference_unet.conv_in.out_channels, out_h, out_w)
|
|
||||||
)
|
|
||||||
|
|
||||||
# down_blocks
|
|
||||||
for down_block in reference_unet.down_blocks:
|
|
||||||
additional_residuals_shapes += [
|
|
||||||
(batch_size, resnet.out_channels, out_h, out_w)
|
|
||||||
for resnet in down_block.resnets
|
|
||||||
]
|
|
||||||
if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
|
|
||||||
for downsampler in down_block.downsamplers:
|
|
||||||
out_h, out_w = conv2d_output_shape(out_h, out_w, downsampler.conv)
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(
|
|
||||||
batch_size,
|
|
||||||
down_block.downsamplers[-1].conv.out_channels,
|
|
||||||
out_h,
|
|
||||||
out_w,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
|
|
||||||
# mid_block
|
|
||||||
additional_residuals_shapes.append(
|
|
||||||
(batch_size, reference_unet.mid_block.resnets[-1].out_channels, out_h, out_w)
|
|
||||||
)
|
|
||||||
|
|
||||||
additional_inputs = {}
|
|
||||||
for i, shape in enumerate(additional_residuals_shapes):
|
|
||||||
sample_residual_input = torch.rand(*shape)
|
|
||||||
additional_inputs[f"additional_residual_{i}"] = sample_residual_input
|
|
||||||
|
|
||||||
return additional_inputs
|
|
||||||
|
|
||||||
|
|
||||||
def convert(
|
|
||||||
out_path: str,
|
|
||||||
batch_size: int = 1,
|
|
||||||
sample_size: tuple[int, int] = (64, 64),
|
|
||||||
controlnet_support: bool = False,
|
|
||||||
lora_paths: list[str] = None,
|
|
||||||
):
|
|
||||||
lora_paths = lora_paths or []
|
|
||||||
coreml_unet, ref_unet = get_unets()
|
|
||||||
|
|
||||||
for lora_path in lora_paths:
|
|
||||||
lora_sd, network_alphas = LoraLoaderMixin.lora_state_dict(lora_path)
|
|
||||||
LoraLoaderMixin.load_lora_into_unet(lora_sd, network_alphas, ref_unet)
|
|
||||||
ref_unet.fuse_lora()
|
|
||||||
|
|
||||||
sample_shape = (
|
|
||||||
batch_size, # B
|
|
||||||
ref_unet.config.in_channels, # C
|
|
||||||
sample_size[0], # H
|
|
||||||
sample_size[1], # W
|
|
||||||
)
|
|
||||||
|
|
||||||
encoder_hidden_states_shape = get_encoder_hidden_states_shape(
|
|
||||||
ref_unet.config, batch_size
|
|
||||||
)
|
|
||||||
|
|
||||||
scheduler = get_scheduler()
|
|
||||||
|
|
||||||
sample_inputs = get_sample_input(
|
|
||||||
batch_size, encoder_hidden_states_shape, sample_shape, scheduler
|
|
||||||
)
|
|
||||||
|
|
||||||
if controlnet_support:
|
|
||||||
sample_inputs |= add_cnet_support(sample_shape, ref_unet)
|
|
||||||
|
|
||||||
sample_inputs_spec = get_unet_inputs_spec(sample_inputs)
|
|
||||||
|
|
||||||
logger.info(f"Sample UNet inputs spec: {sample_inputs_spec}")
|
|
||||||
logger.info("JIT tracing..")
|
|
||||||
traced_unet = torch.jit.trace(
|
|
||||||
coreml_unet, example_inputs=list(sample_inputs.values())
|
|
||||||
)
|
|
||||||
logger.info("Done.")
|
|
||||||
|
|
||||||
coreml_sample_inputs = get_coreml_inputs(sample_inputs)
|
|
||||||
|
|
||||||
coreml_unet = convert_to_coreml(
|
|
||||||
"unet", traced_unet, coreml_sample_inputs, ["noise_pred"], out_path
|
|
||||||
)
|
|
||||||
|
|
||||||
del traced_unet
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
coreml_unet.save(out_path)
|
|
||||||
logger.info(f"Saved unet into {out_path}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
h = 512
|
|
||||||
w = 512
|
|
||||||
sample_size = (h // 8, w // 8)
|
|
||||||
batch_size = 4
|
|
||||||
|
|
||||||
cn_support_str = "_cn" if True else ""
|
|
||||||
|
|
||||||
out_name = f"{MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
|
|
||||||
|
|
||||||
out_path = get_out_path("unet", f"{out_name}")
|
|
||||||
if not os.path.exists(out_path):
|
|
||||||
convert(out_path=out_path, sample_size=sample_size, batch_size=batch_size)
|
|
||||||
@@ -1,70 +0,0 @@
|
|||||||
import os
|
|
||||||
|
|
||||||
from coremltools import ComputeUnit
|
|
||||||
|
|
||||||
from coreml_suite import COREML_NODE
|
|
||||||
from coreml_suite.coreml_model import CoreMLModel
|
|
||||||
|
|
||||||
|
|
||||||
class COREML_CONVERT_LCM(COREML_NODE):
|
|
||||||
"""Converts a LCM model to Core ML."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"height": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
|
||||||
"width": ("INT", {"default": 512, "min": 512, "max": 768, "step": 8}),
|
|
||||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
|
||||||
"compute_unit": (
|
|
||||||
[
|
|
||||||
ComputeUnit.CPU_AND_NE.name,
|
|
||||||
ComputeUnit.CPU_AND_GPU.name,
|
|
||||||
ComputeUnit.ALL.name,
|
|
||||||
ComputeUnit.CPU_ONLY.name,
|
|
||||||
],
|
|
||||||
),
|
|
||||||
"controlnet_support": ("BOOLEAN", {"default": False}),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("COREML_UNET",)
|
|
||||||
RETURN_NAMES = ("coreml_model",)
|
|
||||||
FUNCTION = "convert"
|
|
||||||
|
|
||||||
def convert(self, height, width, batch_size, compute_unit, controlnet_support):
|
|
||||||
"""Converts a LCM model to Core ML.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
height (int): Height of the target image.
|
|
||||||
width (int): Width of the target image.
|
|
||||||
batch_size (int): Batch size.
|
|
||||||
compute_unit (str): Compute unit to use when loading the model.
|
|
||||||
|
|
||||||
Returns:
|
|
||||||
coreml_model: The converted Core ML model.
|
|
||||||
|
|
||||||
The converted model is also saved to "models/unet" directory and
|
|
||||||
can be loaded with the "LCMCoreMLLoaderUNet" node.
|
|
||||||
"""
|
|
||||||
from coreml_suite.lcm import converter as lcm_converter
|
|
||||||
|
|
||||||
h = height
|
|
||||||
w = width
|
|
||||||
sample_size = (h // 8, w // 8)
|
|
||||||
batch_size = batch_size
|
|
||||||
cn_support_str = "_cn" if controlnet_support else ""
|
|
||||||
|
|
||||||
out_name = f"{lcm_converter.MODEL_NAME}_{batch_size}x{w}x{h}{cn_support_str}"
|
|
||||||
|
|
||||||
out_path = lcm_converter.get_out_path("unet", f"{out_name}")
|
|
||||||
|
|
||||||
if not os.path.exists(out_path):
|
|
||||||
lcm_converter.convert(
|
|
||||||
out_path=out_path,
|
|
||||||
sample_size=sample_size,
|
|
||||||
batch_size=batch_size,
|
|
||||||
controlnet_support=controlnet_support,
|
|
||||||
)
|
|
||||||
|
|
||||||
return (CoreMLModel(out_path, compute_unit),)
|
|
||||||
@@ -1,98 +0,0 @@
|
|||||||
from diffusers import UNet2DConditionModel
|
|
||||||
from diffusers.models.embeddings import TimestepEmbedding
|
|
||||||
|
|
||||||
|
|
||||||
class UNet2DConditionModelLCM(UNet2DConditionModel):
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
time_cond_proj_dim=None,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
super().__init__(**kwargs)
|
|
||||||
timestep_input_dim = self.config.block_out_channels[0]
|
|
||||||
time_embed_dim = self.config.block_out_channels[0] * 4
|
|
||||||
|
|
||||||
time_embedding = TimestepEmbedding(
|
|
||||||
timestep_input_dim, time_embed_dim, cond_proj_dim=time_cond_proj_dim
|
|
||||||
)
|
|
||||||
self.time_embedding = time_embedding
|
|
||||||
|
|
||||||
def forward(
|
|
||||||
self,
|
|
||||||
sample,
|
|
||||||
timestep,
|
|
||||||
encoder_hidden_states,
|
|
||||||
timestep_cond,
|
|
||||||
*additional_residuals,
|
|
||||||
):
|
|
||||||
# 0. Project (or look-up) time embeddings
|
|
||||||
t_emb = self.time_proj(timestep)
|
|
||||||
emb = self.time_embedding(t_emb, timestep_cond)
|
|
||||||
|
|
||||||
# 1. center input if necessary
|
|
||||||
if self.config.center_input_sample:
|
|
||||||
sample = 2 * sample - 1.0
|
|
||||||
|
|
||||||
# 2. pre-process
|
|
||||||
sample = self.conv_in(sample)
|
|
||||||
|
|
||||||
# 3. down
|
|
||||||
down_block_res_samples = (sample,)
|
|
||||||
for downsample_block in self.down_blocks:
|
|
||||||
if (
|
|
||||||
hasattr(downsample_block, "attentions")
|
|
||||||
and downsample_block.attentions is not None
|
|
||||||
):
|
|
||||||
sample, res_samples = downsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample, res_samples = downsample_block(hidden_states=sample, temb=emb)
|
|
||||||
|
|
||||||
down_block_res_samples += res_samples
|
|
||||||
|
|
||||||
if additional_residuals:
|
|
||||||
new_down_block_res_samples = ()
|
|
||||||
for i, down_block_res_sample in enumerate(down_block_res_samples):
|
|
||||||
down_block_res_sample = down_block_res_sample + additional_residuals[i]
|
|
||||||
new_down_block_res_samples += (down_block_res_sample,)
|
|
||||||
down_block_res_samples = new_down_block_res_samples
|
|
||||||
|
|
||||||
# 4. mid
|
|
||||||
sample = self.mid_block(
|
|
||||||
sample, emb, encoder_hidden_states=encoder_hidden_states
|
|
||||||
)
|
|
||||||
|
|
||||||
if additional_residuals:
|
|
||||||
sample = sample + additional_residuals[-1]
|
|
||||||
|
|
||||||
# 5. up
|
|
||||||
for upsample_block in self.up_blocks:
|
|
||||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
|
||||||
down_block_res_samples = down_block_res_samples[
|
|
||||||
: -len(upsample_block.resnets)
|
|
||||||
]
|
|
||||||
|
|
||||||
if (
|
|
||||||
hasattr(upsample_block, "attentions")
|
|
||||||
and upsample_block.attentions is not None
|
|
||||||
):
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample,
|
|
||||||
temb=emb,
|
|
||||||
res_hidden_states_tuple=res_samples,
|
|
||||||
encoder_hidden_states=encoder_hidden_states,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
sample = upsample_block(
|
|
||||||
hidden_states=sample, temb=emb, res_hidden_states_tuple=res_samples
|
|
||||||
)
|
|
||||||
|
|
||||||
# 6. post-process
|
|
||||||
sample = self.conv_norm_out(sample)
|
|
||||||
sample = self.conv_act(sample)
|
|
||||||
sample = self.conv_out(sample)
|
|
||||||
|
|
||||||
return (sample,)
|
|
||||||
@@ -1,8 +0,0 @@
|
|||||||
from enum import Enum
|
|
||||||
|
|
||||||
|
|
||||||
class ModelVersion(Enum):
|
|
||||||
SD15 = "sd15"
|
|
||||||
SDXL = "sdxl"
|
|
||||||
SDXL_REFINER = "sdxl_refiner"
|
|
||||||
LCM = "lcm"
|
|
||||||
+51
-31
@@ -4,16 +4,9 @@ from coremltools import ComputeUnit
|
|||||||
|
|
||||||
import folder_paths
|
import folder_paths
|
||||||
from coreml_suite import COREML_NODE
|
from coreml_suite import COREML_NODE
|
||||||
from coreml_suite.attention import ATTENTION_IMPLEMENTATIONS
|
|
||||||
from coreml_suite.coreml_model import CoreMLModel
|
from coreml_suite.coreml_model import CoreMLModel
|
||||||
from coreml_suite.core.naming import (
|
|
||||||
QUANT_NBITS_VALUES,
|
|
||||||
compose_out_name,
|
|
||||||
lora_names_from_params,
|
|
||||||
)
|
|
||||||
from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
|
from coreml_suite.lcm.utils import add_lcm_model_options, lcm_patch, is_lcm
|
||||||
from coreml_suite.logger import logger
|
from coreml_suite.logger import logger
|
||||||
from coreml_suite.model_version import ModelVersion
|
|
||||||
from nodes import KSampler, LoraLoader, KSamplerAdvanced
|
from nodes import KSampler, LoraLoader, KSamplerAdvanced
|
||||||
|
|
||||||
from coreml_suite.models import (
|
from coreml_suite.models import (
|
||||||
@@ -24,6 +17,26 @@ from coreml_suite.models import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _discover(fn_name, fallback):
|
||||||
|
"""Populate a converter dropdown from coreml_diffusion's discovery API.
|
||||||
|
|
||||||
|
Fails soft: if the package is missing, too old to expose ``fn_name``, or
|
||||||
|
errors, the node still registers with the fallback list instead of vanishing
|
||||||
|
from the menu. Evaluated on every INPUT_TYPES call, so installing a newer
|
||||||
|
coreml_diffusion surfaces new conversion types with no Suite change.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
import coreml_diffusion
|
||||||
|
|
||||||
|
return getattr(coreml_diffusion, fn_name)()
|
||||||
|
except Exception as exc: # missing/old package, import error, etc.
|
||||||
|
logger.warning(
|
||||||
|
f"coreml_diffusion.{fn_name} unavailable ({exc}); "
|
||||||
|
f"using fallback {fallback}"
|
||||||
|
)
|
||||||
|
return fallback
|
||||||
|
|
||||||
|
|
||||||
class CoreMLSampler(COREML_NODE, KSampler):
|
class CoreMLSampler(COREML_NODE, KSampler):
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(s):
|
def INPUT_TYPES(s):
|
||||||
@@ -212,24 +225,26 @@ class CoreMLModelAdapter(COREML_NODE):
|
|||||||
|
|
||||||
|
|
||||||
class CoreMLConverter(COREML_NODE):
|
class CoreMLConverter(COREML_NODE):
|
||||||
"""Converts a LCM model to Core ML."""
|
"""Converts a Stable Diffusion checkpoint (UNet) to Core ML.
|
||||||
|
|
||||||
|
The model version (SD15 / SDXL / SDXL refiner / LCM) is auto-detected from
|
||||||
|
the checkpoint's architecture, so there is no version dropdown — one node
|
||||||
|
converts every supported family, including full-distill LCM.
|
||||||
|
"""
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def INPUT_TYPES(cls):
|
def INPUT_TYPES(cls):
|
||||||
return {
|
return {
|
||||||
"required": {
|
"required": {
|
||||||
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
"ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
|
||||||
"model_version": (
|
|
||||||
[
|
|
||||||
ModelVersion.SD15.name,
|
|
||||||
ModelVersion.SDXL.name,
|
|
||||||
],
|
|
||||||
),
|
|
||||||
"height": ("INT", {"default": 512, "min": 8, "step": 8}),
|
"height": ("INT", {"default": 512, "min": 8, "step": 8}),
|
||||||
"width": ("INT", {"default": 512, "min": 8, "step": 8}),
|
"width": ("INT", {"default": 512, "min": 8, "step": 8}),
|
||||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
"batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
|
||||||
"attention_implementation": (
|
"attention_implementation": (
|
||||||
list(ATTENTION_IMPLEMENTATIONS),
|
_discover(
|
||||||
|
"list_attention_impls",
|
||||||
|
["SPLIT_EINSUM", "SPLIT_EINSUM_V2", "ORIGINAL"],
|
||||||
|
),
|
||||||
),
|
),
|
||||||
"compute_unit": (
|
"compute_unit": (
|
||||||
[
|
[
|
||||||
@@ -247,7 +262,10 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
# omits any `required` input. When omitted it defaults to
|
# omits any `required` input. When omitted it defaults to
|
||||||
# "none", identical to unquantized behavior and filename, so
|
# "none", identical to unquantized behavior and filename, so
|
||||||
# existing cached .mlpackages still resolve.
|
# existing cached .mlpackages still resolve.
|
||||||
"quantize_nbits": (list(QUANT_NBITS_VALUES), {"default": "none"}),
|
"quantize_nbits": (
|
||||||
|
_discover("list_quant_modes", ["none", "8", "6", "4"]),
|
||||||
|
{"default": "none"},
|
||||||
|
),
|
||||||
"lora_params": ("LORA_PARAMS",),
|
"lora_params": ("LORA_PARAMS",),
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
@@ -259,7 +277,6 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
def convert(
|
def convert(
|
||||||
self,
|
self,
|
||||||
ckpt_name,
|
ckpt_name,
|
||||||
model_version,
|
|
||||||
height,
|
height,
|
||||||
width,
|
width,
|
||||||
batch_size,
|
batch_size,
|
||||||
@@ -269,9 +286,11 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
quantize_nbits="none",
|
quantize_nbits="none",
|
||||||
lora_params=None,
|
lora_params=None,
|
||||||
):
|
):
|
||||||
"""Converts a LCM model to Core ML.
|
"""Converts a checkpoint's UNet to Core ML.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
|
ckpt_name (str): Checkpoint to convert; its model version is
|
||||||
|
auto-detected from the weights.
|
||||||
height (int): Height of the target image.
|
height (int): Height of the target image.
|
||||||
width (int): Width of the target image.
|
width (int): Width of the target image.
|
||||||
batch_size (int): Batch size.
|
batch_size (int): Batch size.
|
||||||
@@ -281,10 +300,8 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
coreml_model: The converted Core ML model.
|
coreml_model: The converted Core ML model.
|
||||||
|
|
||||||
The converted model is also saved to "models/unet" directory and
|
The converted model is also saved to "models/unet" directory and
|
||||||
can be loaded with the "LCMCoreMLLoaderUNet" node.
|
can be loaded with the "Load Core ML UNet" node.
|
||||||
"""
|
"""
|
||||||
model_version = ModelVersion[model_version]
|
|
||||||
|
|
||||||
lora_params = lora_params or {}
|
lora_params = lora_params or {}
|
||||||
lora_params = [(k, v[0]) for k, v in lora_params.items()]
|
lora_params = [(k, v[0]) for k, v in lora_params.items()]
|
||||||
lora_params = sorted(lora_params, key=lambda lora: lora[0])
|
lora_params = sorted(lora_params, key=lambda lora: lora[0])
|
||||||
@@ -293,14 +310,16 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
h = height
|
h = height
|
||||||
w = width
|
w = width
|
||||||
sample_size = (h // 8, w // 8)
|
sample_size = (h // 8, w // 8)
|
||||||
out_name = compose_out_name(
|
import coreml_diffusion
|
||||||
|
|
||||||
|
out_name = coreml_diffusion.compose_out_name(
|
||||||
ckpt_name=ckpt_name,
|
ckpt_name=ckpt_name,
|
||||||
batch_size=batch_size,
|
batch_size=batch_size,
|
||||||
width=w,
|
width=w,
|
||||||
height=h,
|
height=h,
|
||||||
controlnet_support=controlnet_support,
|
controlnet_support=controlnet_support,
|
||||||
attention_implementation=attention_implementation,
|
attention_implementation=attention_implementation,
|
||||||
lora_names=lora_names_from_params(lora_params),
|
lora_names=coreml_diffusion.lora_names_from_params(lora_params),
|
||||||
quantize_nbits=quantize_nbits,
|
quantize_nbits=quantize_nbits,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -311,13 +330,14 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
logger.info(f"Attention implementation: {attention_implementation}")
|
logger.info(f"Attention implementation: {attention_implementation}")
|
||||||
|
|
||||||
if lora_params:
|
if lora_params:
|
||||||
logger.info(f"LoRAs used:")
|
logger.info("LoRAs used:")
|
||||||
for lora_param in lora_params:
|
for lora_param in lora_params:
|
||||||
logger.info(f" {lora_param[0]} - strength: {lora_param[1]}")
|
logger.info(f" {lora_param[0]} - strength: {lora_param[1]}")
|
||||||
|
|
||||||
from coreml_suite import converter
|
# Resolve the ComfyUI models/unet path here (a node concern); the package
|
||||||
|
# takes the output path as an injected argument.
|
||||||
unet_out_path = converter.get_out_path("unet", f"{out_name}")
|
unet_path = folder_paths.get_folder_paths("unet")[0]
|
||||||
|
unet_out_path = os.path.join(unet_path, f"{out_name}_unet.mlpackage")
|
||||||
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
|
||||||
|
|
||||||
config_filename = ckpt_name.split(".")[0] + ".yaml"
|
config_filename = ckpt_name.split(".")[0] + ".yaml"
|
||||||
@@ -325,10 +345,10 @@ class CoreMLConverter(COREML_NODE):
|
|||||||
if config_path:
|
if config_path:
|
||||||
logger.info(f"Using config file {config_path}")
|
logger.info(f"Using config file {config_path}")
|
||||||
|
|
||||||
converter.convert(
|
coreml_diffusion.convert(
|
||||||
ckpt_path=ckpt_path,
|
ckpt_path,
|
||||||
model_version=model_version,
|
None, # model_version auto-detected from the checkpoint
|
||||||
unet_out_path=unet_out_path,
|
unet_out_path,
|
||||||
sample_size=sample_size,
|
sample_size=sample_size,
|
||||||
batch_size=batch_size,
|
batch_size=batch_size,
|
||||||
controlnet_support=controlnet_support,
|
controlnet_support=controlnet_support,
|
||||||
|
|||||||
@@ -0,0 +1,93 @@
|
|||||||
|
# Conversion
|
||||||
|
|
||||||
|
## Conversion is the only supported path
|
||||||
|
|
||||||
|
You always start from a Stable Diffusion checkpoint (`.safetensors` / `.ckpt`)
|
||||||
|
and convert it with the **Convert Checkpoint to Core ML** node. The model
|
||||||
|
version (SD1.5, SDXL, SDXL refiner, full-distill LCM) is auto-detected from the
|
||||||
|
checkpoint. Pre-converted Core ML models from elsewhere are not supported,
|
||||||
|
because:
|
||||||
|
|
||||||
|
- The suite uses its own input **dimensions**, **naming convention**, and
|
||||||
|
**metadata**, all produced by the
|
||||||
|
[coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) package.
|
||||||
|
- Apple's `ml-stable-diffusion` (which most community Core ML models target) is
|
||||||
|
effectively obsolete, and the layouts differ (see the 2.0.0
|
||||||
|
`encoder_hidden_states` change in the README).
|
||||||
|
- Conversion is cheap and one-time, so there is no value in maintaining
|
||||||
|
backwards compatibility with foreign formats.
|
||||||
|
|
||||||
|
The output is always a **`.mlpackage`**. The suite no longer compiles to
|
||||||
|
`.mlmodelc` (it didn't work with the inference backend), so there is **no Xcode
|
||||||
|
or `coremlcompiler` dependency**.
|
||||||
|
|
||||||
|
## One-time conversion and name-based caching
|
||||||
|
|
||||||
|
Conversion runs **once**, not on every queue. The converter encodes all
|
||||||
|
conversion parameters into the output filename (via `coreml_diffusion.compose_out_name`,
|
||||||
|
called in `coreml_suite/nodes.py`):
|
||||||
|
|
||||||
|
- checkpoint name, `batch_size`, `width`, `height`
|
||||||
|
- `controlnet_support`, `attention_implementation`
|
||||||
|
- baked LoRA names, `quantize_nbits`
|
||||||
|
|
||||||
|
The result is written as `<encoded-name>_unet.mlpackage` in `models/unet`. If a
|
||||||
|
file with that name already exists, it is reused and conversion is skipped. Change
|
||||||
|
any parameter → new name → new conversion; keep them the same → the cached model
|
||||||
|
is loaded instantly.
|
||||||
|
|
||||||
|
This is why the recommended workflow is **convert once, then load**: run the
|
||||||
|
converter a single time, then in day-to-day use load the `.mlpackage` with the
|
||||||
|
**Load Core ML UNet** node. (You can also leave the converter node in the graph;
|
||||||
|
it short-circuits to the cached file.)
|
||||||
|
|
||||||
|
> [!NOTE]
|
||||||
|
> The converter relies on the filename to decide whether to reconvert. If you
|
||||||
|
> rename the `.mlpackage`, it will be converted again. You can otherwise rename
|
||||||
|
> it freely if the auto-generated name is too long.
|
||||||
|
|
||||||
|
## Quantization
|
||||||
|
|
||||||
|
The converter node accepts an optional `quantize_nbits` dropdown that runs
|
||||||
|
k-means weight palettization (`coremltools.optimize.coreml.palettize_weights`) on
|
||||||
|
the UNet before saving.
|
||||||
|
|
||||||
|
Values: `none` (default — no quantization, identical output and filename to
|
||||||
|
before), `8`, `6`, `4`. The number is appended to the `.mlpackage` stem as
|
||||||
|
`_q<bits>`, so quantized and unquantized variants coexist on disk and in cache.
|
||||||
|
|
||||||
|
### SD1.5 1×512×512 SPLIT_EINSUM tradeoffs (M2 Pro, ANE)
|
||||||
|
|
||||||
|
Measured with 20 UNet forward passes at a fixed seed:
|
||||||
|
|
||||||
|
| nbits | size (MB) | size vs none | fwd median (ms) | PSNR vs `none` (dB) |
|
||||||
|
|---|---:|---:|---:|---:|
|
||||||
|
| none | 1641 | 1.000 | 197.1 | — |
|
||||||
|
| 8 | 822 | 0.501 | 186.6 | 53.5 |
|
||||||
|
| 6 | 617 | 0.376 | 183.0 | 40.2 |
|
||||||
|
| 4 | 412 | 0.251 | 179.8 | 27.5 |
|
||||||
|
|
||||||
|
PSNR here is computed on the raw `noise_pred` output of a single UNet forward at a
|
||||||
|
fixed seed, not on the final decoded image — it isolates quantization drift from
|
||||||
|
sampler/VAE noise. Final-image PSNR is comfortably higher (the sampler averages
|
||||||
|
over many steps).
|
||||||
|
|
||||||
|
### Recommended settings per chip / RAM
|
||||||
|
|
||||||
|
- **8 GB (M1/M2 base):** `nbits=4`. ~4× smaller, still loads, 27 dB is visually
|
||||||
|
identical at SD1.5 sizes.
|
||||||
|
- **16 GB (M1/M2/M3 Pro):** `nbits=6` — the sweet spot, ~2.7× smaller, 40 dB, no
|
||||||
|
perceptible quality drop.
|
||||||
|
- **32 GB+ (Max / Ultra):** `nbits=8` for a safety margin, or `none` for
|
||||||
|
bit-identical output (golden testing).
|
||||||
|
|
||||||
|
The default stays `none`, so existing workflows produce byte-for-byte identical
|
||||||
|
output.
|
||||||
|
|
||||||
|
## Where conversion lives
|
||||||
|
|
||||||
|
The conversion engine was extracted into the standalone
|
||||||
|
[coreml-diffusion](https://github.com/aszc-dev/coreml-diffusion) PyPI package.
|
||||||
|
The nodes in this suite resolve ComfyUI paths and call into it; node names,
|
||||||
|
inputs, and outputs are unchanged, so the split has effectively no user-facing
|
||||||
|
impact beyond `pip install` pulling one more dependency.
|
||||||
+77
@@ -0,0 +1,77 @@
|
|||||||
|
# FAQ
|
||||||
|
|
||||||
|
## What's the difference between ANE, GPU, and MPS, and which do I pick?
|
||||||
|
|
||||||
|
ANE is the Neural Engine (Core ML only), GPU is the Metal GPU (Core ML or
|
||||||
|
PyTorch), MPS is PyTorch's GPU backend. This suite uses **Core ML compute units
|
||||||
|
only** and never touches MPS. Short answer: SD1.5 at 512×512 → convert
|
||||||
|
`SPLIT_EINSUM`, load `CPU_AND_NE`; larger sizes or SDXL → convert `ORIGINAL`,
|
||||||
|
load `CPU_AND_GPU`. Full reasoning: [hardware](hardware.md).
|
||||||
|
|
||||||
|
## Do I still need `PYTORCH_ENABLE_MPS_FALLBACK=1`?
|
||||||
|
|
||||||
|
Not for these nodes — Core ML inference doesn't use PyTorch MPS. It may still
|
||||||
|
matter for other parts of your ComfyUI graph, but it has no effect on Core ML
|
||||||
|
sampling.
|
||||||
|
|
||||||
|
## Why is my Core ML SDXL workflow no faster than the default nodes?
|
||||||
|
|
||||||
|
Because **SDXL can't run on the ANE** — the speedup comes from the Neural Engine,
|
||||||
|
and SDXL falls back to the GPU, running at roughly MPS-equivalent speed. This is a
|
||||||
|
known limitation, not a misconfiguration. The ANE benefit is real for SD1.5. See
|
||||||
|
[limitations](limitations.md).
|
||||||
|
|
||||||
|
## Where do I get Core ML models?
|
||||||
|
|
||||||
|
You convert them yourself — that's the only supported path. See
|
||||||
|
[conversion](conversion.md). Downloaded Core ML models (e.g. coreml-community) use
|
||||||
|
different dimensions/metadata and are not supported.
|
||||||
|
|
||||||
|
## Is conversion run every time I queue, or once?
|
||||||
|
|
||||||
|
Once. Parameters are encoded in the output filename, so an already-converted model
|
||||||
|
is reused and conversion is skipped. Convert once, then load the `.mlpackage`. See
|
||||||
|
[conversion → caching](conversion.md#one-time-conversion-and-name-based-caching).
|
||||||
|
|
||||||
|
## Does a converted model produce the same output as the original?
|
||||||
|
|
||||||
|
With the default `quantize_nbits = none`, the converted UNet output matches the
|
||||||
|
source within numerical rounding (the golden test in `tests/m2/test_golden_image.py`
|
||||||
|
gates on PSNR ≥ 20 dB on the decoded image). Quantization (`8`/`6`/`4`) introduces
|
||||||
|
measured, bounded drift — see the [PSNR table](conversion.md#quantization). For
|
||||||
|
bit-identical output, keep `none`.
|
||||||
|
|
||||||
|
## Are `.mlpackage` models safe to use?
|
||||||
|
|
||||||
|
`.mlpackage` is a declarative Core ML model format — it carries weights and a
|
||||||
|
compute graph, not arbitrary executable code or Python pickle, so its safety
|
||||||
|
profile is comparable to `safetensors`. In practice this matters little here,
|
||||||
|
since the only supported models are ones you convert locally from your own
|
||||||
|
checkpoints.
|
||||||
|
|
||||||
|
## Are LoRAs reliable?
|
||||||
|
|
||||||
|
Partially. Some LoRAs convert cleanly; others produce poor or broken output —
|
||||||
|
there's no firm rule, so test per-LoRA. LoRA weights and `strength_model` are
|
||||||
|
baked in at conversion and can't be changed afterward; for some LCM-LoRA cases the
|
||||||
|
[Core ML Adapter](nodes.md#core-ml-adapter-experimental-coremlmodeladapter) path
|
||||||
|
is more reliable. Treat LoRA support as experimental. See
|
||||||
|
[troubleshooting](troubleshooting.md).
|
||||||
|
|
||||||
|
## Does the experimental Adapter cost performance vs the Core ML Sampler?
|
||||||
|
|
||||||
|
Yes, a little. The Adapter wraps the model in a ComfyUI `ModelPatcher` so standard
|
||||||
|
samplers work, which adds per-step interface overhead the native Core ML Sampler
|
||||||
|
avoids. Use the native sampler unless you specifically need a `MODEL` (e.g.
|
||||||
|
`ModelSamplingDiscrete` for LCM LoRAs).
|
||||||
|
|
||||||
|
## Which Python versions work?
|
||||||
|
|
||||||
|
Python 3.12 or newer (`requires-python >=3.12`). Older 3.12 install failures
|
||||||
|
came from the now-removed `ml-stable-diffusion` build, not from this suite.
|
||||||
|
|
||||||
|
## Long prompts crash my workflow
|
||||||
|
|
||||||
|
Core ML has a hard **77-token** prompt limit and doesn't auto-chunk long prompts.
|
||||||
|
Split the prompt across multiple CLIP Text Encode nodes and merge with Conditioning
|
||||||
|
(Combine). See [troubleshooting](troubleshooting.md).
|
||||||
@@ -0,0 +1,95 @@
|
|||||||
|
# Hardware & Compute Units
|
||||||
|
|
||||||
|
This page explains how the suite maps to Apple Silicon hardware, the difference
|
||||||
|
between ANE, GPU, and MPS, and how to choose a compute unit and attention
|
||||||
|
implementation.
|
||||||
|
|
||||||
|
## ANE vs GPU vs MPS
|
||||||
|
|
||||||
|
Three terms get conflated:
|
||||||
|
|
||||||
|
- **ANE (Apple Neural Engine)** — a dedicated ML accelerator on Apple Silicon.
|
||||||
|
Only Core ML can target it; PyTorch cannot. This is the whole reason the suite
|
||||||
|
exists.
|
||||||
|
- **GPU** — the Metal GPU. Reachable both by Core ML (as a compute unit) and by
|
||||||
|
PyTorch (via MPS).
|
||||||
|
- **MPS (Metal Performance Shaders)** — PyTorch's GPU backend on macOS. This is
|
||||||
|
the path standard ComfyUI nodes use.
|
||||||
|
|
||||||
|
**This suite uses Core ML compute units only — it never runs the UNet through
|
||||||
|
PyTorch/MPS.** Consequently `PYTORCH_ENABLE_MPS_FALLBACK` has no effect on these
|
||||||
|
nodes. It may still matter for the rest of your ComfyUI graph (CLIP, VAE,
|
||||||
|
samplers on non-Core ML models), but not for Core ML inference itself.
|
||||||
|
|
||||||
|
Rough performance picture (SD1.5, maintainer- and user-reported):
|
||||||
|
|
||||||
|
- ANE is meaningfully faster than MPS — on the order of **50–100%** for SD1.5.
|
||||||
|
- Core ML on the GPU is only marginally faster than PyTorch/MPS.
|
||||||
|
|
||||||
|
So the speedup comes from the Neural Engine, which means it depends on being able
|
||||||
|
to actually run on the ANE (see [attention implementations](#attention-implementations)
|
||||||
|
and the [SDXL caveat](#sdxl-and-the-ane)).
|
||||||
|
|
||||||
|
## Compute units
|
||||||
|
|
||||||
|
The **compute unit** is set on the loader/converter node and tells Core ML which
|
||||||
|
hardware to use. It is applied when the model is loaded
|
||||||
|
(`coreml_suite/coreml_model.py:22`), not during conversion.
|
||||||
|
|
||||||
|
| Value | Hardware | Best paired with |
|
||||||
|
|---|---|---|
|
||||||
|
| `CPU_AND_NE` (default) | CPU + Neural Engine | `SPLIT_EINSUM` / `SPLIT_EINSUM_V2` |
|
||||||
|
| `CPU_AND_GPU` | CPU + Metal GPU | `ORIGINAL` |
|
||||||
|
| `CPU_ONLY` | CPU only | fallback / debugging |
|
||||||
|
| `ALL` | all available hardware | rarely optimal — see below |
|
||||||
|
|
||||||
|
Notes:
|
||||||
|
|
||||||
|
- Every option includes the CPU; there is no GPU-and-ANE-without-CPU combination.
|
||||||
|
- `NE` in `CPU_AND_NE` is the Neural Engine (Apple's enum spells it `NE`, not
|
||||||
|
`ANE`).
|
||||||
|
- **`CPU_AND_NE` is often faster than `ALL`.** Letting Core ML use everything can
|
||||||
|
be *slower* on non-Max chips, where memory bandwidth is the bottleneck. Try
|
||||||
|
`CPU_AND_NE` first for SD1.5.
|
||||||
|
|
||||||
|
## Attention implementations
|
||||||
|
|
||||||
|
Chosen at conversion time on the **Convert Checkpoint to Core ML** node. It
|
||||||
|
decides whether the model can run on the ANE:
|
||||||
|
|
||||||
|
- **`SPLIT_EINSUM`** — ANE-friendly attention. Use for the Neural Engine.
|
||||||
|
- **`SPLIT_EINSUM_V2`** — a variant; in practice ≈ `SPLIT_EINSUM` for most users.
|
||||||
|
- **`ORIGINAL`** — standard attention. Runs on the GPU, not the ANE.
|
||||||
|
|
||||||
|
The implementation and the compute unit must agree: a `SPLIT_EINSUM` model wants
|
||||||
|
`CPU_AND_NE`; an `ORIGINAL` model wants `CPU_AND_GPU`.
|
||||||
|
|
||||||
|
## Which should I pick?
|
||||||
|
|
||||||
|
| Scenario | Attention | Compute unit |
|
||||||
|
|---|---|---|
|
||||||
|
| SD1.5 at 512×512 | `SPLIT_EINSUM` | `CPU_AND_NE` |
|
||||||
|
| SD1.5 at larger sizes (e.g. 768) | `ORIGINAL` | `CPU_AND_GPU` |
|
||||||
|
| SDXL / SDXL Turbo | `ORIGINAL` | `CPU_AND_GPU` |
|
||||||
|
|
||||||
|
### Resolution crossover
|
||||||
|
|
||||||
|
ANE shines at small latents; the GPU scales better as resolution grows. In user
|
||||||
|
benchmarks:
|
||||||
|
|
||||||
|
- At **512×512**, ANE + `SPLIT_EINSUM` wins by roughly **10%** over the GPU path.
|
||||||
|
- At **768×512**, GPU + `ORIGINAL` pulls ahead by roughly **10%**, and the larger
|
||||||
|
image is about 2× slower overall.
|
||||||
|
|
||||||
|
If you mostly work at 512×512, convert with `SPLIT_EINSUM` and load on
|
||||||
|
`CPU_AND_NE`. If you routinely go larger, an `ORIGINAL` + GPU model may be
|
||||||
|
faster.
|
||||||
|
|
||||||
|
### SDXL and the ANE
|
||||||
|
|
||||||
|
SDXL (and SDXL Turbo) **cannot run on the ANE** — the dual-text-encoder UNet
|
||||||
|
exceeds what the Neural Engine path supports. SDXL therefore runs at roughly
|
||||||
|
MPS-equivalent speed with no ANE speedup. If a Core ML SDXL workflow feels no
|
||||||
|
faster than the standard nodes, this is why. Convert SDXL with `ORIGINAL` and
|
||||||
|
load with `CPU_AND_GPU` or `CPU_ONLY`. See
|
||||||
|
[limitations](limitations.md) for the full picture.
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
# Limitations & Support Matrix
|
||||||
|
|
||||||
|
## Support matrix
|
||||||
|
|
||||||
|
| Feature | Status | Notes |
|
||||||
|
|---|---|---|
|
||||||
|
| SD1.5 | ✅ Full | ANE via `SPLIT_EINSUM`; the primary, fastest path |
|
||||||
|
| SDXL / SDXL Turbo | ⚠️ Partial | GPU only (no ANE), no speedup; possible quality loss vs source. Don't run Turbo at 1024² |
|
||||||
|
| SD2.1 | ❌ Unsupported | |
|
||||||
|
| Inpainting checkpoints (9-channel) | ❌ Unsupported | |
|
||||||
|
| ControlNet | ✅ Supported | Convert the checkpoint with `controlnet_support = True` |
|
||||||
|
| LoRA | ⚠️ Experimental | Inconsistent per-LoRA; baked at conversion, immutable afterward |
|
||||||
|
| LCM | ⚠️ Experimental | Full-distill LCM checkpoints auto-detected by the converter |
|
||||||
|
| SVD | ❌ Not supported | |
|
||||||
|
| AnimateDiff | ❌ Not supported | Motion modules need pre-conversion injection; not feasible today |
|
||||||
|
| IPAdapter | ❌ Not supported | Needs a real `MODEL` the Core ML wrapper can't provide |
|
||||||
|
| Core ML Adapter | ⚠️ Experimental | Works for many nodes; fails for merges/IPAdapter/etc. |
|
||||||
|
|
||||||
|
## Fixed input/output shapes
|
||||||
|
|
||||||
|
A Core ML model is converted for one specific resolution and batch size. To work
|
||||||
|
at a different size, re-convert with the new width/height (conversion is cheap and
|
||||||
|
cached by name). This is also why detailers and latent-upscale workflows that
|
||||||
|
rescale mid-graph break — see [troubleshooting](troubleshooting.md).
|
||||||
|
|
||||||
|
There is experimental support for flexible shapes via
|
||||||
|
[EnumeratedShapes](https://apple.github.io/coremltools/docs-guides/source/flexible-inputs.html#select-from-predetermined-shapes),
|
||||||
|
but it is **much slower** — user benchmarks show roughly **5×** the per-iteration
|
||||||
|
time on every run, not just the first. Fixed-shape models per resolution are the
|
||||||
|
practical choice.
|
||||||
|
|
||||||
|
## SDXL on the Neural Engine
|
||||||
|
|
||||||
|
SDXL and SDXL Turbo cannot run on the ANE — the dual-text-encoder UNet exceeds the
|
||||||
|
supported Neural Engine path. They run on the GPU at roughly MPS-equivalent speed,
|
||||||
|
so Core ML offers no speed advantage for SDXL, and converted output may look
|
||||||
|
degraded versus the safetensors original (an upstream conversion artifact). Use
|
||||||
|
`ORIGINAL` + `CPU_AND_GPU`. See [hardware](hardware.md).
|
||||||
|
|
||||||
|
## Experimental Core ML Adapter
|
||||||
|
|
||||||
|
The Adapter wraps a Core ML model to look like a standard ComfyUI `MODEL`, which
|
||||||
|
covers many standard and custom nodes. But it can't fully emulate a real model:
|
||||||
|
operations that need genuine `MODEL` internals — model merges, IPAdapter, some
|
||||||
|
LoRA flows, detailers — generally won't work, and the model's
|
||||||
|
fixed input shapes aren't validated, so mismatches error at runtime. Prefer the
|
||||||
|
native Core ML Sampler when you don't need the `MODEL` type.
|
||||||
|
|
||||||
|
## Prompt length
|
||||||
|
|
||||||
|
Core ML enforces a hard 77-token prompt limit with no auto-chunking. Split long
|
||||||
|
prompts across multiple CLIP Text Encode nodes and merge with Conditioning
|
||||||
|
(Combine).
|
||||||
+157
@@ -0,0 +1,157 @@
|
|||||||
|
# Node Reference
|
||||||
|
|
||||||
|
All nodes live in the **Core ML Suite** category. Right-click the canvas →
|
||||||
|
**Add Node → Core ML Suite**, or double-click and search.
|
||||||
|
|
||||||
|
| Display name | Class | Purpose |
|
||||||
|
|---|---|---|
|
||||||
|
| Load Core ML UNet | `CoreMLUNetLoader` | Load a converted `.mlpackage` |
|
||||||
|
| Core ML Sampler | `CoreMLSampler` | Sample (KSampler-style) |
|
||||||
|
| Core ML Sampler (Advanced) | `CoreMLSamplerAdvanced` | Sample (KSamplerAdvanced-style) |
|
||||||
|
| Core ML Adapter (Experimental) | `CoreMLModelAdapter` | Wrap as a standard `MODEL` |
|
||||||
|
| Load LoRA to use with Core ML | `Core ML LoRA Loader` | Bake LoRA(s) at conversion |
|
||||||
|
| Convert Checkpoint to Core ML | `Core ML Converter` | Convert a checkpoint |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## Load Core ML UNet (`CoreMLUNetLoader`)
|
||||||
|
|
||||||
|

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

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

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

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

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

|
||||||
|
|
||||||
|
**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`:
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
## 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)).
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
## Checkpoint conversion
|
||||||
|
|
||||||
|
The minimal conversion graph. See
|
||||||
|
[Convert Checkpoint to Core ML](nodes.md#convert-checkpoint-to-core-ml-core-ml-converter).
|
||||||
|
|
||||||
|

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

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

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

|
||||||
|
|
||||||
|
## 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`.
|
||||||
|
|
||||||
|

|
||||||
|
|
||||||
|
## 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).
|
||||||
|
|
||||||
|

|
||||||
+16
-9
@@ -1,27 +1,34 @@
|
|||||||
|
[build-system]
|
||||||
|
requires = ["hatchling"]
|
||||||
|
build-backend = "hatchling.build"
|
||||||
|
|
||||||
[project]
|
[project]
|
||||||
name = "comfyui-coremlsuite"
|
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."
|
description = "This extension contains a set of custom nodes for ComfyUI that allow you to use Core ML models in your ComfyUI workflows."
|
||||||
version = "2.0.2"
|
version = "2.1.2"
|
||||||
license = "MIT"
|
license = "MIT"
|
||||||
requires-python = ">=3.12,<3.13"
|
requires-python = ">=3.12"
|
||||||
packages = [{ include = "coreml_suite" }]
|
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"torch>=2.7,<2.8",
|
# torch is provided by the host (ComfyUI) and intentionally left unpinned
|
||||||
|
# here: a hard torch cap would downgrade the host's torch and break its
|
||||||
|
# torchvision/torchaudio ABI. coreml-diffusion pulls torch>=2.7 transitively.
|
||||||
|
# >=0.1.6: model-version auto-detection (convert(model_version=None)) and the
|
||||||
|
# dropped <3.13 Python cap (kept in sync with this package's requires-python).
|
||||||
|
"coreml-diffusion>=0.1.6,<0.2",
|
||||||
"coremltools>=9,<10",
|
"coremltools>=9,<10",
|
||||||
"numpy>=2,<3",
|
"numpy>=2,<3",
|
||||||
"diffusers>=0.30",
|
|
||||||
"peft>=0.13",
|
|
||||||
"omegaconf>=2.3",
|
|
||||||
"transformers>=4.44",
|
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite"
|
Repository = "https://github.com/aszc-dev/ComfyUI-CoreMLSuite"
|
||||||
|
|
||||||
|
[tool.hatch.build.targets.wheel]
|
||||||
|
packages = ["coreml_suite"]
|
||||||
|
|
||||||
[tool.comfy]
|
[tool.comfy]
|
||||||
PublisherId = "aszc-dev"
|
PublisherId = "aszc-dev"
|
||||||
DisplayName = "ComfyUI-CoreMLSuite"
|
DisplayName = "ComfyUI-CoreMLSuite"
|
||||||
Icon = ""
|
Icon = "https://raw.githubusercontent.com/aszc-dev/ComfyUI-CoreMLSuite/main/assets/snake.png"
|
||||||
requires-comfyui = ">=0.3.27"
|
requires-comfyui = ">=0.3.27"
|
||||||
|
|
||||||
[dependency-groups]
|
[dependency-groups]
|
||||||
|
|||||||
+1
-4
@@ -1,7 +1,4 @@
|
|||||||
torch>=2.7,<2.8
|
coreml-diffusion>=0.1.4,<0.2
|
||||||
coremltools>=9,<10
|
coremltools>=9,<10
|
||||||
numpy>=2,<3
|
numpy>=2,<3
|
||||||
diffusers>=0.30
|
diffusers>=0.30
|
||||||
peft>=0.13
|
|
||||||
omegaconf>=2.3
|
|
||||||
transformers>=4.44
|
|
||||||
|
|||||||
@@ -107,7 +107,6 @@
|
|||||||
"10": {
|
"10": {
|
||||||
"inputs": {
|
"inputs": {
|
||||||
"ckpt_name": "dreamshaper_8.safetensors",
|
"ckpt_name": "dreamshaper_8.safetensors",
|
||||||
"model_version": "SD15",
|
|
||||||
"height": 512,
|
"height": 512,
|
||||||
"width": 512,
|
"width": 512,
|
||||||
"batch_size": 1,
|
"batch_size": 1,
|
||||||
|
|||||||
@@ -1,41 +0,0 @@
|
|||||||
import platform
|
|
||||||
|
|
||||||
import pytest
|
|
||||||
import torch
|
|
||||||
from diffusers.models.attention_processor import Attention, AttnProcessor
|
|
||||||
|
|
||||||
from coreml_suite.conversion.attention import (
|
|
||||||
SplitEinsumAttnProcessor,
|
|
||||||
SplitEinsumV2AttnProcessor,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
pytestmark = pytest.mark.skipif(
|
|
||||||
platform.system() != "Darwin" or platform.machine() != "arm64",
|
|
||||||
reason="Tier 1 requires macOS on Apple Silicon",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"processor",
|
|
||||||
[
|
|
||||||
SplitEinsumAttnProcessor(),
|
|
||||||
SplitEinsumV2AttnProcessor(),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_split_einsum_processor_matches_diffusers_attention(processor):
|
|
||||||
torch.manual_seed(0)
|
|
||||||
reference = Attention(query_dim=32, heads=4, dim_head=8, dropout=0.0)
|
|
||||||
reference.set_processor(AttnProcessor())
|
|
||||||
|
|
||||||
candidate = Attention(query_dim=32, heads=4, dim_head=8, dropout=0.0)
|
|
||||||
candidate.load_state_dict(reference.state_dict())
|
|
||||||
candidate.set_processor(processor)
|
|
||||||
|
|
||||||
hidden_states = torch.randn(2, 17, 32)
|
|
||||||
encoder_hidden_states = torch.randn(2, 11, 32)
|
|
||||||
|
|
||||||
expected = reference(hidden_states, encoder_hidden_states=encoder_hidden_states)
|
|
||||||
actual = candidate(hidden_states, encoder_hidden_states=encoder_hidden_states)
|
|
||||||
|
|
||||||
assert torch.allclose(actual, expected, atol=1e-5)
|
|
||||||
@@ -1,138 +0,0 @@
|
|||||||
"""Tier 1 smoke: convert a synthetic micro-UNet through coremltools and load
|
|
||||||
it back with CoreMLSuite's runtime CoreMLModel wrapper.
|
|
||||||
|
|
||||||
Purpose: catch API breakage in coremltools *without* needing a real SD
|
|
||||||
checkpoint, the ANE, or a converted .mlmodelc on disk.
|
|
||||||
Runs in minutes on a hosted macOS-ARM runner (no Apple internal stuff).
|
|
||||||
|
|
||||||
What it asserts:
|
|
||||||
- coremltools.convert still accepts the call shape we use today
|
|
||||||
- the resulting .mlpackage round-trips through CoreMLSuite's CoreMLModel
|
|
||||||
- expected_inputs exposes the input names/shapes we declared
|
|
||||||
- calling the model returns the named output (`noise_pred`)
|
|
||||||
|
|
||||||
Auto-skips on non-Apple-Silicon hosts so Tier 0 CI on Linux ignores it.
|
|
||||||
"""
|
|
||||||
import platform
|
|
||||||
import shutil
|
|
||||||
from types import SimpleNamespace
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import pytest
|
|
||||||
import torch
|
|
||||||
import torch.nn as nn
|
|
||||||
|
|
||||||
from coreml_suite.conversion.unet import CoreMLUNetWrapper
|
|
||||||
|
|
||||||
|
|
||||||
pytestmark = pytest.mark.skipif(
|
|
||||||
platform.system() != "Darwin" or platform.machine() != "arm64",
|
|
||||||
reason="Tier 1 requires macOS on Apple Silicon",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# Tiny shapes — large enough to exercise conv2d + linear + addition kernels in
|
|
||||||
# coremltools, small enough that conversion finishes in seconds on CPU.
|
|
||||||
SAMPLE_SHAPE = (1, 4, 8, 8)
|
|
||||||
TIMESTEP_SHAPE = (1,)
|
|
||||||
ENCODER_SHAPE = (1, 4, 64) # native diffusers encoder_hidden_states (batch, tokens, hidden)
|
|
||||||
OUT_NAME = "noise_pred"
|
|
||||||
|
|
||||||
|
|
||||||
class TinyUNet(nn.Module):
|
|
||||||
"""Minimal UNet-shaped graph: conv -> add(time+context) -> conv.
|
|
||||||
|
|
||||||
Not a real diffusion model. Just enough op variety to exercise the
|
|
||||||
PyTorch -> MIL frontend in coremltools and confirm we can still wire
|
|
||||||
the inputs/outputs the way CoreMLSuite's runtime expects.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self):
|
|
||||||
super().__init__()
|
|
||||||
self.conv_in = nn.Conv2d(4, 8, kernel_size=3, padding=1)
|
|
||||||
self.conv_out = nn.Conv2d(8, 4, kernel_size=3, padding=1)
|
|
||||||
self.time_proj = nn.Linear(1, 8)
|
|
||||||
self.text_proj = nn.Linear(64, 8)
|
|
||||||
|
|
||||||
def forward(
|
|
||||||
self,
|
|
||||||
sample,
|
|
||||||
timestep,
|
|
||||||
encoder_hidden_states,
|
|
||||||
timestep_cond=None,
|
|
||||||
added_cond_kwargs=None,
|
|
||||||
down_block_additional_residuals=None,
|
|
||||||
mid_block_additional_residual=None,
|
|
||||||
return_dict=True,
|
|
||||||
):
|
|
||||||
h = self.conv_in(sample)
|
|
||||||
t_emb = self.time_proj(timestep.unsqueeze(-1)).view(1, 8, 1, 1)
|
|
||||||
c_emb = self.text_proj(encoder_hidden_states.mean(1)).view(1, 8, 1, 1)
|
|
||||||
h = h + t_emb + c_emb
|
|
||||||
return (self.conv_out(h),)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.fixture(scope="module")
|
|
||||||
def tiny_mlpackage(tmp_path_factory):
|
|
||||||
"""Convert TinyUNet once per test session and reuse the .mlpackage."""
|
|
||||||
import coremltools as ct
|
|
||||||
|
|
||||||
torch.manual_seed(0)
|
|
||||||
model = CoreMLUNetWrapper(
|
|
||||||
TinyUNet().eval(),
|
|
||||||
SimpleNamespace(name="SD15"),
|
|
||||||
)
|
|
||||||
example = (
|
|
||||||
torch.randn(*SAMPLE_SHAPE),
|
|
||||||
torch.randn(*TIMESTEP_SHAPE),
|
|
||||||
torch.randn(*ENCODER_SHAPE),
|
|
||||||
)
|
|
||||||
traced = torch.jit.trace(model, example)
|
|
||||||
|
|
||||||
mlmodel = ct.convert(
|
|
||||||
traced,
|
|
||||||
inputs=[
|
|
||||||
ct.TensorType(name="sample", shape=SAMPLE_SHAPE, dtype=np.float16),
|
|
||||||
ct.TensorType(name="timestep", shape=TIMESTEP_SHAPE, dtype=np.float16),
|
|
||||||
ct.TensorType(name="encoder_hidden_states", shape=ENCODER_SHAPE, dtype=np.float16),
|
|
||||||
],
|
|
||||||
outputs=[ct.TensorType(name=OUT_NAME, dtype=np.float16)],
|
|
||||||
compute_units=ct.ComputeUnit.CPU_ONLY,
|
|
||||||
compute_precision=ct.precision.FLOAT16,
|
|
||||||
convert_to="mlprogram",
|
|
||||||
minimum_deployment_target=ct.target.macOS13,
|
|
||||||
)
|
|
||||||
|
|
||||||
out_dir = tmp_path_factory.mktemp("tiny_unet")
|
|
||||||
pkg_path = out_dir / "tiny.mlpackage"
|
|
||||||
mlmodel.save(str(pkg_path))
|
|
||||||
yield pkg_path
|
|
||||||
shutil.rmtree(out_dir, ignore_errors=True)
|
|
||||||
|
|
||||||
|
|
||||||
def test_coremltools_convert_round_trips_via_coreml_model(tiny_mlpackage):
|
|
||||||
from coreml_suite.coreml_model import CoreMLModel
|
|
||||||
|
|
||||||
model = CoreMLModel(str(tiny_mlpackage), "CPU_ONLY")
|
|
||||||
|
|
||||||
# expected_inputs is the contract our wrappers depend on. Lock the shape
|
|
||||||
# of the dict + a sample entry.
|
|
||||||
expected = dict(model.expected_inputs)
|
|
||||||
assert set(expected.keys()) == {"sample", "timestep", "encoder_hidden_states"}
|
|
||||||
assert tuple(expected["sample"]["shape"]) == SAMPLE_SHAPE
|
|
||||||
assert tuple(expected["timestep"]["shape"]) == TIMESTEP_SHAPE
|
|
||||||
assert tuple(expected["encoder_hidden_states"]["shape"]) == ENCODER_SHAPE
|
|
||||||
|
|
||||||
# Forward pass: drive the model the way CoreMLModelWrapper does.
|
|
||||||
rng = np.random.default_rng(0)
|
|
||||||
inputs = {
|
|
||||||
"sample": rng.standard_normal(SAMPLE_SHAPE).astype(np.float16),
|
|
||||||
"timestep": rng.standard_normal(TIMESTEP_SHAPE).astype(np.float16),
|
|
||||||
"encoder_hidden_states": rng.standard_normal(ENCODER_SHAPE).astype(np.float16),
|
|
||||||
}
|
|
||||||
out = model(**inputs)
|
|
||||||
assert isinstance(out, dict), f"unexpected output type: {type(out)}"
|
|
||||||
assert OUT_NAME in out, f"missing output {OUT_NAME!r}; got {sorted(out)}"
|
|
||||||
assert out[OUT_NAME].shape == SAMPLE_SHAPE, (
|
|
||||||
f"output shape drift: got {out[OUT_NAME].shape}, expected {SAMPLE_SHAPE}"
|
|
||||||
)
|
|
||||||
@@ -1,197 +0,0 @@
|
|||||||
"""Characterization tests for the .mlpackage filename composition.
|
|
||||||
|
|
||||||
The filename composition is the pure
|
|
||||||
coreml_suite.core.naming.compose_out_name function. CoreMLConverter.convert
|
|
||||||
calls it; testing the pure function avoids monkey-patching heavy converter
|
|
||||||
internals just to capture the string.
|
|
||||||
"""
|
|
||||||
import pytest
|
|
||||||
|
|
||||||
from coreml_suite.core.naming import compose_out_name, lora_names_from_params
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- attention suffixes ----------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"attn_name,suffix",
|
|
||||||
[
|
|
||||||
("SPLIT_EINSUM", "se"),
|
|
||||||
("SPLIT_EINSUM_V2", "se2"),
|
|
||||||
("ORIGINAL", "orig"),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_attention_suffix(attn_name, suffix):
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation=attn_name,
|
|
||||||
)
|
|
||||||
assert out == f"dreamshaper_8_1x512x512_{suffix}"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- batch / size ----------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_includes_batch_and_size():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=4, width=768, height=1024,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_4x768x1024_se"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- ControlNet ------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_appends_cn_suffix_when_controlnet_support_true():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=True,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_1x512x512_cn_se"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- ckpt name massage -----------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_drops_extension_at_first_period():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="my.checkpoint.v2.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
)
|
|
||||||
assert out == "my_1x512x512_se"
|
|
||||||
|
|
||||||
|
|
||||||
def test_replaces_spaces_with_underscores():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dream shaper 8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
)
|
|
||||||
assert out == "dream_shaper_8_1x512x512_se"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- LoRA suffixes ---------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_single_lora():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
lora_names=["epi_noiseoffset.safetensors"],
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_epi_noiseoffset_1x512x512_se"
|
|
||||||
|
|
||||||
|
|
||||||
def test_multiple_loras_sorted():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
lora_names=["zoom.safetensors", "alpha.safetensors", "moody.safetensors"],
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_alpha_moody_zoom_1x512x512_se"
|
|
||||||
|
|
||||||
|
|
||||||
def test_lora_plus_controlnet():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=True,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
lora_names=["a.safetensors"],
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_a_1x512x512_cn_se"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- sdxl combinations -----------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_sdxl_1024_original_gpu():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="sd_xl_base_1.0.safetensors",
|
|
||||||
batch_size=1, width=1024, height=1024,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="ORIGINAL",
|
|
||||||
)
|
|
||||||
assert out == "sd_xl_base_1_1x1024x1024_orig"
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- lora_names_from_params helper ----------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_lora_names_from_params_sorts_by_name():
|
|
||||||
names = lora_names_from_params([
|
|
||||||
("zebra.safetensors", 1.0),
|
|
||||||
("apple.safetensors", 0.5),
|
|
||||||
("mango.safetensors", 0.7),
|
|
||||||
])
|
|
||||||
assert names == ["apple.safetensors", "mango.safetensors", "zebra.safetensors"]
|
|
||||||
|
|
||||||
|
|
||||||
def test_lora_names_from_params_empty_list():
|
|
||||||
assert lora_names_from_params([]) == []
|
|
||||||
|
|
||||||
|
|
||||||
# ---------- quantize_nbits suffix ------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def test_quantize_nbits_none_appends_nothing():
|
|
||||||
"""'none' is the default and must keep the unquantized filename so
|
|
||||||
existing cached .mlpackages still resolve."""
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
quantize_nbits="none",
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_1x512x512_se"
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("nbits,suffix", [("4", "_q4"), ("6", "_q6"), ("8", "_q8")])
|
|
||||||
def test_quantize_nbits_appends_q_suffix(nbits, suffix):
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
quantize_nbits=nbits,
|
|
||||||
)
|
|
||||||
assert out == f"dreamshaper_8_1x512x512_se{suffix}"
|
|
||||||
|
|
||||||
|
|
||||||
def test_quantize_nbits_with_controlnet_and_lora():
|
|
||||||
out = compose_out_name(
|
|
||||||
ckpt_name="dreamshaper_8.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=True,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
lora_names=["a.safetensors"],
|
|
||||||
quantize_nbits="6",
|
|
||||||
)
|
|
||||||
assert out == "dreamshaper_8_a_1x512x512_cn_se_q6"
|
|
||||||
|
|
||||||
|
|
||||||
def test_quantize_nbits_invalid_raises():
|
|
||||||
import pytest as _pytest
|
|
||||||
with _pytest.raises(ValueError, match="quantize_nbits"):
|
|
||||||
compose_out_name(
|
|
||||||
ckpt_name="x.safetensors",
|
|
||||||
batch_size=1, width=512, height=512,
|
|
||||||
controlnet_support=False,
|
|
||||||
attention_implementation="SPLIT_EINSUM",
|
|
||||||
quantize_nbits="16", # not in {none, 8, 6, 4}
|
|
||||||
)
|
|
||||||
@@ -1,183 +0,0 @@
|
|||||||
from types import SimpleNamespace
|
|
||||||
|
|
||||||
import torch
|
|
||||||
|
|
||||||
from coreml_suite.conversion.attention import (
|
|
||||||
SplitEinsumAttnProcessor,
|
|
||||||
SplitEinsumV2AttnProcessor,
|
|
||||||
apply_attention_implementation,
|
|
||||||
split_einsum,
|
|
||||||
split_einsum_v2,
|
|
||||||
)
|
|
||||||
from coreml_suite.conversion.shapes import conv2d_output_shape
|
|
||||||
from coreml_suite.conversion.unet import CoreMLUNetWrapper
|
|
||||||
|
|
||||||
|
|
||||||
class RecordingUNet(torch.nn.Module):
|
|
||||||
def __init__(self):
|
|
||||||
super().__init__()
|
|
||||||
self.call = None
|
|
||||||
|
|
||||||
def forward(
|
|
||||||
self,
|
|
||||||
sample,
|
|
||||||
timestep,
|
|
||||||
encoder_hidden_states,
|
|
||||||
timestep_cond=None,
|
|
||||||
added_cond_kwargs=None,
|
|
||||||
down_block_additional_residuals=None,
|
|
||||||
mid_block_additional_residual=None,
|
|
||||||
return_dict=True,
|
|
||||||
**kwargs,
|
|
||||||
):
|
|
||||||
self.call = {
|
|
||||||
"sample": sample,
|
|
||||||
"timestep": timestep,
|
|
||||||
"encoder_hidden_states": encoder_hidden_states,
|
|
||||||
"timestep_cond": timestep_cond,
|
|
||||||
"added_cond_kwargs": added_cond_kwargs,
|
|
||||||
"down_block_additional_residuals": down_block_additional_residuals,
|
|
||||||
"mid_block_additional_residual": mid_block_additional_residual,
|
|
||||||
"return_dict": return_dict,
|
|
||||||
}
|
|
||||||
return (sample + 1,)
|
|
||||||
|
|
||||||
|
|
||||||
def test_conv2d_output_shape_matches_torch_conv2d_contract():
|
|
||||||
conv = torch.nn.Conv2d(
|
|
||||||
4,
|
|
||||||
8,
|
|
||||||
kernel_size=(3, 5),
|
|
||||||
stride=(2, 3),
|
|
||||||
padding=(1, 2),
|
|
||||||
dilation=(1, 2),
|
|
||||||
)
|
|
||||||
|
|
||||||
assert conv2d_output_shape(17, 19, conv) == (9, 5)
|
|
||||||
|
|
||||||
|
|
||||||
def test_unet_wrapper_passes_context_through_for_sd15():
|
|
||||||
unet = RecordingUNet()
|
|
||||||
wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="SD15"))
|
|
||||||
|
|
||||||
sample = torch.randn(2, 4, 8, 8)
|
|
||||||
timestep = torch.randn(2)
|
|
||||||
context = torch.randn(2, 77, 768)
|
|
||||||
|
|
||||||
out = wrapper(sample, timestep, context)
|
|
||||||
|
|
||||||
assert torch.equal(out, sample + 1)
|
|
||||||
assert unet.call["encoder_hidden_states"] is context
|
|
||||||
assert unet.call["return_dict"] is False
|
|
||||||
|
|
||||||
|
|
||||||
def test_unet_wrapper_routes_lcm_sdxl_and_controlnet_inputs():
|
|
||||||
unet = RecordingUNet()
|
|
||||||
wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="LCM"))
|
|
||||||
|
|
||||||
sample = torch.randn(1, 4, 8, 8)
|
|
||||||
timestep = torch.randn(1)
|
|
||||||
context = torch.randn(1, 77, 768)
|
|
||||||
timestep_cond = torch.randn(1, 256)
|
|
||||||
down_residual = torch.randn(1, 320, 8, 8)
|
|
||||||
mid_residual = torch.randn(1, 1280, 1, 1)
|
|
||||||
|
|
||||||
wrapper(sample, timestep, context, timestep_cond, down_residual, mid_residual)
|
|
||||||
|
|
||||||
assert unet.call["timestep_cond"] is timestep_cond
|
|
||||||
assert len(unet.call["down_block_additional_residuals"]) == 1
|
|
||||||
assert unet.call["down_block_additional_residuals"][0] is down_residual
|
|
||||||
assert unet.call["mid_block_additional_residual"] is mid_residual
|
|
||||||
|
|
||||||
|
|
||||||
def test_unet_wrapper_routes_sdxl_added_conditioning():
|
|
||||||
unet = RecordingUNet()
|
|
||||||
wrapper = CoreMLUNetWrapper(unet, SimpleNamespace(name="SDXL"))
|
|
||||||
|
|
||||||
sample = torch.randn(1, 4, 8, 8)
|
|
||||||
timestep = torch.randn(1)
|
|
||||||
context = torch.randn(1, 77, 2048)
|
|
||||||
time_ids = torch.randn(1, 6)
|
|
||||||
text_embeds = torch.randn(1, 1280)
|
|
||||||
|
|
||||||
wrapper(sample, timestep, context, time_ids, text_embeds)
|
|
||||||
|
|
||||||
assert unet.call["added_cond_kwargs"]["time_ids"] is time_ids
|
|
||||||
assert unet.call["added_cond_kwargs"]["text_embeds"] is text_embeds
|
|
||||||
|
|
||||||
|
|
||||||
def test_split_einsum_matches_original_attention_math():
|
|
||||||
torch.manual_seed(0)
|
|
||||||
batch = 2
|
|
||||||
heads = 3
|
|
||||||
dim_head = 4
|
|
||||||
sequence = 16
|
|
||||||
channels = heads * dim_head
|
|
||||||
q = torch.randn(batch, channels, 1, sequence)
|
|
||||||
k = torch.randn(batch, channels, 1, sequence)
|
|
||||||
v = torch.randn(batch, channels, 1, sequence)
|
|
||||||
|
|
||||||
expected = _original_attention(q, k, v, None, heads, dim_head)
|
|
||||||
|
|
||||||
# split-einsum reorders the float32 reductions vs the reference, so equality
|
|
||||||
# only holds up to rounding; the drift exceeds allclose's default atol on
|
|
||||||
# some BLAS backends (e.g. Linux x86 CI).
|
|
||||||
assert torch.allclose(split_einsum(q, k, v, None, heads, dim_head), expected, atol=1e-6)
|
|
||||||
assert torch.allclose(split_einsum_v2(q, k, v, None, heads, dim_head), expected, atol=1e-6)
|
|
||||||
|
|
||||||
|
|
||||||
def test_split_einsum_v2_chunked_path_matches_original_attention_math():
|
|
||||||
torch.manual_seed(0)
|
|
||||||
batch = 1
|
|
||||||
heads = 2
|
|
||||||
dim_head = 2
|
|
||||||
sequence = 512
|
|
||||||
channels = heads * dim_head
|
|
||||||
q = torch.randn(batch, channels, 1, sequence)
|
|
||||||
k = torch.randn(batch, channels, 1, sequence)
|
|
||||||
v = torch.randn(batch, channels, 1, sequence)
|
|
||||||
|
|
||||||
expected = _original_attention(q, k, v, None, heads, dim_head)
|
|
||||||
|
|
||||||
assert torch.allclose(
|
|
||||||
split_einsum_v2(q, k, v, None, heads, dim_head),
|
|
||||||
expected,
|
|
||||||
atol=1e-6,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def test_apply_attention_implementation_sets_split_processors():
|
|
||||||
unet = RecordingProcessorUNet()
|
|
||||||
|
|
||||||
assert apply_attention_implementation(unet, "ORIGINAL") is unet
|
|
||||||
assert unet.processor is None
|
|
||||||
|
|
||||||
apply_attention_implementation(unet, "SPLIT_EINSUM")
|
|
||||||
assert isinstance(unet.processor, SplitEinsumAttnProcessor)
|
|
||||||
|
|
||||||
apply_attention_implementation(unet, "SPLIT_EINSUM_V2")
|
|
||||||
assert isinstance(unet.processor, SplitEinsumV2AttnProcessor)
|
|
||||||
|
|
||||||
|
|
||||||
class RecordingProcessorUNet:
|
|
||||||
def __init__(self):
|
|
||||||
self.processor = None
|
|
||||||
|
|
||||||
def set_attn_processor(self, processor):
|
|
||||||
self.processor = processor
|
|
||||||
|
|
||||||
|
|
||||||
def _original_attention(q, k, v, mask, heads, dim_head):
|
|
||||||
batch = q.size(0)
|
|
||||||
mh_q = q.view(batch, heads, dim_head, -1)
|
|
||||||
mh_k = k.view(batch, heads, dim_head, -1)
|
|
||||||
mh_v = v.view(batch, heads, dim_head, -1)
|
|
||||||
|
|
||||||
weights = torch.einsum("bhcq,bhck->bhqk", mh_q, mh_k)
|
|
||||||
weights = weights * (dim_head**-0.5)
|
|
||||||
if mask is not None:
|
|
||||||
weights = weights + mask
|
|
||||||
weights = weights.softmax(dim=3)
|
|
||||||
|
|
||||||
attn = torch.einsum("bhqk,bhck->bhcq", weights, mh_v)
|
|
||||||
return attn.contiguous().view(batch, heads * dim_head, 1, -1)
|
|
||||||
Reference in New Issue
Block a user