Migrated from GitHub

.gitattributes

@@ -33,3 +33,11 @@ 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
+img/overview.png filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/img2img/GIOApFAWDOc_160.wav filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/img2img/_GI7meqlYZk_30.wav filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/inpainting/14ekd4nkpwc_28.wav filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/inpainting/3ek-xLwr05Q_30.wav filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/inpainting/9z8XIRyUq9Q_30.wav filter=lfs diff=lfs merge=lfs -text
+notebooks/examples/inpainting/IvfaKPDWC00_160.wav filter=lfs diff=lfs merge=lfs -text
+prompt2prompt/Roboto-Regular.ttf filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,437 @@
+Attribution-NonCommercial-ShareAlike 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International Public License
+("Public License"). To the extent this Public License may be
+interpreted as a contract, You are granted the Licensed Rights in
+consideration of Your acceptance of these terms and conditions, and the
+Licensor grants You such rights in consideration of benefits the
+Licensor receives from making the Licensed Material available under
+these terms and conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. BY-NC-SA Compatible License means a license listed at
+     creativecommons.org/compatiblelicenses, approved by Creative
+     Commons as essentially the equivalent of this Public License.
+
+  d. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+
+  e. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  f. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  g. License Elements means the license attributes listed in the name
+     of a Creative Commons Public License. The License Elements of this
+     Public License are Attribution, NonCommercial, and ShareAlike.
+
+  h. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  i. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  j. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  k. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  l. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  m. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  n. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. Additional offer from the Licensor -- Adapted Material.
+               Every recipient of Adapted Material from You
+               automatically receives an offer from the Licensor to
+               exercise the Licensed Rights in the Adapted Material
+               under the conditions of the Adapter's License You apply.
+
+            c. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+  b. ShareAlike.
+
+     In addition to the conditions in Section 3(a), if You Share
+     Adapted Material You produce, the following conditions also apply.
+
+       1. The Adapter's License You apply must be a Creative Commons
+          license with the same License Elements, this version or
+          later, or a BY-NC-SA Compatible License.
+
+       2. You must include the text of, or the URI or hyperlink to, the
+          Adapter's License You apply. You may satisfy this condition
+          in any reasonable manner based on the medium, means, and
+          context in which You Share Adapted Material.
+
+       3. You may not offer or impose any additional or different terms
+          or conditions on, or apply any Effective Technological
+          Measures to, Adapted Material that restrict exercise of the
+          rights granted under the Adapter's License You apply.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material,
+     including for purposes of Section 3(b); and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the 鈥淟icensor.鈥� The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

ORIGINAL_README.md

