update lyrasd2

.gitattributes

@@ -25,6 +25,7 @@
 *.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
@@ -32,19 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-lyraSD/muse_trt/libnvinfer_plugin.so filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/clip.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/superx4.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/unet_fp16.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/vae-decoder.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/vae-encoder.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/scheduler filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/superx4-512-512.plan filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/text_encoder filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/tokenizer filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/vae filter=lfs diff=lfs merge=lfs -text
-sd1.5-engine/feature_extractor filter=lfs diff=lfs merge=lfs -text
-sd1.4-engine/clip.plan filter=lfs diff=lfs merge=lfs -text
-sd1.4-engine/superx4-512-512.plan filter=lfs diff=lfs merge=lfs -text
-sd1.4-engine/unet_fp16.plan filter=lfs diff=lfs merge=lfs -text
-sd1.4-engine/vae-decoder.plan filter=lfs diff=lfs merge=lfs -text
+lyrasd_model/lyrasd_lib/libth_lyrasd_cu11_sm80.so filter=lfs diff=lfs merge=lfs -text
+lyrasd_model/lyrasd_lib/libth_lyrasd_cu11_sm86.so filter=lfs diff=lfs merge=lfs -text
+lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm80.so filter=lfs diff=lfs merge=lfs -text
+lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,3 +1,5 @@
-*.un~
+.idea
 *.pyc
-__pycache__/
+.vscode
+__pycache__
+models/lyrasd*

CHANGELOG.md

