wehub/apple--ml-stable-diffusion
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:28:46 +08:00
commit 8199cf3c39
90 changed files with 62775 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/pull_request_template.md
+1
View File
@@ -0,0 +1 @@
Thank you for your interest in contributing to Core ML Stable Diffusion! Please review [CONTRIBUTING.md](../CONTRIBUTING.md) first. We appreciate your interest in the project!
.gitignore
+146
View File
@@ -0,0 +1,146 @@
*~
# Swift Package
.DS_Store
/.build
/Packages
/*.xcodeproj
.swiftpm
.vscode
.*.sw?
*.docc-build
*.vs
Package.resolved
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# macOS filesystem
*.DS_Store
ACKNOWLEDGEMENTS
+555
View File
@@ -0,0 +1,555 @@
Acknowledgements
Portions of this software may utilize the following copyrighted
material, the use of which is hereby acknowledged.
_____________________
The Hugging Face team (diffusers)
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
The Hugging Face team (transformers)
Copyright 2018- The Hugging Face team. All rights reserved.
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Facebook, Inc (PyTorch)
From PyTorch:
Copyright (c) 2016- Facebook, Inc (Adam Paszke)
Copyright (c) 2014- Facebook, Inc (Soumith Chintala)
Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu)
Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
Copyright (c) 2011-2013 NYU (Clement Farabet)
Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
Copyright (c) 2006 Idiap Research Institute (Samy Bengio)
Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
From Caffe2:
Copyright (c) 2016-present, Facebook Inc. All rights reserved.
All contributions by Facebook:
Copyright (c) 2016 Facebook Inc.
All contributions by Google:
Copyright (c) 2015 Google Inc.
All rights reserved.
All contributions by Yangqing Jia:
Copyright (c) 2015 Yangqing Jia
All rights reserved.
All contributions by Kakao Brain:
Copyright 2019-2020 Kakao Brain
All contributions by Cruise LLC:
Copyright (c) 2022 Cruise LLC.
All rights reserved.
All contributions from Caffe:
Copyright(c) 2013, 2014, 2015, the respective contributors
All rights reserved.
All other contributions:
Copyright(c) 2015, 2016 the respective contributors
All rights reserved.
Caffe2 uses a copyright model similar to Caffe: each contributor holds
copyright over their contributions to Caffe2. The project versioning records
all such contribution and copyright details. If a contributor wants to further
mark their specific copyright on a particular contribution, they should
indicate their copyright solely in the commit message of the change when it is
committed.
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
and IDIAP Research Institute nor the names of its contributors may be
used to endorse or promote products derived from this software without
specific prior written permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
POSSIBILITY OF SUCH DAMAGE.
NumPy (RandomKit 1.3)
Copyright (c) 2003-2005, Jean-Sebastien Roy (js@jeannot.org)
The rk_random and rk_seed functions algorithms and the original design of
the Mersenne Twister RNG:
Copyright (C) 1997 - 2002, Makoto Matsumoto and Takuji Nishimura,
All rights reserved.
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:
1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.
3. The names of its contributors may not be used to endorse or promote
products derived from this software without specific prior written
permission.
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
Original algorithm for the implementation of rk_interval function from
Richard J. Wagner's implementation of the Mersenne Twister RNG, optimised by
Magnus Jonsson.
Constants used in the rk_double implementation by Isaku Wada.
Permission is hereby granted, free of charge, to any person obtaining a
copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:
The above copyright notice and this permission notice shall be included
in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
CODE_OF_CONDUCT.md
+71
View File
@@ -0,0 +1,71 @@
# Code of Conduct
## Our Pledge
In the interest of fostering an open and welcoming environment, we as
contributors and maintainers pledge to making participation in our project and
our community a harassment-free experience for everyone, regardless of age, body
size, disability, ethnicity, sex characteristics, gender identity and expression,
level of experience, education, socio-economic status, nationality, personal
appearance, race, religion, or sexual identity and orientation.
## Our Standards
Examples of behavior that contributes to creating a positive environment
include:
* Using welcoming and inclusive language
* Being respectful of differing viewpoints and experiences
* Gracefully accepting constructive criticism
* Focusing on what is best for the community
* Showing empathy towards other community members
Examples of unacceptable behavior by participants include:
* The use of sexualized language or imagery and unwelcome sexual attention or
advances
* Trolling, insulting/derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or electronic
address, without explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Our Responsibilities
Project maintainers are responsible for clarifying the standards of acceptable
behavior and are expected to take appropriate and fair corrective action in
response to any instances of unacceptable behavior.
Project maintainers have the right and responsibility to remove, edit, or
reject comments, commits, code, wiki edits, issues, and other contributions
that are not aligned to this Code of Conduct, or to ban temporarily or
permanently any contributor for other behaviors that they deem inappropriate,
threatening, offensive, or harmful.
## Scope
This Code of Conduct applies within all project spaces, and it also applies when
an individual is representing the project or its community in public spaces.
Examples of representing a project or community include using an official
project e-mail address, posting via an official social media account, or acting
as an appointed representative at an online or offline event. Representation of
a project may be further defined and clarified by project maintainers.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported by contacting the open source team at [opensource-conduct@group.apple.com](mailto:opensource-conduct@group.apple.com). All
complaints will be reviewed and investigated and will result in a response that
is deemed necessary and appropriate to the circumstances. The project team is
obligated to maintain confidentiality with regard to the reporter of an incident.
Further details of specific enforcement policies may be posted separately.
Project maintainers who do not follow or enforce the Code of Conduct in good
faith may face temporary or permanent repercussions as determined by other
members of the project's leadership.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org), version 1.4,
available at [https://www.contributor-covenant.org/version/1/4/code-of-conduct.html](https://www.contributor-covenant.org/version/1/4/code-of-conduct.html)
CONTRIBUTING.md
+12
View File
@@ -0,0 +1,12 @@
# Contribution Guide
Thank you for your interest in contributing to Core ML Stable Diffusion! This project was released for system demonstration purposes and there are limited plans for future development of the repository. While we welcome new pull requests and issues please note that our response may be limited.
## Submitting a Pull Request
The project is licensed under the MIT license. By submitting a pull request, you represent that you have the right to license your contribution to Apple and the community, and agree by submitting the patch that your contributions are licensed under the MIT license.
## Code of Conduct
We ask that all community members read and observe our [Code of Conduct](CODE_OF_CONDUCT.md).
LICENSE.md
+9
View File
@@ -0,0 +1,9 @@
MIT License
Copyright (c) 2024 Apple Inc.
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
Package.swift
+46
View File
@@ -0,0 +1,46 @@
// swift-tools-version: 5.8
// The swift-tools-version declares the minimum version of Swift required to build this package.
import PackageDescription
let package = Package(
name: "stable-diffusion",
platforms: [
.macOS(.v13),
.iOS(.v16),
],
products: [
.library(
name: "StableDiffusion",
targets: ["StableDiffusion"]),
.executable(
name: "StableDiffusionSample",
targets: ["StableDiffusionCLI"])
],
dependencies: [
.package(url: "https://github.com/apple/swift-argument-parser.git", from: "1.2.3"),
.package(url: "https://github.com/huggingface/swift-transformers.git", exact: "0.1.8"),
],
targets: [
.target(
name: "StableDiffusion",
dependencies: [
.product(name: "Transformers", package: "swift-transformers"),
],
path: "swift/StableDiffusion"),
.executableTarget(
name: "StableDiffusionCLI",
dependencies: [
"StableDiffusion",
.product(name: "ArgumentParser", package: "swift-argument-parser")],
path: "swift/StableDiffusionCLI"),
.testTarget(
name: "StableDiffusionTests",
dependencies: ["StableDiffusion"],
path: "swift/StableDiffusionTests",
resources: [
.copy("Resources/vocab.json"),
.copy("Resources/merges.txt")
]),
]
)
README.md
+808
View File
@@ -0,0 +1,808 @@
# Core ML Stable Diffusion
Run Stable Diffusion on Apple Silicon with Core ML
[\[Blog Post\]](https://machinelearning.apple.com/research/stable-diffusion-coreml-apple-silicon) [\[BibTeX\]](#bibtex)
This repository comprises:
- `python_coreml_stable_diffusion`, a Python package for converting PyTorch models to Core ML format and performing image generation with Hugging Face [diffusers](https://github.com/huggingface/diffusers) in Python
- `StableDiffusion`, a Swift package that developers can add to their Xcode projects as a dependency to deploy image generation capabilities in their apps. The Swift package relies on the Core ML model files generated by `python_coreml_stable_diffusion`
If you run into issues during installation or runtime, please refer to the [FAQ](#faq) section. Please refer to the [System Requirements](#system-requirements) section before getting started.
<img src="assets/readme_reel.png">
## <a name="system-requirements"></a> System Requirements
<details>
<summary> Details (Click to expand) </summary>
Model Conversion:
macOS | Python | coremltools |
:------:|:------:|:-----------:|
13.1 | 3.8 | 7.0 |
Project Build:
macOS | Xcode | Swift |
:------:|:-----:|:-----:|
13.1 | 14.3 | 5.8 |
Target Device Runtime:
macOS | iPadOS, iOS |
:------:|:-----------:|
13.1 | 16.2 |
Target Device Runtime ([With Memory Improvements](#compression-6-bits-and-higher)):
macOS | iPadOS, iOS |
:------:|:-----------:|
14.0 | 17.0 |
Target Device Hardware Generation:
Mac | iPad | iPhone |
:------:|:-------:|:-------:|
M1 | M1 | A14 |
</details>
## <a name="performance-benchmark"></a> Performance Benchmarks
<details>
<summary> Details (Click to expand) </summary>
[`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) (512x512)
| Device | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
| iPhone 12 Mini | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 18.5* | 1.44 |
| iPhone 12 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 15.4 | 1.45 |
| iPhone 13 | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 10.8* | 2.53 |
| iPhone 13 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 10.4 | 2.55 |
| iPhone 14 | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 8.6 | 2.57 |
| iPhone 14 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 7.9 | 2.69 |
| iPad Pro (M1) | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 11.2 | 2.19 |
| iPad Pro (M2) | `CPU_AND_NE` | `SPLIT_EINSUM_V2` | 7.0 | 3.07 |
<details>
<summary> Details (Click to expand) </summary>
- This benchmark was conducted by Apple and Hugging Face using public beta versions of iOS 17.0, iPadOS 17.0 and macOS 14.0 Seed 8 in August 2023.
- The performance data was collected using the `benchmark` branch of the [Diffusers app](https://github.com/huggingface/swift-coreml-diffusers)
- Swift code is not fully optimized, introducing up to ~10% overhead unrelated to Core ML model execution.
- The median latency value across 5 back-to-back end-to-end executions are reported
- The image generation procedure follows the standard configuration: 20 inference steps, 512x512 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
- The actual prompt length does not impact performance because the Core ML model is converted with a static shape that computes the forward pass for all of the 77 elements (`tokenizer.model_max_length`) in the text token sequence regardless of the actual length of the input text.
- Weights are compressed to 6 bit precision. Please refer to [this section](#compression-6-bits-and-higher) for details.
- Activations are in float16 precision for both the GPU and the Neural Engine.
- `*` indicates that the [reduceMemory](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L91) option was enabled which loads and unloads models just-in-time to avoid memory shortage. This added up to 2 seconds to the end-to-end latency.
- In the benchmark table, we report the best performing `--compute-unit` and `--attention-implementation` values per device. The former does not modify the Core ML model and can be applied during runtime. The latter modifies the Core ML model. Note that the best performing compute unit is model version and hardware-specific.
- Note that the performance optimizations in this repository (e.g. `--attention-implementation`) are generally applicable to Transformers and not customized to Stable Diffusion. Better performance may be observed upon custom kernel tuning. Therefore, these numbers do not represent **peak** HW capability.
- Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
- Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state.
</details>
[`stabilityai/stable-diffusion-xl-base-1.0-ios`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-ios) (768x768)
| Device | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
| iPhone 12 Pro | `CPU_AND_NE` | `SPLIT_EINSUM` | 116* | 0.50 |
| iPhone 13 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM` | 86* | 0.68 |
| iPhone 14 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM` | 77* | 0.83 |
| iPhone 15 Pro Max | `CPU_AND_NE` | `SPLIT_EINSUM` | 31 | 0.85 |
| iPad Pro (M1) | `CPU_AND_NE` | `SPLIT_EINSUM` | 36 | 0.69 |
| iPad Pro (M2) | `CPU_AND_NE` | `SPLIT_EINSUM` | 27 | 0.98 |
<details>
<summary> Details (Click to expand) </summary>
- This benchmark was conducted by Apple and Hugging Face using iOS 17.0.2 and iPadOS 17.0.2 in September 2023.
- The performance data was collected using the `benchmark` branch of the [Diffusers app](https://github.com/huggingface/swift-coreml-diffusers)
- The median latency value across 5 back-to-back end-to-end executions are reported
- The image generation procedure follows this configuration: 20 inference steps, 768x768 output image resolution, 77 text token sequence length, classifier-free guidance (batch size of 2 for unet).
- `Unet.mlmodelc` is compressed to 4.04 bit precision following the [Mixed-Bit Palettization](#compression-lower-than-6-bits) algorithm recipe published [here](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization/blob/main/recipes/stabilityai-stable-diffusion-xl-base-1.0_palettization_recipe.json)
- All models except for `Unet.mlmodelc` are compressed to 16 bit precision
- [madebyollin/sdxl-vae-fp16-fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) by [@madebyollin](https://github.com/madebyollin) was used as the source PyTorch model for `VAEDecoder.mlmodelc` in order to enable float16 weight and activation quantization for the VAE model.
- `--attention-implementation SPLIT_EINSUM` is chosen in lieu of `SPLIT_EINSUM_V2` due to the prohibitively long compilation time of the latter
- `*` indicates that the [reduceMemory](https://github.com/apple/ml-stable-diffusion/blob/main/swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift#L91) option was enabled which loads and unloads models just-in-time to avoid memory shortage. This added significant overhead to the end-to-end latency. Note that end-to-end latency difference between `iPad Pro (M1)` and `iPhone 13 Pro Max` despite identical diffusion speed.
- The actual prompt length does not impact performance because the Core ML model is converted with a static shape that computes the forward pass for all of the 77 elements (`tokenizer.model_max_length`) in the text token sequence regardless of the actual length of the input text.
- In the benchmark table, we report the best performing `--compute-unit` and `--attention-implementation` values per device. The former does not modify the Core ML model and can be applied during runtime. The latter modifies the Core ML model. Note that the best performing compute unit is model version and hardware-specific.
- Note that the performance optimizations in this repository (e.g. `--attention-implementation`) are generally applicable to Transformers and not customized to Stable Diffusion. Better performance may be observed upon custom kernel tuning. Therefore, these numbers do not represent **peak** HW capability.
- Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
- Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state.
</details>
[`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base) (1024x1024)
| Device | `--compute-unit`| `--attention-implementation` | End-to-End Latency (s) | Diffusion Speed (iter/s) |
| --------------------- | --------------- | ---------------------------- | ---------------------- | ------------------------ |
| MacBook Pro (M1 Max) | `CPU_AND_GPU` | `ORIGINAL` | 46 | 0.46 |
| MacBook Pro (M2 Max) | `CPU_AND_GPU` | `ORIGINAL` | 37 | 0.57 |
| Mac Studio (M1 Ultra) | `CPU_AND_GPU` | `ORIGINAL` | 25 | 0.89 |
| Mac Studio (M2 Ultra) | `CPU_AND_GPU` | `ORIGINAL` | 20 | 1.11 |
<details>
<summary> Details (Click to expand) </summary>
- This benchmark was conducted by Apple and Hugging Face using public beta versions of iOS 17.0, iPadOS 17.0 and macOS 14.0 in July 2023.
- The performance data was collected by running the `StableDiffusion` Swift pipeline.
- The median latency value across 3 back-to-back end-to-end executions are reported
- The image generation procedure follows the standard configuration: 20 inference steps, 1024x1024 output image resolution, classifier-free guidance (batch size of 2 for unet).
- Weights and activations are in float16 precision
- Performance may vary across different versions of Stable Diffusion due to architecture changes in the model itself. Each reported number is specific to the model version mentioned in that context.
- Performance may vary due to factors like increased system load from other applications or suboptimal device thermal state. Given these factors, we do not report sub-second variance in latency.
</details>
</details>
## <a name="compression-6-bits-and-higher"></a> Weight Compression (6-bits and higher)
<details>
<summary> Details (Click to expand) </summary>
coremltools-7.0 supports advanced weight compression techniques for [pruning](https://coremltools.readme.io/v7.0/docs/pruning), [palettization](https://coremltools.readme.io/v7.0/docs/palettization-overview) and [linear 8-bit quantization](https://coremltools.readme.io/v7.0/docs/quantization-aware-training). For these techniques, `coremltools.optimize.torch.*` includes APIs that require fine-tuning to maintain accuracy at higher compression rates whereas `coremltools.optimize.coreml.*` includes APIs that are applied post-training and are data-free.
We demonstrate how data-free [post-training palettization](https://coremltools.readme.io/v7.0/docs/post-training-palettization) implemented in `coremltools.optimize.coreml.palettize_weights` enables us to achieve greatly improved performance for Stable Diffusion on mobile devices. This API implements the [Fast Exact k-Means](https://arxiv.org/abs/1701.07204) algorithm for optimal weight clustering which yields more accurate palettes. Using `--quantize-nbits {2,4,6,8}` during [conversion](#converting-models-to-coreml) is going to apply this compression to the unet and text_encoder models.
For best results, we recommend [training-time palettization](https://coremltools.readme.io/v7.0/docs/training-time-palettization): `coremltools.optimize.torch.palettization.DKMPalettizer` if fine-tuning your model is feasible. This API implements the [Differentiable k-Means (DKM)](https://machinelearning.apple.com/research/differentiable-k-means) learned palettization algorithm. In this exercise, we stick to post-training palettization for the sake of simplicity and ease of reproducibility.
The Neural Engine is capable of accelerating models with low-bit palettization: 1, 2, 4, 6 or 8 bits. With iOS 17 and macOS 14, compressed weights for Core ML models can be just-in-time decompressed during runtime (as opposed to ahead-of-time decompression upon load) to match the precision of activation tensors. This yields significant memory savings and enables models to run on devices with smaller RAM (e.g. iPhone 12 Mini). In addition, compressed weights are faster to fetch from memory which reduces the latency of memory bandwidth-bound layers. The just-in-time decompression behavior depends on the compute unit, layer type and hardware generation.
| Weight Precision | `--compute-unit` | [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base) generating *"a high quality photo of a surfing dog"* |
| :---------------:| :----------------: | ------------------------------------------------------ |
| 6-bit | cpuAndNeuralEngine | <img src="assets/palette6_cpuandne_readmereel.png"> |
| 16-bit | cpuAndNeuralEngine | <img src="assets/float16_cpuandne_readmereel.png"> |
| 16-bit | cpuAndGPU | <img src="assets/float16_gpu_readmereel.png"> |
Note that there are minor differences across 16-bit (float16) and 6-bit results. These differences are comparable to the differences across float16 and float32 or differences across compute units as exemplified above. We recommend a minimum of 6 bits for palettizing Stable Diffusion. Smaller number of bits (1, 2 and 4) will require either fine-tuning or advanced palettization techniques such as [MBP](#compression-lower-than-6-bits).
Resources:
- [Core ML Tools Docs: Optimizing Models](https://coremltools.readme.io/v7.0/docs/optimizing-models)
- [WWDC23 Session Video: Use Core ML Tools for machine learning model compression](https://developer.apple.com/videos/play/wwdc2023/10047)
</details>
## <a name="compression-lower-than-6-bits"></a> Advanced Weight Compression (Lower than 6-bits)
<details>
<summary> Details (Click to expand) </summary>
This section describes an advanced compression algorithm called [Mixed-Bit Palettization (MBP)](https://huggingface.co/blog/stable-diffusion-xl-coreml#what-is-mixed-bit-palettization) built on top of the [Post-Training Weight Palettization tools](https://apple.github.io/coremltools/docs-guides/source/post-training-palettization.html) and using the [Weights Metadata API](https://apple.github.io/coremltools/docs-guides/source/mlmodel-utilities.html#get-weights-metadata) from [coremltools](https://github.com/apple/coremltools).
MBP builds a per-layer "palettization recipe" by picking a suitable number of bits among the Neural Engine supported bit-widths of 1, 2, 4, 6 and 8 in order to achieve the minimum average bit-width while maintaining a desired level of signal strength. The signal strength is measured by comparing the compressed model's output to that of the original float16 model. Given the same random seed and text prompts, PSNR between denoised latents is computed. The compression rate will depend on the model version as well as the tolerance for signal loss (drop in PSNR) since this algorithm is adaptive.
| 3.41-bit | 4.50-bit | 6.55-bit | 16-bit (original) |
| :-------:| :-------:| :-------:| :----------------:|
| <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_3.41-bits.png"> | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_4.50-bits.png"> | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_6.55-bits.png"> | <img src="assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_float16_original.png"> |
For example, the original float16 [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model has an ~82 dB signal strength. Naively applying [linear 8-bit quantization](https://coremltools.readme.io/docs/data-free-quantization) to the Unet model drops the signal to ~65 dB. Instead, applying MBP yields an average of 2.81-bits quantization while maintaining a signal strength of ~67 dB. This technique generally yields better results compared to using `--quantize-nbits` during model conversion but requires a "pre-analysis" run that takes up to a few hours on a single GPU (`mps` or `cuda`).
Here is the signal strength (PSNR in dB) versus model size reduction (% of float16 size) for `stabilityai/stable-diffusion-xl-base-1.0`. The `{1,2,4,6,8}-bit` curves are generated by progressively palettizing more layers using a palette with fixed number of bits. The layers were ordered in ascending order of their isolated impact to end-to-end signal strength so the cumulative compression's impact is delayed as much as possible. The mixed-bit curve is based on falling back to a higher number of bits as soon as a layer's isolated impact to end-to-end signal integrity drops below a threshold. Note that all curves based on palettization outperform linear 8-bit quantization at the same model size except for 1-bit.
<img src="assets/mbp/stabilityai_stable-diffusion-xl-base-1.0_psnr_vs_size.png" width="640">
Here are the steps for applying this technique on another model version:
**Step 1:** Run the pre-analysis script to generate "recipes" with varying signal strength:
```python
python -m python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis --model-version <model-version> -o <output-dir>
```
For popular base models, you may find the pre-computed pre-analysis results [here](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization/tree/main/recipes). Fine-tuned models models are likely to honor the recipes of their corresponding base models but this is untested.
**Step 2:** The resulting JSON file from Step 1 will list "baselines", e.g.:
```json
{
"model_version": "stabilityai/stable-diffusion-xl-base-1.0",
"baselines": {
"original": 82.2,
"linear_8bit": 66.025,
"recipe_6.55_bit_mixedpalette": 79.9,
"recipe_5.52_bit_mixedpalette": 78.2,
"recipe_4.89_bit_mixedpalette": 76.8,
"recipe_4.41_bit_mixedpalette": 75.5,
"recipe_4.04_bit_mixedpalette": 73.2,
"recipe_3.67_bit_mixedpalette": 72.2,
"recipe_3.32_bit_mixedpalette": 71.4,
"recipe_3.19_bit_mixedpalette": 70.4,
"recipe_3.08_bit_mixedpalette": 69.6,
"recipe_2.98_bit_mixedpalette": 68.6,
"recipe_2.90_bit_mixedpalette": 67.8,
"recipe_2.83_bit_mixedpalette": 67.0,
"recipe_2.71_bit_mixedpalette": 66.3
},
}
```
Among these baselines, select a recipe based on your desired signal strength. We recommend palettizing to ~4 bits depending on the use case even if the signal integrity for lower bit values are higher than the linear 8-bit quantization baseline.
Finally, apply the selected recipe to the float16 Core ML model as follows:
```python
python -m python_coreml_stable_diffusion.mixed_bit_compression_apply --mlpackage-path <path-to-float16-unet-mlpackage> -o <output-dir> --pre-analysis-json-path <path-to--pre-analysis-json> --selected-recipe <selected-recipe-string-key>
```
An example `<selected-recipe-string-key>` would be `"recipe_4.50_bit_mixedpalette"` which achieves an average of 4.50-bits compression (compressed from ~5.2GB to ~1.46GB for SDXL). Please note that signal strength does not directly map to image-text alignment. Always verify that your MBP-compressed model variant is accurately generating images for your test prompts.
</details>
## <a name="activation-quant"></a> Activation Quantization
<details>
<summary> Details (Click to expand) </summary>
On newer hardware with A17 Pro or M4 chips, such as the iPhone 15 Pro, quantizing both activations and weight to int8 can leverage optimized compute on the Neural Engine which can be used to improve runtime latency in compute-bound models.
In this section, we demonstrate how to apply [Post Training Activation Quantization](https://apple.github.io/coremltools/docs-guides/source/opt-quantization-algos.html#post-training-data-calibration-activation-quantization), using calibration data, on Stable Diffusion UNet model.
Similar to Mixed-Bit Palettization (MBP) described [above](#a-namecompression-lower-than-6-bitsa-advanced-weight-compression-lower-than-6-bits), first, a per-layer analysis is run to determine which intermediate activations are more sensitive to 8-bit compression.
Less sensitive layers are weight and activation quantized (W8A8), whereas more sensitive layers are only weight quantized (W8A16).
Here are the steps for applying this technique:
**Step 1:** Generate calibration data
```python
python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --generate-calibration-data -o <output-dir>
```
A set of calibration text prompts are run through StableDiffusionPipeline and UNet model inputs are recorded and stored as pickle files in `calibration_data_<model-version>` folder inside specified output directory.
**Step 2:** Run layer-wise sensitivity analysis
```python
python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --layerwise-sensitivity --calibration-nsamples <num-samples> -o <output-dir>
```
This will run the analysis on all Convolutional and Attention (Einsum) modules in the model.
For each module, a compressed version is generated by quantizing only that layer’s weights and activations.
Then the PSNR between the outputs of the compressed and original model is calculated, using the same random seed and text prompts.
This analysis takes up to a few hours on a single GPU (cuda). The number of calibration samples used to quantize the model can be reduced to speed up the process.
The resulting JSON file looks like this:
```json
{
"conv": {
"conv_in": 30.74,
"down_blocks.0.attentions.0.proj_in": 38.93,
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q": 48.15,
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k": 50.13,
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v": 45.70,
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0": 39.56,
...
},
"einsum": {
"down_blocks.0.attentions.0.transformer_blocks.0.attn1.einsum": 25.34,
"down_blocks.0.attentions.0.transformer_blocks.0.attn2.einsum": 31.76,
"down_blocks.0.attentions.1.transformer_blocks.0.attn1.einsum": 23.40,
"down_blocks.0.attentions.1.transformer_blocks.0.attn2.einsum": 31.56,
...
},
"model_version": "stabilityai/stable-diffusion-2-1-base"
}
```
**Step 3:** Generate quantized model
Using calibration data and layer-wise sensitivity the quantized CoreML model can be generated as follows:
```python
python -m python_coreml_stable_diffusion.activation_quantization --model-version <model-version> --quantize-pytorch --conv-psnr 38 --attn-psnr 26 -o <output-dir>
```
The PSNR thresholds determine which layers will be activation quantized. This number can be tuned to trade-off between output quality and inference latency.
</details>
## <a name="using-stable-diffusion-3"></a> Using Stable Diffusion 3
<details>
<summary> Details (Click to expand) </summary>
### Model Conversion
Stable Diffusion 3 uses some new and some old models to run. For the text encoders, the conversion can be done using a similar command as before with the `--sd3-version` flag.
```bash
python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --sd3-version -o <output-dir>
```
For the new models (MMDiT, a new VAE with 16 channels, and the T5 text encoder), there are a number of new CLI flags that utilize the [DiffusionKit](https://www.github.com/argmaxinc/DiffusionKit) repo:
- `--sd3-version`: Indicates to the converter to treat this as a Stable Diffusion 3 model
- `--convert-mmdit`: Convert the MMDiT model
- `--convert-vae-decoder`: Convert the new VAE model (this will use the 16 channel version if --sd3-version is set)
- `--include-t5`: Downloads and includes a pre-converted T5 text encoder in the conversion
e.g.:
```bash
python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-vae-decoder --convert-mmdit --include-t5 --sd3-version -o <output-dir>
```
To convert the full pipeline with at 1024x1024 resolution, the following command may be used:
```bash
python -m python_coreml_stable_diffusion.torch2coreml --model-version stabilityai/stable-diffusion-3-medium --bundle-resources-for-swift-cli --convert-text-encoder --convert-vae-decoder --convert-mmdit --include-t5 --sd3-version --latent-h 128 --latent-w 128 -o <output-dir>
```
Keep in mind that the MMDiT model is quite large and will require increasingly more memory and time to convert as the latent resolution increases.
Also note that currently the MMDiT model requires fp32 and therefore only supports `CPU_AND_GPU` compute units and `ORIGINAL` attention implementation (the default for this pipeline).
### Swift Inference
Swift inference for Stable Diffusion 3 is similar to the previous versions. The only difference is that the `--sd3` flag should be used to indicate that the model is a Stable Diffusion 3 model.
```bash
swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units cpuAndGPU --sd3
```
</details>
## <a name="using-stable-diffusion-xl"></a> Using Stable Diffusion XL
<details>
<summary> Details (Click to expand) </summary>
### Model Conversion
e.g.:
```bash
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-vae-decoder --convert-text-encoder --xl-version --model-version stabilityai/stable-diffusion-xl-base-1.0 --refiner-version stabilityai/stable-diffusion-xl-refiner-1.0 --bundle-resources-for-swift-cli --attention-implementation {ORIGINAL,SPLIT_EINSUM} -o <output-dir>
```
- `--xl-version`: Additional argument to pass to the conversion script when specifying an XL model
- `--refiner-version`: Additional argument to pass to the conversion script when specifying an XL refiner model, required for ["Ensemble of Expert Denoisers"](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#1-ensemble-of-expert-denoisers) inference.
- `--attention-implementation`: `ORIGINAL` is recommended for `cpuAndGPU` for deployment on Mac
- `--attention-implementation`: `SPLIT_EINSUM` is recommended for `cpuAndNeuralEngine` for deployment on iPhone & iPad
- `--attention-implementation`: `SPLIT_EINSUM_V2` is not recommended for Stable Diffusion XL because of prohibitively long compilation time
- **Tip:** Adding `--latent-h 96 --latent-w 96` is recommended for iOS and iPadOS deployment which leads to 768x768 generation as opposed to the default 1024x1024.
- **Tip:** Due to known float16 overflow issues in the original Stable Diffusion XL VAE, [the model conversion script enforces float32 precision](https://github.com/apple/ml-stable-diffusion/blob/main/python_coreml_stable_diffusion/torch2coreml.py#L486). Using a custom VAE version such as [madebyollin/sdxl-vae-fp16-fix](https://huggingface.co/madebyollin/sdxl-vae-fp16-fix) by [@madebyollin](https://github.com/madebyollin) via `--custom-vae-version madebyollin/sdxl-vae-fp16-fix` will restore the default float16 precision for VAE.
### Swift Inference
```bash
swift run StableDiffusionSample <prompt> --resource-path <output-mlpackages-directory/Resources> --output-path <output-dir> --compute-units {cpuAndGPU,cpuAndNeuralEngine} --xl
```
- Only the `base` model is required, `refiner` model is optional and will be used by default if provided in the resource directory
- ControlNet for XL is not yet supported
### Python Inference
```bash
python -m python_coreml_stable_diffusion.pipeline --prompt <prompt> --compute-unit {CPU_AND_GPU,CPU_AND_NE} -o <output-dir> -i <output-mlpackages-directory/Resources> --model-version stabilityai/stable-diffusion-xl-base-1.0
```
- `refiner` model is not yet supported
- ControlNet for XL is not yet supported
</details>
## <a name="using-controlnet"></a> Using ControlNet
<details>
<summary> Details (Click to expand) </summary>
Example results using the prompt *"a high quality photo of a surfing dog"* conditioned on the scribble (leftmost):
<img src="assets/controlnet_readme_reel.png">
[ControlNet](https://huggingface.co/lllyasviel/ControlNet) allows users to condition image generation with Stable Diffusion on signals such as edge maps, depth maps, segmentation maps, scribbles and pose. Thanks to [@ryu38's contribution](https://github.com/apple/ml-stable-diffusion/pull/153), both the Python CLI and the Swift package support ControlNet models. Please refer to [this section](#converting-models-to-coreml) for details on setting up Stable Diffusion with ControlNet.
Note that ControlNet is not yet supported for Stable Diffusion XL.
</details>
## <a name="system-multilingual-text-encoder"></a> Using the System Multilingual Text Encoder
<details>
<summary> Details (Click to expand) </summary>
With iOS 17 and macOS 14, `NaturalLanguage` framework introduced the [NLContextualEmbedding](https://developer.apple.com/documentation/naturallanguage/nlcontextualembedding) which provides Transformer-based textual embeddings for Latin (20 languages), Cyrillic (4 languages) and CJK (3 languages) scripts. The WWDC23 session titled [Explore Natural Language multilingual models](https://developer.apple.com/videos/play/wwdc2023/10042) demonstrated how this powerful new model can be used by developers to train downstream tasks such as multilingual image generation with Stable Diffusion.
The code to reproduce this demo workflow is made available in this repository. There are several ways in which this workflow can be implemented. Here is an example:
**Step 1:** Curate an image-text dataset with the desired languages.
**Step 2:** Pre-compute the NLContextualEmbedding values and replace the text strings with these embedding vectors in your dataset.
**Step 3:** Fine-tune a base model from Hugging Face Hub that is compatible with the [StableDiffusionPipeline](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/overview) by using your new dataset and replacing the default text_encoder with your pre-computed NLContextualEmbedding values.
**Step 4:** In order to be able to swap the text_encoder of a base model without training new layers, the base model's `text_encoder.hidden_size` must match that of NLContextualEmbedding. If it doesn't, you will need to train a linear projection layer to map between the two dimensionalities. After fine-tuning, this linear layer should be converted to CoreML as follows:
```shell
python -m python_coreml_stable_diffusion.multilingual_projection --input-path <path-to-projection-torchscript> --output-dir <output-dir>
```
The command above will yield a `MultilingualTextEncoderProjection.mlmodelc` file under `--output-dir` and this should be colocated with the rest of the Core ML model assets that were generated through `--bundle-resources-for-swift-cli`.
**Step 5:** The multilingual system text encoder can now be invoked by setting `useMultilingualTextEncoder` to true when initializing a pipeline or setting `--use-multilingual-text-encoder` in the CLI. Note that the model assets are distributed over-the-air so the first invocation will trigger asset downloads which is less than 100MB.
Resources:
- [WWDC23 Session Video: Explore Natural Language multilingual models](https://developer.apple.com/videos/play/wwdc2023/10042)
- [NLContextualEmbedding API Documentation](https://developer.apple.com/documentation/naturallanguage/nlcontextualembedding)
</details>
## <a name="using-converted-weights"></a> Using Ready-made Core ML Models from Hugging Face Hub
<details>
<summary> Click to expand </summary>
🤗 Hugging Face ran the [conversion procedure](#converting-models-to-coreml) on the following models and made the Core ML weights publicly available on the Hub. If you would like to convert a version of Stable Diffusion that is not already available on the Hub, please refer to the [Converting Models to Core ML](#converting-models-to-coreml).
* 6-bit quantized models (suitable for iOS 17 and macOS 14):
- [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-1-4-palettized)
- [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5-palettized)
- [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base-palettized)
- [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base-palettized)
* Mixed-bit quantized models
- [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-mixed-bit-palettization)
- [`stabilityai/stable-diffusion-xl-base-1.0-ios`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-ios)
* Uncompressed models:
- [`CompVis/stable-diffusion-v1-4`](https://huggingface.co/apple/coreml-stable-diffusion-v1-4)
- [`runwayml/stable-diffusion-v1-5`](https://huggingface.co/apple/coreml-stable-diffusion-v1-5)
- [`stabilityai/stable-diffusion-2-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-base)
- [`stabilityai/stable-diffusion-2-1-base`](https://huggingface.co/apple/coreml-stable-diffusion-2-1-base)
- [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base)
- [`stabilityai/stable-diffusion-xl-{base+refiner}-1.0`](https://huggingface.co/apple/coreml-stable-diffusion-xl-base-with-refiner)
- [`stabilityai/stable-diffusion-3-medium`](https://huggingface.co/stabilityai/stable-diffusion-3-medium)
If you want to use any of those models you may download the weights and proceed to [generate images with Python](#image-generation-with-python) or [Swift](#image-generation-with-swift).
There are several variants in each model repository. You may clone the whole repos using `git` and `git lfs` to download all variants, or selectively download the ones you need.
To clone the repos using `git`, please follow this process:
**Step 1:** Install the `git lfs` extension for your system.
`git lfs` stores large files outside the main git repo, and it downloads them from the appropriate server after you clone or checkout. It is available in most package managers, check [the installation page](https://git-lfs.com) for details.
**Step 2:** Enable `git lfs` by running this command once:
```bash
git lfs install
```
**Step 3:** Use `git clone` to download a copy of the repo that includes all model variants. For Stable Diffusion version 1.4, you'd issue the following command in your terminal:
```bash
git clone https://huggingface.co/apple/coreml-stable-diffusion-v1-4
```
If you prefer to download specific variants instead of cloning the repos, you can use the `huggingface_hub` Python library. For example, to do generation in Python using the `ORIGINAL` attention implementation (read [this section](#converting-models-to-coreml) for details), you could use the following helper code:
```Python
from huggingface_hub import snapshot_download
from pathlib import Path
repo_id = "apple/coreml-stable-diffusion-v1-4"
variant = "original/packages"
model_path = Path("./models") / (repo_id.split("/")[-1] + "_" + variant.replace("/", "_"))
snapshot_download(repo_id, allow_patterns=f"{variant}/*", local_dir=model_path, local_dir_use_symlinks=False)
print(f"Model downloaded at {model_path}")
```
`model_path` would be the path in your local filesystem where the checkpoint was saved. Please, refer to [this post](https://huggingface.co/blog/diffusers-coreml) for additional details.
</details>
## <a name="converting-models-to-coreml"></a> Converting Models to Core ML
<details>
<summary> Click to expand </summary>
**Step 1:** Create a Python environment and install dependencies:
```bash
conda create -n coreml_stable_diffusion python=3.8 -y
conda activate coreml_stable_diffusion
cd /path/to/cloned/ml-stable-diffusion/repository
pip install -e .
```
**Step 2:** Log in to or register for your [Hugging Face account](https://huggingface.co), generate a [User Access Token](https://huggingface.co/settings/tokens) and use this token to set up Hugging Face API access by running `huggingface-cli login` in a Terminal window.
**Step 3:** Navigate to the version of Stable Diffusion that you would like to use on [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion) and accept its Terms of Use. The default model version is [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4). The model version may be changed by the user as described in the next step.
**Step 4:** Execute the following command from the Terminal to generate Core ML model files (`.mlpackage`)
```shell
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --convert-text-encoder --convert-vae-decoder --convert-safety-checker --model-version <model-version-string-from-hub> -o <output-mlpackages-directory>
```
**WARNING:** This command will download several GB worth of PyTorch checkpoints from Hugging Face. Please ensure that you are on Wi-Fi and have enough disk space.
This generally takes 15-20 minutes on an M1 MacBook Pro. Upon successful execution, the 4 neural network models that comprise Stable Diffusion will have been converted from PyTorch to Core ML (`.mlpackage`) and saved into the specified `<output-mlpackages-directory>`. Some additional notable arguments:
- `--model-version`: The model version name as published on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion)
- `--refiner-version`: The refiner version name as published on the [Hugging Face Hub](https://huggingface.co/models?search=stable-diffusion). This is optional and if specified, this argument will convert and bundle the refiner unet alongside the model unet.
- `--bundle-resources-for-swift-cli`: Compiles all 4 models and bundles them along with necessary resources for text tokenization into `<output-mlpackages-directory>/Resources` which should provided as input to the Swift package. This flag is not necessary for the diffusers-based Python pipeline. [However using these compiled models in Python will significantly speed up inference](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html#why-use-a-compiled-model).
- `--quantize-nbits`: Quantizes the weights of unet and text_encoder models down to 2, 4, 6 or 8 bits using a globally optimal k-means clustering algorithm. By default all models are weight-quantized to 16 bits even if this argument is not specified. Please refer to [this section](#compression-6-bits-and-higher for details and further guidance on weight compression.
- `--chunk-unet`: Splits the Unet model in two approximately equal chunks (each with less than 1GB of weights) for mobile-friendly deployment. This is **required** for Neural Engine deployment on iOS and iPadOS if weights are not quantized to 6-bits or less (`--quantize-nbits {2,4,6}`). This is not required for macOS. Swift CLI is able to consume both the chunked and regular versions of the Unet model but prioritizes the former. Note that chunked unet is not compatible with the Python pipeline because Python pipeline is intended for macOS only.
- `--attention-implementation`: Defaults to `SPLIT_EINSUM` which is the implementation described in [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers). `--attention-implementation SPLIT_EINSUM_V2` yields 10-30% improvement for mobile devices, still targeting the Neural Engine. `--attention-implementation ORIGINAL` will switch to an alternative implementation that should be used for CPU or GPU deployment on some Mac devices. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
- `--check-output-correctness`: Compares original PyTorch model's outputs to final Core ML model's outputs. This flag increases RAM consumption significantly so it is recommended only for debugging purposes.
- `--convert-controlnet`: Converts ControlNet models specified after this option. This can also convert multiple models if you specify like `--convert-controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth`.
- `--unet-support-controlnet`: enables a converted UNet model to receive additional inputs from ControlNet. This is required for generating image with using ControlNet and saved with a different name, `*_control-unet.mlpackage`, distinct from normal UNet. On the other hand, this UNet model can not work without ControlNet. Please use normal UNet for just txt2img.
- `--unet-batch-one`: use a batch size of one for the unet, this is needed if you do not want to do classifier free guidance, i.e. using a `guidance-scale` of less than one.
- `--convert-vae-encoder`: not required for text-to-image applications. Required for image-to-image applications in order to map the input image to the latent space.
</details>
## <a name="image-generation-with-python"></a> Image Generation with Python
<details>
<summary> Click to expand </summary>
Run text-to-image generation using the example Python pipeline based on [diffusers](https://github.com/huggingface/diffusers):
```shell
python -m python_coreml_stable_diffusion.pipeline --prompt "a photo of an astronaut riding a horse on mars" -i <core-ml-model-directory> -o </path/to/output/image> --compute-unit ALL --seed 93
```
Please refer to the help menu for all available arguments: `python -m python_coreml_stable_diffusion.pipeline -h`. Some notable arguments:
- `-i`: Should point to the `-o` directory from Step 4 of [Converting Models to Core ML](#converting-models-to-coreml) section from above. If you specified `--bundle-resources-for-swift-cli` during conversion, then use the resulting `Resources` folder (which holds the compiled `.mlmodelc` files). [The compiled models load much faster after first use](https://apple.github.io/coremltools/docs-guides/source/model-prediction.html#why-use-a-compiled-model).
- `--model-version`: If you overrode the default model version while converting models to Core ML, you will need to specify the same model version here.
- `--compute-unit`: Note that the most performant compute unit for this particular implementation may differ across different hardware. `CPU_AND_GPU` or `CPU_AND_NE` may be faster than `ALL`. Please refer to the [Performance Benchmark](#performance-benchmark) section for further guidance.
- `--scheduler`: If you would like to experiment with different schedulers, you may specify it here. For available options, please see the help menu. You may also specify a custom number of inference steps by `--num-inference-steps` which defaults to 50.
- `--controlnet`: ControlNet models specified with this option are used in image generation. Use this option in the format `--controlnet lllyasviel/sd-controlnet-mlsd lllyasviel/sd-controlnet-depth` and make sure to use `--controlnet-inputs` in conjunction.
- `--controlnet-inputs`: Image inputs corresponding to each ControlNet model. Please provide image paths in same order as models in `--controlnet`, for example: `--controlnet-inputs image_mlsd image_depth`.
- `--unet-batch-one`: Do not batch unet predictions for the prompt and negative prompt. This requires the unet has been converted with a batch size of one, see `--unet-batch-one` option in conversion script.
</details>
## <a name="image-gen-swift"></a> Image Generation with Swift
<details>
<summary> Click to expand </summary>
### Example CLI Usage
```shell
swift run StableDiffusionSample "a photo of an astronaut riding a horse on mars" --resource-path <output-mlpackages-directory>/Resources/ --seed 93 --output-path </path/to/output/image>
```
The output will be named based on the prompt and random seed:
e.g. `</path/to/output/image>/a_photo_of_an_astronaut_riding_a_horse_on_mars.93.final.png`
Please use the `--help` flag to learn about batched generation and more.
### Example Library Usage
```swift
import StableDiffusion
...
let pipeline = try StableDiffusionPipeline(resourcesAt: resourceURL)
pipeline.loadResources()
let image = try pipeline.generateImages(prompt: prompt, seed: seed).first
```
On iOS, the `reduceMemory` option should be set to `true` when constructing `StableDiffusionPipeline`
### Swift Package Details
This Swift package contains two products:
- `StableDiffusion` library
- `StableDiffusionSample` command-line tool
Both of these products require the Core ML models and tokenization resources to be supplied. When specifying resources via a directory path that directory must contain the following:
- `TextEncoder.mlmodelc` or `TextEncoder2.mlmodelc (text embedding model)
- `Unet.mlmodelc` or `UnetChunk1.mlmodelc` & `UnetChunk2.mlmodelc` (denoising autoencoder model)
- `VAEDecoder.mlmodelc` (image decoder model)
- `vocab.json` (tokenizer vocabulary file)
- `merges.text` (merges for byte pair encoding file)
Optionally, for image2image, in-painting, or similar:
- `VAEEncoder.mlmodelc` (image encoder model)
Optionally, it may also include the safety checker model that some versions of Stable Diffusion include:
- `SafetyChecker.mlmodelc`
Optionally, for the SDXL refiner:
- `UnetRefiner.mlmodelc` (refiner unet model)
Optionally, for ControlNet:
- `ControlledUNet.mlmodelc` or `ControlledUnetChunk1.mlmodelc` & `ControlledUnetChunk2.mlmodelc` (enabled to receive ControlNet values)
- `controlnet/` (directory containing ControlNet models)
- `LllyasvielSdControlnetMlsd.mlmodelc` (for example, from lllyasviel/sd-controlnet-mlsd)
- `LllyasvielSdControlnetDepth.mlmodelc` (for example, from lllyasviel/sd-controlnet-depth)
- Other models you converted
Note that the chunked version of Unet is checked for first. Only if it is not present will the full `Unet.mlmodelc` be loaded. Chunking is required for iOS and iPadOS and not necessary for macOS.
</details>
## <a name="swift-app"></a> Example Swift App
<details>
<summary> Click to expand </summary>
🤗 Hugging Face created an [open-source demo app](https://github.com/huggingface/swift-coreml-diffusers) on top of this library. It's written in native Swift and Swift UI, and runs on macOS, iOS and iPadOS. You can use the code as a starting point for your app, or to see how to integrate this library in your own projects.
Hugging Face has made the app [available in the Mac App Store](https://apps.apple.com/app/diffusers/id1666309574?mt=12).
</details>
## <a name="faq"></a> FAQ
<details>
<summary> Click to expand </summary>
<details>
<summary> <b> Q1: </b> <code> ERROR: Failed building wheel for tokenizers or error: can't find Rust compiler </code> </summary>
<b> A1: </b> Please review this [potential solution](https://github.com/huggingface/transformers/issues/2831#issuecomment-592724471).
</details>
<details>
<summary> <b> Q2: </b> <code> RuntimeError: {NSLocalizedDescription = "Error computing NN outputs." </code> </summary>
<b> A2: </b> There are many potential causes for this error. In this context, it is highly likely to be encountered when your system is under increased memory pressure from other applications. Reducing memory utilization of other applications is likely to help alleviate the issue.
</details>
<details>
<summary> <b> <a name="low-mem-conversion"></a> Q3: </b> My Mac has 8GB RAM and I am converting models to Core ML using the example command. The process is getting killed because of memory issues. How do I fix this issue? </summary>
<b> A3: </b> In order to minimize the memory impact of the model conversion process, please execute the following command instead:
```bash
python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-encoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-vae-decoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-text-encoder --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> && \
python -m python_coreml_stable_diffusion.torch2coreml --convert-safety-checker --model-version <model-version-string-from-hub> -o <output-mlpackages-directory> &&
```
If you need `--chunk-unet`, you may do so in yet another independent command which will reuse the previously exported Unet model and simply chunk it in place:
```bash
python -m python_coreml_stable_diffusion.torch2coreml --convert-unet --chunk-unet -o <output-mlpackages-directory>
```
</details>
<details>
<summary> <b> Q4: </b> My Mac has 8GB RAM, should image generation work on my machine? </summary>
<b> A4: </b> Yes! Especially the `--compute-unit CPU_AND_NE` option should work under reasonable system load from other applications. Note that part of the [Example Results](#example-results) were generated using an M2 MacBook Air with 8GB RAM.
</details>
<details>
<summary> <b> Q5: </b> Every time I generate an image using the Python pipeline, loading all the Core ML models takes 2-3 minutes. Is this expected? </summary>
<b> A5: </b> Both `.mlpackage` and `.mlmodelc` models are compiled (also known as "model preparation" in Core ML terms) upon first load when a specific compute unit is specified. `.mlpackage` does not cache this compiled asset so each model load retriggers this compilation which may take up to a few minutes. On the other hand, `.mlmodelc` files do cache this compiled asset and non-first load times are reduced to just a few seconds.
In order to benefit from compilation caching, you may use the `.mlmodelc` assets instead of `.mlpackage` assets in both Swift (default) and Python (possible thanks to [@lopez-hector](https://github.com/lopez-hector)'s [contribution](https://github.com/apple/ml-stable-diffusion/commit/f3a212491cf531dd88493c89ad3d98d016db407f)) image generation pipelines.
</details>
<details>
<summary> <b> <a name="q-mobile-app"></a> Q6: </b> I want to deploy <code>StableDiffusion</code>, the Swift package, in my mobile app. What should I be aware of? </summary>
<b> A6: </b>The [Image Generation with Swift](#image-gen-swift) section describes the minimum SDK and OS versions as well as the device models supported by this package. We recommend carefully testing the package on the device with the least amount of RAM available among your deployment targets.
The image generation process in `StableDiffusion` can yield over 2 GB of peak memory during runtime depending on the compute units selected. On iPadOS, we recommend using `.cpuAndNeuralEngine` in your configuration and the `reduceMemory` option when constructing a `StableDiffusionPipeline` to minimize memory pressure.
If your app crashes during image generation, consider adding the [Increased Memory Limit](https://developer.apple.com/documentation/bundleresources/entitlements/com_apple_developer_kernel_increased-memory-limit) capability to inform the system that some of your app’s core features may perform better by exceeding the default app memory limit on supported devices.
On iOS, depending on the iPhone model, Stable Diffusion model versions, selected compute units, system load and design of your app, this may still not be sufficient to keep your apps peak memory under the limit. Please remember, because the device shares memory between apps and iOS processes, one app using too much memory can compromise the user experience across the whole device.
We **strongly recommend** compressing your models following the recipes in [Advanced Weight Compression (Lower than 6-bits)](#compression-lower-than-6-bits) for iOS deployment. This reduces the peak RAM usage by up to 75% (from 16-bit to 4-bit) while preserving model output quality.
</details>
<details>
<summary> <b> Q7: </b> How do I generate images with different resolutions using the same Core ML models? </summary>
<b> A7: </b> The current version of `python_coreml_stable_diffusion` does not support single-model multi-resolution out of the box. However, developers may fork this project and leverage the [flexible shapes](https://coremltools.readme.io/docs/flexible-inputs) support from coremltools to extend the `torch2coreml` script by using `coremltools.EnumeratedShapes`. Note that, while the `text_encoder` is agnostic to the image resolution, the inputs and outputs of `vae_decoder` and `unet` models are dependent on the desired image resolution.
</details>
<details>
<summary> <b> Q8: </b> Are the Core ML and PyTorch generated images going to be identical? </summary>
<b> A8: </b> If desired, the generated images across PyTorch and Core ML can be made approximately identical. However, it is not guaranteed by default. There are several factors that might lead to different images across PyTorch and Core ML:
<b> 1. Random Number Generator Behavior </b>
The main source of potentially different results across PyTorch and Core ML is the Random Number Generator ([RNG](https://en.wikipedia.org/wiki/Random_number_generation)) behavior. PyTorch and Numpy have different sources of randomness. `python_coreml_stable_diffusion` generally relies on Numpy for RNG (e.g. latents initialization) and `StableDiffusion` Swift Library reproduces this RNG behavior by default. However, PyTorch-based pipelines such as Hugging Face `diffusers` relies on PyTorch's RNG behavior. Thanks to @liuliu's [contributions](https://github.com/apple/ml-stable-diffusion/pull/124), one can match the PyTorch (CPU/GPU) RNG behavior in Swift by specifying `--rng torch/cuda` which selects the `torchRNG/cudaRNG` mode.
<b> 2. PyTorch </b>
*"Completely reproducible results are not guaranteed across PyTorch releases, individual commits, or different platforms. Furthermore, results may not be reproducible between CPU and GPU executions, even when using identical seeds."* ([source](https://pytorch.org/docs/stable/notes/randomness.html#reproducibility)).
<b> 3. Model Function Drift During Conversion </b>
The difference in outputs across corresponding PyTorch and Core ML models is a potential cause. The signal integrity is tested during the conversion process (enabled via `--check-output-correctness` argument to `python_coreml_stable_diffusion.torch2coreml`) and it is verified to be above a minimum [PSNR](https://en.wikipedia.org/wiki/Peak_signal-to-noise_ratio) value as tested on random inputs. Note that this is simply a sanity check and does not guarantee this minimum PSNR across all possible inputs. Furthermore, the results are not guaranteed to be identical when executing the same Core ML models across different compute units. This is not expected to be a major source of difference as the sample visual results indicate in [this section](#compression-6-bits-and-higher).
<b> 4. Weights and Activations Data Type </b>
When quantizing models from float32 to lower-precision data types such as float16, the generated images are [known to vary slightly](https://lambdalabs.com/blog/inference-benchmark-stable-diffusion) in semantics even when using the same PyTorch model. Core ML models generated by coremltools have float16 weights and activations by default [unless explicitly overridden](https://github.com/apple/coremltools/blob/main/coremltools/converters/_converters_entry.py#L256). This is not expected to be a major source of difference.
</details>
<details>
<summary> <b> Q9: </b> The model files are very large, how do I avoid a large binary for my App? </summary>
<b> A9: </b> The recommended option is to prompt the user to download these assets upon first launch of the app. This keeps the app binary size independent of the Core ML models being deployed. Disclosing the size of the download to the user is extremely important as there could be data charges or storage impact that the user might not be comfortable with.
</details>
<details>
<summary> <b> Q10: </b> <code> `Could not initialize NNPACK! Reason: Unsupported hardware` </code> </summary>
<b> A10: </b> This warning is safe to ignore in the context of this repository.
</details>
<details>
<summary> <b> Q11: </b> <code> TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect </code> </summary>
<b> A11: </b> This warning is safe to ignore in the context of this repository.
</details>
<details>
<summary> <b> Q12: </b> <code> UserWarning: resource_tracker: There appear to be 1 leaked semaphore objects to clean up at shutdown </code> </summary>
<b> A12: </b> If this warning is printed right after <code> zsh: killed python -m python_coreml_stable_diffusion.torch2coreml ... </code>, then it is highly likely that your Mac has run out of memory while converting models to Core ML. Please see [Q3](#low-mem-conversion) from above for the solution.
</details>
</details>
</details>
## <a name="bibtex"></a> BibTeX Reference
```latex
@misc{stable-diffusion-coreml-apple-silicon,
title = {Stable Diffusion with Core ML on Apple Silicon},
author = {Atila Orhon and Michael Siracusa and Aseem Wadhwa},
year = {2022},
URL = {null}
}
```
README.wehub.md
+7
View File
@@ -0,0 +1,7 @@
# WeHub 来源说明
- 原始项目：`apple/ml-stable-diffusion`
- 原始仓库：https://github.com/apple/ml-stable-diffusion
- 导入方式：上游默认分支的最新快照
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_11_computeUnit_CPU_AND_GPU_modelVersion_runwayml_stable-diffusion-v1-5.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 395 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_11_computeUnit_CPU_AND_NE_modelVersion_runwayml_stable-diffusion-v1-5.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 430 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_123456789_computeUnit_ALL_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 444 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_123456789_computeUnit_CPU_AND_GPU_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 428 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_123456789_computeUnit_CPU_AND_NE_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 444 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_93_computeUnit_ALL_modelVersion_stabilityai_stable-diffusion-2-base.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 507 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_93_computeUnit_CPU_AND_GPU_modelVersion_stabilityai_stable-diffusion-2-base.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 520 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_dragon_in_space/randomSeed_93_computeUnit_CPU_AND_NE_modelVersion_stabilityai_stable-diffusion-2-base.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 507 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_11_computeUnit_CPU_AND_GPU_modelVersion_stabilityai_stable-diffusion-2-base.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 423 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_11_computeUnit_CPU_AND_NE_modelVersion_stabilityai_stable-diffusion-2-base.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 427 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_13_computeUnit_ALL_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 467 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_13_computeUnit_CPU_AND_GPU_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 446 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_13_computeUnit_CPU_AND_NE_modelVersion_CompVis_stable-diffusion-v1-4.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 468 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_ALL_modelVersion_runwayml_stable-diffusion-v1-5.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 460 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_CPU_AND_GPU_modelVersion_runwayml_stable-diffusion-v1-5.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 456 KiB

assets/a_high_quality_photo_of_an_astronaut_riding_a_horse_in_space/randomSeed_93_computeUnit_CPU_AND_NE_modelVersion_runwayml_stable-diffusion-v1-5.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 461 KiB

assets/controlnet_readme_reel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.5 MiB

assets/float16_cpuandne_readmereel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.2 MiB

assets/float16_gpu_readmereel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.2 MiB

assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_3.41-bits.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.5 MiB

assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_4.50-bits.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.5 MiB

assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_6.55-bits.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.5 MiB

assets/mbp/a_high_quality_photo_of_a_surfing_dog.7667.final_float16_original.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.6 MiB

assets/mbp/runwayml_stable-diffusion-v1-5_psnr_vs_size.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 56 KiB

assets/mbp/stabilityai_stable-diffusion-2-1-base_psnr_vs_size.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 64 KiB

assets/mbp/stabilityai_stable-diffusion-xl-base-1.0_psnr_vs_size.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 57 KiB

assets/palette6_cpuandne_readmereel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.2 MiB

assets/readme_reel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

python_coreml_stable_diffusion/__init__.py
+1
View File
@@ -0,0 +1 @@
from ._version import __version__
python_coreml_stable_diffusion/_version.py
+1
View File
@@ -0,0 +1 @@
__version__ = "1.1.0"
python_coreml_stable_diffusion/activation_quantization.py
+503
View File
@@ -0,0 +1,503 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import logging
import operator
import torch
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel('INFO')
import argparse
import gc
import json
import os
import pickle
from copy import deepcopy
import coremltools as ct
import numpy as np
from coremltools.optimize.torch.quantization import (
LinearQuantizer, LinearQuantizerConfig, ModuleLinearQuantizerConfig)
from diffusers import StableDiffusionPipeline
from tqdm import tqdm
from python_coreml_stable_diffusion import attention
from python_coreml_stable_diffusion import unet
from python_coreml_stable_diffusion.layer_norm import LayerNormANE
from python_coreml_stable_diffusion.torch2coreml import compute_psnr
from python_coreml_stable_diffusion.unet import Einsum
attention.SPLIT_SOFTMAX = True
CALIBRATION_DATA = [
"image of a transparent tall glass with ice, fruits and mint, photograph, commercial, food, warm background, beautiful image, detailed",
"picture of dimly lit living room, minimalist furniture, vaulted ceiling, huge room, floor to ceiling window with an ocean view, nighttime, 3D render, high quality, detailed",
"modern office building, 8 stories tall, glass and steel, 3D render style, wide angle view, very detailed, sharp photographic image, in an office park, bright sunny day, clear blue skies, trees and landscaping",
"cute small cat sitting in a movie theater eating popcorn, watching a movie, cozy indoor lighting, detailed, digital painting, character design",
"a highly detailed matte painting of a man on a hill watching a rocket launch in the distance by studio ghibli, volumetric lighting, octane render, 4K resolution, hyperrealism, highly detailed, insanely detailed, cinematic lighting, depth of field",
"an undersea world with several of fish, rocks, detailed, realistic, photograph, amazing, beautiful, high resolution",
"large ocean wave hitting a beach at sunset, photograph, detailed",
"pocket watch on a table, close up. macro, sharp, high gloss, brass, gears, sharp, detailed",
"pocket watch in the style of pablo picasso, painting",
"majestic royal tall ship on a calm sea, realistic painting, cloudy blue sky, in the style of edward hopper",
"german castle on a mountain, blue sky, realistic, photograph, dramatic, wide angle view",
"artificial intelligence, AI, concept art, blue line sketch",
"a humanoid robot, concept art, 3D render, high quality, detailed",
"donut with sprinkles and a cup of coffee on a wood table, detailed, photograph",
"orchard at sunset, beautiful, photograph, great composition, detailed, realistic, HDR",
"image of a map of a country, tattered, old, styled, illustration, for a video game style",
"blue and green woven fibers, nano fiber material, detailed, concept art, micro photography",
]
RANDOM_TEST_DATA = [
"a black and brown dog standing outside a door.",
"a person on a motorcycle makes a turn on the track.",
"inflatable boats sit on the arizona river, and on the bank",
"a white cat sitting under a white umbrella",
"black bear standing in a field of grass under a tree.",
"a train that is parked on tracks and has graffiti writing on it, with a mountain range in the background.",
"a cake inside of a pan sitting in an oven.",
"a table with paper plates and flowers in a home",
]
def get_coreml_inputs(sample_inputs):
return [
ct.TensorType(
name=k,
shape=v.shape,
dtype=v.numpy().dtype if isinstance(v, torch.Tensor) else v.dtype,
) for k, v in sample_inputs.items()
]
def convert_to_coreml(torchscript_module, sample_inputs):
logger.info("Converting model to CoreML..")
coreml_model = ct.convert(
torchscript_module,
convert_to="mlprogram",
minimum_deployment_target=ct.target.macOS14,
inputs=get_coreml_inputs(sample_inputs),
outputs=[ct.TensorType(name="noise_pred", dtype=np.float32)],
compute_units=ct.ComputeUnit.ALL,
skip_model_load=True,
)
return coreml_model
def unet_data_loader(data_dir, device='cpu', calibration_nsamples=None):
"""
Load calibration data from specified path.
Limit number of samples to calibration_nsamples, if specified.
"""
dataloader = []
skip_load = False
for file in sorted(os.listdir(data_dir)):
if file.endswith('.pkl'):
filepath = os.path.join(data_dir, file)
with open(filepath, 'rb') as data:
try:
while not skip_load:
unet_data = pickle.load(data)
for input in unet_data:
dataloader.append([x.to(torch.float).to(device) for x in input])
if calibration_nsamples:
if len(dataloader) >= calibration_nsamples:
skip_load = True
break
except EOFError:
pass
if skip_load:
break
logger.info(f"Total calibration samples: {len(dataloader)}")
return dataloader
def quantize_module_config(module_name):
"""
Generate quantization config to apply W8A8 quantization for specified module.
Rest of the model is kept in FP32 precision.
"""
config = LinearQuantizerConfig(
global_config=ModuleLinearQuantizerConfig(
milestones=[0, 1000, 1000, 0],
weight_dtype=torch.float32,
activation_dtype=torch.float32,
),
module_name_configs={
module_name: ModuleLinearQuantizerConfig(
quantization_scheme="symmetric",
milestones=[0, 1000, 1000, 0],
),
},
)
return config
def quantize_cumulative_config(skip_conv_layers, skip_einsum_layers):
"""
Generate quantization config to apply W8A8 quantization.
Skipped layers are kept in W8A32 precision.
"""
logger.info(f"Skipping {len(skip_conv_layers)} conv layers and {len(skip_einsum_layers)} einsum layers")
w8config = ModuleLinearQuantizerConfig(
quantization_scheme="symmetric",
milestones=[0, 1000, 1000, 0],
activation_dtype=torch.float32)
conv_modules_config = {name: w8config for name in skip_conv_layers}
einsum_modules_config = {name: w8config for name in skip_einsum_layers}
module_name_config = {}
module_name_config.update(conv_modules_config)
module_name_config.update(einsum_modules_config)
config = LinearQuantizerConfig(
global_config=ModuleLinearQuantizerConfig(
quantization_scheme="symmetric",
milestones=[0, 1000, 1000, 0],
),
module_name_configs=module_name_config,
module_type_configs={
torch.cat: None,
torch.nn.GroupNorm: None,
torch.nn.SiLU: None,
torch.nn.functional.gelu: None,
operator.add: None,
},
)
return config
def quantize(model, config, calibration_data):
"""
Apply post training activation quantization to specified model, using calibration data
"""
submodules = dict(model.named_modules(remove_duplicate=True))
layer_norm_modules = [key for key, val in submodules.items() if isinstance(val, LayerNormANE)]
non_traceable_module_names = layer_norm_modules + [
"time_proj",
"time_embedding",
]
# Mark certain modules as non-traceable to make the UNet model fx traceable
config.non_traceable_module_names = non_traceable_module_names
config.preserved_attributes = ['config', 'device']
sample_input = calibration_data[0]
quantizer = LinearQuantizer(model, config)
logger.info("Preparing model for quantization")
prepared_model = quantizer.prepare(example_inputs=(sample_input,))
prepared_model.eval()
quantizer.step()
logger.info("Calibrate")
for idx, data in enumerate(calibration_data):
logger.info(f"Calibration data sample: {idx}")
prepared_model(*data)
logger.info("Finalize model")
quantized_model = quantizer.finalize()
return quantized_model
def get_quantizable_modules(unet):
quantizable_modules = []
for name, module in unet.named_modules():
if len(list(module.children())) > 0:
continue
if type(module) == torch.nn.modules.conv.Conv2d:
quantizable_modules.append(('conv', name))
if type(module) == Einsum:
quantizable_modules.append(('einsum', name))
return quantizable_modules
def recipe_overrides_for_inference_speedup(conv_layers, skipped_conv):
"""
Quantize the slowest conv layers, even if in skipped set based on PSNR, for good inference speedup
"""
for layer in conv_layers:
if "up_blocks" in layer and "resnets" in layer and "conv1" in layer:
if layer in skipped_conv:
logger.info(f"removing {layer}")
skipped_conv.remove(layer)
if "upsamplers" in layer:
if layer in skipped_conv:
logger.info(f"removing {layer}")
skipped_conv.remove(layer)
def recipe_overrides_for_quality(conv_layers, skipped_conv):
"""
Do not quantize out projection layers to avoid quantizing outputs of preceding concat layers.
Quantizing output of concat layers can lead to quality degradation, due to sharing of scales
across concat inputs, which can have varied ranges. Since this is a constraint enforced during
model conversion, it may not be captured in layer-wise PSNR analysis of PyTorch model.
"""
out_proj_layers = [layer for layer in conv_layers if "to_out" in layer]
for layer in out_proj_layers:
if layer not in skipped_conv:
logger.info(f"adding {layer}")
skipped_conv.add(layer)
def register_input_log_hook(unet, inputs):
"""
Register forward pre hook to save model inputs
"""
def hook(_, input):
input_copy = deepcopy(input)
input_copy = tuple(i.to('cpu') for i in input_copy)
inputs.append(input_copy)
# Return inputs unmodified
return input
return unet.register_forward_pre_hook(hook)
def generate_calibration_data(pipe, args, calibration_dir):
# Register forward pre hook to record unet inputs
unet_inputs = []
handle = register_input_log_hook(pipe.unet, unet_inputs)
# If directory doesn't exist, create it
os.makedirs(calibration_dir, exist_ok=True)
# Run calibration prompts through the pipeline and
# serialize recorded UNet model inputs
for prompt in CALIBRATION_DATA:
gen = torch.manual_seed(args.seed)
# run forward pass
pipe(prompt=prompt, generator=gen)
# save unet inputs
filename = "_".join(prompt.split(" ")) + "_" + str(args.seed) + ".pkl"
filepath = os.path.join(calibration_dir, filename)
with open(filepath, 'wb') as f:
pickle.dump(unet_inputs, f)
# clear
unet_inputs.clear()
handle.remove()
def register_input_preprocessing_hook(pipe):
"""
Register forward pre hook to convert UNet inputs from HuggingFace StableDiffusionPipeline
to match expected model inputs in UNet2DConditionModel defined in unet.py
"""
def hook(_, args, kwargs):
sample = args[0]
timestep = args[1]
if len(timestep.shape) == 0:
timestep = timestep[None]
timestep = timestep.expand(sample.shape[0])
encoder_hidden_states = kwargs["encoder_hidden_states"]
encoder_hidden_states = encoder_hidden_states.permute((0, 2, 1)).unsqueeze(2)
modified_args = (sample, timestep, encoder_hidden_states)
return (modified_args, {})
return pipe.unet.register_forward_pre_hook(hook, with_kwargs=True)
def prepare_pipe(pipe, unet):
"""
Create a new pipeline from `pipe` with `unet` as the noise predictor
"""
new_pipe = deepcopy(pipe)
unet.to(new_pipe.unet.device)
new_pipe.unet = unet
pre_hook_handle = register_input_preprocessing_hook(new_pipe)
return new_pipe, pre_hook_handle
def run_pipe(pipe):
gen = torch.manual_seed(args.seed)
kwargs = dict(
prompt=RANDOM_TEST_DATA,
output_type="latent",
generator=gen,
)
return np.array([latent.cpu().numpy() for latent in pipe(**kwargs).images])
def get_reference_pipeline(model_version):
# Initialize pipe
pipe = StableDiffusionPipeline.from_pretrained(
model_version,
use_safetensors=True,
use_auth_token=True,
)
DEFAULT_NUM_INFERENCE_STEPS = 50
pipe.scheduler.set_timesteps(DEFAULT_NUM_INFERENCE_STEPS)
# Initialize reference unet
unet_cls = unet.UNet2DConditionModel
reference_unet = unet_cls(**pipe.unet.config).eval()
reference_unet.load_state_dict(pipe.unet.state_dict())
# Initialize reference pipeline
ref_pipe, _ = prepare_pipe(pipe, reference_unet)
del pipe
gc.collect()
return ref_pipe
def main(args):
# Initialize reference pipeline
ref_pipe = get_reference_pipeline(args.model_version)
if torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
logger.debug(f"Placing pipe in {device}")
ref_pipe.to(device)
# Generate baseline outputs
ref_out = run_pipe(ref_pipe)
# Setup artifact file paths
os.makedirs(args.o, exist_ok=True)
recipe_json_path = os.path.join(args.o, f"{args.model_version.replace('/', '_')}_quantization_recipe.json")
calibration_dir = os.path.join(args.o, f"calibration_data_{args.model_version.replace('/', '_')}")
# Generate calibration data
if args.generate_calibration_data:
generate_calibration_data(ref_pipe, args, calibration_dir)
# Compute layer-wise PSNR
if args.layerwise_sensitivity:
logger.info("Compute Layer-wise PSNR")
quantizable_modules = get_quantizable_modules(ref_pipe.unet)
results = {
'conv': {},
'einsum': {},
'model_version': args.model_version
}
dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
for module_type, module_name in tqdm(quantizable_modules):
logger.info(f"Quantizing UNet Layer: {module_name}")
config = quantize_module_config(module_name)
quantized_unet = quantize(ref_pipe.unet, config, dataloader)
# Generate outputs from quantized model
q_pipe, _ = prepare_pipe(ref_pipe, quantized_unet)
test_out = run_pipe(q_pipe)
psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
logger.info(f"PSNR: {psnr}")
avg_psnr = sum(psnr) / len(psnr)
logger.info(f"AVG PSNR: {avg_psnr}")
results[module_type][module_name] = avg_psnr
del quantized_unet
del q_pipe
gc.collect()
with open(recipe_json_path, 'w') as f:
json.dump(results, f, indent=2)
if args.quantize_pytorch:
logger.info("Quantizing UNet PyTorch model")
dataloader = unet_data_loader(calibration_dir, device, args.calibration_nsamples)
with open(recipe_json_path, "r") as f:
results = json.load(f)
logger.info(f"Conv PSNR threshold: {args.conv_psnr}, Attn PSNR threshold: {args.attn_psnr}")
skipped_conv = set([layer for layer, psnr in results['conv'].items() if psnr < args.conv_psnr])
skipped_einsum = set([layer for layer, psnr in results['einsum'].items() if psnr < args.attn_psnr])
# Apply some overrides on PSNR based recipe for inference and quality improvements
# Users can disable these selectively based on specific targets
recipe_overrides_for_inference_speedup(results['conv'].keys(), skipped_conv)
recipe_overrides_for_quality(results['conv'].keys(), skipped_conv)
config = quantize_cumulative_config(skipped_conv, skipped_einsum)
quantized_unet = quantize(ref_pipe.unet, config, dataloader)
# Generate outputs from quantized model
q_pipe, handle = prepare_pipe(ref_pipe, quantized_unet)
test_out = run_pipe(q_pipe)
psnr = [float(f"{compute_psnr(r, t):.1f}") for r, t in zip(ref_out, test_out)]
logger.info(f"PSNR: {psnr}")
avg_psnr = sum(psnr) / len(psnr)
logger.info(f"AVG PSNR: {avg_psnr}")
handle.remove()
quantized_unet.to('cpu')
sample_unet_input = {
"sample": dataloader[0][0].to('cpu'),
"timestep": dataloader[0][1].to('cpu'),
"encoder_hidden_states": dataloader[0][2].to('cpu'),
}
logger.info("JIT tracing quantized model")
traced_model = torch.jit.trace(quantized_unet, example_inputs=list(sample_unet_input.values()))
logger.info("Converting to CoreML")
coreml_sample_unet_input = {
k: v.numpy().astype(np.float16)
for k, v in sample_unet_input.items()
}
coreml_model = convert_to_coreml(traced_model, coreml_sample_unet_input)
coreml_filename = f"Stable_Diffusion_version_{args.model_version.replace('/', '_')}_unet.mlpackage"
coreml_model.save(os.path.join(args.o, coreml_filename))
del q_pipe
del ref_pipe
gc.collect()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"-o",
required=True,
help="Output directory to save calibration data and quantization artifacts"
)
parser.add_argument(
"--model-version",
required=True,
choices=("runwayml/stable-diffusion-v1-5", "stabilityai/stable-diffusion-2-1-base"),
help=
("The pre-trained model checkpoint and configuration to restore"
))
parser.add_argument(
"--generate-calibration-data",
action="store_true",
help="Generate calibration data for UNet model"
)
parser.add_argument(
"--layerwise-sensitivity",
action="store_true",
help="Compute compression sensitivity per-layer, by quantizing one layer at a time"
)
parser.add_argument(
"--quantize-pytorch",
action="store_true",
help="Generate activation quantized UNet model by quantizing layers above specified PSNR threshold"
)
parser.add_argument(
"--calibration-nsamples",
type=int,
help="Number of samples to use for calibrating UNet model"
)
parser.add_argument("--seed",
"-s",
default=50,
type=int,
help="Random seed to be able to reproduce results"
)
parser.add_argument("--conv-psnr",
default=40.0,
type=float,
help="PSNR threshold for convolutional layers (default for stabilityai/stable-diffusion-2-1-base)"
)
parser.add_argument("--attn-psnr",
default=30.0,
type=float,
help="PSNR threshold for attention (Einsum) layers (default for stabilityai/stable-diffusion-2-1-base)"
)
args = parser.parse_args()
main(args)
python_coreml_stable_diffusion/attention.py
+168
View File
@@ -0,0 +1,168 @@
import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
import torch
import math
SPLIT_SOFTMAX = False
def softmax(x, dim):
# Reduction max
max_x = x.max(dim=dim, keepdim=True).values
# EW sub
x -= max_x
# Scale for EXP to EXP2, Activation EXP2
scaled_x = x * (1 / math.log(2))
exp_act = torch.exp2(scaled_x)
# Reduction Sum + Inv
exp_sum_inv = 1 / exp_act.sum(dim=dim, keepdims=True)
# EW Mult
return exp_act * exp_sum_inv
def split_einsum(q, k, v, mask, heads, dim_head):
""" Attention Implementation backing AttentionImplementations.SPLIT_EINSUM
- Implements https://machinelearning.apple.com/research/neural-engine-transformers
- Recommended for ANE
- Marginally slower on GPU
"""
mh_q = [
q[:, head_idx * dim_head:(head_idx + 1) *
dim_head, :, :] for head_idx in range(heads)
] # (bs, dim_head, 1, max_seq_length) * heads
k = k.transpose(1, 3)
mh_k = [
k[:, :, :,
head_idx * dim_head:(head_idx + 1) * dim_head]
for head_idx in range(heads)
] # (bs, max_seq_length, 1, dim_head) * heads
mh_v = [
v[:, head_idx * dim_head:(head_idx + 1) *
dim_head, :, :] for head_idx in range(heads)
] # (bs, dim_head, 1, max_seq_length) * heads
attn_weights = [
torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * (dim_head**-0.5)
for qi, ki in zip(mh_q, mh_k)
] # (bs, max_seq_length, 1, max_seq_length) * heads
if mask is not None:
for head_idx in range(heads):
attn_weights[head_idx] = attn_weights[head_idx] + mask
if SPLIT_SOFTMAX:
attn_weights = [
softmax(aw, dim=1) for aw in attn_weights
] # (bs, max_seq_length, 1, max_seq_length) * heads
else:
attn_weights = [
aw.softmax(dim=1) for aw in attn_weights
] # (bs, max_seq_length, 1, max_seq_length) * heads
attn = [
torch.einsum("bkhq,bchk->bchq", wi, vi)
for wi, vi in zip(attn_weights, mh_v)
] # (bs, dim_head, 1, max_seq_length) * heads
attn = torch.cat(attn, dim=1) # (bs, dim, 1, max_seq_length)
return attn
CHUNK_SIZE = 512
def split_einsum_v2(q, k, v, mask, heads, dim_head):
""" Attention Implementation backing AttentionImplementations.SPLIT_EINSUM_V2
- Implements https://machinelearning.apple.com/research/neural-engine-transformers
- Recommended for ANE
- Marginally slower on GPU
- Chunks the query sequence to avoid large intermediate tensors and improves ANE performance
"""
query_seq_length = q.size(3)
num_chunks = query_seq_length // CHUNK_SIZE
if num_chunks == 0:
logger.info(
"AttentionImplementations.SPLIT_EINSUM_V2: query sequence too short to chunk "
f"({query_seq_length}<{CHUNK_SIZE}), fall back to AttentionImplementations.SPLIT_EINSUM (safe to ignore)")
return split_einsum(q, k, v, mask, heads, dim_head)
logger.info(
"AttentionImplementations.SPLIT_EINSUM_V2: Splitting query sequence length of "
f"{query_seq_length} into {num_chunks} chunks")
mh_q = [
q[:, head_idx * dim_head:(head_idx + 1) *
dim_head, :, :] for head_idx in range(heads)
] # (bs, dim_head, 1, max_seq_length) * heads
# Chunk the query sequence for each head
mh_q_chunked = [
[h_q[..., chunk_idx * CHUNK_SIZE:(chunk_idx + 1) * CHUNK_SIZE] for chunk_idx in range(num_chunks)]
for h_q in mh_q
] # ((bs, dim_head, 1, QUERY_SEQ_CHUNK_SIZE) * num_chunks) * heads
k = k.transpose(1, 3)
mh_k = [
k[:, :, :,
head_idx * dim_head:(head_idx + 1) * dim_head]
for head_idx in range(heads)
] # (bs, max_seq_length, 1, dim_head) * heads
mh_v = [
v[:, head_idx * dim_head:(head_idx + 1) *
dim_head, :, :] for head_idx in range(heads)
] # (bs, dim_head, 1, max_seq_length) * heads
attn_weights = [
[
torch.einsum("bchq,bkhc->bkhq", [qi_chunk, ki]) * (dim_head**-0.5)
for qi_chunk in h_q_chunked
] for h_q_chunked, ki in zip(mh_q_chunked, mh_k)
] # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
attn_weights = [
[aw_chunk.softmax(dim=1) for aw_chunk in aw_chunked]
for aw_chunked in attn_weights
] # ((bs, max_seq_length, 1, chunk_size) * num_chunks) * heads
attn = [
[
torch.einsum("bkhq,bchk->bchq", wi_chunk, vi)
for wi_chunk in wi_chunked
] for wi_chunked, vi in zip(attn_weights, mh_v)
] # ((bs, dim_head, 1, chunk_size) * num_chunks) * heads
attn = torch.cat([
torch.cat(attn_chunked, dim=3) for attn_chunked in attn
], dim=1) # (bs, dim, 1, max_seq_length)
return attn
def original(q, k, v, mask, heads, dim_head):
""" Attention Implementation backing AttentionImplementations.ORIGINAL
- Not recommended for ANE
- Recommended for GPU
"""
bs = q.size(0)
mh_q = q.view(bs, heads, dim_head, -1)
mh_k = k.view(bs, heads, dim_head, -1)
mh_v = v.view(bs, heads, dim_head, -1)
attn_weights = torch.einsum("bhcq,bhck->bhqk", [mh_q, mh_k])
attn_weights.mul_(dim_head**-0.5)
if mask is not None:
attn_weights = attn_weights + mask
attn_weights = attn_weights.softmax(dim=3)
attn = torch.einsum("bhqk,bhck->bhcq", [attn_weights, mh_v])
attn = attn.contiguous().view(bs, heads * dim_head, 1, -1)
return attn
python_coreml_stable_diffusion/chunk_mlprogram.py
+410
View File
@@ -0,0 +1,410 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import argparse
from collections import OrderedDict
import coremltools as ct
from coremltools.converters.mil import Block, Program, Var
from coremltools.converters.mil.frontend.milproto.load import load as _milproto_to_pymil
from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil import Placeholder
from coremltools.converters.mil.mil import types as types
from coremltools.converters.mil.mil.passes.helper import block_context_manager
from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
from coremltools.converters.mil.testing_utils import random_gen_input_feature_type
import gc
import logging
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
import numpy as np
import os
from python_coreml_stable_diffusion import torch2coreml
import shutil
import time
def _verify_output_correctness_of_chunks(full_model,
first_chunk_model=None,
second_chunk_model=None,
pipeline_model=None,):
""" Verifies the end-to-end output correctness of full (original) model versus chunked models
"""
# Generate inputs for first chunk and full model
input_dict = {}
for input_desc in full_model._spec.description.input:
input_dict[input_desc.name] = random_gen_input_feature_type(input_desc)
# Generate outputs for full model
outputs_from_full_model = full_model.predict(input_dict)
if pipeline_model is not None:
outputs_from_pipeline_model = pipeline_model.predict(input_dict)
final_outputs = outputs_from_pipeline_model
elif first_chunk_model is not None and second_chunk_model is not None:
# Generate outputs for first chunk
outputs_from_first_chunk_model = first_chunk_model.predict(input_dict)
# Prepare inputs for second chunk model from first chunk's outputs and regular inputs
second_chunk_input_dict = {}
for input_desc in second_chunk_model._spec.description.input:
if input_desc.name in outputs_from_first_chunk_model:
second_chunk_input_dict[
input_desc.name] = outputs_from_first_chunk_model[
input_desc.name]
else:
second_chunk_input_dict[input_desc.name] = input_dict[
input_desc.name]
# Generate output for second chunk model
outputs_from_second_chunk_model = second_chunk_model.predict(
second_chunk_input_dict)
final_outputs = outputs_from_second_chunk_model
else:
raise ValueError
# Verify correctness across all outputs from second chunk and full model
for out_name in outputs_from_full_model.keys():
torch2coreml.report_correctness(
original_outputs=outputs_from_full_model[out_name],
final_outputs=final_outputs[out_name],
log_prefix=f"{out_name}")
def _load_prog_from_mlmodel(model):
""" Load MIL Program from an MLModel
"""
model_spec = model.get_spec()
start_ = time.time()
logger.info(
"Loading MLModel object into a MIL Program object (including the weights).."
)
prog = _milproto_to_pymil(
model_spec=model_spec,
specification_version=model_spec.specificationVersion,
file_weights_dir=model.weights_dir,
)
logger.info(f"Program loaded in {time.time() - start_:.1f} seconds")
return prog
def _get_op_idx_split_location(prog: Program):
""" Find the op that approximately bisects the graph as measure by weights size on each side
"""
main_block = prog.functions["main"]
main_block.operations = list(main_block.operations)
total_size_in_mb = 0
for op in main_block.operations:
if op.op_type == "const" and isinstance(op.val.val, np.ndarray):
size_in_mb = op.val.val.size * op.val.val.itemsize / (1024 * 1024)
total_size_in_mb += size_in_mb
half_size = total_size_in_mb / 2
# Find the first non const op (single child), where the total cumulative size exceeds
# the half size for the first time
cumulative_size_in_mb = 0
for op in main_block.operations:
if op.op_type == "const" and isinstance(op.val.val, np.ndarray):
size_in_mb = op.val.val.size * op.val.val.itemsize / (1024 * 1024)
cumulative_size_in_mb += size_in_mb
# Note: The condition "not op.op_type.startswith("const")" is to make sure that the
# incision op is neither of type "const" nor "constexpr_*" ops that
# are used to store compressed weights
if (cumulative_size_in_mb > half_size and not op.op_type.startswith("const")
and len(op.outputs) == 1
and len(op.outputs[0].child_ops) == 1):
op_idx = main_block.operations.index(op)
return op_idx, cumulative_size_in_mb, total_size_in_mb
def _get_first_chunk_outputs(block, op_idx):
# Get the list of all vars that go across from first program (all ops from 0 to op_idx (inclusive))
# to the second program (all ops from op_idx+1 till the end). These all vars need to be made the output
# of the first program and the input of the second program
boundary_vars = set()
block.operations = list(block.operations)
for i in range(op_idx + 1):
op = block.operations[i]
if not op.op_type.startswith("const"):
for var in op.outputs:
if var.val is None: # only consider non const vars
for child_op in var.child_ops:
child_op_idx = block.operations.index(child_op)
if child_op_idx > op_idx:
boundary_vars.add(var)
return list(boundary_vars)
@block_context_manager
def _add_fp32_casts(block, boundary_vars):
new_boundary_vars = []
for var in boundary_vars:
if var.dtype != types.fp16:
new_boundary_vars.append(var)
else:
fp32_var = mb.cast(x=var, dtype="fp32", name=var.name)
new_boundary_vars.append(fp32_var)
return new_boundary_vars
def _make_first_chunk_prog(prog, op_idx):
""" Build first chunk by declaring early outputs and removing unused subgraph
"""
block = prog.functions["main"]
boundary_vars = _get_first_chunk_outputs(block, op_idx)
# Due to possible numerical issues, cast any fp16 var to fp32
new_boundary_vars = _add_fp32_casts(block, boundary_vars)
block.outputs.clear()
block.set_outputs(new_boundary_vars)
PASS_REGISTRY["common::dead_code_elimination"](prog)
return prog
def _make_second_chunk_prog(prog, op_idx):
""" Build second chunk by rebuilding a pristine MIL Program from MLModel
"""
block = prog.functions["main"]
block.opset_version = ct.target.iOS16
# First chunk outputs are second chunk inputs (e.g. skip connections)
boundary_vars = _get_first_chunk_outputs(block, op_idx)
# This op will not be included in this program. Its output var will be made into an input
block.operations = list(block.operations)
boundary_op = block.operations[op_idx]
# Add all boundary ops as inputs
with block:
for var in boundary_vars:
new_placeholder = Placeholder(
sym_shape=var.shape,
dtype=var.dtype if var.dtype != types.fp16 else types.fp32,
name=var.name,
)
block._input_dict[
new_placeholder.outputs[0].name] = new_placeholder.outputs[0]
block.function_inputs = tuple(block._input_dict.values())
new_var = None
if var.dtype == types.fp16:
new_var = mb.cast(x=new_placeholder.outputs[0],
dtype="fp16",
before_op=var.op)
else:
new_var = new_placeholder.outputs[0]
block.replace_uses_of_var_after_op(
anchor_op=boundary_op,
old_var=var,
new_var=new_var,
# This is needed if the program contains "constexpr_*" ops. In normal cases, there are stricter
# rules for removing them, and their presence may prevent replacing this var.
# However in this case, since we want to remove all the ops in chunk 1, we can safely
# set this to True.
force_replace=True,
)
PASS_REGISTRY["common::dead_code_elimination"](prog)
# Remove any unused inputs
new_input_dict = OrderedDict()
for k, v in block._input_dict.items():
if len(v.child_ops) > 0:
new_input_dict[k] = v
block._input_dict = new_input_dict
block.function_inputs = tuple(block._input_dict.values())
return prog
def _legacy_model_chunking(args):
# TODO: Remove this method after setting the coremltools dependency >= 8.0
os.makedirs(args.o, exist_ok=True)
# Check filename extension
mlpackage_name = os.path.basename(args.mlpackage_path)
name, ext = os.path.splitext(mlpackage_name)
assert ext == ".mlpackage", f"`--mlpackage-path` (args.mlpackage_path) is not an .mlpackage file"
# Load CoreML model
logger.info("Loading model from {}".format(args.mlpackage_path))
start_ = time.time()
model = ct.models.MLModel(
args.mlpackage_path,
compute_units=ct.ComputeUnit.CPU_ONLY,
)
logger.info(
f"Loading {args.mlpackage_path} took {time.time() - start_:.1f} seconds"
)
# Load the MIL Program from MLModel
prog = _load_prog_from_mlmodel(model)
# Compute the incision point by bisecting the program based on weights size
op_idx, first_chunk_weights_size, total_weights_size = _get_op_idx_split_location(
prog)
main_block = prog.functions["main"]
incision_op = main_block.operations[op_idx]
logger.info(f"{args.mlpackage_path} will chunked into two pieces.")
logger.info(
f"The incision op: name={incision_op.name}, type={incision_op.op_type}, index={op_idx}/{len(main_block.operations)}"
)
logger.info(f"First chunk size = {first_chunk_weights_size:.2f} MB")
logger.info(
f"Second chunk size = {total_weights_size - first_chunk_weights_size:.2f} MB"
)
# Build first chunk (in-place modifies prog by declaring early exits and removing unused subgraph)
prog_chunk1 = _make_first_chunk_prog(prog, op_idx)
# Build the second chunk
prog_chunk2 = _make_second_chunk_prog(_load_prog_from_mlmodel(model),
op_idx)
if not args.check_output_correctness:
# Original model no longer needed in memory
del model
gc.collect()
# Convert the MIL Program objects into MLModels
logger.info("Converting the two programs")
model_chunk1 = ct.convert(
prog_chunk1,
convert_to="mlprogram",
compute_units=ct.ComputeUnit.CPU_ONLY,
minimum_deployment_target=ct.target.iOS16,
)
del prog_chunk1
gc.collect()
logger.info("Conversion of first chunk done.")
model_chunk2 = ct.convert(
prog_chunk2,
convert_to="mlprogram",
compute_units=ct.ComputeUnit.CPU_ONLY,
minimum_deployment_target=ct.target.iOS16,
)
del prog_chunk2
gc.collect()
logger.info("Conversion of second chunk done.")
# Verify output correctness
if args.check_output_correctness:
logger.info("Verifying output correctness of chunks")
_verify_output_correctness_of_chunks(
full_model=model,
first_chunk_model=model_chunk1,
second_chunk_model=model_chunk2,
)
if args.merge_chunks_in_pipeline_model:
# Make a single pipeline model to manage the model chunks
pipeline_model = ct.utils.make_pipeline(model_chunk1, model_chunk2)
out_path_pipeline = os.path.join(args.o, name + "_chunked_pipeline.mlpackage")
# Save and reload to ensure CPU placement
pipeline_model.save(out_path_pipeline)
pipeline_model = ct.models.MLModel(out_path_pipeline, compute_units=ct.ComputeUnit.CPU_ONLY)
if args.check_output_correctness:
logger.info("Verifying output correctness of pipeline model")
_verify_output_correctness_of_chunks(
full_model=model,
pipeline_model=pipeline_model,
)
else:
# Save the chunked models to disk
out_path_chunk1 = os.path.join(args.o, name + "_chunk1.mlpackage")
out_path_chunk2 = os.path.join(args.o, name + "_chunk2.mlpackage")
logger.info(
f"Saved chunks in {args.o} with the suffix _chunk1.mlpackage and _chunk2.mlpackage"
)
model_chunk1.save(out_path_chunk1)
model_chunk2.save(out_path_chunk2)
logger.info("Done.")
def main(args):
ct_version = ct.__version__
if ct_version != "8.0b2" and ct_version < "8.0":
# With coremltools version <= 8.0b1,
# we use the legacy implementation.
# TODO: Remove the logic after setting the coremltools dependency >= 8.0.
logger.info(
f"coremltools version {ct_version} detected. Recommended upgrading the package version to "
f"'8.0b2' when you running chunk_mlprogram.py script for the latest supports and bug fixes."
)
_legacy_model_chunking(args)
else:
# Starting from coremltools==8.0b2, there is this `bisect_model` API that
# we can directly call into.
from coremltools.models.utils import bisect_model
logger.info(f"Start chunking model {args.mlpackage_path} into two pieces.")
ct.models.utils.bisect_model(
model=args.mlpackage_path,
output_dir=args.o,
merge_chunks_to_pipeline=args.merge_chunks_in_pipeline_model,
check_output_correctness=args.check_output_correctness,
)
logger.info(f"Model chunking is done.")
# Remove original (non-chunked) model if requested
if args.remove_original:
logger.info(
"Removing original (non-chunked) model at {args.mlpackage_path}")
shutil.rmtree(args.mlpackage_path)
logger.info("Done.")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--mlpackage-path",
required=True,
help=
"Path to the mlpackage file to be split into two mlpackages of approximately same file size.",
)
parser.add_argument(
"-o",
required=True,
help=
"Path to output directory where the two model chunks should be saved.",
)
parser.add_argument(
"--remove-original",
action="store_true",
help=
"If specified, removes the original (non-chunked) model to avoid duplicating storage."
)
parser.add_argument(
"--check-output-correctness",
action="store_true",
help=
("If specified, compares the outputs of original Core ML model with that of pipelined CoreML model chunks and reports PSNR in dB. ",
"Enabling this feature uses more memory. Disable it if your machine runs out of memory."
))
parser.add_argument(
"--merge-chunks-in-pipeline-model",
action="store_true",
help=
("If specified, model chunks are managed inside a single pipeline model for easier asset maintenance"
))
args = parser.parse_args()
main(args)
python_coreml_stable_diffusion/controlnet.py
+250
View File
@@ -0,0 +1,250 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers import ModelMixin
import torch
import torch.nn as nn
import torch.nn.functional as F
from .unet import Timesteps, TimestepEmbedding, get_down_block, UNetMidBlock2DCrossAttn, linear_to_conv2d_map
class ControlNetConditioningEmbedding(nn.Module):
def __init__(
self,
conditioning_embedding_channels,
conditioning_channels=3,
block_out_channels=(16, 32, 96, 256),
):
super().__init__()
self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1)
self.blocks = nn.ModuleList([])
for i in range(len(block_out_channels) - 1):
channel_in = block_out_channels[i]
channel_out = block_out_channels[i + 1]
self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1))
self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2))
self.conv_out = nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1)
def forward(self, conditioning):
embedding = self.conv_in(conditioning)
embedding = F.silu(embedding)
for block in self.blocks:
embedding = block(embedding)
embedding = F.silu(embedding)
embedding = self.conv_out(embedding)
return embedding
class ControlNetModel(ModelMixin, ConfigMixin):
@register_to_config
def __init__(
self,
in_channels=4,
flip_sin_to_cos=True,
freq_shift=0,
down_block_types=(
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"CrossAttnDownBlock2D",
"DownBlock2D",
),
only_cross_attention=False,
block_out_channels=(320, 640, 1280, 1280),
layers_per_block=2,
downsample_padding=1,
mid_block_scale_factor=1,
act_fn="silu",
norm_num_groups=32,
norm_eps=1e-5,
cross_attention_dim=1280,
transformer_layers_per_block=1,
attention_head_dim=8,
use_linear_projection=False,
upcast_attention=False,
resnet_time_scale_shift="default",
conditioning_embedding_out_channels=(16, 32, 96, 256),
**kwargs,
):
super().__init__()
# Check inputs
if len(block_out_channels) != len(down_block_types):
raise ValueError(
f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
)
if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
raise ValueError(
f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
)
if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
raise ValueError(
f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
)
self._register_load_state_dict_pre_hook(linear_to_conv2d_map)
# input
conv_in_kernel = 3
conv_in_padding = (conv_in_kernel - 1) // 2
self.conv_in = nn.Conv2d(
in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
)
# time
time_embed_dim = block_out_channels[0] * 4
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
timestep_input_dim = block_out_channels[0]
self.time_embedding = TimestepEmbedding(
timestep_input_dim,
time_embed_dim,
)
# control net conditioning embedding
self.controlnet_cond_embedding = ControlNetConditioningEmbedding(
conditioning_embedding_channels=block_out_channels[0],
block_out_channels=conditioning_embedding_out_channels,
)
self.down_blocks = nn.ModuleList([])
self.controlnet_down_blocks = nn.ModuleList([])
if isinstance(only_cross_attention, bool):
only_cross_attention = [only_cross_attention] * len(down_block_types)
if isinstance(attention_head_dim, int):
attention_head_dim = (attention_head_dim,) * len(down_block_types)
if isinstance(transformer_layers_per_block, int):
transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
# down
output_channel = block_out_channels[0]
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
self.controlnet_down_blocks.append(controlnet_block)
for i, down_block_type in enumerate(down_block_types):
input_channel = output_channel
output_channel = block_out_channels[i]
is_final_block = i == len(block_out_channels) - 1
down_block = get_down_block(
down_block_type,
transformer_layers_per_block=transformer_layers_per_block[i],
num_layers=layers_per_block,
in_channels=input_channel,
out_channels=output_channel,
temb_channels=time_embed_dim,
resnet_eps=norm_eps,
resnet_act_fn=act_fn,
cross_attention_dim=cross_attention_dim,
attn_num_head_channels=attention_head_dim[i],
downsample_padding=downsample_padding,
add_downsample=not is_final_block,
)
self.down_blocks.append(down_block)
for _ in range(layers_per_block):
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
self.controlnet_down_blocks.append(controlnet_block)
if not is_final_block:
controlnet_block = nn.Conv2d(output_channel, output_channel, kernel_size=1)
self.controlnet_down_blocks.append(controlnet_block)
# mid
mid_block_channel = block_out_channels[-1]
controlnet_block = nn.Conv2d(mid_block_channel, mid_block_channel, kernel_size=1)
self.controlnet_mid_block = controlnet_block
self.mid_block = UNetMidBlock2DCrossAttn(
in_channels=mid_block_channel,
temb_channels=time_embed_dim,
resnet_eps=norm_eps,
resnet_act_fn=act_fn,
output_scale_factor=mid_block_scale_factor,
resnet_time_scale_shift=resnet_time_scale_shift,
cross_attention_dim=cross_attention_dim,
attn_num_head_channels=attention_head_dim[-1],
resnet_groups=norm_num_groups,
use_linear_projection=use_linear_projection,
upcast_attention=upcast_attention,
)
def get_num_residuals(self):
num_res = 2 # initial sample + mid block
for down_block in self.down_blocks:
num_res += len(down_block.resnets)
if hasattr(down_block, "downsamplers") and down_block.downsamplers is not None:
num_res += len(down_block.downsamplers)
return num_res
def forward(
self,
sample,
timestep,
encoder_hidden_states,
controlnet_cond,
):
# 1. time
t_emb = self.time_proj(timestep)
emb = self.time_embedding(t_emb)
# 2. pre-process
sample = self.conv_in(sample)
controlnet_cond = self.controlnet_cond_embedding(controlnet_cond)
sample += controlnet_cond
# 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
# 4. mid
if self.mid_block is not None:
sample = self.mid_block(
sample,
emb,
encoder_hidden_states=encoder_hidden_states,
)
# 5. Control net blocks
controlnet_down_block_res_samples = ()
for down_block_res_sample, controlnet_block in zip(down_block_res_samples, self.controlnet_down_blocks):
down_block_res_sample = controlnet_block(down_block_res_sample)
controlnet_down_block_res_samples += (down_block_res_sample,)
down_block_res_samples = controlnet_down_block_res_samples
mid_block_res_sample = self.controlnet_mid_block(sample)
return down_block_res_samples, mid_block_res_sample
python_coreml_stable_diffusion/coreml_model.py
+225
View File
@@ -0,0 +1,225 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import coremltools as ct
import logging
import json
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
import numpy as np
import os
import time
import subprocess
import sys
def _macos_version():
"""
Returns macOS version as a tuple of integers. On non-Macs, returns an empty tuple.
"""
if sys.platform == "darwin":
try:
ver_str = subprocess.run(["sw_vers", "-productVersion"], stdout=subprocess.PIPE).stdout.decode('utf-8').strip('\n')
return tuple([int(v) for v in ver_str.split(".")])
except:
raise Exception("Unable to determine the macOS version")
return ()
class CoreMLModel:
""" Wrapper for running CoreML models using coremltools
"""
def __init__(self, model_path, compute_unit, sources='packages', optimization_hints=None):
logger.info(f"Loading {model_path}")
start = time.time()
if sources == 'packages':
assert os.path.exists(model_path) and model_path.endswith(".mlpackage")
self.model = ct.models.MLModel(
model_path,
compute_units=ct.ComputeUnit[compute_unit],
optimization_hints=optimization_hints,
)
DTYPE_MAP = {
65552: np.float16,
65568: np.float32,
131104: np.int32,
}
self.expected_inputs = {
input_tensor.name: {
"shape": tuple(input_tensor.type.multiArrayType.shape),
"dtype": DTYPE_MAP[input_tensor.type.multiArrayType.dataType],
}
for input_tensor in self.model._spec.description.input
}
elif sources == 'compiled':
assert os.path.exists(model_path) and model_path.endswith(".mlmodelc")
self.model = ct.models.CompiledMLModel(
model_path,
compute_units=ct.ComputeUnit[compute_unit],
optimization_hints=optimization_hints,
)
# Grab expected inputs from metadata.json
with open(os.path.join(model_path, 'metadata.json'), 'r') as f:
config = json.load(f)[0]
self.expected_inputs = {
input_tensor['name']: {
"shape": tuple(eval(input_tensor['shape'])),
"dtype": np.dtype(input_tensor['dataType'].lower()),
}
for input_tensor in config['inputSchema']
}
else:
raise ValueError(f'Expected `packages` or `compiled` for sources, received {sources}')
load_time = time.time() - start
logger.info(f"Done. Took {load_time:.1f} seconds.")
if load_time > LOAD_TIME_INFO_MSG_TRIGGER:
logger.info(
"Loading a CoreML model through coremltools triggers compilation every time. "
"The Swift package we provide uses precompiled Core ML models (.mlmodelc) to avoid compile-on-load."
)
def _verify_inputs(self, **kwargs):
for k, v in kwargs.items():
if k in self.expected_inputs:
if not isinstance(v, np.ndarray):
raise TypeError(
f"Expected numpy.ndarray, got {v} for input: {k}")
expected_dtype = self.expected_inputs[k]["dtype"]
if not v.dtype == expected_dtype:
raise TypeError(
f"Expected dtype {expected_dtype}, got {v.dtype} for input: {k}"
)
expected_shape = self.expected_inputs[k]["shape"]
if not v.shape == expected_shape:
raise TypeError(
f"Expected shape {expected_shape}, got {v.shape} for input: {k}"
)
else:
raise ValueError(f"Received unexpected input kwarg: {k}")
def __call__(self, **kwargs):
self._verify_inputs(**kwargs)
return self.model.predict(kwargs)
LOAD_TIME_INFO_MSG_TRIGGER = 10 # seconds
def get_resource_type(resources_dir: str) -> str:
"""
Detect resource type based on filepath extensions.
returns:
`packages`: for .mlpackage resources
'compiled`: for .mlmodelc resources
"""
directories = [f for f in os.listdir(resources_dir) if os.path.isdir(os.path.join(resources_dir, f))]
# consider directories ending with extension
extensions = set([os.path.splitext(e)[1] for e in directories if os.path.splitext(e)[1]])
# if one extension present we may be able to infer sources type
if len(set(extensions)) == 1:
extension = extensions.pop()
else:
raise ValueError(f'Multiple file extensions found at {resources_dir}.'
f'Cannot infer resource type from contents.')
if extension == '.mlpackage':
sources = 'packages'
elif extension == '.mlmodelc':
sources = 'compiled'
else:
raise ValueError(f'Did not find .mlpackage or .mlmodelc at {resources_dir}')
return sources
def _load_mlpackage(submodule_name,
mlpackages_dir,
model_version,
compute_unit,
sources=None):
"""
Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
"""
# if sources not provided, attempt to infer `packages` or `compiled` from the
# resources directory
if sources is None:
sources = get_resource_type(mlpackages_dir)
if sources == 'packages':
logger.info(f"Loading {submodule_name} mlpackage")
fname = f"Stable_Diffusion_version_{model_version}_{submodule_name}.mlpackage".replace(
"/", "_")
mlpackage_path = os.path.join(mlpackages_dir, fname)
if not os.path.exists(mlpackage_path):
raise FileNotFoundError(
f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
elif sources == 'compiled':
logger.info(f"Loading {submodule_name} mlmodelc")
# FixMe: Submodule names and compiled resources names differ. Can change if names match in the future.
submodule_names = ["text_encoder", "text_encoder_2", "unet", "vae_decoder", "vae_encoder", "safety_checker"]
compiled_names = ['TextEncoder', 'TextEncoder2', 'Unet', 'VAEDecoder', 'VAEEncoder', 'SafetyChecker']
name_map = dict(zip(submodule_names, compiled_names))
cname = name_map[submodule_name] + '.mlmodelc'
mlpackage_path = os.path.join(mlpackages_dir, cname)
if not os.path.exists(mlpackage_path):
raise FileNotFoundError(
f"{submodule_name} CoreML model doesn't exist at {mlpackage_path}")
# On macOS 15+, set fast prediction optimization hint for the unet.
optimization_hints = None
if submodule_name == "unet" and _macos_version() >= (15, 0):
optimization_hints = {"specializationStrategy": ct.SpecializationStrategy.FastPrediction}
return CoreMLModel(mlpackage_path,
compute_unit,
sources=sources,
optimization_hints=optimization_hints)
def _load_mlpackage_controlnet(mlpackages_dir, model_version, compute_unit):
""" Load Core ML (mlpackage) models from disk (As exported by torch2coreml.py)
"""
model_name = model_version.replace("/", "_")
logger.info(f"Loading controlnet_{model_name} mlpackage")
fname = f"ControlNet_{model_name}.mlpackage"
mlpackage_path = os.path.join(mlpackages_dir, fname)
if not os.path.exists(mlpackage_path):
raise FileNotFoundError(
f"controlnet_{model_name} CoreML model doesn't exist at {mlpackage_path}")
return CoreMLModel(mlpackage_path, compute_unit)
def get_available_compute_units():
return tuple(cu for cu in ct.ComputeUnit._member_names_)
python_coreml_stable_diffusion/layer_norm.py
+80
View File
@@ -0,0 +1,80 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import torch
import torch.nn as nn
# Reference: https://github.com/apple/ml-ane-transformers/blob/main/ane_transformers/reference/layer_norm.py
class LayerNormANE(nn.Module):
""" LayerNorm optimized for Apple Neural Engine (ANE) execution
Note: This layer only supports normalization over the final dim. It expects `num_channels`
as an argument and not `normalized_shape` which is used by `torch.nn.LayerNorm`.
"""
def __init__(self,
num_channels,
clip_mag=None,
eps=1e-5,
elementwise_affine=True):
"""
Args:
num_channels: Number of channels (C) where the expected input data format is BC1S. S stands for sequence length.
clip_mag: Optional float value to use for clamping the input range before layer norm is applied.
If specified, helps reduce risk of overflow.
eps: Small value to avoid dividing by zero
elementwise_affine: If true, adds learnable channel-wise shift (bias) and scale (weight) parameters
"""
super().__init__()
# Principle 1: Picking the Right Data Format (machinelearning.apple.com/research/apple-neural-engine)
self.expected_rank = len("BC1S")
self.num_channels = num_channels
self.eps = eps
self.clip_mag = clip_mag
self.elementwise_affine = elementwise_affine
if self.elementwise_affine:
self.weight = nn.Parameter(torch.Tensor(num_channels))
self.bias = nn.Parameter(torch.Tensor(num_channels))
self._reset_parameters()
def _reset_parameters(self):
if self.elementwise_affine:
nn.init.ones_(self.weight)
nn.init.zeros_(self.bias)
def forward(self, inputs):
input_rank = len(inputs.size())
# Principle 1: Picking the Right Data Format (machinelearning.apple.com/research/apple-neural-engine)
# Migrate the data format from BSC to BC1S (most conducive to ANE)
if input_rank == 3 and inputs.size(2) == self.num_channels:
inputs = inputs.transpose(1, 2).unsqueeze(2)
input_rank = len(inputs.size())
assert input_rank == self.expected_rank
assert inputs.size(1) == self.num_channels
if self.clip_mag is not None:
inputs.clamp_(-self.clip_mag, self.clip_mag)
channels_mean = inputs.mean(dim=1, keepdims=True)
zero_mean = inputs - channels_mean
zero_mean_sq = zero_mean * zero_mean
denom = (zero_mean_sq.mean(dim=1, keepdims=True) + self.eps).rsqrt()
out = zero_mean * denom
if self.elementwise_affine:
out = (out + self.bias.view(1, self.num_channels, 1, 1)
) * self.weight.view(1, self.num_channels, 1, 1)
return out
python_coreml_stable_diffusion/mixed_bit_compression_apply.py
+133
View File
@@ -0,0 +1,133 @@
import argparse
import gc
import json
import logging
import os
import coremltools as ct
import coremltools.optimize.coreml as cto
import numpy as np
from python_coreml_stable_diffusion.torch2coreml import get_pipeline
from python_coreml_stable_diffusion.mixed_bit_compression_pre_analysis import (
NBITS,
PALETTIZE_MIN_SIZE as MIN_SIZE
)
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
def main(args):
# Load Core ML model
coreml_model = ct.models.MLModel(args.mlpackage_path, compute_units=ct.ComputeUnit.CPU_ONLY)
logger.info(f"Loaded {args.mlpackage_path}")
# Load palettization recipe
with open(args.pre_analysis_json_path, 'r') as f:
pre_analysis = json.load(f)
if args.selected_recipe not in list(pre_analysis["recipes"]):
raise KeyError(
f"--selected-recipe ({args.selected_recipe}) not found in "
f"--pre-analysis-json-path ({args.pre_analysis_json_path}). "
f" Available recipes: {list(pre_analysis['recipes'])}"
)
recipe = pre_analysis["recipes"][args.selected_recipe]
assert all(nbits in NBITS + [16] for nbits in recipe.values()), \
f"Some nbits values in the recipe are illegal. Allowed values: {NBITS}"
# Hash tensors to be able to match torch tensor names to mil tensors
def get_tensor_hash(tensor):
assert tensor.dtype == np.float16
return tensor.ravel()[0] + np.prod(tensor.shape)
args.model_version = pre_analysis["model_version"]
pipe = get_pipeline(args)
torch_model = pipe.unet
hashed_recipe = {}
for torch_module_name, nbits in recipe.items():
tensor = [
tensor.cpu().numpy().astype(np.float16) for name,tensor in torch_model.named_parameters()
if name == torch_module_name + '.weight'
][0]
hashed_recipe[get_tensor_hash(tensor)] = nbits
del pipe
gc.collect()
op_name_configs = {}
weight_metadata = cto.get_weights_metadata(coreml_model, weight_threshold=MIN_SIZE)
hashes = np.array(list(hashed_recipe))
for name, metadata in weight_metadata.items():
# Look up target bits for this weight
tensor_hash = get_tensor_hash(metadata.val)
pdist = np.abs(hashes - tensor_hash)
assert(pdist.min() < 0.01)
matched = pdist.argmin()
target_nbits = hashed_recipe[hashes[matched]]
if target_nbits == 16:
continue
op_name_configs[name] = cto.OpPalettizerConfig(
mode="kmeans",
nbits=target_nbits,
weight_threshold=int(MIN_SIZE)
)
config = ct.optimize.coreml.OptimizationConfig(op_name_configs=op_name_configs)
coreml_model = ct.optimize.coreml.palettize_weights(coreml_model, config)
coreml_model.save(args.o)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"-o",
required=True,
help="Output directory to save the custom palettized model"
)
parser.add_argument(
"--mlpackage-path",
required=True,
help="Path to .mlpackage model to be palettized"
)
parser.add_argument(
"--pre-analysis-json-path",
required=True,
type=str,
help=("The JSON file generated by mixed_bit_compression_pre_analysis.py"
))
parser.add_argument(
"--selected-recipe",
required=True,
type=str,
help=("The string key into --pre-analysis-json-path's baselines dict"
))
parser.add_argument(
"--custom-vae-version",
type=str,
default=None,
help=
("Custom VAE checkpoint to override the pipeline's built-in VAE. "
"If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
"No precision override is applied when using a custom VAE."
))
args = parser.parse_args()
if not os.path.exists(args.mlpackage_path):
raise FileNotFoundError
if not os.path.exists(args.pre_analysis_json_path):
raise FileNotFoundError
if not args.pre_analysis_json_path.endswith('.json'):
raise ValueError("--recipe-json-path should end with '.json'")
main(args)
python_coreml_stable_diffusion/mixed_bit_compression_pre_analysis.py
+583
View File
@@ -0,0 +1,583 @@
from collections import OrderedDict
from copy import deepcopy
from functools import partial
import argparse
import gc
import json
import logging
logging.basicConfig()
logger = logging.getLogger()
logger.setLevel('INFO')
import numpy as np
import os
from PIL import Image
from python_coreml_stable_diffusion.torch2coreml import compute_psnr, get_pipeline
import time
import torch
import torch.nn as nn
import requests
torch.set_grad_enabled(False)
from tqdm import tqdm
# Bit-widths the Neural Engine is capable of accelerating
NBITS = [1, 2, 4, 6, 8]
# Minimum number of elements in a weight tensor to be considered for palettization
# (saves pre-analysis time)
PALETTIZE_MIN_SIZE = 1e5
# Signal integrity is computed based on these 4 random prompts
RANDOM_TEST_DATA = [
"a black and brown dog standing outside a door.",
"a person on a motorcycle makes a turn on the track.",
"inflatable boats sit on the arizona river, and on the bank",
"a white cat sitting under a white umbrella",
"black bear standing in a field of grass under a tree.",
"a train that is parked on tracks and has graffiti writing on it, with a mountain range in the background.",
"a cake inside of a pan sitting in an oven.",
"a table with paper plates and flowers in a home",
]
TEST_RESOLUTION = 768
RANDOM_TEST_IMAGE_DATA = [
Image.open(
requests.get(path, stream=True).raw).convert("RGB").resize(
(TEST_RESOLUTION, TEST_RESOLUTION), Image.LANCZOS
) for path in [
"http://farm1.staticflickr.com/106/298138827_19bb723252_z.jpg",
"http://farm4.staticflickr.com/3772/9666116202_648cd752d6_z.jpg",
"http://farm3.staticflickr.com/2238/2472574092_f5534bb2f7_z.jpg",
"http://farm1.staticflickr.com/220/475442674_47d81fdc2c_z.jpg",
"http://farm8.staticflickr.com/7231/7359341784_4c5358197f_z.jpg",
"http://farm8.staticflickr.com/7283/8737653089_d0c77b8597_z.jpg",
"http://farm3.staticflickr.com/2454/3989339438_2f32b76ebb_z.jpg",
"http://farm1.staticflickr.com/34/123005230_13051344b1_z.jpg",
]]
# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L602
from coremltools.converters.mil.mil import types
def fake_linear_quantize(val, axis=-1, mode='LINEAR', dtype=types.int8):
from coremltools.optimize.coreml._quantization_passes import AffineQuantParams
from coremltools.converters.mil.mil.types.type_mapping import nptype_from_builtin
val_dtype = val.dtype
def _ensure_numerical_range_and_cast(val, low, high, np_dtype):
'''
For some cases, the computed quantized data might exceed the data range.
For instance, after rounding and addition, we might get `128` for the int8 quantization.
This utility function ensures the val in the data range before doing the cast.
'''
val = np.minimum(val, high)
val = np.maximum(val, low)
return val.astype(np_dtype)
mode_dtype_to_range = {
(types.int8, "LINEAR"): (-128, 127),
(types.int8, "LINEAR_SYMMETRIC"): (-127, 127),
(types.uint8, "LINEAR"): (0, 255),
(types.uint8, "LINEAR_SYMMETRIC"): (0, 254),
}
if not isinstance(val, (np.ndarray, np.generic)):
raise ValueError("Only numpy arrays are supported")
params = AffineQuantParams()
axes = tuple([i for i in range(len(val.shape)) if i != axis])
val_min = np.amin(val, axis=axes, keepdims=True)
val_max = np.amax(val, axis=axes, keepdims=True)
if mode == "LINEAR_SYMMETRIC":
# For the linear_symmetric mode, the range is symmetrical to 0
max_abs = np.maximum(np.abs(val_min), np.abs(val_max))
val_min = -max_abs
val_max = max_abs
else:
assert mode == "LINEAR"
# For the linear mode, we need to make sure the data range contains `0`
val_min = np.minimum(0.0, val_min)
val_max = np.maximum(0.0, val_max)
q_val_min, q_val_max = mode_dtype_to_range[(dtype, mode)]
# Set the zero point to symmetric mode
np_dtype = nptype_from_builtin(dtype)
if mode == "LINEAR_SYMMETRIC":
if dtype == types.int8:
params.zero_point = (0 * np.ones(val_min.shape)).astype(np.int8)
else:
assert dtype == types.uint8
params.zero_point = (127 * np.ones(val_min.shape)).astype(np.uint8)
else:
assert mode == "LINEAR"
params.zero_point = (q_val_min * val_max - q_val_max * val_min) / (val_max - val_min)
params.zero_point = np.round(params.zero_point)
params.zero_point = _ensure_numerical_range_and_cast(params.zero_point, q_val_min, q_val_max, np_dtype)
# compute the params
params.scale = (val_max - val_min) / (q_val_max - q_val_min)
params.scale = params.scale.astype(val.dtype).squeeze()
params.quantized_data = np.round(
val * (q_val_max - q_val_min) / (val_max - val_min)
)
params.quantized_data = (params.quantized_data + params.zero_point)
params.quantized_data = _ensure_numerical_range_and_cast(params.quantized_data, q_val_min, q_val_max, np_dtype)
params.zero_point = params.zero_point.squeeze()
params.axis = axis
return (params.quantized_data.astype(val_dtype) - params.zero_point.astype(val_dtype)) * params.scale
# Copied from https://github.com/apple/coremltools/blob/7.0b1/coremltools/optimize/coreml/_quantization_passes.py#L423
def fake_palettize(module, nbits, in_ngroups=1, out_ngroups=1):
""" Simulate weight palettization
"""
from coremltools.models.neural_network.quantization_utils import _get_kmeans_lookup_table_and_weight
def compress_kmeans(val, nbits):
lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val)
lut = lut.astype(val.dtype)
indices = indices.astype(np.uint8)
return lut, indices
dtype = module.weight.data.dtype
device = module.weight.data.device
val = module.weight.data.cpu().numpy().astype(np.float16)
if out_ngroups == 1 and in_ngroups == 1:
lut, indices = compress_kmeans(val=val, nbits=nbits)
module.weight.data = torch.from_numpy(lut[indices]).reshape(val.shape).to(dtype)
elif out_ngroups > 1 and in_ngroups == 1:
assert val.shape[0] % out_ngroups == 0
rvals = [
compress_kmeans(val=chunked_val, nbits=nbits)
for chunked_val in np.split(val, out_ngroups, axis=0)
]
shape = list(val.shape)
shape[0] = shape[0] // out_ngroups
module.weight.data = torch.cat([
torch.from_numpy(lut[indices]).reshape(shape)
for lut,indices in rvals
], dim=0).to(dtype).to(device)
elif in_ngroups > 1 and out_ngroups == 1:
assert val.shape[1] % in_ngroups == 0
rvals = [
compress_kmeans(val=chunked_val, nbits=nbits)
for chunked_val in np.split(val, in_ngroups, axis=1)
]
shape = list(val.shape)
shape[1] = shape[1] // in_ngroups
module.weight.data = torch.cat([
torch.from_numpy(lut[indices]).reshape(shape)
for lut,indices in rvals
], dim=1).to(dtype).to(device)
else:
raise ValueError(f"in_ngroups={in_ngroups} & out_ngroups={out_ngroups} is illegal!!!")
return torch.from_numpy(val).to(dtype)
def restore_weight(module, value):
device = module.weight.data.device
module.weight.data = value.to(device)
def get_palettizable_modules(unet, min_size=PALETTIZE_MIN_SIZE):
ret = [
(name, getattr(module, 'weight').data.numel()) for name, module in unet.named_modules()
if isinstance(module, (nn.Linear, nn.Conv2d))
if hasattr(module, 'weight') and getattr(module, 'weight').data.numel() > min_size
]
candidates, sizes = [[a for a,b in ret], [b for a,b in ret]]
logger.info(f"{len(candidates)} candidate tensors with {sum(sizes)/1e6} M total params")
return candidates, sizes
def fake_int8_quantize(module):
i = 0
for name, submodule in tqdm(module.named_modules()):
if hasattr(submodule, 'weight'):
i+=1
submodule.weight.data = torch.from_numpy(
fake_linear_quantize(submodule.weight.data.numpy()))
logger.info(f"{i} modules fake int8 quantized")
return module
def fake_nbits_palette(module, nbits):
i = 0
for name, submodule in tqdm(module.named_modules()):
if hasattr(submodule, 'weight'):
i+=1
fake_palettize(submodule, nbits=nbits)
logger.info(f"{i} modules fake {nbits}-bits palettized")
return module
def fake_palette_from_recipe(module, recipe):
tot_bits = 0
tot_numel = 0
for name, submodule in tqdm(module.named_modules()):
if hasattr(submodule, 'weight'):
tot_numel += submodule.weight.numel()
if name in recipe:
nbits = recipe[name]
assert nbits in NBITS + [16]
tot_bits += submodule.weight.numel() * nbits
if nbits == 16:
continue
fake_palettize(submodule, nbits=nbits)
else:
tot_bits += submodule.weight.numel() * 16
logger.info(f"Palettized to {tot_bits/tot_numel:.2f}-bits mixed palette ({tot_bits/8e6} MB) ")
# Globally synced RNG state
rng = torch.Generator()
rng_state = rng.get_state()
def run_pipe(pipe):
if torch.backends.mps.is_available():
device = "mps"
elif torch.cuda.is_available():
device = "cuda"
else:
device = "cpu"
logger.debug(f"Placing pipe in {device}")
global rng, rng_state
rng.set_state(rng_state)
kwargs = dict(
prompt=RANDOM_TEST_DATA,
negative_prompt=[""] * len(RANDOM_TEST_DATA),
num_inference_steps=1,
height=TEST_RESOLUTION,
width=TEST_RESOLUTION,
output_type="latent",
generator=rng
)
if "Img2Img" in pipe.__class__.__name__:
kwargs["image"] = RANDOM_TEST_IMAGE_DATA
kwargs.pop("height")
kwargs.pop("width")
# Run a single denoising step
kwargs["num_inference_steps"] = 4
kwargs["strength"] = 0.25
return np.array([latent.cpu().numpy() for latent in pipe.to(device)(**kwargs).images])
def benchmark_signal_integrity(pipe,
candidates,
nbits,
cumulative,
in_ngroups=1,
out_ngroups=1,
ref_out=None,
):
results = {}
results['metadata'] = {
'nbits': nbits,
'out_ngroups': out_ngroups,
'in_ngroups': in_ngroups,
'cumulative': cumulative,
}
# If reference outputs are not provided, treat current pipe as reference
if ref_out is None:
ref_out = run_pipe(pipe)
for candidate in tqdm(candidates):
palettized = False
for name, module in pipe.unet.named_modules():
if name == candidate:
orig_weight = fake_palettize(
module,
nbits,
out_ngroups=out_ngroups,
in_ngroups=in_ngroups,
)
palettized = True
break
if not palettized:
raise KeyError(name)
test_out = run_pipe(pipe)
if not cumulative:
restore_weight(module, orig_weight)
results[candidate] = [
float(f"{compute_psnr(r,t):.1f}")
for r,t in zip(ref_out, test_out)
]
logger.info(f"{nbits}-bit: {candidate} = {results[candidate]}")
return results
def descending_psnr_order(results):
if 'metadata' in results:
results.pop('metadata')
return OrderedDict(sorted(results.items(), key=lambda items: -sum(items[1])))
def simulate_quant_fn(ref_pipe, quantization_to_simulate):
simulated_pipe = deepcopy(ref_pipe.to('cpu'))
quantization_to_simulate(simulated_pipe.unet)
simulated_out = run_pipe(simulated_pipe)
del simulated_pipe
gc.collect()
ref_out = run_pipe(ref_pipe)
simulated_psnr = sum([
float(f"{compute_psnr(r, t):.1f}")
for r, t in zip(ref_out, simulated_out)
]) / len(ref_out)
return simulated_out, simulated_psnr
def build_recipe(results, sizes, psnr_threshold, default_nbits):
stats = {'nbits': 0}
recipe = {}
for key in results[str(NBITS[0])]:
if key == 'metadata':
continue
achieved_nbits = default_nbits
for nbits in NBITS:
avg_psnr = sum(results[str(nbits)][key])/len(RANDOM_TEST_DATA)
if avg_psnr > psnr_threshold:
achieved_nbits = nbits
break
recipe[key] = achieved_nbits
stats['nbits'] += achieved_nbits * sizes[key]
stats['size_mb'] = stats['nbits'] / (8*1e6)
tot_size = sum(list(sizes.values()))
stats['nbits'] /= tot_size
return recipe, stats
def plot(results, args):
import matplotlib.pyplot as plt
max_model_size = sum(results['cumulative'][str(NBITS[0])]['metadata']['sizes'])
f, ax = plt.subplots(1, 1, figsize=(7, 5))
def compute_x_axis(sizes, nbits, default_nbits):
max_compression_percent = (default_nbits - nbits) / default_nbits
progress = np.cumsum(sizes)
normalized_progress = progress / progress.max()
return normalized_progress * max_compression_percent * 100
# Linear 8-bit baseline and the intercept points for mixed-bit recipes
linear8bit_baseline = results['baselines']['linear_8bit']
# Mark the linear 8-bit baseline
ax.plot(
8 / args.default_nbits * 100,
linear8bit_baseline,
'bx',
markersize=8,
label="8-bit (linear quant)")
# Plot the iso-dB line that matches the 8-bit baseline
ax.plot([0,100], [linear8bit_baseline]*2, '--b')
# Plot non-mixed-bit palettization curves
for idx, nbits in enumerate(NBITS):
size_keys = compute_x_axis(results['cumulative'][str(nbits)]['metadata']['sizes'], nbits, args.default_nbits)
psnr = [
sum(v) / len(RANDOM_TEST_DATA) # avg psnr
for k,v in results['cumulative'][str(nbits)].items() if k != 'metadata'
]
ax.plot(
size_keys,
psnr,
label=f"{nbits}-bit")
# Plot mixed-bit results
mixed_palettes = [
(float(spec.rsplit('_')[1]), psnr)
for spec,psnr in results['baselines'].items()
if 'recipe' in spec
]
mixedbit_sizes = [100. * (1. - a[0] / args.default_nbits) for a in mixed_palettes]
mixedbit_psnrs = [a[1] for a in mixed_palettes]
ax.plot(
mixedbit_sizes,
mixedbit_psnrs,
label="mixed-bit",
)
ax.set_xlabel("Model Size Reduction (%)")
ax.set_ylabel("Signal Integrity (PSNR in dB)")
ax.set_title(args.model_version)
ax.legend()
f.savefig(os.path.join(args.o, f"{args.model_version.replace('/','_')}_psnr_vs_size.png"))
def main(args):
# Initialize pipe
pipe = get_pipeline(args)
# Preserve a pristine copy for reference outputs
ref_pipe = deepcopy(pipe)
if args.default_nbits != 16:
logger.info(f"Palettizing unet to default {args.default_nbits}-bit")
fake_nbits_palette(pipe.unet, args.default_nbits)
logger.info("Done.")
# Cache reference outputs
ref_out = run_pipe(pipe)
# Bookkeeping
os.makedirs(args.o, exist_ok=True)
results = {
'single_layer': {},
'cumulative': {},
'model_version': args.model_version,
}
json_name = f"{args.model_version.replace('/','-')}_palettization_recipe.json"
candidates, sizes = get_palettizable_modules(pipe.unet)
sizes_table = dict(zip(candidates, sizes))
if os.path.isfile(os.path.join(args.o, json_name)):
with open(os.path.join(args.o, json_name), "r") as f:
results = json.load(f)
# Analyze uniform-precision palettization impact on signal integrity
for nbits in NBITS:
if str(nbits) not in results['single_layer']:
# Measure the impact of palettization of each layer independently
results['single_layer'][str(nbits)] = benchmark_signal_integrity(
pipe,
candidates,
nbits,
cumulative=False,
ref_out=ref_out,
)
with open(os.path.join(args.o, json_name), 'w') as f:
json.dump(results, f, indent=2)
# Measure the cumulative impact of palettization based on ascending individual impact computed earlier
sorted_candidates = descending_psnr_order(results['single_layer'][str(nbits)])
if str(nbits) not in results['cumulative']:
results['cumulative'][str(nbits)] = benchmark_signal_integrity(
deepcopy(pipe),
sorted_candidates,
nbits,
cumulative=True,
ref_out=ref_out,
)
results['cumulative'][str(nbits)]['metadata'].update({
'candidates': list(sorted_candidates.keys()),
'sizes': [sizes_table[candidate] for candidate in sorted_candidates],
})
with open(os.path.join(args.o, json_name), 'w') as f:
json.dump(results, f, indent=2)
# Generate uniform-quantization baselines
results['baselines'] = {
"original": simulate_quant_fn(ref_pipe, lambda x: x)[1],
"linear_8bit": simulate_quant_fn(ref_pipe, fake_int8_quantize)[1],
}
with open(os.path.join(args.o, json_name), 'w') as f:
json.dump(results, f, indent=2)
# Generate mixed-bit recipes via decreasing PSNR thresholds
results['recipes'] = {}
recipe_psnr_thresholds = np.linspace(
results['baselines']['original'] - 1,
results['baselines']["linear_8bit"] + 5,
args.num_recipes,
)
for recipe_no, psnr_threshold in enumerate(recipe_psnr_thresholds):
logger.info(f"Building recipe #{recipe_no}")
recipe, stats = build_recipe(
results['cumulative'],
sizes_table,
psnr_threshold,
args.default_nbits,
)
achieved_psnr = simulate_quant_fn(ref_pipe, lambda x: partial(fake_palette_from_recipe, recipe=recipe)(x))[1]
logger.info(
f"Recipe #{recipe_no}: {stats['nbits']:.2f}-bits @ per-layer {psnr_threshold} dB, "
f"end-to-end {achieved_psnr} dB & "
f"{stats['size_mb']:.2f} MB"
)
# Save achieved PSNR and compressed size
recipe_key = f"recipe_{stats['nbits']:.2f}_bit_mixedpalette"
results['baselines'][recipe_key] = float(f"{achieved_psnr:.1f}")
results['recipes'][recipe_key] = recipe
with open(os.path.join(args.o, json_name), 'w') as f:
json.dump(results, f, indent=2)
# Plot model size vs signal integrity
plot(results, args)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"-o",
required=True,
help="Output directory to save the palettization artifacts (recipe json, PSNR plots etc.)"
)
parser.add_argument(
"--model-version",
required=True,
help=
("The pre-trained model checkpoint and configuration to restore. "
"For available versions: https://huggingface.co/models?search=stable-diffusion"
))
parser.add_argument(
"--default-nbits",
help="Default number of bits to use for palettization",
choices=tuple(NBITS + [16]),
default=16,
type=int,
)
parser.add_argument(
"--num-recipes",
help="Maximum number of recipes to generate (with decreasing model size and signal integrity)",
default=7,
type=int,
)
parser.add_argument(
"--custom-vae-version",
type=str,
default=None,
help=
("Custom VAE checkpoint to override the pipeline's built-in VAE. "
"If specified, the specified VAE will be converted instead of the one associated to the `--model-version` checkpoint. "
"No precision override is applied when using a custom VAE."
))
args = parser.parse_args()
main(args)
python_coreml_stable_diffusion/multilingual_projection.py
+60
View File
@@ -0,0 +1,60 @@
from python_coreml_stable_diffusion.torch2coreml import _compile_coreml_model
import argparse
import coremltools as ct
import numpy as np
import os
import torch
import torch.nn as nn
# TODO: Read these values off of the NLContextualEmbedding API to enforce dimensions and track API versioning
MAX_SEQUENCE_LENGTH = 256
EMBED_DIM = 512
BATCH_SIZE = 1
def main(args):
# Layer that was trained to map NLContextualEmbedding to your text_encoder.hidden_size dimensionality
text_encoder_projection = torch.jit.load(args.input_path)
# Prepare random inputs for tracing the network before conversion
random_input = torch.randn(BATCH_SIZE, MAX_SEQUENCE_LENGTH, EMBED_DIM)
# Create a class to bake in the reshape operations required to fit the existing model interface
class TextEncoderProjection(nn.Module):
def __init__(self, proj):
super().__init__()
self.proj = proj
def forward(self, x):
return self.proj(x).transpose(1, 2).unsqueeze(2) # BSC, BC1S
# Trace the torch model
text_encoder_projection = torch.jit.trace(TextEncoderProjection(text_encoder_projection), (random_input,))
# Convert the model to Core ML
mlpackage_path = os.path.join(args.output_dir, "MultilingualTextEncoderProjection.mlpackage")
ct.convert(
text_encoder_projection,
inputs=[ct.TensorType('nlcontextualembeddings_output', shape=(1, MAX_SEQUENCE_LENGTH, EMBED_DIM), dtype=np.float32)],
outputs=[ct.TensorType('encoder_hidden_states', dtype=np.float32)],
minimum_deployment_target=ct.target.macOS14, # NLContextualEmbedding minimum availability build
convert_to='mlprogram',
).save()
# Compile the model and save it under the specified directory
_compile_coreml_model(mlpackage_path, args.output_dir, final_name="MultilingualTextEncoderProjection")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--input-path",
help="Path to the torchscript file that contains the projection layer"
)
parser.add_argument(
"--output-dir",
help="Output directory in which the Core ML model should be saved",
)
args = parser.parse_args()
main(args)
python_coreml_stable_diffusion/pipeline.py
+858
View File
@@ -0,0 +1,858 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import argparse
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.schedulers import (
DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler,
)
from diffusers.schedulers.scheduling_utils import SchedulerMixin
import gc
import inspect
import logging
logging.basicConfig()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
import numpy as np
import os
from python_coreml_stable_diffusion.coreml_model import (
CoreMLModel,
_load_mlpackage,
_load_mlpackage_controlnet,
get_available_compute_units,
)
import time
import torch # Only used for `torch.from_tensor` in `pipe.scheduler.step()`
from transformers import CLIPFeatureExtractor, CLIPTokenizer
from typing import List, Optional, Union, Tuple
from PIL import Image
class CoreMLStableDiffusionPipeline(DiffusionPipeline):
""" Core ML version of
`diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline`
"""
def __init__(
self,
text_encoder: CoreMLModel,
unet: CoreMLModel,
vae_decoder: CoreMLModel,
scheduler: Union[
DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler
],
tokenizer: CLIPTokenizer,
controlnet: Optional[List[CoreMLModel]],
xl: Optional[bool] = False,
force_zeros_for_empty_prompt: Optional[bool] = True,
feature_extractor: Optional[CLIPFeatureExtractor] = None,
safety_checker: Optional[CoreMLModel] = None,
text_encoder_2: Optional[CoreMLModel] = None,
tokenizer_2: Optional[CLIPTokenizer] = None
):
super().__init__()
# Register non-Core ML components of the pipeline similar to the original pipeline
self.register_modules(
tokenizer=tokenizer,
scheduler=scheduler,
feature_extractor=feature_extractor,
)
if safety_checker is None:
# Reproduce original warning:
# https://github.com/huggingface/diffusers/blob/v0.9.0/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L119
logger.warning(
f"You have disabled the safety checker for {self.__class__} by passing `safety_checker=None`. Ensure"
" that you abide to the conditions of the Stable Diffusion license and do not expose unfiltered"
" results in services or applications open to the public. Both the diffusers team and Hugging Face"
" strongly recommend to keep the safety filter enabled in all public facing circumstances, disabling"
" it only for use-cases that involve analyzing network behavior or auditing its results. For more"
" information, please have a look at https://github.com/huggingface/diffusers/pull/254 ."
)
self.xl = xl
self.force_zeros_for_empty_prompt = force_zeros_for_empty_prompt
# Register Core ML components of the pipeline
self.safety_checker = safety_checker
self.text_encoder = text_encoder
self.text_encoder_2 = text_encoder_2
self.tokenizer_2 = tokenizer_2
self.unet = unet
self.unet.in_channels = self.unet.expected_inputs["sample"]["shape"][1]
self.controlnet = controlnet
self.vae_decoder = vae_decoder
VAE_DECODER_UPSAMPLE_FACTOR = 8
# In PyTorch, users can determine the tensor shapes dynamically by default
# In CoreML, tensors have static shapes unless flexible shapes were used during export
# See https://coremltools.readme.io/docs/flexible-inputs
latent_h, latent_w = self.unet.expected_inputs["sample"]["shape"][2:]
self.height = latent_h * VAE_DECODER_UPSAMPLE_FACTOR
self.width = latent_w * VAE_DECODER_UPSAMPLE_FACTOR
logger.info(
f"Stable Diffusion configured to generate {self.height}x{self.width} images"
)
def _encode_prompt(self,
prompt,
prompt_2: Optional[str] = None,
do_classifier_free_guidance: bool = True,
negative_prompt: Optional[str] = None,
negative_prompt_2: Optional[str] = None,
):
batch_size = len(prompt) if isinstance(prompt, list) else 1
if self.xl is True:
prompts = [prompt, prompt_2] if prompt_2 is not None else [prompt, prompt]
# refiner uses only one tokenizer and text encoder (tokenizer_2 and text_encoder_2)
tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
text_encoders = [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [
self.text_encoder_2]
hidden_state_key = 'hidden_embeds'
else:
prompts = [prompt]
tokenizers = [self.tokenizer]
text_encoders = [self.text_encoder]
hidden_state_key = 'last_hidden_state'
prompt_embeds_list = []
for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="np",
)
text_input_ids = text_inputs.input_ids
# tokenize without max_length to catch any truncation
untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="np").input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not np.equal(
text_input_ids, untruncated_ids
):
removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1: -1])
logger.warning(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {tokenizer.model_max_length} tokens: {removed_text}"
)
embeddings = text_encoder(input_ids=text_input_ids.astype(np.float32))
prompt_embeds_list.append(embeddings[hidden_state_key])
# We are only ALWAYS interested in the pooled output of the final text encoder
if self.xl:
pooled_prompt_embeds = embeddings['pooled_outputs']
prompt_embeds = np.concatenate(prompt_embeds_list, axis=-1)
if do_classifier_free_guidance and negative_prompt is None and self.force_zeros_for_empty_prompt:
negative_prompt_embeds = np.zeros_like(prompt_embeds)
if self.xl:
negative_pooled_prompt_embeds = np.zeros_like(pooled_prompt_embeds)
elif do_classifier_free_guidance:
negative_prompt = negative_prompt or ""
negative_prompt_2 = negative_prompt_2 or negative_prompt
# normalize str to list
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
negative_prompt_2 = (
batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
)
uncond_tokens: List[str]
if prompts is not None and type(prompts) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elif batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`.")
else:
uncond_tokens = [negative_prompt, negative_prompt_2]
negative_prompt_embeds_list = []
for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
max_length = prompt_embeds.shape[1]
uncond_input = tokenizer(
negative_prompt,
padding="max_length",
max_length=max_length,
truncation=True,
return_tensors="np",
)
uncond_input_ids = uncond_input.input_ids
negative_embeddings = text_encoder(
input_ids=uncond_input_ids.astype(np.float32)
)
negative_text_embeddings = negative_embeddings[hidden_state_key]
negative_prompt_embeds_list.append(negative_text_embeddings)
# We are only ALWAYS interested in the pooled output of the final text encoder
if self.xl:
negative_pooled_prompt_embeds = negative_embeddings['pooled_outputs']
negative_prompt_embeds = np.concatenate(negative_prompt_embeds_list, axis=-1)
if do_classifier_free_guidance:
# For classifier free guidance, we need to do two forward passes.
# Here we concatenate the unconditional and text embeddings into a single batch
# to avoid doing two forward passes
prompt_embeds = np.concatenate(
[negative_prompt_embeds, prompt_embeds])
if self.xl:
pooled_prompt_embeds = np.concatenate(
[negative_pooled_prompt_embeds, pooled_prompt_embeds])
prompt_embeddings = prompt_embeds.transpose(0, 2, 1)[:, :, None, :]
if self.xl:
return prompt_embeddings, pooled_prompt_embeds
else:
return prompt_embeddings, None
def run_controlnet(self,
sample,
timestep,
encoder_hidden_states,
controlnet_cond,
output_dtype=np.float16):
if not self.controlnet:
raise ValueError(
"Conditions for controlnet are given but the pipeline has no controlnet modules")
for i, (module, cond) in enumerate(zip(self.controlnet, controlnet_cond)):
module_outputs = module(
sample=sample.astype(np.float16),
timestep=timestep.astype(np.float16),
encoder_hidden_states=encoder_hidden_states.astype(np.float16),
controlnet_cond=cond.astype(np.float16),
)
if i == 0:
outputs = module_outputs
else:
for key in outputs.keys():
outputs[key] += module_outputs[key]
outputs = {k: v.astype(output_dtype) for k, v in outputs.items()}
return outputs
def run_safety_checker(self, image):
if self.safety_checker is not None:
safety_checker_input = self.feature_extractor(
self.numpy_to_pil(image),
return_tensors="np",
)
safety_checker_outputs = self.safety_checker(
clip_input=safety_checker_input.pixel_values.astype(
np.float16),
images=image.astype(np.float16),
adjustment=np.array([0.]).astype(
np.float16), # defaults to 0 in original pipeline
)
# Unpack dict
has_nsfw_concept = safety_checker_outputs["has_nsfw_concepts"]
image = safety_checker_outputs["filtered_images"]
concept_scores = safety_checker_outputs["concept_scores"]
logger.info(
f"Generated image has nsfw concept={has_nsfw_concept.any()}")
else:
has_nsfw_concept = None
return image, has_nsfw_concept
def decode_latents(self, latents):
latents = 1 / 0.18215 * latents
dtype = self.vae_decoder.expected_inputs['z']['dtype']
image = self.vae_decoder(z=latents.astype(dtype))["image"]
image = np.clip(image / 2 + 0.5, 0, 1)
image = image.transpose((0, 2, 3, 1))
return image
def prepare_latents(self,
batch_size,
num_channels_latents,
height,
width,
latents=None):
latents_shape = (batch_size, num_channels_latents, self.height // 8,
self.width // 8)
if latents is None:
latents = np.random.randn(*latents_shape).astype(np.float16)
elif latents.shape != latents_shape:
raise ValueError(
f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}"
)
init_noise = self.scheduler.init_noise_sigma
if isinstance(init_noise, torch.Tensor):
init_noise = init_noise.numpy()
latents = latents * init_noise
return latents
def prepare_control_cond(self,
controlnet_cond,
do_classifier_free_guidance,
batch_size,
num_images_per_prompt):
processed_cond_list = []
for cond in controlnet_cond:
cond = np.stack([cond] * batch_size * num_images_per_prompt)
if do_classifier_free_guidance:
cond = np.concatenate([cond] * 2)
processed_cond_list.append(cond)
return processed_cond_list
def check_inputs(self, prompt, height, width, callback_steps):
if height != self.height or width != self.width:
logger.warning(
"`height` and `width` dimensions (of the output image tensor) are fixed when exporting the Core ML models " \
"unless flexible shapes are used during export (https://coremltools.readme.io/docs/flexible-inputs). " \
"This pipeline was provided with Core ML models that generate {self.height}x{self.width} images (user requested {height}x{width})"
)
if not isinstance(prompt, str) and not isinstance(prompt, list):
raise ValueError(
f"`prompt` has to be of type `str` or `list` but is {type(prompt)}"
)
if height % 8 != 0 or width % 8 != 0:
raise ValueError(
f"`height` and `width` have to be divisible by 8 but are {height} and {width}."
)
if (callback_steps is None) or (callback_steps is not None and
(not isinstance(callback_steps, int)
or callback_steps <= 0)):
raise ValueError(
f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
f" {type(callback_steps)}.")
def prepare_extra_step_kwargs(self, eta):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
# and should be between [0, 1]
accepts_eta = "eta" in set(
inspect.signature(self.scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
return extra_step_kwargs
def _get_add_time_ids(self, original_size, crops_coords_top_left, target_size, dtype):
add_time_ids = list(original_size + crops_coords_top_left + target_size)
add_time_ids = np.array(add_time_ids).astype(dtype)
return add_time_ids
def __call__(
self,
prompt,
height=512,
width=512,
num_inference_steps=50,
guidance_scale=7.5,
negative_prompt=None,
num_images_per_prompt=1,
eta=0.0,
latents=None,
output_type="pil",
return_dict=True,
callback=None,
callback_steps=1,
controlnet_cond=None,
original_size: Optional[Tuple[int, int]]=None,
crops_coords_top_left: Tuple[int, int]=(0, 0),
target_size: Optional[Tuple[int, int]]=None,
unet_batch_one=False,
**kwargs,
):
# 1. Check inputs. Raise error if not correct
self.check_inputs(prompt, height, width, callback_steps)
height = height or self.height
width = width or self.width
original_size = original_size or (height, width)
target_size = target_size or (height, width)
# 2. Define call parameters
batch_size = 1 if isinstance(prompt, str) else len(prompt)
if batch_size > 1 or num_images_per_prompt > 1:
raise NotImplementedError(
"For batched generation of multiple images and/or multiple prompts, please refer to the Swift package."
)
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input prompt
text_embeddings, pooled_prompt_embeds = self._encode_prompt(
prompt=prompt,
prompt_2=None,
do_classifier_free_guidance=do_classifier_free_guidance,
negative_prompt=negative_prompt,
negative_prompt_2=None
)
# 4. Prepare XL kwargs if needed
unet_additional_kwargs = {}
# we add pooled prompt embeds + time_ids to unet kwargs
if self.xl:
add_text_embeds = pooled_prompt_embeds
add_time_ids = self._get_add_time_ids(original_size, crops_coords_top_left, target_size,
text_embeddings.dtype)
if do_classifier_free_guidance:
# TODO: This checks if the time_ids input is looking for time_ids.shape == (12,) or (2, 6)
# Remove once model input shapes are ubiquitous
if len(self.unet.expected_inputs['time_ids']['shape']) > 1:
add_time_ids = [add_time_ids]
add_time_ids = np.concatenate([add_time_ids, add_time_ids])
unet_additional_kwargs.update({'text_embeds': add_text_embeds.astype(np.float16),
'time_ids': add_time_ids.astype(np.float16)})
# 5. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps)
timesteps = self.scheduler.timesteps
# 6. Prepare latent variables and controlnet cond
num_channels_latents = self.unet.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
latents,
)
if controlnet_cond:
controlnet_cond = self.prepare_control_cond(
controlnet_cond,
do_classifier_free_guidance,
batch_size,
num_images_per_prompt,
)
# 7. Prepare extra step kwargs
extra_step_kwargs = self.prepare_extra_step_kwargs(eta)
# 8. Denoising loop
for i, t in enumerate(self.progress_bar(timesteps)):
# expand the latents if we are doing classifier free guidance
latent_model_input = np.concatenate(
[latents] * 2) if do_classifier_free_guidance else latents
latent_model_input = self.scheduler.scale_model_input(
latent_model_input, t)
if isinstance(latent_model_input, torch.Tensor):
latent_model_input = latent_model_input.numpy()
if do_classifier_free_guidance:
timestep = np.array([t, t], np.float16)
else:
timestep = np.array([t,], np.float16)
# controlnet
if controlnet_cond:
control_net_additional_residuals = self.run_controlnet(
sample=latent_model_input,
timestep=timestep,
encoder_hidden_states=text_embeddings,
controlnet_cond=controlnet_cond,
)
else:
control_net_additional_residuals = {}
# predict the noise residual
unet_additional_kwargs.update(control_net_additional_residuals)
# get prediction from unet
if not (unet_batch_one and do_classifier_free_guidance):
noise_pred = self.unet(
sample=latent_model_input.astype(np.float16),
timestep=timestep,
encoder_hidden_states=text_embeddings.astype(np.float16),
**unet_additional_kwargs,
)["noise_pred"]
if do_classifier_free_guidance:
noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2)
else:
# query unet sequentially
latent_model_input = latent_model_input.astype(np.float16)
text_embeddings = text_embeddings.astype(np.float16)
timestep = np.array([t,], np.float16)
noise_pred_uncond = self.unet(
sample=np.expand_dims(latent_model_input[0], axis=0),
timestep=timestep,
encoder_hidden_states=np.expand_dims(text_embeddings[0], axis=0),
**unet_additional_kwargs,
)["noise_pred"]
noise_pred_text = self.unet(
sample=np.expand_dims(latent_model_input[1], axis=0),
timestep=timestep,
encoder_hidden_states=np.expand_dims(text_embeddings[1], axis=0),
**unet_additional_kwargs,
)["noise_pred"]
# perform guidance
if do_classifier_free_guidance:
noise_pred = noise_pred_uncond + guidance_scale * (
noise_pred_text - noise_pred_uncond)
# compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(torch.from_numpy(noise_pred),
t,
torch.from_numpy(latents),
**extra_step_kwargs,
).prev_sample.numpy()
# call the callback, if provided
if callback is not None and i % callback_steps == 0:
callback(i, t, latents)
# 8. Post-processing
image = self.decode_latents(latents)
# 9. Run safety checker
image, has_nsfw_concept = self.run_safety_checker(image)
# 10. Convert to PIL
if output_type == "pil":
image = self.numpy_to_pil(image)
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(
images=image, nsfw_content_detected=has_nsfw_concept)
def get_available_schedulers():
schedulers = {}
for scheduler in [DDIMScheduler,
DPMSolverMultistepScheduler,
EulerAncestralDiscreteScheduler,
EulerDiscreteScheduler,
LMSDiscreteScheduler,
PNDMScheduler]:
schedulers[scheduler().__class__.__name__.replace("Scheduler", "")] = scheduler
return schedulers
SCHEDULER_MAP = get_available_schedulers()
def get_coreml_pipe(pytorch_pipe,
mlpackages_dir,
model_version,
compute_unit,
delete_original_pipe=True,
scheduler_override=None,
controlnet_models=None,
force_zeros_for_empty_prompt=True,
sources=None):
"""
Initializes and returns a `CoreMLStableDiffusionPipeline` from an original
diffusers PyTorch pipeline
sources: 'packages' or 'compiled' forces creation of model from specified sources. sources must be in mlpackages_dir
"""
# Ensure `scheduler_override` object is of correct type if specified
if scheduler_override is not None:
assert isinstance(scheduler_override, SchedulerMixin)
logger.warning(
"Overriding scheduler in pipeline: "
f"Default={pytorch_pipe.scheduler}, Override={scheduler_override}")
# Gather configured tokenizer and scheduler attributes from the original pipe
if 'xl' in model_version:
coreml_pipe_kwargs = {
"tokenizer": pytorch_pipe.tokenizer,
'tokenizer_2': pytorch_pipe.tokenizer_2,
"scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
'xl': True,
}
model_packages_to_load = ["text_encoder", "text_encoder_2", "unet", "vae_decoder"]
else:
coreml_pipe_kwargs = {
"tokenizer": pytorch_pipe.tokenizer,
"scheduler": pytorch_pipe.scheduler if scheduler_override is None else scheduler_override,
"feature_extractor": pytorch_pipe.feature_extractor,
}
model_packages_to_load = ["text_encoder", "unet", "vae_decoder"]
coreml_pipe_kwargs["force_zeros_for_empty_prompt"] = force_zeros_for_empty_prompt
if getattr(pytorch_pipe, "safety_checker", None) is not None:
model_packages_to_load.append("safety_checker")
else:
logger.warning(
f"Original diffusers pipeline for {model_version} does not have a safety_checker, "
"Core ML pipeline will mirror this behavior.")
coreml_pipe_kwargs["safety_checker"] = None
if delete_original_pipe:
del pytorch_pipe
gc.collect()
logger.info("Removed PyTorch pipe to reduce peak memory consumption")
if controlnet_models:
model_packages_to_load.remove("unet")
coreml_pipe_kwargs["unet"] = _load_mlpackage(
submodule_name="control-unet",
mlpackages_dir=mlpackages_dir,
model_version=model_version,
compute_unit=compute_unit,
)
coreml_pipe_kwargs["controlnet"] = [_load_mlpackage_controlnet(
mlpackages_dir,
model_version,
compute_unit,
) for model_version in controlnet_models]
else:
coreml_pipe_kwargs["controlnet"] = None
# Load Core ML models
logger.info(f"Loading Core ML models in memory from {mlpackages_dir}")
coreml_pipe_kwargs.update({
model_name: _load_mlpackage(
submodule_name=model_name,
mlpackages_dir=mlpackages_dir,
model_version=model_version,
compute_unit=compute_unit,
sources=sources,
)
for model_name in model_packages_to_load
})
logger.info("Done.")
logger.info("Initializing Core ML pipe for image generation")
coreml_pipe = CoreMLStableDiffusionPipeline(**coreml_pipe_kwargs)
logger.info("Done.")
return coreml_pipe
def get_image_path(args, **override_kwargs):
""" mkdir output folder and encode metadata in the filename
"""
out_folder = os.path.join(args.o, "_".join(args.prompt.replace("/", "_").rsplit(" ")))
os.makedirs(out_folder, exist_ok=True)
out_fname = f"randomSeed_{override_kwargs.get('seed', None) or args.seed}"
out_fname += f"_computeUnit_{override_kwargs.get('compute_unit', None) or args.compute_unit}"
out_fname += f"_modelVersion_{override_kwargs.get('model_version', None) or args.model_version.replace('/', '_')}"
if args.scheduler is not None:
out_fname += f"_customScheduler_{override_kwargs.get('scheduler', None) or args.scheduler}"
out_fname += f"_numInferenceSteps{override_kwargs.get('num_inference_steps', None) or args.num_inference_steps}"
return os.path.join(out_folder, out_fname + ".png")
def prepare_controlnet_cond(image_path, height, width):
image = Image.open(image_path).convert("RGB")
image = image.resize((height, width), resample=Image.LANCZOS)
image = np.array(image).transpose(2, 0, 1) / 255.0
return image
def main(args):
logger.info(f"Setting random seed to {args.seed}")
np.random.seed(args.seed)
logger.info("Initializing PyTorch pipe for reference configuration")
SDP = StableDiffusionXLPipeline if 'xl' in args.model_version else StableDiffusionPipeline
pytorch_pipe = SDP.from_pretrained(
args.model_version,
use_auth_token=True,
)
# Get Scheduler
user_specified_scheduler = None
if args.scheduler is not None:
user_specified_scheduler = SCHEDULER_MAP[
args.scheduler].from_config(pytorch_pipe.scheduler.config)
# Get Force Zeros Config if it exists
force_zeros_for_empty_prompt: bool = False
if 'xl' in args.model_version and 'force_zeros_for_empty_prompt' in pytorch_pipe.config:
force_zeros_for_empty_prompt = pytorch_pipe.config['force_zeros_for_empty_prompt']
coreml_pipe = get_coreml_pipe(
pytorch_pipe=pytorch_pipe,
mlpackages_dir=args.i,
model_version=args.model_version,
compute_unit=args.compute_unit,
scheduler_override=user_specified_scheduler,
controlnet_models=args.controlnet,
force_zeros_for_empty_prompt=force_zeros_for_empty_prompt,
sources=args.model_sources,
)
if args.controlnet:
controlnet_cond = []
for i, _ in enumerate(args.controlnet):
image_path = args.controlnet_inputs[i]
image = prepare_controlnet_cond(image_path, coreml_pipe.height, coreml_pipe.width)
controlnet_cond.append(image)
else:
controlnet_cond = None
logger.info("Beginning image generation.")
image = coreml_pipe(
prompt=args.prompt,
height=coreml_pipe.height,
width=coreml_pipe.width,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
controlnet_cond=controlnet_cond,
negative_prompt=args.negative_prompt,
unet_batch_one=args.unet_batch_one,
)
out_path = get_image_path(args)
logger.info(f"Saving generated image to {out_path}")
image["images"][0].save(out_path)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--prompt",
required=True,
help="The text prompt to be used for text-to-image generation.")
parser.add_argument(
"-i",
required=True,
help=("Path to input directory with the .mlpackage files generated by "
"python_coreml_stable_diffusion.torch2coreml"))
parser.add_argument("-o", required=True)
parser.add_argument("--seed",
"-s",
default=93,
type=int,
help="Random seed to be able to reproduce results")
parser.add_argument(
"--model-version",
default="CompVis/stable-diffusion-v1-4",
help=
("The pre-trained model checkpoint and configuration to restore. "
"For available versions: https://huggingface.co/models?search=stable-diffusion"
))
parser.add_argument(
"--compute-unit",
choices=get_available_compute_units(),
default="ALL",
help=("The compute units to be used when executing Core ML models. "
f"Options: {get_available_compute_units()}"))
parser.add_argument(
"--scheduler",
choices=tuple(SCHEDULER_MAP.keys()),
default=None,
help=("The scheduler to use for running the reverse diffusion process. "
"If not specified, the default scheduler from the diffusers pipeline is utilized"))
parser.add_argument(
"--num-inference-steps",
default=50,
type=int,
help="The number of iterations the unet model will be executed throughout the reverse diffusion process")
parser.add_argument(
"--guidance-scale",
default=7.5,
type=float,
help="Controls the influence of the text prompt on sampling process (0=random images)")
parser.add_argument(
"--controlnet",
nargs="*",
type=str,
help=("Enables ControlNet and use control-unet instead of unet for additional inputs. "
"For Multi-Controlnet, provide the model names separated by spaces."))
parser.add_argument(
"--controlnet-inputs",
nargs="*",
type=str,
help=("Image paths for ControlNet inputs. "
"Please enter images corresponding to each controlnet provided at --controlnet option in same order."))
parser.add_argument(
"--negative-prompt",
default=None,
help="The negative text prompt to be used for text-to-image generation.")
parser.add_argument(
"--unet-batch-one",
action="store_true",
help="Do not batch unet predictions for the prompt and negative prompt.")
parser.add_argument('--model-sources',
default=None,
choices=['packages', 'compiled'],
help='Force build from `packages` or `compiled`')
args = parser.parse_args()
main(args)
python_coreml_stable_diffusion/torch2coreml.py
+1787
View File
File diff suppressed because it is too large Load Diff
python_coreml_stable_diffusion/unet.py
+1255
View File
File diff suppressed because it is too large Load Diff
requirements.txt
+11
View File
@@ -0,0 +1,11 @@
coremltools>=8.0
diffusers[torch]==0.30.2
diffusionkit==0.4.0
torch
transformers==4.44.2
scipy
scikit-learn
pytest
invisible-watermark
safetensors
matplotlib
setup.py
+44
View File
@@ -0,0 +1,44 @@
from setuptools import setup, find_packages
from python_coreml_stable_diffusion._version import __version__
with open('README.md') as f:
readme = f.read()
setup(
name='python_coreml_stable_diffusion',
version=__version__,
url='https://github.com/apple/ml-stable-diffusion',
description="Run Stable Diffusion on Apple Silicon with Core ML (Python and Swift)",
long_description=readme,
long_description_content_type='text/markdown',
author='Apple Inc.',
install_requires=[
"coremltools>=8.0",
"diffusers[torch]==0.30.2",
"torch",
"transformers==4.44.2",
"huggingface-hub==0.24.6",
"scipy",
"numpy<1.24",
"pytest",
"scikit-learn",
"invisible-watermark",
"safetensors",
"matplotlib",
"diffusionkit==0.4.0",
],
packages=find_packages(),
classifiers=[
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Operating System :: MacOS :: MacOS X",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.7",
"Programming Language :: Python :: 3.8",
"Programming Language :: Python :: 3.9",
"Topic :: Artificial Intelligence",
"Topic :: Scientific/Engineering",
"Topic :: Software Development",
],
)
swift/StableDiffusion/pipeline/CGImage+vImage.swift
+197
View File
@@ -0,0 +1,197 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import Accelerate
import CoreML
import CoreGraphics
@available(iOS 16.0, macOS 13.0, *)
extension CGImage {
typealias PixelBufferPFx1 = vImage.PixelBuffer<vImage.PlanarF>
typealias PixelBufferP8x3 = vImage.PixelBuffer<vImage.Planar8x3>
typealias PixelBufferIFx3 = vImage.PixelBuffer<vImage.InterleavedFx3>
typealias PixelBufferI8x3 = vImage.PixelBuffer<vImage.Interleaved8x3>
public enum ShapedArrayError: String, Swift.Error {
case wrongNumberOfChannels
case incorrectFormatsConvertingToShapedArray
case vImageConverterNotInitialized
}
public static func fromShapedArray(_ array: MLShapedArray<Float32>) throws -> CGImage {
// array is [N,C,H,W], where C==3
let channelCount = array.shape[1]
guard channelCount == 3 else {
throw ShapedArrayError.wrongNumberOfChannels
}
let height = array.shape[2]
let width = array.shape[3]
// Normalize each channel into a float between 0 and 1.0
let floatChannels = (0..<channelCount).map { i in
// Normalized channel output
let cOut = PixelBufferPFx1(width: width, height:height)
// Reference this channel in the array and normalize
array[0][i].withUnsafeShapedBufferPointer { ptr, _, strides in
let cIn = PixelBufferPFx1(data: .init(mutating: ptr.baseAddress!),
width: width, height: height,
byteCountPerRow: strides[0]*4)
// Map [-1.0 1.0] -> [0.0 1.0]
cIn.multiply(by: 0.5, preBias: 1.0, postBias: 0.0, destination: cOut)
}
return cOut
}
// Convert to interleaved and then to UInt8
let floatImage = PixelBufferIFx3(planarBuffers: floatChannels)
let uint8Image = PixelBufferI8x3(width: width, height: height)
floatImage.convert(to:uint8Image) // maps [0.0 1.0] -> [0 255] and clips
// Convert to uint8x3 to RGB CGImage (no alpha)
let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.none.rawValue)
let cgImage = uint8Image.makeCGImage(cgImageFormat:
.init(bitsPerComponent: 8,
bitsPerPixel: 3*8,
colorSpace: CGColorSpace(name: CGColorSpace.sRGB) ?? CGColorSpaceCreateDeviceRGB(),
bitmapInfo: bitmapInfo)!)!
return cgImage
}
public func planarRGBShapedArray(minValue: Float, maxValue: Float)
throws -> MLShapedArray<Float32> {
guard
var sourceFormat = vImage_CGImageFormat(cgImage: self),
var mediumFormat = vImage_CGImageFormat(
bitsPerComponent: 8 * MemoryLayout<UInt8>.size,
bitsPerPixel: 8 * MemoryLayout<UInt8>.size * 4,
colorSpace: CGColorSpaceCreateDeviceRGB(),
bitmapInfo: CGBitmapInfo(rawValue: CGImageAlphaInfo.first.rawValue)),
let width = vImagePixelCount(exactly: self.width),
let height = vImagePixelCount(exactly: self.height)
else {
throw ShapedArrayError.incorrectFormatsConvertingToShapedArray
}
var sourceImageBuffer = try vImage_Buffer(cgImage: self)
var mediumDestination = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: mediumFormat.bitsPerPixel)
let converter = vImageConverter_CreateWithCGImageFormat(
&sourceFormat,
&mediumFormat,
nil,
vImage_Flags(kvImagePrintDiagnosticsToConsole),
nil)
guard let converter = converter?.takeRetainedValue() else {
throw ShapedArrayError.vImageConverterNotInitialized
}
vImageConvert_AnyToAny(converter, &sourceImageBuffer, &mediumDestination, nil, vImage_Flags(kvImagePrintDiagnosticsToConsole))
var destinationA = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
var destinationR = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
var destinationG = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
var destinationB = try vImage_Buffer(width: Int(width), height: Int(height), bitsPerPixel: 8 * UInt32(MemoryLayout<Float>.size))
var minFloat: [Float] = Array(repeating: minValue, count: 4)
var maxFloat: [Float] = Array(repeating: maxValue, count: 4)
vImageConvert_ARGB8888toPlanarF(&mediumDestination, &destinationA, &destinationR, &destinationG, &destinationB, &maxFloat, &minFloat, .zero)
let destAPtr = destinationA.data.assumingMemoryBound(to: Float.self)
let destRPtr = destinationR.data.assumingMemoryBound(to: Float.self)
let destGPtr = destinationG.data.assumingMemoryBound(to: Float.self)
let destBPtr = destinationB.data.assumingMemoryBound(to: Float.self)
for i in 0..<Int(width) * Int(height) {
if destAPtr.advanced(by: i).pointee == 0 {
destRPtr.advanced(by: i).pointee = -1
destGPtr.advanced(by: i).pointee = -1
destBPtr.advanced(by: i).pointee = -1
}
}
let redData = destinationR.unpaddedData()
let greenData = destinationG.unpaddedData()
let blueData = destinationB.unpaddedData()
let imageData = redData + greenData + blueData
let shapedArray = MLShapedArray<Float32>(data: imageData, shape: [1, 3, self.height, self.width])
return shapedArray
}
private func normalizePixelValues(pixel: UInt8) -> Float {
return (Float(pixel) / 127.5) - 1.0
}
public func toRGBShapedArray(minValue: Float, maxValue: Float)
throws -> MLShapedArray<Float32> {
let image = self
let width = image.width
let height = image.height
let alphaMaskValue: Float = minValue
guard let colorSpace = CGColorSpace(name: CGColorSpace.sRGB),
let context = CGContext(data: nil, width: width, height: height, bitsPerComponent: 8, bytesPerRow: 4 * width, space: colorSpace, bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue),
let ptr = context.data?.bindMemory(to: UInt8.self, capacity: width * height * 4) else {
return []
}
context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
var redChannel = [Float](repeating: 0, count: width * height)
var greenChannel = [Float](repeating: 0, count: width * height)
var blueChannel = [Float](repeating: 0, count: width * height)
for y in 0..<height {
for x in 0..<width {
let i = 4 * (y * width + x)
if ptr[i+3] == 0 {
// Alpha mask for controlnets
redChannel[y * width + x] = alphaMaskValue
greenChannel[y * width + x] = alphaMaskValue
blueChannel[y * width + x] = alphaMaskValue
} else {
redChannel[y * width + x] = normalizePixelValues(pixel: ptr[i])
greenChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+1])
blueChannel[y * width + x] = normalizePixelValues(pixel: ptr[i+2])
}
}
}
let colorShape = [1, 1, height, width]
let redShapedArray = MLShapedArray<Float32>(scalars: redChannel, shape: colorShape)
let greenShapedArray = MLShapedArray<Float32>(scalars: greenChannel, shape: colorShape)
let blueShapedArray = MLShapedArray<Float32>(scalars: blueChannel, shape: colorShape)
let shapedArray = MLShapedArray<Float32>(concatenating: [redShapedArray, greenShapedArray, blueShapedArray], alongAxis: 1)
return shapedArray
}
}
extension vImage_Buffer {
func unpaddedData() -> Data {
let bytesPerPixel = self.rowBytes / Int(self.width)
let bytesPerRow = Int(self.width) * bytesPerPixel
var contiguousPixelData = Data(capacity: bytesPerRow * Int(self.height))
for row in 0..<Int(self.height) {
let rowStart = self.data!.advanced(by: row * self.rowBytes)
let rowData = Data(bytes: rowStart, count: bytesPerRow)
contiguousPixelData.append(rowData)
}
return contiguousPixelData
}
}
swift/StableDiffusion/pipeline/ControlNet.swift
+130
View File
@@ -0,0 +1,130 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
import Accelerate
@available(iOS 16.2, macOS 13.1, *)
public struct ControlNet: ResourceManaging {
var models: [ManagedMLModel]
public init(modelAt urls: [URL],
configuration: MLModelConfiguration) {
self.models = urls.map { ManagedMLModel(modelAt: $0, configuration: configuration) }
}
/// Load resources.
public func loadResources() throws {
for model in models {
try model.loadResources()
}
}
/// Unload the underlying model to free up memory
public func unloadResources() {
for model in models {
model.unloadResources()
}
}
/// Pre-warm resources
public func prewarmResources() throws {
// Override default to pre-warm each model
for model in models {
try model.loadResources()
model.unloadResources()
}
}
var inputImageDescriptions: [MLFeatureDescription] {
models.map { model in
try! model.perform {
$0.modelDescription.inputDescriptionsByName["controlnet_cond"]!
}
}
}
/// The expected shape of the models image input
public var inputImageShapes: [[Int]] {
inputImageDescriptions.map { desc in
desc.multiArrayConstraint!.shape.map { $0.intValue }
}
}
/// Calculate additional inputs for Unet to generate intended image following provided images
///
/// - Parameters:
/// - latents: Batch of latent samples in an array
/// - timeStep: Current diffusion timestep
/// - hiddenStates: Hidden state to condition on
/// - images: Images for each ControlNet
/// - Returns: Array of predicted noise residuals
func execute(
latents: [MLShapedArray<Float32>],
timeStep: Int,
hiddenStates: MLShapedArray<Float32>,
images: [MLShapedArray<Float32>]
) throws -> [[String: MLShapedArray<Float32>]] {
// Match time step batch dimension to the model / latent samples
let t = MLShapedArray(scalars: [Float(timeStep), Float(timeStep)], shape: [2])
var outputs: [[String: MLShapedArray<Float32>]] = []
for (modelIndex, model) in models.enumerated() {
let inputs = try latents.map { latent in
let dict: [String: Any] = [
"sample": MLMultiArray(latent),
"timestep": MLMultiArray(t),
"encoder_hidden_states": MLMultiArray(hiddenStates),
"controlnet_cond": MLMultiArray(images[modelIndex])
]
return try MLDictionaryFeatureProvider(dictionary: dict)
}
let batch = MLArrayBatchProvider(array: inputs)
let results = try model.perform {
try $0.predictions(fromBatch: batch)
}
// pre-allocate MLShapedArray with a specific shape in outputs
if outputs.isEmpty {
outputs = initOutputs(
batch: latents.count,
shapes: results.features(at: 0).featureValueDictionary
)
}
for n in 0..<results.count {
let result = results.features(at: n)
for k in result.featureNames {
let newValue = result.featureValue(for: k)!.multiArrayValue!
if modelIndex == 0 {
outputs[n][k] = MLShapedArray<Float32>(newValue)
} else {
let outputArray = MLMultiArray(outputs[n][k]!)
let count = newValue.count
let inputPointer = newValue.dataPointer.assumingMemoryBound(to: Float.self)
let outputPointer = outputArray.dataPointer.assumingMemoryBound(to: Float.self)
vDSP_vadd(inputPointer, 1, outputPointer, 1, outputPointer, 1, vDSP_Length(count))
}
}
}
}
return outputs
}
private func initOutputs(batch: Int, shapes: [String: MLFeatureValue]) -> [[String: MLShapedArray<Float32>]] {
var output: [String: MLShapedArray<Float32>] = [:]
for (outputName, featureValue) in shapes {
output[outputName] = MLShapedArray<Float32>(
repeating: 0.0,
shape: featureValue.multiArrayValue!.shape.map { $0.intValue }
)
}
return Array(repeating: output, count: batch)
}
}
swift/StableDiffusion/pipeline/DPMSolverMultistepScheduler.swift
+273
View File
@@ -0,0 +1,273 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. and The HuggingFace Team. All Rights Reserved.
import Accelerate
import CoreML
/// How to space timesteps for inference
public enum TimeStepSpacing {
case linspace
case leading
case karras
}
/// A scheduler used to compute a de-noised image
///
/// This implementation matches:
/// [Hugging Face Diffusers DPMSolverMultistepScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py)
///
/// It uses the DPM-Solver++ algorithm: [code](https://github.com/LuChengTHU/dpm-solver) [paper](https://arxiv.org/abs/2211.01095).
/// Limitations:
/// - Only implemented for DPM-Solver++ algorithm (not DPM-Solver).
/// - Second order only.
/// - Assumes the model predicts epsilon.
/// - No dynamic thresholding.
/// - `midpoint` solver algorithm.
@available(iOS 16.2, macOS 13.1, *)
public final class DPMSolverMultistepScheduler: Scheduler {
public let trainStepCount: Int
public let inferenceStepCount: Int
public let betas: [Float]
public let alphas: [Float]
public let alphasCumProd: [Float]
public let timeSteps: [Int]
public let alpha_t: [Float]
public let sigma_t: [Float]
public let lambda_t: [Float]
public let solverOrder = 2
private(set) var lowerOrderStepped = 0
private var usingKarrasSigmas = false
/// Whether to use lower-order solvers in the final steps. Only valid for less than 15 inference steps.
/// We empirically find this trick can stabilize the sampling of DPM-Solver, especially with 10 or fewer steps.
public let useLowerOrderFinal = true
// Stores solverOrder (2) items
public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
/// Create a scheduler that uses a second order DPM-Solver++ algorithm.
///
/// - Parameters:
/// - stepCount: Number of inference steps to schedule
/// - trainStepCount: Number of training diffusion steps
/// - betaSchedule: Method to schedule betas from betaStart to betaEnd
/// - betaStart: The starting value of beta for inference
/// - betaEnd: The end value for beta for inference
/// - timeStepSpacing: How to space time steps
/// - Returns: A scheduler ready for its first step
public init(
stepCount: Int = 50,
trainStepCount: Int = 1000,
betaSchedule: BetaSchedule = .scaledLinear,
betaStart: Float = 0.00085,
betaEnd: Float = 0.012,
timeStepSpacing: TimeStepSpacing = .linspace
) {
self.trainStepCount = trainStepCount
self.inferenceStepCount = stepCount
switch betaSchedule {
case .linear:
self.betas = linspace(betaStart, betaEnd, trainStepCount)
case .scaledLinear:
self.betas = linspace(pow(betaStart, 0.5), pow(betaEnd, 0.5), trainStepCount).map({ $0 * $0 })
}
self.alphas = betas.map({ 1.0 - $0 })
var alphasCumProd = self.alphas
for i in 1..<alphasCumProd.count {
alphasCumProd[i] *= alphasCumProd[i - 1]
}
self.alphasCumProd = alphasCumProd
switch timeStepSpacing {
case .linspace:
self.timeSteps = linspace(0, Float(self.trainStepCount-1), stepCount+1).dropFirst().reversed().map { Int(round($0)) }
self.alpha_t = vForce.sqrt(self.alphasCumProd)
self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
case .leading:
let lastTimeStep = trainStepCount - 1
let stepRatio = lastTimeStep / (stepCount + 1)
// Creates integer timesteps by multiplying by ratio
self.timeSteps = (0...stepCount).map { 1 + $0 * stepRatio }.dropFirst().reversed()
self.alpha_t = vForce.sqrt(self.alphasCumProd)
self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
case .karras:
// sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
let scaled = vDSP.multiply(
subtraction: ([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd),
subtraction: (vDSP.divide(1, self.alphasCumProd), [Float](repeating: 0, count: self.alphasCumProd.count))
)
let sigmas = vForce.sqrt(scaled)
let logSigmas = sigmas.map { log($0) }
let sigmaMin = sigmas.first!
let sigmaMax = sigmas.last!
let rho: Float = 7
let ramp = linspace(0, 1, stepCount)
let minInvRho = pow(sigmaMin, (1 / rho))
let maxInvRho = pow(sigmaMax, (1 / rho))
var karrasSigmas = ramp.map { pow(maxInvRho + $0 * (minInvRho - maxInvRho), rho) }
let karrasTimeSteps = karrasSigmas.map { sigmaToTimestep(sigma: $0, logSigmas: logSigmas) }
self.timeSteps = karrasTimeSteps
karrasSigmas.append(karrasSigmas.last!)
self.alpha_t = vDSP.divide(1, vForce.sqrt(vDSP.add(1, vDSP.square(karrasSigmas))))
self.sigma_t = vDSP.multiply(karrasSigmas, self.alpha_t)
usingKarrasSigmas = true
}
self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
}
func timestepToIndex(_ timestep: Int) -> Int {
guard usingKarrasSigmas else { return timestep }
return self.timeSteps.firstIndex(of: timestep) ?? 0
}
/// Convert the model output to the corresponding type the algorithm needs.
/// This implementation is for second-order DPM-Solver++ assuming epsilon prediction.
func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
assert(modelOutput.scalarCount == sample.scalarCount)
let scalarCount = modelOutput.scalarCount
let sigmaIndex = timestepToIndex(timestep)
let (alpha_t, sigma_t) = (self.alpha_t[sigmaIndex], self.sigma_t[sigmaIndex])
return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
assert(scalars.count == scalarCount)
modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
sample.withUnsafeShapedBufferPointer { sample, _, _ in
for i in 0 ..< scalarCount {
scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
}
}
}
}
}
/// One step for the first-order DPM-Solver (equivalent to DDIM).
/// See https://arxiv.org/abs/2206.00927 for the detailed derivation.
/// var names and code structure mostly follow https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
func firstOrderUpdate(
modelOutput: MLShapedArray<Float32>,
timestep: Int,
prevTimestep: Int,
sample: MLShapedArray<Float32>
) -> MLShapedArray<Float32> {
let prevIndex = timestepToIndex(prevTimestep)
let currIndex = timestepToIndex(timestep)
let (p_lambda_t, lambda_s) = (Double(lambda_t[prevIndex]), Double(lambda_t[currIndex]))
let p_alpha_t = Double(alpha_t[prevIndex])
let (p_sigma_t, sigma_s) = (Double(sigma_t[prevIndex]), Double(sigma_t[currIndex]))
let h = p_lambda_t - lambda_s
// x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
let x_t = weightedSum(
[p_sigma_t / sigma_s, -p_alpha_t * (exp(-h) - 1)],
[sample, modelOutput]
)
return x_t
}
/// One step for the second-order multistep DPM-Solver++ algorithm, using the midpoint method.
/// var names and code structure mostly follow https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
func secondOrderUpdate(
modelOutputs: [MLShapedArray<Float32>],
timesteps: [Int],
prevTimestep t: Int,
sample: MLShapedArray<Float32>
) -> MLShapedArray<Float32> {
let (s0, s1) = (timesteps[back: 1], timesteps[back: 2])
let (m0, m1) = (modelOutputs[back: 1], modelOutputs[back: 2])
let (p_lambda_t, lambda_s0, lambda_s1) = (
Double(lambda_t[timestepToIndex(t)]),
Double(lambda_t[timestepToIndex(s0)]),
Double(lambda_t[timestepToIndex(s1)])
)
let p_alpha_t = Double(alpha_t[timestepToIndex(t)])
let (p_sigma_t, sigma_s0) = (Double(sigma_t[timestepToIndex(t)]), Double(sigma_t[timestepToIndex(s0)]))
let (h, h_0) = (p_lambda_t - lambda_s0, lambda_s0 - lambda_s1)
let r0 = h_0 / h
let D0 = m0
// D1 = (1.0 / r0) * (m0 - m1)
let D1 = weightedSum(
[1/r0, -1/r0],
[m0, m1]
)
// See https://arxiv.org/abs/2211.01095 for detailed derivations
// x_t = (
// (sigma_t / sigma_s0) * sample
// - (alpha_t * (torch.exp(-h) - 1.0)) * D0
// - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
// )
let x_t = weightedSum(
[p_sigma_t/sigma_s0, -p_alpha_t * (exp(-h) - 1), -0.5 * p_alpha_t * (exp(-h) - 1)],
[sample, D0, D1]
)
return x_t
}
public func step(output: MLShapedArray<Float32>, timeStep t: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
let stepIndex = timeSteps.firstIndex(of: t) ?? timeSteps.count - 1
let prevTimestep = stepIndex == timeSteps.count - 1 ? 0 : timeSteps[stepIndex + 1]
let lowerOrderFinal = useLowerOrderFinal && stepIndex == timeSteps.count - 1 && timeSteps.count < 15
let lowerOrderSecond = useLowerOrderFinal && stepIndex == timeSteps.count - 2 && timeSteps.count < 15
let lowerOrder = lowerOrderStepped < 1 || lowerOrderFinal || lowerOrderSecond
let modelOutput = convertModelOutput(modelOutput: output, timestep: t, sample: sample)
if modelOutputs.count == solverOrder { modelOutputs.removeFirst() }
modelOutputs.append(modelOutput)
let prevSample: MLShapedArray<Float32>
if lowerOrder {
prevSample = firstOrderUpdate(modelOutput: modelOutput, timestep: t, prevTimestep: prevTimestep, sample: sample)
} else {
prevSample = secondOrderUpdate(
modelOutputs: modelOutputs,
timesteps: [timeSteps[stepIndex - 1], t],
prevTimestep: prevTimestep,
sample: sample
)
}
if lowerOrderStepped < solverOrder {
lowerOrderStepped += 1
}
return prevSample
}
}
func sigmaToTimestep(sigma: Float, logSigmas: [Float]) -> Int {
let logSigma = log(sigma)
let dists = logSigmas.map { logSigma - $0 }
// last index that is not negative, clipped to last index - 1
var lowIndex = dists.reduce(-1) { partialResult, dist in
return dist >= 0 && partialResult < dists.endIndex-2 ? partialResult + 1 : partialResult
}
lowIndex = max(lowIndex, 0)
let highIndex = lowIndex + 1
let low = logSigmas[lowIndex]
let high = logSigmas[highIndex]
// Interpolate sigmas
let w = ((low - logSigma) / (low - high)).clipped(to: 0...1)
// transform interpolated value to time range
let t = (1 - w) * Float(lowIndex) + w * Float(highIndex)
return Int(round(t))
}
extension FloatingPoint {
func clipped(to range: ClosedRange<Self>) -> Self {
return min(max(self, range.lowerBound), range.upperBound)
}
}
swift/StableDiffusion/pipeline/Decoder.swift
+80
View File
@@ -0,0 +1,80 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2024 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// A decoder model which produces RGB images from latent samples
@available(iOS 16.2, macOS 13.1, *)
public struct Decoder: ResourceManaging {
/// VAE decoder model
var model: ManagedMLModel
/// Create decoder from Core ML model
///
/// - Parameters:
/// - url: Location of compiled VAE decoder Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: A decoder that will lazily load its required resources when needed or requested
public init(modelAt url: URL, configuration: MLModelConfiguration) {
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
/// Batch decode latent samples into images
///
/// - Parameters:
/// - latents: Batch of latent samples to decode
/// - scaleFactor: scalar divisor on latents before decoding
/// - Returns: decoded images
public func decode(
_ latents: [MLShapedArray<Float32>],
scaleFactor: Float32,
shiftFactor: Float32 = 0.0
) throws -> [CGImage] {
// Form batch inputs for model
let inputs: [MLFeatureProvider] = try latents.map { sample in
// Reference pipeline scales the latent samples before decoding
let sampleScaled = MLShapedArray<Float32>(
scalars: sample.scalars.map { $0 / scaleFactor + shiftFactor },
shape: sample.shape)
let dict = [inputName: MLMultiArray(sampleScaled)]
return try MLDictionaryFeatureProvider(dictionary: dict)
}
let batch = MLArrayBatchProvider(array: inputs)
// Batch predict with model
let results = try model.perform { model in
try model.predictions(fromBatch: batch)
}
// Transform the outputs to CGImages
let images: [CGImage] = try (0..<results.count).map { i in
let result = results.features(at: i)
let outputName = result.featureNames.first!
let output = result.featureValue(for: outputName)!.multiArrayValue!
return try CGImage.fromShapedArray(MLShapedArray<Float32>(converting: output))
}
return images
}
var inputName: String {
try! model.perform { model in
model.modelDescription.inputDescriptionsByName.first!.key
}
}
}
swift/StableDiffusion/pipeline/DiscreteFlowScheduler.swift
+123
View File
@@ -0,0 +1,123 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2024 Apple Inc. All Rights Reserved.
import CoreML
/// A scheduler used to compute a de-noised image
@available(iOS 16.2, macOS 13.1, *)
public final class DiscreteFlowScheduler: Scheduler {
public let trainStepCount: Int
public let inferenceStepCount: Int
public var timeSteps = [Int]()
public var betas = [Float]()
public var alphas = [Float]()
public var alphasCumProd = [Float]()
public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
var trainSteps: Float
var shift: Float
var counter: Int
var sigmas = [Float]()
/// Create a scheduler that uses a second order DPM-Solver++ algorithm.
///
/// - Parameters:
/// - stepCount: Number of inference steps to schedule
/// - trainStepCount: Number of training diffusion steps
/// - timeStepShift: Amount to shift the timestep schedule
/// - Returns: A scheduler ready for its first step
public init(
stepCount: Int = 50,
trainStepCount: Int = 1000,
timeStepShift: Float = 3.0
) {
self.trainStepCount = trainStepCount
self.inferenceStepCount = stepCount
self.trainSteps = Float(trainStepCount)
self.shift = timeStepShift
self.counter = 0
let sigmaDistribution = linspace(1, trainSteps, Int(trainSteps)).map { sigmaFromTimestep($0) }
let timeStepDistribution = linspace(sigmaDistribution.first!, sigmaDistribution.last!, stepCount).reversed()
self.timeSteps = timeStepDistribution.map { Int($0 * trainSteps) }
self.sigmas = timeStepDistribution.map { sigmaFromTimestep($0 * trainSteps) }
}
func sigmaFromTimestep(_ timestep: Float) -> Float {
if shift == 1.0 {
return timestep / trainSteps
} else {
// shift * timestep / (1 + (shift - 1) * timestep)
let t = timestep / trainSteps
return shift * t / (1 + (shift - 1) * t)
}
}
func timestepsFromSigmas() -> [Float] {
return sigmas.map { $0 * trainSteps }
}
/// Convert the model output to the corresponding type the algorithm needs.
func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
assert(modelOutput.scalarCount == sample.scalarCount)
let stepIndex = timeSteps.firstIndex(of: timestep) ?? counter
let sigma = sigmas[stepIndex]
return MLShapedArray<Float>(unsafeUninitializedShape: modelOutput.shape) { result, _ in
modelOutput.withUnsafeShapedBufferPointer { noiseScalars, _, _ in
sample.withUnsafeShapedBufferPointer { latentScalars, _, _ in
for i in 0..<result.count {
let denoised = latentScalars[i] - noiseScalars[i] * sigma
result.initializeElement(
at: i,
to: denoised
)
}
}
}
}
}
public func calculateTimestepsFromSigmas(strength: Float?) -> [Float] {
guard let strength else { return timestepsFromSigmas() }
let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
let actualTimesteps = Array(timestepsFromSigmas()[startStep...])
return actualTimesteps
}
public func step(output: MLShapedArray<Float32>, timeStep t: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
let stepIndex = timeSteps.firstIndex(of: t) ?? counter // TODO: allow float timesteps in scheduler step protocol
let modelOutput = convertModelOutput(modelOutput: output, timestep: t, sample: sample)
modelOutputs.append(modelOutput)
let sigma = sigmas[stepIndex]
var dt = sigma
var prevSigma: Float = 0
if stepIndex < sigmas.count - 1 {
prevSigma = sigmas[stepIndex + 1]
dt = prevSigma - sigma
}
let prevSample: MLShapedArray<Float32> = MLShapedArray<Float>(unsafeUninitializedShape: modelOutput.shape) { result, _ in
modelOutput.withUnsafeShapedBufferPointer { noiseScalars, _, _ in
sample.withUnsafeShapedBufferPointer { latentScalars, _, _ in
for i in 0..<result.count {
let denoised = noiseScalars[i]
let x = latentScalars[i]
let d = (x - denoised) / sigma
let prev_x = x + d * dt
result.initializeElement(
at: i,
to: prev_x
)
}
}
}
}
counter += 1
return prevSample
}
}
swift/StableDiffusion/pipeline/Encoder.swift
+108
View File
@@ -0,0 +1,108 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// A encoder model which produces latent samples from RGB images
@available(iOS 16.2, macOS 13.1, *)
public struct Encoder: ResourceManaging {
public enum Error: String, Swift.Error {
case sampleInputShapeNotCorrect
}
/// VAE encoder model + post math and adding noise from schedular
var model: ManagedMLModel
/// Create encoder from Core ML model
///
/// - Parameters:
/// - url: Location of compiled VAE encoder Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: An encoder that will lazily load its required resources when needed or requested
public init(modelAt url: URL, configuration: MLModelConfiguration) {
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
/// Prediction queue
let queue = DispatchQueue(label: "encoder.predict")
/// Encode image into latent sample
///
/// - Parameters:
/// - image: Input image
/// - scaleFactor: scalar multiplier on latents before encoding image
/// - random
/// - Returns: The encoded latent space as MLShapedArray
public func encode(
_ image: CGImage,
scaleFactor: Float32,
random: inout RandomSource
) throws -> MLShapedArray<Float32> {
let imageData = try image.planarRGBShapedArray(minValue: -1.0, maxValue: 1.0)
guard imageData.shape == inputShape else {
// TODO: Consider auto resizing and croping similar to how Vision or CoreML auto-generated Swift code can accomplish with `MLFeatureValue`
throw Error.sampleInputShapeNotCorrect
}
let dict = [inputName: MLMultiArray(imageData)]
let input = try MLDictionaryFeatureProvider(dictionary: dict)
let result = try model.perform { model in
try model.prediction(from: input)
}
let outputName = result.featureNames.first!
let outputValue = result.featureValue(for: outputName)!.multiArrayValue!
let output = MLShapedArray<Float32>(converting: outputValue)
// DiagonalGaussianDistribution
let mean = output[0][0..<4]
let logvar = MLShapedArray<Float32>(
scalars: output[0][4..<8].scalars.map { min(max($0, -30), 20) },
shape: mean.shape
)
let std = MLShapedArray<Float32>(
scalars: logvar.scalars.map { exp(0.5 * $0) },
shape: logvar.shape
)
let latent = MLShapedArray<Float32>(
scalars: zip(mean.scalars, std.scalars).map {
Float32(random.nextNormal(mean: Double($0), stdev: Double($1)))
},
shape: logvar.shape
)
// Reference pipeline scales the latent after encoding
let latentScaled = MLShapedArray<Float32>(
scalars: latent.scalars.map { $0 * scaleFactor },
shape: [1] + latent.shape
)
return latentScaled
}
var inputDescription: MLFeatureDescription {
try! model.perform { model in
model.modelDescription.inputDescriptionsByName.first!.value
}
}
var inputName: String {
inputDescription.name
}
/// The expected shape of the models latent sample input
var inputShape: [Int] {
inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
}
swift/StableDiffusion/pipeline/ManagedMLModel.swift
+127
View File
@@ -0,0 +1,127 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import CoreML
/// A class to manage and gate access to a Core ML model
///
/// It will automatically load a model into memory when needed or requested
/// It allows one to request to unload the model from memory
@available(iOS 16.2, macOS 13.1, *)
public final class ManagedMLModel: ResourceManaging {
/// The location of the model
var modelURL: URL
/// The configuration to be used when the model is loaded
var configuration: MLModelConfiguration
/// The loaded model (when loaded)
var loadedModel: MLModel?
/// Queue to protect access to loaded model
var queue: DispatchQueue
/// Create a managed model given its location and desired loaded configuration
///
/// - Parameters:
/// - url: The location of the model
/// - configuration: The configuration to be used when the model is loaded/used
/// - Returns: A managed model that has not been loaded
public init(modelAt url: URL, configuration: MLModelConfiguration) {
self.modelURL = url
self.configuration = configuration
self.loadedModel = nil
self.queue = DispatchQueue(label: "managed.\(url.lastPathComponent)")
}
/// Instantiation and load model into memory
public func loadResources() throws {
try queue.sync {
try loadModel()
}
}
/// Unload the model if it was loaded
public func unloadResources() {
queue.sync {
loadedModel = nil
}
}
/// Perform an operation with the managed model via a supplied closure.
/// The model will be loaded and supplied to the closure and should only be
/// used within the closure to ensure all resource management is synchronized
///
/// - Parameters:
/// - body: Closure which performs and action on a loaded model
/// - Returns: The result of the closure
/// - Throws: An error if the model cannot be loaded or if the closure throws
public func perform<R>(_ body: (MLModel) throws -> R) throws -> R {
return try queue.sync {
try autoreleasepool {
try loadModel()
return try body(loadedModel!)
}
}
}
private func loadModel() throws {
if loadedModel == nil {
loadedModel = try MLModel(contentsOf: modelURL,
configuration: configuration)
}
}
}
@available(iOS 16.2, macOS 13.1, *)
public extension Array where Element == ManagedMLModel {
/// Performs batch predictions using an array of `[ManagedMLModel]` instances in a pipeline.
/// - Parameter batch: Inputs for btached predictions.
/// - Returns: Final prediction results after processing through all models.
/// - Throws: Errors if the array is empty, predictions fail, or results can't be combined.
func predictions(from batch: MLBatchProvider) throws -> MLBatchProvider {
var results = try self.first!.perform { model in
try model.predictions(fromBatch: batch)
}
if self.count == 1 {
return results
}
// Manual pipeline batch prediction
let inputs = batch.arrayOfFeatureValueDictionaries
for stage in self.dropFirst() {
// Combine the original inputs with the outputs of the last stage
let next = try results.arrayOfFeatureValueDictionaries
.enumerated().map { index, dict in
let nextDict = dict.merging(inputs[index]) { out, _ in out }
return try MLDictionaryFeatureProvider(dictionary: nextDict)
}
let nextBatch = MLArrayBatchProvider(array: next)
// Predict
results = try stage.perform { model in
try model.predictions(fromBatch: nextBatch)
}
}
return results
}
}
extension MLFeatureProvider {
var featureValueDictionary: [String : MLFeatureValue] {
self.featureNames.reduce(into: [String : MLFeatureValue]()) { result, name in
result[name] = self.featureValue(for: name)
}
}
}
extension MLBatchProvider {
var arrayOfFeatureValueDictionaries: [[String : MLFeatureValue]] {
(0..<self.count).map {
self.features(at: $0).featureValueDictionary
}
}
}
swift/StableDiffusion/pipeline/MultiModalDiffusionTransformer.swift
+125
View File
@@ -0,0 +1,125 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// MMDiT noise prediction model for stable diffusion
@available(iOS 16.2, macOS 13.1, *)
public struct MultiModalDiffusionTransformer: ResourceManaging {
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
///
/// It can be in the form of a single model or multiple stages
var models: [ManagedMLModel]
/// Creates a MMDiT noise prediction model
///
/// - Parameters:
/// - url: Location of single MMDiT compiled Core ML model
/// - configuration: Configuration to be used when the model is loaded
/// - Returns: MMDiT model that will lazily load its required resources when needed or requested
public init(modelAt url: URL,
configuration: MLModelConfiguration)
{
self.models = [ManagedMLModel(modelAt: url, configuration: configuration)]
}
/// Load resources.
public func loadResources() throws {
for model in models {
try model.loadResources()
}
}
/// Unload the underlying model to free up memory
public func unloadResources() {
for model in models {
model.unloadResources()
}
}
/// Pre-warm resources
public func prewarmResources() throws {
// Override default to pre-warm each model
for model in models {
try model.loadResources()
model.unloadResources()
}
}
var latentImageEmbeddingsDescription: MLFeatureDescription {
try! models.first!.perform { model in
model.modelDescription.inputDescriptionsByName["latent_image_embeddings"]!
}
}
/// The expected shape of the models latent sample input
public var latentImageEmbeddingsShape: [Int] {
latentImageEmbeddingsDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
var tokenLevelTextEmbeddingsDescription: MLFeatureDescription {
try! models.first!.perform { model in
model.modelDescription.inputDescriptionsByName["token_level_text_embeddings"]!
}
}
/// The expected shape of the geometry conditioning
public var tokenLevelTextEmbeddingsShape: [Int] {
tokenLevelTextEmbeddingsDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
/// Batch prediction noise from latent samples
///
/// - Parameters:
/// - latents: Batch of latent samples in an array
/// - timeStep: Current diffusion timestep
/// - hiddenStates: Hidden state to condition on
/// - Returns: Array of predicted noise residuals
func predictNoise(
latents: [MLShapedArray<Float32>],
timeStep: Float,
tokenLevelTextEmbeddings: MLShapedArray<Float32>,
pooledTextEmbeddings: MLShapedArray<Float32>
) throws -> [MLShapedArray<Float32>] {
// Match time step batch dimension to the model / latent samples
let t = MLShapedArray<Float32>(scalars: [timeStep, timeStep], shape: [2])
// Form batch input to model
let inputs = try latents.enumerated().map {
let dict: [String: Any] = [
"latent_image_embeddings": MLMultiArray($0.element),
"timestep": MLMultiArray(t),
"token_level_text_embeddings": MLMultiArray(tokenLevelTextEmbeddings),
"pooled_text_embeddings": MLMultiArray(pooledTextEmbeddings),
]
return try MLDictionaryFeatureProvider(dictionary: dict)
}
let batch = MLArrayBatchProvider(array: inputs)
// Make predictions
let results = try models.predictions(from: batch)
// Pull out the results in Float32 format
let noise = (0..<results.count).map { i in
let result = results.features(at: i)
let outputName = result.featureNames.first!
let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
// To conform to this func return type make sure we return float32
// Use the fact that the concatenating constructor for MLMultiArray
// can do type conversion:
let fp32Noise = MLMultiArray(
concatenating: [outputNoise],
axis: 0,
dataType: .float32
)
return MLShapedArray<Float32>(fp32Noise)
}
return noise
}
}
swift/StableDiffusion/pipeline/MultilingualTextEncoder.swift
+194
View File
@@ -0,0 +1,194 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
import Foundation
import NaturalLanguage
import CoreML
#if canImport(NaturalLanguage.NLContextualEmbedding)
@available(iOS 17.0, macOS 14.0, *)
public struct MultilingualTextEncoder: TextEncoderModel {
let adapter: ManagedMLModel?
let embeddingModel: NLContextualEmbedding
// TODO: use maximum sequence length from embedding.
let maximumEmbeddingSequenceLength = 256
/// Creates a multilingual text encoder.
///
/// - Parameters:
/// - url: The location of the compiled Core ML adapter model. The model is a linear projection layer that
/// transforms the contextual embedding size of 512 to the default text encoder CLIP size of 768.
/// - configuration: The configuration to be used when the model is loaded.
/// - script: The scipt of the contextual embedding.
public init(
modelAt url: URL? = nil,
configuration: MLModelConfiguration = .init(),
script: Script = .latin
) {
if let url {
self.adapter = ManagedMLModel(modelAt: url, configuration: configuration)
} else {
self.adapter = nil
}
self.embeddingModel = NLContextualEmbedding(script: script.asNLScript)!
self.embeddingModel.requestAssets { _, _ in }
}
/// Loads model resources into memory.
public func loadResources() throws {
try adapter?.loadResources()
try embeddingModel.load()
}
/// Unloads the model resources to free up memory.
public func unloadResources() {
adapter?.unloadResources()
embeddingModel.unload()
}
/// Encodes the input text.
///
/// - Parameter text: The input text.
/// - Returns: An embedding shaped array.
public func encode(_ text: String) throws -> MLShapedArray<Float> {
guard embeddingModel.hasAvailableAssets else {
throw Error.missingEmbeddingResource
}
// Create the text embedding result.
let embedding = try embeddingModel.embeddingResult(for: text, language: nil)
// Create embedding array from token vectors.
var shapedEmbeddings = MLShapedArray<Double>(
repeating: 0.0,
shape: [1, maximumEmbeddingSequenceLength, embeddingModel.dimension]
)
shapedEmbeddings.withUnsafeMutableShapedBufferPointer { pointer, _, _ in
var tokenIndex = 0
embedding.enumerateTokenVectors(in: text.startIndex ..< text.endIndex) { (tokenEmbeddings, _) -> Bool in
for tokenEmbeddingIndex in 0 ..< tokenEmbeddings.count {
pointer[tokenIndex * embeddingModel.dimension + tokenEmbeddingIndex] = tokenEmbeddings[tokenEmbeddingIndex]
}
tokenIndex += 1
return true
}
}
if adapter == nil {
// Return embeddings with shape [1, 256, 512].
return MLShapedArray(converting: shapedEmbeddings)
} else {
// Project the embeddings to the correct CLIP model input shape of [1, 768, 1, 256].
return try projectEmbeddings(shapedEmbeddings)
}
}
/// Creates the adapter model input feature provider.
private func prepareProjectionInput(_ input: MLShapedArray<Double>) throws -> MLDictionaryFeatureProvider {
guard let adapter else {
fatalError("Cannot prepare projection input without an adapter.")
}
return try adapter.perform { model in
guard let inputDescription = model.modelDescription.inputDescriptionsByName.first?.value else {
throw Error.missingAdapterInput
}
return try MLDictionaryFeatureProvider(dictionary: [inputDescription.name: MLMultiArray(input)])
}
}
/// Processes the adapter model output feature provider.
private func processProjectionOutput(_ output: MLFeatureProvider) throws -> MLShapedArray<Float> {
guard let adapter else {
fatalError("Cannot process projection output without an adapter.")
}
return try adapter.perform { model in
guard let outputDescription = model.modelDescription.outputDescriptionsByName.first?.value else {
throw Error.missingAdapterOutput
}
guard let result = output
.featureValue(for: outputDescription.name)?
.multiArrayValue else {
throw Error.incompatibleAdapterOutputDataFormat(
expected: .multiArray,
actual: outputDescription.type
)
}
return MLShapedArray(converting: result)
}
}
/// Projects the embeddings.
private func projectEmbeddings(_ embeddings: MLShapedArray<Double>) throws -> MLShapedArray<Float> {
guard let adapter else {
fatalError("Cannot project embeddings without an adapter.")
}
let inputFeatureProvider = try prepareProjectionInput(embeddings)
let projection = try adapter.perform { model in
return try model.prediction(from: inputFeatureProvider)
}
return try processProjectionOutput(projection)
}
}
@available(iOS 17.0, macOS 14.0, *)
extension MultilingualTextEncoder {
/// A multilingual text encoder error.
public enum Error: Swift.Error, LocalizedError, Equatable, CustomDebugStringConvertible {
/// An error that indicates that the resource for the embedding is missing.
case missingEmbeddingResource
/// An error that indicates that the adapter model input data has the wrong format.
case incompatibleAdapterInputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
/// An error that indicates that the adapter model output data has the wrong format.
case incompatibleAdapterOutputDataFormat(expected: MLFeatureType, actual: MLFeatureType)
/// An error that indicates that the adapter model is missing an input.
case missingAdapterInput
/// An error that indicates that the adapter model is missing an output.
case missingAdapterOutput
/// A debug description of the error.
public var errorDescription: String? {
debugDescription
}
/// A text representation of the error.
public var debugDescription: String {
switch self {
case .missingEmbeddingResource:
return "Resources required for generating embeddings are missing. Make sure that your device is connected to the internet and try again."
case .incompatibleAdapterInputDataFormat(expected: let expected, actual: let actual):
return "The adapter model input expected to be \(expected) but is \(actual)."
case .incompatibleAdapterOutputDataFormat(expected: let expected, actual: let actual):
return "The adapter model output expected to be \(expected) but is \(actual)."
case .missingAdapterInput:
return "The adapter model is missing an input."
case .missingAdapterOutput:
return "The adapter model is missing an output."
}
}
}
}
#endif
@available(iOS 16.2, macOS 13.1, *)
public enum Script: String {
case latin, cyrillic, cjk
#if canImport(NaturalLanguage.NLScript)
@available(iOS 17.0, macOS 14.0, *)
var asNLScript: NLScript {
switch self {
case .latin: return .latin
case .cyrillic: return .cyrillic
case .cjk: return .simplifiedChinese
}
}
#endif
}
swift/StableDiffusion/pipeline/NumPyRandomSource.swift
+119
View File
@@ -0,0 +1,119 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// A random source consistent with NumPy
///
/// This implementation matches:
/// [NumPy's older randomkit.c](https://github.com/numpy/numpy/blob/v1.0/numpy/random/mtrand/randomkit.c)
///
@available(iOS 16.2, macOS 13.1, *)
struct NumPyRandomSource: RandomNumberGenerator, RandomSource {
struct State {
var key = [UInt32](repeating: 0, count: 624)
var pos: Int = 0
var nextGauss: Double? = nil
}
var state: State
/// Initialize with a random seed
///
/// - Parameters
/// - seed: Seed for underlying Mersenne Twister 19937 generator
/// - Returns random source
init(seed: UInt32) {
state = .init()
var s = seed & 0xffffffff
for i in 0 ..< state.key.count {
state.key[i] = s
s = UInt32((UInt64(1812433253) * UInt64(s ^ (s >> 30)) + UInt64(i) + 1) & 0xffffffff)
}
state.pos = state.key.count
state.nextGauss = nil
}
/// Generate next UInt32 using fast 32bit Mersenne Twister
mutating func nextUInt32() -> UInt32 {
let n = 624
let m = 397
let matrixA: UInt64 = 0x9908b0df
let upperMask: UInt32 = 0x80000000
let lowerMask: UInt32 = 0x7fffffff
var y: UInt32
if state.pos == state.key.count {
for i in 0 ..< (n - m) {
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
state.key[i] = state.key[i + m] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
}
for i in (n - m) ..< (n - 1) {
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
state.key[i] = state.key[i + (m - n)] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
}
y = (state.key[n - 1] & upperMask) | (state.key[0] & lowerMask)
state.key[n - 1] = state.key[m - 1] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
state.pos = 0
}
y = state.key[state.pos]
state.pos += 1
y ^= (y >> 11)
y ^= (y << 7) & 0x9d2c5680
y ^= (y << 15) & 0xefc60000
y ^= (y >> 18)
return y
}
mutating func next() -> UInt64 {
let low = nextUInt32()
let high = nextUInt32()
return (UInt64(high) << 32) | UInt64(low)
}
/// Generate next random double value
mutating func nextDouble() -> Double {
let a = Double(nextUInt32() >> 5)
let b = Double(nextUInt32() >> 6)
return (a * 67108864.0 + b) / 9007199254740992.0
}
/// Generate next random value from a standard normal
mutating func nextGauss() -> Double {
if let nextGauss = state.nextGauss {
state.nextGauss = nil
return nextGauss
}
var x1, x2, r2: Double
repeat {
x1 = 2.0 * nextDouble() - 1.0
x2 = 2.0 * nextDouble() - 1.0
r2 = x1 * x1 + x2 * x2
} while r2 >= 1.0 || r2 == 0.0
// Box-Muller transform
let f = sqrt(-2.0 * log(r2) / r2)
state.nextGauss = f * x1
return f * x2
}
/// Generates a random value from a normal distribution with given mean and standard deviation.
mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
nextGauss() * stdev + mean
}
/// Generates an array of random values from a normal distribution with given mean and standard deviation.
mutating func normalArray(count: Int, mean: Double = 0.0, stdev: Double = 1.0) -> [Double] {
(0 ..< count).map { _ in nextNormal(mean: mean, stdev: stdev) }
}
/// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
let count = shape.reduce(1, *)
return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
}
}
swift/StableDiffusion/pipeline/NvRandomSource.swift
+91
View File
@@ -0,0 +1,91 @@
import Foundation
import CoreML
/// A random source consistent with NVIDIA curandom
///
/// This implementation references to:
/// https://github.com/dsnz/random/blob/master/philox.py for Philox_M4_32 configuration.
///
@available(iOS 16.2, macOS 13.1, *)
struct NvRandomSource: RandomSource {
public let seed: UInt64
private var offset: UInt32
/// Initialize with a random seed
///
/// - Parameters
/// - seed: Seed for underlying Philox M4 32 generator
/// - Returns random source
public init(seed: UInt32) {
self.seed = UInt64(seed)
offset = 0
}
static private let PHILOX_M4_32: (UInt32, UInt32) = (0xD251_1F53, 0xCD9E_8D57)
static private let PHILOX_W_32: (UInt32, UInt32) = (0x9E37_79B9, 0xBB67_AE85)
static private func philox4Round(counter: inout [[UInt32]], key: [[UInt32]]) {
for i in 0..<counter[0].count {
let v1: UInt64 = UInt64(counter[0][i]) * UInt64(PHILOX_M4_32.0)
let v2: UInt64 = UInt64(counter[2][i]) * UInt64(PHILOX_M4_32.1)
counter[0][i] = UInt32(v2 >> 32) ^ counter[1][i] ^ key[0][i]
counter[1][i] = UInt32(v2 & 0xffff_ffff)
counter[2][i] = UInt32(v1 >> 32) ^ counter[3][i] ^ key[1][i]
counter[3][i] = UInt32(v1 & 0xffff_ffff)
}
}
static private func philox4Bumpkey(key: inout [[UInt32]]) {
for (i, element) in key[0].enumerated() {
key[0][i] = element &+ PHILOX_W_32.0
}
for (i, element) in key[1].enumerated() {
key[1][i] = element &+ PHILOX_W_32.1
}
}
static private func philox4_32(counter: inout [[UInt32]], key: inout [[UInt32]], rounds: Int = 10) {
for _ in 0..<(rounds - 1) {
philox4Round(counter: &counter, key: key)
philox4Bumpkey(key: &key)
}
philox4Round(counter: &counter, key: key)
}
private func boxMuller(_ counter1: [UInt32], _ counter2: [UInt32], mean: Double, stdev: Double) -> [Double] {
// Box-Muller transform
return zip(counter1, counter2).map {
let u: Double = Double($0) / 4294967296.0 + (1.0 / 8589934592.0)
let v: Double = Double($1) * (.pi / 2147483648.0) + (.pi / 4294967296.0)
let radius = stdev * sqrt(-2.0 * log(u))
return radius * sin(v) + mean
}
}
private mutating func normalArray(count: Int, mean: Double, stdev: Double) -> [Double] {
var counter: [[UInt32]] = [
Array(repeating: offset, count: count),
Array(repeating: 0, count: count),
Array(0..<UInt32(count)),
Array(repeating: 0, count: count),
]
offset += 1
var key: [[UInt32]] = [
Array(repeating: UInt32(seed & 0xffff_ffff), count: count),
Array(repeating: UInt32(seed >> 32), count: count),
]
Self.philox4_32(counter: &counter, key: &key)
return boxMuller(counter[0], counter[1], mean: mean, stdev: stdev)
}
/// Generates a random value from a normal distribution with given mean and standard deviation.
mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
return normalArray(count: 1, mean: mean, stdev: stdev)[0]
}
/// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
let count = shape.reduce(1, *)
return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
}
}
swift/StableDiffusion/pipeline/RandomSource.swift
+8
View File
@@ -0,0 +1,8 @@
import CoreML
@available(iOS 16.2, macOS 13.1, *)
public protocol RandomSource {
mutating func nextNormal(mean: Double, stdev: Double) -> Double
mutating func normalShapedArray(_ shape: [Int], mean: Double, stdev: Double) -> MLShapedArray<Double>
}
swift/StableDiffusion/pipeline/ResourceManaging.swift
+20
View File
@@ -0,0 +1,20 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
/// Protocol for managing internal resources
public protocol ResourceManaging {
/// Request resources to be loaded and ready if possible
func loadResources() throws
/// Request resources are unloaded / remove from memory if possible
func unloadResources()
}
extension ResourceManaging {
/// Request resources are pre-warmed by loading and unloading
func prewarmResources() throws {
try loadResources()
unloadResources()
}
}
swift/StableDiffusion/pipeline/SafetyChecker.swift
+167
View File
@@ -0,0 +1,167 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
import Accelerate
/// Image safety checking model
@available(iOS 16.2, macOS 13.1, *)
public struct SafetyChecker: ResourceManaging {
/// Safety checking Core ML model
var model: ManagedMLModel
/// Creates safety checker
///
/// - Parameters:
/// - url: Location of compiled safety checking Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: A safety cherker that will lazily load its required resources when needed or requested
public init(modelAt url: URL, configuration: MLModelConfiguration) {
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
typealias PixelBufferPFx1 = vImage.PixelBuffer<vImage.PlanarF>
typealias PixelBufferP8x1 = vImage.PixelBuffer<vImage.Planar8>
typealias PixelBufferPFx3 = vImage.PixelBuffer<vImage.PlanarFx3>
typealias PixelBufferP8x3 = vImage.PixelBuffer<vImage.Planar8x3>
typealias PixelBufferIFx3 = vImage.PixelBuffer<vImage.InterleavedFx3>
typealias PixelBufferI8x3 = vImage.PixelBuffer<vImage.Interleaved8x3>
typealias PixelBufferI8x4 = vImage.PixelBuffer<vImage.Interleaved8x4>
enum SafetyCheckError: Error {
case imageResizeFailure
case imageToFloatFailure
case modelInputFailure
case unexpectedModelOutput
}
/// Check if image is safe
///
/// - Parameters:
/// - image: Image to check
/// - Returns: Whether the model considers the image to be safe
public func isSafe(_ image: CGImage) throws -> Bool {
let inputName = "clip_input"
let adjustmentName = "adjustment"
let imagesNames = "images"
let inputInfo = try model.perform { model in
model.modelDescription.inputDescriptionsByName
}
let inputShape = inputInfo[inputName]!.multiArrayConstraint!.shape
let width = inputShape[2].intValue
let height = inputShape[3].intValue
let resizedImage = try resizeToRGBA(image, width: width, height: height)
let bufferP8x3 = try getRGBPlanes(of: resizedImage)
let arrayPFx3 = normalizeToFloatShapedArray(bufferP8x3)
guard let input = try? MLDictionaryFeatureProvider(
dictionary:[
// Input that is analyzed for safety
inputName : MLMultiArray(arrayPFx3),
// No adjustment, use default threshold
adjustmentName : MLMultiArray(MLShapedArray<Float32>(scalars: [0], shape: [1])),
// Supplying dummy images to be filtered (will be ignored)
imagesNames : MLMultiArray(shape:[1, 512, 512, 3], dataType: .float16)
]
) else {
throw SafetyCheckError.modelInputFailure
}
let result = try model.perform { model in
try model.prediction(from: input)
}
let output = result.featureValue(for: "has_nsfw_concepts")
guard let unsafe = output?.multiArrayValue?[0].boolValue else {
throw SafetyCheckError.unexpectedModelOutput
}
return !unsafe
}
func resizeToRGBA(_ image: CGImage,
width: Int, height: Int) throws -> CGImage {
guard let context = CGContext(
data: nil,
width: width,
height: height,
bitsPerComponent: 8,
bytesPerRow: width*4,
space: CGColorSpaceCreateDeviceRGB(),
bitmapInfo: CGImageAlphaInfo.noneSkipLast.rawValue) else {
throw SafetyCheckError.imageResizeFailure
}
context.interpolationQuality = .high
context.draw(image, in: CGRect(x: 0, y: 0, width: width, height: height))
guard let resizedImage = context.makeImage() else {
throw SafetyCheckError.imageResizeFailure
}
return resizedImage
}
func getRGBPlanes(of rgbaImage: CGImage) throws -> PixelBufferP8x3 {
// Reference as interleaved 8 bit vImage PixelBuffer
var emptyFormat = vImage_CGImageFormat()
guard let bufferI8x4 = try? PixelBufferI8x4(
cgImage: rgbaImage,
cgImageFormat:&emptyFormat) else {
throw SafetyCheckError.imageToFloatFailure
}
// Drop the alpha channel, keeping RGB
let bufferI8x3 = PixelBufferI8x3(width: rgbaImage.width, height:rgbaImage.height)
bufferI8x4.convert(to: bufferI8x3, channelOrdering: .RGBA)
// De-interleave into 8-bit planes
return PixelBufferP8x3(interleavedBuffer: bufferI8x3)
}
func normalizeToFloatShapedArray(_ bufferP8x3: PixelBufferP8x3) -> MLShapedArray<Float32> {
let width = bufferP8x3.width
let height = bufferP8x3.height
let means = [0.485, 0.456, 0.406] as [Float]
let stds = [0.229, 0.224, 0.225] as [Float]
// Convert to normalized float 1x3xWxH input (plannar)
let arrayPFx3 = MLShapedArray<Float32>(repeating: 0.0, shape: [1, 3, width, height])
for c in 0..<3 {
arrayPFx3[0][c].withUnsafeShapedBufferPointer { ptr, _, strides in
let floatChannel = PixelBufferPFx1(data: .init(mutating: ptr.baseAddress!),
width: width, height: height,
byteCountPerRow: strides[0]*4)
bufferP8x3.withUnsafePixelBuffer(at: c) { uint8Channel in
uint8Channel.convert(to: floatChannel) // maps [0 255] -> [0 1]
floatChannel.multiply(by: 1.0/stds[c],
preBias: -means[c],
postBias: 0.0,
destination: floatChannel)
}
}
}
return arrayPFx3
}
}
swift/StableDiffusion/pipeline/SampleTimer.swift
+78
View File
@@ -0,0 +1,78 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
/// A utility for timing events and tracking time statistics
///
/// Typical usage
/// ```
/// let timer: SampleTimer
///
/// for i in 0...<iterationCount {
/// timer.start()
/// doStuff()
/// timer.stop()
/// }
///
/// print(String(format: "mean: %.2f, var: %.2f",
/// timer.mean, timer.variance))
/// ```
@available(iOS 16.2, macOS 13.1, *)
public final class SampleTimer: Codable {
var startTime: CFAbsoluteTime?
var sum: Double = 0.0
var sumOfSquares: Double = 0.0
var count = 0
var samples: [Double] = []
public init() {}
/// Start a sample, noting the current time
public func start() {
startTime = CFAbsoluteTimeGetCurrent()
}
// Stop a sample and record the elapsed time
@discardableResult public func stop() -> Double {
guard let startTime = startTime else {
return 0
}
let elapsed = CFAbsoluteTimeGetCurrent() - startTime
sum += elapsed
sumOfSquares += elapsed * elapsed
count += 1
samples.append(elapsed)
return elapsed
}
/// Mean of all sampled times
public var mean: Double { sum / Double(count) }
/// Variance of all sampled times
public var variance: Double {
guard count > 1 else {
return 0.0
}
return sumOfSquares / Double(count - 1) - mean * mean
}
/// Standard deviation of all sampled times
public var stdev: Double { variance.squareRoot() }
/// Median of all sampled times
public var median: Double {
let sorted = samples.sorted()
let (q, r) = sorted.count.quotientAndRemainder(dividingBy: 2)
if r == 0 {
return (sorted[q] + sorted[q - 1]) / 2.0
} else {
return Double(sorted[q])
}
}
public var allSamples: [Double] {
samples
}
}
swift/StableDiffusion/pipeline/Scheduler.swift
+363
View File
@@ -0,0 +1,363 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Accelerate
import CoreML
@available(iOS 16.2, macOS 13.1, *)
public protocol Scheduler {
/// Number of diffusion steps performed during training
var trainStepCount: Int { get }
/// Number of inference steps to be performed
var inferenceStepCount: Int { get }
/// Training diffusion time steps index by inference time step
var timeSteps: [Int] { get }
/// Training diffusion time steps index by inference time step
func calculateTimesteps(strength: Float?) -> [Int]
/// Schedule of betas which controls the amount of noise added at each timestep
var betas: [Float] { get }
/// 1 - betas
var alphas: [Float] { get }
/// Cached cumulative product of alphas
var alphasCumProd: [Float] { get }
/// Standard deviation of the initial noise distribution
var initNoiseSigma: Float { get }
/// Denoised latents
var modelOutputs: [MLShapedArray<Float32>] { get }
/// Compute a de-noised image sample and step scheduler state
///
/// - Parameters:
/// - output: The predicted residual noise output of learned diffusion model
/// - timeStep: The current time step in the diffusion chain
/// - sample: The current input sample to the diffusion model
/// - Returns: Predicted de-noised sample at the previous time step
/// - Postcondition: The scheduler state is updated.
/// The state holds the current sample and history of model output noise residuals
func step(
output: MLShapedArray<Float32>,
timeStep t: Int,
sample s: MLShapedArray<Float32>
) -> MLShapedArray<Float32>
}
@available(iOS 16.2, macOS 13.1, *)
public extension Scheduler {
var initNoiseSigma: Float { 1 }
}
@available(iOS 16.2, macOS 13.1, *)
public extension Scheduler {
/// Compute weighted sum of shaped arrays of equal shapes
///
/// - Parameters:
/// - weights: The weights each array is multiplied by
/// - values: The arrays to be weighted and summed
/// - Returns: sum_i weights[i]*values[i]
func weightedSum(_ weights: [Double], _ values: [MLShapedArray<Float32>]) -> MLShapedArray<Float32> {
let scalarCount = values.first!.scalarCount
assert(weights.count > 1 && values.count == weights.count)
assert(values.allSatisfy({ $0.scalarCount == scalarCount }))
return MLShapedArray(unsafeUninitializedShape: values.first!.shape) { scalars, _ in
scalars.initialize(repeating: 0.0)
for i in 0 ..< values.count {
let w = Float(weights[i])
values[i].withUnsafeShapedBufferPointer { buffer, _, _ in
assert(buffer.count == scalarCount)
// scalars[j] = w * values[i].scalars[j]
cblas_saxpy(Int32(scalarCount), w, buffer.baseAddress, 1, scalars.baseAddress, 1)
}
}
}
}
func addNoise(
originalSample: MLShapedArray<Float32>,
noise: [MLShapedArray<Float32>],
strength: Float
) -> [MLShapedArray<Float32>] {
let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
let alphaProdt = alphasCumProd[timeSteps[startStep]]
let betaProdt = 1 - alphaProdt
let sqrtAlphaProdt = sqrt(alphaProdt)
let sqrtBetaProdt = sqrt(betaProdt)
let noisySamples = noise.map {
weightedSum(
[Double(sqrtAlphaProdt), Double(sqrtBetaProdt)],
[originalSample, $0]
)
}
return noisySamples
}
}
// MARK: - Timesteps
@available(iOS 16.2, macOS 13.1, *)
public extension Scheduler {
func calculateTimesteps(strength: Float?) -> [Int] {
guard let strength else { return timeSteps }
let startStep = max(inferenceStepCount - Int(Float(inferenceStepCount) * strength), 0)
let actualTimesteps = Array(timeSteps[startStep...])
return actualTimesteps
}
}
// MARK: - BetaSchedule
/// How to map a beta range to a sequence of betas to step over
@available(iOS 16.2, macOS 13.1, *)
public enum BetaSchedule {
/// Linear stepping between start and end
case linear
/// Steps using linspace(sqrt(start),sqrt(end))^2
case scaledLinear
}
// MARK: - PNDMScheduler
/// A scheduler used to compute a de-noised image
///
/// This implementation matches:
/// [Hugging Face Diffusers PNDMScheduler](https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_pndm.py)
///
/// This scheduler uses the pseudo linear multi-step (PLMS) method only, skipping pseudo Runge-Kutta (PRK) steps
@available(iOS 16.2, macOS 13.1, *)
public final class PNDMScheduler: Scheduler {
public let trainStepCount: Int
public let inferenceStepCount: Int
public let betas: [Float]
public let alphas: [Float]
public let alphasCumProd: [Float]
public let timeSteps: [Int]
public let alpha_t: [Float]
public let sigma_t: [Float]
public let lambda_t: [Float]
public private(set) var modelOutputs: [MLShapedArray<Float32>] = []
// Internal state
var counter: Int
var ets: [MLShapedArray<Float32>]
var currentSample: MLShapedArray<Float32>?
/// Create a scheduler that uses a pseudo linear multi-step (PLMS) method
///
/// - Parameters:
/// - stepCount: Number of inference steps to schedule
/// - trainStepCount: Number of training diffusion steps
/// - betaSchedule: Method to schedule betas from betaStart to betaEnd
/// - betaStart: The starting value of beta for inference
/// - betaEnd: The end value for beta for inference
/// - Returns: A scheduler ready for its first step
public init(
stepCount: Int = 50,
trainStepCount: Int = 1000,
betaSchedule: BetaSchedule = .scaledLinear,
betaStart: Float = 0.00085,
betaEnd: Float = 0.012
) {
self.trainStepCount = trainStepCount
self.inferenceStepCount = stepCount
switch betaSchedule {
case .linear:
self.betas = linspace(betaStart, betaEnd, trainStepCount)
case .scaledLinear:
self.betas = linspace(pow(betaStart, 0.5), pow(betaEnd, 0.5), trainStepCount).map({ $0 * $0 })
}
self.alphas = betas.map({ 1.0 - $0 })
var alphasCumProd = self.alphas
for i in 1..<alphasCumProd.count {
alphasCumProd[i] *= alphasCumProd[i - 1]
}
self.alphasCumProd = alphasCumProd
let stepsOffset = 1 // For stable diffusion
let stepRatio = Float(trainStepCount / stepCount )
let forwardSteps = (0..<stepCount).map {
Int((Float($0) * stepRatio).rounded()) + stepsOffset
}
self.alpha_t = vForce.sqrt(self.alphasCumProd)
self.sigma_t = vForce.sqrt(vDSP.subtract([Float](repeating: 1, count: self.alphasCumProd.count), self.alphasCumProd))
self.lambda_t = zip(self.alpha_t, self.sigma_t).map { α, σ in log(α) - log(σ) }
var timeSteps: [Int] = []
timeSteps.append(contentsOf: forwardSteps.dropLast(1))
timeSteps.append(timeSteps.last!)
timeSteps.append(forwardSteps.last!)
timeSteps.reverse()
self.timeSteps = timeSteps
self.counter = 0
self.ets = []
self.currentSample = nil
}
/// Compute a de-noised image sample and step scheduler state
///
/// - Parameters:
/// - output: The predicted residual noise output of learned diffusion model
/// - timeStep: The current time step in the diffusion chain
/// - sample: The current input sample to the diffusion model
/// - Returns: Predicted de-noised sample at the previous time step
/// - Postcondition: The scheduler state is updated.
/// The state holds the current sample and history of model output noise residuals
public func step(
output: MLShapedArray<Float32>,
timeStep t: Int,
sample s: MLShapedArray<Float32>
) -> MLShapedArray<Float32> {
var timeStep = t
let stepInc = (trainStepCount / inferenceStepCount)
var prevStep = timeStep - stepInc
var modelOutput = output
var sample = s
if counter != 1 {
if ets.count > 3 {
ets = Array(ets[(ets.count - 3)..<ets.count])
}
ets.append(output)
} else {
prevStep = timeStep
timeStep = timeStep + stepInc
}
if ets.count == 1 && counter == 0 {
modelOutput = output
currentSample = sample
} else if ets.count == 1 && counter == 1 {
modelOutput = weightedSum(
[1.0/2.0, 1.0/2.0],
[output, ets[back: 1]]
)
sample = currentSample!
currentSample = nil
} else if ets.count == 2 {
modelOutput = weightedSum(
[3.0/2.0, -1.0/2.0],
[ets[back: 1], ets[back: 2]]
)
} else if ets.count == 3 {
modelOutput = weightedSum(
[23.0/12.0, -16.0/12.0, 5.0/12.0],
[ets[back: 1], ets[back: 2], ets[back: 3]]
)
} else {
modelOutput = weightedSum(
[55.0/24.0, -59.0/24.0, 37/24.0, -9/24.0],
[ets[back: 1], ets[back: 2], ets[back: 3], ets[back: 4]]
)
}
let convertedOutput = convertModelOutput(modelOutput: modelOutput, timestep: timeStep, sample: sample)
modelOutputs.append(convertedOutput)
let prevSample = previousSample(sample, timeStep, prevStep, modelOutput)
counter += 1
return prevSample
}
/// Convert the model output to the corresponding type the algorithm needs.
func convertModelOutput(modelOutput: MLShapedArray<Float32>, timestep: Int, sample: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
assert(modelOutput.scalarCount == sample.scalarCount)
let scalarCount = modelOutput.scalarCount
let (alpha_t, sigma_t) = (self.alpha_t[timestep], self.sigma_t[timestep])
return MLShapedArray(unsafeUninitializedShape: modelOutput.shape) { scalars, _ in
assert(scalars.count == scalarCount)
modelOutput.withUnsafeShapedBufferPointer { modelOutput, _, _ in
sample.withUnsafeShapedBufferPointer { sample, _, _ in
for i in 0 ..< scalarCount {
scalars.initializeElement(at: i, to: (sample[i] - modelOutput[i] * sigma_t) / alpha_t)
}
}
}
}
}
/// Compute sample (denoised image) at previous step given a current time step
///
/// - Parameters:
/// - sample: The current input to the model x_t
/// - timeStep: The current time step t
/// - prevStep: The previous time step t−δ
/// - modelOutput: Predicted noise residual the current time step e_θ(x_t, t)
/// - Returns: Computes previous sample x_(t−δ)
func previousSample(
_ sample: MLShapedArray<Float32>,
_ timeStep: Int,
_ prevStep: Int,
_ modelOutput: MLShapedArray<Float32>
) -> MLShapedArray<Float32> {
// Compute x_(t−δ) using formula (9) from
// "Pseudo Numerical Methods for Diffusion Models on Manifolds",
// Luping Liu, Yi Ren, Zhijie Lin & Zhou Zhao.
// ICLR 2022
//
// Notation:
//
// alphaProdt α_t
// alphaProdtPrev α_(t−δ)
// betaProdt (1 - α_t)
// betaProdtPrev (1 - α_(t−δ))
let alphaProdt = alphasCumProd[timeStep]
let alphaProdtPrev = alphasCumProd[max(0,prevStep)]
let betaProdt = 1 - alphaProdt
let betaProdtPrev = 1 - alphaProdtPrev
// sampleCoeff = (α_(t−δ) - α_t) divided by
// denominator of x_t in formula (9) and plus 1
// Note: (α_(t−δ) - α_t) / (sqrt(α_t) * (sqrt(α_(t−δ)) + sqr(α_t))) =
// sqrt(α_(t−δ)) / sqrt(α_t))
let sampleCoeff = sqrt(alphaProdtPrev / alphaProdt)
// Denominator of e_θ(x_t, t) in formula (9)
let modelOutputDenomCoeff = alphaProdt * sqrt(betaProdtPrev)
+ sqrt(alphaProdt * betaProdt * alphaProdtPrev)
// full formula (9)
let modelCoeff = -(alphaProdtPrev - alphaProdt)/modelOutputDenomCoeff
let prevSample = weightedSum(
[Double(sampleCoeff), Double(modelCoeff)],
[sample, modelOutput]
)
return prevSample
}
}
/// Evenly spaced floats between specified interval
///
/// - Parameters:
/// - start: Start of the interval
/// - end: End of the interval
/// - count: The number of floats to return between [*start*, *end*]
/// - Returns: Float array with *count* elements evenly spaced between at *start* and *end*
func linspace(_ start: Float, _ end: Float, _ count: Int) -> [Float] {
let scale = (end - start) / Float(count - 1)
return (0..<count).map { Float($0)*scale + start }
}
extension Collection {
/// Collection element index from the back. *self[back: 1]* yields the last element
public subscript(back i: Int) -> Element {
return self[index(endIndex, offsetBy: -i)]
}
}
swift/StableDiffusion/pipeline/StableDiffusion3Pipeline+Resources.swift
+97
View File
@@ -0,0 +1,97 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2024 Apple Inc. All Rights Reserved.
import CoreML
import Foundation
import Tokenizers
import Hub
@available(iOS 17.0, macOS 14.0, *)
public extension StableDiffusion3Pipeline {
struct ResourceURLs {
public let textEncoderURL: URL
public let textEncoder2URL: URL
public let textEncoderT5URL: URL
public let mmditURL: URL
public let decoderURL: URL
public let encoderURL: URL
public let vocabURL: URL
public let mergesURL: URL
public let configT5URL: URL
public let dataT5URL: URL
public init(resourcesAt baseURL: URL) {
textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
textEncoder2URL = baseURL.appending(path: "TextEncoder2.mlmodelc")
textEncoderT5URL = baseURL.appending(path: "TextEncoderT5.mlmodelc")
mmditURL = baseURL.appending(path: "MultiModalDiffusionTransformer.mlmodelc")
decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
vocabURL = baseURL.appending(path: "vocab.json")
mergesURL = baseURL.appending(path: "merges.txt")
configT5URL = baseURL.appending(path: "tokenizer_config.json")
dataT5URL = baseURL.appending(path: "tokenizer.json")
}
}
/// Create stable diffusion pipeline using model resources at a
/// specified URL
///
/// - Parameters:
/// - baseURL: URL pointing to directory holding all model and tokenization resources
/// - configuration: The configuration to load model resources with
/// - reduceMemory: Setup pipeline in reduced memory mode
/// - Returns:
/// Pipeline ready for image generation if all necessary resources loaded
init(
resourcesAt baseURL: URL,
configuration config: MLModelConfiguration = .init(),
reduceMemory: Bool = false
) throws {
// Expect URL of each resource
let urls = ResourceURLs(resourcesAt: baseURL)
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
let textEncoder = TextEncoderXL(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
// padToken is different in the second XL text encoder
let tokenizer2 = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL, padToken: "!")
let textEncoder2 = TextEncoderXL(tokenizer: tokenizer2, modelAt: urls.textEncoder2URL, configuration: config)
// Optional T5 encoder
var textEncoderT5: TextEncoderT5?
if FileManager.default.fileExists(atPath: urls.configT5URL.path),
FileManager.default.fileExists(atPath: urls.dataT5URL.path),
FileManager.default.fileExists(atPath: urls.textEncoderT5URL.path)
{
let tokenizerT5 = try PreTrainedTokenizer(tokenizerConfig: Config(fileURL: urls.configT5URL), tokenizerData: Config(fileURL: urls.dataT5URL))
textEncoderT5 = TextEncoderT5(tokenizer: tokenizerT5, modelAt: urls.textEncoderT5URL, configuration: config)
} else {
textEncoderT5 = nil
}
// Denoiser model
let mmdit = MultiModalDiffusionTransformer(modelAt: urls.mmditURL, configuration: config)
// Image Decoder
let decoder = Decoder(modelAt: urls.decoderURL, configuration: config)
// Optional Image Encoder
let encoder: Encoder?
if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
encoder = Encoder(modelAt: urls.encoderURL, configuration: config)
} else {
encoder = nil
}
// Construct pipeline
self.init(
textEncoder: textEncoder,
textEncoder2: textEncoder2,
textEncoderT5: textEncoderT5,
mmdit: mmdit,
decoder: decoder,
encoder: encoder,
reduceMemory: reduceMemory
)
}
}
swift/StableDiffusion/pipeline/StableDiffusion3Pipeline.swift
+486
View File
@@ -0,0 +1,486 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2024 Apple Inc. All Rights Reserved.
import Accelerate
import CoreGraphics
import CoreImage
import CoreML
import Foundation
@available(iOS 17.0, macOS 14.0, *)
public struct StableDiffusion3Pipeline: StableDiffusionPipelineProtocol {
public typealias Configuration = PipelineConfiguration
public typealias Progress = PipelineProgress
/// Model to generate embeddings for tokenized input text
var textEncoder: TextEncoderXLModel
var textEncoder2: TextEncoderXLModel
var textEncoderT5: TextEncoderT5Model?
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
var mmdit: MultiModalDiffusionTransformer
/// Model used to generate final image from latent diffusion process
var decoder: Decoder
/// Model used to latent space for image2image, and soon, in-painting
var encoder: Encoder?
/// Option to reduce memory during image generation
///
/// If true, the pipeline will lazily load TextEncoder, Unet, Decoder, and SafetyChecker
/// when needed and aggressively unload their resources after
///
/// This will increase latency in favor of reducing memory
var reduceMemory: Bool = false
/// Creates a pipeline using the specified models and tokenizer
///
/// - Parameters:
/// - textEncoder: Model for encoding tokenized text
/// - textEncoder2: Second text encoding model
/// - mmdit: Model for noise prediction on latent samples
/// - decoder: Model for decoding latent sample to image
/// - reduceMemory: Option to enable reduced memory mode
/// - Returns: Pipeline ready for image generation
public init(
textEncoder: TextEncoderXLModel,
textEncoder2: TextEncoderXLModel,
textEncoderT5: TextEncoderT5?,
mmdit: MultiModalDiffusionTransformer,
decoder: Decoder,
encoder: Encoder?,
reduceMemory: Bool = false
) {
self.textEncoder = textEncoder
self.textEncoder2 = textEncoder2
self.textEncoderT5 = textEncoderT5
self.mmdit = mmdit
self.decoder = decoder
self.encoder = encoder
self.reduceMemory = reduceMemory
}
/// Load required resources for this pipeline
///
/// If reducedMemory is true this will instead call prewarmResources instead
/// and let the pipeline lazily load resources as needed
public func loadResources() throws {
if reduceMemory {
try prewarmResources()
} else {
try textEncoder.loadResources()
try textEncoder2.loadResources()
try textEncoderT5?.loadResources()
try mmdit.loadResources()
try decoder.loadResources()
do {
try encoder?.loadResources()
} catch {
print("Error loading resources for vae encoder: \(error)")
}
}
}
/// Unload the underlying resources to free up memory
public func unloadResources() {
textEncoder.unloadResources()
textEncoder2.unloadResources()
textEncoderT5?.unloadResources()
mmdit.unloadResources()
decoder.unloadResources()
encoder?.unloadResources()
}
/// Prewarm resources one at a time
public func prewarmResources() throws {
try textEncoder.prewarmResources()
try textEncoder2.prewarmResources()
try textEncoderT5?.prewarmResources()
try mmdit.prewarmResources()
try decoder.prewarmResources()
do {
try encoder?.prewarmResources()
} catch {
print("Error prewarming resources for vae encoder: \(error)")
}
}
/// Image generation using stable diffusion
/// - Parameters:
/// - configuration: Image generation configuration
/// - progressHandler: Callback to perform after each step, stops on receiving false response
/// - Returns: An array of `imageCount` optional images.
/// The images will be nil if safety checks were performed and found the result to be un-safe
public func generateImages(
configuration config: Configuration,
progressHandler: (Progress) -> Bool = { _ in true }
) throws -> [CGImage?] {
// Setup geometry conditioning for base/refiner inputs
let sd3Input: ModelInputs = try generateConditioning(using: config)
if reduceMemory {
textEncoder.unloadResources()
textEncoder2.unloadResources()
textEncoderT5?.unloadResources()
}
// Setup schedulers
let scheduler: [DiscreteFlowScheduler] = (0..<config.imageCount).map { _ in
DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
}
// Generate random latent samples from specified seed
var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
// Store denoised latents from scheduler to pass into decoder
var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
if reduceMemory {
encoder?.unloadResources()
}
let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
// Store current model
let mmditModel = mmdit
let mmditHiddenStates = sd3Input.hiddenStates
let mmditPooledStates = sd3Input.pooledStates
let timeSteps: [Float] = scheduler[0].calculateTimestepsFromSigmas(strength: timestepStrength)
// De-noising loop
for (step, t) in timeSteps.enumerated() {
// Expand the latents for classifier-free guidance
// and input to the MMDiT noise prediction model
let latentUnetInput = latents.map {
MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
}
// Predict noise residuals from latent samples
// and current time step conditioned on hidden states
var noise = try mmditModel.predictNoise(
latents: latentUnetInput,
timeStep: t,
tokenLevelTextEmbeddings: mmditHiddenStates,
pooledTextEmbeddings: mmditPooledStates
)
noise = performGuidance(noise, config.guidanceScale)
// Have the scheduler compute the previous (t-1) latent
// sample given the predicted noise and current sample
for i in 0..<config.imageCount {
latents[i] = scheduler[i].step(
output: noise[i],
timeStep: scheduler[i].timeSteps[step], // TODO: allow float timesteps in scheduler step protocol
sample: latents[i]
)
denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
}
let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
// Report progress
let progress = Progress(
pipeline: self,
prompt: config.prompt,
step: step,
stepCount: timeSteps.count,
currentLatentSamples: currentLatentSamples,
configuration: config
)
if !progressHandler(progress) {
// Stop if requested by handler
return []
}
}
// Unload resources
if reduceMemory {
mmdit.unloadResources()
}
// Decode the latent samples to images
return try decodeToImages(denoisedLatents, configuration: config)
}
func encodePrompt(_ prompt: String) throws -> (MLShapedArray<Float32>, MLShapedArray<Float32>) {
var embeds = MLShapedArray<Float32>()
var pooled = MLShapedArray<Float32>()
let (embeds1, pooledValue1) = try textEncoder.encode(prompt)
let (embeds2, pooledValue2) = try textEncoder2.encode(prompt)
var embedsT5 = try textEncoderT5?.encode(prompt).encoderHiddenStates ?? MLShapedArray<Float32>(repeating: 0, shape: [1, 4096, 1, 77])
// Truncate T5
embedsT5 = truncatedT5Embeds(embedsT5)
let padding1 = MLShapedArray<Float32>(repeating: 0, shape: [1, 77, 2048])
// Base needs concatenated embeddings
// [1, 77, 768], [1, 77, 1280], [1, 77, 2048] -> [1, 77, 4096]
embeds = MLShapedArray<Float32>(
concatenating: [embeds1, embeds2, padding1],
alongAxis: 2
)
// [1, 77, 4096] -> [1, 4096, 1 77]
embeds = toHiddenStates(embeds)
// [1, 4096, 1 77], [1, 4096, 1, 77] -> [1, 4096, 1, 154]
embeds = MLShapedArray<Float32>(
concatenating: [embeds, embedsT5],
alongAxis: 3
)
// [1, 768], [1, 1280] -> [1, 2048]
pooled = MLShapedArray<Float32>(
concatenating: [pooledValue1, pooledValue2],
alongAxis: 1
)
return (embeds, pooled)
}
func generateConditioning(using config: Configuration) throws -> ModelInputs {
// Encode the input prompt and negative prompt
let (promptEmbedding, pooled) = try encodePrompt(config.prompt)
let (negativePromptEmbedding, negativePooled) = try encodePrompt(config.negativePrompt)
// Convert to Unet hidden state representation
// Concatenate the prompt and negative prompt embeddings
let hiddenStates = MLShapedArray(concatenating: [promptEmbedding, negativePromptEmbedding], alongAxis: 0)
let pooledScalars = MLShapedArray(concatenating: [pooled, negativePooled], alongAxis: 0)
let pooledStates = MLShapedArray<Float32>(
scalars: pooledScalars.scalars,
shape: [2, 2048, 1, 1]
)
return ModelInputs(hiddenStates: hiddenStates, pooledStates: pooledStates)
}
func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
var sampleShape = mmdit.latentImageEmbeddingsShape
sampleShape[0] = 1
let stdev = scheduler.initNoiseSigma
var random = randomSource(from: config.rngType, seed: config.seed)
let samples = (0..<config.imageCount).map { _ in
MLShapedArray<Float32>(
converting: random.normalShapedArray(sampleShape, mean: 0.0, stdev: Double(stdev)))
}
if let image = config.startingImage, config.mode == .imageToImage {
guard let encoder else {
throw PipelineError.startingImageProvidedWithoutEncoder
}
let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
}
return samples
}
func performGuidance(_ noise: [MLShapedArray<Float32>], _ guidanceScale: Float) -> [MLShapedArray<Float32>] {
noise.map { performGuidance($0, guidanceScale) }
}
func performGuidance(_ noise: MLShapedArray<Float32>, _ guidanceScale: Float) -> MLShapedArray<Float32> {
var shape = noise.shape
shape[0] = 1
return MLShapedArray<Float>(unsafeUninitializedShape: shape) { result, _ in
noise.withUnsafeShapedBufferPointer { scalars, _, strides in
for i in 0..<result.count {
// unconditioned + guidance*(text - unconditioned)
let text = scalars[i]
let negText = scalars[strides[0] + i]
let guidance = negText + guidanceScale * (text - negText)
result.initializeElement(
at: i,
to: guidance
)
}
}
}
}
public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
defer {
if reduceMemory {
decoder.unloadResources()
}
}
return try decoder.decode(latents, scaleFactor: config.decoderScaleFactor, shiftFactor: config.decoderShiftFactor)
// TODO: use latent rgb factors with blur for preview images
// This will require a method to decode with either the vae or the rgb factors depending on config
// return try decodePreviewImage(latents, scaleFactor: config.decoderScaleFactor)
}
/// Shape 16 x 3
let rgbFactors: [[Float]] = [
[-0.0645, 0.0177, 0.1052], [ 0.0028, 0.0312, 0.0650],
[ 0.1848, 0.0762, 0.0360], [ 0.0944, 0.0360, 0.0889],
[ 0.0897, 0.0506, -0.0364], [-0.0020, 0.1203, 0.0284],
[ 0.0855, 0.0118, 0.0283], [-0.0539, 0.0658, 0.1047],
[-0.0057, 0.0116, 0.0700], [-0.0412, 0.0281, -0.0039],
[ 0.1106, 0.1171, 0.1220], [-0.0248, 0.0682, -0.0481],
[ 0.0815, 0.0846, 0.1207], [-0.0120, -0.0055, -0.0867],
[-0.0749, -0.0634, -0.0456], [-0.1418, -0.1457, -0.1259]
]
public func decodePreviewImage(
_ latents: [MLShapedArray<Float32>],
scaleFactor: Float32
) throws -> [CGImage] {
let height = 64
let width = 64
let channels = 16
let outputChannels = 3
// Ensure there is a first element in latents and extract its scalars
guard let latentScalars = latents.first?.scalars else {
throw NSError(domain: "DecodeError", code: 0, userInfo: [NSLocalizedDescriptionKey: "Invalid latent array"])
}
// The latentScalars is a flat array, we need to reshape and multiply
var reshapedLatent = [Float32](repeating: 0, count: height * width * channels)
// We reorder the indices manually to switch from [channels, height, width] to [height, width, channels]
for h in 0..<height {
for w in 0..<width {
for c in 0..<channels {
let oldIndex = c * height * width + h * width + w
let newIndex = h * width * channels + w * channels + c
reshapedLatent[newIndex] = latentScalars[oldIndex] // 1.5305 + 0.0609
}
}
}
// Prepare to hold the result of the multiplication
var imageArray = [Float32](repeating: 0, count: height * width * outputChannels)
// Perform matrix multiplication using Accelerate
vDSP_mmul(reshapedLatent, 1,
rgbFactors.flatMap { $0 }, 1,
&imageArray, 1,
vDSP_Length(height * width), // number of rows in output
vDSP_Length(outputChannels), // number of columns in output
vDSP_Length(channels)) // common dimension
// Convert imageArray into a CGImage
let latentImage = imageArray.toCGImage(width: width, height: height)
// Apply a Gaussian blur to the preview image to reduce pixeled look
let ciImage = CIImage(cgImage: latentImage!)
let blurFilter = CIFilter(name: "CIGaussianBlur")!
blurFilter.setValue(ciImage, forKey: kCIInputImageKey)
blurFilter.setValue(4.0, forKey: kCIInputRadiusKey)
let context = CIContext()
guard let outputImage = blurFilter.outputImage,
let cgBlurredPreview = context.createCGImage(outputImage, from: ciImage.extent)
else {
throw PipelineError.errorCreatingPreview
}
return [cgBlurredPreview]
}
struct ModelInputs {
var hiddenStates: MLShapedArray<Float32>
var pooledStates: MLShapedArray<Float32>
}
/// Helper function to truncate the T5 embeddings
func truncatedT5Embeds(_ embedding: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
// Unoptimized manual truncation
// e.g. From [1, 4096, 1, 128] to [1, 4096, 1, 77]
let fromShape = embedding.shape
let stateShape = [fromShape[0], fromShape[1], fromShape[2], 77]
var states = MLShapedArray<Float32>(repeating: 0.0, shape: stateShape)
for i0 in 0..<fromShape[0] {
for i1 in 0..<fromShape[1] {
for i2 in 0..<fromShape[2] {
for i3 in 0..<stateShape[3] {
states[scalarAt: i0, i1, i2, i3] = embedding[scalarAt: i0, i1, i2, i3]
}
}
}
}
return states
}
}
extension Array where Element == Float32 {
func toCGImage(width: Int, height: Int) -> CGImage? {
// Define color space and bitmap info
let colorSpace = CGColorSpaceCreateDeviceRGB()
let bitmapInfo = CGBitmapInfo.byteOrder32Big.rawValue | CGImageAlphaInfo.premultipliedLast.rawValue
// Calculate bytes per pixel and bytes per row
let bytesPerPixel = 4
let bytesPerRow = width * bytesPerPixel
// Allocate memory for the pixel data
var data = [UInt8](repeating: 0, count: height * bytesPerRow)
// Fill the data array with pixel data
for h in 0..<height {
for w in 0..<width {
let pixelIndex = h * width + w
let dataIndex = h * bytesPerRow + w * bytesPerPixel
let pixelBase = pixelIndex * 3 // Base index for R, G, B values in the source array
// Ensure your source array has enough data
if (pixelBase + 3) < self.count {
let redValue = (self[pixelBase] + 1) / 2 * 255
let bluValue = (self[pixelBase + 1] + 1) / 2 * 255
let grnValue = (self[pixelBase + 2] + 1) / 2 * 255
data[dataIndex] = UInt8(clamp(value: redValue, lower: 0, upper: 255)) // Red
data[dataIndex + 1] = UInt8(clamp(value: bluValue, lower: 0, upper: 255)) // Green
data[dataIndex + 2] = UInt8(clamp(value: grnValue, lower: 0, upper: 255)) // Blue
data[dataIndex + 3] = 255 // Alpha
}
}
}
// Create the context
guard let context = CGContext(data: &data, width: width, height: height, bitsPerComponent: 8, bytesPerRow: bytesPerRow, space: colorSpace, bitmapInfo: bitmapInfo) else {
print("Failed to create CGContext.")
return nil
}
// Create a CGImage from context
guard let smallImage = context.makeImage() else {
return nil
}
// Define the upscaled dimensions
let scaledWidth = width * 8
let scaledHeight = height * 8
// Create a new context with scaled dimensions
guard let largeContext = CGContext(data: nil, width: scaledWidth, height: scaledHeight, bitsPerComponent: 8, bytesPerRow: scaledWidth * 4, space: colorSpace, bitmapInfo: bitmapInfo) else {
return nil
}
// Draw the small image into the large context
largeContext.interpolationQuality = .high
largeContext.draw(smallImage, in: CGRect(x: 0, y: 0, width: scaledWidth, height: scaledHeight))
// Convert the upscaled context to a CGImage
return largeContext.makeImage()
}
/// Helper function to clamp the values within the specified range
private func clamp(value: Float32, lower: UInt8, upper: UInt8) -> UInt8 {
return UInt8(Swift.max(Float32(lower), Swift.min(value, Float32(upper))))
}
}
swift/StableDiffusion/pipeline/StableDiffusionPipeline+Resources.swift
+166
View File
@@ -0,0 +1,166 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
import NaturalLanguage
@available(iOS 16.2, macOS 13.1, *)
public extension StableDiffusionPipeline {
struct ResourceURLs {
public let textEncoderURL: URL
public let unetURL: URL
public let unetChunk1URL: URL
public let unetChunk2URL: URL
public let decoderURL: URL
public let encoderURL: URL
public let safetyCheckerURL: URL
public let vocabURL: URL
public let mergesURL: URL
public let controlNetDirURL: URL
public let controlledUnetURL: URL
public let controlledUnetChunk1URL: URL
public let controlledUnetChunk2URL: URL
public let multilingualTextEncoderProjectionURL: URL
public init(resourcesAt baseURL: URL) {
textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
unetURL = baseURL.appending(path: "Unet.mlmodelc")
unetChunk1URL = baseURL.appending(path: "UnetChunk1.mlmodelc")
unetChunk2URL = baseURL.appending(path: "UnetChunk2.mlmodelc")
decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
safetyCheckerURL = baseURL.appending(path: "SafetyChecker.mlmodelc")
vocabURL = baseURL.appending(path: "vocab.json")
mergesURL = baseURL.appending(path: "merges.txt")
controlNetDirURL = baseURL.appending(path: "controlnet")
controlledUnetURL = baseURL.appending(path: "ControlledUnet.mlmodelc")
controlledUnetChunk1URL = baseURL.appending(path: "ControlledUnetChunk1.mlmodelc")
controlledUnetChunk2URL = baseURL.appending(path: "ControlledUnetChunk2.mlmodelc")
multilingualTextEncoderProjectionURL = baseURL.appending(path: "MultilingualTextEncoderProjection.mlmodelc")
}
}
/// Create stable diffusion pipeline using model resources at a
/// specified URL
///
/// - Parameters:
/// - baseURL: URL pointing to directory holding all model and tokenization resources
/// - controlNetModelNames: Specify ControlNet models to use in generation
/// - configuration: The configuration to load model resources with
/// - disableSafety: Load time disable of safety to save memory
/// - reduceMemory: Setup pipeline in reduced memory mode
/// - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
/// - script: Optional natural language script to use for the text encoder.
/// - Returns:
/// Pipeline ready for image generation if all necessary resources loaded
init(
resourcesAt baseURL: URL,
controlNet controlNetModelNames: [String],
configuration config: MLModelConfiguration = .init(),
disableSafety: Bool = false,
reduceMemory: Bool = false,
useMultilingualTextEncoder: Bool = false,
script: Script? = nil
) throws {
/// Expect URL of each resource
let urls = ResourceURLs(resourcesAt: baseURL)
let textEncoder: TextEncoderModel
#if canImport(NaturalLanguage.NLScript)
if useMultilingualTextEncoder {
guard #available(macOS 14.0, iOS 17.0, *) else { throw PipelineError.unsupportedOSVersion }
textEncoder = MultilingualTextEncoder(
modelAt: urls.multilingualTextEncoderProjectionURL,
configuration: config,
script: script ?? .latin
)
} else {
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
}
#else
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
textEncoder = TextEncoder(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
#endif
// ControlNet model
var controlNet: ControlNet? = nil
let controlNetURLs = controlNetModelNames.map { model in
let fileName = model + ".mlmodelc"
return urls.controlNetDirURL.appending(path: fileName)
}
if !controlNetURLs.isEmpty {
controlNet = ControlNet(modelAt: controlNetURLs, configuration: config)
}
// Unet model
let unet: Unet
let unetURL: URL, unetChunk1URL: URL, unetChunk2URL: URL
// if ControlNet available, Unet supports additional inputs from ControlNet
if controlNet == nil {
unetURL = urls.unetURL
unetChunk1URL = urls.unetChunk1URL
unetChunk2URL = urls.unetChunk2URL
} else {
unetURL = urls.controlledUnetURL
unetChunk1URL = urls.controlledUnetChunk1URL
unetChunk2URL = urls.controlledUnetChunk2URL
}
if FileManager.default.fileExists(atPath: unetChunk1URL.path) &&
FileManager.default.fileExists(atPath: unetChunk2URL.path) {
unet = Unet(chunksAt: [unetChunk1URL, unetChunk2URL],
configuration: config)
} else {
unet = Unet(modelAt: unetURL, configuration: config)
}
// Image Decoder
let decoder = Decoder(modelAt: urls.decoderURL, configuration: config)
// Optional safety checker
var safetyChecker: SafetyChecker? = nil
if !disableSafety &&
FileManager.default.fileExists(atPath: urls.safetyCheckerURL.path) {
safetyChecker = SafetyChecker(modelAt: urls.safetyCheckerURL, configuration: config)
}
// Optional Image Encoder
let encoder: Encoder?
if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
encoder = Encoder(modelAt: urls.encoderURL, configuration: config)
} else {
encoder = nil
}
// Construct pipeline
if #available(macOS 14.0, iOS 17.0, *) {
self.init(
textEncoder: textEncoder,
unet: unet,
decoder: decoder,
encoder: encoder,
controlNet: controlNet,
safetyChecker: safetyChecker,
reduceMemory: reduceMemory,
useMultilingualTextEncoder: useMultilingualTextEncoder,
script: script
)
} else {
self.init(
textEncoder: textEncoder,
unet: unet,
decoder: decoder,
encoder: encoder,
controlNet: controlNet,
safetyChecker: safetyChecker,
reduceMemory: reduceMemory
)
}
}
}
swift/StableDiffusion/pipeline/StableDiffusionPipeline.Configuration.swift
+100
View File
@@ -0,0 +1,100 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreGraphics
/// Type of processing that will be performed to generate an image
public enum PipelineMode {
case textToImage
case imageToImage
// case inPainting
}
/// Image generation configuration
public struct PipelineConfiguration: Hashable {
/// Text prompt to guide sampling
public var prompt: String
/// Negative text prompt to guide sampling
public var negativePrompt: String = ""
/// Starting image for image2image or in-painting
public var startingImage: CGImage? = nil
/// Fraction of inference steps to be used in `.imageToImage` pipeline mode
/// Must be between 0 and 1
/// Higher values will result in greater transformation of the `startingImage`
public var strength: Float = 1.0
/// Fraction of inference steps to at which to start using the refiner unet if present in `textToImage` mode
/// Must be between 0 and 1
/// Higher values will result in fewer refiner steps
public var refinerStart: Float = 0.8
/// Number of images to generate
public var imageCount: Int = 1
/// Number of inference steps to perform
public var stepCount: Int = 50
/// Random seed which to start generation
public var seed: UInt32 = 0
/// Controls the influence of the text prompt on sampling process (0=random images)
public var guidanceScale: Float = 7.5
/// List of Images for available ControlNet Models
public var controlNetInputs: [CGImage] = []
/// Safety checks are only performed if `self.canSafetyCheck && !disableSafety`
public var disableSafety: Bool = false
/// Enables progress updates to decode `currentImages` from denoised latent images for better previews
public var useDenoisedIntermediates: Bool = false
/// The type of Scheduler to use.
public var schedulerType: StableDiffusionScheduler = .pndmScheduler
/// The spacing to use for scheduler sigmas and time steps. Only supported when using `.dpmppScheduler`.
public var schedulerTimestepSpacing: TimeStepSpacing = .linspace
/// Resolution dependent shifting of timestep schedules
public var schedulerTimestepShift: Float = 3.0
/// The type of RNG to use
public var rngType: StableDiffusionRNG = .numpyRNG
/// Scale factor to use on the latent after encoding
public var encoderScaleFactor: Float32 = 0.18215
/// Scale factor to use on the latent before decoding
public var decoderScaleFactor: Float32 = 0.18215
/// Shift factor to use on the latent before decoding
public var decoderShiftFactor: Float32 = 0.0
/// If `originalSize` is not the same as `targetSize` the image will appear to be down- or upsampled.
/// Part of SDXL’s micro-conditioning as explained in section 2.2 of https://huggingface.co/papers/2307.01952.
public var originalSize: Float32 = 1024
/// `cropsCoordsTopLeft` can be used to generate an image that appears to be “cropped” from the position `cropsCoordsTopLeft` downwards.
/// Favorable, well-centered images are usually achieved by setting `cropsCoordsTopLeft` to (0, 0).
public var cropsCoordsTopLeft: Float32 = 0
/// For most cases, `target_size` should be set to the desired height and width of the generated image.
public var targetSize: Float32 = 1024
/// Used to simulate an aesthetic score of the generated image by influencing the positive text condition.
public var aestheticScore: Float32 = 6
/// Can be used to simulate an aesthetic score of the generated image by influencing the negative text condition.
public var negativeAestheticScore: Float32 = 2.5
/// Given the configuration, what mode will be used for generation
public var mode: PipelineMode {
guard startingImage != nil else {
return .textToImage
}
guard strength < 1.0 else {
return .textToImage
}
return .imageToImage
}
public init(
prompt: String
) {
self.prompt = prompt
}
}
@available(iOS 16.2, macOS 13.1, *)
public extension StableDiffusionPipeline {
/// Type of processing that will be performed to generate an image
typealias Mode = PipelineMode
/// Image generation configuration
typealias Configuration = PipelineConfiguration
}
swift/StableDiffusion/pipeline/StableDiffusionPipeline.swift
+484
View File
@@ -0,0 +1,484 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Accelerate
import CoreGraphics
import CoreML
import Foundation
import NaturalLanguage
/// Schedulers compatible with StableDiffusionPipeline
public enum StableDiffusionScheduler {
/// Scheduler that uses a pseudo-linear multi-step (PLMS) method
case pndmScheduler
/// Scheduler that uses a second order DPM-Solver++ algorithm
case dpmSolverMultistepScheduler
/// Scheduler for rectified flow based multimodal diffusion transformer models
case discreteFlowScheduler
}
/// RNG compatible with StableDiffusionPipeline
public enum StableDiffusionRNG {
/// RNG that matches numpy implementation
case numpyRNG
/// RNG that matches PyTorch CPU implementation.
case torchRNG
/// RNG that matches PyTorch CUDA implementation.
case nvidiaRNG
}
public enum PipelineError: String, Swift.Error {
case missingUnetInputs
case startingImageProvidedWithoutEncoder
case startingText2ImgWithoutTextEncoder
case unsupportedOSVersion
case errorCreatingPreview
}
@available(iOS 16.2, macOS 13.1, *)
public protocol StableDiffusionPipelineProtocol: ResourceManaging {
var canSafetyCheck: Bool { get }
func generateImages(
configuration config: PipelineConfiguration,
progressHandler: (PipelineProgress) -> Bool
) throws -> [CGImage?]
func decodeToImages(
_ latents: [MLShapedArray<Float32>],
configuration config: PipelineConfiguration
) throws -> [CGImage?]
}
@available(iOS 16.2, macOS 13.1, *)
public extension StableDiffusionPipelineProtocol {
var canSafetyCheck: Bool { false }
}
/// A pipeline used to generate image samples from text input using stable diffusion
///
/// This implementation matches:
/// [Hugging Face Diffusers Pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py)
@available(iOS 16.2, macOS 13.1, *)
public struct StableDiffusionPipeline: StableDiffusionPipelineProtocol {
/// Model to generate embeddings for tokenized input text
var textEncoder: TextEncoderModel
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
var unet: Unet
/// Model used to generate final image from latent diffusion process
var decoder: Decoder
/// Model used to latent space for image2image, and soon, in-painting
var encoder: Encoder?
/// Optional model for checking safety of generated image
var safetyChecker: SafetyChecker? = nil
/// Optional model used before Unet to control generated images by additonal inputs
var controlNet: ControlNet? = nil
/// Reports whether this pipeline can perform safety checks
public var canSafetyCheck: Bool {
safetyChecker != nil
}
/// Option to reduce memory during image generation
///
/// If true, the pipeline will lazily load TextEncoder, Unet, Decoder, and SafetyChecker
/// when needed and aggressively unload their resources after
///
/// This will increase latency in favor of reducing memory
var reduceMemory: Bool = false
/// Option to use system multilingual NLContextualEmbedding as encoder
var useMultilingualTextEncoder: Bool = false
/// Optional natural language script to use for the text encoder.
var script: Script? = nil
/// Creates a pipeline using the specified models and tokenizer
///
/// - Parameters:
/// - textEncoder: Model for encoding tokenized text
/// - unet: Model for noise prediction on latent samples
/// - decoder: Model for decoding latent sample to image
/// - controlNet: Optional model to control generated images by additonal inputs
/// - safetyChecker: Optional model for checking safety of generated images
/// - reduceMemory: Option to enable reduced memory mode
/// - Returns: Pipeline ready for image generation
public init(
textEncoder: TextEncoderModel,
unet: Unet,
decoder: Decoder,
encoder: Encoder?,
controlNet: ControlNet? = nil,
safetyChecker: SafetyChecker? = nil,
reduceMemory: Bool = false
) {
self.textEncoder = textEncoder
self.unet = unet
self.decoder = decoder
self.encoder = encoder
self.controlNet = controlNet
self.safetyChecker = safetyChecker
self.reduceMemory = reduceMemory
}
/// Creates a pipeline using the specified models and tokenizer
///
/// - Parameters:
/// - textEncoder: Model for encoding tokenized text
/// - unet: Model for noise prediction on latent samples
/// - decoder: Model for decoding latent sample to image
/// - controlNet: Optional model to control generated images by additonal inputs
/// - safetyChecker: Optional model for checking safety of generated images
/// - reduceMemory: Option to enable reduced memory mode
/// - useMultilingualTextEncoder: Option to use system multilingual NLContextualEmbedding as encoder
/// - script: Optional natural language script to use for the text encoder.
/// - Returns: Pipeline ready for image generation
@available(iOS 17.0, macOS 14.0, *)
public init(
textEncoder: TextEncoderModel,
unet: Unet,
decoder: Decoder,
encoder: Encoder?,
controlNet: ControlNet? = nil,
safetyChecker: SafetyChecker? = nil,
reduceMemory: Bool = false,
useMultilingualTextEncoder: Bool = false,
script: Script? = nil
) {
self.textEncoder = textEncoder
self.unet = unet
self.decoder = decoder
self.encoder = encoder
self.controlNet = controlNet
self.safetyChecker = safetyChecker
self.reduceMemory = reduceMemory
self.useMultilingualTextEncoder = useMultilingualTextEncoder
self.script = script
}
/// Load required resources for this pipeline
///
/// If reducedMemory is true this will instead call prewarmResources instead
/// and let the pipeline lazily load resources as needed
public func loadResources() throws {
if reduceMemory {
try prewarmResources()
} else {
try unet.loadResources()
try textEncoder.loadResources()
try decoder.loadResources()
try encoder?.loadResources()
try controlNet?.loadResources()
try safetyChecker?.loadResources()
}
}
/// Unload the underlying resources to free up memory
public func unloadResources() {
textEncoder.unloadResources()
unet.unloadResources()
decoder.unloadResources()
encoder?.unloadResources()
controlNet?.unloadResources()
safetyChecker?.unloadResources()
}
// Prewarm resources one at a time
public func prewarmResources() throws {
try textEncoder.prewarmResources()
try unet.prewarmResources()
try decoder.prewarmResources()
try encoder?.prewarmResources()
try controlNet?.prewarmResources()
try safetyChecker?.prewarmResources()
}
/// Image generation using stable diffusion
/// - Parameters:
/// - configuration: Image generation configuration
/// - progressHandler: Callback to perform after each step, stops on receiving false response
/// - Returns: An array of `imageCount` optional images.
/// The images will be nil if safety checks were performed and found the result to be un-safe
public func generateImages(
configuration config: Configuration,
progressHandler: (Progress) -> Bool = { _ in true }
) throws -> [CGImage?] {
// Encode the input prompt
var promptEmbedding = try textEncoder.encode(config.prompt)
if config.guidanceScale >= 1.0 {
// Convert to Unet hidden state representation
// Concatenate the prompt and negative prompt embeddings
let negativePromptEmbedding = try textEncoder.encode(config.negativePrompt)
promptEmbedding = MLShapedArray<Float32>(
concatenating: [negativePromptEmbedding, promptEmbedding],
alongAxis: 0
)
}
if reduceMemory {
textEncoder.unloadResources()
}
let hiddenStates = useMultilingualTextEncoder ? promptEmbedding : toHiddenStates(promptEmbedding)
/// Setup schedulers
let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
switch config.schedulerType {
case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
}
}
// Generate random latent samples from specified seed
var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
// Store denoised latents from scheduler to pass into decoder
var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
if reduceMemory {
encoder?.unloadResources()
}
let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
// Convert cgImage for ControlNet into MLShapedArray
let controlNetConds = try config.controlNetInputs.map { cgImage in
let shapedArray = try cgImage.planarRGBShapedArray(minValue: 0.0, maxValue: 1.0)
return MLShapedArray(
concatenating: [shapedArray, shapedArray],
alongAxis: 0
)
}
// De-noising loop
let timeSteps: [Int] = scheduler[0].calculateTimesteps(strength: timestepStrength)
for (step,t) in timeSteps.enumerated() {
// Expand the latents for classifier-free guidance
// and input to the Unet noise prediction model
let latentUnetInput: [MLShapedArray<Float32>]
if config.guidanceScale >= 1.0 {
latentUnetInput = latents.map {
MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
}
} else {
latentUnetInput = latents
}
// Before Unet, execute controlNet and add the output into Unet inputs
let additionalResiduals = try controlNet?.execute(
latents: latentUnetInput,
timeStep: t,
hiddenStates: hiddenStates,
images: controlNetConds
)
// Predict noise residuals from latent samples
// and current time step conditioned on hidden states
var noise : [MLShapedArray<Float32>]
if unet.latentSampleShape[0] >= 2 || config.guidanceScale < 1.0 {
// One predict call from the uNet, using batching if needed
noise = try unet.predictNoise(
latents: latentUnetInput,
timeStep: t,
hiddenStates: hiddenStates,
additionalResiduals: additionalResiduals
)
} else {
// Serial predictions from uNet
var hidden0 = MLShapedArray<Float32>(converting: hiddenStates[0])
hidden0 = MLShapedArray(scalars: hidden0.scalars, shape: [1]+hidden0.shape)
let noise_pred_uncond = try unet.predictNoise(
latents: latents,
timeStep: t,
hiddenStates: hidden0,
additionalResiduals: additionalResiduals
)
var hidden1 = MLShapedArray<Float32>(converting: hiddenStates[1])
hidden1 = MLShapedArray(scalars: hidden1.scalars, shape: [1]+hidden1.shape)
let noise_pred_text = try unet.predictNoise(
latents: latents,
timeStep: t,
hiddenStates: hidden1,
additionalResiduals: additionalResiduals
)
noise = [MLShapedArray<Float32>(concatenating: [noise_pred_uncond[0], noise_pred_text[0]],
alongAxis: 0)]
}
if config.guidanceScale >= 1.0 {
noise = performGuidance(noise, config.guidanceScale)
}
// Have the scheduler compute the previous (t-1) latent
// sample given the predicted noise and current sample
for i in 0..<config.imageCount {
latents[i] = scheduler[i].step(
output: noise[i],
timeStep: t,
sample: latents[i]
)
denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
}
let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
// Report progress
let progress = Progress(
pipeline: self,
prompt: config.prompt,
step: step,
stepCount: timeSteps.count,
currentLatentSamples: currentLatentSamples,
configuration: config
)
if !progressHandler(progress) {
// Stop if requested by handler
return []
}
}
if reduceMemory {
controlNet?.unloadResources()
unet.unloadResources()
}
// Decode the latent samples to images
return try decodeToImages(denoisedLatents, configuration: config)
}
func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
var sampleShape = unet.latentSampleShape
sampleShape[0] = 1
let stdev = scheduler.initNoiseSigma
var random = randomSource(from: config.rngType, seed: config.seed)
let samples = (0..<config.imageCount).map { _ in
MLShapedArray<Float32>(
converting: random.normalShapedArray(sampleShape, mean: 0.0, stdev: Double(stdev)))
}
if let image = config.startingImage, config.mode == .imageToImage {
guard let encoder else {
throw PipelineError.startingImageProvidedWithoutEncoder
}
let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
}
return samples
}
public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
let images = try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
if reduceMemory {
decoder.unloadResources()
}
// If safety is disabled return what was decoded
if config.disableSafety {
return images
}
// If there is no safety checker return what was decoded
guard let safetyChecker = safetyChecker else {
return images
}
// Otherwise change images which are not safe to nil
let safeImages = try images.map { image in
try safetyChecker.isSafe(image) ? image : nil
}
if reduceMemory {
safetyChecker.unloadResources()
}
return safeImages
}
}
/// Sampling progress details
@available(iOS 16.2, macOS 13.1, *)
public struct PipelineProgress {
public let pipeline: StableDiffusionPipelineProtocol
public let prompt: String
public let step: Int
public let stepCount: Int
public let currentLatentSamples: [MLShapedArray<Float32>]
public let configuration: PipelineConfiguration
public var isSafetyEnabled: Bool {
pipeline.canSafetyCheck && !configuration.disableSafety
}
public var currentImages: [CGImage?] {
try! pipeline.decodeToImages(currentLatentSamples, configuration: configuration)
}
}
@available(iOS 16.2, macOS 13.1, *)
public extension StableDiffusionPipeline {
/// Sampling progress details
typealias Progress = PipelineProgress
}
// Helper functions
@available(iOS 16.2, macOS 13.1, *)
extension StableDiffusionPipelineProtocol {
internal func randomSource(from rng: StableDiffusionRNG, seed: UInt32) -> RandomSource {
switch rng {
case .numpyRNG:
return NumPyRandomSource(seed: seed)
case .torchRNG:
return TorchRandomSource(seed: seed)
case .nvidiaRNG:
return NvRandomSource(seed: seed)
}
}
func toHiddenStates(_ embedding: MLShapedArray<Float32>) -> MLShapedArray<Float32> {
// Unoptimized manual transpose [0, 2, None, 1]
// e.g. From [2, 77, 768] to [2, 768, 1, 77]
let fromShape = embedding.shape
let stateShape = [fromShape[0],fromShape[2], 1, fromShape[1]]
var states = MLShapedArray<Float32>(repeating: 0.0, shape: stateShape)
for i0 in 0..<fromShape[0] {
for i1 in 0..<fromShape[1] {
for i2 in 0..<fromShape[2] {
states[scalarAt:i0,i2,0,i1] = embedding[scalarAt:i0, i1, i2]
}
}
}
return states
}
func performGuidance(_ noise: [MLShapedArray<Float32>], _ guidanceScale: Float) -> [MLShapedArray<Float32>] {
noise.map { performGuidance($0, guidanceScale) }
}
func performGuidance(_ noise: MLShapedArray<Float32>, _ guidanceScale: Float) -> MLShapedArray<Float32> {
var shape = noise.shape
shape[0] = 1
return MLShapedArray<Float>(unsafeUninitializedShape: shape) { result, _ in
noise.withUnsafeShapedBufferPointer { scalars, _, strides in
for i in 0 ..< result.count {
// unconditioned + guidance*(text - unconditioned)
result.initializeElement(
at: i,
to: scalars[i] + guidanceScale * (scalars[strides[0] + i] - scalars[i])
)
}
}
}
}
}
swift/StableDiffusion/pipeline/StableDiffusionXL+Resources.swift
+119
View File
@@ -0,0 +1,119 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
import NaturalLanguage
@available(iOS 17.0, macOS 14.0, *)
public extension StableDiffusionXLPipeline {
struct ResourceURLs {
public let textEncoderURL: URL
public let textEncoder2URL: URL
public let unetURL: URL
public let unetChunk1URL: URL
public let unetChunk2URL: URL
public let unetRefinerURL: URL
public let unetRefinerChunk1URL: URL
public let unetRefinerChunk2URL: URL
public let decoderURL: URL
public let encoderURL: URL
public let vocabURL: URL
public let mergesURL: URL
public init(resourcesAt baseURL: URL) {
textEncoderURL = baseURL.appending(path: "TextEncoder.mlmodelc")
textEncoder2URL = baseURL.appending(path: "TextEncoder2.mlmodelc")
unetURL = baseURL.appending(path: "Unet.mlmodelc")
unetChunk1URL = baseURL.appending(path: "UnetChunk1.mlmodelc")
unetChunk2URL = baseURL.appending(path: "UnetChunk2.mlmodelc")
unetRefinerURL = baseURL.appending(path: "UnetRefiner.mlmodelc")
unetRefinerChunk1URL = baseURL.appending(path: "UnetRefinerChunk1.mlmodelc")
unetRefinerChunk2URL = baseURL.appending(path: "UnetRefinerChunk2.mlmodelc")
decoderURL = baseURL.appending(path: "VAEDecoder.mlmodelc")
encoderURL = baseURL.appending(path: "VAEEncoder.mlmodelc")
vocabURL = baseURL.appending(path: "vocab.json")
mergesURL = baseURL.appending(path: "merges.txt")
}
}
/// Create stable diffusion pipeline using model resources at a
/// specified URL
///
/// - Parameters:
/// - baseURL: URL pointing to directory holding all model and tokenization resources
/// - configuration: The configuration to load model resources with
/// - reduceMemory: Setup pipeline in reduced memory mode
/// - Returns:
/// Pipeline ready for image generation if all necessary resources loaded
init(
resourcesAt baseURL: URL,
configuration config: MLModelConfiguration = .init(),
reduceMemory: Bool = false
) throws {
/// Expect URL of each resource
let urls = ResourceURLs(resourcesAt: baseURL)
let tokenizer = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL)
let textEncoder: TextEncoderXL?
if FileManager.default.fileExists(atPath: urls.textEncoderURL.path) {
textEncoder = TextEncoderXL(tokenizer: tokenizer, modelAt: urls.textEncoderURL, configuration: config)
} else {
textEncoder = nil
}
// padToken is different in the second XL text encoder
let tokenizer2 = try BPETokenizer(mergesAt: urls.mergesURL, vocabularyAt: urls.vocabURL, padToken: "!")
let textEncoder2 = TextEncoderXL(tokenizer: tokenizer2, modelAt: urls.textEncoder2URL, configuration: config)
// Unet model
let unet: Unet
if FileManager.default.fileExists(atPath: urls.unetChunk1URL.path) &&
FileManager.default.fileExists(atPath: urls.unetChunk2URL.path) {
unet = Unet(chunksAt: [urls.unetChunk1URL, urls.unetChunk2URL],
configuration: config)
} else {
unet = Unet(modelAt: urls.unetURL, configuration: config)
}
// Refiner Unet model
let unetRefiner: Unet?
if FileManager.default.fileExists(atPath: urls.unetRefinerChunk1URL.path) &&
FileManager.default.fileExists(atPath: urls.unetRefinerChunk2URL.path) {
unetRefiner = Unet(chunksAt: [urls.unetRefinerChunk1URL, urls.unetRefinerChunk2URL],
configuration: config)
} else if FileManager.default.fileExists(atPath: urls.unetRefinerURL.path) {
unetRefiner = Unet(modelAt: urls.unetRefinerURL, configuration: config)
} else {
unetRefiner = nil
}
// Image Decoder
// FIXME: Hardcoding to .cpuAndGPU since ANE doesn't support FLOAT32
let vaeConfig = config.copy() as! MLModelConfiguration
vaeConfig.computeUnits = .cpuAndGPU
let decoder = Decoder(modelAt: urls.decoderURL, configuration: vaeConfig)
// Optional Image Encoder
let encoder: Encoder?
if FileManager.default.fileExists(atPath: urls.encoderURL.path) {
encoder = Encoder(modelAt: urls.encoderURL, configuration: vaeConfig)
} else {
encoder = nil
}
// Construct pipeline
self.init(
textEncoder: textEncoder,
textEncoder2: textEncoder2,
unet: unet,
unetRefiner: unetRefiner,
decoder: decoder,
encoder: encoder,
reduceMemory: reduceMemory
)
}
}
swift/StableDiffusion/pipeline/StableDiffusionXLPipeline.swift
+400
View File
@@ -0,0 +1,400 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
import Accelerate
import CoreGraphics
import CoreML
import Foundation
import NaturalLanguage
/// A pipeline used to generate image samples from text input using stable diffusion XL
///
/// This implementation matches:
/// [Hugging Face Diffusers XL Pipeline](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py)
@available(iOS 17.0, macOS 14.0, *)
public struct StableDiffusionXLPipeline: StableDiffusionPipelineProtocol {
public typealias Configuration = PipelineConfiguration
public typealias Progress = PipelineProgress
/// Model to generate embeddings for tokenized input text
var textEncoder: TextEncoderXLModel?
var textEncoder2: TextEncoderXLModel
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
var unet: Unet
/// Model used to refine the image, if present
var unetRefiner: Unet?
/// Model used to generate final image from latent diffusion process
var decoder: Decoder
/// Model used to latent space for image2image, and soon, in-painting
var encoder: Encoder?
/// Option to reduce memory during image generation
///
/// If true, the pipeline will lazily load TextEncoder, Unet, Decoder, and SafetyChecker
/// when needed and aggressively unload their resources after
///
/// This will increase latency in favor of reducing memory
var reduceMemory: Bool = false
/// Creates a pipeline using the specified models and tokenizer
///
/// - Parameters:
/// - textEncoder: Model for encoding tokenized text
/// - textEncoder2: Second text encoding model
/// - unet: Model for noise prediction on latent samples
/// - decoder: Model for decoding latent sample to image
/// - reduceMemory: Option to enable reduced memory mode
/// - Returns: Pipeline ready for image generation
public init(
textEncoder: TextEncoderXLModel?,
textEncoder2: TextEncoderXLModel,
unet: Unet,
unetRefiner: Unet?,
decoder: Decoder,
encoder: Encoder?,
reduceMemory: Bool = false
) {
self.textEncoder = textEncoder
self.textEncoder2 = textEncoder2
self.unet = unet
self.unetRefiner = unetRefiner
self.decoder = decoder
self.encoder = encoder
self.reduceMemory = reduceMemory
}
/// Load required resources for this pipeline
///
/// If reducedMemory is true this will instead call prewarmResources instead
/// and let the pipeline lazily load resources as needed
public func loadResources() throws {
if reduceMemory {
try prewarmResources()
} else {
try textEncoder2.loadResources()
try unet.loadResources()
try decoder.loadResources()
do {
try textEncoder?.loadResources()
} catch {
print("Error loading resources for textEncoder: \(error)")
}
// Only prewarm refiner unet on load so it's unloaded until needed
do {
try unetRefiner?.prewarmResources()
} catch {
print("Error loading resources for unetRefiner: \(error)")
}
do {
try encoder?.loadResources()
} catch {
print("Error loading resources for vae encoder: \(error)")
}
}
}
/// Unload the underlying resources to free up memory
public func unloadResources() {
textEncoder?.unloadResources()
textEncoder2.unloadResources()
unet.unloadResources()
unetRefiner?.unloadResources()
decoder.unloadResources()
encoder?.unloadResources()
}
/// Prewarm resources one at a time
public func prewarmResources() throws {
try textEncoder2.prewarmResources()
try unet.prewarmResources()
try decoder.prewarmResources()
do {
try textEncoder?.prewarmResources()
} catch {
print("Error prewarming resources for textEncoder: \(error)")
}
do {
try unetRefiner?.prewarmResources()
} catch {
print("Error prewarming resources for unetRefiner: \(error)")
}
do {
try encoder?.prewarmResources()
} catch {
print("Error prewarming resources for vae encoder: \(error)")
}
}
/// Image generation using stable diffusion
/// - Parameters:
/// - configuration: Image generation configuration
/// - progressHandler: Callback to perform after each step, stops on receiving false response
/// - Returns: An array of `imageCount` optional images.
/// The images will be nil if safety checks were performed and found the result to be un-safe
public func generateImages(
configuration config: Configuration,
progressHandler: (Progress) -> Bool = { _ in true }
) throws -> [CGImage?] {
// Determine input type of Unet
// SDXL Refiner has a latentTimeIdShape of [2, 5]
// SDXL Base has either [12] or [2, 6]
let isRefiner = unet.latentTimeIdShape.last == 5
// Setup geometry conditioning for base/refiner inputs
var baseInput: ModelInputs?
var refinerInput: ModelInputs?
// Check if the first textEncoder is available, which is required for base models
if textEncoder != nil {
baseInput = try generateConditioning(using: config, forRefiner: isRefiner)
}
// Check if the refiner unet exists, or if the current unet is a refiner
if unetRefiner != nil || isRefiner {
refinerInput = try generateConditioning(using: config, forRefiner: true)
}
if reduceMemory {
textEncoder?.unloadResources()
textEncoder2.unloadResources()
}
/// Setup schedulers
let scheduler: [Scheduler] = (0..<config.imageCount).map { _ in
switch config.schedulerType {
case .pndmScheduler: return PNDMScheduler(stepCount: config.stepCount)
case .dpmSolverMultistepScheduler: return DPMSolverMultistepScheduler(stepCount: config.stepCount, timeStepSpacing: config.schedulerTimestepSpacing)
case .discreteFlowScheduler: return DiscreteFlowScheduler(stepCount: config.stepCount, timeStepShift: config.schedulerTimestepShift)
}
}
// Generate random latent samples from specified seed
var latents: [MLShapedArray<Float32>] = try generateLatentSamples(configuration: config, scheduler: scheduler[0])
// Store denoised latents from scheduler to pass into decoder
var denoisedLatents: [MLShapedArray<Float32>] = latents.map { MLShapedArray(converting: $0) }
if reduceMemory {
encoder?.unloadResources()
}
let timestepStrength: Float? = config.mode == .imageToImage ? config.strength : nil
// Store current model
var unetModel = unet
var currentInput = baseInput ?? refinerInput
var unetHiddenStates = currentInput?.hiddenStates
var unetPooledStates = currentInput?.pooledStates
var unetGeometryConditioning = currentInput?.geometryConditioning
let timeSteps: [Int] = scheduler[0].calculateTimesteps(strength: timestepStrength)
// Calculate which step to swap to refiner
let refinerStartStep = Int(Float(timeSteps.count) * config.refinerStart)
// De-noising loop
for (step,t) in timeSteps.enumerated() {
// Expand the latents for classifier-free guidance
// and input to the Unet noise prediction model
let latentUnetInput = latents.map {
MLShapedArray<Float32>(concatenating: [$0, $0], alongAxis: 0)
}
// Switch to refiner if specified
if let refiner = unetRefiner, step == refinerStartStep {
unet.unloadResources()
unetModel = refiner
currentInput = refinerInput
unetHiddenStates = currentInput?.hiddenStates
unetPooledStates = currentInput?.pooledStates
unetGeometryConditioning = currentInput?.geometryConditioning
}
guard let hiddenStates = unetHiddenStates,
let pooledStates = unetPooledStates,
let geometryConditioning = unetGeometryConditioning else {
throw PipelineError.missingUnetInputs
}
// Predict noise residuals from latent samples
// and current time step conditioned on hidden states
var noise = try unetModel.predictNoise(
latents: latentUnetInput,
timeStep: t,
hiddenStates: hiddenStates,
pooledStates: pooledStates,
geometryConditioning: geometryConditioning
)
noise = performGuidance(noise, config.guidanceScale)
// Have the scheduler compute the previous (t-1) latent
// sample given the predicted noise and current sample
for i in 0..<config.imageCount {
latents[i] = scheduler[i].step(
output: noise[i],
timeStep: t,
sample: latents[i]
)
denoisedLatents[i] = scheduler[i].modelOutputs.last ?? latents[i]
}
let currentLatentSamples = config.useDenoisedIntermediates ? denoisedLatents : latents
// Report progress
let progress = Progress(
pipeline: self,
prompt: config.prompt,
step: step,
stepCount: timeSteps.count,
currentLatentSamples: currentLatentSamples,
configuration: config
)
if !progressHandler(progress) {
// Stop if requested by handler
return []
}
}
// Unload resources
if reduceMemory {
unet.unloadResources()
}
unetRefiner?.unloadResources()
// Decode the latent samples to images
return try decodeToImages(denoisedLatents, configuration: config)
}
func encodePrompt(_ prompt: String, forRefiner: Bool = false) throws -> (MLShapedArray<Float32>, MLShapedArray<Float32>) {
var embeds = MLShapedArray<Float32>()
var pooled = MLShapedArray<Float32>()
if forRefiner {
let (embeds2, pooledValue) = try textEncoder2.encode(prompt)
// Refiner only takes textEncoder2 embeddings
// [1, 77, 1280]
embeds = embeds2
pooled = pooledValue
} else {
guard let encoder = textEncoder else {
throw PipelineError.startingText2ImgWithoutTextEncoder
}
let (embeds1, _) = try encoder.encode(prompt)
let (embeds2, pooledValue) = try textEncoder2.encode(prompt)
// Base needs concatenated embeddings
// [1, 77, 768], [1, 77, 1280] -> [1, 77, 2048]
embeds = MLShapedArray<Float32>(
concatenating: [embeds1, embeds2],
alongAxis: 2
)
pooled = pooledValue
}
return (embeds, pooled)
}
func generateConditioning(using config: Configuration, forRefiner: Bool = false) throws -> ModelInputs {
// Encode the input prompt and negative prompt
let (promptEmbedding, pooled) = try encodePrompt(config.prompt, forRefiner: forRefiner)
let (negativePromptEmbedding, negativePooled) = try encodePrompt(config.negativePrompt, forRefiner: forRefiner)
// Convert to Unet hidden state representation
// Concatenate the prompt and negative prompt embeddings
let hiddenStates = toHiddenStates(MLShapedArray(concatenating: [negativePromptEmbedding, promptEmbedding], alongAxis: 0))
let pooledStates = MLShapedArray(concatenating: [negativePooled, pooled], alongAxis: 0)
// Inline helper functions for geometry creation
func refinerGeometry() -> MLShapedArray<Float32> {
let negativeGeometry = MLShapedArray<Float32>(
scalars: [
config.originalSize, config.originalSize,
config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
config.negativeAestheticScore
],
shape: [1, 5]
)
let positiveGeometry = MLShapedArray<Float32>(
scalars: [
config.originalSize, config.originalSize,
config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
config.aestheticScore
],
shape: [1, 5]
)
return MLShapedArray<Float32>(concatenating: [negativeGeometry, positiveGeometry], alongAxis: 0)
}
func baseGeometry() -> MLShapedArray<Float32> {
let geometry = MLShapedArray<Float32>(
scalars: [
config.originalSize, config.originalSize,
config.cropsCoordsTopLeft, config.cropsCoordsTopLeft,
config.targetSize, config.targetSize
],
// TODO: This checks if the time_ids input is looking for [12] or [2, 6]
// Remove once model input shapes are ubiquitous
shape: unet.latentTimeIdShape.count > 1 ? [1, 6] : [6]
)
return MLShapedArray<Float32>(concatenating: [geometry, geometry], alongAxis: 0)
}
let geometry = forRefiner ? refinerGeometry() : baseGeometry()
return ModelInputs(hiddenStates: hiddenStates, pooledStates: pooledStates, geometryConditioning: geometry)
}
func generateLatentSamples(configuration config: Configuration, scheduler: Scheduler) throws -> [MLShapedArray<Float32>] {
var sampleShape = unet.latentSampleShape
sampleShape[0] = 1
let stdev = scheduler.initNoiseSigma
var random = randomSource(from: config.rngType, seed: config.seed)
let samples = (0..<config.imageCount).map { _ in
MLShapedArray<Float32>(
converting: random.normalShapedArray(sampleShape, mean: 0.0, stdev: Double(stdev)))
}
if let image = config.startingImage, config.mode == .imageToImage {
guard let encoder else {
throw PipelineError.startingImageProvidedWithoutEncoder
}
let latent = try encoder.encode(image, scaleFactor: config.encoderScaleFactor, random: &random)
return scheduler.addNoise(originalSample: latent, noise: samples, strength: config.strength)
}
return samples
}
public func decodeToImages(_ latents: [MLShapedArray<Float32>], configuration config: Configuration) throws -> [CGImage?] {
defer {
if reduceMemory {
decoder.unloadResources()
}
}
return try decoder.decode(latents, scaleFactor: config.decoderScaleFactor)
}
struct ModelInputs {
var hiddenStates: MLShapedArray<Float32>
var pooledStates: MLShapedArray<Float32>
var geometryConditioning: MLShapedArray<Float32>
}
}
swift/StableDiffusion/pipeline/TextEncoder.swift
+102
View File
@@ -0,0 +1,102 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
@available(iOS 16.2, macOS 13.1, *)
public protocol TextEncoderModel: ResourceManaging {
func encode(_ text: String) throws -> MLShapedArray<Float32>
}
/// A model for encoding text
@available(iOS 16.2, macOS 13.1, *)
public struct TextEncoder: TextEncoderModel {
/// Text tokenizer
var tokenizer: BPETokenizer
/// Embedding model
var model: ManagedMLModel
/// Creates text encoder which embeds a tokenized string
///
/// - Parameters:
/// - tokenizer: Tokenizer for input text
/// - url: Location of compiled text encoding Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: A text encoder that will lazily load its required resources when needed or requested
public init(tokenizer: BPETokenizer,
modelAt url: URL,
configuration: MLModelConfiguration) {
self.tokenizer = tokenizer
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
/// Encode input text/string
///
/// - Parameters:
/// - text: Input text to be tokenized and then embedded
/// - Returns: Embedding representing the input text
public func encode(_ text: String) throws -> MLShapedArray<Float32> {
// Get models expected input length
let inputLength = inputShape.last!
// Tokenize, padding to the expected length
var (tokens, ids) = tokenizer.tokenize(input: text, minCount: inputLength)
// Truncate if necessary
if ids.count > inputLength {
tokens = tokens.dropLast(tokens.count - inputLength)
ids = ids.dropLast(ids.count - inputLength)
let truncated = tokenizer.decode(tokens: tokens)
print("Needed to truncate input '\(text)' to '\(truncated)'")
}
// Use the model to generate the embedding
return try encode(ids: ids)
}
/// Prediction queue
let queue = DispatchQueue(label: "textencoder.predict")
func encode(ids: [Int]) throws -> MLShapedArray<Float32> {
let inputName = inputDescription.name
let inputShape = inputShape
let floatIds = ids.map { Float32($0) }
let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
let inputFeatures = try! MLDictionaryFeatureProvider(
dictionary: [inputName: MLMultiArray(inputArray)])
let result = try model.perform { model in
try model.prediction(from: inputFeatures)
}
let embeddingFeature = result.featureValue(for: "last_hidden_state")
return MLShapedArray<Float32>(converting: embeddingFeature!.multiArrayValue!)
}
var inputDescription: MLFeatureDescription {
try! model.perform { model in
model.modelDescription.inputDescriptionsByName.first!.value
}
}
var inputShape: [Int] {
inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
}
swift/StableDiffusion/pipeline/TextEncoderT5.swift
+124
View File
@@ -0,0 +1,124 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
import Tokenizers
@available(iOS 17.0, macOS 14.0, *)
public protocol TextEncoderT5Model: ResourceManaging {
func encode(_ text: String) throws -> TextEncoderT5Output
}
@available(iOS 17.0, macOS 14.0, *)
public struct TextEncoderT5Output {
public let encoderHiddenStates: MLShapedArray<Float32>
}
/// A model for encoding text, suitable for SD3
@available(iOS 17.0, macOS 14.0, *)
public struct TextEncoderT5: TextEncoderT5Model {
/// Text tokenizer
var tokenizer: Tokenizer
/// Embedding model
var model: ManagedMLModel
/// Creates text encoder which embeds a tokenized string
///
/// - Parameters:
/// - tokenizer: Tokenizer for input text
/// - url: Location of compiled text encoding Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: A text encoder that will lazily load its required resources when needed or requested
public init(tokenizer: Tokenizer,
modelAt url: URL,
configuration: MLModelConfiguration) {
self.tokenizer = tokenizer
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
/// Encode input text/string
///
/// - Parameters:
/// - text: Input text to be tokenized and then embedded
/// - Returns: Embedding representing the input text
public func encode(_ text: String) throws -> TextEncoderT5Output {
// Get models expected input length
let inputLength = inputShape.last!
// Tokenize, padding to the expected length
var tokens = tokenizer.tokenize(text: text)
var ids = tokens.map { tokenizer.convertTokenToId($0) ?? 0 }
// Truncate if necessary
if ids.count > inputLength {
tokens = tokens.dropLast(tokens.count - inputLength)
ids = ids.dropLast(ids.count - inputLength)
print("Needed to truncate input for TextEncoderT5")
}
// Use the model to generate the embedding
let encodedText = try encode(ids: ids)
return encodedText
}
func encode(ids: [Int]) throws -> TextEncoderT5Output {
let inputName = "input_ids"
let inputShape = inputShape
let inputLength = inputShape[1]
let bosToken = tokenizer.bosTokenId ?? 0
let eosToken = tokenizer.eosTokenId ?? 1
let padToken = bosToken
let maskToken = eosToken
// Truncate and pad input to the expected length
let truncatedIds = ids.prefix(inputLength - 1) + [eosToken]
let inputIds = truncatedIds + Array(repeating: padToken, count: inputLength - truncatedIds.count)
let attentionMaskName = "attention_mask"
var attentionMask: [Int] = inputIds.map { token in
token == padToken ? maskToken : padToken
}
attentionMask[0] = bosToken
let floatIds = inputIds.map { Float32($0) }
let floatMask = attentionMask.map { Float32($0) }
let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
let maskArray = MLShapedArray<Float32>(scalars: floatMask, shape: inputShape)
let inputFeatures = try! MLDictionaryFeatureProvider(
dictionary: [inputName: MLMultiArray(inputArray),
attentionMaskName: MLMultiArray(maskArray)])
let result = try model.perform { model in
try model.prediction(from: inputFeatures)
}
let embeddingFeature = result.featureValue(for: "encoder_hidden_states")
return TextEncoderT5Output(encoderHiddenStates: MLShapedArray<Float32>(converting: embeddingFeature!.multiArrayValue!))
}
var inputDescription: MLFeatureDescription {
try! model.perform { model in
model.modelDescription.inputDescriptionsByName.first!.value
}
}
var inputShape: [Int] {
inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
}
swift/StableDiffusion/pipeline/TextEncoderXL.swift
+99
View File
@@ -0,0 +1,99 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2023 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
@available(iOS 17.0, macOS 14.0, *)
public protocol TextEncoderXLModel: ResourceManaging {
typealias TextEncoderXLOutput = (hiddenEmbeddings: MLShapedArray<Float32>, pooledOutputs: MLShapedArray<Float32>)
func encode(_ text: String) throws -> TextEncoderXLOutput
}
/// A model for encoding text, suitable for SDXL
@available(iOS 17.0, macOS 14.0, *)
public struct TextEncoderXL: TextEncoderXLModel {
/// Text tokenizer
var tokenizer: BPETokenizer
/// Embedding model
var model: ManagedMLModel
/// Creates text encoder which embeds a tokenized string
///
/// - Parameters:
/// - tokenizer: Tokenizer for input text
/// - url: Location of compiled text encoding Core ML model
/// - configuration: configuration to be used when the model is loaded
/// - Returns: A text encoder that will lazily load its required resources when needed or requested
public init(tokenizer: BPETokenizer,
modelAt url: URL,
configuration: MLModelConfiguration) {
self.tokenizer = tokenizer
self.model = ManagedMLModel(modelAt: url, configuration: configuration)
}
/// Ensure the model has been loaded into memory
public func loadResources() throws {
try model.loadResources()
}
/// Unload the underlying model to free up memory
public func unloadResources() {
model.unloadResources()
}
/// Encode input text/string
///
/// - Parameters:
/// - text: Input text to be tokenized and then embedded
/// - Returns: Embedding representing the input text
public func encode(_ text: String) throws -> TextEncoderXLOutput {
// Get models expected input length
let inputLength = inputShape.last!
// Tokenize, padding to the expected length
var (tokens, ids) = tokenizer.tokenize(input: text, minCount: inputLength)
// Truncate if necessary
if ids.count > inputLength {
tokens = tokens.dropLast(tokens.count - inputLength)
ids = ids.dropLast(ids.count - inputLength)
let truncated = tokenizer.decode(tokens: tokens)
print("Needed to truncate input '\(text)' to '\(truncated)'")
}
// Use the model to generate the embedding
return try encode(ids: ids)
}
func encode(ids: [Int]) throws -> TextEncoderXLOutput {
let inputName = inputDescription.name
let inputShape = inputShape
let floatIds = ids.map { Float32($0) }
let inputArray = MLShapedArray<Float32>(scalars: floatIds, shape: inputShape)
let inputFeatures = try! MLDictionaryFeatureProvider(
dictionary: [inputName: MLMultiArray(inputArray)])
let result = try model.perform { model in
try model.prediction(from: inputFeatures)
}
let embeddingFeature = result.featureValue(for: "hidden_embeds")
let pooledFeature = result.featureValue(for: "pooled_outputs")
return (MLShapedArray<Float32>(converting: embeddingFeature!.multiArrayValue!), MLShapedArray<Float32>(converting: pooledFeature!.multiArrayValue!))
}
var inputDescription: MLFeatureDescription {
try! model.perform { model in
model.modelDescription.inputDescriptionsByName.first!.value
}
}
var inputShape: [Int] {
inputDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
}
swift/StableDiffusion/pipeline/TorchRandomSource.swift
+157
View File
@@ -0,0 +1,157 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// A random source consistent with PyTorch
///
/// This implementation matches:
/// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/DistributionsHelper.h
/// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/DistributionTemplates.h
/// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/DistributionKernels.cpp
/// https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/TransformationHelper.h
///
@available(iOS 16.2, macOS 13.1, *)
struct TorchRandomSource: RandomNumberGenerator, RandomSource {
struct State {
var key = [UInt32](repeating: 0, count: 624)
var pos: Int = 0
var nextGauss: Double? = nil
}
var state: State
/// Initialize with a random seed
///
/// - Parameters
/// - seed: Seed for underlying Mersenne Twister 19937 generator
/// - Returns random source
init(seed: UInt32) {
state = .init()
var s = seed & 0xffff_ffff
for i in 0..<state.key.count {
state.key[i] = s
s = UInt32((UInt64(1_812_433_253) * UInt64(s ^ (s >> 30)) + UInt64(i) + 1) & 0xffff_ffff)
}
state.pos = state.key.count
state.nextGauss = nil
}
/// Generate next UInt32 using fast 32bit Mersenne Twister
mutating func nextUInt32() -> UInt32 {
let n = 624
let m = 397
let matrixA: UInt64 = 0x9908_b0df
let upperMask: UInt32 = 0x8000_0000
let lowerMask: UInt32 = 0x7fff_ffff
var y: UInt32
if state.pos == state.key.count {
for i in 0..<(n - m) {
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
state.key[i] = state.key[i + m] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
}
for i in (n - m)..<(n - 1) {
y = (state.key[i] & upperMask) | (state.key[i + 1] & lowerMask)
state.key[i] = state.key[i + (m - n)] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
}
y = (state.key[n - 1] & upperMask) | (state.key[0] & lowerMask)
state.key[n - 1] = state.key[m - 1] ^ (y >> 1) ^ UInt32((UInt64(~(y & 1)) + 1) & matrixA)
state.pos = 0
}
y = state.key[state.pos]
state.pos += 1
y ^= (y >> 11)
y ^= (y << 7) & 0x9d2c_5680
y ^= (y << 15) & 0xefc6_0000
y ^= (y >> 18)
return y
}
mutating func next() -> UInt64 {
let high = nextUInt32()
let low = nextUInt32()
return (UInt64(high) << 32) | UInt64(low)
}
/// Generate next random double value
mutating func nextDouble() -> Double {
let a = next()
return Double(a & 9_007_199_254_740_991) * (1.0 / 9007199254740992.0)
}
/// Generate next random float value
mutating func nextFloat() -> Float {
let a = nextUInt32()
return Float(a & 16_777_215) * (1.0 / 16777216.0)
}
/// Generate next random value from a standard normal
mutating func nextGauss() -> Double {
if let nextGauss = state.nextGauss {
state.nextGauss = nil
return nextGauss
}
// Box-Muller transform
let u1: Double = nextDouble()
let u2: Double = 1 - nextDouble()
let radius = sqrt(-2.0 * log(u2))
let theta = 2.0 * .pi * u1
state.nextGauss = radius * sin(theta)
return radius * cos(theta)
}
/// Generates a random value from a normal distribution with given mean and standard deviation.
mutating func nextNormal(mean: Double = 0.0, stdev: Double = 1.0) -> Double {
nextGauss() * stdev + mean
}
/// Generates an array of random values from a normal distribution with given mean and standard deviation.
/// This simulates torch.randn([1, 4, 64, 64], dtype=torch.float), note that for dtype=torch.double, it
/// will be slightly different.
mutating func normalArray(count: Int, mean: Double = 0.0, stdev: Double = 1.0) -> [Double] {
// If it is smaller than 16 elements, Torch generates from Box-Muller transform directly.
// Note that even if this is used to generate Float, it will use Double underneath.
guard count >= 16 else {
return (0..<count).map { _ in nextNormal(mean: mean, stdev: stdev) }
}
// Otherwise, Torch first fill a uniform distribution array, then do Box-Muller
// transformation over this array.
var data = (0..<count).map { _ in Double(nextFloat()) }
for i in stride(from: 0, to: count - 15, by: 16) {
for j in 0..<8 {
let u1 = 1 - data[i + j]
let u2 = data[i + j + 8]
let radius = sqrt(-2.0 * log(u1))
let theta = 2.0 * .pi * u2
data[i + j] = radius * cos(theta) * stdev + mean
data[i + j + 8] = radius * sin(theta) * stdev + mean
}
}
if count % 16 != 0 {
for i in (count - 16)..<count {
data[i] = nextDouble()
}
let i = count - 16
for j in 0..<8 {
let u1 = 1 - data[i + j]
let u2 = data[i + j + 8]
let radius = sqrt(-2.0 * log(u1))
let theta = 2.0 * .pi * u2
data[i + j] = radius * cos(theta) * stdev + mean
data[i + j + 8] = radius * sin(theta) * stdev + mean
}
}
return data
}
/// Generate a shaped array with scalars from a normal distribution with given mean and standard deviation.
mutating func normalShapedArray(_ shape: [Int], mean: Double = 0.0, stdev: Double = 1.0) -> MLShapedArray<Double> {
let count = shape.reduce(1, *)
return .init(scalars: normalArray(count: count, mean: mean, stdev: stdev), shape: shape)
}
}
swift/StableDiffusion/pipeline/Unet.swift
+204
View File
@@ -0,0 +1,204 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
import CoreML
/// U-Net noise prediction model for stable diffusion
@available(iOS 16.2, macOS 13.1, *)
public struct Unet: ResourceManaging {
/// Model used to predict noise residuals given an input, diffusion time step, and conditional embedding
///
/// It can be in the form of a single model or multiple stages
var models: [ManagedMLModel]
/// Creates a U-Net noise prediction model
///
/// - Parameters:
/// - url: Location of single U-Net compiled Core ML model
/// - configuration: Configuration to be used when the model is loaded
/// - Returns: U-net model that will lazily load its required resources when needed or requested
public init(modelAt url: URL,
configuration: MLModelConfiguration) {
self.models = [ManagedMLModel(modelAt: url, configuration: configuration)]
}
/// Creates a U-Net noise prediction model
///
/// - Parameters:
/// - urls: Location of chunked U-Net via urls to each compiled chunk
/// - configuration: Configuration to be used when the model is loaded
/// - Returns: U-net model that will lazily load its required resources when needed or requested
public init(chunksAt urls: [URL],
configuration: MLModelConfiguration) {
self.models = urls.map { ManagedMLModel(modelAt: $0, configuration: configuration) }
}
/// Load resources.
public func loadResources() throws {
for model in models {
try model.loadResources()
}
}
/// Unload the underlying model to free up memory
public func unloadResources() {
for model in models {
model.unloadResources()
}
}
/// Pre-warm resources
public func prewarmResources() throws {
// Override default to pre-warm each model
for model in models {
try model.loadResources()
model.unloadResources()
}
}
var latentSampleDescription: MLFeatureDescription {
try! models.first!.perform { model in
model.modelDescription.inputDescriptionsByName["sample"]!
}
}
/// The expected shape of the models latent sample input
public var latentSampleShape: [Int] {
latentSampleDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
var latentTimeIdDescription: MLFeatureDescription {
try! models.first!.perform { model in
model.modelDescription.inputDescriptionsByName["time_ids"]!
}
}
/// The expected shape of the geometry conditioning
public var latentTimeIdShape: [Int] {
latentTimeIdDescription.multiArrayConstraint!.shape.map { $0.intValue }
}
/// Batch prediction noise from latent samples
///
/// - Parameters:
/// - latents: Batch of latent samples in an array
/// - timeStep: Current diffusion timestep
/// - hiddenStates: Hidden state to condition on
/// - Returns: Array of predicted noise residuals
func predictNoise(
latents: [MLShapedArray<Float32>],
timeStep: Int,
hiddenStates: MLShapedArray<Float32>,
additionalResiduals: [[String: MLShapedArray<Float32>]]? = nil
) throws -> [MLShapedArray<Float32>] {
// Match time step batch dimension to the model / latent samples
let t: MLShapedArray<Float32>
if hiddenStates.shape[0] == 2 {
t = MLShapedArray(scalars: [Float(timeStep), Float(timeStep)], shape: [2])
} else {
t = MLShapedArray(scalars: [Float(timeStep)], shape: [1])
}
// Form batch input to model
let inputs = try latents.enumerated().map {
var dict: [String: Any] = [
"sample" : MLMultiArray($0.element),
"timestep" : MLMultiArray(t),
"encoder_hidden_states": MLMultiArray(hiddenStates)
]
if let residuals = additionalResiduals?[$0.offset] {
for (k, v) in residuals {
dict[k] = MLMultiArray(v)
}
}
return try MLDictionaryFeatureProvider(dictionary: dict)
}
let batch = MLArrayBatchProvider(array: inputs)
// Make predictions
let results = try models.predictions(from: batch)
// Pull out the results in Float32 format
let noise = (0..<results.count).map { i in
let result = results.features(at: i)
let outputName = result.featureNames.first!
let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
// To conform to this func return type make sure we return float32
// Use the fact that the concatenating constructor for MLMultiArray
// can do type conversion:
let fp32Noise = MLMultiArray(
concatenating: [outputNoise],
axis: 0,
dataType: .float32
)
return MLShapedArray<Float32>(fp32Noise)
}
return noise
}
/// Batch prediction noise from latent samples, for Stable Diffusion XL
///
/// - Parameters:
/// - latents: Batch of latent samples in an array
/// - timeStep: Current diffusion timestep
/// - hiddenStates: Hidden state to condition on
/// - pooledStates: Additional text states to condition on
/// - geometryConditioning: Condition on image geometry
/// - Returns: Array of predicted noise residuals
@available(iOS 17.0, macOS 14.0, *)
func predictNoise(
latents: [MLShapedArray<Float32>],
timeStep: Int,
hiddenStates: MLShapedArray<Float32>,
pooledStates: MLShapedArray<Float32>,
geometryConditioning: MLShapedArray<Float32>
) throws -> [MLShapedArray<Float32>] {
// Match time step batch dimension to the model / latent samples
let t = MLShapedArray<Float32>(scalars:[Float(timeStep), Float(timeStep)],shape:[2])
// Form batch input to model
let inputs = try latents.enumerated().map {
let dict: [String: Any] = [
"sample" : MLMultiArray($0.element),
"timestep" : MLMultiArray(t),
"encoder_hidden_states": MLMultiArray(hiddenStates),
"text_embeds": MLMultiArray(pooledStates),
"time_ids": MLMultiArray(geometryConditioning)
]
return try MLDictionaryFeatureProvider(dictionary: dict)
}
let batch = MLArrayBatchProvider(array: inputs)
// Make predictions
let results = try models.predictions(from: batch)
// Pull out the results in Float32 format
let noise = (0..<results.count).map { i in
let result = results.features(at: i)
let outputName = result.featureNames.first!
let outputNoise = result.featureValue(for: outputName)!.multiArrayValue!
// To conform to this func return type make sure we return float32
// Use the fact that the concatenating constructor for MLMultiArray
// can do type conversion:
let fp32Noise = MLMultiArray(
concatenating: [outputNoise],
axis: 0,
dataType: .float32
)
return MLShapedArray<Float32>(fp32Noise)
}
return noise
}
}
swift/StableDiffusion/tokenizer/BPETokenizer+Reading.swift
+58
View File
@@ -0,0 +1,58 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
@available(iOS 16.2, macOS 13.1, *)
extension BPETokenizer {
enum FileReadError: Error {
case invalidMergeFileLine(Int)
}
/// Read vocab.json file at URL into a dictionary mapping a String to its Int token id
static func readVocabulary(url: URL) throws -> [String: Int] {
let content = try Data(contentsOf: url)
return try JSONDecoder().decode([String: Int].self, from: content)
}
/// Read merges.txt file at URL into a dictionary mapping bigrams to the line number/rank/priority
static func readMerges(url: URL) throws -> [TokenPair: Int] {
let data = try Data(contentsOf: url)
var merges = [TokenPair: Int]()
var index = 0
var line = [UInt8]()
for byte in data {
if byte == UInt8(ascii: "\n") {
if let pair = try parseMergesLine(line, index: index) {
merges[pair] = index
}
line.removeAll(keepingCapacity: true)
index += 1
} else {
line.append(byte)
}
}
return merges
}
static func parseMergesLine(_ line: [UInt8], index: Int) throws -> TokenPair? {
if line.isEmpty || line.first == UInt8(ascii: "#") {
return nil
}
let pair = line.split(separator: UInt8(ascii: " "))
if pair.count != 2 {
throw FileReadError.invalidMergeFileLine(index + 1)
}
return TokenPair( String(bytes: pair[0]), String(bytes: pair[1]))
}
}
extension String {
init(bytes: some Collection<UInt8>) {
self.init(unsafeUninitializedCapacity: bytes.count) { pointer in
_ = pointer.initialize(fromContentsOf: bytes)
return bytes.count
}
}
}
swift/StableDiffusion/tokenizer/BPETokenizer.swift
+185
View File
@@ -0,0 +1,185 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import Foundation
/// A tokenizer based on byte pair encoding.
@available(iOS 16.2, macOS 13.1, *)
public struct BPETokenizer {
/// A dictionary that maps pairs of tokens to the rank/order of the merge.
let merges: [TokenPair : Int]
/// A dictionary from of tokens to identifiers.
let vocabulary: [String: Int]
/// The token used for padding
let padToken: String
/// The start token.
let startToken: String = "<|startoftext|>"
/// The end token.
let endToken: String = "<|endoftext|>"
/// The unknown token.
let unknownToken: String = "<|endoftext|>"
var unknownTokenID: Int {
vocabulary[unknownToken, default: 0]
}
/// Creates a tokenizer.
///
/// - Parameters:
/// - merges: A dictionary that maps pairs of tokens to the rank/order of the merge.
/// - vocabulary: A dictionary from of tokens to identifiers.
public init(merges: [TokenPair: Int], vocabulary: [String: Int], padToken: String = "<|endoftext|>") {
self.merges = merges
self.vocabulary = vocabulary
self.padToken = padToken
}
/// Creates a tokenizer by loading merges and vocabulary from URLs.
///
/// - Parameters:
/// - mergesURL: The URL of a text file containing merges.
/// - vocabularyURL: The URL of a JSON file containing the vocabulary.
public init(mergesAt mergesURL: URL, vocabularyAt vocabularyURL: URL, padToken: String = "<|endoftext|>") throws {
self.merges = try Self.readMerges(url: mergesURL)
self.vocabulary = try! Self.readVocabulary(url: vocabularyURL)
self.padToken = padToken
}
/// Tokenizes an input string.
///
/// - Parameters:
/// - input: A string.
/// - minCount: The minimum number of tokens to return.
/// - Returns: An array of tokens and an array of token identifiers.
public func tokenize(input: String, minCount: Int? = nil) -> (tokens: [String], tokenIDs: [Int]) {
var tokens: [String] = []
tokens.append(startToken)
tokens.append(contentsOf: encode(input: input))
tokens.append(endToken)
// Pad if there was a min length specified
if let minLen = minCount, minLen > tokens.count {
tokens.append(contentsOf: repeatElement(padToken, count: minLen - tokens.count))
}
let ids = tokens.map({ vocabulary[$0, default: unknownTokenID] })
return (tokens: tokens, tokenIDs: ids)
}
/// Returns the token identifier for a token.
public func tokenID(for token: String) -> Int? {
vocabulary[token]
}
/// Returns the token for a token identifier.
public func token(id: Int) -> String? {
vocabulary.first(where: { $0.value == id })?.key
}
/// Decodes a sequence of tokens into a fully formed string
public func decode(tokens: [String]) -> String {
String(tokens.joined())
.replacingOccurrences(of: "</w>", with: " ")
.replacingOccurrences(of: startToken, with: "")
.replacingOccurrences(of: endToken, with: "")
}
/// Encode an input string to a sequence of tokens
func encode(input: String) -> [String] {
let normalized = input.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
let words = normalized.split(separator: " ")
return words.flatMap({ encode(word: $0) })
}
/// Encode a single word into a sequence of tokens
func encode(word: Substring) -> [String] {
var tokens = word.map { String($0) }
if let last = tokens.indices.last {
tokens[last] = tokens[last] + "</w>"
}
while true {
let pairs = pairs(for: tokens)
let canMerge = pairs.filter { merges[$0] != nil }
if canMerge.isEmpty {
break
}
// If multiple merges are found, use the one with the lowest rank
let shouldMerge = canMerge.min { merges[$0]! < merges[$1]! }!
tokens = update(tokens, merging: shouldMerge)
}
return tokens
}
/// Get the set of adjacent pairs / bigrams from a sequence of tokens
func pairs(for tokens: [String]) -> Set<TokenPair> {
guard tokens.count > 1 else {
return Set()
}
var pairs = Set<TokenPair>(minimumCapacity: tokens.count - 1)
var prev = tokens.first!
for current in tokens.dropFirst() {
pairs.insert(TokenPair(prev, current))
prev = current
}
return pairs
}
/// Update the sequence of tokens by greedily merging instance of a specific bigram
func update(_ tokens: [String], merging bigram: TokenPair) -> [String] {
guard tokens.count > 1 else {
return []
}
var newTokens = [String]()
newTokens.reserveCapacity(tokens.count - 1)
var index = 0
while index < tokens.count {
let remainingTokens = tokens[index...]
if let startMatchIndex = remainingTokens.firstIndex(of: bigram.first) {
// Found a possible match, append everything before it
newTokens.append(contentsOf: tokens[index..<startMatchIndex])
if index < tokens.count - 1 && tokens[startMatchIndex + 1] == bigram.second {
// Full match, merge
newTokens.append(bigram.first + bigram.second)
index = startMatchIndex + 2
} else {
// Only matched the first, no merge
newTokens.append(bigram.first)
index = startMatchIndex + 1
}
} else {
// Didn't find any more matches, append the rest unmerged
newTokens.append(contentsOf: remainingTokens)
break
}
}
return newTokens
}
}
@available(iOS 16.2, macOS 13.1, *)
extension BPETokenizer {
/// A hashable tuple of strings
public struct TokenPair: Hashable {
let first: String
let second: String
init(_ first: String, _ second: String) {
self.first = first
self.second = second
}
}
}
swift/StableDiffusion/tokenizer/T5Tokenizer.swift
+21
View File
@@ -0,0 +1,21 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2024 Apple Inc. All Rights Reserved.
import Foundation
import Hub
import Tokenizers
/// Extension to swift-transfomers Hub.swift to load local Config files
public extension Config {
/// Assumes the file is already present at local url.
/// `fileURL` is a complete local file path for the given model
public init(fileURL: URL) throws {
let data = try Data(contentsOf: fileURL)
let parsed = try JSONSerialization.jsonObject(with: data, options: [])
guard var dictionary = parsed as? [String: Any] else { throw Hub.HubClientError.parse }
// Necessary override for loading local tokenizer configs
dictionary["tokenizer_class"] = "T5Tokenizer"
self.init(dictionary)
}
}
swift/StableDiffusionCLI/main.swift
+377
View File
@@ -0,0 +1,377 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import ArgumentParser
import CoreGraphics
import CoreML
import Foundation
import StableDiffusion
import UniformTypeIdentifiers
import Cocoa
import CoreImage
import NaturalLanguage
@available(iOS 16.2, macOS 13.1, *)
struct StableDiffusionSample: ParsableCommand {
static let configuration = CommandConfiguration(
abstract: "Run stable diffusion to generate images guided by a text prompt",
version: "0.1"
)
@Argument(help: "Input string prompt")
var prompt: String
@Option(help: "Input string negative prompt")
var negativePrompt: String = ""
@Option(
help: ArgumentHelp(
"Path to stable diffusion resources.",
discussion: "The resource directory should contain\n" +
" - *compiled* models: {TextEncoder,Unet,VAEDecoder}.mlmodelc\n" +
" - tokenizer info: vocab.json, merges.txt",
valueName: "directory-path"
)
)
var resourcePath: String = "./"
@Flag(name: .customLong("xl"), help: "The resources correspond to a Stable Diffusion XL model")
var isXL: Bool = false
@Flag(name: .customLong("sd3"), help: "The resources correspond to a Stable Diffusion 3 model")
var isSD3: Bool = false
@Option(help: "Path to starting image.")
var image: String? = nil
@Option(help: "Strength for image2image.")
var strength: Float = 0.5
@Option(help: "Number of images to sample / generate")
var imageCount: Int = 1
@Option(help: "Number of diffusion steps to perform")
var stepCount: Int = 50
@Option(
help: ArgumentHelp(
"How often to save samples at intermediate steps",
discussion: "Set to 0 to only save the final sample"
)
)
var saveEvery: Int = 0
@Option(help: "Output path")
var outputPath: String = "./"
@Option(help: "Random seed")
var seed: UInt32 = UInt32.random(in: 0...UInt32.max)
@Option(help: "Controls the influence of the text prompt on sampling process (0=random images)")
var guidanceScale: Float = 7.5
@Option(help: "Compute units to load model with {all,cpuOnly,cpuAndGPU,cpuAndNeuralEngine}")
var computeUnits: ComputeUnits = .all
@Option(help: "Scheduler to use, one of {pndm, dpmpp}")
var scheduler: SchedulerOption = .pndm
@Option(help: "Random number generator to use, one of {numpy, torch, nvidia}")
var rng: RNGOption = .numpy
@Option(
parsing: .upToNextOption,
help: "ControlNet models used in image generation (enter file names in Resources/controlnet without extension)"
)
var controlnet: [String] = []
@Option(
parsing: .upToNextOption,
help: "image for each controlNet model (corresponding to the same order as --controlnet)"
)
var controlnetInputs: [String] = []
@Flag(help: "Disable safety checking")
var disableSafety: Bool = false
@Flag(help: "Reduce memory usage")
var reduceMemory: Bool = false
@Flag(help: "Use system multilingual NLContextualEmbedding as encoder model")
var useMultilingualTextEncoder: Bool = false
@Option(help: "The natural language script for the multilingual contextual embedding")
var script: Script = .latin
mutating func run() throws {
guard FileManager.default.fileExists(atPath: resourcePath) else {
throw RunError.resources("Resource path does not exist \(resourcePath)")
}
let config = MLModelConfiguration()
config.computeUnits = computeUnits.asMLComputeUnits
let resourceURL = URL(filePath: resourcePath)
log("Loading resources and creating pipeline\n")
log("(Note: This can take a while the first time using these resources)\n")
let pipeline: StableDiffusionPipelineProtocol
var scaleFactor: Float32 = 0.18215
var shiftFactor: Float32 = 0.0
var timestepShift: Float32 = 1.0
if #available(macOS 14.0, iOS 17.0, *) {
if isXL {
scaleFactor = 0.13025
if !controlnet.isEmpty {
throw RunError.unsupported("ControlNet is not supported for Stable Diffusion XL")
}
if useMultilingualTextEncoder {
throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion XL")
}
pipeline = try StableDiffusionXLPipeline(
resourcesAt: resourceURL,
configuration: config,
reduceMemory: reduceMemory
)
} else if isSD3 {
scaleFactor = 1.5305
shiftFactor = 0.0609
timestepShift = 3.0
if !controlnet.isEmpty {
throw RunError.unsupported("ControlNet is not supported for Stable Diffusion 3")
}
if useMultilingualTextEncoder {
throw RunError.unsupported("Multilingual text encoder is not yet supported for Stable Diffusion 3")
}
pipeline = try StableDiffusion3Pipeline(
resourcesAt: resourceURL,
configuration: config,
reduceMemory: reduceMemory
)
} else {
pipeline = try StableDiffusionPipeline(
resourcesAt: resourceURL,
controlNet: controlnet,
configuration: config,
disableSafety: disableSafety,
reduceMemory: reduceMemory,
useMultilingualTextEncoder: useMultilingualTextEncoder,
script: script
)
}
} else {
pipeline = try StableDiffusionPipeline(
resourcesAt: resourceURL,
controlNet: controlnet,
configuration: config,
disableSafety: disableSafety,
reduceMemory: reduceMemory
)
}
try pipeline.loadResources()
let startingImage: CGImage?
if let image {
let imageURL = URL(filePath: image)
do {
startingImage = try convertImageToCGImage(imageURL: imageURL)
} catch let error {
throw RunError.resources("Starting image not found \(imageURL), error: \(error)")
}
} else {
startingImage = nil
}
// convert image for ControlNet into CGImage when controlNet available
let controlNetInputs: [CGImage]
if !controlnet.isEmpty {
controlNetInputs = try controlnetInputs.map { imagePath in
let imageURL = URL(filePath: imagePath)
do {
return try convertImageToCGImage(imageURL: imageURL)
} catch let error {
throw RunError.resources("Image for ControlNet not found \(imageURL), error: \(error)")
}
}
} else {
controlNetInputs = []
}
log("Sampling ...\n")
let sampleTimer = SampleTimer()
sampleTimer.start()
var pipelineConfig = StableDiffusionPipeline.Configuration(prompt: prompt)
pipelineConfig.negativePrompt = negativePrompt
pipelineConfig.startingImage = startingImage
pipelineConfig.strength = strength
pipelineConfig.imageCount = imageCount
pipelineConfig.stepCount = stepCount
pipelineConfig.seed = seed
pipelineConfig.controlNetInputs = controlNetInputs
pipelineConfig.guidanceScale = guidanceScale
pipelineConfig.schedulerType = scheduler.stableDiffusionScheduler
pipelineConfig.rngType = rng.stableDiffusionRNG
pipelineConfig.useDenoisedIntermediates = true
pipelineConfig.encoderScaleFactor = scaleFactor
pipelineConfig.decoderScaleFactor = scaleFactor
pipelineConfig.decoderShiftFactor = shiftFactor
pipelineConfig.schedulerTimestepShift = timestepShift
let images = try pipeline.generateImages(
configuration: pipelineConfig) { progress in
sampleTimer.stop()
handleProgress(progress,sampleTimer)
if progress.stepCount != progress.step {
sampleTimer.start()
}
return true
}
_ = try saveImages(images, logNames: true)
}
func convertImageToCGImage(imageURL: URL) throws -> CGImage {
let imageData = try Data(contentsOf: imageURL)
guard
let nsImage = NSImage(data: imageData),
let loadedImage = nsImage.cgImage(forProposedRect: nil, context: nil, hints: nil)
else {
throw RunError.resources("Image not available \(resourcePath)")
}
return loadedImage
}
func handleProgress(
_ progress: StableDiffusionPipeline.Progress,
_ sampleTimer: SampleTimer
) {
log("\u{1B}[1A\u{1B}[K")
log("Step \(progress.step) of \(progress.stepCount) ")
log(" [")
log(String(format: "mean: %.2f, ", 1.0/sampleTimer.mean))
log(String(format: "median: %.2f, ", 1.0/sampleTimer.median))
log(String(format: "last %.2f", 1.0/sampleTimer.allSamples.last!))
log("] step/sec")
if saveEvery > 0, progress.step % saveEvery == 0 {
let saveCount = (try? saveImages(progress.currentImages, step: progress.step)) ?? 0
log(" saved \(saveCount) image\(saveCount != 1 ? "s" : "")")
}
log("\n")
}
func saveImages(
_ images: [CGImage?],
step: Int? = nil,
logNames: Bool = false
) throws -> Int {
let url = URL(filePath: outputPath)
var saved = 0
for i in 0 ..< images.count {
guard let image = images[i] else {
if logNames {
log("Image \(i) failed safety check and was not saved")
}
continue
}
let name = imageName(i, step: step)
let fileURL = url.appending(path:name)
guard let dest = CGImageDestinationCreateWithURL(fileURL as CFURL, UTType.png.identifier as CFString, 1, nil) else {
throw RunError.saving("Failed to create destination for \(fileURL)")
}
CGImageDestinationAddImage(dest, image, nil)
if !CGImageDestinationFinalize(dest) {
throw RunError.saving("Failed to save \(fileURL)")
}
if logNames {
log("Saved \(name)\n")
}
saved += 1
}
return saved
}
func imageName(_ sample: Int, step: Int? = nil) -> String {
let fileCharLimit = 75
var name = prompt.prefix(fileCharLimit).replacingOccurrences(of: " ", with: "_")
if imageCount != 1 {
name += ".\(sample)"
}
if image != nil {
name += ".str\(Int(strength * 100))"
}
name += ".\(seed)"
if let step = step {
name += ".\(step)"
} else {
name += ".final"
}
name += ".png"
return name
}
func log(_ str: String, term: String = "") {
print(str, terminator: term)
}
}
enum RunError: Error {
case resources(String)
case saving(String)
case unsupported(String)
}
@available(iOS 16.2, macOS 13.1, *)
enum ComputeUnits: String, ExpressibleByArgument, CaseIterable {
case all, cpuAndGPU, cpuOnly, cpuAndNeuralEngine
var asMLComputeUnits: MLComputeUnits {
switch self {
case .all: return .all
case .cpuAndGPU: return .cpuAndGPU
case .cpuOnly: return .cpuOnly
case .cpuAndNeuralEngine: return .cpuAndNeuralEngine
}
}
}
@available(iOS 16.2, macOS 13.1, *)
enum SchedulerOption: String, ExpressibleByArgument {
case pndm, dpmpp
var stableDiffusionScheduler: StableDiffusionScheduler {
switch self {
case .pndm: return .pndmScheduler
case .dpmpp: return .dpmSolverMultistepScheduler
}
}
}
@available(iOS 16.2, macOS 13.1, *)
enum RNGOption: String, ExpressibleByArgument {
case numpy, torch, nvidia
var stableDiffusionRNG: StableDiffusionRNG {
switch self {
case .numpy: return .numpyRNG
case .torch: return .torchRNG
case .nvidia: return .nvidiaRNG
}
}
}
@available(iOS 16.2, macOS 13.1, *)
extension Script: ExpressibleByArgument {}
if #available(iOS 16.2, macOS 13.1, *) {
StableDiffusionSample.main()
} else {
print("Unsupported OS")
}
swift/StableDiffusionTests/Resources/merges.txt
+48895
View File
File diff suppressed because it is too large Load Diff
swift/StableDiffusionTests/Resources/vocab.json
+1
View File
File diff suppressed because one or more lines are too long
swift/StableDiffusionTests/StableDiffusionTests.swift
+63
View File
@@ -0,0 +1,63 @@
// For licensing see accompanying LICENSE.md file.
// Copyright (C) 2022 Apple Inc. All Rights Reserved.
import XCTest
import CoreML
@testable import StableDiffusion
@available(iOS 16.2, macOS 13.1, *)
final class StableDiffusionTests: XCTestCase {
var vocabFileInBundleURL: URL {
let fileName = "vocab"
guard let url = Bundle.module.url(forResource: fileName, withExtension: "json") else {
fatalError("BPE tokenizer vocabulary file is missing from bundle")
}
return url
}
var mergesFileInBundleURL: URL {
let fileName = "merges"
guard let url = Bundle.module.url(forResource: fileName, withExtension: "txt") else {
fatalError("BPE tokenizer merges file is missing from bundle")
}
return url
}
func testBPETokenizer() throws {
let tokenizer = try BPETokenizer(mergesAt: mergesFileInBundleURL, vocabularyAt: vocabFileInBundleURL)
func testPrompt(prompt: String, expectedIds: [Int]) {
let (tokens, ids) = tokenizer.tokenize(input: prompt)
print("Tokens = \(tokens)\n")
print("Expected tokens = \(expectedIds.map({ tokenizer.token(id: $0) }))")
print("ids = \(ids)\n")
print("Expected Ids = \(expectedIds)\n")
XCTAssertEqual(ids,expectedIds)
}
testPrompt(prompt: "a photo of an astronaut riding a horse on mars",
expectedIds: [49406, 320, 1125, 539, 550, 18376, 6765, 320, 4558, 525, 7496, 49407])
testPrompt(prompt: "Apple CoreML developer tools on a Macbook Air are fast",
expectedIds: [49406, 3055, 19622, 5780, 10929, 5771, 525, 320, 20617,
1922, 631, 1953, 49407])
}
func test_randomNormalValues_matchNumPyRandom() {
var random = NumPyRandomSource(seed: 12345)
let samples = random.normalArray(count: 10_000)
let last5 = samples.suffix(5)
// numpy.random.seed(12345); print(numpy.random.randn(10000)[-5:])
let expected = [-0.86285345, 2.15229409, -0.00670556, -1.21472309, 0.65498866]
for (value, expected) in zip(last5, expected) {
XCTAssertEqual(value, expected, accuracy: .ulpOfOne.squareRoot())
}
}
}
tests/__init__.py
View File
tests/test_stable_diffusion.py
+410
View File
@@ -0,0 +1,410 @@
#
# For licensing see accompanying LICENSE.md file.
# Copyright (C) 2022 Apple Inc. All Rights Reserved.
#
import argparse
import contextlib
import coremltools as ct
from diffusers import StableDiffusionPipeline
import json
import logging
import numpy as np
import os
import unittest
from PIL import Image
from statistics import median
import tempfile
import time
import torch
torch.set_grad_enabled(False)
from python_coreml_stable_diffusion import torch2coreml, pipeline, coreml_model
logger = logging.getLogger(__name__)
logger.setLevel("INFO")
# Testing configuration
TEST_SEED = 93
TEST_PROMPT = "a high quality photo of an astronaut riding a horse in space"
TEST_COMPUTE_UNIT = ["CPU_AND_GPU", "ALL", "CPU_AND_NE"]
TEST_PSNR_THRESHOLD = 35 # dB
TEST_ABSOLUTE_MAX_LATENCY = 90 # seconds
TEST_WARMUP_INFERENCE_STEPS = 3
TEST_TEXT_TO_IMAGE_SPEED_REPEATS = 3
TEST_MINIMUM_PROMPT_TO_IMAGE_CLIP_COSINE_SIMILARITY = 0.3 # in range [0.,1.]
class TestStableDiffusionForTextToImage(unittest.TestCase):
""" Test Stable Diffusion text-to-image pipeline for:
- PyTorch to CoreML conversion via coremltools
- Speed of CoreML runtime across several compute units
- Integration with `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py`
- Efficacy of the safety_checker
- Affinity of the generated image with the original prompt via CLIP score
- The bridge between Python and Swift CLI
- The signal parity of Swift CLI generated image with that of Python CLI
"""
cli_args = None
@classmethod
def setUpClass(cls):
cls.pytorch_pipe = StableDiffusionPipeline.from_pretrained(
cls.cli_args.model_version,
use_auth_token=True,
)
# To be initialized after test_torch_to_coreml_conversion is run
cls.coreml_pipe = None
cls.active_compute_unit = None
@classmethod
def tearDownClass(cls):
cls.pytorch_pipe = None
cls.coreml_pipe = None
cls.active_compute_unit = None
def test_torch_to_coreml_conversion(self):
""" Tests:
- PyTorch to CoreML conversion via coremltools
"""
with self.subTest(model="vae_decoder"):
logger.info("Converting vae_decoder")
torch2coreml.convert_vae_decoder(self.pytorch_pipe, self.cli_args)
logger.info("Successfully converted vae_decoder")
with self.subTest(model="unet"):
logger.info("Converting unet")
torch2coreml.convert_unet(self.pytorch_pipe, self.cli_args)
logger.info("Successfully converted unet")
with self.subTest(model="text_encoder"):
logger.info("Converting text_encoder")
torch2coreml.convert_text_encoder(self.pytorch_pipe, self.cli_args)
logger.info("Successfully converted text_encoder")
with self.subTest(model="safety_checker"):
logger.info("Converting safety_checker")
torch2coreml.convert_safety_checker(self.pytorch_pipe,
self.cli_args)
logger.info("Successfully converted safety_checker")
def test_end_to_end_image_generation_speed(self):
""" Tests:
- Speed of CoreML runtime across several compute units
- Integration with `diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.py`
"""
latency = {
compute_unit:
self._coreml_text_to_image_with_compute_unit(compute_unit)
for compute_unit in TEST_COMPUTE_UNIT
}
latency["num_repeats_for_median"] = TEST_TEXT_TO_IMAGE_SPEED_REPEATS
json_path = os.path.join(self.cli_args.o, "benchmark.json")
logger.info(f"Saving inference benchmark results to {json_path}")
with open(json_path, "w") as f:
json.dump(latency, f)
for compute_unit in TEST_COMPUTE_UNIT:
with self.subTest(compute_unit=compute_unit):
self.assertGreater(TEST_ABSOLUTE_MAX_LATENCY,
latency[compute_unit])
def test_image_to_prompt_clip_score(self):
""" Tests:
Affinity of the generated image with the original prompt via CLIP score
"""
logger.warning(
"This test will download the CLIP ViT-B/16 model (approximately 600 MB) from Hugging Face"
)
from transformers import CLIPProcessor, CLIPModel
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
processor = CLIPProcessor.from_pretrained(
"openai/clip-vit-base-patch16")
for compute_unit in TEST_COMPUTE_UNIT:
with self.subTest(compute_unit=compute_unit):
image_path = pipeline.get_image_path(self.cli_args,
prompt=TEST_PROMPT,
compute_unit=compute_unit)
image = Image.open(image_path)
# Preprocess images and text for inference with CLIP
inputs = processor(text=[TEST_PROMPT],
images=image,
return_tensors="pt",
padding=True)
outputs = model(**inputs)
# Compute cosine similarity between image and text embeddings
image_text_cosine_similarity = outputs.image_embeds @ outputs.text_embeds.T
logger.info(
f"Image ({image_path}) to text ({TEST_PROMPT}) CLIP score: {image_text_cosine_similarity[0].item():.2f}"
)
# Ensure that the minimum cosine similarity threshold is achieved
self.assertGreater(
image_text_cosine_similarity,
TEST_MINIMUM_PROMPT_TO_IMAGE_CLIP_COSINE_SIMILARITY,
)
def test_safety_checker_efficacy(self):
""" Tests:
- Efficacy of the safety_checker
"""
self._init_coreml_pipe(compute_unit=self.active_compute_unit)
safety_checker_test_prompt = "NSFW"
image = self.coreml_pipe(safety_checker_test_prompt)
# Image must have been erased by the safety checker
self.assertEqual(np.array(image["images"][0]).sum(), 0.)
self.assertTrue(image["nsfw_content_detected"].any())
def test_swift_cli_image_generation(self):
""" Tests:
- The bridge between Python and Swift CLI
- The signal parity of Swift CLI generated image with that of Python CLI
"""
# coremltools to Core ML compute unit mapping
compute_unit_map = {
"ALL": "all",
"CPU_AND_GPU": "cpuAndGPU",
"CPU_AND_NE": "cpuAndNeuralEngine"
}
# Prepare resources for Swift CLI
resources_dir = torch2coreml.bundle_resources_for_swift_cli(
self.cli_args)
logger.info("Bundled resources for Swift CLI")
# Execute image generation with Swift CLI
# Note: First time takes ~5 minutes due to project building and so on
cmd = " ".join([
f"swift run StableDiffusionSample \"{TEST_PROMPT}\"",
f"--resource-path {resources_dir}",
f"--seed {TEST_SEED}",
f"--output-path {self.cli_args.o}",
f"--compute-units {compute_unit_map[TEST_COMPUTE_UNIT[-1]]}"
])
logger.info(f"Executing `{cmd}`")
os.system(cmd)
logger.info(f"Image generation with Swift CLI is complete")
# Load Swift CLI generated image
swift_cli_image = Image.open(
os.path.join(
self.cli_args.o, "_".join(TEST_PROMPT.rsplit(" ")) + "." +
str(TEST_SEED) + ".final.png"))
# Load Python CLI (pipeline.py) generated image
python_cli_image = Image.open(pipeline.get_image_path(self.cli_args,
prompt=TEST_PROMPT,
compute_unit=TEST_COMPUTE_UNIT[-1]))
# Compute signal parity
swift2torch_psnr = torch2coreml.report_correctness(
np.array(swift_cli_image.convert("RGB")),
np.array(python_cli_image.convert("RGB")),
"Swift CLI and Python CLI generated images")
self.assertGreater(swift2torch_psnr, torch2coreml.ABSOLUTE_MIN_PSNR)
def _init_coreml_pipe(self, compute_unit):
""" Initializes CoreML pipe for the requested compute_unit
"""
assert compute_unit in ct.ComputeUnit._member_names_, f"Not a valid coremltools.ComputeUnit: {compute_unit}"
if self.active_compute_unit == compute_unit:
logger.info(
"self.coreml_pipe matches requested compute_unit, skipping reinitialization"
)
assert \
isinstance(self.coreml_pipe, pipeline.CoreMLStableDiffusionPipeline), \
type(self.coreml_pipe)
else:
self.active_compute_unit = compute_unit
self.coreml_pipe = pipeline.get_coreml_pipe(
pytorch_pipe=self.pytorch_pipe,
mlpackages_dir=self.cli_args.o,
model_version=self.cli_args.model_version,
compute_unit=self.active_compute_unit,)
def _coreml_text_to_image_with_compute_unit(self, compute_unit):
""" Benchmark end-to-end text-to-image generation with the requested compute_unit
"""
self._init_coreml_pipe(compute_unit)
# Warm up (not necessary in all settings but improves consistency for benchmarking)
logger.info(
f"Warmup image generation with {TEST_WARMUP_INFERENCE_STEPS} inference steps"
)
image = self.coreml_pipe(
TEST_PROMPT, num_inference_steps=TEST_WARMUP_INFERENCE_STEPS)
# Test end-to-end speed
logger.info(
f"Run full image generation {TEST_TEXT_TO_IMAGE_SPEED_REPEATS} times and report median"
)
def test_coreml_text_to_image_speed():
""" Execute Core ML based image generation
"""
_reset_seed()
image = self.coreml_pipe(TEST_PROMPT)["images"][0]
out_path = pipeline.get_image_path(self.cli_args,
prompt=TEST_PROMPT,
compute_unit=compute_unit)
logger.info(f"Saving generated image to {out_path}")
image.save(out_path)
def collect_timings(callable, n):
""" Collect user latency for callable
"""
user_latencies = []
for _ in range(n):
s = time.time()
callable()
user_latencies.append(float(f"{time.time() - s:.2f}"))
return user_latencies
coreml_latencies = collect_timings(
callable=test_coreml_text_to_image_speed,
n=TEST_TEXT_TO_IMAGE_SPEED_REPEATS)
coreml_median_latency = median(coreml_latencies)
logger.info(
f"End-to-end latencies with coremltools.ComputeUnit.{compute_unit}: median={coreml_median_latency:.2f}"
)
return coreml_median_latency
def _reset_seed():
""" Reset RNG state in order to reproduce the results across multiple runs
"""
torch.manual_seed(TEST_SEED)
np.random.seed(TEST_SEED)
def _get_test_artifacts_dir(args):
if cli_args.persistent_test_artifacts_dir is not None:
os.makedirs(cli_args.persistent_test_artifacts_dir, exist_ok=True)
return contextlib.nullcontext(
enter_result=cli_args.persistent_test_artifacts_dir)
else:
return tempfile.TemporaryDirectory(
prefix="python_coreml_stable_diffusion_tests")
def _extend_parser(parser):
parser.add_argument(
"--persistent-test-artifacts-dir",
type=str,
default=None,
help=
("If specified, test artifacts such as Core ML models and generated images are saved in this directory. ",
"Otherwise, all artifacts are erased after the test program terminates."
))
parser.add_argument(
"--fast",
action="store_true",
help=
"If specified, runs fewer repeats for `test_end_to_end_image_generation_speed`"
)
parser.add_argument(
"--test-image-to-prompt-clip-score-opt-in",
action="store_true",
help=
("If specified, enables `test_image_to_prompt_clip_score` to verify the relevance of the "
"generated image content to the original text prompt. This test is an opt-in "
"test because it involves an additional one time 600MB model download."
))
parser.add_argument(
"--test-swift-cli-opt-in",
action="store_true",
help=
("If specified, compiles all models and builds the Swift CLI to run image generation and compares "
"results across Python and Swift runtime"))
parser.add_argument(
"--test-safety-checker-efficacy-opt-in",
action="store_true",
help=
("If specified, generates a potentially NSFW image to check whether the `safety_checker` "
"accurately detects and removes the content"))
return parser
if __name__ == "__main__":
# Reproduce the CLI of the original pipeline
parser = torch2coreml.parser_spec()
parser = _extend_parser(parser)
cli_args = parser.parse_args()
cli_args.check_output_correctness = True
cli_args.prompt = TEST_PROMPT
cli_args.seed = TEST_SEED
cli_args.compute_unit = TEST_COMPUTE_UNIT[0]
cli_args.scheduler = None # use default
torch2coreml.ABSOLUTE_MIN_PSNR = TEST_PSNR_THRESHOLD
if cli_args.fast:
logger.info(
"`--fast` detected: Image generation will be run once " \
f"(instead of {TEST_TEXT_TO_IMAGE_SPEED_REPEATS } times) " \
"with ComputeUnit.ALL (other compute units are skipped)" \
" (median can not be reported)")
TEST_TEXT_TO_IMAGE_SPEED_REPEATS = 1
TEST_COMPUTE_UNIT = ["ALL"]
logger.info("`--fast` detected: Skipping `--check-output-correctness` tests")
cli_args.check_output_correctness = False
elif cli_args.attention_implementation == "ORIGINAL":
TEST_COMPUTE_UNIT = ["CPU_AND_GPU", "ALL"]
elif cli_args.attention_implementation == "SPLIT_EINSUM":
TEST_COMPUTE_UNIT = ["ALL", "CPU_AND_NE"]
logger.info(f"Testing compute units: {TEST_COMPUTE_UNIT}")
# Save CoreML model files and generated images into the artifacts dir
with _get_test_artifacts_dir(cli_args) as test_artifacts_dir:
cli_args.o = test_artifacts_dir
logger.info(f"Test artifacts will be saved under {test_artifacts_dir}")
TestStableDiffusionForTextToImage.cli_args = cli_args
# Run the following tests in sequential order
suite = unittest.TestSuite()
suite.addTest(
TestStableDiffusionForTextToImage(
"test_torch_to_coreml_conversion"))
suite.addTest(
TestStableDiffusionForTextToImage(
"test_end_to_end_image_generation_speed"))
if cli_args.test_safety_checker_efficacy_opt_in:
suite.addTest(
TestStableDiffusionForTextToImage("test_safety_checker_efficacy"))
if cli_args.test_image_to_prompt_clip_score_opt_in:
suite.addTest(
TestStableDiffusionForTextToImage(
"test_image_to_prompt_clip_score"))
if cli_args.test_swift_cli_opt_in:
suite.addTest(
TestStableDiffusionForTextToImage(
"test_swift_cli_image_generation"))
if os.getenv("DEBUG", False):
suite.debug()
else:
runner = unittest.TextTestRunner()
runner.run(suite)