@@ -0,0 +1,167 @@
+# Auffusion: Leveraging the Power of Diffusion and Large Language Models for Text-to-Audio Generation
+
+
+[Paper](https://arxiv.org/pdf/2401.01044) | [Model](https://huggingface.co/auffusion/auffusion) | [Website and Examples](https://auffusion.github.io) | [Audio Manipulation Notebooks](https://github.com/happylittlecat2333/Auffusion/tree/main/notebooks/README.md) | [Hugging Face Models](https://huggingface.co/auffusion) | [Google Colab](https://colab.research.google.com/drive/1JEPHT_AvHZxvlaZAsetkBnMrzCGMRKaf?usp=sharing)
+
+
+## Description
+
+**Auffusion** is a latent diffusion model (LDM) for text-to-audio (TTA) generation. **Auffusion** can generate realistic audios including human sounds, animal sounds, natural and artificial sounds and sound effects from textual prompts. We introduce Auffusion, a TTA system adapting T2I model frameworks to TTA task, by effectively leveraging their inherent generative strengths and precise cross-modal alignment. Our objective and subjective evaluations demonstrate that Auffusion surpasses previous TTA approaches using limited data and computational resource. We release our model, inference code, and pre-trained checkpoints for the research community.
+
+
+<p align="center">
+  <img src=img/overview.png />
+</p>
+
+
+## 🚀 News
+- **2024/01/02**: 📣 [Colab notebooks](notebooks/README.md) for audio manipulation is released. Feel free to try!
+
+- **2023/12/31**: 📣 Auffusion release.[Demo website](https://auffusion.github.io/) and 3 models are released in [Hugging Face](https://huggingface.co/auffusion). 
+
+
+## Auffusion Model Family
+
+| Model Name                 | Model Path                                                                                                              |
+|----------------------------|------------------------------------------------------------------------------------------------------------------------ |
+| Auffusion                  | [https://huggingface.co/auffusion/auffusion](https://huggingface.co/auffusion/auffusion)                                |
+| Auffusion-Full             | [https://huggingface.co/auffusion/auffusion-full](https://huggingface.co/auffusion/auffusion-full)                      |
+| Auffusion-Full-no-adapter  | [https://huggingface.co/auffusion/auffusion-full-no-adapter](https://huggingface.co/auffusion/auffusion-full-no-adapter)|
+
+
+
+## 📀 Prerequisites
+
+Our code is built on pytorch version 2.0.1. We mention `torch==2.0.1` in the requirements file but you might need to install a specific cuda version of torch depending on your GPU device type. We also depend on `diffusers==0.18.2`.
+
+Install `requirements.txt`.
+
+```bash
+git clone https://github.com/happylittlecat2333/Auffusion/
+cd Auffusion
+pip install -r requirements.txt
+```
+
+You might also need to install `libsndfile1` for soundfile to work properly in linux:
+
+```bash
+(sudo) apt-get install libsndfile1
+```
+
+
+## ⭐ Quickstart Guide
+
+Download the **Auffusion** model and generate audio from a text prompt:
+
+```python
+import IPython, torch
+import soundfile as sf
+from auffusion_pipeline import AuffusionPipeline
+
+pipeline = AuffusionPipeline.from_pretrained("auffusion/auffusion")
+
+prompt = "Birds singing sweetly in a blooming garden"
+output = pipeline(prompt=prompt)
+audio = output.audios[0]
+sf.write(f"{prompt}.wav", audio, samplerate=16000)
+IPython.display.Audio(data=audio, rate=16000)
+```
+
+The auffusion model will be automatically downloaded from Hugging Face and saved in cache. Subsequent runs will load the model directly from cache.
+
+The `generate` function uses 100 steps and 7.5 guidance_scale by default to sample from the latent diffusion model. You can also vary parameters for different results.
+
+```python
+prompt = "Rolling thunder with lightning strikes"
+output = pipeline(prompt=prompt, num_inference_steps=100, guidance_scale=7.5)
+audio = output.audios[0]
+IPython.display.Audio(data=audio, rate=16000)
+```
+
+
+More generated samples are shown [here](https://auffusion.github.io). You can also try out the [colab notebook](https://colab.research.google.com/drive/1JEPHT_AvHZxvlaZAsetkBnMrzCGMRKaf?usp=sharing) to generate your own audio samples.
+
+
+## 🐍 How to make inferences?
+
+### From our released checkpoints in Hugging Face Hub
+
+To perform audio generation in AudioCaps test set from our Hugging Face checkpoints:
+
+```bash
+python inference.py \
+--pretrained_model_name_or_path="auffusion/auffusion" \
+--test_data_dir="./data/test_audiocaps.raw.json" \
+--output_dir="./output/auffusion_hf" \
+--enable_xformers_memory_efficient_attention \
+```
+
+### Note
+
+We use the evaluation tools from [https://github.com/haoheliu/audioldm_eval](https://github.com/haoheliu/audioldm_eval) to evaluate our models, and we adopt [https://huggingface.co/laion/clap-htsat-unfused](https://huggingface.co/laion/clap-htsat-unfused) to compute CLAP score.
+
+Some data instances originally released in AudioCaps have since been removed from YouTube and are no longer available. We thus evaluated our models on all the instances which were available as June, 2023.
+
+## Audio Manipulation
+
+We show some examples of audio manipulation using Auffusion. Current audio manipulation methods include:
+
+- Text-to-audio generation: [notebook](notebooks/text_to_audio.ipynb) or [colab](https://colab.research.google.com/drive/1JEPHT_AvHZxvlaZAsetkBnMrzCGMRKaf?usp=sharing)
+- Text-guided style transfer: [notebook](notebooks/img2img.ipynb) or [colab](https://colab.research.google.com/drive/1VjgryIz7kSXDzgCClqtqVgoDXIECeG0M?usp=sharing)
+- Audio inpainting: [notebook](notebooks/inpainting.ipynb) or [colab](https://colab.research.google.com/drive/1NsqeiutoAynhtaZnlhzBdTXtZ27tQxVc?usp=sharing)
+- attention-based word swap control: [notebook](notebooks/word_swap.ipynb) or [colab](https://colab.research.google.com/drive/18CtUoBMsPbgzeI-o0wHDYTtaErnq9KoI?usp=sharing)
+- attention-based reweight control: [notebook](notebooks/reweight.ipynb) or [colab](https://colab.research.google.com/drive/18CtUoBMsPbgzeI-o0wHDYTtaErnq9KoI?usp=sharing)
+
+The audio manipulation code examples can all be found in [notebooks](notebooks/README.md).
+
+# TODO
+
+[!["Buy Me A Coffee"](https://www.buymeacoffee.com/assets/img/custom_images/orange_img.png)](https://www.buymeacoffee.com/happylittlecat)
+
+- [x] Publish demo website and arxiv link.
+- [x] Publish Auffusion and Auffusion-Full checkpoints.
+- [x] Add text-guided style transfer.
+- [x] Add audio-to-audio generation.
+- [x] Add audio inpainting.
+- [x] Add word_swap and reweight prompt2prompt-based control.
+- [ ] Add audio super-resolution.
+- [ ] Build Gradio web application.
+- [ ] Add audio-to-audio, inpainting into Gradio web application.
+- [ ] Add style-transfer into Gradio web application.
+- [ ] Add audio super-resolution into Gradio web application.
+- [ ] Add prompt2prompt-based control into Gradio web application.
+- [ ] Add data preprocess and training code.
+
+
+
+## 📚 Citation
+Please consider citing the following article if you found our work useful:
+
+```bibtex
+@article{xue2024auffusion,
+  title={Auffusion: Leveraging the Power of Diffusion and Large Language Models for Text-to-Audio Generation}, 
+  author={Jinlong Xue and Yayue Deng and Yingming Gao and Ya Li},
+  journal={arXiv preprint arXiv:2401.01044},
+  year={2024}
+}
+```
+
+## 🙏 Acknowledgement
+Part of the code is borrowed from the following repos. We would like to thank the authors of these repos for their contribution.
+
+- https://github.com/huggingface/diffusers
+
+- https://github.com/huggingface/transformers
+
+- https://github.com/google/prompt-to-prompt
+
+- https://github.com/declare-lab/tango
+
+- https://github.com/riffusion/riffusion
+
+-  https://github.com/haoheliu/audioldm_eval
+
+
+## Contact
+
+If you have any problems regarding the paper, code, models, or the project itself, please feel free to open an issue or contact [Jinlong Xue](mailto:[email protected]) directly :)

auffusion_pipeline.py

@@ -0,0 +1,1039 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import inspect
+import warnings
+from typing import Any, Callable, Dict, List, Optional, Union
+from dataclasses import dataclass
+
+import torch
+from packaging import version
+from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+    deprecate,
+    is_accelerate_available,
+    is_accelerate_version,
+    logging,
+    randn_tensor,
+    replace_example_docstring,
+)
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from huggingface_hub import snapshot_download
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, PNDMScheduler
+from transformers import PretrainedConfig, AutoTokenizer
+import torch.nn as nn
+import os, json, PIL
+import numpy as np
+import torch.nn.functional as F
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from diffusers.utils.outputs import BaseOutput
+import matplotlib.pyplot as plt
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+
+def json_dump(data_json, json_save_path):
+    with open(json_save_path, 'w') as f:
+        json.dump(data_json, f, indent=4)
+        f.close()
+
+
+def json_load(json_path):
+    with open(json_path, 'r') as f:
+        data = json.load(f)
+        f.close()
+    return data                
+
+
+def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+        return CLIPTextModel
+    if "t5" in model_class.lower():
+        from transformers import T5EncoderModel
+        return T5EncoderModel
+    if "clap" in model_class.lower():
+        from transformers import ClapTextModelWithProjection
+        return ClapTextModelWithProjection
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+
+
+class ConditionAdapter(nn.Module):
+    def __init__(self, config):
+        super(ConditionAdapter, self).__init__()
+        self.config = config
+        self.proj = nn.Linear(self.config["condition_dim"], self.config["cross_attention_dim"])
+        self.norm = torch.nn.LayerNorm(self.config["cross_attention_dim"])
+        print(f"INITIATED: ConditionAdapter: {self.config}")
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.norm(x)
+        return x
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path):
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt")
+        config = json.loads(open(config_path).read())
+        instance = cls(config)
+        instance.load_state_dict(torch.load(ckpt_path))
+        print(f"LOADED: ConditionAdapter from {pretrained_model_name_or_path}")
+        return instance
+
+    def save_pretrained(self, pretrained_model_name_or_path):
+        os.makedirs(pretrained_model_name_or_path, exist_ok=True)
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt")        
+        json_dump(self.config, config_path)
+        torch.save(self.state_dict(), ckpt_path)
+        print(f"SAVED: ConditionAdapter {self.config['model_name']} to {pretrained_model_name_or_path}")
+
+
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    return noise_cfg
+
+
+
+LRELU_SLOPE = 0.1
+MAX_WAV_VALUE = 32768.0
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def get_config(config_path):
+    config = json.loads(open(config_path).read())
+    config = AttrDict(config)
+    return config
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def apply_weight_norm(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        weight_norm(m)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        # self.conv_pre = weight_norm(Conv1d(80, h.upsample_initial_channel, 7, 1, padding=3))
+        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) # change: 80 --> 512
+        resblock = ResBlock1 if h.resblock == '1' else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            if (k-u) % 2 == 0:
+                self.ups.append(weight_norm(
+                    ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+                                    k, u, padding=(k-u)//2)))
+            else:
+                self.ups.append(weight_norm(
+                    ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+                                    k, u, padding=(k-u)//2+1, output_padding=1)))
+            
+            # self.ups.append(weight_norm(
+            #     ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+            #                     k, u, padding=(k-u)//2)))
+            
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        print('Removing weight norm...')
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None):
+        if subfolder is not None:
+            pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, subfolder)
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "vocoder.pt")
+
+        config = get_config(config_path)
+        vocoder = cls(config)
+
+        state_dict_g = torch.load(ckpt_path)
+        vocoder.load_state_dict(state_dict_g["generator"])
+        vocoder.eval()
+        vocoder.remove_weight_norm()
+        return vocoder
+    
+    @torch.no_grad()
+    def inference(self, mels, lengths=None):
+        self.eval()
+        with torch.no_grad():
+            wavs = self(mels).squeeze(1)
+
+        wavs = (wavs.cpu().numpy() * MAX_WAV_VALUE).astype("int16")
+
+        if lengths is not None:
+            wavs = wavs[:, :lengths]
+
+        return wavs
+
+
+
+def normalize_spectrogram(
+    spectrogram: torch.Tensor,
+    max_value: float = 200, 
+    min_value: float = 1e-5, 
+    power: float = 1., 
+) -> torch.Tensor:
+    
+    # Rescale to 0-1
+    max_value = np.log(max_value) # 5.298317366548036
+    min_value = np.log(min_value) # -11.512925464970229
+    spectrogram = torch.clamp(spectrogram, min=min_value, max=max_value)
+    data = (spectrogram - min_value) / (max_value - min_value)
+    # Apply the power curve
+    data = torch.pow(data, power)
+    # 1D -> 3D
+    data = data.repeat(3, 1, 1)
+    # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner
+    data = torch.flip(data, [1])
+
+    return data
+
+
+def denormalize_spectrogram(
+    data: torch.Tensor,
+    max_value: float = 200, 
+    min_value: float = 1e-5, 
+    power: float = 1, 
+) -> torch.Tensor:
+    
+    assert len(data.shape) == 3, "Expected 3 dimensions, got {}".format(len(data.shape))
+
+    max_value = np.log(max_value)
+    min_value = np.log(min_value)
+    # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner
+    data = torch.flip(data, [1])    
+    if data.shape[0] == 1:
+        data = data.repeat(3, 1, 1)        
+    assert data.shape[0] == 3, "Expected 3 channels, got {}".format(data.shape[0])
+    data = data[0]
+    # Reverse the power curve
+    data = torch.pow(data, 1 / power)
+    # Rescale to max value
+    spectrogram = data * (max_value - min_value) + min_value
+
+    return spectrogram
+
+@staticmethod
+def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray:
+    """
+    Convert a PyTorch tensor to a NumPy image.
+    """
+    images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+    return images
+
+@staticmethod
+def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
+    """
+    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 = [PIL.Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [PIL.Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def image_add_color(spec_img):
+    cmap = plt.get_cmap('viridis')
+    cmap_r = cmap.reversed()
+    image = cmap(np.array(spec_img)[:,:,0])[:, :, :3]  # 省略透明度通道
+    image = (image - image.min()) / (image.max() - image.min())
+    image = PIL.Image.fromarray(np.uint8(image*255))
+    return image
+
+
+@dataclass
+class PipelineOutput(BaseOutput):
+    """
+    Output class for audio pipelines.
+
+    Args:
+        audios (`np.ndarray`)
+            List of denoised audio samples of a NumPy array of shape `(batch_size, num_channels, sample_rate)`.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+    spectrograms: Union[List[np.ndarray], np.ndarray]
+    audios: Union[List[np.ndarray], np.ndarray]
+
+
+
+class AuffusionPipeline(DiffusionPipeline):
+
+    r"""
+    Pipeline for text-to-image generation using Stable Diffusion.
+
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    In addition the pipeline inherits the following loading methods:
+        - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`]
+        - *LoRA*: [`loaders.LoraLoaderMixin.load_lora_weights`]
+        - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
+
+    as well as the following saving methods:
+        - *LoRA*: [`loaders.LoraLoaderMixin.save_lora_weights`]
+
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder. Stable Diffusion uses the text portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.
+        feature_extractor ([`CLIPImageProcessor`]):
+            Model that extracts features from generated images to be used as inputs for the `safety_checker`.
+    """
+    _optional_components = ["safety_checker", "feature_extractor", "text_encoder_list", "tokenizer_list", "adapter_list", "vocoder"]
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        unet: UNet2DConditionModel,        
+        scheduler: KarrasDiffusionSchedulers,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPImageProcessor,
+        text_encoder_list: Optional[List[Callable]] = None,
+        tokenizer_list: Optional[List[Callable]] = None,
+        vocoder: Generator = None,
+        requires_safety_checker: bool = False,        
+        adapter_list: Optional[List[Callable]] = None,
+        tokenizer_model_max_length: Optional[int] = 77, # 77 is the default value for the CLIPTokenizer(and set for other models)
+    ):
+        super().__init__()
+
+        self.text_encoder_list = text_encoder_list
+        self.tokenizer_list = tokenizer_list
+        self.vocoder = vocoder
+        self.adapter_list = adapter_list
+        self.tokenizer_model_max_length = tokenizer_model_max_length
+
+        self.register_modules(
+            vae=vae,
+            unet=unet,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str = "auffusion/auffusion",
+        dtype: torch.dtype = torch.float16,
+        device: str = "cuda",
+    ):
+        if not os.path.isdir(pretrained_model_name_or_path):
+            pretrained_model_name_or_path = snapshot_download(pretrained_model_name_or_path) 
+        
+        vae = AutoencoderKL.from_pretrained(pretrained_model_name_or_path, subfolder="vae")
+        unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet")
+        feature_extractor = CLIPImageProcessor.from_pretrained(pretrained_model_name_or_path, subfolder="feature_extractor")
+        scheduler = PNDMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+
+        vocoder = Generator.from_pretrained(pretrained_model_name_or_path, subfolder="vocoder").to(device, dtype)
+
+        text_encoder_list, tokenizer_list, adapter_list = [], [], []
+        
+        condition_json_path = os.path.join(pretrained_model_name_or_path, "condition_config.json")
+        condition_json_list = json.loads(open(condition_json_path).read())
+        
+        for i, condition_item in enumerate(condition_json_list):
+            
+            # Load Condition Adapter
+            text_encoder_path = os.path.join(pretrained_model_name_or_path, condition_item["text_encoder_name"])
+            tokenizer = AutoTokenizer.from_pretrained(text_encoder_path)
+            tokenizer_list.append(tokenizer)
+            text_encoder_cls = import_model_class_from_model_name_or_path(text_encoder_path)
+            text_encoder = text_encoder_cls.from_pretrained(text_encoder_path).to(device, dtype)
+            text_encoder_list.append(text_encoder)
+            print(f"LOADING CONDITION ENCODER {i}")
+
+            # Load Condition Adapter
+            adapter_path = os.path.join(pretrained_model_name_or_path, condition_item["condition_adapter_name"])
+            adapter = ConditionAdapter.from_pretrained(adapter_path).to(device, dtype)
+            adapter_list.append(adapter)
+            print(f"LOADING CONDITION ADAPTER {i}")
+
+
+        pipeline = cls(
+            vae=vae,
+            unet=unet,
+            text_encoder_list=text_encoder_list,
+            tokenizer_list=tokenizer_list,
+            vocoder=vocoder,
+            adapter_list=adapter_list,
+            scheduler=scheduler,
+            safety_checker=None,
+            feature_extractor=feature_extractor,
+        )
+        pipeline = pipeline.to(device, dtype)
+
+        return pipeline
+    
+
+    def to(self, device, dtype=None):
+        super().to(device, dtype)
+
+        self.vocoder.to(device, dtype)
+
+        for text_encoder in self.text_encoder_list:
+            text_encoder.to(device, dtype)
+       
+        if self.adapter_list is not None:
+            for adapter in self.adapter_list:
+                adapter.to(device, dtype)
+
+        return self
+
+
+    def enable_vae_slicing(self):
+        r"""
+        Enable sliced VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
+        steps. This is useful to save some memory and allow larger batch sizes.
+        """
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        r"""
+        Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_slicing()
+
+    def enable_vae_tiling(self):
+        r"""
+        Enable tiled VAE decoding.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
+        several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
+        """
+        self.vae.enable_tiling()
+
+    def disable_vae_tiling(self):
+        r"""
+        Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
+        computing decoding in one step.
+        """
+        self.vae.disable_tiling()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
+        text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
+        `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
+        Note that offloading happens on a submodule basis. Memory savings are higher than with
+        `enable_model_cpu_offload`, but performance is lower.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+            from accelerate import cpu_offload
+        else:
+            raise ImportError("`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            cpu_offload(cpu_offloaded_model, device)
+
+        if self.safety_checker is not None:
+            cpu_offload(self.safety_checker, execution_device=device, offload_buffers=True)
+
+    def enable_model_cpu_offload(self, gpu_id=0):
+        r"""
+        Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
+        to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
+        method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
+        `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
+        """
+        if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
+            from accelerate import cpu_offload_with_hook
+        else:
+            raise ImportError("`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        if self.device.type != "cpu":
+            self.to("cpu", silence_dtype_warnings=True)
+            torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)
+
+        hook = None
+        for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
+            _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
+
+        if self.safety_checker is not None:
+            _, hook = cpu_offload_with_hook(self.safety_checker, device, prev_module_hook=hook)
+
+        # We'll offload the last model manually.
+        self.final_offload_hook = hook
+
+    @property
+    def _execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    
+    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,
+    ):
+
+        assert len(self.text_encoder_list) == len(self.tokenizer_list), "Number of text_encoders must match number of tokenizers"
+        if self.adapter_list is not None:
+            assert len(self.text_encoder_list) == len(self.adapter_list), "Number of text_encoders must match number of adapters"
+
+        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]
+
+        def get_prompt_embeds(prompt_list, device):
+            if isinstance(prompt_list, str):
+                prompt_list = [prompt_list]
+
+            prompt_embeds_list = []
+            for prompt in prompt_list:
+                encoder_hidden_states_list = []
+
+                # Generate condition embedding
+                for j in range(len(self.text_encoder_list)):
+                    # get condition embedding using condition encoder
+                    input_ids = self.tokenizer_list[j](prompt, return_tensors="pt").input_ids.to(device)            
+                    cond_embs = self.text_encoder_list[j](input_ids).last_hidden_state # [bz, text_len, text_dim]
+                    # padding to max_length
+                    if cond_embs.shape[1] < self.tokenizer_model_max_length: 
+                        cond_embs = torch.functional.F.pad(cond_embs, (0, 0, 0, self.tokenizer_model_max_length - cond_embs.shape[1]), value=0)
+                    else:
+                        cond_embs = cond_embs[:, :self.tokenizer_model_max_length, :]
+
+                    # use condition adapter
+                    if self.adapter_list is not None:
+                        cond_embs = self.adapter_list[j](cond_embs)
+                        encoder_hidden_states_list.append(cond_embs)
+
+                prompt_embeds = torch.cat(encoder_hidden_states_list, dim=1)
+                prompt_embeds_list.append(prompt_embeds)
+
+            prompt_embeds = torch.cat(prompt_embeds_list, dim=0)
+            return prompt_embeds
+
+
+        if prompt_embeds is None:           
+            prompt_embeds = get_prompt_embeds(prompt, device)
+
+        prompt_embeds = prompt_embeds.to(dtype=self.unet.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)
+
+        if do_classifier_free_guidance and negative_prompt_embeds is None:
+
+            if negative_prompt is None:
+                negative_prompt_embeds = torch.zeros_like(prompt_embeds).to(dtype=prompt_embeds.dtype, device=device)
+
+            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):
+                negative_prompt = [negative_prompt]
+                negative_prompt_embeds = get_prompt_embeds(negative_prompt, device)
+            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:
+                negative_prompt_embeds = get_prompt_embeds(negative_prompt, device)
+
+        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.unet.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 run_safety_checker(self, image, device, dtype):
+        if self.safety_checker is None:
+            has_nsfw_concept = None
+        else:
+            if torch.is_tensor(image):
+                feature_extractor_input = self.image_processor.postprocess(image, output_type="pil")
+            else:
+                feature_extractor_input = self.image_processor.numpy_to_pil(image)
+            safety_checker_input = self.feature_extractor(feature_extractor_input, return_tensors="pt").to(device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
+            )
+        return image, has_nsfw_concept
+
+    def decode_latents(self, latents):
+        warnings.warn(
+            "The decode_latents method is deprecated and will be removed in a future version. Please"
+            " use VaeImageProcessor instead",
+            FutureWarning,
+        )
+        latents = 1 / self.vae.config.scaling_factor * latents
+        image = self.vae.decode(latents, return_dict=False)[0]
+        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 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,
+        height,
+        width,
+        callback_steps,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+    ):
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        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
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = 256,
+        width: Optional[int] = 1024,
+        num_inference_steps: int = 100,
+        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,
+        output_type: Optional[str] = "pt",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: int = 1,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        guidance_rescale: float = 0.0,
+        duration: Optional[float] = 10,
+    ):
+       
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        audio_length = int(duration * 16000)
+
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt, height, width, 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._execution_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.config.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
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            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
+                noise_pred = self.unet(
+                    latent_model_input,
+                    t,
+                    encoder_hidden_states=prompt_embeds,
+                    cross_attention_kwargs=cross_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                # 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)
+
+                if do_classifier_free_guidance and guidance_rescale > 0.0:
+                    # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                    noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        # Offload last model to CPU
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.final_offload_hook.offload()
+
+
+        # Generate audio
+        spectrograms, audios = [], []
+        for img in image:
+            spectrogram = denormalize_spectrogram(img)
+            audio = self.vocoder.inference(spectrogram, lengths=audio_length)[0]
+            audios.append(audio)
+            spectrograms.append(spectrogram)
+
+        # Convert to PIL
+        images = pt_to_numpy(image)    
+        images = numpy_to_pil(images)
+        images = [image_add_color(image) for image in images]
+
+        if not return_dict:
+            return (images, audios, spectrograms)
+                    
+
+        return PipelineOutput(images=images, audios=audios, spectrograms=spectrograms)
+
+
+
+

config/clip_condition_config.json

@@ -0,0 +1,11 @@
+[
+    {
+        "text_encoder_name": "text_encoder_0",
+        "condition_adapter_name": "condition_adapter_0",
+        "condition_type": "clip-vit-large-patch14_text",
+        "pretrained_model_name_or_path": "openai/clip-vit-large-patch14",
+        "condition_max_length": 77,
+        "condition_dim": 768,
+        "cross_attention_dim": 768
+    }
+]

converter.py

@@ -0,0 +1,375 @@
+
+
+import numpy as np
+from PIL import Image
+
+import math
+import os
+import random
+import torch
+import json
+import torch.utils.data
+import numpy as np
+import librosa
+from librosa.util import normalize
+from scipy.io.wavfile import read
+from librosa.filters import mel as librosa_mel_fn
+
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
+from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+
+MAX_WAV_VALUE = 32768.0
+
+
+def load_wav(full_path):
+    sampling_rate, data = read(full_path)
+    return data, sampling_rate
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    return np.exp(x) / C
+
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression_torch(x, C=1):
+    return torch.exp(x) / C
+
+
+def spectral_normalize_torch(magnitudes):
+    output = dynamic_range_compression_torch(magnitudes)
+    return output
+
+
+def spectral_de_normalize_torch(magnitudes):
+    output = dynamic_range_decompression_torch(magnitudes)
+    return output
+
+
+mel_basis = {}
+hann_window = {}
+
+
+def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
+    if torch.min(y) < -1.:
+        print('min value is ', torch.min(y))
+    if torch.max(y) > 1.:
+        print('max value is ', torch.max(y))
+
+    global mel_basis, hann_window
+    if fmax not in mel_basis:
+        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
+        mel_basis[str(fmax)+'_'+str(y.device)] = torch.from_numpy(mel).float().to(y.device)
+        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    y = y.squeeze(1)
+
+    # complex tensor as default, then use view_as_real for future pytorch compatibility
+    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
+    spec = torch.view_as_real(spec)
+    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
+
+    spec = torch.matmul(mel_basis[str(fmax)+'_'+str(y.device)], spec)
+    spec = spectral_normalize_torch(spec)
+
+    return spec
+
+
+def spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
+    if torch.min(y) < -1.:
+        print('min value is ', torch.min(y))
+    if torch.max(y) > 1.:
+        print('max value is ', torch.max(y))
+
+    global hann_window
+    hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+
+    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
+    y = y.squeeze(1)
+
+    # complex tensor as default, then use view_as_real for future pytorch compatibility
+    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
+                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
+    spec = torch.view_as_real(spec)
+    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
+
+    return spec
+
+
+def normalize_spectrogram(
+    spectrogram: torch.Tensor,
+    max_value: float = 200, 
+    min_value: float = 1e-5, 
+    power: float = 1., 
+    inverse: bool = False
+) -> torch.Tensor:
+    
+    # Rescale to 0-1
+    max_value = np.log(max_value) # 5.298317366548036
+    min_value = np.log(min_value) # -11.512925464970229
+
+    assert spectrogram.max() <= max_value and spectrogram.min() >= min_value
+
+    data = (spectrogram - min_value) / (max_value - min_value)
+
+    # Invert
+    if inverse:
+        data = 1 - data
+
+    # Apply the power curve
+    data = torch.pow(data, power)  
+    
+    # 1D -> 3D
+    data = data.repeat(3, 1, 1)
+
+    # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner
+    data = torch.flip(data, [1])
+
+    return data
+
+
+
+def denormalize_spectrogram(
+    data: torch.Tensor,
+    max_value: float = 200, 
+    min_value: float = 1e-5, 
+    power: float = 1, 
+    inverse: bool = False,
+) -> torch.Tensor:
+    
+    max_value = np.log(max_value)
+    min_value = np.log(min_value)
+
+    # Flip Y axis: image origin at the top-left corner, spectrogram origin at the bottom-left corner
+    data = torch.flip(data, [1])
+
+    assert len(data.shape) == 3, "Expected 3 dimensions, got {}".format(len(data.shape))
+    
+    if data.shape[0] == 1:
+        data = data.repeat(3, 1, 1)
+        
+    assert data.shape[0] == 3, "Expected 3 channels, got {}".format(data.shape[0])
+    data = data[0]
+
+    # Reverse the power curve
+    data = torch.pow(data, 1 / power)
+
+    # Invert
+    if inverse:
+        data = 1 - data
+
+    # Rescale to max value
+    spectrogram = data * (max_value - min_value) + min_value
+
+    return spectrogram
+
+
+def get_mel_spectrogram_from_audio(audio, device="cuda"):
+    audio = audio / MAX_WAV_VALUE
+    audio = librosa.util.normalize(audio) * 0.95
+        
+    audio = torch.FloatTensor(audio)
+    audio = audio.unsqueeze(0)    
+
+    waveform = audio.to(device)
+    spec = mel_spectrogram(waveform, n_fft=2048, num_mels=256, sampling_rate=16000, hop_size=160, win_size=1024, fmin=0, fmax=8000, center=False)
+    return audio, spec
+
+
+
+LRELU_SLOPE = 0.1
+MAX_WAV_VALUE = 32768.0
+
+
+class AttrDict(dict):
+    def __init__(self, *args, **kwargs):
+        super(AttrDict, self).__init__(*args, **kwargs)
+        self.__dict__ = self
+
+
+def get_config(config_path):
+    config = json.loads(open(config_path).read())
+    config = AttrDict(config)
+    return config
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def apply_weight_norm(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        weight_norm(m)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+class ResBlock1(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlock1, self).__init__()
+        self.h = h
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlock2(torch.nn.Module):
+    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3)):
+        super(ResBlock2, self).__init__()
+        self.h = h
+        self.convs = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1])))
+        ])
+        self.convs.apply(init_weights)
+
+    def forward(self, x):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            xt = c(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+
+class Generator(torch.nn.Module):
+    def __init__(self, h):
+        super(Generator, self).__init__()
+        self.h = h
+        self.num_kernels = len(h.resblock_kernel_sizes)
+        self.num_upsamples = len(h.upsample_rates)
+        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) # change: 80 --> 512
+        resblock = ResBlock1 if h.resblock == '1' else ResBlock2
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+            if (k-u) % 2 == 0:
+                self.ups.append(weight_norm(
+                    ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+                                    k, u, padding=(k-u)//2)))
+            else:
+                self.ups.append(weight_norm(
+                    ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+                                    k, u, padding=(k-u)//2+1, output_padding=1)))
+            
+            # self.ups.append(weight_norm(
+            #     ConvTranspose1d(h.upsample_initial_channel//(2**i), h.upsample_initial_channel//(2**(i+1)),
+            #                     k, u, padding=(k-u)//2)))
+            
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = h.upsample_initial_channel//(2**(i+1))
+            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+                self.resblocks.append(resblock(h, ch, k, d))
+
+        self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3))
+        self.ups.apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+    def forward(self, x):
+        x = self.conv_pre(x)
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i*self.num_kernels+j](x)
+                else:
+                    xs += self.resblocks[i*self.num_kernels+j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.ups:
+            remove_weight_norm(l)
+        for l in self.resblocks:
+            l.remove_weight_norm()
+        remove_weight_norm(self.conv_pre)
+        remove_weight_norm(self.conv_post)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, subfolder=None):
+        if subfolder is not None:
+            pretrained_model_name_or_path = os.path.join(pretrained_model_name_or_path, subfolder)
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "vocoder.pt")
+
+        config = get_config(config_path)
+        vocoder = cls(config)
+
+        state_dict_g = torch.load(ckpt_path)
+        vocoder.load_state_dict(state_dict_g["generator"])
+        vocoder.eval()
+        vocoder.remove_weight_norm()
+        return vocoder    
+    
+    
+    @torch.no_grad()
+    def inference(self, mels, lengths=None):
+        self.eval()
+        with torch.no_grad():
+            wavs = self(mels).squeeze(1)
+
+        wavs = (wavs.cpu().numpy() * MAX_WAV_VALUE).astype("int16")
+
+        if lengths is not None:
+            wavs = wavs[:, :lengths]
+
+        return wavs