@@ -1,4 +0,0 @@
-## V1.0
-
-- Add accelerated Stable Diffusion pretrained model v1.4 (from: https://huggingface.co/CompVis/stable-diffusion-v1-4)
-- Add accelerated Real-ESRGAN(4x) (from https://github.com/xinntao/Real-ESRGAN)

LISENCE

@@ -1,494 +0,0 @@
-CreativeML Open RAIL-M License
-
-Copyright (c) 2023 Tencent Music Entertainment
-
-
-Terms of the CreativeML Open RAIL-M:
---------------------------------------------------------------------
-Copyright (c) 2022 Robin Rombach and Patrick Esser and contributors
-
-CreativeML Open RAIL-M
-dated August 22, 2022
-
-Section I: PREAMBLE
-
-Multimodal generative models are being widely adopted and used, and have the potential to transform the way artists, among other individuals, conceive and benefit from AI or ML technologies as a tool for content creation.
-
-Notwithstanding the current and potential benefits that these artifacts can bring to society at large, there are also concerns about potential misuses of them, either due to their technical limitations or ethical considerations.
-
-In short, this license strives for both the open and responsible downstream use of the accompanying model. When it comes to the open character, we took inspiration from open source permissive licenses regarding the grant of IP rights. Referring to the downstream responsible use, we added use-based restrictions not permitting the use of the Model in very specific scenarios, in order for the licensor to be able to enforce the license in case potential misuses of the Model may occur. At the same time, we strive to promote open and responsible research on generative models for art and content generation.
-
-Even though downstream derivative versions of the model could be released under different licensing terms, the latter will always have to include - at minimum - the same use-based restrictions as the ones in the original license (this license). We believe in the intersection between open and responsible AI development; thus, this License aims to strike a balance between both in order to enable responsible open-science in the field of AI.
-
-This License governs the use of the model (and its derivatives) and is informed by the model card associated with the model.
-
-NOW THEREFORE, You and Licensor agree as follows:
-
-1. Definitions
-
-- "License" means the terms and conditions for use, reproduction, and Distribution as defined in this document.
-- "Data" means a collection of information and/or content extracted from the dataset used with the Model, including to train, pretrain, or otherwise evaluate the Model. The Data is not licensed under this License.
-- "Output" means the results of operating a Model as embodied in informational content resulting therefrom.
-- "Model" means any accompanying machine-learning based assemblies (including checkpoints), consisting of learnt weights, parameters (including optimizer states), corresponding to the model architecture as embodied in the Complementary Material, that have been trained or tuned, in whole or in part on the Data, using the Complementary Material.
-- "Derivatives of the Model" means all modifications to the Model, works based on the Model, or any other model which is created or initialized by transfer of patterns of the weights, parameters, activations or output of the Model, to the other model, in order to cause the other model to perform similarly to the Model, including - but not limited to - distillation methods entailing the use of intermediate data representations or methods based on the generation of synthetic data by the Model for training the other model.
-- "Complementary Material" means the accompanying source code and scripts used to define, run, load, benchmark or evaluate the Model, and used to prepare data for training or evaluation, if any. This includes any accompanying documentation, tutorials, examples, etc, if any.
-- "Distribution" means any transmission, reproduction, publication or other sharing of the Model or Derivatives of the Model to a third party, including providing the Model as a hosted service made available by electronic or other remote means - e.g. API-based or web access.
-- "Licensor" means the copyright owner or entity authorized by the copyright owner that is granting the License, including the persons or entities that may have rights in the Model and/or distributing the Model.
-- "You" (or "Your") means an individual or Legal Entity exercising permissions granted by this License and/or making use of the Model for whichever purpose and in any field of use, including usage of the Model in an end-use application - e.g. chatbot, translator, image generator.
-- "Third Parties" means individuals or legal entities that are not under common control with Licensor or You.
-- "Contribution" means any work of authorship, including the original version of the Model and any modifications or additions to that Model or Derivatives of the Model thereof, that is intentionally submitted to Licensor for inclusion in the Model 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 Model, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
-- "Contributor" means Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Model.
-
-Section II: INTELLECTUAL PROPERTY RIGHTS
-
-Both copyright and patent grants apply to the Model, Derivatives of the Model and Complementary Material. The Model and Derivatives of the Model are subject to additional terms as described in Section III.
-
-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, publicly display, publicly perform, sublicense, and distribute the Complementary Material, the Model, and Derivatives of the Model.
-3. Grant of Patent License. Subject to the terms and conditions of this License and where and as applicable, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this paragraph) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Model and the Complementary Material, 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 Model 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 Model and/or Complementary Material or a Contribution incorporated within the Model and/or Complementary Material constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for the Model and/or Work shall terminate as of the date such litigation is asserted or filed.
-
-Section III: CONDITIONS OF USAGE, DISTRIBUTION AND REDISTRIBUTION
-
-4. Distribution and Redistribution. You may host for Third Party remote access purposes (e.g. software-as-a-service), reproduce and distribute copies of the Model or Derivatives of the Model thereof in any medium, with or without modifications, provided that You meet the following conditions:
-Use-based restrictions as referenced in paragraph 5 MUST be included as an enforceable provision by You in any type of legal agreement (e.g. a license) governing the use and/or distribution of the Model or Derivatives of the Model, and You shall give notice to subsequent users You Distribute to, that the Model or Derivatives of the Model are subject to paragraph 5. This provision does not apply to the use of Complementary Material.
-You must give any Third Party recipients of the Model or Derivatives of the Model a copy of this License;
-You must cause any modified files to carry prominent notices stating that You changed the files;
-You must retain all copyright, patent, trademark, and attribution notices excluding those notices that do not pertain to any part of the Model, Derivatives of the Model.
-You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions - respecting paragraph 4.a. - for use, reproduction, or Distribution of Your modifications, or for any such Derivatives of the Model as a whole, provided Your use, reproduction, and Distribution of the Model otherwise complies with the conditions stated in this License.
-5. Use-based restrictions. The restrictions set forth in Attachment A are considered Use-based restrictions. Therefore You cannot use the Model and the Derivatives of the Model for the specified restricted uses. You may use the Model subject to this License, including only for lawful purposes and in accordance with the License. Use may include creating any content with, finetuning, updating, running, training, evaluating and/or reparametrizing the Model. You shall require all of Your users who use the Model or a Derivative of the Model to comply with the terms of this paragraph (paragraph 5).
-6. The Output You Generate. Except as set forth herein, Licensor claims no rights in the Output You generate using the Model. You are accountable for the Output you generate and its subsequent uses. No use of the output can contravene any provision as stated in the License.
-
-Section IV: OTHER PROVISIONS
-
-7. Updates and Runtime Restrictions. To the maximum extent permitted by law, Licensor reserves the right to restrict (remotely or otherwise) usage of the Model in violation of this License, update the Model through electronic means, or modify the Output of the Model based on updates. You shall undertake reasonable efforts to use the latest version of the Model.
-8. Trademarks and related. Nothing in this License permits You to make use of Licensors’ trademarks, trade names, logos or to otherwise suggest endorsement or misrepresent the relationship between the parties; and any rights not expressly granted herein are reserved by the Licensors.
-9. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Model and the Complementary Material (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 Model, Derivatives of the Model, and the Complementary Material and assume any risks associated with Your exercise of permissions under this License.
-10. 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 Model and the Complementary Material (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.
-11. Accepting Warranty or Additional Liability. While redistributing the Model, Derivatives of the Model and the Complementary Material 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.
-12. If any provision of this License is held to be invalid, illegal or unenforceable, the remaining provisions shall be unaffected thereby and remain valid as if such provision had not been set forth herein.
-
-END OF TERMS AND CONDITIONS
-
-
-
-Attachment A
-
-Use Restrictions
-
-You agree not to use the Model or Derivatives of the Model:
-- In any way that violates any applicable national, federal, state, local or international law or regulation;
-- For the purpose of exploiting, harming or attempting to exploit or harm minors in any way;
-- To generate or disseminate verifiably false information and/or content with the purpose of harming others;
-- To generate or disseminate personal identifiable information that can be used to harm an individual;
-- To defame, disparage or otherwise harass others;
-- For fully automated decision making that adversely impacts an individual’s legal rights or otherwise creates or modifies a binding, enforceable obligation;
-- For any use intended to or which has the effect of discriminating against or harming individuals or groups based on online or offline social behavior or known or predicted personal or personality characteristics;
-- To exploit any of the vulnerabilities of a specific group of persons based on their age, social, physical or mental characteristics, in order to materially distort the behavior of a person pertaining to that group in a manner that causes or is likely to cause that person or another person physical or psychological harm;
-- For any use intended to or which has the effect of discriminating against individuals or groups based on legally protected characteristics or categories;
-- To provide medical advice and medical results interpretation;
-- To generate or disseminate information for the purpose to be used 
-
-
-
-
-Other dependencies and licenses:
-
-
-Open Source Software Licensed under the CreativeML Open RAIL-M License:
---------------------------------------------------------------------
-1. stable-diffuison 
-Files：https://huggingface.co/CompVis/stable-diffusion-v1-4
-License：CreativeML Open RAIL-M
-For details：https://huggingface.co/spaces/CompVis/stable-diffusion-license
-
-A copy of the MIT License is included in this file.
-
-
-Open Source Software Licensed under the Apache License Version 2.0:
---------------------------------------------------------------------
-1. huggingface/diffusers
-File：https://github.com/huggingface/diffusers
-License：Apache License V2
-For details：https://github.com/huggingface/diffusers/blob/main/LICENSE
-
-2. huggingface/transformers
-Copyright 2018- The Hugging Face team. All rights reserved.
-
-3. NVIDIA/TensorRT 
- Copyright 2021 NVIDIA Corporation
-
-   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.
-
-4. TensorRT/tools/Polygraphy
-Copyright 2020 NVIDIA Corporation
-
-   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.
-
-
-Terms of the Apache License Version 2.0:
---------------------------------------------------------------------
-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:
-
-You must give any other recipients of the Work or Derivative Works a copy of this License; and 
-
-You must cause any modified files to carry prominent notices stating that You changed the files; and 
-
-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 
-
-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
- 
-
-Open Source Software Licensed under the Modified BSD License:
---------------------------------------------------------------------
-1. Numpy
-Copyright (c) 2005-2023, NumPy Developers.
-All rights reserved.
-
-Terms of the Modified BSD License:
--------------------------------------------------------------------
-This project is licensed under the terms of the Modified BSD License, as follows:
-
-Copyright (c) 2005-2023, NumPy Developers.
-All rights reserved.
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-Redistributions of source code must retain the above copyright notice, this
-list of conditions and the following disclaimer.
-
-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.
-
-Neither the name of the NumPy Developers nor the names of any 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.
-
-2. 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.
-
-
-Terms of the Modified BSD License:
--------------------------------------------------------------------
-This project is licensed under the terms of the Modified BSD License, as follows:
-
-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.
-
-
-
-Open Source Software Licensed under the BSD 3-Clause License:
---------------------------------------------------------------------
-1. scipy
-Copyright (c) 2001-2002 Enthought, Inc. 2003-2023, SciPy Developers.
-
-Terms of the BSD 3-Clause License:
---------------------------------------------------------------------
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-* Redistributions of source code must retain the above copyright notice, this
-  list of conditions and the following disclaimer.
-
-* 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.
-
-* Neither the name of the copyright holder 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 HOLDER 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.
-
-
-Open Source Software Licensed under the Python Software Foundation License Version 2:
---------------------------------------------------------------------------
-1. Python/cpython
-Copyright © 2001-2023 Python Software Foundation. All rights reserved
-
-
-A. HISTORY OF THE SOFTWARE
-==========================
-
-Python was created in the early 1990s by Guido van Rossum at Stichting
-Mathematisch Centrum (CWI, see https://www.cwi.nl) in the Netherlands
-as a successor of a language called ABC.  Guido remains Python's
-principal author, although it includes many contributions from others.
-
-In 1995, Guido continued his work on Python at the Corporation for
-National Research Initiatives (CNRI, see https://www.cnri.reston.va.us)
-in Reston, Virginia where he released several versions of the
-software.
-
-In May 2000, Guido and the Python core development team moved to
-BeOpen.com to form the BeOpen PythonLabs team.  In October of the same
-year, the PythonLabs team moved to Digital Creations, which became
-Zope Corporation.  In 2001, the Python Software Foundation (PSF, see
-https://www.python.org/psf/) was formed, a non-profit organization
-created specifically to own Python-related Intellectual Property.
-Zope Corporation was a sponsoring member of the PSF.
-
-All Python releases are Open Source (see https://opensource.org for
-the Open Source Definition).  Historically, most, but not all, Python
-releases have also been GPL-compatible; the table below summarizes
-the various releases.
-
-    Release         Derived     Year        Owner       GPL-
-                    from                                compatible? (1)
-
-    0.9.0 thru 1.2              1991-1995   CWI         yes
-    1.3 thru 1.5.2  1.2         1995-1999   CNRI        yes
-    1.6             1.5.2       2000        CNRI        no
-    2.0             1.6         2000        BeOpen.com  no
-    1.6.1           1.6         2001        CNRI        yes (2)
-    2.1             2.0+1.6.1   2001        PSF         no
-    2.0.1           2.0+1.6.1   2001        PSF         yes
-    2.1.1           2.1+2.0.1   2001        PSF         yes
-    2.1.2           2.1.1       2002        PSF         yes
-    2.1.3           2.1.2       2002        PSF         yes
-    2.2 and above   2.1.1       2001-now    PSF         yes
-
-Footnotes:
-
-(1) GPL-compatible doesn't mean that we're distributing Python under
-    the GPL.  All Python licenses, unlike the GPL, let you distribute
-    a modified version without making your changes open source.  The
-    GPL-compatible licenses make it possible to combine Python with
-    other software that is released under the GPL; the others don't.
-
-(2) According to Richard Stallman, 1.6.1 is not GPL-compatible,
-    because its license has a choice of law clause.  According to
-    CNRI, however, Stallman's lawyer has told CNRI's lawyer that 1.6.1
-    is "not incompatible" with the GPL.
-
-Thanks to the many outside volunteers who have worked under Guido's
-direction to make these releases possible.
-
-
-B. TERMS AND CONDITIONS FOR ACCESSING OR OTHERWISE USING PYTHON
-===============================================================
-
-Python software and documentation are licensed under the
-Python Software Foundation License Version 2.
-
-Starting with Python 3.8.6, examples, recipes, and other code in
-the documentation are dual licensed under the PSF License Version 2
-and the Zero-Clause BSD license.
-
-Some software incorporated into Python is under different licenses.
-The licenses are listed with code falling under that license.
-
-
-PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2
---------------------------------------------
-
-1. This LICENSE AGREEMENT is between the Python Software Foundation
-("PSF"), and the Individual or Organization ("Licensee") accessing and
-otherwise using this software ("Python") in source or binary form and
-its associated documentation.
-
-2. Subject to the terms and conditions of this License Agreement, PSF hereby
-grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce,
-analyze, test, perform and/or display publicly, prepare derivative works,
-distribute, and otherwise use Python alone or in any derivative version,
-provided, however, that PSF's License Agreement and PSF's notice of copyright,
-i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010,
-2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022, 2023 Python Software Foundation;
-All Rights Reserved" are retained in Python alone or in any derivative version
-prepared by Licensee.
-
-3. In the event Licensee prepares a derivative work that is based on
-or incorporates Python or any part thereof, and wants to make
-the derivative work available to others as provided herein, then
-Licensee hereby agrees to include in any such work a brief summary of
-the changes made to Python.
-
-4. PSF is making Python available to Licensee on an "AS IS"
-basis.  PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR
-IMPLIED.  BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND
-DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS
-FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT
-INFRINGE ANY THIRD PARTY RIGHTS.
-
-5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON
-FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS
-A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON,
-OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF.
-
-6. This License Agreement will automatically terminate upon a material
-breach of its terms and conditions.
-
-7. Nothing in this License Agreement shall be deemed to create any
-relationship of agency, partnership, or joint venture between PSF and
-Licensee.  This License Agreement does not grant permission to use PSF
-trademarks or trade name in a trademark sense to endorse or promote
-products or services of Licensee, or any third party.
-
-8. By copying, installing or otherwise using Python, Licensee
-agrees to be bound by the terms and conditions of this License
-Agreement.
-
-
-
-

README.md

@@ -4,19 +4,23 @@ language:
 - en
 tags:
 - art
-- tensorRT
+- Stable Diffusion
 ---
-## Model Card for lyraSD
+## Model Card for lyraSD2
 
-lyraSD is currently the **fastest Stable Diffusion model** available, boasting an inference cost of only **0.435 seconds** for a 512x512 image, accelerating the process up to **10 times faster** than the original version. 
+lyraSD2 is currently the **fastest Stable Diffusion model** that can 100% align the outputs of **diffusers** available, boasting an inference cost of only **0.52 seconds** for a 512x512 image, accelerating the process up to **80% faster** than the original version. 
 
 Among its main features are:
 
-- weights: original Stable Diffusion 1.4 weights
-- input image: 512x512 (in img2img mode)
-- output image: 512x512
-- device requirements: Nvidia Ampere architecture (A100) or compatable
-- super-resultion: 4x by default, optional.
+- 4 Commonly used Pipelines
+- - Text2Img
+- - Img2Img
+- - ControlNetText2Img
+- - ControlNetImg2Img
+- 100% likeness to diffusers output
+- ControlNet Hot Swap: Can hot swap a ControlNet model weights within 0.4s (0s if cached)
+- Lora How Swap: Can hot swap a Lora within 0.5s (0.1s if cached)
+- device requirements: Nvidia Ampere architecture (A100, A10) or compatable
 
 ## Speed
 
@@ -25,75 +29,125 @@ Among its main features are:
 - device: Nvidia A100 40G
 - img size: 512x512
 - percision:fp16
-- steps: 30
-- solver: LMSD
-
-### text2img
-|model|time cost(ms)|memory(MB)|
-|:-:|:-:|:-:|
-|Pytorch SD|~5000ms|~10240|
-|lyraSD|~435ms|~4026|
-
-### superResolution(SR)
-|model|time cost(ms)|memory(MB)|
-|:-:|:-:|:-:|
-|Pytorch SR|~720ms|~6650|
-|lyraSD|~26ms|~1600|
-
-
+- steps: 20
+- sampler: EulerA
 
+### Text2Img
+|model|time cost(ms)|
+|:-:|:-:|
+|torch2.0.1 + diffusers|~667ms|
+|lyraSD|~528ms|
 
+### ControlNet-Text2Img
+|model|time cost(ms)|
+|:-:|:-:|
+|torch2.0.1 + diffusers|~930ms|
+|lyraSD2|~745ms|
 
 ## Model Sources
 
-- **Repository:** https://huggingface.co/CompVis/stable-diffusion-v1-4
+- **Checkpoint:** https://civitai.com/models/7371/rev-animated
+- **ControlNet:** https://huggingface.co/lllyasviel/sd-controlnet-canny
+- **Lora:** https://civitai.com/models/18323?modelVersionId=46846
 
-## Uses
+## Text2Img Uses
 
 ```python
-from lyraSD import LyraSD
-
-t2imodel = LyraSD("text2img", "./sd1.4-engine")
-t2imodel.inference(prompt="A fantasy landscape, trending on artstation", use_super=False)
-
+import torch
+import time
+
+from lyrasd_model import LyraSdTxt2ImgPipeline
+
+# 存放模型文件的路径，应该包含一下结构：
+#   1. clip 模型
+#   2. 转换好的优化后的 unet 模型，放入其中的 unet_bins 文件夹
+#   3. vae 模型
+#   4. scheduler 配置
+
+# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
+lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu11_sm80.so"
+model_path = "./models/lyrasd_rev_animated"
+lora_path = "./models/lyrasd_xiaorenshu_lora"
+
+# 构建 Txt2Img 的 Pipeline
+model = LyraSdTxt2ImgPipeline(model_path, lib_path)
+
+# load lora
+# 参数分别为 lora 存放位置，名字，lora 强度，lora模型精度
+model.load_lora(lora_path, "xiaorenshu", 0.4, "fp32")
+
+# 准备应用的输入和超参数
+prompt = "a cat, cute, cartoon, concise, traditional, chinese painting, Tang and Song Dynasties, masterpiece, 4k, 8k, UHD, best quality"
+negative_prompt = "(((horrible))), (((scary))), (((naked))), (((large breasts))), high saturation, colorful, human:2, body:2, low quality, bad quality, lowres, out of frame, duplicate, watermark, signature, text, frames, cut, cropped, malformed limbs, extra limbs, (((missing arms))), (((missing legs)))"
+height, width = 512, 512
+steps = 30
+guidance_scale = 7
+generator = torch.Generator().manual_seed(123)
+num_images = 1
+
+start = time.perf_counter()
+# 推理生成
+images = model(prompt, height, width, steps,
+        guidance_scale, negative_prompt, num_images,
+        generator=generator)
+print("image gen cost: ",time.perf_counter() - start)
+# 存储生成的图片
+for i, image in enumerate(images):
+    image.save(f"outputs/res_txt2img_lora_{i}.png")
+
+# unload lora，参数为 lora 的名字，是否清除 lora 缓存
+# model.unload_lora("xiaorenshu", True)
 
-from PIL import Image
-i2imodel = LyraSD("img2img", "./sd1.4-engine")
-demo_img = Image.open("output/text2img_demo.jpg")
-i2imodel.inference(prompt="A fantasy landscape, trending on artstation", image=demo_img)
 
 ```
 ## Demo output
 
-### text2img
-![text2img_demo](./output/text2img_demo.jpg)
+### Text2Img
+#### Text2Img without Lora
+![text2img_demo](./outputs/res_txt2img_0.png)
 
-### img2img
+#### Text2Img with Lora
+![text2img_demo](./outputs/res_txt2img_lora_0.png)
 
-![text2img_demo](./output/img2img_input.jpg)
+### Img2Img
 
-![text2img_demo](./output/img2img_demo.jpg)
+#### Img2Img input
+<img src="https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/seaside_town.png?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1692601590;1865401590&q-key-time=1692601590;1865401590&q-header-list=&q-url-param-list=&q-signature=ca04ca92d990d94813029c0d9ef29537e5f4637c" alt="img2img input" width="512"/>
 
+#### Img2Img output
+![text2img_demo](./outputs/res_img2img_0.png)
 
+### ControlNet Text2Img
 
-## Environment
+#### Control Image
+![text2img_demo](./control_bird_canny.png)
 
-- hardware: Nvidia Ampere architecture (A100) or compatable
-- docker image avaible: https://hub.docker.com/r/bigmoyan/lyra_aigc/tags
-```
-docker pull bigmoyan/lyra_aigc:v0.1
+#### ControlNet Text2Img Output
+![text2img_demo](./outputs/res_controlnet_txt2img_0.png)
+
+## Docker Environment Recommendation
+
+- For Cuda 11.X: we recommend ```nvcr.io/nvidia/pytorch:22.12-py3```
+- For Cuda 12.0: we recommend ```nvcr.io/nvidia/pytorch:23.02-py3```
+
+```bash
+docker pull nvcr.io/nvidia/pytorch:23.02-py3
+docker run --rm -it --gpus all -v ./:/lyraSD2 nvcr.io/nvidia/pytorch:23.02-py3
+
+pip install -r requirements.txt
+python txt2img_demo.py
 ```
 
 ## Citation
 ``` bibtex
-@Misc{lyraSD2023,
-  author =       {Kangjian Wu, Zhengtao Wang, Bin Wu},
-  title =        {lyraSD: Accelerating Stable Diffusion by 10x+},
-  howpublished = {\url{https://huggingface.co/TMElyralab/lyraSD}},
+@Misc{lyraSD2_2023,
+  author =       {Kangjian Wu, Zhengtao Wang, Yibo Lu, Haoxiong Su, Bin Wu},
+  title =        {lyraSD2: Accelerating Stable Diffusion with best flexibility},
+  howpublished = {\url{https://huggingface.co/TMElyralab/lyraSD2}},
   year =         {2023}
 }
 ```
 
 ## Report bug
-- start a discussion to report any bugs!--> https://huggingface.co/TMElyralab/lyraSD/discussions
+- start a discussion to report any bugs!--> https://huggingface.co/TMElyralab/lyraSD2/discussions
 - report bug with a `[bug]` mark in the title.

controlnet_img2img_demo.py

@@ -0,0 +1,62 @@
+from io import BytesIO
+
+import requests
+import torch
+from PIL import Image
+
+from lyrasd_model import LyraSdControlnetImg2ImgPipeline
+
+# 存放模型文件的路径，应该包含一下结构：
+#   1. clip 模型
+#   2. 转换好的优化后的 unet 模型
+#   3. vae 模型
+#   4. scheduler 配置
+
+# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
+lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so"
+model_path = "./models/lyrasd_rev_animated"
+canny_controlnet_path = "./models/lyrasd_canny"
+
+# 构建 Img2Img 的 Pipeline
+model = LyraSdControlnetImg2ImgPipeline(model_path, lib_path)
+
+# load Controlnet 模型，最多load 3个
+model.load_controlnet_model("canny", canny_controlnet_path, "fp32")
+
+control_img = Image.open("control_bird_canny.png")
+
+# 准备应用的输入和超参数
+prompt = "a bird"
+negative_prompt = "NSFW"
+height, width = 512, 512
+steps = 20
+guidance_scale = 7.5
+generator = torch.Generator().manual_seed(123)
+num_images = 1
+
+# 可以一次性load 3 个 Controlnets，达到multi Controlnet的效果，这里的参数的长度需要对其
+# Controlnet 所输入的img list 长度应该和 controlnet scale 与 Controlnet name 一致，而内部的list长度需要和batch size一致
+# 对应的index 可以对其
+controlnet_images = [[control_img]] 
+controlnet_scale= [0.5]
+controlnet_names= ['canny'] 
+
+# 从 cos 上拿个图作为初始化图片
+init_image_url = "https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/seaside_town.png?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1692601590;1865401590&q-key-time=1692601590;1865401590&q-header-list=&q-url-param-list=&q-signature=ca04ca92d990d94813029c0d9ef29537e5f4637c"
+init_image = BytesIO(requests.get(init_image_url).content)
+init_image = Image.open(init_image).convert('RGB')
+init_image = init_image.resize((width, height), Image.Resampling.LANCZOS)
+guess_mode = False
+strength = 0.8
+
+# 推理生成
+images = model(prompt, init_image, strength, height, width, steps,
+               guidance_scale, negative_prompt, num_images,
+               generator=generator, controlnet_images=controlnet_images,
+               controlnet_scale=controlnet_scale, controlnet_names=controlnet_names,
+               guess_mode=guess_mode
+               )
+
+# 存储生成的图片
+for i, image in enumerate(images):
+    image.save(f"outputs/res_controlnet_img2img_{i}.png")

controlnet_txt2img_demo.py

@@ -0,0 +1,63 @@
+import torch
+import time
+from PIL import Image
+
+from lyrasd_model import LyraSdControlnetTxt2ImgPipeline
+
+# 存放模型文件的路径，应该包含一下结构：
+#   1. clip 模型
+#   2. 转换好的优化后的 unet 模型
+#   3. 转换好的优化后的 controlnet 模型
+#   4. vae 模型
+#   5. scheduler 配置
+
+# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
+lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so"
+model_path = "./models/lyrasd_rev_animated"
+canny_controlnet_path = "./models/lyrasd_canny"
+# 构建 Txt2Img 的 Pipeline
+pipe = LyraSdControlnetTxt2ImgPipeline(model_path, lib_path)
+
+# load Controlnet 模型，最多load 3个
+start = time.perf_counter()
+pipe.load_controlnet_model("canny", canny_controlnet_path, "fp32")
+print(f"controlnet load cost: {time.perf_counter() - start}")
+# 可以通过 get_loaded_controlnet 方法获取目前已经load 好的Controlnet list
+print(pipe.get_loaded_controlnet())
+
+# 可以通过unload_controlnet_model 方法unload Controlnet
+# pipe.unload_controlnet_model("canny")
+
+control_img = Image.open("control_bird_canny.png")
+
+# 准备应用的输入和超参数
+prompt = "a blue bird"
+negative_prompt = "NSFW"
+height, width = 512, 512
+steps = 20
+guidance_scale = 7.5
+generator = torch.Generator().manual_seed(123)
+num_images = 1
+guess_mode = False
+
+# 可以一次性load 3 个 Controlnets，达到multi Controlnet的效果，这里的参数的长度需要对其
+# Controlnet 所输入的img list 长度应该和 controlnet scale 与 Controlnet name 一致，而内部的list长度需要和batch size一致
+# 对应的index 可以对其
+controlnet_images = [[control_img]]
+controlnet_scale = [0.5]
+controlnet_names = ['canny']
+
+# 推理生成，返回结果都是生成好的 PIL.Image
+
+start = time.perf_counter()
+images = pipe(prompt, height, width, steps,
+                guidance_scale, negative_prompt, num_images,
+                generator=generator, controlnet_images=controlnet_images,
+                controlnet_scale=controlnet_scale, controlnet_names=controlnet_names,
+                guess_mode=guess_mode
+                )
+print("cur cost: ",time.perf_counter() - start)
+
+# 存储生成的图片
+for i, image in enumerate(images):
+    image.save(f"./outputs/res_controlnet_txt2img_{i}.png")

demo.py

@@ -1,12 +0,0 @@
-from lyraSD import LyraSD
-
-t2imodel = LyraSD("text2img", "./sd1.4-engine")
-t2imodel.inference(prompt="A fantasy landscape, trending on artstation", use_super=True)
-
-
-from PIL import Image
-i2imodel = LyraSD("img2img", "./sd1.4-engine")
-demo_img = Image.open("output/img2img_input.jpg")
-i2imodel.inference(prompt="A fantasy landscape, trending on artstation", 
-                   image=demo_img, use_super=True)
-

img2img_demo.py

@@ -0,0 +1,47 @@
+from io import BytesIO
+
+import requests
+import torch
+from PIL import Image
+
+from lyrasd_model import LyraSDImg2ImgPipeline
+
+# 存放模型文件的路径，应该包含一下结构：
+#   1. clip 模型
+#   2. 转换好的优化后的 unet 模型
+#   3. vae 模型
+#   4. scheduler 配置
+
+# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
+lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so"
+model_path = "./models/lyrasd_rev_animated"
+
+# 构建 Img2Img 的 Pipeline
+model = LyraSDImg2ImgPipeline(model_path, lib_path)
+
+# 准备应用的输入和超参数
+prompt = "a cat, cartoon style"
+negative_prompt = "NSFW"
+height, width = 512, 512
+steps = 20
+guidance_scale = 7.5
+generator = torch.Generator().manual_seed(123)
+num_images = 1
+
+# 从 cos 上拿个图作为初始化图片
+init_image_url = "https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/seaside_town.png?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1692601590;1865401590&q-key-time=1692601590;1865401590&q-header-list=&q-url-param-list=&q-signature=ca04ca92d990d94813029c0d9ef29537e5f4637c"
+init_image = BytesIO(requests.get(init_image_url).content)
+init_image = Image.open(init_image).convert('RGB')
+init_image = init_image.resize((width, height), Image.Resampling.LANCZOS)
+
+strength = 0.8
+
+# 推理生成
+images = model(prompt, init_image, strength, steps,
+               guidance_scale, negative_prompt, num_images,
+               generator=generator
+               )
+
+# 存储生成的图片
+for i, image in enumerate(images):
+    image.save(f"./outputs/res_img2img_{i}.png")

lyraSD/__init__.py

@@ -1 +0,0 @@
-from .inference import LyraSD

lyraSD/inference.py

@@ -1,85 +0,0 @@
-import os
-from PIL import Image
-from .muse_trt import TRTStableDiffusionText2ImgPipeline
-from .muse_trt import TRTStableDiffusionImg2ImgPipeline
-import numpy as np
-
-
-class LyraSD(object):
-    def __init__(self, sd_mode, engine_dir,o_height=512, o_width=512, device="cuda:0"):
-        self.sd_mode = sd_mode
-        self.device = device
-        self.o_height = o_height
-        self.o_width = o_width
-        if self.sd_mode == "text2img":
-            self.pipeline = TRTStableDiffusionText2ImgPipeline(
-                engine_dir = engine_dir,
-                o_height = o_height,
-                o_width = o_width,
-                device=device
-            )
-        elif self.sd_mode == "img2img":
-            self.pipeline = TRTStableDiffusionImg2ImgPipeline(
-                engine_dir = engine_dir,
-                o_height = o_height,
-                o_width = o_width,
-                device=device
-            )
-        else:
-            raise ValueError("Invalid sd_mode: {}".format(self.sd_mode))
-
-    
-
-    def inference(self, prompt, 
-                  image=None, 
-                  save_dir="./output", 
-                  save_basename="sd-",
-                  negative_prompts='', 
-                  strength=0.8,
-                  height=None,
-                  width =None,
-                  num_images_per_prompt=1,
-                  num_inference_steps=50,
-                  guidance_scale=7.5,
-                  use_super=False,
-                  ):
-        
-        if self.sd_mode=="text2img" and prompt is None:
-            raise ValueError("prompt must be set on text2img mode")
-        
-        if self.sd_mode=="img2img" and image is None:
-            raise ValueError("image must be set on img2img mode")
-
-        save_basename += f"{self.sd_mode}"
-        if height is None:
-            height = self.o_height
-        if width is None:
-            width = self.o_width
-        
-        if self.sd_mode=="text2img":
-            result_image = self.pipeline(prompt=prompt, negative_prompt=negative_prompts,
-                                    num_inference_steps= num_inference_steps,
-                                    num_images_per_prompt=num_images_per_prompt,
-                                    guidance_scale=guidance_scale,
-                                    height=height,
-                                    width=width,
-                                    use_super=use_super)
-        elif self.sd_mode=="img2img":
-            result_image = self.pipeline(prompt=prompt, 
-                                    image=image,
-                                    negative_prompt=negative_prompts,
-                                    strength = strength,
-                                    num_inference_steps= num_inference_steps,
-                                    num_images_per_prompt=num_images_per_prompt,
-                                    guidance_scale=guidance_scale,
-                                    height=height,
-                                    width=width,
-                                    use_super=use_super)
-
-
-        for i in range(result_image.shape[0]):
-            result_image =  Image.fromarray(result_image[0])
-            result_image = result_image.resize((512, 512))
-            result_image.save(os.path.join(save_dir, save_basename + "-{}.jpg".format(i)))
-                
-        return result_image

lyraSD/muse_trt/__init__.py

@@ -1,10 +0,0 @@
-import ctypes
-import os
-
-current_workdir = os.path.dirname(__file__)
-
-ctypes.cdll.LoadLibrary(os.path.join(current_workdir, "libnvinfer_plugin.so"))
-
-from .sd_img2img import TRTStableDiffusionImg2ImgPipeline
-from .sd_text2img import TRTStableDiffusionText2ImgPipeline
-from .super import SuperX4TRTInfer

lyraSD/muse_trt/models.py

@@ -1,149 +0,0 @@
-r"""models components"""
-from collections import OrderedDict
-from copy import deepcopy
-from typing import Any, Dict, Optional, Union
-
-import numpy as np
-import torch
-from cuda import cudart
-from diffusers import ControlNetModel
-from diffusers.models import AutoencoderKL, UNet2DConditionModel
-from torch import nn
-from torch.nn import functional as F
-from transformers import CLIPTextModel
-
-
-class BaseModel():
-    def __init__(
-        self,
-        local_model_path=None,
-        hf_token=None,
-        text_maxlen=77,
-        embedding_dim=768,
-        fp16=False,
-        device='cuda',
-        verbose=True,
-        max_batch_size=16
-    ):
-        self.fp16 = fp16
-        self.device = device
-        self.verbose = verbose
-        self.hf_token = hf_token
-        self.local_model_path = local_model_path
-
-        # Defaults
-        self.text_maxlen = text_maxlen
-        self.embedding_dim = embedding_dim
-        self.min_batch = 1
-        self.max_batch = max_batch_size
-        self.min_latent_shape = 256 // 8  # min image resolution: 256x256
-        self.max_latent_shape = 1024 // 8  # max image resolution: 1024x1024
-
-    def get_model(self):
-        pass
-    
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        return None
-
-    def check_dims(self, batch_size, image_height, image_width):
-        assert batch_size >= self.min_batch and batch_size <= self.max_batch
-        assert image_height % 8 == 0 or image_width % 8 == 0
-        latent_height = image_height // 8
-        latent_width = image_width // 8
-        assert latent_height >= self.min_latent_shape and latent_height <= self.max_latent_shape
-        assert latent_width >= self.min_latent_shape and latent_width <= self.max_latent_shape
-        return (latent_height, latent_width)
-  
-
-class CLIP(BaseModel):
-    def get_model(self):
-        if self.hf_token is None and self.local_model_path is not None:
-            clip_model = CLIPTextModel.from_pretrained(
-                self.local_model_path, subfolder="text_encoder").to(self.device)
-        else:
-            clip_model = CLIPTextModel.from_pretrained(
-                "openai/clip-vit-large-patch14").to(self.device)
-        return clip_model
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        self.check_dims(batch_size, image_height, image_width)
-        return {
-            'input_ids': (batch_size, self.text_maxlen),
-            'text_embeddings': (batch_size, self.text_maxlen, self.embedding_dim)
-        }
- 
-
-class UNet(BaseModel):
-    def get_model(self):
-        model_opts = {'revision': 'fp16',
-                      'torch_dtype': torch.float16} if self.fp16 else {}
-        print(model_opts)
-        if self.hf_token is None and self.local_model_path is not None:
-            unet_model = UNet2DConditionModel.from_pretrained(
-                self.local_model_path, subfolder="unet",
-                **model_opts
-            ).to(self.device)
-        else:
-            unet_model = UNet2DConditionModel.from_pretrained(
-                "CompVis/stable-diffusion-v1-4",
-                subfolder="unet",
-                use_auth_token=self.hf_token,
-                **model_opts).to(self.device)
-        return unet_model
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'sample': (2*batch_size, 4, latent_height, latent_width),
-            'encoder_hidden_states': (2*batch_size, self.text_maxlen, self.embedding_dim),
-            'latent': (2*batch_size, 4, latent_height, latent_width)
-        }    
-
-class VAEEncoderModule(nn.Module):
-    def __init__(self, local_model_path, device) -> None:
-        super().__init__()
-        self.vae = AutoencoderKL.from_pretrained(
-            local_model_path, subfolder="vae"
-        ).to(device)
-
-    def forward(self, x):
-        h = self.vae.encoder(x)
-        moments = self.vae.quant_conv(h)
-        return moments
-
-
-class VAEEncoder(BaseModel):
-    def get_model(self):
-        vae_encoder = VAEEncoderModule(self.local_model_path, self.device)
-        return vae_encoder
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        image_height, image_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'images': (batch_size, 3, image_height, image_width),
-            'latent': (batch_size, 8, image_height//8, image_width//8),
-        }
-
-
-class VAEDecoder(BaseModel):
-    def get_model(self):
-        if self.hf_token is None and self.local_model_path is not None:
-            vae = AutoencoderKL.from_pretrained(
-                self.local_model_path, subfolder="vae"
-            ).to(self.device)
-        else:
-            vae = AutoencoderKL.from_pretrained(
-                "CompVis/stable-diffusion-v1-4",
-                subfolder="vae",
-                use_auth_token=self.hf_token).to(self.device)
-        vae.forward = vae.decode
-        return vae
-
-    def get_shape_dict(self, batch_size, image_height, image_width):
-        latent_height, latent_width = self.check_dims(
-            batch_size, image_height, image_width)
-        return {
-            'latent': (batch_size, 4, latent_height, latent_width),
-            'images': (batch_size, 3, image_height, image_width)
-        }

lyraSD/muse_trt/sd_img2img.py

@@ -1,368 +0,0 @@
-r"""
-StableDiffusion Img2Img Pipeline by TensorRT.
-It has included SuperResolutionX4 TensorRT Engine. 
-
-Inspired by: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
-            https://developer.nvidia.com/tensorrt
-"""
-
-import inspect
-import os
-from typing import List, Optional, Union
-
-import numpy as np
-import PIL.Image
-import tensorrt as trt
-import torch
-import time
-from diffusers import AutoencoderKL
-from diffusers.schedulers import DPMSolverMultistepScheduler
-from diffusers.models.vae import DiagonalGaussianDistribution
-from diffusers.utils import PIL_INTERPOLATION, randn_tensor
-from polygraphy import cuda
-from transformers import CLIPTokenizer
-
-from .models import CLIP, UNet, VAEDecoder, VAEEncoder
-from .super import SuperX4TRTInfer
-from .utilities import TRT_LOGGER, Engine
-
-
-def preprocess(image):
-    if isinstance(image, torch.Tensor):
-        return image
-    elif isinstance(image, PIL.Image.Image):
-        image = [image]
-
-    if isinstance(image[0], PIL.Image.Image):
-        w, h = image[0].size
-        w, h = map(lambda x: x - x % 8, (w, h))  # resize to integer multiple of 8
-
-        image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image]
-        image = np.concatenate(image, axis=0)
-        image = np.array(image).astype(np.float32) / 255.0
-        image = image.transpose(0, 3, 1, 2)
-        image = 2.0 * image - 1.0
-        image = torch.from_numpy(image)
-    elif isinstance(image[0], torch.Tensor):
-        image = torch.cat(image, dim=0)
-    return image
-
-
-class TRTStableDiffusionImg2ImgPipeline:
-    def __init__(self, engine_dir: str, o_height: int = 1300, o_width: int = 750, device: str = 'cuda:0'):
-        self.device = torch.device(device)
-        super().__init__()
-        self.vae = AutoencoderKL.from_pretrained(
-            os.path.join(engine_dir, 'vae'),
-            torch_dtype=torch.float16
-        ).to(self.device)
-
-        self.tokenizer = CLIPTokenizer.from_pretrained(
-            os.path.join(engine_dir, 'tokenizer')
-        )
-        self.scheduler = DPMSolverMultistepScheduler.from_pretrained(
-            os.path.join(engine_dir, 'scheduler')
-        )
-
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.trt_torch_models_cls = {
-            'clip': CLIP(),
-            'unet_fp16': UNet(),
-            #'vae-encoder': VAEEncoder(),
-            'vae-decoder': VAEDecoder()
-        }
-
-        self.engine = {}
-        # Build engines
-        for model_name, _ in self.trt_torch_models_cls.items():
-            engine = Engine(model_name, engine_dir)
-            self.engine[model_name] = engine
-        # Separate iteration to activate engines
-        for model_name, _ in self.trt_torch_models_cls.items():
-            self.engine[model_name].activate()
-        self.stream = cuda.Stream()
-
-        self.super = SuperX4TRTInfer(
-            engine_dir,
-            model_name='superx4-512-512.plan',
-            fp16=True,
-        )
-
-    def runEngine(self, model_name, feed_dict):
-        engine = self.engine[model_name]
-        return engine.infer(feed_dict, self.stream)
-
-    def _torch_decode_latents(self, latents):
-        latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents).sample
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        image = (image * 255).round()
-        return image
-
-    def _trt_decode_latents(self, latents):
-        latents = 1 / self.vae.config.scaling_factor * latents
-        sample_inp = cuda.DeviceView(
-            ptr=latents.data_ptr(), shape=latents.shape, dtype=np.float32)
-        image = self.runEngine('vae-decoder', {"latent": sample_inp})['images']
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        image = (image * 255).round()
-
-        return image
-
-    def prepare_extra_step_kwargs(self, generator, 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
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def get_timesteps(self, num_inference_steps, strength, device):
-        # get the original timestep using init_timestep
-        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
-
-        t_start = max(num_inference_steps - init_timestep, 0)
-        timesteps = self.scheduler.timesteps[t_start:]
-
-        return timesteps, num_inference_steps - t_start
-
-    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
-        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
-            raise ValueError(
-                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
-            )
-
-        image = image.to(device=device, dtype=dtype)
-
-        batch_size = batch_size * num_images_per_prompt
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if isinstance(generator, list):
-            init_latents = [
-                self.vae.encode(image[i: i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
-            ]
-            init_latents = torch.cat(init_latents, dim=0)
-        else:
-            init_latents = self.vae.encode(image).latent_dist.sample(generator)
-
-        init_latents = self.vae.config.scaling_factor * init_latents
-
-        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
-            raise ValueError(
-                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
-            )
-        else:
-            init_latents = torch.cat([init_latents], dim=0)
-
-        shape = init_latents.shape
-        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-
-        # get latents
-        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
-        latents = init_latents
-
-        return latents
-
-    def _default_height_width(self, height, width, image):
-        if isinstance(image, list):
-            image = image[0]
-
-        if height is None:
-            if isinstance(image, PIL.Image.Image):
-                height = image.height
-            elif isinstance(image, torch.Tensor):
-                height = image.shape[3]
-
-            height = (height // 8) * 8  # round down to nearest multiple of 8
-
-        if width is None:
-            if isinstance(image, PIL.Image.Image):
-                width = image.width
-            elif isinstance(image, torch.Tensor):
-                width = image.shape[2]
-
-            width = (width // 8) * 8  # round down to nearest multiple of 8
-
-        return height, width
-
-    def _trt_encode_prompt(self, prompt, negative_prompt, num_images_per_prompt,):
-        # Tokenize input
-        text_input_ids = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            return_tensors="pt",
-        ).input_ids.type(torch.int32).to(self.device)
-
-        # CLIP text encoder
-        text_input_ids_inp = cuda.DeviceView(
-            ptr=text_input_ids.data_ptr(), shape=text_input_ids.shape, dtype=np.int32
-        )
-        text_embeddings = self.runEngine('clip', {"input_ids": text_input_ids_inp})['text_embeddings']
-
-        # Duplicate text embeddings for each generation per prompt
-        bs_embed, seq_len, _ = text_embeddings.shape
-        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
-        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        max_length = text_input_ids.shape[-1]
-        uncond_input_ids = self.tokenizer(
-            negative_prompt,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt",
-        ).input_ids.type(torch.int32).to(self.device)
-        uncond_input_ids_inp = cuda.DeviceView(
-            ptr=uncond_input_ids.data_ptr(), shape=uncond_input_ids.shape, dtype=np.int32)
-        uncond_embeddings = self.runEngine('clip', {"input_ids": uncond_input_ids_inp})['text_embeddings']
-
-        # Duplicate unconditional embeddings for each generation per prompt
-        seq_len = uncond_embeddings.shape[1]
-        uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)
-        uncond_embeddings = uncond_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
-        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
-        text_embeddings = text_embeddings.to(dtype=torch.float16)
-
-        return text_embeddings
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        image: Union[torch.Tensor, PIL.Image.Image] = None,
-        strength: float = 0.8,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 7.5,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        use_super: bool = True,
-    ):
-        # 1. Default height and width to unet
-        height, width = self._default_height_width(height, width, image)
-
-        # 2. Define call parameters and Allocate the cuda buffers for TRT Engine bindings.
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # Allocate buffers for TensorRT engine bindings
-        for model_name, obj in self.trt_torch_models_cls.items():
-            self.engine[model_name].allocate_buffers(
-                shape_dict=obj.get_shape_dict(batch_size, height, width),
-                device=self.device
-            )
-
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        with trt.Runtime(TRT_LOGGER) as runtime:
-            torch.cuda.synchronize()
-
-            # 3. Encode input prompt. TRT Clip model.
-            prompt_embeds = self._trt_encode_prompt(
-                prompt, negative_prompt, num_images_per_prompt
-            )
-
-            # 4. Prepare mask, image, and controlnet_conditioning_image
-            image = preprocess(image)
-
-            # 5. Prepare timesteps.
-            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
-            timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, self.device)
-            latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
-
-            # 6. Prepare latent variables. It will use VAE-Enoder(currently the encoder is torch model, not trt)
-            latents = self.prepare_latents(
-                image,
-                latent_timestep,
-                batch_size,
-                num_images_per_prompt,
-                prompt_embeds.dtype,
-                self.device,
-                generator,
-            )
-
-            # 7. Prepare extra step kwargs and Set lantens/controlnet_conditioning_image/prompt_embeds to special dtype.
-            #   The dytpe must be equal to the following to ensure that the NAN can not be issued in trt engine.
-            extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-            latents = latents.to(dtype=torch.float32)
-            prompt_embeds = prompt_embeds.to(dtype=torch.float16)
-
-            # 8. Denoising loop
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                # predict the noise residual
-
-                dtype = np.float16
-                if t.dtype != torch.float32:
-                    timestep_float = t.float()
-                else:
-                    timestep_float = t
-
-                sample_inp = cuda.DeviceView(
-                    ptr=latent_model_input.data_ptr(), shape=latent_model_input.shape, dtype=np.float32
-                )
-                timestep_inp = cuda.DeviceView(
-                    ptr=timestep_float.data_ptr(), shape=timestep_float.shape, dtype=np.float32
-                )
-                embeddings_inp = cuda.DeviceView(
-                    ptr=prompt_embeds.data_ptr(), shape=prompt_embeds.shape, dtype=dtype
-                )
-
-                noise_pred = self.engine['unet_fp16'].infer(
-                    {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp},
-                    self.stream)['latent']
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
-
-            # 9. Use VAE-Decoder to decode the latents
-            image = self._trt_decode_latents(latents)
-
-        # 10. SuperX4 Resolution, Optional.
-        if use_super:
-            image = np.ascontiguousarray(np.transpose(image, (0, 3, 1, 2))).astype(np.float16)
-            #image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
-            image = self.super.infer(image)
-            image = np.uint8(np.transpose(image, (0, 2, 3, 1)))
-        else:
-            image = np.uint8(image)
-
-        return image
-

lyraSD/muse_trt/sd_text2img.py

@@ -1,292 +0,0 @@
-r"""
-StableDiffusion Text2Img Pipeline by TensorRT.
-It has included SuperResolutionX4 TensorRT Engine. 
-
-Inspired by: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
-            https://developer.nvidia.com/tensorrt
-"""
-
-import inspect
-import os
-from typing import List, Optional, Union
-
-import numpy as np
-import tensorrt as trt
-import torch
-from diffusers import AutoencoderKL
-from diffusers.schedulers import DPMSolverMultistepScheduler
-from diffusers.utils import PIL_INTERPOLATION, randn_tensor
-from polygraphy import cuda
-from transformers import CLIPTokenizer
-
-from .models import CLIP, UNet, VAEDecoder
-from .super import SuperX4TRTInfer
-from .utilities import TRT_LOGGER, Engine
-
-
-class TRTStableDiffusionText2ImgPipeline:
-    def __init__(self, engine_dir: str, o_height: int = 512, o_width: int = 512, device: str = 'cuda:0'):
-        self.device = torch.device(device)
-        super().__init__()
-        self.vae = AutoencoderKL.from_pretrained(
-            os.path.join(engine_dir, 'vae'),
-            torch_dtype=torch.float16
-        ).to(self.device)
-
-        self.tokenizer = CLIPTokenizer.from_pretrained(
-            os.path.join(engine_dir, 'tokenizer')
-        )
-        self.scheduler = DPMSolverMultistepScheduler.from_pretrained(
-            os.path.join(engine_dir, 'scheduler')
-        )
-
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.trt_torch_models_cls = {
-            'clip': CLIP(),
-            'unet_fp16': UNet(),
-            'vae-decoder': VAEDecoder()
-        }
-
-        self.engine = {}
-        # Build engines
-        for model_name, _ in self.trt_torch_models_cls.items():
-            engine = Engine(model_name, engine_dir)
-            self.engine[model_name] = engine
-        # Separate iteration to activate engines
-        for model_name, _ in self.trt_torch_models_cls.items():
-            self.engine[model_name].activate()
-        self.stream = cuda.Stream()
-
-        self.super = SuperX4TRTInfer(
-            engine_dir=engine_dir,
-            model_name='superx4-512-512.plan',
-            fp16=True
-        )
-
-    def runEngine(self, model_name, feed_dict):
-        engine = self.engine[model_name]
-        return engine.infer(feed_dict, self.stream)
-
-    def _torch_decode_latents(self, latents):
-        latents = 1 / self.vae.config.scaling_factor * latents
-        image = self.vae.decode(latents).sample
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        image = (image * 255).round()
-        return image
-
-    def _trt_decode_latents(self, latents):
-        latents = 1 / self.vae.config.scaling_factor * latents
-        sample_inp = cuda.DeviceView(
-            ptr=latents.data_ptr(), shape=latents.shape, dtype=np.float32)
-        image = self.runEngine('vae-decoder', {"latent": sample_inp})['images']
-        image = (image / 2 + 0.5).clamp(0, 1)
-        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
-        image = (image * 255).round()
-
-        return image
-
-    def prepare_extra_step_kwargs(self, generator, 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
-
-        # check if the scheduler accepts generator
-        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
-        if accepts_generator:
-            extra_step_kwargs["generator"] = generator
-        return extra_step_kwargs
-
-    def get_timesteps(self, num_inference_steps, strength, device):
-        # get the original timestep using init_timestep
-        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
-
-        t_start = max(num_inference_steps - init_timestep, 0)
-        timesteps = self.scheduler.timesteps[t_start:]
-
-        return timesteps, num_inference_steps - t_start
-
-    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
-        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        else:
-            latents = latents.to(device)
-
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    def _trt_encode_prompt(self, prompt, negative_prompt, num_images_per_prompt,):
-        # Tokenize input
-        text_input_ids = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length,
-            return_tensors="pt",
-        ).input_ids.type(torch.int32).to(self.device)
-
-        # CLIP text encoder
-        text_input_ids_inp = cuda.DeviceView(
-            ptr=text_input_ids.data_ptr(), shape=text_input_ids.shape, dtype=np.int32
-        )
-        text_embeddings = self.runEngine('clip', {"input_ids": text_input_ids_inp})['text_embeddings']
-
-        # Duplicate text embeddings for each generation per prompt
-        bs_embed, seq_len, _ = text_embeddings.shape
-        text_embeddings = text_embeddings.repeat(1, num_images_per_prompt, 1)
-        text_embeddings = text_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        max_length = text_input_ids.shape[-1]
-        uncond_input_ids = self.tokenizer(
-            negative_prompt,
-            padding="max_length",
-            max_length=max_length,
-            truncation=True,
-            return_tensors="pt",
-        ).input_ids.type(torch.int32).to(self.device)
-        uncond_input_ids_inp = cuda.DeviceView(
-            ptr=uncond_input_ids.data_ptr(), shape=uncond_input_ids.shape, dtype=np.int32)
-        uncond_embeddings = self.runEngine('clip', {"input_ids": uncond_input_ids_inp})['text_embeddings']
-
-        # Duplicate unconditional embeddings for each generation per prompt
-        seq_len = uncond_embeddings.shape[1]
-        uncond_embeddings = uncond_embeddings.repeat(1, num_images_per_prompt, 1)
-        uncond_embeddings = uncond_embeddings.view(bs_embed * num_images_per_prompt, seq_len, -1)
-
-        # Concatenate the unconditional and text embeddings into a single batch to avoid doing two forward passes for classifier free guidance
-        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
-        text_embeddings = text_embeddings.to(dtype=torch.float16)
-
-        return text_embeddings
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 7.5,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        num_images_per_prompt: Optional[int] = 1,
-        eta: float = 0.0,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        use_super: bool = True,
-    ):
-        # 1. Default height and width to unet
-        assert height is not None, "height can not be None"
-        assert width is not None, "width can not be None"
-
-        # 2. Define call parameters and Allocate the cuda buffers for TRT Engine bindings.
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        # Allocate buffers for TensorRT engine bindings
-        for model_name, obj in self.trt_torch_models_cls.items():
-            self.engine[model_name].allocate_buffers(
-                shape_dict=obj.get_shape_dict(batch_size, height, width),
-                device=self.device
-            )
-
-        do_classifier_free_guidance = guidance_scale > 1.0
-
-        with trt.Runtime(TRT_LOGGER) as runtime:
-            torch.cuda.synchronize()
-
-            # 3. Encode input prompt. TRT Clip model.
-            prompt_embeds = self._trt_encode_prompt(
-                prompt, negative_prompt, num_images_per_prompt
-            )
-
-            # 4. Prepare timesteps.
-            self.scheduler.set_timesteps(num_inference_steps, device=self.device)
-            timesteps = self.scheduler.timesteps
-
-            # 5. Prepare latent variables. It will use VAE-Enoder(currently the encoder is torch model, not trt)
-            num_channels_latents = 4
-            latents = self.prepare_latents(
-                batch_size*num_images_per_prompt,
-                num_channels_latents,
-                height,
-                width,
-                prompt_embeds.dtype,
-                self.device,
-                generator,
-                latents
-            )
-
-            # 6. Prepare extra step kwargs and Set lantens/controlnet_conditioning_image/prompt_embeds to special dtype.
-            #   The dytpe must be equal to the following to ensure that the NAN can not be issued in trt engine.
-            extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
-            latents = latents.to(dtype=torch.float32)
-            prompt_embeds = prompt_embeds.to(dtype=torch.float16)
-
-            # 7. Denoising loop
-            for i, t in enumerate(timesteps):
-                # expand the latents if we are doing classifier free guidance
-                latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
-
-                latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
-
-                # predict the noise residual
-
-                dtype = np.float16
-                if t.dtype != torch.float32:
-                    timestep_float = t.float()
-                else:
-                    timestep_float = t
-
-                sample_inp = cuda.DeviceView(
-                    ptr=latent_model_input.data_ptr(), shape=latent_model_input.shape, dtype=np.float32
-                )
-                timestep_inp = cuda.DeviceView(
-                    ptr=timestep_float.data_ptr(), shape=timestep_float.shape, dtype=np.float32
-                )
-                embeddings_inp = cuda.DeviceView(
-                    ptr=prompt_embeds.data_ptr(), shape=prompt_embeds.shape, dtype=dtype
-                )
-
-                noise_pred = self.engine['unet_fp16'].infer(
-                    {"sample": sample_inp, "timestep": timestep_inp, "encoder_hidden_states": embeddings_inp},
-                    self.stream)['latent']
-                # perform guidance
-                if do_classifier_free_guidance:
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
-
-            # 8. Use VAE-Decoder to decode the latents
-            image = self._trt_decode_latents(latents)
-
-        # 9. SuperX4 Resolution, Optional.
-        if use_super:
-            image = np.ascontiguousarray(np.transpose(image, (0, 3, 1, 2))).astype(np.float16)
-            #image = self.super.infer(np.transpose(image.astype(np.float16), (0, 3, 1, 2)))
-            image = self.super.infer(image)
-            image = np.uint8(np.transpose(image, (0, 2, 3, 1)))
-        else:
-            image = np.uint8(image)
-        return image

lyraSD/muse_trt/super.py

@@ -1,64 +0,0 @@
-r"""use tensorrt engine to infer, a useful pipeline"""
-
-import os
-
-import numpy as np
-from polygraphy import cuda
-from polygraphy.backend.common import bytes_from_path
-from polygraphy.backend.trt import engine_from_bytes
-
-
-class SuperX4TRTInfer:
-    def __init__(self, engine_dir,
-                 model_name='superx4.plan',
-                 o_height=None,
-                 o_width=None,
-                 fp16=True,
-                 ) -> None:
-        engine_path = os.path.join(engine_dir, model_name)
-        self.engine = engine_from_bytes(bytes_from_path(engine_path))
-        self.context = self.engine.create_execution_context()
-
-        self.o_height = o_height
-        self.o_width = o_width
-        self.fp = fp16
-        self.dtype = np.float16 if fp16 else np.float32
-
-        self.stream = cuda.Stream()
-
-    def infer(self, x):
-        batch_size, channel, height, width = x.shape
-        if self.o_height is None or self.o_width is None:
-            o_height = height*4
-            o_width = width*4
-        else:
-            o_height = self.o_height
-            o_width = self.o_width
-
-        h_output = np.empty([batch_size, channel, o_height, o_width], dtype=self.dtype)
-
-        # allocate device memory
-        d_input = cuda.wrapper().malloc(1 * x.nbytes)
-        d_output = cuda.wrapper().malloc(1*h_output.nbytes)
-
-        bindings = [int(d_input), int(d_output)]
-
-        # transfer input data to device
-        cuda.wrapper().memcpy(d_input, x.ctypes.data, x.nbytes, cuda.MemcpyKind.HostToDevice, self.stream.ptr)
-
-        # execute model
-        noerror = self.context.execute_async_v2(bindings, self.stream.ptr)
-        if not noerror:
-            raise ValueError(f"ERROR: inference failed.")
-
-        # transfer predictions back
-        cuda.wrapper().memcpy(h_output.ctypes.data, d_output, h_output.nbytes, cuda.MemcpyKind.DeviceToHost, self.stream.ptr)
-        cuda.wrapper().free(d_input)
-        cuda.wrapper().free(d_output)
-
-        return h_output
-
-    def teardown(self):
-        del self.engine
-        self.stream.free()
-        del self.stream

lyraSD/muse_trt/utilities.py

@@ -1,538 +0,0 @@
-r"""utils components"""
-
-from collections import OrderedDict
-from copy import copy
-import numpy as np
-import os
-import math
-from PIL import Image
-from polygraphy.backend.common import bytes_from_path
-from polygraphy.backend.trt import CreateConfig, Profile
-from polygraphy.backend.trt import engine_from_bytes, engine_from_network, network_from_onnx_path, save_engine
-from polygraphy.backend.trt import util as trt_util
-from polygraphy import cuda
-import random
-from scipy import integrate
-import tensorrt as trt
-import torch
-
-TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
-
-
-class Engine():
-    def __init__(
-        self,
-        model_name,
-        engine_dir,
-        memory_pool_size=None
-    ):
-        self.engine_path = os.path.join(engine_dir, model_name+'.plan')
-        self.engine = None
-        self.context = None
-        self.buffers = OrderedDict()
-        self.tensors = OrderedDict()
-        self.memory_pool_size = memory_pool_size
-
-    def __del__(self):
-        [buf.free() for buf in self.buffers.values() if isinstance(buf, cuda.DeviceArray)]
-        del self.engine
-        del self.context
-        del self.buffers
-        del self.tensors
-
-    def build(self, onnx_path, fp16, input_profile=None, enable_preview=False):
-        print(f"Building TensorRT engine for {onnx_path}: {self.engine_path}")
-        p = Profile()
-        if input_profile:
-            for name, dims in input_profile.items():
-                assert len(dims) == 3
-                p.add(name, min=dims[0], opt=dims[1], max=dims[2])
-
-        preview_features = []
-        if enable_preview:
-            trt_version = [int(i) for i in trt.__version__.split(".")]
-            # FASTER_DYNAMIC_SHAPES_0805 should only be used for TRT 8.5.1 or above.
-            if trt_version[0] > 8 or \
-                    (trt_version[0] == 8 and (trt_version[1] > 5 or (trt_version[1] == 5 and trt_version[2] >= 1))):
-                preview_features = [trt.PreviewFeature.FASTER_DYNAMIC_SHAPES_0805]
-
-        if self.memory_pool_size is not None:
-            memory_pool_limits = {trt.MemoryPoolType.WORKSPACE: (self.memory_pool_size*(2 ** 30))}
-            print(memory_pool_limits)
-        else:
-            memory_pool_limits = None
-        engine = engine_from_network(
-            network_from_onnx_path(onnx_path),
-            config=CreateConfig(
-                fp16=fp16, profiles=[p], preview_features=preview_features, memory_pool_limits=memory_pool_limits
-            )
-        )
-        save_engine(engine, path=self.engine_path)
-
-    def activate(self):
-        print(f"Loading TensorRT engine: {self.engine_path}")
-        self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
-        self.context = self.engine.create_execution_context()
-
-    def allocate_buffers(self, shape_dict=None, device='cuda'):
-        for idx in range(trt_util.get_bindings_per_profile(self.engine)):
-            binding = self.engine[idx]
-            if shape_dict and binding in shape_dict:
-                shape = shape_dict[binding]
-            else:
-                shape = self.engine.get_binding_shape(binding)
-            dtype = trt_util.np_dtype_from_trt(self.engine.get_binding_dtype(binding))
-            if self.engine.binding_is_input(binding):
-                self.context.set_binding_shape(idx, shape)
-            # Workaround to convert np dtype to torch
-            np_type_tensor = np.empty(shape=[], dtype=dtype)
-            torch_type_tensor = torch.from_numpy(np_type_tensor)
-            tensor = torch.empty(tuple(shape), dtype=torch_type_tensor.dtype).to(device=device)
-            self.tensors[binding] = tensor
-            self.buffers[binding] = cuda.DeviceView(ptr=tensor.data_ptr(), shape=shape, dtype=dtype)
-
-    def infer(self, feed_dict, stream):
-        start_binding, end_binding = trt_util.get_active_profile_bindings(self.context)
-        # shallow copy of ordered dict
-        device_buffers = copy(self.buffers)
-        for name, buf in feed_dict.items():
-            assert isinstance(buf, cuda.DeviceView)
-            device_buffers[name] = buf
-        bindings = [0] * start_binding + [buf.ptr for buf in device_buffers.values()]
-        noerror = self.context.execute_async_v2(bindings=bindings, stream_handle=stream.ptr)
-        if not noerror:
-            raise ValueError(f"ERROR: inference failed.")
-
-        return self.tensors
-
-
-class LMSDiscreteScheduler():
-    def __init__(
-        self,
-        device='cuda',
-        beta_start=0.00085,
-        beta_end=0.012,
-        num_train_timesteps=1000,
-    ):
-        self.num_train_timesteps = num_train_timesteps
-        self.order = 4
-
-        self.beta_start = beta_start
-        self.beta_end = beta_end
-        betas = (torch.linspace(beta_start**0.5, beta_end**0.5, self.num_train_timesteps, dtype=torch.float32) ** 2)
-        alphas = 1.0 - betas
-        self.alphas_cumprod = torch.cumprod(alphas, dim=0)
-
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
-        sigmas = np.concatenate([sigmas[::-1], [0.0]]).astype(np.float32)
-        self.sigmas = torch.from_numpy(sigmas)
-
-        # standard deviation of the initial noise distribution
-        self.init_noise_sigma = self.sigmas.max()
-
-        self.device = device
-
-    def set_timesteps(self, steps):
-        self.num_inference_steps = steps
-
-        timesteps = np.linspace(0, self.num_train_timesteps - 1, steps, dtype=float)[::-1].copy()
-        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
-        sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
-        sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32)
-        self.sigmas = torch.from_numpy(sigmas).to(device=self.device)
-
-        # Move all timesteps to correct device beforehand
-        self.timesteps = torch.from_numpy(timesteps).to(device=self.device).float()
-        self.derivatives = []
-
-    def scale_model_input(self, sample: torch.FloatTensor, idx, *args, **kwargs) -> torch.FloatTensor:
-        return sample * self.latent_scales[idx]
-
-    def configure(self):
-        order = self.order
-        self.lms_coeffs = []
-        self.latent_scales = [1./((sigma**2 + 1) ** 0.5) for sigma in self.sigmas]
-
-        def get_lms_coefficient(order, t, current_order):
-            """
-            Compute a linear multistep coefficient.
-            """
-            def lms_derivative(tau):
-                prod = 1.0
-                for k in range(order):
-                    if current_order == k:
-                        continue
-                    prod *= (tau - self.sigmas[t - k]) / (self.sigmas[t - current_order] - self.sigmas[t - k])
-                return prod
-            integrated_coeff = integrate.quad(lms_derivative, self.sigmas[t], self.sigmas[t + 1], epsrel=1e-4)[0]
-            return integrated_coeff
-
-        for step_index in range(self.num_inference_steps):
-            order = min(step_index + 1, order)
-            self.lms_coeffs.append([get_lms_coefficient(order, step_index, curr_order) for curr_order in range(order)])
-
-    def step(self, output, latents, idx, timestep):
-        # compute the previous noisy sample x_t -> x_t-1
-        # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise
-        sigma = self.sigmas[idx]
-        pred_original_sample = latents - sigma * output
-        # 2. Convert to an ODE derivative
-        derivative = (latents - pred_original_sample) / sigma
-        self.derivatives.append(derivative)
-        if len(self.derivatives) > self.order:
-            self.derivatives.pop(0)
-        # 3. Compute previous sample based on the derivatives path
-        prev_sample = latents + sum(
-            coeff * derivative for coeff, derivative in zip(self.lms_coeffs[idx], reversed(self.derivatives))
-        )
-
-        return prev_sample
-
-
-class DPMScheduler():
-    def __init__(
-        self,
-        beta_start=0.00085,
-        beta_end=0.012,
-        num_train_timesteps=1000,
-        solver_order=2,
-        predict_epsilon=True,
-        thresholding=False,
-        dynamic_thresholding_ratio=0.995,
-        sample_max_value=1.0,
-        algorithm_type="dpmsolver++",
-        solver_type="midpoint",
-        lower_order_final=True,
-        device='cuda',
-    ):
-        # this schedule is very specific to the latent diffusion model.
-        self.betas = (
-            torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
-        )
-
-        self.device = device
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
-        # Currently we only support VP-type noise schedule
-        self.alpha_t = torch.sqrt(self.alphas_cumprod)
-        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
-        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
-
-        # standard deviation of the initial noise distribution
-        self.init_noise_sigma = 1.0
-
-        self.algorithm_type = algorithm_type
-        self.predict_epsilon = predict_epsilon
-        self.thresholding = thresholding
-        self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
-        self.sample_max_value = sample_max_value
-        self.lower_order_final = lower_order_final
-
-        # settings for DPM-Solver
-        if algorithm_type not in ["dpmsolver", "dpmsolver++"]:
-            raise NotImplementedError(f"{algorithm_type} does is not implemented for {self.__class__}")
-        if solver_type not in ["midpoint", "heun"]:
-            raise NotImplementedError(f"{solver_type} does is not implemented for {self.__class__}")
-
-        # setable values
-        self.num_inference_steps = None
-        self.solver_order = solver_order
-        self.num_train_timesteps = num_train_timesteps
-        self.solver_type = solver_type
-
-        self.first_order_first_coef = []
-        self.first_order_second_coef = []
-
-        self.second_order_first_coef = []
-        self.second_order_second_coef = []
-        self.second_order_third_coef = []
-
-        self.third_order_first_coef = []
-        self.third_order_second_coef = []
-        self.third_order_third_coef = []
-        self.third_order_fourth_coef = []
-
-    def scale_model_input(self, sample: torch.FloatTensor, *args, **kwargs) -> torch.FloatTensor:
-        return sample
-
-    def configure(self):
-        lower_order_nums = 0
-        for step_index in range(self.num_inference_steps):
-            step_idx = step_index
-            timestep = self.timesteps[step_idx]
-
-            prev_timestep = 0 if step_idx == len(self.timesteps) - 1 else self.timesteps[step_idx + 1]
-
-            self.dpm_solver_first_order_coefs_precompute(timestep, prev_timestep)
-
-            timestep_list = [self.timesteps[step_index - 1], timestep]
-            self.multistep_dpm_solver_second_order_coefs_precompute(timestep_list, prev_timestep)
-
-            timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep]
-            self.multistep_dpm_solver_third_order_coefs_precompute(timestep_list, prev_timestep)
-
-            if lower_order_nums < self.solver_order:
-                lower_order_nums += 1
-
-    def dpm_solver_first_order_coefs_precompute(self, timestep, prev_timestep):
-        lambda_t, lambda_s = self.lambda_t[prev_timestep], self.lambda_t[timestep]
-        alpha_t, alpha_s = self.alpha_t[prev_timestep], self.alpha_t[timestep]
-        sigma_t, sigma_s = self.sigma_t[prev_timestep], self.sigma_t[timestep]
-        h = lambda_t - lambda_s
-        if self.algorithm_type == "dpmsolver++":
-            self.first_order_first_coef.append(sigma_t / sigma_s)
-            self.first_order_second_coef.append(alpha_t * (torch.exp(-h) - 1.0))
-        elif self.algorithm_type == "dpmsolver":
-            self.first_order_first_coef.append(alpha_t / alpha_s)
-            self.first_order_second_coef.append(sigma_t * (torch.exp(h) - 1.0))
-
-    def multistep_dpm_solver_second_order_coefs_precompute(self, timestep_list, prev_timestep):
-        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
-        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
-        h = lambda_t - lambda_s0
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2211.01095 for detailed derivations
-            if self.solver_type == "midpoint":
-                self.second_order_first_coef.append(sigma_t / sigma_s0)
-                self.second_order_second_coef.append((alpha_t * (torch.exp(-h) - 1.0)))
-                self.second_order_third_coef.append(0.5 * (alpha_t * (torch.exp(-h) - 1.0)))
-            elif self.solver_type == "heun":
-                self.second_order_first_coef.append(sigma_t / sigma_s0)
-                self.second_order_second_coef.append((alpha_t * (torch.exp(-h) - 1.0)))
-                self.second_order_third_coef.append(alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0))
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            if self.solver_type == "midpoint":
-                self.second_order_first_coef.append(alpha_t / alpha_s0)
-                self.second_order_second_coef.append((sigma_t * (torch.exp(h) - 1.0)))
-                self.second_order_third_coef.append(0.5 * (sigma_t * (torch.exp(h) - 1.0)))
-            elif self.solver_type == "heun":
-                self.second_order_first_coef.append(alpha_t / alpha_s0)
-                self.second_order_second_coef.append((sigma_t * (torch.exp(h) - 1.0)))
-                self.second_order_third_coef.append((sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)))
-
-    def multistep_dpm_solver_third_order_coefs_precompute(self, timestep_list, prev_timestep):
-        t, s0 = prev_timestep, timestep_list[-1]
-        lambda_t, lambda_s0 = (
-            self.lambda_t[t],
-            self.lambda_t[s0]
-        )
-        alpha_t, alpha_s0 = self.alpha_t[t], self.alpha_t[s0]
-        sigma_t, sigma_s0 = self.sigma_t[t], self.sigma_t[s0]
-        h = lambda_t - lambda_s0
-        if self.algorithm_type == "dpmsolver++":
-            self.third_order_first_coef.append(sigma_t / sigma_s0)
-            self.third_order_second_coef.append(alpha_t * (torch.exp(-h) - 1.0))
-            self.third_order_third_coef.append(alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0))
-            self.third_order_fourth_coef.append(alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5))
-        elif self.algorithm_type == "dpmsolver":
-            self.third_order_first_coef.append(alpha_t / alpha_s0)
-            self.third_order_second_coef.append(sigma_t * (torch.exp(h) - 1.0))
-            self.third_order_third_coef.append(sigma_t * ((torch.exp(h) - 1.0) / h - 1.0))
-            self.third_order_fourth_coef.append(sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5))
-
-    def set_timesteps(self, num_inference_steps):
-        self.num_inference_steps = num_inference_steps
-        timesteps = (
-            np.linspace(0, self.num_train_timesteps - 1, num_inference_steps + 1)
-            .round()[::-1][:-1]
-            .copy()
-            .astype(np.int32)
-        )
-        self.timesteps = torch.from_numpy(timesteps).to(self.device)
-        self.model_outputs = [
-            None,
-        ] * self.solver_order
-        self.lower_order_nums = 0
-
-    def convert_model_output(
-        self, model_output, timestep, sample
-    ):
-        # DPM-Solver++ needs to solve an integral of the data prediction model.
-        if self.algorithm_type == "dpmsolver++":
-            if self.predict_epsilon:
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
-                x0_pred = (sample - sigma_t * model_output) / alpha_t
-            else:
-                x0_pred = model_output
-            if self.thresholding:
-                # Dynamic thresholding in https://arxiv.org/abs/2205.11487
-                dynamic_max_val = torch.quantile(
-                    torch.abs(x0_pred).reshape((x0_pred.shape[0], -1)), self.dynamic_thresholding_ratio, dim=1
-                )
-                dynamic_max_val = torch.maximum(
-                    dynamic_max_val,
-                    self.sample_max_value * torch.ones_like(dynamic_max_val).to(dynamic_max_val.device),
-                )[(...,) + (None,) * (x0_pred.ndim - 1)]
-                x0_pred = torch.clamp(x0_pred, -dynamic_max_val, dynamic_max_val) / dynamic_max_val
-            return x0_pred
-        # DPM-Solver needs to solve an integral of the noise prediction model.
-        elif self.algorithm_type == "dpmsolver":
-            if self.predict_epsilon:
-                return model_output
-            else:
-                alpha_t, sigma_t = self.alpha_t[timestep], self.sigma_t[timestep]
-                epsilon = (sample - alpha_t * model_output) / sigma_t
-                return epsilon
-
-    def dpm_solver_first_order_update(
-        self,
-        idx,
-        model_output,
-        sample
-    ):
-        first_coef = self.first_order_first_coef[idx]
-        second_coef = self.first_order_second_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            x_t = first_coef * sample - second_coef * model_output
-        elif self.algorithm_type == "dpmsolver":
-            x_t = first_coef * sample - second_coef * model_output
-        return x_t
-
-    def multistep_dpm_solver_second_order_update(
-        self,
-        idx,
-        model_output_list,
-        timestep_list,
-        prev_timestep,
-        sample
-    ):
-        t, s0, s1 = prev_timestep, timestep_list[-1], timestep_list[-2]
-        m0, m1 = model_output_list[-1], model_output_list[-2]
-        lambda_t, lambda_s0, lambda_s1 = self.lambda_t[t], self.lambda_t[s0], self.lambda_t[s1]
-        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
-        r0 = h_0 / h
-        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
-
-        first_coef = self.second_order_first_coef[idx]
-        second_coef = self.second_order_second_coef[idx]
-        third_coef = self.second_order_third_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2211.01095 for detailed derivations
-            if self.solver_type == "midpoint":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-            elif self.solver_type == "heun":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    + third_coef * D1
-                )
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            if self.solver_type == "midpoint":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-            elif self.solver_type == "heun":
-                x_t = (
-                    first_coef * sample
-                    - second_coef * D0
-                    - third_coef * D1
-                )
-        return x_t
-
-    def multistep_dpm_solver_third_order_update(
-        self,
-        idx,
-        model_output_list,
-        timestep_list,
-        prev_timestep,
-        sample
-    ):
-        t, s0, s1, s2 = prev_timestep, timestep_list[-1], timestep_list[-2], timestep_list[-3]
-        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
-        lambda_t, lambda_s0, lambda_s1, lambda_s2 = (
-            self.lambda_t[t],
-            self.lambda_t[s0],
-            self.lambda_t[s1],
-            self.lambda_t[s2],
-        )
-        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
-        r0, r1 = h_0 / h, h_1 / h
-        D0 = m0
-        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
-        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
-        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
-
-        first_coef = self.third_order_first_coef[idx]
-        second_coef = self.third_order_second_coef[idx]
-        third_coef = self.third_order_third_coef[idx]
-        fourth_coef = self.third_order_fourth_coef[idx]
-
-        if self.algorithm_type == "dpmsolver++":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            x_t = (
-                first_coef * sample
-                - second_coef * D0
-                + third_coef * D1
-                - fourth_coef * D2
-            )
-        elif self.algorithm_type == "dpmsolver":
-            # See https://arxiv.org/abs/2206.00927 for detailed derivations
-            x_t = (
-                first_coef * sample
-                - second_coef * D0
-                - third_coef * D1
-                - fourth_coef * D2
-            )
-        return x_t
-
-    def step(self, output, latents, step_index, timestep):
-        if self.num_inference_steps is None:
-            raise ValueError(
-                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
-            )
-
-        prev_timestep = 0 if step_index == len(self.timesteps) - 1 else self.timesteps[step_index + 1]
-        lower_order_final = (
-            (step_index == len(self.timesteps) - 1) and self.lower_order_final and len(self.timesteps) < 15
-        )
-        lower_order_second = (
-            (step_index == len(self.timesteps) - 2) and self.lower_order_final and len(self.timesteps) < 15
-        )
-
-        output = self.convert_model_output(output, timestep, latents)
-        for i in range(self.solver_order - 1):
-            self.model_outputs[i] = self.model_outputs[i + 1]
-        self.model_outputs[-1] = output
-
-        if self.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
-            prev_sample = self.dpm_solver_first_order_update(step_index, output, latents)
-        elif self.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
-            timestep_list = [self.timesteps[step_index - 1], timestep]
-            prev_sample = self.multistep_dpm_solver_second_order_update(
-                step_index, self.model_outputs, timestep_list, prev_timestep, latents
-            )
-        else:
-            timestep_list = [self.timesteps[step_index - 2], self.timesteps[step_index - 1], timestep]
-            prev_sample = self.multistep_dpm_solver_third_order_update(
-                step_index, self.model_outputs, timestep_list, prev_timestep, latents
-            )
-
-        if self.lower_order_nums < self.solver_order:
-            self.lower_order_nums += 1
-
-        return prev_sample
-
-
-def save_image(images, image_path_dir, image_name_prefix):
-    """
-    Save the generated images to png files.
-    """
-    images = ((images + 1) * 255 / 2).clamp(0, 255).detach().permute(0, 2, 3, 1).round().type(torch.uint8).cpu().numpy()
-    for i in range(images.shape[0]):
-        image_path = os.path.join(image_path_dir, image_name_prefix+str(i+1)+'-'+str(random.randint(1000, 9999))+'.png')
-        print(f"Saving image {i+1} / {images.shape[0]} to: {image_path}")
-        Image.fromarray(images[i]).save(image_path)
-
-

lyrasd_model/__init__.py

@@ -0,0 +1,5 @@
+from . import lyrasd_img2img_pipeline, lyrasd_txt2img_pipeline, lyrasd_controlnet_txt2img_pipeline, lyrasd_controlnet_img2img_pipeline
+from .lyrasd_txt2img_pipeline import LyraSdTxt2ImgPipeline
+from .lyrasd_img2img_pipeline import LyraSDImg2ImgPipeline
+from .lyrasd_controlnet_txt2img_pipeline import LyraSdControlnetTxt2ImgPipeline
+from .lyrasd_controlnet_img2img_pipeline import LyraSdControlnetImg2ImgPipeline

lyrasd_model/lora_util.py

@@ -0,0 +1,54 @@
+import os
+import torch
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+import numpy as np
+
+def add_text_lora_layer(clip_model, lora_model_path="Misaka.safetensors", alpha=1.0, lora_file_format="fp32", device="cuda:0"):
+    if lora_file_format == "fp32":
+        model_dtype = np.float32
+    elif lora_file_format == "fp16":
+        model_dtype = np.float16
+    else:
+        raise Exception(f"unsupported model dtype: {lora_file_format}")
+    all_files = os.scandir(lora_model_path)
+    unload_dict = []
+    # directly update weight in diffusers model
+    for file in all_files:
+            
+        if 'text' in file.name:
+            layer_infos = file.name.split('.')[0].split('text_model_')[-1].split('_')
+            curr_layer = clip_model.text_model
+        else:
+            continue
+
+        # find the target layer
+        temp_name = layer_infos.pop(0)
+        while len(layer_infos) > -1:
+            try:
+                curr_layer = curr_layer.__getattr__(temp_name)
+                if len(layer_infos) > 0:
+                    temp_name = layer_infos.pop(0)
+                    # if temp_name == "self":
+                    #     temp_name += "_" + layer_infos.pop(0)
+                    # elif temp_name != "mlp" and len(layer_infos) == 1:
+                    #     temp_name += "_" + layer_infos.pop(0)
+                elif len(layer_infos) == 0:
+                    break
+            except Exception:
+                if len(temp_name) > 0:
+                    temp_name += '_'+layer_infos.pop(0)
+                else:
+                    temp_name = layer_infos.pop(0)
+        data = torch.from_numpy(np.fromfile(file.path, dtype=model_dtype)).to(clip_model.dtype).to(clip_model.device).reshape(curr_layer.weight.data.shape) 
+        if len(curr_layer.weight.data) == 4:
+            adding_weight = alpha * data.permute(0,3,1,2)
+        else:
+            adding_weight = alpha * data
+        curr_layer.weight.data += adding_weight
+
+        curr_layer_unload_data = {
+            "layer": curr_layer,
+            "added_weight": adding_weight
+        }
+        unload_dict.append(curr_layer_unload_data)
+    return unload_dict

lyrasd_model/lyrasd_controlnet_img2img_pipeline.py

@@ -0,0 +1,637 @@
+import torch
+from typing import Any, Callable, Dict, List, Optional, Union
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.utils import randn_tensor, logging
+from diffusers.schedulers import EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler
+from diffusers.utils import PIL_INTERPOLATION
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+import os
+import numpy as np
+import warnings
+from .lora_util import add_text_lora_layer
+import gc
+
+from PIL import Image
+import PIL
+
+import inspect
+
+import time
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def preprocess(image):
+    warnings.warn(
+        "The preprocess method is deprecated and will be removed in a future version. Please"
+        " use VaeImageProcessor.preprocess instead",
+        FutureWarning,
+    )
+    if isinstance(image, torch.Tensor):
+        return image
+    elif isinstance(image, PIL.Image.Image):
+        image = [image]
+
+    if isinstance(image[0], PIL.Image.Image):
+        w, h = image[0].size
+        w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 8
+
+        image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image]
+        image = np.concatenate(image, axis=0)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image.transpose(0, 3, 1, 2)
+        image = 2.0 * image - 1.0
+        image = torch.from_numpy(image)
+    elif isinstance(image[0], torch.Tensor):
+        image = torch.cat(image, dim=0)
+    return image
+
+class LyraSdControlnetImg2ImgPipeline(TextualInversionLoaderMixin):
+    def __init__(self, model_path, lib_so_path, model_dtype='fp32', device=torch.device("cuda"), dtype=torch.float16) -> None:
+        self.device = device
+        self.dtype = dtype
+
+        torch.classes.load_library(lib_so_path)
+
+        self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(dtype).to(device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+        self.text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="text_encoder").to(dtype).to(device)
+
+        self.unet_in_channels = 4
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.vae.enable_tiling()
+        self.unet = torch.classes.lyrasd.Unet2dConditionalModelOp(
+            3,     # max num of controlnets
+            "fp16" # inference dtype (can only use fp16 for now)
+        )
+
+        unet_path = os.path.join(model_path, "unet_bins/")
+
+        self.reload_unet_model(unet_path, model_dtype)
+
+        self.scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
+
+    def load_controlnet_model(self, model_name, controlnet_path, model_dtype="fp32"):
+        if len(controlnet_path) > 0 and controlnet_path[-1] != "/":
+            controlnet_path = controlnet_path + "/"
+        self.unet.load_controlnet_model(model_name, controlnet_path, model_dtype)
+
+    def unload_controlnet_model(self, model_name):
+        self.unet.unload_controlnet_model(model_name, True)
+
+    def get_loaded_controlnet(self):
+        return self.unet.get_loaded_controlnet()
+
+    def reload_unet_model(self, unet_path, unet_file_format='fp32'):
+        if len(unet_path) > 0 and unet_path[-1] != "/":
+            unet_path = unet_path + "/"
+        return self.unet.reload_unet_model(unet_path, unet_file_format)
+
+    def load_lora(self, lora_model_path, lora_name, lora_strength, lora_file_format='fp32'):
+        if len(lora_model_path) > 0 and lora_model_path[-1] != "/":
+            lora_model_path = lora_model_path + "/"
+        lora = add_text_lora_layer(self.text_encoder, lora_model_path, lora_strength, lora_file_format)
+        self.loaded_lora[lora_name] = lora
+        self.unet.load_lora(lora_model_path, lora_name, lora_strength, lora_file_format)
+
+    def unload_lora(self, lora_name, clean_cache=False):
+        for layer_data in self.loaded_lora[lora_name]:
+            layer = layer_data['layer']
+            added_weight = layer_data['added_weight']
+            layer.weight.data -= added_weight
+        self.unet.unload_lora(lora_name, clean_cache)
+        del self.loaded_lora[lora_name]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def clean_lora_cache(self):
+        self.unet.clean_lora_cache()
+    
+    def get_loaded_lora(self):
+        return self.unet.get_loaded_lora()
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.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" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) 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 isinstance(negative_prompt, str):
+                uncond_tokens = [negative_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
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # 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 = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+    
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 64 != 0 or width % 64 != 0: # 初版暂时只支持 64 的倍数的 height 和 width
+            raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.")
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (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 negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            elif isinstance(generator, list):
+                init_latents = [
+                    self.vae.encode(image[i: i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = self.vae.encode(image).latent_dist.sample(generator)
+
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            # expand init_latents for batch_size
+            deprecation_message = (
+                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"
+                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+                " your script to pass as many initial images as text prompts to suppress this warning."
+            )
+            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        shape = init_latents.shape
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        # get latents
+        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+        latents = init_latents
+
+        return latents
+
+    def prepare_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance=False,
+        guess_mode=False,
+    ):
+        if not isinstance(image, torch.Tensor):
+            if isinstance(image, PIL.Image.Image):
+                image = [image]
+
+            if isinstance(image[0], PIL.Image.Image):
+                images = []
+
+                for image_ in image:
+                    image_ = image_.convert("RGB")
+                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])
+                    image_ = np.array(image_)
+                    image_ = image_[None, :]
+                    images.append(image_)
+
+                image = images
+
+                image = np.concatenate(image, axis=0)
+                image = np.array(image).astype(np.float32) / 255.0
+                image = torch.from_numpy(image)
+            elif isinstance(image[0], torch.Tensor):
+                image = torch.cat(image, dim=0)
+
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        if do_classifier_free_guidance and not guess_mode:
+            image = torch.cat([image] * 2)
+
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, 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
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def get_timesteps(self, num_inference_steps, strength, device):
+        # get the original timestep using init_timestep
+        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+        t_start = max(num_inference_steps - init_timestep, 0)
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order:]
+
+        return timesteps, num_inference_steps - t_start
+
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        image: Union[
+            torch.FloatTensor,
+            PIL.Image.Image,
+            np.ndarray,
+            List[torch.FloatTensor],
+            List[PIL.Image.Image],
+            List[np.ndarray],
+        ] = None,
+        strength: float = 0.8,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        controlnet_images: Optional[List[PIL.Image.Image]] = None,
+        controlnet_scale: Optional[List[float]] = None,
+        controlnet_names: Optional[List[str]] = None,
+        guess_mode = False,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+
+        """
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self.device
+
+        # 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
+        start = time.perf_counter()
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+        control_images = []
+
+        # 4 prepare controlnet images
+        for image_ in controlnet_images:
+            image_ = self.prepare_image(
+                image=image_,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=prompt_embeds.dtype,
+                do_classifier_free_guidance=do_classifier_free_guidance
+            )
+
+            control_images.append(image_)
+
+        control_scales = []
+
+        scales = [1.0, ] * 13
+        if guess_mode:
+            scales = torch.logspace(-1, 0, 13).tolist()
+        
+        for scale in controlnet_scale:
+            scales_ = [d * scale for d in scales]
+            control_scales.append(scales_)
+
+        image = preprocess(image)
+        
+        # 5. set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        # 6. Prepare latent variables
+        latents = self.prepare_latents(
+            image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator
+        )
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        start_unet = time.perf_counter()
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = latent_model_input.permute(0, 2, 3, 1).contiguous()
+
+            # 后边三个 None 是给到controlnet 的参数，暂时给到 None 当 placeholder
+            noise_pred = self.unet.forward(latent_model_input, prompt_embeds, t, controlnet_names, control_images, control_scales, guess_mode)
+
+            noise_pred = noise_pred.permute(0, 3, 1, 2)
+            # perform guidance
+
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+        torch.cuda.synchronize()
+
+        start = time.perf_counter()
+        image = self.decode_latents(latents)
+        torch.cuda.synchronize()
+        image = numpy_to_pil(image)
+
+        return image
+        # return None

lyrasd_model/lyrasd_controlnet_txt2img_pipeline.py

@@ -0,0 +1,547 @@
+import torch
+from typing import Any, Callable, Dict, List, Optional, Union
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.utils import randn_tensor, logging
+from diffusers.schedulers import EulerAncestralDiscreteScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler
+from diffusers.utils import PIL_INTERPOLATION
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+import os
+import numpy as np
+from .lora_util import add_text_lora_layer
+import gc
+from PIL import Image
+import PIL
+
+import inspect
+
+import time
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+class LyraSdControlnetTxt2ImgPipeline(TextualInversionLoaderMixin):
+    def __init__(self, model_path, lib_so_path, model_dtype='fp32', device=torch.device("cuda"), dtype=torch.float16) -> None:
+        self.device = device
+        self.dtype = dtype
+
+        torch.classes.load_library(lib_so_path)
+
+        self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(dtype).to(device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+        self.text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="text_encoder").to(dtype).to(device)
+        self.unet_in_channels = 4
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.vae.enable_tiling()
+        self.unet = torch.classes.lyrasd.Unet2dConditionalModelOp(
+            3,     # max num of controlnets
+            "fp16" # inference dtype (can only use fp16 for now)
+        )
+
+        unet_path = os.path.join(model_path, "unet_bins/")
+        self.reload_unet_model(unet_path, model_dtype)
+
+        self.scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
+
+    def load_controlnet_model(self, model_name, controlnet_path, model_dtype="fp32"):
+        if len(controlnet_path) > 0 and controlnet_path[-1] != "/":
+            controlnet_path = controlnet_path + "/"
+        self.unet.load_controlnet_model(model_name, controlnet_path, model_dtype)
+
+    def unload_controlnet_model(self, model_name):
+        self.unet.unload_controlnet_model(model_name, True)
+
+    def get_loaded_controlnet(self):
+        return self.unet.get_loaded_controlnet()
+
+    def reload_unet_model(self, unet_path, unet_file_format='fp32'):
+        if len(unet_path) > 0 and unet_path[-1] != "/":
+            unet_path = unet_path + "/"
+        return self.unet.reload_unet_model(unet_path, unet_file_format)
+
+    def load_lora(self, lora_model_path, lora_name, lora_strength, lora_file_format='fp32'):
+        if len(lora_model_path) > 0 and lora_model_path[-1] != "/":
+            lora_model_path = lora_model_path + "/"
+        lora = add_text_lora_layer(self.text_encoder, lora_model_path, lora_strength, lora_file_format)
+        self.loaded_lora[lora_name] = lora
+        self.unet.load_lora(lora_model_path, lora_name, lora_strength, lora_file_format)
+
+    def unload_lora(self, lora_name, clean_cache=False):
+        for layer_data in self.loaded_lora[lora_name]:
+            layer = layer_data['layer']
+            added_weight = layer_data['added_weight']
+            layer.weight.data -= added_weight
+        self.unet.unload_lora(lora_name, clean_cache)
+        del self.loaded_lora[lora_name]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def clean_lora_cache(self):
+        self.unet.clean_lora_cache()
+    
+    def get_loaded_lora(self):
+        return self.unet.get_loaded_lora()
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.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" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) 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 isinstance(negative_prompt, str):
+                uncond_tokens = [negative_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
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # 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 = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+    
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 64 != 0 or width % 64 != 0: # 初版暂时只支持 64 的倍数的 height 和 width
+            raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.")
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (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 negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def prepare_image(
+        self,
+        image,
+        width,
+        height,
+        batch_size,
+        num_images_per_prompt,
+        device,
+        dtype,
+        do_classifier_free_guidance=False,
+        guess_mode=False,
+    ):
+        if not isinstance(image, torch.Tensor):
+            if isinstance(image, PIL.Image.Image):
+                image = [image]
+
+            if isinstance(image[0], PIL.Image.Image):
+                images = []
+
+                for image_ in image:
+                    image_ = image_.convert("RGB")
+                    image_ = image_.resize((width, height), resample=PIL_INTERPOLATION["lanczos"])
+                    image_ = np.array(image_)
+                    image_ = image_[None, :]
+                    images.append(image_)
+
+                image = images
+
+                image = np.concatenate(image, axis=0)
+                image = np.array(image).astype(np.float32) / 255.0
+                image = torch.from_numpy(image)
+            elif isinstance(image[0], torch.Tensor):
+                image = torch.cat(image, dim=0)
+
+        image_batch_size = image.shape[0]
+
+        if image_batch_size == 1:
+            repeat_by = batch_size
+        else:
+            # image batch size is the same as prompt batch size
+            repeat_by = num_images_per_prompt
+
+        image = image.repeat_interleave(repeat_by, dim=0)
+
+        image = image.to(device=device, dtype=dtype)
+
+        if do_classifier_free_guidance and not guess_mode:
+            image = torch.cat([image] * 2)
+
+        return image
+
+    def prepare_extra_step_kwargs(self, generator, 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
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        controlnet_images: Optional[List[PIL.Image.Image]] = None,
+        controlnet_scale: Optional[List[float]] = None,
+        controlnet_names: Optional[List[str]] = None,
+        guess_mode = False,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+
+        """
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self.device
+
+        # 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
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+        control_images = []
+
+        for image_ in controlnet_images:
+            image_ = self.prepare_image(
+                image=image_,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=prompt_embeds.dtype,
+                do_classifier_free_guidance=do_classifier_free_guidance
+            )
+
+            control_images.append(image_)
+
+        control_scales = []
+
+        scales = [1.0, ] * 13
+        if guess_mode:
+            scales = torch.logspace(-1, 0, 13).tolist()
+        
+        for scale in controlnet_scale:
+            scales_ = [d * scale for d in scales]
+            control_scales.append(scales_)
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet_in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        start_unet = time.perf_counter()
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = latent_model_input.permute(0, 2, 3, 1).contiguous()
+
+            # 后边三个 None 是给到controlnet 的参数，暂时给到 None 当 placeholder
+            noise_pred = self.unet.forward(latent_model_input, prompt_embeds, t, controlnet_names, control_images, control_scales, guess_mode)
+
+            noise_pred = noise_pred.permute(0, 3, 1, 2)
+            # perform guidance
+
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+        image = self.decode_latents(latents)
+        image = numpy_to_pil(image)
+
+        return image
+        # return None

lyrasd_model/lyrasd_img2img_pipeline.py

@@ -0,0 +1,554 @@
+import inspect
+import logging
+import os
+import warnings
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import PIL
+import torch
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import EulerAncestralDiscreteScheduler
+from diffusers.utils import PIL_INTERPOLATION, deprecate, logging, randn_tensor
+from PIL import Image
+from transformers import CLIPTextModel, CLIPTokenizer
+from .lora_util import add_text_lora_layer
+import gc
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def preprocess(image):
+    warnings.warn(
+        "The preprocess method is deprecated and will be removed in a future version. Please"
+        " use VaeImageProcessor.preprocess instead",
+        FutureWarning,
+    )
+    if isinstance(image, torch.Tensor):
+        return image
+    elif isinstance(image, PIL.Image.Image):
+        image = [image]
+
+    if isinstance(image[0], PIL.Image.Image):
+        w, h = image[0].size
+        w, h = (x - x % 8 for x in (w, h))  # resize to integer multiple of 8
+
+        image = [np.array(i.resize((w, h), resample=PIL_INTERPOLATION["lanczos"]))[None, :] for i in image]
+        image = np.concatenate(image, axis=0)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image.transpose(0, 3, 1, 2)
+        image = 2.0 * image - 1.0
+        image = torch.from_numpy(image)
+    elif isinstance(image[0], torch.Tensor):
+        image = torch.cat(image, dim=0)
+    return image
+
+
+class LyraSDImg2ImgPipeline(TextualInversionLoaderMixin):
+    def __init__(self, model_path, lib_so_path, model_dtype='fp32', device=torch.device("cuda"), dtype=torch.float16) -> None:
+        self.device = device
+        self.dtype = dtype
+
+        torch.classes.load_library(lib_so_path)
+
+        self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(dtype).to(device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+        self.text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="text_encoder").to(dtype).to(device)
+        unet_path = os.path.join(model_path, "unet_bins/")
+
+        self.unet_in_channels = 4
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.vae.enable_tiling()
+        self.unet = torch.classes.lyrasd.Unet2dConditionalModelOp(
+            3,     # max num of controlnets
+            "fp16" # inference dtype (can only use fp16 for now)
+        )
+
+        self.reload_unet_model(unet_path, model_dtype)
+
+        self.scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
+
+    def reload_unet_model(self, unet_path, unet_file_format='fp32'):
+        if len(unet_path) > 0 and unet_path[-1] != "/":
+            unet_path = unet_path + "/"
+        return self.unet.reload_unet_model(unet_path, unet_file_format)
+
+    def load_lora(self, lora_model_path, lora_name, lora_strength, lora_file_format='fp32'):
+        if len(lora_model_path) > 0 and lora_model_path[-1] != "/":
+            lora_model_path = lora_model_path + "/"
+        lora = add_text_lora_layer(self.text_encoder, lora_model_path, lora_strength, lora_file_format)
+        self.loaded_lora[lora_name] = lora
+        self.unet.load_lora(lora_model_path, lora_name, lora_strength, lora_file_format)
+
+    def unload_lora(self, lora_name, clean_cache=False):
+        for layer_data in self.loaded_lora[lora_name]:
+            layer = layer_data['layer']
+            added_weight = layer_data['added_weight']
+            layer.weight.data -= added_weight
+        self.unet.unload_lora(lora_name, clean_cache)
+        del self.loaded_lora[lora_name]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def clean_lora_cache(self):
+        self.unet.clean_lora_cache()
+    
+    def get_loaded_lora(self):
+        return self.unet.get_loaded_lora()
+
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._encode_prompt
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            """
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.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" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        if self.text_encoder is not None:
+            prompt_embeds_dtype = self.text_encoder.dtype
+        elif self.unet is not None:
+            prompt_embeds_dtype = self.unet.dtype
+        else:
+            prompt_embeds_dtype = prompt_embeds.dtype
+
+        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif prompt is not None and type(prompt) 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 isinstance(negative_prompt, str):
+                uncond_tokens = [negative_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
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # 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 = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.decode_latents
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+    def prepare_extra_step_kwargs(self, generator, 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
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(
+        self, prompt, strength, callback_steps, negative_prompt=None, prompt_embeds=None, negative_prompt_embeds=None
+    ):
+        if strength < 0 or strength > 1:
+            raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
+
+        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)}."
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (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 negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def get_timesteps(self, num_inference_steps, strength, device):
+        # get the original timestep using init_timestep
+        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+
+        t_start = max(num_inference_steps - init_timestep, 0)
+        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order:]
+
+        return timesteps, num_inference_steps - t_start
+
+    def prepare_latents(self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
+        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
+            raise ValueError(
+                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+            )
+
+        image = image.to(device=device, dtype=dtype)
+
+        batch_size = batch_size * num_images_per_prompt
+
+        if image.shape[1] == 4:
+            init_latents = image
+
+        else:
+            if isinstance(generator, list) and len(generator) != batch_size:
+                raise ValueError(
+                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+                )
+
+            elif isinstance(generator, list):
+                init_latents = [
+                    self.vae.encode(image[i: i + 1]).latent_dist.sample(generator[i]) for i in range(batch_size)
+                ]
+                init_latents = torch.cat(init_latents, dim=0)
+            else:
+                init_latents = self.vae.encode(image).latent_dist.sample(generator)
+
+            init_latents = self.vae.config.scaling_factor * init_latents
+
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            # expand init_latents for batch_size
+            deprecation_message = (
+                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"
+                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
+                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
+                " your script to pass as many initial images as text prompts to suppress this warning."
+            )
+            deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents], dim=0)
+
+        shape = init_latents.shape
+        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+
+        # get latents
+        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+        latents = init_latents
+
+        return latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        image: Union[
+            torch.FloatTensor,
+            PIL.Image.Image,
+            np.ndarray,
+            List[torch.FloatTensor],
+            List[PIL.Image.Image],
+            List[np.ndarray],
+        ] = None,
+        strength: float = 0.8,
+        num_inference_steps: Optional[int] = 50,
+        guidance_scale: Optional[float] = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: Optional[float] = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+    ):
+        r"""
+        The call function to the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide image generation. If not defined, you need to pass `prompt_embeds`.
+            image (`torch.FloatTensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.FloatTensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+                `Image` or tensor representing an image batch to be used as the starting point. Can also accept image
+                latents as `image`, but if passing latents directly it is not encoded again.
+            strength (`float`, *optional*, defaults to 0.8):
+                Indicates extent to transform the reference `image`. Must be between 0 and 1. `image` is used as a
+                starting point and more noise is added the higher the `strength`. The number of denoising steps depends
+                on the amount of noise initially added. When `strength` is 1, added noise is maximum and the denoising
+                process runs for the full number of iterations specified in `num_inference_steps`. A value of 1
+                essentially ignores `image`.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference. This parameter is modulated by `strength`.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                A higher guidance scale value encourages the model to generate images closely linked to the text
+                `prompt` at the expense of lower image quality. Guidance scale is enabled when `guidance_scale > 1`.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide what to not include in image generation. If not defined, you need to
+                pass `negative_prompt_embeds` instead. Ignored when not using guidance (`guidance_scale < 1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) from the [DDIM](https://arxiv.org/abs/2010.02502) paper. Only applies
+                to the [`~schedulers.DDIMScheduler`], and is ignored in other schedulers.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                A [`torch.Generator`](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make
+                generation deterministic.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs (prompt weighting). If not
+                provided, text embeddings are generated from the `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs (prompt weighting). If
+                not provided, `negative_prompt_embeds` are generated from the `negative_prompt` input argument.
+            callback (`Callable`, *optional*):
+                A function that calls every `callback_steps` steps during inference. The function is called with the
+                following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
+            callback_steps (`int`, *optional*, defaults to 1):
+                The frequency at which the `callback` function is called. If not specified, the callback is called at
+                every step.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+                If `return_dict` is `True`, [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] is returned,
+                otherwise a `tuple` is returned where the first element is a list with the generated images and the
+                second element is a list of `bool`s indicating whether the corresponding generated image contains
+                "not-safe-for-work" (nsfw) content.
+        """
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(prompt, strength, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds)
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self.device
+        # 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
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        # 4. Preprocess image
+        image = preprocess(image)
+
+        # 5. set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+        # 6. Prepare latent variables
+        latents = self.prepare_latents(
+            image, latent_timestep, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator
+        )
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = latent_model_input.permute(0, 2, 3, 1).contiguous()
+
+            # predict the noise residual
+            # 后边4个 None 是给到controlnet 的参数，暂时给到 None 当 placeholder
+            noise_pred = self.unet.forward(latent_model_input, prompt_embeds, t, None, None, None, None)
+
+            noise_pred = noise_pred.permute(0, 3, 1, 2)
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+        image = self.decode_latents(latents)
+        image = numpy_to_pil(image)
+
+        return image

lyraSD/muse_trt/libnvinfer_plugin.so → lyrasd_model/lyrasd_lib/libth_lyrasd_cu11_sm80.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:53cbcc8a47524652bb8e0399a2fbbcfc0b785f11bdc491bbb6a71e4b888ee124
-size 85198184
+oid sha256:0689ed5d3b55f5033a8869d5f23ce900793aa0ab7fdc4a3e3c0a0f3a243c83da
+size 65441456

sd1.4-engine/superx4-512-512.plan → lyrasd_model/lyrasd_lib/libth_lyrasd_cu11_sm86.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b4c37b1fa1a9966104975398cf7b4b4a3ed722335ccde47145c9eb8316030797
-size 10211800
+oid sha256:b8e27e715fa3a17ce25bf23b772e0dd355d0780c1bd93cfeeb12ef45b0ba2444
+size 65389176

sd1.4-engine/clip.plan → lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm80.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dc71dac2ad328d4e49d99fad6bbdcb1ce201d2cd376da34f55565c6a60188d1a
-size 247251489
+oid sha256:c2eaa9067ad8eb1d20872afa71ed9497f62d930819704d15e5e8bf559623eca7
+size 65498752

sd1.4-engine/vae-decoder.plan → lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4fd093053ffd431b8331ac11e179a31122b8636c5f2f18b8ae9614b714706d4b
-size 100138684
+oid sha256:7d0c909ff2498934c6d1ed8f46af6cdc7812872177c0a4e7ca0ee99bf88fcb65
+size 65519232

lyrasd_model/lyrasd_txt2img_pipeline.py

@@ -0,0 +1,458 @@
+import inspect
+import os
+import time
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL
+from diffusers.schedulers import (DPMSolverMultistepScheduler,
+                                  EulerAncestralDiscreteScheduler,
+                                  EulerDiscreteScheduler,
+                                  KarrasDiffusionSchedulers)
+from diffusers.utils import logging, randn_tensor
+from PIL import Image
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+import gc
+import numpy as np
+
+from .lora_util import add_text_lora_layer
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+class LyraSdTxt2ImgPipeline(TextualInversionLoaderMixin):
+    def __init__(self, model_path, lib_so_path, model_dtype="fp32", device=torch.device("cuda"), dtype=torch.float16) -> None:
+        self.device = device
+        self.dtype = dtype
+
+        torch.classes.load_library(lib_so_path)
+
+        self.vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae").to(dtype).to(device)
+        self.tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+        self.text_encoder = CLIPTextModel.from_pretrained(model_path, subfolder="text_encoder").to(dtype).to(device)
+        unet_path = os.path.join(model_path, "unet_bins/")
+
+        self.unet_in_channels = 4
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.vae.enable_tiling()
+        self.unet = torch.classes.lyrasd.Unet2dConditionalModelOp(
+            3,     # max num of controlnets
+            "fp16" # inference dtype (can only use fp16 for now)
+        )
+
+        unet_path = os.path.join(model_path, "unet_bins/")
+
+        self.reload_unet_model(unet_path, model_dtype)
+
+        self.scheduler = EulerAncestralDiscreteScheduler.from_pretrained(model_path, subfolder="scheduler")
+
+        self.loaded_lora = {}
+
+    def reload_unet_model(self, unet_path, unet_file_format='fp32'):
+        if len(unet_path) > 0 and unet_path[-1] != "/":
+            unet_path = unet_path + "/"
+        return self.unet.reload_unet_model(unet_path, unet_file_format)
+
+    def load_lora(self, lora_model_path, lora_name, lora_strength, lora_file_format='fp32'):
+        if len(lora_model_path) > 0 and lora_model_path[-1] != "/":
+            lora_model_path = lora_model_path + "/"
+        lora = add_text_lora_layer(self.text_encoder, lora_model_path, lora_strength, lora_file_format)
+        self.loaded_lora[lora_name] = lora
+        self.unet.load_lora(lora_model_path, lora_name, lora_strength, lora_file_format)
+
+    def unload_lora(self, lora_name, clean_cache=False):
+        for layer_data in self.loaded_lora[lora_name]:
+            layer = layer_data['layer']
+            added_weight = layer_data['added_weight']
+            layer.weight.data -= added_weight
+        self.unet.unload_lora(lora_name, clean_cache)
+        del self.loaded_lora[lora_name]
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def clean_lora_cache(self):
+        self.unet.clean_lora_cache()
+    
+    def get_loaded_lora(self):
+        return self.unet.get_loaded_lora()
+
+    def _encode_prompt(
+        self,
+        prompt,
+        device,
+        num_images_per_prompt,
+        do_classifier_free_guidance,
+        negative_prompt=None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Encodes the prompt into text encoder hidden states.
+
+        Args:
+             prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            do_classifier_free_guidance (`bool`):
+                whether to use classifier free guidance or not
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+        """
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+            text_inputs = self.tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input_ids = text_inputs.input_ids
+            untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+            if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+                text_input_ids, untruncated_ids
+            ):
+                removed_text = self.tokenizer.batch_decode(
+                    untruncated_ids[:, self.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" {self.tokenizer.model_max_length} tokens: {removed_text}"
+                )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = text_inputs.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            prompt_embeds = self.text_encoder(
+                text_input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            prompt_embeds = prompt_embeds[0]
+
+        prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+        bs_embed, seq_len, _ = prompt_embeds.shape
+        # duplicate text embeddings for each generation per prompt, using mps friendly method
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+        # get unconditional embeddings for classifier free guidance
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+            uncond_tokens: List[str]
+            if negative_prompt is None:
+                uncond_tokens = [""] * batch_size
+            elif type(prompt) 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 isinstance(negative_prompt, str):
+                uncond_tokens = [negative_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
+
+            # textual inversion: procecss multi-vector tokens if necessary
+            if isinstance(self, TextualInversionLoaderMixin):
+                uncond_tokens = self.maybe_convert_prompt(uncond_tokens, self.tokenizer)
+
+            max_length = prompt_embeds.shape[1]
+            uncond_input = self.tokenizer(
+                uncond_tokens,
+                padding="max_length",
+                max_length=max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+
+            if hasattr(self.text_encoder.config, "use_attention_mask") and self.text_encoder.config.use_attention_mask:
+                attention_mask = uncond_input.attention_mask.to(device)
+            else:
+                attention_mask = None
+
+            negative_prompt_embeds = self.text_encoder(
+                uncond_input.input_ids.to(device),
+                attention_mask=attention_mask,
+            )
+            negative_prompt_embeds = negative_prompt_embeds[0]
+
+        if do_classifier_free_guidance:
+            # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+            seq_len = negative_prompt_embeds.shape[1]
+
+            negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+            negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+            negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+            # 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 = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        return prompt_embeds
+
+    def decode_latents(self, latents):
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 64 != 0 or width % 64 != 0:  # 初版暂时只支持 64 的倍数的 height 和 width
+            raise ValueError(f"`height` and `width` have to be divisible by 64 but are {height} and {width}.")
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (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 negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and negative_prompt_embeds is not None:
+            if prompt_embeds.shape != negative_prompt_embeds.shape:
+                raise ValueError(
+                    "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+                    f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+                    f" {negative_prompt_embeds.shape}."
+                )
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def prepare_extra_step_kwargs(self, generator, 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
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            guidance_scale (`float`, *optional*, defaults to 7.5):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+                less than `1`).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            eta (`float`, *optional*, defaults to 0.0):
+                Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+                [`schedulers.DDIMScheduler`], will be ignored for others.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+
+        """
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, negative_prompt, prompt_embeds, negative_prompt_embeds
+        )
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self.device
+        # 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
+        prompt_embeds = self._encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+        )
+
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet_in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+            latent_model_input = latent_model_input.permute(0, 2, 3, 1).contiguous()
+
+            # 后边4个 None 是给到controlnet 的参数，暂时给到 None 当 placeholder
+            noise_pred = self.unet.forward(latent_model_input, prompt_embeds, t, None, None, None, None)
+
+            noise_pred = noise_pred.permute(0, 3, 1, 2)
+            # perform guidance
+
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                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(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+        image = self.decode_latents(latents)
+        image = numpy_to_pil(image)
+
+        return image

models/README.md

@@ -0,0 +1,12 @@
+# Models
+### This is the place where you should download the checkpoints, and unzip them
+
+```bash
+wget -O lyrasd_rev_animated.tar.gz "https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/lyrasd/lyrasd_rev_animated.tar.gz?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1694078210;1866878210&q-key-time=1694078210;1866878210&q-header-list=&q-url-param-list=&q-signature=6046546135631dee9e8be7d8e061a77e8790e675"
+wget -O lyrasd_canny.tar.gz "https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/lyrasd/lyrasd_canny.tar.gz?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1694078194;1866878194&q-key-time=1694078194;1866878194&q-header-list=&q-url-param-list=&q-signature=efb713ee650a0ee3c954fb3a0e148c37ef13cd3b"
+wget -O lyrasd_xiaorenshu_lora.tar.gz "https://chuangxin-research-1258344705.cos.ap-guangzhou.myqcloud.com/share/files/lyrasd/lyrasd_xiaorenshu_lora.tar.gz?q-sign-algorithm=sha1&q-ak=AKIDBF6i7GCtKWS8ZkgOtACzX3MQDl37xYty&q-sign-time=1694078234;1866878234&q-key-time=1694078234;1866878234&q-header-list=&q-url-param-list=&q-signature=fb9a577a54ea6dedd9be696e40b96b71a1b23b5d"
+
+tar -xvf lyrasd_rev_animated.tar.gz
+tar -xvf lyrasd_canny.tar.gz
+tar -xvf lyrasd_xiaorenshu_lora.tar.gz
+```

requirements.txt

@@ -0,0 +1,2 @@
+diffusers
+transformers

sd1.4-engine/feature_extractor/preprocessor_config.json

@@ -1,28 +0,0 @@
-{
-  "crop_size": {
-    "height": 224,
-    "width": 224
-  },
-  "do_center_crop": true,
-  "do_convert_rgb": true,
-  "do_normalize": true,
-  "do_rescale": true,
-  "do_resize": true,
-  "feature_extractor_type": "CLIPFeatureExtractor",
-  "image_mean": [
-    0.48145466,
-    0.4578275,
-    0.40821073
-  ],
-  "image_processor_type": "CLIPImageProcessor",
-  "image_std": [
-    0.26862954,
-    0.26130258,
-    0.27577711
-  ],
-  "resample": 3,
-  "rescale_factor": 0.00392156862745098,
-  "size": {
-    "shortest_edge": 224
-  }
-}

sd1.4-engine/scheduler/scheduler_config.json

@@ -1,14 +0,0 @@
-{
-  "_class_name": "PNDMScheduler",
-  "_diffusers_version": "0.14.0",
-  "beta_end": 0.012,
-  "beta_schedule": "scaled_linear",
-  "beta_start": 0.00085,
-  "clip_sample": false,
-  "num_train_timesteps": 1000,
-  "prediction_type": "epsilon",
-  "set_alpha_to_one": false,
-  "skip_prk_steps": true,
-  "steps_offset": 1,
-  "trained_betas": null
-}

sd1.4-engine/text_encoder/config.json

@@ -1,25 +0,0 @@
-{
-  "_name_or_path": "openai/clip-vit-large-patch14",
-  "architectures": [
-    "CLIPTextModel"
-  ],
-  "attention_dropout": 0.0,
-  "bos_token_id": 0,
-  "dropout": 0.0,
-  "eos_token_id": 2,
-  "hidden_act": "quick_gelu",
-  "hidden_size": 768,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 3072,
-  "layer_norm_eps": 1e-05,
-  "max_position_embeddings": 77,
-  "model_type": "clip_text_model",
-  "num_attention_heads": 12,
-  "num_hidden_layers": 12,
-  "pad_token_id": 1,
-  "projection_dim": 768,
-  "torch_dtype": "float32",
-  "transformers_version": "4.25.1",
-  "vocab_size": 49408
-}

sd1.4-engine/tokenizer/special_tokens_map.json

@@ -1,24 +0,0 @@
-{
-  "bos_token": {
-    "content": "<|startoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": "<|endoftext|>",
-  "unk_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  }
-}

sd1.4-engine/tokenizer/tokenizer_config.json

@@ -1,34 +0,0 @@
-{
-  "add_prefix_space": false,
-  "bos_token": {
-    "__type": "AddedToken",
-    "content": "<|startoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "do_lower_case": true,
-  "eos_token": {
-    "__type": "AddedToken",
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "errors": "replace",
-  "model_max_length": 77,
-  "name_or_path": "openai/clip-vit-large-patch14",
-  "pad_token": "<|endoftext|>",
-  "special_tokens_map_file": "./special_tokens_map.json",
-  "tokenizer_class": "CLIPTokenizer",
-  "unk_token": {
-    "__type": "AddedToken",
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  }
-}

sd1.4-engine/unet_fp16.plan

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:0373e858a69bb44fe0f21c4990de3ae18d415b4d99ca44e4809ea48fc3482e5a
-size 1725864976

sd1.4-engine/vae/config.json

@@ -1,31 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.14.0",
-  "_name_or_path": "stabilityai/sd-vae-ft-mse",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "in_channels": 3,
-  "latent_channels": 4,
-  "layers_per_block": 2,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 256,
-  "scaling_factor": 0.18215,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ]
-}

sd1.4-engine/vae/diffusion_pytorch_model.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a4302e1efa25f3a47ceb7536bc335715ad9d1f203e90c2d25507600d74006e89
-size 334715313

txt2img_demo.py

@@ -0,0 +1,44 @@
+import torch
+import time
+
+from lyrasd_model import LyraSdTxt2ImgPipeline
+
+# 存放模型文件的路径，应该包含一下结构：
+#   1. clip 模型
+#   2. 转换好的优化后的 unet 模型，放入其中的 unet_bins 文件夹
+#   3. vae 模型
+#   4. scheduler 配置
+
+# LyraSD 的 C++ 编译动态链接库，其中包含 C++ CUDA 计算的细节
+lib_path = "./lyrasd_model/lyrasd_lib/libth_lyrasd_cu12_sm86.so"
+model_path = "./models/lyrasd_rev_animated"
+lora_path = "./models/lyrasd_xiaorenshu_lora"
+
+# 构建 Txt2Img 的 Pipeline
+model = LyraSdTxt2ImgPipeline(model_path, lib_path)
+
+# load lora
+# 参数分别为 lora 存放位置，名字，lora 强度，lora模型精度
+model.load_lora(lora_path, "xiaorenshu", 0.4, "fp32")
+
+# 准备应用的输入和超参数
+prompt = "a cat, cute, cartoon, concise, traditional, chinese painting, Tang and Song Dynasties, masterpiece, 4k, 8k, UHD, best quality"
+negative_prompt = "(((horrible))), (((scary))), (((naked))), (((large breasts))), high saturation, colorful, human:2, body:2, low quality, bad quality, lowres, out of frame, duplicate, watermark, signature, text, frames, cut, cropped, malformed limbs, extra limbs, (((missing arms))), (((missing legs)))"
+height, width = 512, 512
+steps = 30
+guidance_scale = 7
+generator = torch.Generator().manual_seed(123)
+num_images = 1
+
+start = time.perf_counter()
+# 推理生成
+images = model(prompt, height, width, steps,
+        guidance_scale, negative_prompt, num_images,
+        generator=generator)
+print("image gen cost: ",time.perf_counter() - start)
+# 存储生成的图片
+for i, image in enumerate(images):
+    image.save(f"outputs/res_txt2img_lora_{i}.png")
+
+# unload lora，参数为 lora 的名字，是否清除 lora 缓存
+# model.unload_lora("xiaorenshu", True)