inference.py

@@ -0,0 +1,129 @@
+
+import os
+import argparse
+import torch
+from diffusers.utils.import_utils import is_xformers_available
+from datasets import load_dataset
+from tqdm.auto import tqdm
+from scipy.io.wavfile import write
+from auffusion_pipeline import AuffusionPipeline
+
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Simple example of a inference script.")
+    parser.add_argument(
+        "--pretrained_model_name_or_path",
+        type=str,
+        default="auffusion/auffusion",
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+    )    
+    parser.add_argument(
+        "--test_data_dir",
+        type=str,
+        default="./data/test_audiocaps.raw.json",
+        help="Path to test dataset in json file",
+    )    
+    parser.add_argument(
+        "--audio_column", type=str, default="audio_path", help="The column of the dataset containing an audio."
+    )
+    parser.add_argument(
+        "--caption_column", type=str, default="text", help="The column of the dataset containing a caption."
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="./output/auffusion_hf",
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--sample_rate", type=int, default=16000, help="The sample rate of audio."
+    )
+    parser.add_argument(
+        "--duration", type=int, default=10, help="The duration(s) of audio."
+    ) 
+    parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible inference.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="fp16",
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
+        ),
+    )    
+    parser.add_argument(
+        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
+    )
+    parser.add_argument(
+        "--guidance_scale", type=float, default=7.5, help="The scale of guidance."
+    )
+    parser.add_argument(
+        "--num_inference_steps", type=int, default=100, help="Number of inference steps to perform."
+    )
+    parser.add_argument(
+        "--width", type=int, default=1024, help="Width of the image."
+    )
+    parser.add_argument(
+        "--height", type=int, default=256, help="Height of the image."
+    ) 
+    args = parser.parse_args()
+    
+    return args
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)   
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    weight_dtype = torch.float16 if args.mixed_precision == "fp16" else torch.float32
+
+    pipeline = AuffusionPipeline.from_pretrained(args.pretrained_model_name_or_path)
+    pipeline = pipeline.to(device, weight_dtype)
+    pipeline.set_progress_bar_config(disable=True)
+    
+    if is_xformers_available() and args.enable_xformers_memory_efficient_attention:
+        pipeline.enable_xformers_memory_efficient_attention()
+
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+
+    # load dataset
+    audio_column, caption_column = args.audio_column, args.caption_column
+    data_files = {"test": args.test_data_dir}
+    dataset = load_dataset("json", data_files=data_files, split="test")
+
+    # output dir
+    audio_output_dir = os.path.join(args.output_dir, "audios")
+    os.makedirs(audio_output_dir, exist_ok=True)
+
+    # generating    
+    audio_length = args.sample_rate * args.duration
+    for i in tqdm(range(len(dataset)), desc="Generating"):
+
+        prompt = dataset[i][caption_column]        
+        audio_name = os.path.basename(dataset[i][audio_column])
+
+        audio_path = os.path.join(audio_output_dir, audio_name)
+
+        if os.path.exists(audio_path):
+            continue
+
+        with torch.autocast("cuda"):
+            output = pipeline(prompt=prompt, num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, generator=generator, width=args.width, height=args.height)        
+
+        audio = output.audios[0][:audio_length]
+
+        write(audio_path, args.sample_rate, audio)
+
+
+if __name__ == "__main__":
+    main()
+
+
+
+

inference.sh

@@ -0,0 +1,24 @@
+
+MODEL_NAME="auffusion/auffusion"
+test_data_dir="./data/test_audiocaps.raw.json"  
+output_dir="./output/auffusion"
+audio_column="spec_path"
+caption_column="text"
+num_inference_steps=100 
+guidance_scale=7.5
+
+
+training_params="--pretrained_model_name_or_path=$MODEL_NAME \
+    --test_data_dir=$test_data_dir \
+    --output_dir=$output_dir \
+    --audio_column=$audio_column \
+    --caption_column=$caption_column \
+    --sample_rate=16000 \
+    --duration=10 \
+    --num_inference_steps=$num_inference_steps \
+    --guidance_scale=$guidance_scale \
+    --mixed_precision="fp16" \
+    --enable_xformers_memory_efficient_attention \
+    "
+
+python inference.py $training_params

notebooks/README.md

@@ -0,0 +1,9 @@
+# Audio Manipulation
+
+We show examples of audio manipulation. Current audio manipulation methods include:
+
+- Text-to-audio generation: [text_to_audio.ipynb](text_to_audio.ipynb) or [colab notebook](https://colab.research.google.com/drive/1JEPHT_AvHZxvlaZAsetkBnMrzCGMRKaf?usp=sharing)
+- Text-guided style transfer: [img2img.ipynb](img2img.ipynb) or [colab notebook](https://colab.research.google.com/drive/1VjgryIz7kSXDzgCClqtqVgoDXIECeG0M?usp=sharing)
+- Audio inpainting: [inpainting.ipynb](inpainting.ipynb) or [colab notebook](https://colab.research.google.com/drive/1NsqeiutoAynhtaZnlhzBdTXtZ27tQxVc?usp=sharing)
+- attention-based word swap control: [word_swap.ipynb](word_swap.ipynb) or [colab notebook](https://colab.research.google.com/drive/18CtUoBMsPbgzeI-o0wHDYTtaErnq9KoI?usp=sharing)
+- attention-based reweight control: [reweight.ipynb](reweight.ipynb) or [colab notebook](https://colab.research.google.com/drive/18CtUoBMsPbgzeI-o0wHDYTtaErnq9KoI?usp=sharing)

notebooks/examples/img2img/GIOApFAWDOc_160.wav

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

notebooks/examples/img2img/_GI7meqlYZk_30.wav

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

notebooks/examples/inpainting/14ekd4nkpwc_28.wav

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

notebooks/examples/inpainting/3ek-xLwr05Q_30.wav

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

notebooks/examples/inpainting/9z8XIRyUq9Q_30.wav

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

notebooks/examples/inpainting/IvfaKPDWC00_160.wav

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

prompt2prompt/Roboto-Regular.ttf

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

prompt2prompt/pipeline_prompt2prompt.py

@@ -0,0 +1,353 @@
+from typing import Callable, List, Optional, Union
+import numpy as np
+import torch
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.models.cross_attention import CrossAttention
+
+#from pipeline_sd import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipeline
+import matplotlib.pyplot as plt
+
+from prompt2prompt.ptp_utils import AttentionStore
+import prompt2prompt.ptp_utils as ptp_utils
+from PIL import Image
+
+
+class Prompt2PromptPipeline(StableDiffusionPipeline):
+    r"""
+    Pipeline for text-to-image generation using Stable Diffusion.
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+    Args:
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            Frozen text-encoder. Stable Diffusion uses the text portion of
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+        unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+        scheduler ([`SchedulerMixin`]):
+            A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+            [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+        safety_checker ([`StableDiffusionSafetyChecker`]):
+            Classification module that estimates whether generated images could be considered offensive or harmful.
+            Please, refer to the [model card](https://huggingface.co/runwayml/stable-diffusion-v1-5) for details.
+        feature_extractor ([`CLIPFeatureExtractor`]):
+            Model that extracts features from generated images to be used as inputs for the `safety_checker`.
+    """
+    _optional_components = ["safety_checker", "feature_extractor"]
+
+    
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        controller: AttentionStore = None, # 传入attention_store作为p2p的控制。
+        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,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`):
+                The prompt or prompts to guide the image generation.
+            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. 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`, *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`.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
+                plain tuple.
+            callback (`Callable`, *optional*):
+                A function that will be called every `callback_steps` steps during inference. The function will be
+                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 will be called. If not specified, the callback will be
+                called at every step.
+
+        Returns:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
+            [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
+            When returning a tuple, the first element is a list with the generated images, and the second element is a
+            list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
+            (nsfw) content, according to the `safety_checker`.
+        """
+
+        self.register_attention_control(controller) # add attention controller
+
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(prompt, height, width, callback_steps)
+
+        # 2. Define call parameters
+        batch_size = 1 if isinstance(prompt, str) else len(prompt)
+        device = self._execution_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
+        text_embeddings = self._encode_prompt(
+            prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
+        )
+
+        # 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,
+            text_embeddings.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
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            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
+                noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
+
+                # 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
+
+                # step callback
+                latents = controller.step_callback(latents)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+                    if callback is not None and i % callback_steps == 0:
+                        callback(i, t, latents)
+
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+
+        # 9. Run safety checker
+        image, has_nsfw_concept = self.run_safety_checker(image, device, text_embeddings.dtype)
+
+        # 10. Convert to PIL
+        if output_type == "pil":
+            image = self.numpy_to_pil(image)
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+    def register_attention_control(self, controller):
+        attn_procs = {}
+        cross_att_count = 0
+        for name in self.unet.attn_processors.keys():
+            cross_attention_dim = None if name.endswith("attn1.processor") else self.unet.config.cross_attention_dim
+            if name.startswith("mid_block"):
+                hidden_size = self.unet.config.block_out_channels[-1]
+                place_in_unet = "mid"
+            elif name.startswith("up_blocks"):
+                block_id = int(name[len("up_blocks.")])
+                hidden_size = list(reversed(self.unet.config.block_out_channels))[block_id]
+                place_in_unet = "up"
+            elif name.startswith("down_blocks"):
+                block_id = int(name[len("down_blocks.")])
+                hidden_size = self.unet.config.block_out_channels[block_id]
+                place_in_unet = "down"
+            else:
+                continue
+            cross_att_count += 1
+            attn_procs[name] = P2PCrossAttnProcessor(
+                controller=controller, place_in_unet=place_in_unet
+            )
+
+        self.unet.set_attn_processor(attn_procs)
+        controller.num_att_layers = cross_att_count
+
+
+    # def aggregate_attention(self, prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
+    #     out = []
+    #     attention_maps = attention_store.get_average_attention()
+    #     num_pixels = res ** 2
+    #     for location in from_where:
+    #         for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
+    #             if item.shape[1] == num_pixels:
+    #                 cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
+    #                 out.append(cross_maps)
+    #     out = torch.cat(out, dim=0)
+    #     out = out.sum(0) / out.shape[0]
+    #     return out.cpu()
+
+    def aggregate_attention(self, prompts, attention_store: AttentionStore, res: List[int], from_where: List[str], is_cross: bool, select: int):
+        out = []
+        attention_maps = attention_store.get_average_attention()
+        # num_pixels = res ** 2
+        num_pixels = res[0] * res[1]
+        for location in from_where:
+            for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
+                if item.shape[1] == num_pixels:
+                    cross_maps = item.reshape(len(prompts), -1, res[0], res[1], item.shape[-1])[select]
+                    out.append(cross_maps)
+        out = torch.cat(out, dim=0)
+        out = out.sum(0) / out.shape[0]
+        return out.cpu()
+
+    def show_cross_attention(self, prompts, attention_store: AttentionStore, res: List[int], from_where: List[str], select: int = 0, image_size: List[int]=[1024, 256], num_rows: int = 1, font_scale=2, thickness=4, cmap_name="plasma"):
+        tokens = self.tokenizer.encode(prompts[select])
+        decoder = self.tokenizer.decode
+        attention_maps = self.aggregate_attention(prompts, attention_store, res, from_where, True, select)
+        images = []
+
+        cmap = plt.get_cmap(cmap_name)
+        cmap_r = cmap.reversed()
+
+        for i in range(len(tokens)):
+            image = attention_maps[:, :, i]
+            image = 255 * image / image.max()
+            image = image.unsqueeze(-1).expand(*image.shape, 3)
+            image = image.numpy().astype(np.uint8)
+            # image = np.array(Image.fromarray(image).resize((256, 256)))
+            # image = np.array(Image.fromarray(image).resize((512, 128)))
+
+            image = cmap(np.array(image)[:,:,0])[:, :, :3]  # 省略透明度通道
+            # image = image ** 2
+            image = (image - image.min()) / (image.max() - image.min())
+            image = Image.fromarray(np.uint8(image*255))
+            # image = np.array(image.resize((1024, 256)))
+            image = np.array(image.resize(image_size))
+
+
+            image = ptp_utils.text_under_image(image, decoder(int(tokens[i])), font_scale=font_scale, thickness=thickness)
+            images.append(image)
+        return ptp_utils.view_images(np.stack(images, axis=0), num_rows=num_rows)
+
+    # def show_cross_attention(self, prompts, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
+    #     tokens = self.tokenizer.encode(prompts[select])
+    #     decoder = self.tokenizer.decode
+    #     attention_maps = self.aggregate_attention(prompts, attention_store, res, from_where, True, select)
+    #     images = []
+    #     for i in range(len(tokens)):
+    #         image = attention_maps[:, :, i]
+    #         image = 255 * image / image.max()
+    #         image = image.unsqueeze(-1).expand(*image.shape, 3)
+    #         image = image.numpy().astype(np.uint8)
+    #         image = np.array(Image.fromarray(image).resize((256, 256)))
+    #         image = ptp_utils.text_under_image(image, decoder(int(tokens[i])))
+    #         images.append(image)
+    #     ptp_utils.view_images(np.stack(images, axis=0))
+        
+
+    def show_self_attention_comp(self, prompts, attention_store: AttentionStore, res: int, from_where: List[str],
+                            max_com=10, select: int = 0):
+        attention_maps = self.aggregate_attention(prompts, attention_store, res, from_where, False, select).numpy().reshape((res ** 2, res ** 2))
+        u, s, vh = np.linalg.svd(attention_maps - np.mean(attention_maps, axis=1, keepdims=True))
+        images = []
+        for i in range(max_com):
+            image = vh[i].reshape(res, res)
+            image = image - image.min()
+            image = 255 * image / image.max()
+            image = np.repeat(np.expand_dims(image, axis=2), 3, axis=2).astype(np.uint8)
+            image = Image.fromarray(image).resize((256, 256))
+            image = np.array(image)
+            images.append(image)
+        ptp_utils.view_images(np.concatenate(images, axis=1))
+
+
+class P2PCrossAttnProcessor:
+
+    def __init__(self, controller, place_in_unet):
+        super().__init__()
+        self.controller = controller
+        self.place_in_unet = place_in_unet
+
+    def __call__(self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None):
+        batch_size, sequence_length, _ = hidden_states.shape
+        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size=batch_size)
+
+        query = attn.to_q(hidden_states)
+
+        is_cross = encoder_hidden_states is not None
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query = attn.head_to_batch_dim(query)
+        key = attn.head_to_batch_dim(key)
+        value = attn.head_to_batch_dim(value)
+
+        attention_probs = attn.get_attention_scores(query, key, attention_mask)
+
+        # one line change
+        self.controller(attention_probs, is_cross, self.place_in_unet)
+
+        hidden_states = torch.bmm(attention_probs, value)
+        hidden_states = attn.batch_to_head_dim(hidden_states)
+
+        # linear proj
+        hidden_states = attn.to_out[0](hidden_states)
+        # dropout
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states

prompt2prompt/ptp_utils.py

@@ -0,0 +1,557 @@
+import abc
+
+import cv2
+import numpy as np
+import torch
+from IPython.display import display
+from PIL import Image
+from typing import Union, Tuple, List, Dict, Optional
+import torch.nn.functional as nnf
+from PIL import Image, ImageDraw, ImageFont
+
+
+# def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0)) -> np.ndarray:
+#     h, w, c = image.shape
+#     offset = int(h * .2)
+#     img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
+#     font = cv2.FONT_HERSHEY_SIMPLEX
+#     img[:h] = image
+#     textsize = cv2.getTextSize(text, font, 1, 2)[0]
+#     text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
+#     cv2.putText(img, text, (text_x, text_y), font, 1, text_color, 2)
+#     return img
+
+def text_under_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0), font_scale: float = 1.0, thickness: int = 3) -> np.ndarray:
+    h, w, c = image.shape
+    # offset = int(h * .3)
+    offset = int(h * .2)
+    img = np.ones((h + offset, w, c), dtype=np.uint8) * 255
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    img[:h] = image
+    textsize = cv2.getTextSize(text, font, font_scale, thickness)[0]
+    text_x, text_y = (w - textsize[0]) // 2, h + offset - textsize[1] // 2
+    cv2.putText(img, text, (text_x, text_y), font, font_scale, text_color, 2)
+    return img
+
+
+def text_under_image_pil(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0), font_scale: float = 1.0) -> np.ndarray:
+    image_pil = Image.fromarray(image)
+    draw = ImageDraw.Draw(image_pil)
+    
+    font_size = int(font_scale * image.shape[0] / 20)
+    # font = ImageFont.truetype("arial.ttf", font_size)
+    font_path = "./Roboto-Regular.ttf"
+    font = ImageFont.truetype(font_path, font_size)
+    
+    textsize = draw.textsize(text, font=font)
+    text_x = (image.shape[1] - textsize[0]) // 2
+    text_y = image.shape[0]
+    
+    draw.text((text_x, text_y), text, font=font, fill=text_color)
+    
+    return np.array(image_pil)
+
+
+def view_images(images: Union[np.ndarray, List],
+                num_rows: int = 1,
+                offset_ratio: float = 0.02,
+                display_image: bool = True) -> Image.Image:
+    """ Displays a list of images in a grid. """
+    if type(images) is list:
+        num_empty = len(images) % num_rows
+    elif images.ndim == 4:
+        num_empty = images.shape[0] % num_rows
+    else:
+        images = [images]
+        num_empty = 0
+    
+    empty_images = np.ones(images[0].shape, dtype=np.uint8) * 255
+    images = [image.astype(np.uint8) for image in images] + [empty_images] * num_empty
+    num_items = len(images)
+
+    h, w, c = images[0].shape
+    offset = int(h * offset_ratio)
+    num_cols = num_items // num_rows
+    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
+                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
+    for i in range(num_rows):
+        for j in range(num_cols):
+            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images[
+                i * num_cols + j]
+
+    pil_img = Image.fromarray(image_)
+    if display_image:
+        display(pil_img)
+    return pil_img
+
+
+def view_images_with_texts(images: Union[np.ndarray, List],
+                           texts: Union[str, List[str]],
+                           num_rows: int = 1,
+                           offset_ratio: float = 0.02,
+                           font_scale: float = 1.0,
+                           display_image: bool = True) -> Image.Image:
+    """ Displays a list of images in a grid with texts below them. """
+    
+    # Ensure texts is a list
+    if isinstance(texts, str):
+        texts = [texts] * len(images)
+    
+    # Add text under each image
+    images_with_texts = [text_under_image(img, txt, font_scale=font_scale) for img, txt in zip(images, texts)]
+    
+    if type(images_with_texts) is list:
+        num_empty = len(images_with_texts) % num_rows
+    elif images_with_texts.ndim == 4:
+        num_empty = images_with_texts.shape[0] % num_rows
+    else:
+        images_with_texts = [images_with_texts]
+        num_empty = 0
+    
+    empty_images = np.ones(images_with_texts[0].shape, dtype=np.uint8) * 255
+    images_with_texts = [image.astype(np.uint8) for image in images_with_texts] + [empty_images] * num_empty
+    num_items = len(images_with_texts)
+
+    h, w, c = images_with_texts[0].shape
+    offset = int(h * offset_ratio)
+    num_cols = num_items // num_rows
+    image_ = np.ones((h * num_rows + offset * (num_rows - 1),
+                      w * num_cols + offset * (num_cols - 1), 3), dtype=np.uint8) * 255
+    for i in range(num_rows):
+        for j in range(num_cols):
+            image_[i * (h + offset): i * (h + offset) + h:, j * (w + offset): j * (w + offset) + w] = images_with_texts[
+                i * num_cols + j]
+
+    pil_img = Image.fromarray(image_)
+    if display_image:
+        display(pil_img)
+    return pil_img
+
+
+
+class AttentionControl(abc.ABC):
+    
+    def step_callback(self, x_t):
+        return x_t
+    
+    def between_steps(self):
+        return
+    
+    @property
+    def num_uncond_att_layers(self):
+        return 0
+    
+    @abc.abstractmethod
+    def forward (self, attn, is_cross: bool, place_in_unet: str):
+        raise NotImplementedError
+
+    def __call__(self, attn, is_cross: bool, place_in_unet: str):
+        if self.cur_att_layer >= self.num_uncond_att_layers:
+            h = attn.shape[0]   
+            attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet) 
+        self.cur_att_layer += 1
+        if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
+            self.cur_att_layer = 0
+            self.cur_step += 1
+            self.between_steps()
+        return attn
+    
+    def reset(self):
+        self.cur_step = 0
+        self.cur_att_layer = 0
+
+    def __init__(self):
+        self.cur_step = 0
+        self.num_att_layers = -1
+        self.cur_att_layer = 0
+
+
+class EmptyControl(AttentionControl):
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        return attn
+
+
+class AttentionStore(AttentionControl):
+
+    @staticmethod
+    def get_empty_store():
+        return {"down_cross": [], "mid_cross": [], "up_cross": [],
+                "down_self": [],  "mid_self": [],  "up_self": []}
+
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
+        if attn.shape[1] <= 32 ** 2:  # avoid memory overhead
+            self.step_store[key].append(attn)
+        return attn
+
+    def between_steps(self):
+        if len(self.attention_store) == 0:
+            self.attention_store = self.step_store
+        else:
+            for key in self.attention_store:
+                for i in range(len(self.attention_store[key])):
+                    self.attention_store[key][i] += self.step_store[key][i]
+        self.step_store = self.get_empty_store()
+
+    def get_average_attention(self):
+        average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store}
+        return average_attention
+
+
+    def reset(self):
+        super(AttentionStore, self).reset()
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+
+    def __init__(self):
+        super(AttentionStore, self).__init__()
+        self.step_store = self.get_empty_store()
+        self.attention_store = {}
+
+class LocalBlend:
+
+    def __call__(self, x_t, attention_store):
+        k = 1
+        maps = attention_store["down_cross"][2:4] + attention_store["up_cross"][:3]
+        maps = [item.reshape(self.alpha_layers.shape[0], -1, 1, 16, 16, self.max_num_words) for item in maps]
+        maps = torch.cat(maps, dim=1)
+        maps = (maps * self.alpha_layers).sum(-1).mean(1)
+        mask = nnf.max_pool2d(maps, (k * 2 + 1, k * 2 +1), (1, 1), padding=(k, k))
+        mask = nnf.interpolate(mask, size=(x_t.shape[2:]))
+        mask = mask / mask.max(2, keepdims=True)[0].max(3, keepdims=True)[0]
+        mask = mask.gt(self.threshold)
+        mask = (mask[:1] + mask[1:]).float()
+        x_t = x_t[:1] + mask * (x_t - x_t[:1])
+        return x_t
+       
+    def __init__(self, prompts: List[str], words: [List[List[str]]], tokenizer, device, dtype=torch.float32, threshold=.3, max_num_words=77):
+        self.max_num_words = 77
+        
+        alpha_layers = torch.zeros(len(prompts),  1, 1, 1, 1, self.max_num_words)
+        for i, (prompt, words_) in enumerate(zip(prompts, words)):
+            if type(words_) is str:
+                words_ = [words_]
+            for word in words_:
+                ind = get_word_inds(prompt, word, tokenizer)
+                alpha_layers[i, :, :, :, :, ind] = 1
+        self.alpha_layers = alpha_layers.to(device, dtype)
+        self.threshold = threshold
+
+class AttentionControlEdit(AttentionStore, abc.ABC):
+    
+    def step_callback(self, x_t):
+        if self.local_blend is not None:
+            x_t = self.local_blend(x_t, self.attention_store)
+        return x_t
+        
+    def replace_self_attention(self, attn_base, att_replace):
+        if att_replace.shape[2] <= 16 ** 2:
+            return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
+        else:
+            return att_replace
+    
+    @abc.abstractmethod
+    def replace_cross_attention(self, attn_base, att_replace):
+        raise NotImplementedError
+    
+    def forward(self, attn, is_cross: bool, place_in_unet: str):
+        super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
+        # FIXME not replace correctly
+        if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]):
+            h = attn.shape[0] // (self.batch_size)
+            attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
+            attn_base, attn_repalce = attn[0], attn[1:]
+            if is_cross:
+                alpha_words = self.cross_replace_alpha[self.cur_step]
+                attn_repalce_new = self.replace_cross_attention(attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce
+                attn[1:] = attn_repalce_new
+            else:
+                attn[1:] = self.replace_self_attention(attn_base, attn_repalce)
+            attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
+        return attn
+    
+    def __init__(self, prompts, num_steps: int,
+                 cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]],
+                 self_replace_steps: Union[float, Tuple[float, float]],
+                 local_blend: Optional[LocalBlend],
+                 tokenizer,
+                 device,
+                 dtype):
+        super(AttentionControlEdit, self).__init__()
+        # add tokenizer and device here
+
+        self.tokenizer = tokenizer
+        self.device = device
+        self.dtype = dtype
+
+        self.batch_size = len(prompts)
+        self.cross_replace_alpha = get_time_words_attention_alpha(prompts, num_steps, cross_replace_steps, self.tokenizer).to(self.device, self.dtype)
+        if type(self_replace_steps) is float:
+            self_replace_steps = 0, self_replace_steps
+        self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1])
+        self.local_blend = local_blend  # 在外面定义后传进来
+
+class AttentionReplace(AttentionControlEdit):
+
+    def replace_cross_attention(self, attn_base, att_replace):
+        return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper)
+      
+    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
+                 local_blend: Optional[LocalBlend] = None, tokenizer=None, device=None, dtype=torch.float32):
+        super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend, tokenizer, device, dtype)
+        self.mapper = get_replacement_mapper(prompts, self.tokenizer).to(self.device, dtype=dtype)
+        
+
+class AttentionRefine(AttentionControlEdit):
+
+    def replace_cross_attention(self, attn_base, att_replace):
+        attn_base_replace = attn_base[:, :, self.mapper].permute(2, 0, 1, 3)
+        attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
+        return attn_replace
+
+    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float,
+                 local_blend: Optional[LocalBlend] = None, tokenizer=None, device=None, dtype=torch.float32):
+        super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend, tokenizer, device, dtype)
+        self.mapper, alphas = get_refinement_mapper(prompts, self.tokenizer)
+        self.mapper, alphas = self.mapper.to(self.device, self.dtype), alphas.to(self.device, self.dtype)
+        self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
+
+class AttentionReweight(AttentionControlEdit):
+
+    def replace_cross_attention(self, attn_base, att_replace):
+        if self.prev_controller is not None:
+            attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace)
+        attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
+        return attn_replace
+
+    def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer,
+                local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None, tokenizer=None, device=None, dtype=torch.float32):
+        super(AttentionReweight, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend, tokenizer, device, dtype)
+        self.equalizer = equalizer.to(self.device, self.dtype)
+        self.prev_controller = controller
+
+
+def get_equalizer(text: str, word_select: Union[int, Tuple[int, ...]], values: Union[List[float], Tuple[float, ...]], tokenizer):
+    if type(word_select) is int or type(word_select) is str:
+        word_select = (word_select,)
+    equalizer = torch.ones(len(values), 77)
+    values = torch.tensor(values, dtype=torch.float32)
+    for word in word_select:
+        inds = get_word_inds(text, word, tokenizer)
+        equalizer[:, inds] = values
+    return equalizer
+
+
+def update_alpha_time_word(alpha, bounds: Union[float, Tuple[float, float]], prompt_ind: int,
+                           word_inds: Optional[torch.Tensor]=None):
+    if type(bounds) is float:
+        bounds = 0, bounds
+    start, end = int(bounds[0] * alpha.shape[0]), int(bounds[1] * alpha.shape[0])
+    if word_inds is None:
+        word_inds = torch.arange(alpha.shape[2])
+    alpha[: start, prompt_ind, word_inds] = 0
+    alpha[start: end, prompt_ind, word_inds] = 1
+    alpha[end:, prompt_ind, word_inds] = 0
+    return alpha
+
+def get_time_words_attention_alpha(prompts, num_steps,
+                                   cross_replace_steps: Union[float, Dict[str, Tuple[float, float]]],
+                                   tokenizer, max_num_words=77):
+    if type(cross_replace_steps) is not dict:
+        cross_replace_steps = {"default_": cross_replace_steps}
+    if "default_" not in cross_replace_steps:
+        cross_replace_steps["default_"] = (0., 1.)
+    alpha_time_words = torch.zeros(num_steps + 1, len(prompts) - 1, max_num_words)
+    for i in range(len(prompts) - 1):
+        alpha_time_words = update_alpha_time_word(alpha_time_words, cross_replace_steps["default_"],
+                                                  i)
+    for key, item in cross_replace_steps.items():
+        if key != "default_":
+             inds = [get_word_inds(prompts[i], key, tokenizer) for i in range(1, len(prompts))]
+             for i, ind in enumerate(inds):
+                 if len(ind) > 0:
+                    alpha_time_words = update_alpha_time_word(alpha_time_words, item, i, ind)
+    alpha_time_words = alpha_time_words.reshape(num_steps + 1, len(prompts) - 1, 1, 1, max_num_words)
+    return alpha_time_words
+
+
+
+# seg_alinger
+class ScoreParams:
+
+    def __init__(self, gap, match, mismatch):
+        self.gap = gap
+        self.match = match
+        self.mismatch = mismatch
+
+    def mis_match_char(self, x, y):
+        if x != y:
+            return self.mismatch
+        else:
+            return self.match
+        
+    
+def get_matrix(size_x, size_y, gap):
+    matrix = []
+    for i in range(len(size_x) + 1):
+        sub_matrix = []
+        for j in range(len(size_y) + 1):
+            sub_matrix.append(0)
+        matrix.append(sub_matrix)
+    for j in range(1, len(size_y) + 1):
+        matrix[0][j] = j*gap
+    for i in range(1, len(size_x) + 1):
+        matrix[i][0] = i*gap
+    return matrix
+
+
+def get_matrix(size_x, size_y, gap):
+    matrix = np.zeros((size_x + 1, size_y + 1), dtype=np.int32)
+    matrix[0, 1:] = (np.arange(size_y) + 1) * gap
+    matrix[1:, 0] = (np.arange(size_x) + 1) * gap
+    return matrix
+
+
+def get_traceback_matrix(size_x, size_y):
+    matrix = np.zeros((size_x + 1, size_y +1), dtype=np.int32)
+    matrix[0, 1:] = 1
+    matrix[1:, 0] = 2
+    matrix[0, 0] = 4
+    return matrix
+
+
+def global_align(x, y, score):
+    matrix = get_matrix(len(x), len(y), score.gap)
+    trace_back = get_traceback_matrix(len(x), len(y))
+    for i in range(1, len(x) + 1):
+        for j in range(1, len(y) + 1):
+            left = matrix[i, j - 1] + score.gap
+            up = matrix[i - 1, j] + score.gap
+            diag = matrix[i - 1, j - 1] + score.mis_match_char(x[i - 1], y[j - 1])
+            matrix[i, j] = max(left, up, diag)
+            if matrix[i, j] == left:
+                trace_back[i, j] = 1
+            elif matrix[i, j] == up:
+                trace_back[i, j] = 2
+            else:
+                trace_back[i, j] = 3
+    return matrix, trace_back
+
+
+def get_aligned_sequences(x, y, trace_back):
+    x_seq = []
+    y_seq = []
+    i = len(x)
+    j = len(y)
+    mapper_y_to_x = []
+    while i > 0 or j > 0:
+        if trace_back[i, j] == 3:
+            x_seq.append(x[i-1])
+            y_seq.append(y[j-1])
+            i = i-1
+            j = j-1
+            mapper_y_to_x.append((j, i))
+        elif trace_back[i][j] == 1:
+            x_seq.append('-')
+            y_seq.append(y[j-1])
+            j = j-1
+            mapper_y_to_x.append((j, -1))
+        elif trace_back[i][j] == 2:
+            x_seq.append(x[i-1])
+            y_seq.append('-')
+            i = i-1
+        elif trace_back[i][j] == 4:
+            break
+    mapper_y_to_x.reverse()
+    return x_seq, y_seq, torch.tensor(mapper_y_to_x, dtype=torch.int64)
+
+
+def get_mapper(x: str, y: str, tokenizer, max_len=77):
+    x_seq = tokenizer.encode(x)
+    y_seq = tokenizer.encode(y)
+    score = ScoreParams(0, 1, -1)
+    matrix, trace_back = global_align(x_seq, y_seq, score)
+    mapper_base = get_aligned_sequences(x_seq, y_seq, trace_back)[-1]
+    alphas = torch.ones(max_len)
+    alphas[: mapper_base.shape[0]] = mapper_base[:, 1].ne(-1).float()
+    mapper = torch.zeros(max_len, dtype=torch.int64)
+    mapper[:mapper_base.shape[0]] = mapper_base[:, 1]
+    mapper[mapper_base.shape[0]:] = len(y_seq) + torch.arange(max_len - len(y_seq))
+    return mapper, alphas
+
+
+def get_refinement_mapper(prompts, tokenizer, max_len=77):
+    x_seq = prompts[0]
+    mappers, alphas = [], []
+    for i in range(1, len(prompts)):
+        mapper, alpha = get_mapper(x_seq, prompts[i], tokenizer, max_len)
+        mappers.append(mapper)
+        alphas.append(alpha)
+    return torch.stack(mappers), torch.stack(alphas)
+
+
+def get_word_inds(text: str, word_place: int, tokenizer):
+    split_text = text.split(" ")
+    if type(word_place) is str:
+        word_place = [i for i, word in enumerate(split_text) if word_place == word]
+    elif type(word_place) is int:
+        word_place = [word_place]
+    out = []
+    if len(word_place) > 0:
+        words_encode = [tokenizer.decode([item]).strip("#") for item in tokenizer.encode(text)][1:-1]
+        cur_len, ptr = 0, 0
+
+        for i in range(len(words_encode)):
+            cur_len += len(words_encode[i])
+            if ptr in word_place:
+                out.append(i + 1)
+            if cur_len >= len(split_text[ptr]):
+                ptr += 1
+                cur_len = 0
+    return np.array(out)
+
+
+def get_replacement_mapper_(x: str, y: str, tokenizer, max_len=77):
+    words_x = x.split(' ')
+    words_y = y.split(' ')
+    if len(words_x) != len(words_y):
+        raise ValueError(f"attention replacement edit can only be applied on prompts with the same length"
+                         f" but prompt A has {len(words_x)} words and prompt B has {len(words_y)} words.")
+    inds_replace = [i for i in range(len(words_y)) if words_y[i] != words_x[i]]
+    inds_source = [get_word_inds(x, i, tokenizer) for i in inds_replace]
+    inds_target = [get_word_inds(y, i, tokenizer) for i in inds_replace]
+    mapper = np.zeros((max_len, max_len))
+    i = j = 0
+    cur_inds = 0
+    while i < max_len and j < max_len:
+        if cur_inds < len(inds_source) and inds_source[cur_inds][0] == i:
+            inds_source_, inds_target_ = inds_source[cur_inds], inds_target[cur_inds]
+            if len(inds_source_) == len(inds_target_):
+                mapper[inds_source_, inds_target_] = 1
+            else:
+                ratio = 1 / len(inds_target_)
+                for i_t in inds_target_:
+                    mapper[inds_source_, i_t] = ratio
+            cur_inds += 1
+            i += len(inds_source_)
+            j += len(inds_target_)
+        elif cur_inds < len(inds_source):
+            mapper[i, j] = 1
+            i += 1
+            j += 1
+        else:
+            mapper[j, j] = 1
+            i += 1
+            j += 1
+
+    return torch.from_numpy(mapper).float()
+
+
+def get_replacement_mapper(prompts, tokenizer, max_len=77):
+    x_seq = prompts[0]
+    mappers = []
+    for i in range(1, len(prompts)):
+        mapper = get_replacement_mapper_(x_seq, prompts[i], tokenizer, max_len)
+        mappers.append(mapper)
+    return torch.stack(mappers)

requirements.txt

@@ -0,0 +1,19 @@
+
+accelerate
+diffusers==0.18.2
+einops
+huggingface_hub==0.16.4
+ipykernel
+librosa==0.9.2
+matplotlib
+numpy
+tqdm
+scikit-learn
+scipy
+soundfile==0.12.1
+tokenizers==0.13.3
+torch==2.0.1
+torchaudio==2.0.2
+torchvision==0.15.2
+transformers
+xformers==0.0.20

utils.py

@@ -0,0 +1,319 @@
+import os, json
+import math, random
+from multiprocessing import Pool
+from tqdm import tqdm
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+import matplotlib.pyplot as plt
+from torchvision import transforms
+from transformers import CLIPTextModel
+from transformers import PretrainedConfig
+
+
+def pad_spec(spec, spec_length, pad_value=0, random_crop=True): # spec: [3, mel_dim, spec_len]
+    assert spec_length % 8 == 0, "spec_length must be divisible by 8"
+    if spec.shape[-1] < spec_length:
+        # pad spec to spec_length
+        spec = F.pad(spec, (0, spec_length - spec.shape[-1]), value=pad_value)
+    else:
+        # random crop
+        if random_crop:
+            start = random.randint(0, spec.shape[-1] - spec_length)
+            spec = spec[:, :, start:start+spec_length]
+        else:
+            spec = spec[:, :, :spec_length]
+    return spec
+
+
+def load_spec(spec_path):
+    if spec_path.endswith(".pt"):
+        spec = torch.load(spec_path, map_location="cpu")
+    elif spec_path.endswith(".npy"):
+        spec = torch.from_numpy(np.load(spec_path))
+    else:
+        raise ValueError(f"Unknown spec file type {spec_path}")
+    assert len(spec.shape) == 3, f"spec shape must be [3, mel_dim, spec_len], got {spec.shape}"
+    if spec.size(0) == 1:
+        spec = spec.repeat(3, 1, 1)
+    return spec
+
+
+def random_crop_spec(spec, target_spec_length, pad_value=0, frame_per_sec=100, time_step=5): # spec: [3, mel_dim, spec_len]
+    assert target_spec_length % 8 == 0, "spec_length must be divisible by 8"
+
+    spec_length = spec.shape[-1]
+    full_s = math.ceil(spec_length / frame_per_sec / time_step) * time_step # get full seconds(ceil)
+    start_s = random.randint(0, math.floor(spec_length / frame_per_sec / time_step)) * time_step # random get start seconds
+
+    end_s = min(start_s + math.ceil(target_spec_length / frame_per_sec), full_s) # get end seconds
+
+    spec = spec[:, :, start_s * frame_per_sec : end_s * frame_per_sec] # get spec in seconds(crop more than target_spec_length because ceiling)
+
+    if spec.shape[-1] < target_spec_length:
+        spec = F.pad(spec, (0, target_spec_length - spec.shape[-1]), value=pad_value) # pad to target_spec_length
+    else:     
+        spec = spec[:, :, :target_spec_length] # crop to target_spec_length
+
+    return spec, int(start_s), int(end_s), int(full_s)
+
+
+
+def load_condion_embed(text_embed_path):
+    if text_embed_path.endswith(".pt"):
+        text_embed_list = torch.load(text_embed_path, map_location="cpu")
+    elif text_embed_path.endswith(".npy"):
+        text_embed_list = torch.from_numpy(np.load(text_embed_path))
+    else:
+        raise ValueError(f"Unknown text embedding file type {text_embed_path}")
+    if type(text_embed_list) == list:
+        text_embed = random.choice(text_embed_list)
+    if len(text_embed.shape) == 3: # [1, text_len, text_dim]
+        text_embed = text_embed.squeeze(0) # random choice and return text_emb: [text_len, text_dim]
+    return text_embed.detach().cpu()
+    
+
+def process_condition_embed(cond_emb, max_length): # [text_len, text_dim], Padding 0 and random drop by CFG
+    if cond_emb.shape[0] < max_length:
+        cond_emb = F.pad(cond_emb, (0, 0, 0, max_length - cond_emb.shape[0]), value=0)
+    else:
+        cond_emb = cond_emb[:max_length, :]
+    return cond_emb
+                
+
+def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+        return CLIPTextModel
+    if "t5" in model_class.lower():
+        from transformers import T5EncoderModel
+        return T5EncoderModel
+    if "clap" in model_class.lower():
+        from transformers import ClapTextModelWithProjection
+        return ClapTextModelWithProjection
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+    
+
+
+def str2bool(string):
+    str2val = {"True": True, "False": False, "true": True, "false": False, "none": False, "None": False}
+    if string in str2val:
+        return str2val[string]
+    else:
+        raise ValueError(f"Expected one of {set(str2val.keys())}, got {string}")
+    
+
+def str2str(string):
+    if string.lower() == "none" or string.lower() == "null" or string.lower() == "false" or string == "":
+        return None
+    else:
+        return string    
+
+
+def json_dump(data_json, json_save_path):
+    with open(json_save_path, 'w') as f:
+        json.dump(data_json, f, indent=4)
+        f.close()
+
+
+def json_load(json_path):
+    with open(json_path, 'r') as f:
+        data = json.load(f)
+        f.close()
+    return data
+
+
+def load_json_list(path):
+    with open(path, 'r', encoding='utf-8') as f:
+        return [json.loads(line) for line in f.readlines()]
+    
+
+def save_json_list(data, path):
+    with open(path, 'w', encoding='utf-8') as f:
+        for d in data:
+            f.write(json.dumps(d) + '\n')
+    
+
+def multiprocess_function(func, func_args, n_jobs=32):  
+    with Pool(processes=n_jobs) as p:
+            with tqdm(total=len(func_args)) as pbar:
+                for i, _ in enumerate(p.imap_unordered(func, func_args)):
+                    pbar.update()
+
+
+def image_add_color(spec_img):
+    cmap = plt.get_cmap('viridis')
+    cmap_r = cmap.reversed()
+    image = cmap(np.array(spec_img)[:,:,0])[:, :, :3]  # 省略透明度通道
+    image = (image - image.min()) / (image.max() - image.min())
+    image = Image.fromarray(np.uint8(image*255))
+    return image
+
+
+@staticmethod
+def pt_to_numpy(images: torch.FloatTensor) -> np.ndarray:
+    """
+    Convert a PyTorch tensor to a NumPy image.
+    """
+    images = images.cpu().permute(0, 2, 3, 1).float().numpy()
+    return images
+
+
+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
+
+### CODE FOR INPAITING ###
+def normalize(images):
+    """
+    Normalize an image array to [-1,1].
+    """
+    if images.min() >= 0:
+        return 2.0 * images - 1.0
+    else:
+        return images
+
+def denormalize(images):
+    """
+    Denormalize an image array to [0,1].
+    """
+    if images.min() < 0:
+        return (images / 2 + 0.5).clamp(0, 1)
+    else:
+        return images.clamp(0, 1)     
+    
+
+def prepare_mask_and_masked_image(image, mask):
+    """
+    Prepare a binary mask and the masked image.
+    
+    Parameters:
+    - image (torch.Tensor): The input image tensor of shape [3, height, width] with values in the range [0, 1].
+    - mask (torch.Tensor): The input mask tensor of shape [1, height, width].
+    
+    Returns:
+    - tuple: A tuple containing the binary mask and the masked image.
+    """
+    # Noralize image to [0,1]
+    if image.max() > 1:
+        image = (image - image.min()) / (image.max() - image.min())
+    # Normalize image from [0,1] to [-1,1]
+    if image.min() >= 0:
+        image = normalize(image)    
+    # Apply the mask to the image
+    masked_image = image * (mask < 0.5)
+    
+    return mask, masked_image
+
+
+def torch_to_pil(image):
+    """
+    Convert a torch tensor to a PIL image.
+    """
+    if image.min() < 0:
+        image = denormalize(image)
+
+    return transforms.ToPILImage()(image.cpu().detach().squeeze())
+
+
+
+# class TextEncoderAdapter(nn.Module):    
+#     def __init__(self, hidden_size, cross_attention_dim=768):
+#         super(TextEncoderAdapter, self).__init__()
+#         self.hidden_size = hidden_size
+#         self.cross_attention_dim = cross_attention_dim
+#         self.proj = nn.Linear(self.hidden_size, self.cross_attention_dim)
+#         self.norm = torch.nn.LayerNorm(self.cross_attention_dim)
+
+#     def forward(self, x):
+#         x = self.proj(x)
+#         x = self.norm(x)
+#         return x
+    
+#     def save_pretrained(self, save_directory, subfolder=""):
+#         if subfolder:
+#             save_directory = os.path.join(save_directory, subfolder)
+#         os.makedirs(save_directory, exist_ok=True)
+#         ckpt_path = os.path.join(save_directory, "adapter.pt")
+#         config_path = os.path.join(save_directory, "config.json")
+#         config = {"hidden_size": self.hidden_size, "cross_attention_dim": self.cross_attention_dim}
+#         json_dump(config, config_path)
+#         torch.save(self.state_dict(), ckpt_path)
+#         print(f"Saving adapter model to {ckpt_path}")
+
+#     @classmethod
+#     def from_pretrained(cls, load_directory, subfolder=""):
+#         if subfolder:
+#             load_directory = os.path.join(load_directory, subfolder)
+#         ckpt_path = os.path.join(load_directory, "adapter.pt")
+#         config_path = os.path.join(load_directory, "config.json")
+#         config = json_load(config_path)
+#         instance = cls(**config)
+#         instance.load_state_dict(torch.load(ckpt_path))
+#         print(f"Loading adapter model from {ckpt_path}")
+#         return instance          
+
+
+
+class ConditionAdapter(nn.Module):
+    def __init__(self, config):
+        super(ConditionAdapter, self).__init__()
+        self.config = config
+        self.proj = nn.Linear(self.config["condition_dim"], self.config["cross_attention_dim"])
+        self.norm = torch.nn.LayerNorm(self.config["cross_attention_dim"])
+        print(f"INITIATED: ConditionAdapter: {self.config}")
+
+    def forward(self, x):
+        x = self.proj(x)
+        x = self.norm(x)
+        return x
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path):
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt")
+        config = json_load(config_path)
+        instance = cls(config)
+        instance.load_state_dict(torch.load(ckpt_path))
+        print(f"LOADED: ConditionAdapter from {pretrained_model_name_or_path}")
+        return instance
+
+    def save_pretrained(self, pretrained_model_name_or_path):
+        os.makedirs(pretrained_model_name_or_path, exist_ok=True)
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        ckpt_path = os.path.join(pretrained_model_name_or_path, "condition_adapter.pt")        
+        json_dump(self.config, config_path)
+        torch.save(self.state_dict(), ckpt_path)
+        print(f"SAVED: ConditionAdapter {self.config['condition_adapter_name']} to {pretrained_model_name_or_path}")
+
+
+# class TextEncoderWrapper(CLIPTextModel):
+#     def __init__(self, text_encoder, text_encoder_adapter):
+#         super().__init__(text_encoder.config)
+#         self.text_encoder = text_encoder
+#         self.adapter = text_encoder_adapter
+
+#     def forward(self, input_ids, **kwargs):
+#         outputs = self.text_encoder(input_ids, **kwargs)
+#         adapted_output = self.adapter(outputs[0])
+#         return [adapted_output] # to compatible with last_hidden_state
+