added model

CITATION.cff

@@ -0,0 +1,10 @@
+cff-version: 1.3.0
+message: "If you use this software, please cite it as below."
+authors:
+- family-names: "Betker"
+  given-names: "James"
+  orcid: "https://orcid.org/my-orcid?orcid=0000-0003-3259-4862"
+title: "TorToiSe text-to-speech"
+version: 2.0
+date-released: 2022-04-28
+url: "https://github.com/neonbjb/tortoise-tts"

LICENSE

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

app.py

@@ -0,0 +1,306 @@
+# AGPL: a notification must be added stating that changes have been made to that file.
+
+import os
+import shutil
+from pathlib import Path
+
+import streamlit as st
+from random import randint
+
+from tortoise.api import MODELS_DIR
+from tortoise.inference import (
+    infer_on_texts,
+    run_and_save_tts,
+    split_and_recombine_text,
+)
+from tortoise.utils.diffusion import SAMPLERS
+from app_utils.filepicker import st_file_selector
+from app_utils.conf import TortoiseConfig
+
+from app_utils.funcs import (
+    timeit,
+    load_model,
+    list_voices,
+    load_voice_conditionings,
+)
+
+
+LATENT_MODES = [
+    "Tortoise original (bad)",
+    "average per 4.27s (broken on small files)",
+    "average per voice file (broken on small files)",
+]
+
+def main():
+    conf = TortoiseConfig()
+    
+    with st.expander("Create New Voice", expanded=True):
+        if "file_uploader_key" not in st.session_state:
+            st.session_state["file_uploader_key"] = str(randint(1000, 100000000))
+            st.session_state["text_input_key"] = str(randint(1000, 100000000))
+
+        uploaded_files = st.file_uploader(
+            "Upload Audio Samples for a New Voice",
+            accept_multiple_files=True,
+            type=["wav"],
+            key=st.session_state["file_uploader_key"]
+        )
+
+        voice_name = st.text_input(
+            "New Voice Name",
+            help="Enter a name for your new voice.",
+            value="",
+            key=st.session_state["text_input_key"]
+        )
+
+        create_voice_button = st.button(
+            "Create Voice",
+            disabled = ((voice_name.strip() == "") | (len(uploaded_files) == 0))
+        )
+        if create_voice_button:
+            st.write(st.session_state)
+            with st.spinner(f"Creating new voice: {voice_name}"):
+                new_voice_name = voice_name.strip().replace(" ", "_")
+
+                voices_dir = f'./tortoise/voices/{new_voice_name}/'
+                if os.path.exists(voices_dir):
+                    shutil.rmtree(voices_dir)
+                os.makedirs(voices_dir)
+
+                for index, uploaded_file in enumerate(uploaded_files):
+                    bytes_data = uploaded_file.read()
+                    with open(f"{voices_dir}voice_sample{index}.wav", "wb") as wav_file:
+                        wav_file.write(bytes_data)
+
+                st.session_state["text_input_key"] = str(randint(1000, 100000000))
+                st.session_state["file_uploader_key"] = str(randint(1000, 100000000))
+                st.experimental_rerun()
+
+    text = st.text_area(
+        "Text",
+        help="Text to speak.",
+        value="The expressiveness of autoregressive transformers is literally nuts! I absolutely adore them.",
+    )
+
+    voices = [v for v in os.listdir("tortoise/voices") if v != "cond_latent_example"]
+
+    voice = st.selectbox(
+        "Voice",
+        voices,
+        help="Selects the voice to use for generation. See options in voices/ directory (and add your own!) "
+        "Use the & character to join two voices together. Use a comma to perform inference on multiple voices.",
+        index=0,
+    )
+    preset = st.selectbox(
+        "Preset",
+        (
+            "single_sample",
+            "ultra_fast",
+            "very_fast",
+            "ultra_fast_old",
+            "fast",
+            "standard",
+            "high_quality",
+        ),
+        help="Which voice preset to use.",
+        index=1,
+    )
+    with st.expander("Advanced"):
+        col1, col2 = st.columns(2)
+        with col1:
+            """#### Model parameters"""
+            candidates = st.number_input(
+                "Candidates",
+                help="How many output candidates to produce per-voice.",
+                value=1,
+            )
+            latent_averaging_mode = st.radio(
+                "Latent averaging mode",
+                LATENT_MODES,
+                help="How voice samples should be averaged together.",
+                index=0,
+            )
+            sampler = st.radio(
+                "Sampler",
+                #SAMPLERS,
+                ["dpm++2m", "p", "ddim"],
+                help="Diffusion sampler. Note that dpm++2m is experimental and typically requires more steps.",
+                index=1,
+            )
+            steps = st.number_input(
+                "Steps",
+                help="Override the steps used for diffusion (default depends on preset)",
+                value=10,
+            )
+            seed = st.number_input(
+                "Seed",
+                help="Random seed which can be used to reproduce results.",
+                value=-1,
+            )
+            if seed == -1:
+                seed = None
+            voice_fixer = st.checkbox(
+                "Voice fixer",
+                help="Use `voicefixer` to improve audio quality. This is a post-processing step which can be applied to any output.",
+                value=True,
+            )
+            """#### Directories"""
+            output_path = st.text_input(
+                "Output Path", help="Where to store outputs.", value="results/"
+            )
+
+        with col2:
+            """#### Optimizations"""
+            high_vram = not st.checkbox(
+                "Low VRAM",
+                help="Re-enable default offloading behaviour of tortoise",
+                value=True,
+            )
+            half = st.checkbox(
+                "Half-Precision",
+                help="Enable autocast to half precision for autoregressive model",
+                value=False,
+            )
+            kv_cache = st.checkbox(
+                "Key-Value Cache",
+                help="Enable kv_cache usage, leading to drastic speedups but worse memory usage",
+                value=True,
+            )
+            cond_free = st.checkbox(
+                "Conditioning Free",
+                help="Force conditioning free diffusion",
+                value=True,
+            )
+            no_cond_free = st.checkbox(
+                "Force Not Conditioning Free",
+                help="Force disable conditioning free diffusion",
+                value=False,
+            )
+
+            """#### Text Splitting"""
+            min_chars_to_split = st.number_input(
+                "Min Chars to Split",
+                help="Minimum number of characters to split text on",
+                min_value=50,
+                value=200,
+                step=1,
+            )
+
+            """#### Debug"""
+            produce_debug_state = st.checkbox(
+                "Produce Debug State",
+                help="Whether or not to produce debug_state.pth, which can aid in reproducing problems. Defaults to true.",
+                value=True,
+            )
+
+    ar_checkpoint = "."
+    diff_checkpoint = "." 
+    if st.button("Update Basic Settings"):
+        conf.update(
+            EXTRA_VOICES_DIR=extra_voices_dir,
+            LOW_VRAM=not high_vram,
+            AR_CHECKPOINT=ar_checkpoint,
+            DIFF_CHECKPOINT=diff_checkpoint,
+        )
+
+    ar_checkpoint = None
+    diff_checkpoint = None
+    tts = load_model(MODELS_DIR, high_vram, kv_cache, ar_checkpoint, diff_checkpoint)
+
+    if st.button("Start"):
+        assert latent_averaging_mode
+        assert preset
+        assert voice
+
+        def show_generation(fp, filename: str):
+            """
+            audio_buffer = BytesIO()
+            save_gen_with_voicefix(g, audio_buffer, squeeze=False)
+            torchaudio.save(audio_buffer, g, 24000, format='wav')
+            """
+            st.audio(str(fp), format="audio/wav")
+            st.download_button(
+                "Download sample",
+                str(fp),
+                file_name=filename,  # this doesn't actually seem to work lol
+            )
+
+        with st.spinner(
+            f"Generating {candidates} candidates for voice {voice} (seed={seed}). You can see progress in the terminal"
+        ):
+            os.makedirs(output_path, exist_ok=True)
+
+            selected_voices = voice.split(",")
+            for k, selected_voice in enumerate(selected_voices):
+                if "&" in selected_voice:
+                    voice_sel = selected_voice.split("&")
+                else:
+                    voice_sel = [selected_voice]
+                voice_samples, conditioning_latents = load_voice_conditionings(
+                    voice_sel, []
+                )
+
+                voice_path = Path(os.path.join(output_path, selected_voice))
+
+                with timeit(
+                    f"Generating {candidates} candidates for voice {selected_voice} (seed={seed})"
+                ):
+                    nullable_kwargs = {
+                        k: v
+                        for k, v in zip(
+                            ["sampler", "diffusion_iterations", "cond_free"],
+                            [sampler, steps, cond_free],
+                        )
+                        if v is not None
+                    }
+
+                    def call_tts(text: str):
+                        return tts.tts_with_preset(
+                            text,
+                            k=candidates,
+                            voice_samples=voice_samples,
+                            conditioning_latents=conditioning_latents,
+                            preset=preset,
+                            use_deterministic_seed=seed,
+                            return_deterministic_state=True,
+                            cvvp_amount=0.0,
+                            half=half,
+                            latent_averaging_mode=LATENT_MODES.index(
+                                latent_averaging_mode
+                            ),
+                            **nullable_kwargs,
+                        )
+
+                    if len(text) < min_chars_to_split:
+                        filepaths = run_and_save_tts(
+                            call_tts,
+                            text,
+                            voice_path,
+                            return_deterministic_state=True,
+                            return_filepaths=True,
+                            voicefixer=voice_fixer,
+                        )
+                        for i, fp in enumerate(filepaths):
+                            show_generation(fp, f"{selected_voice}-text-{i}.wav")
+                    else:
+                        desired_length = int(min_chars_to_split)
+                        texts = split_and_recombine_text(
+                            text, desired_length, desired_length + 100
+                        )
+                        filepaths = infer_on_texts(
+                            call_tts,
+                            texts,
+                            voice_path,
+                            return_deterministic_state=True,
+                            return_filepaths=True,
+                            lines_to_regen=set(range(len(texts))),
+                            voicefixer=voice_fixer,
+                        )
+                        for i, fp in enumerate(filepaths):
+                            show_generation(fp, f"{selected_voice}-text-{i}.wav")
+        if produce_debug_state:
+            """Debug states can be found in the output directory"""
+
+
+if __name__ == "__main__":
+    main()

app_utils/conf.py

@@ -0,0 +1,40 @@
+import shelve
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class PersistentSettings(BaseModel):
+    """
+    This pydantic model will try to initialize itself from
+    the database upon every instantiation
+
+    It further supplies an update function, that allows to write
+    back any changes into the database, under its key.
+    """
+
+    def __init__(self, **data: Any):
+        with shelve.open("config.db") as db:
+            super().__init__(**db.get("settings", default={}), **data)
+
+    def update(self, **data: Any) -> None:
+        """
+        Persist the pydantic-dict that represents the model
+        """
+        with shelve.open("config.db") as db:
+            db["settings"] = {**self.dict(), **data}
+
+
+class TortoiseConfig(PersistentSettings):
+    EXTRA_VOICES_DIR: str = ""
+    AR_CHECKPOINT: str = "."
+    DIFF_CHECKPOINT: str = "."
+    LOW_VRAM: bool = True
+
+    def __init__(self, **data: Any):
+        super().__init__(**data)
+        if not Path(self.AR_CHECKPOINT).is_file():
+            self.AR_CHECKPOINT = "."
+        if not Path(self.DIFF_CHECKPOINT).is_file():
+            self.DIFF_CHECKPOINT = "."

app_utils/filepicker.py

@@ -0,0 +1,71 @@
+# taken from https://gist.github.com/benlansdell/44000c264d1b373c77497c0ea73f0ef2
+# slightly modified
+"""FilePicker for streamlit. 
+
+Still doesn't seem to be a good solution for a way to select files to process from the server Streamlit is running on.
+
+Here's a pretty functional solution. 
+
+Usage:
+
+```
+import streamlit as st
+from filepicker import st_file_selector
+
+tif_file = st_file_selector(st, key = 'tif', label = 'Choose tif file')
+```
+"""
+
+import os
+
+import streamlit as st
+
+
+def update_dir(key):
+    global i_will_regret_this, i_will_regret_this2
+    choice = st.session_state[key]
+    if os.path.isdir(os.path.join(st.session_state[key + "curr_dir"], choice)):
+        st.session_state[key + "index"] = 0
+        st.session_state[key + "curr_dir"] = os.path.normpath(
+            os.path.join(st.session_state[key + "curr_dir"], choice)
+        )
+        files = sorted(os.listdir(st.session_state[key + "curr_dir"]))
+        files.insert(0, "..")
+        files.insert(0, ".")
+        st.session_state[key + "files"] = files
+
+
+def st_file_selector(
+    st_placeholder, path=".", label="Select a file/folder", key="selected"
+):
+    if key + "curr_dir" not in st.session_state:
+        base_path = "." if path is None or path == "" else path
+        base_path = (
+            base_path if os.path.isdir(base_path) else os.path.dirname(base_path)
+        )
+        base_path = "." if base_path is None or base_path == "" else base_path
+
+        files = sorted(os.listdir(base_path))
+        files.insert(0, "..")
+        files.insert(0, ".")
+        st.session_state[key + "files"] = files
+        st.session_state[key + "curr_dir"] = base_path
+        st.session_state[key + "index"] = (
+            st.session_state[key + "files"].index(os.path.basename(path))
+            if os.path.isfile(path) and path[-4:] == '.pth'
+            else 0
+        )
+    else:
+        base_path = st.session_state[key + "curr_dir"]
+
+    selected_file = st_placeholder.selectbox(
+        label=label,
+        options=st.session_state[key + "files"],
+        index=st.session_state[key + "index"],
+        key=key,
+        on_change=lambda: update_dir(key),
+    )
+    selected_path = os.path.normpath(os.path.join(base_path, selected_file))
+    st_placeholder.write(os.path.abspath(selected_path))
+
+    return selected_path

app_utils/funcs.py

@@ -0,0 +1,57 @@
+import gc
+import os
+from contextlib import contextmanager
+from time import time
+from typing import Optional
+
+import streamlit as st
+
+from tortoise.api import TextToSpeech
+from tortoise.utils.audio import load_voices
+
+
+@contextmanager
+def timeit(desc=""):
+    start = time()
+    yield
+    print(f"{desc} took {time() - start:.2f} seconds")
+
+
+@st.cache_resource(max_entries=1)
+def load_model(
+    model_dir,
+    high_vram,
+    kv_cache,
+    ar_checkpoint,
+    diff_checkpoint,
+):
+    gc.collect()
+    return TextToSpeech(
+        models_dir=model_dir,
+        high_vram=high_vram,
+        kv_cache=kv_cache,
+        ar_checkpoint=ar_checkpoint,
+        diff_checkpoint=diff_checkpoint,
+    )
+
+
+@st.cache_data
+def list_voices(extra_voices_dir: Optional[str]):
+    voices = ["random"]
+    if extra_voices_dir and os.path.isdir(extra_voices_dir):
+        voices.extend(os.listdir(extra_voices_dir))
+        extra_voices_ls = [extra_voices_dir]
+    else:
+        extra_voices_ls = []
+    voices.extend(
+        [v for v in os.listdir("tortoise/voices") if v != "cond_latent_example"]
+    )
+    #
+    return voices, extra_voices_ls
+
+
+@st.cache_resource(max_entries=1)
+def load_voice_conditionings(voice, extra_voices_ls):
+    gc.collect()
+    voice_samples, conditioning_latents = load_voices(voice, extra_voices_ls)
+    return voice_samples, conditioning_latents

tortoise/api.py

@@ -0,0 +1,827 @@
+# ## AGPL: a notification must be added stating that changes have been made to that file.
+
+import os
+import random
+from time import time
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from tqdm import tqdm
+
+from tortoise.models.arch_util import TorchMelSpectrogram
+from tortoise.models.autoregressive import UnifiedVoice
+from tortoise.models.classifier import AudioMiniEncoderWithClassifierHead
+from tortoise.models.clvp import CLVP
+from tortoise.models.cvvp import CVVP
+from tortoise.models.diffusion_decoder import DiffusionTts
+from tortoise.models.random_latent_generator import RandomLatentConverter
+from tortoise.models.vocoder import VocConf
+from tortoise.utils.audio import denormalize_tacotron_mel, wav_to_univnet_mel
+from tortoise.utils.diffusion import (
+    SpacedDiffusion,
+    get_named_beta_schedule,
+    space_timesteps,
+)
+from tortoise.utils.tokenizer import VoiceBpeTokenizer
+from tortoise.utils.wav2vec_alignment import Wav2VecAlignment
+
+from tortoise.models.utils import MODELS_DIR, get_model_path
+
+from contextlib import contextmanager
+
+def pad_or_truncate(t, length):
+    """
+    Utility function for forcing <t> to have the specified sequence length, whether by clipping it or padding it with 0s.
+    """
+    if t.shape[-1] == length:
+        return t
+    elif t.shape[-1] < length:
+        return F.pad(t, (0, length - t.shape[-1]))
+    else:
+        return t[..., :length]
+
+
+def load_discrete_vocoder_diffuser(
+    trained_diffusion_steps=4000,
+    desired_diffusion_steps=200,
+    cond_free=True,
+    cond_free_k=1,
+    sampler="ddim",
+):
+    """
+    Helper function to load a GaussianDiffusion instance configured for use as a vocoder.
+    """
+    return SpacedDiffusion(
+        use_timesteps=space_timesteps(
+            trained_diffusion_steps, [desired_diffusion_steps]
+        ),
+        model_mean_type="epsilon",
+        model_var_type="learned_range",
+        loss_type="mse",
+        betas=get_named_beta_schedule("linear", trained_diffusion_steps),
+        conditioning_free=cond_free,
+        conditioning_free_k=cond_free_k,
+        sampler=sampler,
+    )
+
+
+def format_conditioning(clip, cond_length=132300, device="cuda"):
+    """
+    Converts the given conditioning signal to a MEL spectrogram and clips it as expected by the models.
+    """
+    gap = clip.shape[-1] - cond_length
+    if gap < 0:
+        clip = F.pad(clip, pad=(0, abs(gap)))
+    elif gap > 0:
+        rand_start = random.randint(0, gap)
+        clip = clip[:, rand_start : rand_start + cond_length]
+    mel_clip = TorchMelSpectrogram()(clip.unsqueeze(0)).squeeze(0)
+    return mel_clip.unsqueeze(0).to(device)
+
+
+def fix_autoregressive_output(codes, stop_token, complain=True):
+    """
+    This function performs some padding on coded audio that fixes a mismatch issue between what the diffusion model was
+    trained on and what the autoregressive code generator creates (which has no padding or end).
+    This is highly specific to the DVAE being used, so this particular coding will not necessarily work if used with
+    a different DVAE. This can be inferred by feeding a audio clip padded with lots of zeros on the end through the DVAE
+    and copying out the last few codes.
+
+    Failing to do this padding will produce speech with a harsh end that sounds like "BLAH" or similar.
+    """
+    # Strip off the autoregressive stop token and add padding.
+    stop_token_indices = (codes == stop_token).nonzero()
+    if len(stop_token_indices) == 0:
+        if complain:
+            print(
+                "No stop tokens found in one of the generated voice clips. This typically means the spoken audio is "
+                "too long. In some cases, the output will still be good, though. Listen to it and if it is missing words, "
+                "try breaking up your input text."
+            )
+        return codes
+    else:
+        codes[stop_token_indices] = 83
+    stm = stop_token_indices.min().item()
+    codes[stm:] = 83
+    if stm - 3 < codes.shape[0]:
+        codes[-3] = 45
+        codes[-2] = 45
+        codes[-1] = 248
+
+    return codes
+
+
+def do_spectrogram_diffusion(
+    diffusion_model,
+    diffuser,
+    latents,
+    conditioning_latents,
+    temperature=1,
+    verbose=True,
+):
+    """
+    Uses the specified diffusion model to convert discrete codes into a spectrogram.
+    """
+    with torch.no_grad():
+        output_seq_len = (
+            latents.shape[1] * 4 * 24000 // 22050
+        )  # This diffusion model converts from 22kHz spectrogram codes to a 24kHz spectrogram signal.
+        output_shape = (latents.shape[0], 100, output_seq_len)
+        precomputed_embeddings = diffusion_model.timestep_independent(
+            latents, conditioning_latents, output_seq_len, False
+        )
+
+        noise = torch.randn(output_shape, device=latents.device) * temperature
+        mel = diffuser.sample_loop(
+            diffusion_model,
+            output_shape,
+            noise=noise,
+            model_kwargs={"precomputed_aligned_embeddings": precomputed_embeddings},
+            progress=verbose,
+        )
+        return denormalize_tacotron_mel(mel)[:, :, :output_seq_len]
+
+
+def classify_audio_clip(clip):
+    """
+    Returns whether or not Tortoises' classifier thinks the given clip came from Tortoise.
+    :param clip: torch tensor containing audio waveform data (get it from load_audio)
+    :return: True if the clip was classified as coming from Tortoise and false if it was classified as real.
+    """
+    classifier = AudioMiniEncoderWithClassifierHead(
+        2,
+        spec_dim=1,
+        embedding_dim=512,
+        depth=5,
+        downsample_factor=4,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        base_channels=32,
+        dropout=0,
+        kernel_size=5,
+        distribute_zero_label=False,
+    )
+    classifier.load_state_dict(
+        torch.load(get_model_path("classifier.pth"), map_location=torch.device("cpu"))
+    )
+    clip = clip.cpu().unsqueeze(0)
+    results = F.softmax(classifier(clip), dim=-1)
+    return results[0][0]
+
+
+def pick_best_batch_size_for_gpu():
+    """
+    Tries to pick a batch size that will fit in your GPU. These sizes aren't guaranteed to work, but they should give
+    you a good shot.
+    """
+    if torch.cuda.is_available():
+        _, available = torch.cuda.mem_get_info()
+        availableGb = available / (1024**3)
+        if availableGb > 14:
+            return 16
+        elif availableGb > 10:
+            return 8
+        elif availableGb > 7:
+            return 4
+    return 1
+
+
+class TextToSpeech:
+    """
+    Main entry point into Tortoise.
+    """
+
+    def _config(self):
+        raise RuntimeError("This is depreciated")
+        return {
+            "high_vram": self.high_vram,
+            "models_dir": self.models_dir,
+            "kv_cache": self.autoregressive.inference_model.kv_cache,
+            "ar_checkpoint": self.ar_checkpoint,
+        }
+
+    def __init__(
+        self,
+        autoregressive_batch_size=None,
+        models_dir=MODELS_DIR,
+        enable_redaction=True,
+        device=None,
+        high_vram=False,
+        kv_cache=True,
+        ar_checkpoint=None,
+        clvp_checkpoint=None,
+        diff_checkpoint=None,
+        vocoder=VocConf.Univnet,
+    ):
+        """
+        Constructor
+        :param autoregressive_batch_size: Specifies how many samples to generate per batch. Lower this if you are seeing
+                                          GPU OOM errors. Larger numbers generates slightly faster.
+        :param models_dir: Where model weights are stored. This should only be specified if you are providing your own
+                           models, otherwise use the defaults.
+        :param enable_redaction: When true, text enclosed in brackets are automatically redacted from the spoken output
+                                 (but are still rendered by the model). This can be used for prompt engineering.
+                                 Default is true.
+        :param device: Device to use when running the model. If omitted, the device will be automatically chosen.
+        :param high_vram: If true, the model will use more VRAM but will run faster.
+        :param kv_cache: If true, the autoregressive model will cache key value attention pairs to speed up generation.
+        :param ar_checkpoint: Path to a checkpoint file for the autoregressive model. If omitted, uses default
+        :param clvp_checkpoint: Path to a checkpoint file for the CLVP model. If omitted, uses default
+        :param diff_checkpoint: Path to a checkpoint file for the diffusion model. If omitted, uses default
+        """
+        self.ar_checkpoint = ar_checkpoint
+        self.diff_checkpoint = diff_checkpoint  # TODO: check if this is even needed
+        self.models_dir = models_dir
+        self.autoregressive_batch_size = (
+            pick_best_batch_size_for_gpu()
+            if autoregressive_batch_size is None
+            else autoregressive_batch_size
+        )
+        self.enable_redaction = enable_redaction
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if self.enable_redaction:
+            self.aligner = Wav2VecAlignment()
+
+        self.tokenizer = VoiceBpeTokenizer()
+
+        if os.path.exists(f"{models_dir}/autoregressive.ptt"):
+            # Assume this is a traced directory.
+            self.autoregressive = torch.jit.load(f"{models_dir}/autoregressive.ptt")
+            self.diffusion = torch.jit.load(f"{models_dir}/diffusion_decoder.ptt")
+        else:
+            self.autoregressive = (
+                UnifiedVoice(
+                    max_mel_tokens=604,
+                    max_text_tokens=402,
+                    max_conditioning_inputs=2,
+                    layers=30,
+                    model_dim=1024,
+                    heads=16,
+                    number_text_tokens=255,
+                    start_text_token=255,
+                    checkpointing=False,
+                    train_solo_embeddings=False,
+                )
+                .cpu()
+                .eval()
+            )
+            ar_path = ar_checkpoint or get_model_path("autoregressive.pth", models_dir)
+            self.autoregressive.load_state_dict(torch.load(ar_path))
+            self.autoregressive.post_init_gpt2_config(kv_cache)
+
+            diff_path = diff_checkpoint or get_model_path(
+                "diffusion_decoder.pth", models_dir
+            )
+            self.diffusion = (
+                DiffusionTts(
+                    model_channels=1024,
+                    num_layers=10,
+                    in_channels=100,
+                    out_channels=200,
+                    in_latent_channels=1024,
+                    in_tokens=8193,
+                    dropout=0,
+                    use_fp16=False,
+                    num_heads=16,
+                    layer_drop=0,
+                    unconditioned_percentage=0,
+                )
+                .cpu()
+                .eval()
+            )
+            self.diffusion.load_state_dict(torch.load(diff_path))
+
+        self.clvp = (
+            CLVP(
+                dim_text=768,
+                dim_speech=768,
+                dim_latent=768,
+                num_text_tokens=256,
+                text_enc_depth=20,
+                text_seq_len=350,
+                text_heads=12,
+                num_speech_tokens=8192,
+                speech_enc_depth=20,
+                speech_heads=12,
+                speech_seq_len=430,
+                use_xformers=True,
+            )
+            .cpu()
+            .eval()
+        )
+        clvp_path = clvp_checkpoint or get_model_path("clvp2.pth", models_dir)
+        self.clvp.load_state_dict(torch.load(clvp_path))
+        self.cvvp = None  # CVVP model is only loaded if used.
+
+        self.vocoder = vocoder.value.constructor().cpu()
+        self.vocoder.load_state_dict(
+            vocoder.value.optionally_index(
+                torch.load(
+                    get_model_path(vocoder.value.model_path, models_dir),
+                    map_location=torch.device("cpu"),
+                )
+            )
+        )
+        self.vocoder.eval(inference=True)
+
+        # Random latent generators (RLGs) are loaded lazily.
+        self.rlg_auto = None
+        self.rlg_diffusion = None
+
+        if high_vram:
+            self.autoregressive = self.autoregressive.to(self.device)
+            self.diffusion = self.diffusion.to(self.device)
+            self.clvp = self.clvp.to(self.device)
+            self.vocoder = self.vocoder.to(self.device)
+        self.high_vram = high_vram
+
+    @contextmanager
+    def temporary_cuda(self, model):
+        if self.high_vram:
+            yield model
+        else:
+            m = model.to(self.device)
+            yield m
+            m = model.cpu()
+
+    def load_cvvp(self):
+        """Load CVVP model."""
+        self.cvvp = (
+            CVVP(
+                model_dim=512,
+                transformer_heads=8,
+                dropout=0,
+                mel_codes=8192,
+                conditioning_enc_depth=8,
+                cond_mask_percentage=0,
+                speech_enc_depth=8,
+                speech_mask_percentage=0,
+                latent_multiplier=1,
+            )
+            .cpu()
+            .eval()
+        )
+        self.cvvp.load_state_dict(
+            torch.load(get_model_path("cvvp.pth", self.models_dir))
+        )
+
+    def get_conditioning_latents(self, voice_samples, return_mels=False, latent_averaging_mode=0, original_tortoise=False):
+      """
+      Transforms one or more voice_samples into a tuple (autoregressive_conditioning_latent, diffusion_conditioning_latent).
+      These are expressive learned latents that encode aspects of the provided clips like voice, intonation, and acoustic
+      properties.
+      :param voice_samples: List of arbitrary reference clips, which should be *pairs* of torch tensors containing arbitrary kHz waveform data.
+      :param latent_averaging_mode: 0/1/2 for following modes:
+          0 - latents will be generated as in the original tortoise, using ~4.27s from each voice sample, averaging latent across all samples
+          1 - latents will be generated using (almost) entire voice samples, averaged across all the ~4.27s chunks
+          2 - latents will be generated using (almost) entire voice samples, averaged per voice sample
+      """
+      assert latent_averaging_mode in [0, 1, 2], "latent_averaging mode has to be one of (0, 1, 2)"
+      print("mode", latent_averaging_mode)
+      device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+      with torch.no_grad():
+          # Move the entire nested structure to the device
+          voice_samples = [
+              (pair[0].to(device), pair[1].to(device))
+              for pair in voice_samples
+          ]
+
+          auto_conds = []
+          for ls in voice_samples:
+              auto_conds.append(format_conditioning(ls[0], device=device))  # Use device here
+          auto_conds = torch.stack(auto_conds, dim=1)
+          with self.temporary_cuda(self.autoregressive) as ar:
+              auto_latent = ar.get_conditioning(auto_conds)
+
+          diffusion_conds = []
+
+          DURS_CONST = 102400
+          for ls in voice_samples:
+              # The diffuser operates at a sample rate of 24000 (except for the latent inputs)
+              sample = (
+                  torchaudio.functional.resample(ls[0], 22050, 24000)
+                  if original_tortoise
+                  else ls[1]
+              )
+              if latent_averaging_mode == 0:
+                  sample = pad_or_truncate(sample, DURS_CONST)
+                  cond_mel = wav_to_univnet_mel(
+                      sample.to(device),  # Use device here
+                      do_normalization=False,
+                      device=device,
+                  )
+                  diffusion_conds.append(cond_mel)
+              else:
+                  from math import ceil
+
+                  if latent_averaging_mode == 2:
+                      temp_diffusion_conds = []
+                  for chunk in range(ceil(sample.shape[1] / DURS_CONST)):
+                      current_sample = sample[
+                          :, chunk * DURS_CONST : (chunk + 1) * DURS_CONST
+                      ]
+                      current_sample = pad_or_truncate(current_sample, DURS_CONST)
+                      cond_mel = wav_to_univnet_mel(
+                          current_sample.to(device),  # Use device here
+                          do_normalization=False,
+                          device=device,
+                      )
+                      if latent_averaging_mode == 1:
+                          diffusion_conds.append(cond_mel)
+                      elif latent_averaging_mode == 2:
+                          temp_diffusion_conds.append(cond_mel)
+                  if latent_averaging_mode == 2:
+                      diffusion_conds.append(
+                          torch.stack(temp_diffusion_conds).mean(0)
+                      )
+          diffusion_conds = torch.stack(diffusion_conds, dim=1)
+
+          with self.temporary_cuda(self.diffusion) as diffusion:
+              diffusion_latent = diffusion.get_conditioning(diffusion_conds)
+
+      if return_mels:
+          return auto_latent, diffusion_latent, auto_conds, diffusion_conds
+      else:
+          return auto_latent, diffusion_latent
+
+    def get_random_conditioning_latents(self):
+        # Lazy-load the RLG models.
+        if self.rlg_auto is None:
+            self.rlg_auto = RandomLatentConverter(1024).eval()
+            self.rlg_auto.load_state_dict(
+                torch.load(
+                    get_model_path("rlg_auto.pth", self.models_dir),
+                    map_location=torch.device("cpu"),
+                )
+            )
+            self.rlg_diffusion = RandomLatentConverter(2048).eval()
+            self.rlg_diffusion.load_state_dict(
+                torch.load(
+                    get_model_path("rlg_diffuser.pth", self.models_dir),
+                    map_location=torch.device("cpu"),
+                )
+            )
+        with torch.no_grad():
+            return self.rlg_auto(torch.tensor([0.0])), self.rlg_diffusion(
+                torch.tensor([0.0])
+            )
+
+    def tts_with_preset(self, text, preset="fast", **kwargs):
+        """
+        Calls TTS with one of a set of preset generation parameters. Options:
+            'single_sample': Produces speech even faster, but only produces 1 sample.
+            'ultra_fast': Produces speech much faster than the original tortoise repo.
+            'ultra_fast_old': Produces speech at a speed which belies the name of this repo. (Not really, but it's definitely fastest).
+            'fast': Decent quality speech at a decent inference rate. A good choice for mass inference.
+            'standard': Very good quality. This is generally about as good as you are going to get.
+            'high_quality': Use if you want the absolute best. This is not really worth the compute, though.
+        """
+        # Use generally found best tuning knobs for generation.
+        settings = {
+            "temperature": 0.2,
+            "length_penalty": 1.0,
+            "repetition_penalty": 2.0,
+            "top_p": 0.8,
+            "cond_free_k": 2.0,
+            "diffusion_temperature": 1.0,
+        }
+        # Presets are defined here.
+        presets = {
+            "single_sample": {
+                "num_autoregressive_samples": 8,
+                "diffusion_iterations": 10,
+                "sampler": "ddim",
+            },
+            "ultra_fast": {
+                "num_autoregressive_samples": 16,
+                "diffusion_iterations": 10,
+                "sampler": "ddim",
+            },
+            "ultra_fast_old": {
+                "num_autoregressive_samples": 16,
+                "diffusion_iterations": 30,
+                "cond_free": False,
+            },
+            "very_fast": {
+                "num_autoregressive_samples": 32,
+                "diffusion_iterations": 30,
+                "sampler": "dpm++2m",
+            },
+            "fast": {
+                "num_autoregressive_samples": 96,
+                "diffusion_iterations": 20,
+                "sampler": "dpm++2m",
+            },
+            "fast_old": {"num_autoregressive_samples": 96, "diffusion_iterations": 80},
+            "standard": {
+                "num_autoregressive_samples": 256,
+                "diffusion_iterations": 200,
+            },
+            "high_quality": {
+                "num_autoregressive_samples": 256,
+                "diffusion_iterations": 400,
+            },
+        }
+        settings.update(presets[preset])
+        settings.update(kwargs)  # allow overriding of preset settings with kwargs
+        return self.tts(text, **settings)
+
+    def tts(
+        self,
+        text,
+        voice_samples=None,
+        conditioning_latents=None,
+        k=1,
+        verbose=True,
+        use_deterministic_seed=None,
+        return_deterministic_state=False,
+        latent_averaging_mode=0,
+        # autoregressive generation parameters follow
+        num_autoregressive_samples=512,
+        temperature=0.8,
+        length_penalty=1,
+        repetition_penalty=2.0,
+        top_p=0.8,
+        max_mel_tokens=500,
+        # CVVP parameters follow
+        cvvp_amount=0.0,
+        # diffusion generation parameters follow
+        diffusion_iterations=100,
+        cond_free=True,
+        cond_free_k=2,
+        diffusion_temperature=1.0,
+        sampler="ddim",
+        half=True,
+        original_tortoise=False,
+        **hf_generate_kwargs,
+    ):
+        """
+        Produces an audio clip of the given text being spoken with the given reference voice.
+        :param text: Text to be spoken.
+        :param voice_samples: List of an arbitrary number of reference clips, which should be *tuple-pairs* of torch tensors containing arbitrary kHz waveform data.
+        :param conditioning_latents: A tuple of (autoregressive_conditioning_latent, diffusion_conditioning_latent), which
+                                     can be provided in lieu of voice_samples. This is ignored unless voice_samples=None.
+                                     Conditioning latents can be retrieved via get_conditioning_latents().
+        :param k: The number of returned clips. The most likely (as determined by Tortoises' CLVP model) clips are returned.
+        :param latent_averaging_mode: 0/1/2 for following modes:
+            0 - latents will be generated as in original tortoise, using ~4.27s from each voice sample, averaging latent across all samples
+            1 - latents will be generated using (almost) entire voice samples, averaged across all the ~4.27s chunks
+            2 - latents will be generated using (almost) entire voice samples, averaged per voice sample
+        :param verbose: Whether or not to print log messages indicating the progress of creating a clip. Default=true.
+        ~~AUTOREGRESSIVE KNOBS~~
+        :param num_autoregressive_samples: Number of samples taken from the autoregressive model, all of which are filtered using CLVP.
+               As Tortoise is a probabilistic model, more samples means a higher probability of creating something "great".
+        :param temperature: The softmax temperature of the autoregressive model.
+        :param length_penalty: A length penalty applied to the autoregressive decoder. Higher settings causes the model to produce more terse outputs.
+        :param repetition_penalty: A penalty that prevents the autoregressive decoder from repeating itself during decoding. Can be used to reduce the incidence
+                                   of long silences or "uhhhhhhs", etc.
+        :param top_p: P value used in nucleus sampling. (0,1]. Lower values mean the decoder produces more "likely" (aka boring) outputs.
+        :param max_mel_tokens: Restricts the output length. (0,600] integer. Each unit is 1/20 of a second.
+        :param typical_sampling: Turns typical sampling on or off. This sampling mode is discussed in this paper: https://arxiv.org/abs/2202.00666
+                                 I was interested in the premise, but the results were not as good as I was hoping. This is off by default, but
+                                 could use some tuning.
+        :param typical_mass: The typical_mass parameter from the typical_sampling algorithm.
+        ~~CLVP-CVVP KNOBS~~
+        :param cvvp_amount: Controls the influence of the CVVP model in selecting the best output from the autoregressive model.
+                            [0,1]. Values closer to 1 mean the CVVP model is more important, 0 disables the CVVP model.
+        ~~DIFFUSION KNOBS~~
+        :param diffusion_iterations: Number of diffusion steps to perform. [0,4000]. More steps means the network has more chances to iteratively refine
+                                     the output, which should theoretically mean a higher quality output. Generally a value above 250 is not noticeably better,
+                                     however.
+        :param cond_free: Whether or not to perform conditioning-free diffusion. Conditioning-free diffusion performs two forward passes for
+                          each diffusion step: one with the outputs of the autoregressive model and one with no conditioning priors. The output
+                          of the two is blended according to the cond_free_k value below. Conditioning-free diffusion is the real deal, and
+                          dramatically improves realism.
+        :param cond_free_k: Knob that determines how to balance the conditioning free signal with the conditioning-present signal. [0,inf].
+                            As cond_free_k increases, the output becomes dominated by the conditioning-free signal.
+                            Formula is: output=cond_present_output*(cond_free_k+1)-cond_absenct_output*cond_free_k
+        :param diffusion_temperature: Controls the variance of the noise fed into the diffusion model. [0,1]. Values at 0
+                                      are the "mean" prediction of the diffusion network and will sound bland and smeared.
+        ~~OTHER STUFF~~
+        :param hf_generate_kwargs: The huggingface Transformers generate API is used for the autoregressive transformer.
+                                   Extra keyword args fed to this function get forwarded directly to that API. Documentation
+                                   here: https://huggingface.co/docs/transformers/internal/generation_utils
+        :return: Generated audio clip(s) as a torch tensor. Shape 1,S if k=1 else, (k,1,S) where S is the sample length.
+                 Sample rate is 24kHz.
+        """
+        deterministic_seed = self.deterministic_state(seed=use_deterministic_seed)
+
+        text_tokens = (
+            torch.IntTensor(self.tokenizer.encode(text)).unsqueeze(0).to(self.device)
+        )
+        text_tokens = F.pad(text_tokens, (0, 1))  # This may not be necessary.
+        assert (
+            text_tokens.shape[-1] < 400
+        ), "Too much text provided. Break the text up into separate segments and re-try inference."
+
+        auto_conds = None
+        if voice_samples is not None:
+            (
+                auto_conditioning,
+                diffusion_conditioning,
+                auto_conds,
+                _,
+            ) = self.get_conditioning_latents(
+                voice_samples,
+                return_mels=True,
+                latent_averaging_mode=latent_averaging_mode,
+                original_tortoise=original_tortoise,
+            )
+        elif conditioning_latents is not None:
+            auto_conditioning, diffusion_conditioning = conditioning_latents
+        else:
+            (
+                auto_conditioning,
+                diffusion_conditioning,
+            ) = self.get_random_conditioning_latents()
+        auto_conditioning = auto_conditioning.to(self.device)
+        diffusion_conditioning = diffusion_conditioning.to(self.device)
+
+        diffuser = load_discrete_vocoder_diffuser(
+            desired_diffusion_steps=diffusion_iterations,
+            cond_free=cond_free,
+            cond_free_k=cond_free_k,
+            sampler=sampler,
+        )
+
+        # in the case of single_sample,
+        orig_batch_size = self.autoregressive_batch_size
+        while num_autoregressive_samples % self.autoregressive_batch_size:
+            self.autoregressive_batch_size //= 2
+        with torch.no_grad():
+            samples = []
+            num_batches = num_autoregressive_samples // self.autoregressive_batch_size
+            stop_mel_token = self.autoregressive.stop_mel_token
+            calm_token = 83  # This is the token for coding silence, which is fixed in place with "fix_autoregressive_output"
+            self.autoregressive = self.autoregressive.to(self.device)
+            if verbose:
+                print("Generating autoregressive samples..")
+            with self.temporary_cuda(
+                self.autoregressive
+            ) as autoregressive, torch.autocast(
+                device_type="cuda", dtype=torch.float16, enabled=half
+            ):
+                for b in tqdm(range(num_batches), disable=not verbose):
+                    codes = autoregressive.inference_speech(
+                        auto_conditioning,
+                        text_tokens,
+                        do_sample=True,
+                        top_p=top_p,
+                        temperature=temperature,
+                        num_return_sequences=self.autoregressive_batch_size,
+                        length_penalty=length_penalty,
+                        repetition_penalty=repetition_penalty,
+                        max_generate_length=max_mel_tokens,
+                        **hf_generate_kwargs,
+                    )
+                    padding_needed = max_mel_tokens - codes.shape[1]
+                    codes = F.pad(codes, (0, padding_needed), value=stop_mel_token)
+                    samples.append(codes)
+            self.autoregressive_batch_size = (
+                orig_batch_size  # in the case of single_sample
+            )
+
+            clip_results = []
+            with self.temporary_cuda(self.clvp) as clvp, torch.autocast(
+                device_type="cuda", dtype=torch.float16, enabled=half
+            ):
+                if cvvp_amount > 0:
+                    if self.cvvp is None:
+                        self.load_cvvp()
+                    self.cvvp = self.cvvp.to(self.device)
+                if verbose:
+                    if self.cvvp is None:
+                        print("Computing best candidates using CLVP")
+                    else:
+                        print(
+                            f"Computing best candidates using CLVP {((1-cvvp_amount) * 100):2.0f}% and CVVP {(cvvp_amount * 100):2.0f}%"
+                        )
+                for batch in tqdm(samples, disable=not verbose):
+                    for i in range(batch.shape[0]):
+                        batch[i] = fix_autoregressive_output(batch[i], stop_mel_token)
+                    if cvvp_amount != 1:
+                        clvp_res = clvp(
+                            text_tokens.repeat(batch.shape[0], 1),
+                            batch,
+                            return_loss=False,
+                        )
+                    if auto_conds is not None and cvvp_amount > 0:
+                        cvvp_accumulator = 0
+                        for cl in range(auto_conds.shape[1]):
+                            cvvp_accumulator = cvvp_accumulator + self.cvvp(
+                                auto_conds[:, cl].repeat(batch.shape[0], 1, 1),
+                                batch,
+                                return_loss=False,
+                            )
+                        cvvp = cvvp_accumulator / auto_conds.shape[1]
+                        if cvvp_amount == 1:
+                            clip_results.append(cvvp)
+                        else:
+                            clip_results.append(
+                                cvvp * cvvp_amount + clvp_res * (1 - cvvp_amount)
+                            )
+                    else:
+                        clip_results.append(clvp_res)
+                clip_results = torch.cat(clip_results, dim=0)
+                samples = torch.cat(samples, dim=0)
+                best_results = samples[torch.topk(clip_results, k=k).indices]
+            if self.cvvp is not None:
+                self.cvvp = self.cvvp.cpu()
+            del samples
+
+            # The diffusion model actually wants the last hidden layer from the autoregressive model as conditioning
+            # inputs. Re-produce those for the top results. This could be made more efficient by storing all of these
+            # results, but will increase memory usage.
+            with self.temporary_cuda(self.autoregressive) as autoregressive:
+                best_latents = autoregressive(
+                    auto_conditioning.repeat(k, 1),
+                    text_tokens.repeat(k, 1),
+                    torch.tensor([text_tokens.shape[-1]], device=text_tokens.device),
+                    best_results,
+                    torch.tensor(
+                        [
+                            best_results.shape[-1]
+                            * self.autoregressive.mel_length_compression
+                        ],
+                        device=text_tokens.device,
+                    ),
+                    return_latent=True,
+                    clip_inputs=False,
+                )
+            del auto_conditioning
+
+            if verbose:
+                print("Transforming autoregressive outputs into audio..")
+            wav_candidates = []
+            with self.temporary_cuda(self.diffusion) as diffusion, self.temporary_cuda(
+                self.vocoder
+            ) as vocoder:
+                diffusion.enable_fp16 = half  # hacky
+                for b in range(best_results.shape[0]):
+                    codes = best_results[b].unsqueeze(0)
+                    latents = best_latents[b].unsqueeze(0)
+
+                    # Find the first occurrence of the "calm" token and trim the codes to that.
+                    ctokens = 0
+                    for k in range(codes.shape[-1]):
+                        if codes[0, k] == calm_token:
+                            ctokens += 1
+                        else:
+                            ctokens = 0
+                        if (
+                            ctokens > 8
+                        ):  # 8 tokens gives the diffusion model some "breathing room" to terminate speech.
+                            latents = latents[:, :k]
+                            break
+
+                    mel = do_spectrogram_diffusion(
+                        diffusion,
+                        diffuser,
+                        latents,
+                        diffusion_conditioning,
+                        temperature=diffusion_temperature,
+                        verbose=verbose,
+                    )
+                    wav = vocoder.inference(mel)
+                    wav_candidates.append(wav.cpu())
+
+            def potentially_redact(clip, text):
+                if self.enable_redaction:
+                    return self.aligner.redact(clip.squeeze(1), text).unsqueeze(1)
+                return clip
+
+            wav_candidates = [
+                potentially_redact(wav_candidate, text)
+                for wav_candidate in wav_candidates
+            ]
+
+            if len(wav_candidates) > 1:
+                res = wav_candidates
+            else:
+                res = wav_candidates[0]
+
+            if return_deterministic_state:
+                return res, (
+                    deterministic_seed,
+                    text,
+                    voice_samples,
+                    conditioning_latents,
+                )
+            else:
+                return res
+
+    def deterministic_state(self, seed=None):
+        """
+        Sets the random seeds that tortoise uses to the current time() and returns that seed so results can be
+        reproduced.
+        """
+        seed = int(time()) if seed is None else seed
+        torch.manual_seed(seed)
+        random.seed(seed)
+        # Can't currently set this because of CUBLAS. TODO: potentially enable it if necessary.
+        # torch.use_deterministic_algorithms(True)
+
+        return seed

tortoise/data/got.txt

@@ -0,0 +1,276 @@
+Chapter One
+
+
+Bran
+
+
+The morning had dawned clear and cold, with a crispness that hinted at the end of summer. They set forth at daybreak to see a man beheaded, twenty in all, and Bran rode among them, nervous with excitement. This was the first time he had been deemed old enough to go with his lord father and his brothers to see the king's justice done. It was the ninth year of summer, and the seventh of Bran's life.
+
+
+The man had been taken outside a small holdfast in the hills. Robb thought he was a wildling, his sword sworn to Mance Rayder, the King-beyond-the-Wall. It made Bran's skin prickle to think of it. He remembered the hearth tales Old Nan told them. The wildlings were cruel men, she said, slavers and slayers and thieves. They consorted with giants and ghouls, stole girl children in the dead of night, and drank blood from polished horns. And their women lay with the Others in the Long Night to sire terrible half-human children.
+
+
+But the man they found bound hand and foot to the holdfast wall awaiting the king's justice was old and scrawny, not much taller than Robb. He had lost both ears and a finger to frostbite, and he dressed all in black, the same as a brother of the Night's Watch, except that his furs were ragged and greasy.
+
+
+The breath of man and horse mingled, steaming, in the cold morning air as his lord father had the man cut down from the wall and dragged before them. Robb and Jon sat tall and still on their horses, with Bran between them on his pony, trying to seem older than seven, trying to pretend that he'd seen all this before. A faint wind blew through the holdfast gate. Over their heads flapped the banner of the Starks of Winterfell: a grey direwolf racing across an ice-white field.
+
+Bran's father sat solemnly on his horse, long brown hair stirring in the wind. His closely trimmed beard was shot with white, making him look older than his thirty-five years. He had a grim cast to his grey eyes this day, and he seemed not at all the man who would sit before the fire in the evening and talk softly of the age of heroes and the children of the forest. He had taken off Father's face, Bran thought, and donned the face of Lord Stark of Winterfell.
+
+
+There were questions asked and answers given there in the chill of morning, but afterward Bran could not recall much of what had been said. Finally his lord father gave a command, and two of his guardsmen dragged the ragged man to the ironwood stump in the center of the square. They forced his head down onto the hard black wood. Lord Eddard Stark dismounted and his ward Theon Greyjoy brought forth the sword. "Ice," that sword was called. It was as wide across as a man's hand, and taller even than Robb. The blade was Valyrian steel, spell-forged and dark as smoke. Nothing held an edge like Valyrian steel.
+
+
+His father peeled off his gloves and handed them to Jory Cassel, the captain of his household guard. He took hold of Ice with both hands and said, "In the name of Robert of the House Baratheon, the First of his Name, King of the Andals and the Rhoynar and the First Men, Lord of the Seven Kingdoms and Protector of the Realm, by the word of Eddard of the House Stark, Lord of Winterfell and Warden of the North, I do sentence you to die." He lifted the greatsword high above his head.
+
+
+Bran's bastard brother Jon Snow moved closer. "Keep the pony well in hand," he whispered. "And don't look away. Father will know if you do."
+
+
+Bran kept his pony well in hand, and did not look away.
+
+
+His father took off the man's head with a single sure stroke. Blood sprayed out across the snow, as red as surnmerwine. One of the horses reared and had to be restrained to keep from bolting. Bran could not take his eyes off the blood. The snows around the stump drank it eagerly, reddening as he watched.
+
+The head bounced off a thick root and rolled. It came up near Greyjoy's feet. Theon was a lean, dark youth of nineteen who found everything amusing. He laughed, put his boot on the head, and kicked it away.
+
+
+"Ass," Jon muttered, low enough so Greyjoy did not hear. He put a hand on Bran's shoulder, and Bran looked over at his bastard brother. "You did well," Jon told him solemnly. Jon was fourteen, an old hand at justice.
+
+
+It seemed colder on the long ride back to Winterfell, though the wind had died by then and the sun was higher in the sky. Bran rode with his brothers, well ahead of the main party, his pony struggling hard to keep up with their horses.
+
+
+"The deserter died bravely," Robb said. He was big and broad and growing every day, with his mother's coloring, the fair skin, red-brown hair, and blue eyes of the Tullys of Riverrun. "He had courage, at the least."
+
+
+"No," Jon Snow said quietly. "It was not courage. This one was dead of fear. You could see it in his eyes, Stark." Jon's eyes were a grey so dark they seemed almost black, but there was little they did not see. He was of an age with Robb, but they did not look alike. Jon was slender where Robb was muscular, dark where Robb was fair, graceful and quick where his half brother was strong and fast.
+
+
+Robb was not impressed. "The Others take his eyes," he swore. "He died well. Race you to the bridge?"
+
+
+"Done," Jon said, kicking his horse forward. Robb cursed and followed, and they galloped off down the trail, Robb laughing and hooting, Jon silent and intent. The hooves of their horses kicked up showers of snow as they went.
+
+Bran did not try to follow. His pony could not keep up. He had seen the ragged man's eyes, and he was thinking of them now. After a while, the sound of Robb's laughter receded, and the woods grew silent again.
+
+
+So deep in thought was he that he never heard the rest of the party until his father moved up to ride beside him. "Are you well, Bran?" he asked, not unkindly.
+
+
+"Yes, Father," Bran told him. He looked up. Wrapped in his furs and leathers, mounted on his great warhorse, his lord father loomed over him like a giant. "Robb says the man died bravely, but Jon says he was afraid."
+
+
+"What do you think?" his father asked.
+
+
+Bran thought about it. "Can a man still be brave if he's afraid?"
+
+
+"That is the only time a man can be brave," his father told him. "Do you understand why I did it?"
+
+
+"He was a wildling," Bran said. "They carry off women and sell them to the Others."
+
+
+His lord father smiled. "Old Nan has been telling you stories again. In truth, the man was an oathbreaker, a deserter from the Night's Watch. No man is more dangerous. The deserter knows his life is forfeit if he is taken, so he will not flinch from any crime, no matter how vile. But you mistake me. The question was not why the man had to die, but why I must do it."
+
+
+Bran had no answer for that. "King Robert has a headsman," he said, uncertainly.
+
+
+"He does," his father admitted. "As did the Targaryen kings before him. Yet our way is the older way. The blood of the First Men still flows in the veins of the Starks, and we hold to the belief that the man who passes the sentence should swing the sword. If you would take a man's life, you owe it to him to look into his eyes and hear his final words. And if you cannot bear to do that, then perhaps the man does not deserve to die.
+
+
+"One day, Bran, you will be Robb's bannerman, holding a keep of your own for your brother and your king, and justice will fall to you. When that day comes, you must take no pleasure in the task, but neither must you look away. A ruler who hides behind paid executioners soon forgets what death is."
+
+
+That was when Jon reappeared on the crest of the hill before them. He waved and shouted down at them. "Father, Bran, come quickly, see what Robb has found!" Then he was gone again.
+
+
+Jory rode up beside them. "Trouble, my lord?"
+
+
+"Beyond a doubt," his lord father said. "Come, let us see what mischief my sons have rooted out now." He sent his horse into a trot. Jory and Bran and the rest came after.
+
+
+They found Robb on the riverbank north of the bridge, with Jon still mounted beside him. The late summer snows had been heavy this moonturn. Robb stood knee-deep in white, his hood pulled back so the sun shone in his hair. He was cradling something in his arm, while the boys talked in hushed, excited voices.
+
+
+The riders picked their way carefully through the drifts, groping for solid footing on the hidden, uneven ground . Jory Cassel and Theon Greyjoy were the first to reach the boys. Greyjoy was laughing and joking as he rode. Bran heard the breath go out of him. "Gods!" he exclaimed, struggling to keep control of his horse as he reached for his sword.
+
+
+Jory's sword was already out. "Robb, get away from it!" he called as his horse reared under him.
+
+
+Robb grinned and looked up from the bundle in his arms. "She can't hurt you," he said. "She's dead, Jory."
+
+
+Bran was afire with curiosity by then. He would have spurred the pony faster, but his father made them dismount beside the bridge and approach on foot. Bran jumped off and ran.
+
+
+By then Jon, Jory, and Theon Greyjoy had all dismounted as well. "What in the seven hells is it?" Greyjoy was saying.
+
+
+"A wolf," Robb told him.
+
+
+"A freak," Greyjoy said. "Look at the size of it."
+
+
+Bran's heart was thumping in his chest as he pushed through a waist-high drift to his brothers' side.
+
+
+Half-buried in bloodstained snow, a huge dark shape slumped in death. Ice had formed in its shaggy grey fur, and the faint smell of corruption clung to it like a woman's perfume. Bran glimpsed blind eyes crawling with maggots, a wide mouth full of yellowed teeth. But it was the size of it that made him gasp. It was bigger than his pony, twice the size of the largest hound in his father's kennel.
+
+
+"It's no freak," Jon said calmly. "That's a direwolf. They grow larger than the other kind."
+
+
+Theon Greyjoy said, "There's not been a direwolf sighted south of the Wall in two hundred years."
+
+
+"I see one now," Jon replied.
+
+
+Bran tore his eyes away from the monster. That was when he noticed the bundle in Robb's arms. He gave a cry of delight and moved closer. The pup was a tiny ball of grey-black fur, its eyes still closed. It nuzzled blindly against Robb's chest as he cradled it, searching for milk among his leathers, making a sad little whimpery sound. Bran reached out hesitantly. "Go on," Robb told him. "You can touch him."
+
+
+Bran gave the pup a quick nervous stroke, then turned as Jon said, "Here you go." His half brother put a second pup into his arms. "There are five of them." Bran sat down in the snow and hugged the wolf pup to his face. Its fur was soft and warm against his cheek.
+
+
+"Direwolves loose in the realm, after so many years," muttered Hullen, the master of horse. "I like it not."
+
+
+"It is a sign," Jory said.
+
+
+Father frowned. "This is only a dead animal, Jory," he said. Yet he seemed troubled. Snow crunched under his boots as he moved around the body. "Do we know what killed her?"
+
+
+"There's something in the throat," Robb told him, proud to have found the answer before his father even asked. "There, just under the jaw."
+
+
+His father knelt and groped under the beast's head with his hand. He gave a yank and held it up for all to see. A foot of shattered antler, tines snapped off, all wet with blood.
+
+
+A sudden silence descended over the party. The men looked at the antler uneasily, and no one dared to speak. Even Bran could sense their fear, though he did not understand.
+
+
+His father tossed the antler to the side and cleansed his hands in the snow. "I'm surprised she lived long enough to whelp," he said. His voice broke the spell.
+
+
+"Maybe she didn't," Jory said. "I've heard tales . . . maybe the bitch was already dead when the pups came."
+
+
+"Born with the dead," another man put in. "Worse luck."
+
+
+"No matter," said Hullen. "They be dead soon enough too."
+
+
+Bran gave a wordless cry of dismay.
+
+
+"The sooner the better," Theon Greyjoy agreed. He drew his sword. "Give the beast here, Bran."
+
+
+The little thing squirmed against him, as if it heard and understood. "No!" Bran cried out fiercely. "It's mine."
+
+
+"Put away your sword, Greyjoy," Robb said. For a moment he sounded as commanding as their father, like the lord he would someday be. "We will keep these pups."
+
+
+"You cannot do that, boy," said Harwin, who was Hullen's son.
+
+
+"It be a mercy to kill them," Hullen said.
+
+
+Bran looked to his lord father for rescue, but got only a frown, a furrowed brow. "Hullen speaks truly, son. Better a swift death than a hard one from cold and starvation."
+
+
+"No!" He could feel tears welling in his eyes, and he looked away. He did not want to cry in front of his father.
+
+
+Robb resisted stubbornly. "Ser Rodrik's red bitch whelped again last week," he said. "It was a small litter, only two live pups. She'll have milk enough."
+
+
+"She'll rip them apart when they try to nurse."
+
+
+"Lord Stark," Jon said. It was strange to hear him call Father that, so formal. Bran looked at him with desperate hope. "There are five pups," he told Father. "Three male, two female."
+
+
+"What of it, Jon?"
+
+
+"You have five trueborn children," Jon said. "Three sons, two daughters. The direwolf is the sigil of your House. Your children were meant to have these pups, my lord."
+
+
+Bran saw his father's face change, saw the other men exchange glances. He loved Jon with all his heart at that moment. Even at seven, Bran understood what his brother had done. The count had come right only because Jon had omitted himself. He had included the girls, included even Rickon, the baby, but not the bastard who bore the surname Snow, the name that custom decreed be given to all those in the north unlucky enough to be born with no name of their own.
+
+
+Their father understood as well. "You want no pup for yourself, Jon?" he asked softly.
+
+
+"The direwolf graces the banners of House Stark," Jon pointed out. "I am no Stark, Father."
+
+
+Their lord father regarded Jon thoughtfully. Robb rushed into the silence he left. "I will nurse him myself, Father," he promised. "I will soak a towel with warm milk, and give him suck from that."
+
+
+"Me too!" Bran echoed.
+
+
+The lord weighed his sons long and carefully with his eyes. "Easy to say, and harder to do. I will not have you wasting the servants' time with this. If you want these pups, you will feed them yourselves. Is that understood?"
+
+
+Bran nodded eagerly. The pup squirmed in his grasp, licked at his face with a warm tongue.
+
+
+"You must train them as well," their father said. "You must train them. The kennelmaster will have nothing to do with these monsters, I promise you that. And the gods help you if you neglect them, or brutalize them, or train them badly. These are not dogs to beg for treats and slink off at a kick. A direwolf will rip a man's arm off his shoulder as easily as a dog will kill a rat. Are you sure you want this?"
+
+"Yes, Father," Bran said.
+
+
+"Yes," Robb agreed.
+
+
+"The pups may die anyway, despite all you do."
+
+
+"They won't die," Robb said. "We won't let them die."
+
+
+"Keep them, then. Jory, Desmond, gather up the other pups. It's time we were back to Winterfell."
+
+
+It was not until they were mounted and on their way that Bran allowed himself to taste the sweet air of victory. By then, his pup was snuggled inside his leathers, warm against him, safe for the long ride home. Bran was wondering what to name him.
+
+
+Halfway across the bridge, Jon pulled up suddenly.
+
+
+"What is it, Jon?" their lord father asked.
+
+
+"Can't you hear it?"
+
+
+Bran could hear the wind in the trees, the clatter of their hooves on the ironwood planks, the whimpering of his hungry pup, but Jon was listening to something else.
+
+
+"There," Jon said. He swung his horse around and galloped back across the bridge. They watched him dismount where the direwolf lay dead in the snow, watched him kneel. A moment later he was riding back to them, smiling.
+
+
+"He must have crawled away from the others," Jon said.
+
+
+"Or been driven away," their father said, looking at the sixth pup. His fur was white, where the rest of the litter was grey. His eyes were as red as the blood of the ragged man who had died that morning. Bran thought it curious that this pup alone would have opened his eyes while the others were still blind.
+
+
+"An albino," Theon Greyjoy said with wry amusement. "This one will die even faster than the others."
+
+
+Jon Snow gave his father's ward a long, chilling look. "I think not, Greyjoy," he said. "This one belongs to me."

tortoise/data/riding_hood.txt

@@ -0,0 +1,54 @@
+Once upon a time there lived in a certain village a little country girl, the prettiest creature who was ever seen. Her mother was excessively fond of her; and her grandmother doted on her still more. This good woman had a little red riding hood made for her. It suited the girl so extremely well that everybody called her Little Red Riding Hood.
+One day her mother, having made some cakes, said to her, "Go, my dear, and see how your grandmother is doing, for I hear she has been very ill. Take her a cake, and this little pot of butter."
+
+Little Red Riding Hood set out immediately to go to her grandmother, who lived in another village.
+
+As she was going through the wood, she met with a wolf, who had a very great mind to eat her up, but he dared not, because of some woodcutters working nearby in the forest. He asked her where she was going. The poor child, who did not know that it was dangerous to stay and talk to a wolf, said to him, "I am going to see my grandmother and carry her a cake and a little pot of butter from my mother."
+
+"Does she live far off?" said the wolf
+
+"Oh I say," answered Little Red Riding Hood; "it is beyond that mill you see there, at the first house in the village."
+
+"Well," said the wolf, "and I'll go and see her too. I'll go this way and go you that, and we shall see who will be there first."
+
+The wolf ran as fast as he could, taking the shortest path, and the little girl took a roundabout way, entertaining herself by gathering nuts, running after butterflies, and gathering bouquets of little flowers. It was not long before the wolf arrived at the old woman's house. He knocked at the door: tap, tap.
+
+"Who's there?"
+
+"Your grandchild, Little Red Riding Hood," replied the wolf, counterfeiting her voice; "who has brought you a cake and a little pot of butter sent you by mother."
+
+The good grandmother, who was in bed, because she was somewhat ill, cried out, "Pull the bobbin, and the latch will go up."
+
+The wolf pulled the bobbin, and the door opened, and then he immediately fell upon the good woman and ate her up in a moment, for it been more than three days since he had eaten. He then shut the door and got into the grandmother's bed, expecting Little Red Riding Hood, who came some time afterwards and knocked at the door: tap, tap.
+
+"Who's there?"
+
+Little Red Riding Hood, hearing the big voice of the wolf, was at first afraid; but believing her grandmother had a cold and was hoarse, answered, "It is your grandchild Little Red Riding Hood, who has brought you a cake and a little pot of butter mother sends you."
+
+The wolf cried out to her, softening his voice as much as he could, "Pull the bobbin, and the latch will go up."
+
+Little Red Riding Hood pulled the bobbin, and the door opened.
+
+The wolf, seeing her come in, said to her, hiding himself under the bedclothes, "Put the cake and the little pot of butter upon the stool, and come get into bed with me."
+
+Little Red Riding Hood took off her clothes and got into bed. She was greatly amazed to see how her grandmother looked in her nightclothes, and said to her, "Grandmother, what big arms you have!"
+
+"All the better to hug you with, my dear."
+
+"Grandmother, what big legs you have!"
+
+"All the better to run with, my child."
+
+"Grandmother, what big ears you have!"
+
+"All the better to hear with, my child."
+
+"Grandmother, what big eyes you have!"
+
+"All the better to see with, my child."
+
+"Grandmother, what big teeth you have got!"
+
+"All the better to eat you up with."
+
+And, saying these words, this wicked wolf fell upon Little Red Riding Hood, and ate her all up.

tortoise/data/seal_copypasta.txt

@@ -0,0 +1 @@
+What the fuck did you just fucking say about me, you little bitch? I'll have you know I graduated top of my class in the Navy Seals, and I've been involved in numerous secret raids on Al kayda, and I have over 300 confirmed kills. I am trained in gorilla warfare and I'm the top sniper in the entire U S armed forces. You are nothing to me but just another target. I will wipe you the fuck out with precision the likes of which has never been seen before on this Earth, mark my fucking words. You think you can get away with saying that shit to me over the Internet? Think again, fucker. As we speak I am contacting my secret network of spies across the U S A and your IP is being traced right now so you better prepare for the storm, maggot. The storm that wipes out the pathetic little thing you call your life. You're fucking dead, kid. I can be anywhere, anytime, and I can kill you in over seven hundred ways, and that's just with my bare hands. Not only am I extensively trained in unarmed combat, but I have access to the entire arsenal of the United States Marine Corps and I will use it to its full extent to wipe your miserable ass off the face of the continent, you little shit. If only you could have known what unholy retribution your little "clever" comment was about to bring down upon you, maybe you would have held your fucking tongue. But you couldn't, you didn't, and now you're paying the price, you goddamn idiot. I will shit fury all over you and you will drown in it. You're fucking dead, kiddo.

tortoise/data/tokenizer.json

@@ -0,0 +1 @@
+{"version":"1.0","truncation":null,"padding":null,"added_tokens":[{"id":0,"special":true,"content":"[STOP]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":1,"special":true,"content":"[UNK]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false},{"id":2,"special":true,"content":"[SPACE]","single_word":false,"lstrip":false,"rstrip":false,"normalized":false}],"normalizer":null,"pre_tokenizer":{"type":"Whitespace"},"post_processor":null,"decoder":null,"model":{"type":"BPE","dropout":null,"unk_token":"[UNK]","continuing_subword_prefix":null,"end_of_word_suffix":null,"fuse_unk":false,"vocab":{"[STOP]":0,"[UNK]":1,"[SPACE]":2,"!":3,"'":4,"(":5,")":6,",":7,"-":8,".":9,"/":10,":":11,";":12,"?":13,"a":14,"b":15,"c":16,"d":17,"e":18,"f":19,"g":20,"h":21,"i":22,"j":23,"k":24,"l":25,"m":26,"n":27,"o":28,"p":29,"q":30,"r":31,"s":32,"t":33,"u":34,"v":35,"w":36,"x":37,"y":38,"z":39,"th":40,"in":41,"the":42,"an":43,"er":44,"ou":45,"re":46,"on":47,"at":48,"ed":49,"en":50,"to":51,"ing":52,"and":53,"is":54,"as":55,"al":56,"or":57,"of":58,"ar":59,"it":60,"es":61,"he":62,"st":63,"le":64,"om":65,"se":66,"be":67,"ad":68,"ow":69,"ly":70,"ch":71,"wh":72,"that":73,"you":74,"li":75,"ve":76,"ac":77,"ti":78,"ld":79,"me":80,"was":81,"gh":82,"id":83,"ll":84,"wi":85,"ent":86,"for":87,"ay":88,"ro":89,"ver":90,"ic":91,"her":92,"ke":93,"his":94,"no":95,"ut":96,"un":97,"ir":98,"lo":99,"we":100,"ri":101,"ha":102,"with":103,"ght":104,"out":105,"im":106,"ion":107,"all":108,"ab":109,"one":110,"ne":111,"ge":112,"ould":113,"ter":114,"mo":115,"had":116,"ce":117,"she":118,"go":119,"sh":120,"ur":121,"am":122,"so":123,"pe":124,"my":125,"de":126,"are":127,"but":128,"ome":129,"fr":130,"ther":131,"fe":132,"su":133,"do":134,"con":135,"te":136,"ain":137,"ere":138,"po":139,"if":140,"they":141,"us":142,"ag":143,"tr":144,"now":145,"oun":146,"this":147,"have":148,"not":149,"sa":150,"il":151,"up":152,"thing":153,"from":154,"ap":155,"him":156,"ack":157,"ation":158,"ant":159,"our":160,"op":161,"like":162,"ust":163,"ess":164,"bo":165,"ok":166,"ul":167,"ind":168,"ex":169,"com":170,"some":171,"there":172,"ers":173,"co":174,"res":175,"man":176,"ard":177,"pl":178,"wor":179,"way":180,"tion":181,"fo":182,"ca":183,"were":184,"by":185,"ate":186,"pro":187,"ted":188,"ound":189,"own":190,"would":191,"ts":192,"what":193,"qu":194,"ally":195,"ight":196,"ck":197,"gr":198,"when":199,"ven":200,"can":201,"ough":202,"ine":203,"end":204,"per":205,"ous":206,"od":207,"ide":208,"know":209,"ty":210,"very":211,"si":212,"ak":213,"who":214,"about":215,"ill":216,"them":217,"est":218,"red":219,"ye":220,"could":221,"ong":222,"your":223,"their":224,"em":225,"just":226,"other":227,"into":228,"any":229,"whi":230,"um":231,"tw":232,"ast":233,"der":234,"did":235,"ie":236,"been":237,"ace":238,"ink":239,"ity":240,"back":241,"ting":242,"br":243,"more":244,"ake":245,"pp":246,"then":247,"sp":248,"el":249,"use":250,"bl":251,"said":252,"over":253,"get":254},"merges":["t h","i n","th e","a n","e r","o u","r e","o n","a t","e d","e n","t o","in g","an d","i s","a s","a l","o r","o f","a r","i t","e s","h e","s t","l e","o m","s e","b e","a d","o w","l y","c h","w h","th at","y ou","l i","v e","a c","t i","l d","m e","w as","g h","i d","l l","w i","en t","f or","a y","r o","v er","i c","h er","k e","h is","n o","u t","u n","i r","l o","w e","r i","h a","wi th","gh t","ou t","i m","i on","al l","a b","on e","n e","g e","ou ld","t er","m o","h ad","c e","s he","g o","s h","u r","a m","s o","p e","m y","d e","a re","b ut","om e","f r","the r","f e","s u","d o","c on","t e","a in","er e","p o","i f","the y","u s","a g","t r","n ow","ou n","th is","ha ve","no t","s a","i l","u p","th ing","fr om","a p","h im","ac k","at ion","an t","ou r","o p","li ke","u st","es s","b o","o k","u l","in d","e x","c om","s ome","the re","er s","c o","re s","m an","ar d","p l","w or","w ay","ti on","f o","c a","w ere","b y","at e","p ro","t ed","oun d","ow n","w ould","t s","wh at","q u","al ly","i ght","c k","g r","wh en","v en","c an","ou gh","in e","en d","p er","ou s","o d","id e","k now","t y","ver y","s i","a k","wh o","ab out","i ll","the m","es t","re d","y e","c ould","on g","you r","the ir","e m","j ust","o ther","in to","an y","wh i","u m","t w","as t","d er","d id","i e","be en","ac e","in k","it y","b ack","t ing","b r","mo re","a ke","p p","the n","s p","e l","u se","b l","sa id","o ver","ge t"]}}

tortoise/dpm_solver_pytorch.py

@@ -0,0 +1,1653 @@
+import math
+
+import torch
+
+
+class NoiseScheduleVP:
+    def __init__(
+        self,
+        schedule="discrete",
+        betas=None,
+        alphas_cumprod=None,
+        continuous_beta_0=0.1,
+        continuous_beta_1=20.0,
+        dtype=torch.float32,
+    ):
+        """Create a wrapper class for the forward SDE (VP type).
+
+        ***
+        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
+                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
+        ***
+
+        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
+        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
+        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
+
+            log_alpha_t = self.marginal_log_mean_coeff(t)
+            sigma_t = self.marginal_std(t)
+            lambda_t = self.marginal_lambda(t)
+
+        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
+
+            t = self.inverse_lambda(lambda_t)
+
+        ===============================================================
+
+        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
+
+        1. For discrete-time DPMs:
+
+            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
+                t_i = (i + 1) / N
+            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
+            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
+
+            Args:
+                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
+                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
+
+            Note that we always have alphas_cumprod = cumprod(1 - betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+
+            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
+                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
+                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
+                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
+                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
+                and
+                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
+
+
+        2. For continuous-time DPMs:
+
+            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
+            schedule are the default settings in DDPM and improved-DDPM:
+
+            Args:
+                beta_min: A `float` number. The smallest beta for the linear schedule.
+                beta_max: A `float` number. The largest beta for the linear schedule.
+                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
+                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
+                T: A `float` number. The ending time of the forward process.
+
+        ===============================================================
+
+        Args:
+            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
+                    'linear' or 'cosine' for continuous-time DPMs.
+        Returns:
+            A wrapper object of the forward SDE (VP type).
+
+        ===============================================================
+
+        Example:
+
+        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', betas=betas)
+
+        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
+        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
+
+        # For continuous-time DPMs (VPSDE), linear schedule:
+        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
+
+        """
+
+        if schedule not in ["discrete", "linear", "cosine"]:
+            raise ValueError(
+                "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
+                    schedule
+                )
+            )
+
+        self.schedule = schedule
+        if schedule == "discrete":
+            if betas is not None:
+                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
+            else:
+                assert alphas_cumprod is not None
+                log_alphas = 0.5 * torch.log(alphas_cumprod)
+            self.total_N = len(log_alphas)
+            self.T = 1.0
+            self.t_array = (
+                torch.linspace(0.0, 1.0, self.total_N + 1)[1:]
+                .reshape((1, -1))
+                .to(dtype=dtype)
+            )
+            self.log_alpha_array = log_alphas.reshape(
+                (
+                    1,
+                    -1,
+                )
+            ).to(dtype=dtype)
+        else:
+            self.total_N = 1000
+            self.beta_0 = continuous_beta_0
+            self.beta_1 = continuous_beta_1
+            self.cosine_s = 0.008
+            self.cosine_beta_max = 999.0
+            self.cosine_t_max = (
+                math.atan(self.cosine_beta_max * (1.0 + self.cosine_s) / math.pi)
+                * 2.0
+                * (1.0 + self.cosine_s)
+                / math.pi
+                - self.cosine_s
+            )
+            self.cosine_log_alpha_0 = math.log(
+                math.cos(self.cosine_s / (1.0 + self.cosine_s) * math.pi / 2.0)
+            )
+            self.schedule = schedule
+            if schedule == "cosine":
+                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
+                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
+                self.T = 0.9946
+            else:
+                self.T = 1.0
+
+    def marginal_log_mean_coeff(self, t):
+        """
+        Compute log(alpha_t) of a given continuous-time label t in [0, T].
+        """
+        if self.schedule == "discrete":
+            return interpolate_fn(
+                t.reshape((-1, 1)),
+                self.t_array.to(t.device),
+                self.log_alpha_array.to(t.device),
+            ).reshape((-1))
+        elif self.schedule == "linear":
+            return -0.25 * t**2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
+        elif self.schedule == "cosine":
+
+            def log_alpha_fn(s):
+                return torch.log(
+                    torch.cos(
+                        (s + self.cosine_s) / (1.0 + self.cosine_s) * math.pi / 2.0
+                    )
+                )
+
+            log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
+            return log_alpha_t
+
+    def marginal_alpha(self, t):
+        """
+        Compute alpha_t of a given continuous-time label t in [0, T].
+        """
+        return torch.exp(self.marginal_log_mean_coeff(t))
+
+    def marginal_std(self, t):
+        """
+        Compute sigma_t of a given continuous-time label t in [0, T].
+        """
+        return torch.sqrt(1.0 - torch.exp(2.0 * self.marginal_log_mean_coeff(t)))
+
+    def marginal_lambda(self, t):
+        """
+        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
+        """
+        log_mean_coeff = self.marginal_log_mean_coeff(t)
+        log_std = 0.5 * torch.log(1.0 - torch.exp(2.0 * log_mean_coeff))
+        return log_mean_coeff - log_std
+
+    def inverse_lambda(self, lamb):
+        """
+        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
+        """
+        if self.schedule == "linear":
+            tmp = (
+                2.0
+                * (self.beta_1 - self.beta_0)
+                * torch.logaddexp(-2.0 * lamb, torch.zeros((1,)).to(lamb))
+            )
+            Delta = self.beta_0**2 + tmp
+            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
+        elif self.schedule == "discrete":
+            log_alpha = -0.5 * torch.logaddexp(
+                torch.zeros((1,)).to(lamb.device), -2.0 * lamb
+            )
+            t = interpolate_fn(
+                log_alpha.reshape((-1, 1)),
+                torch.flip(self.log_alpha_array.to(lamb.device), [1]),
+                torch.flip(self.t_array.to(lamb.device), [1]),
+            )
+            return t.reshape((-1,))
+        else:
+            log_alpha = -0.5 * torch.logaddexp(-2.0 * lamb, torch.zeros((1,)).to(lamb))
+
+            def t_fn(log_alpha_t):
+                return (
+                    torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0))
+                    * 2.0
+                    * (1.0 + self.cosine_s)
+                    / math.pi
+                    - self.cosine_s
+                )
+
+            t = t_fn(log_alpha)
+            return t
+
+
+def model_wrapper(
+    model,
+    noise_schedule,
+    model_type="noise",
+    model_kwargs={},
+    guidance_type="uncond",
+    condition=None,
+    unconditional_condition=None,
+    guidance_scale=1.0,
+    classifier_fn=None,
+    classifier_kwargs={},
+):
+    """Create a wrapper function for the noise prediction model.
+
+    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
+    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
+
+    We support four types of the diffusion model by setting `model_type`:
+
+        1. "noise": noise prediction model. (Trained by predicting noise).
+
+        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
+
+        3. "v": velocity prediction model. (Trained by predicting the velocity).
+            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
+
+            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
+                arXiv preprint arXiv:2202.00512 (2022).
+            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
+                arXiv preprint arXiv:2210.02303 (2022).
+
+        4. "score": marginal score function. (Trained by denoising score matching).
+            Note that the score function and the noise prediction model follows a simple relationship:
+            ```
+                noise(x_t, t) = -sigma_t * score(x_t, t)
+            ```
+
+    We support three types of guided sampling by DPMs by setting `guidance_type`:
+        1. "uncond": unconditional sampling by DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
+            ``
+
+            The input `classifier_fn` has the following format:
+            ``
+                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
+            ``
+
+            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
+                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
+
+        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
+            The input `model` has the following format:
+            ``
+                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
+            ``
+            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
+
+            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
+                arXiv preprint arXiv:2207.12598 (2022).
+
+
+    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
+    or continuous-time labels (i.e. epsilon to T).
+
+    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
+    ``
+        def model_fn(x, t_continuous) -> noise:
+            t_input = get_model_input_time(t_continuous)
+            return noise_pred(model, x, t_input, **model_kwargs)
+    ``
+    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
+
+    ===============================================================
+
+    Args:
+        model: A diffusion model with the corresponding format described above.
+        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+        model_type: A `str`. The parameterization type of the diffusion model.
+                    "noise" or "x_start" or "v" or "score".
+        model_kwargs: A `dict`. A dict for the other inputs of the model function.
+        guidance_type: A `str`. The type of the guidance for sampling.
+                    "uncond" or "classifier" or "classifier-free".
+        condition: A pytorch tensor. The condition for the guided sampling.
+                    Only used for "classifier" or "classifier-free" guidance type.
+        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
+                    Only used for "classifier-free" guidance type.
+        guidance_scale: A `float`. The scale for the guided sampling.
+        classifier_fn: A classifier function. Only used for the classifier guidance.
+        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
+    Returns:
+        A noise prediction model that accepts the noised data and the continuous time as the inputs.
+    """
+
+    def get_model_input_time(t_continuous):
+        """
+        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
+        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
+        For continuous-time DPMs, we just use `t_continuous`.
+        """
+        if noise_schedule.schedule == "discrete":
+            return (t_continuous - 1.0 / noise_schedule.total_N) * 1000.0
+        else:
+            return t_continuous
+
+    def noise_pred_fn(x, t_continuous, cond=None):
+        t_input = get_model_input_time(t_continuous)
+        if cond is None:
+            output = model(x, t_input, **model_kwargs)
+        else:
+            output = model(x, t_input, cond, **model_kwargs)
+        if model_type == "noise":
+            return output
+        elif model_type == "x_start":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(
+                t_continuous
+            ), noise_schedule.marginal_std(t_continuous)
+            return (x - alpha_t * output) / sigma_t
+        elif model_type == "v":
+            alpha_t, sigma_t = noise_schedule.marginal_alpha(
+                t_continuous
+            ), noise_schedule.marginal_std(t_continuous)
+            return alpha_t * output + sigma_t * x
+        elif model_type == "score":
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            return -sigma_t * output
+
+    def cond_grad_fn(x, t_input):
+        """
+        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
+        """
+        with torch.enable_grad():
+            x_in = x.detach().requires_grad_(True)
+            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
+            return torch.autograd.grad(log_prob.sum(), x_in)[0]
+
+    def model_fn(x, t_continuous):
+        """
+        The noise predicition model function that is used for DPM-Solver.
+        """
+        if guidance_type == "uncond":
+            return noise_pred_fn(x, t_continuous)
+        elif guidance_type == "classifier":
+            assert classifier_fn is not None
+            t_input = get_model_input_time(t_continuous)
+            cond_grad = cond_grad_fn(x, t_input)
+            sigma_t = noise_schedule.marginal_std(t_continuous)
+            noise = noise_pred_fn(x, t_continuous)
+            return noise - guidance_scale * sigma_t * cond_grad
+        elif guidance_type == "classifier-free":
+            if guidance_scale == 1.0 or unconditional_condition is None:
+                return noise_pred_fn(x, t_continuous, cond=condition)
+            else:
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                c_in = torch.cat([unconditional_condition, condition])
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+                return noise_uncond + guidance_scale * (noise - noise_uncond)
+
+    assert model_type in ["noise", "x_start", "v", "score"]
+    assert guidance_type in ["uncond", "classifier", "classifier-free"]
+    return model_fn
+
+
+class DPM_Solver:
+    def __init__(
+        self,
+        model_fn,
+        noise_schedule,
+        algorithm_type="dpmsolver++",
+        correcting_x0_fn=None,
+        correcting_xt_fn=None,
+        thresholding_max_val=1.0,
+        dynamic_thresholding_ratio=0.995,
+    ):
+        """Construct a DPM-Solver.
+
+        We support both DPM-Solver (`algorithm_type="dpmsolver"`) and DPM-Solver++ (`algorithm_type="dpmsolver++"`).
+
+        We also support the "dynamic thresholding" method in Imagen[1]. For pixel-space diffusion models, you
+        can set both `algorithm_type="dpmsolver++"` and `correcting_x0_fn="dynamic_thresholding"` to use the
+        dynamic thresholding. The "dynamic thresholding" can greatly improve the sample quality for pixel-space
+        DPMs with large guidance scales. Note that the thresholding method is **unsuitable** for latent-space
+        DPMs (such as stable-diffusion).
+
+        To support advanced algorithms in image-to-image applications, we also support corrector functions for
+        both x0 and xt.
+
+        Args:
+            model_fn: A noise prediction model function which accepts the continuous-time input (t in [epsilon, T]):
+                ``
+                def model_fn(x, t_continuous):
+                    return noise
+                ``
+                The shape of `x` is `(batch_size, **shape)`, and the shape of `t_continuous` is `(batch_size,)`.
+            noise_schedule: A noise schedule object, such as NoiseScheduleVP.
+            algorithm_type: A `str`. Either "dpmsolver" or "dpmsolver++".
+            correcting_x0_fn: A `str` or a function with the following format:
+                ```
+                def correcting_x0_fn(x0, t):
+                    x0_new = ...
+                    return x0_new
+                ```
+                This function is to correct the outputs of the data prediction model at each sampling step. e.g.,
+                ```
+                x0_pred = data_pred_model(xt, t)
+                if correcting_x0_fn is not None:
+                    x0_pred = correcting_x0_fn(x0_pred, t)
+                xt_1 = update(x0_pred, xt, t)
+                ```
+                If `correcting_x0_fn="dynamic_thresholding"`, we use the dynamic thresholding proposed in Imagen[1].
+            correcting_xt_fn: A function with the following format:
+                ```
+                def correcting_xt_fn(xt, t, step):
+                    x_new = ...
+                    return x_new
+                ```
+                This function is to correct the intermediate samples xt at each sampling step. e.g.,
+                ```
+                xt = ...
+                xt = correcting_xt_fn(xt, t, step)
+                ```
+            thresholding_max_val: A `float`. The max value for thresholding.
+                Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`.
+            dynamic_thresholding_ratio: A `float`. The ratio for dynamic thresholding (see Imagen[1] for details).
+                Valid only when use `dpmsolver++` and `correcting_x0_fn="dynamic_thresholding"`.
+
+        [1] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour,
+            Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models
+            with deep language understanding. arXiv preprint arXiv:2205.11487, 2022b.
+        """
+        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
+        self.noise_schedule = noise_schedule
+        assert algorithm_type in ["dpmsolver", "dpmsolver++"]
+        self.algorithm_type = algorithm_type
+        if correcting_x0_fn == "dynamic_thresholding":
+            self.correcting_x0_fn = self.dynamic_thresholding_fn
+        else:
+            self.correcting_x0_fn = correcting_x0_fn
+        self.correcting_xt_fn = correcting_xt_fn
+        self.dynamic_thresholding_ratio = dynamic_thresholding_ratio
+        self.thresholding_max_val = thresholding_max_val
+
+    def dynamic_thresholding_fn(self, x0, t):
+        """
+        The dynamic thresholding method.
+        """
+        dims = x0.dim()
+        p = self.dynamic_thresholding_ratio
+        s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
+        s = expand_dims(
+            torch.maximum(
+                s, self.thresholding_max_val * torch.ones_like(s).to(s.device)
+            ),
+            dims,
+        )
+        x0 = torch.clamp(x0, -s, s) / s
+        return x0
+
+    def noise_prediction_fn(self, x, t):
+        """
+        Return the noise prediction model.
+        """
+        return self.model(x, t)
+
+    def data_prediction_fn(self, x, t):
+        """
+        Return the data prediction model (with corrector).
+        """
+        noise = self.noise_prediction_fn(x, t)
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(
+            t
+        ), self.noise_schedule.marginal_std(t)
+        x0 = (x - sigma_t * noise) / alpha_t
+        if self.correcting_x0_fn is not None:
+            x0 = self.correcting_x0_fn(x0, t)
+        return x0
+
+    def model_fn(self, x, t):
+        """
+        Convert the model to the noise prediction model or the data prediction model.
+        """
+        if self.algorithm_type == "dpmsolver++":
+            return self.data_prediction_fn(x, t)
+        else:
+            return self.noise_prediction_fn(x, t)
+
+    def get_time_steps(self, skip_type, t_T, t_0, N, device):
+        """Compute the intermediate time steps for sampling.
+
+        Args:
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            N: A `int`. The total number of the spacing of the time steps.
+            device: A torch device.
+        Returns:
+            A pytorch tensor of the time steps, with the shape (N + 1,).
+        """
+        if skip_type == "logSNR":
+            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
+            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
+            logSNR_steps = torch.linspace(
+                lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1
+            ).to(device)
+            return self.noise_schedule.inverse_lambda(logSNR_steps)
+        elif skip_type == "time_uniform":
+            return torch.linspace(t_T, t_0, N + 1).to(device)
+        elif skip_type == "time_quadratic":
+            t_order = 2
+            t = (
+                torch.linspace(t_T ** (1.0 / t_order), t_0 ** (1.0 / t_order), N + 1)
+                .pow(t_order)
+                .to(device)
+            )
+            return t
+        else:
+            raise ValueError(
+                "Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(
+                    skip_type
+                )
+            )
+
+    def get_orders_and_timesteps_for_singlestep_solver(
+        self, steps, order, skip_type, t_T, t_0, device
+    ):
+        """
+        Get the order of each step for sampling by the singlestep DPM-Solver.
+
+        We combine both DPM-Solver-1,2,3 to use all the function evaluations, which is named as "DPM-Solver-fast".
+        Given a fixed number of function evaluations by `steps`, the sampling procedure by DPM-Solver-fast is:
+            - If order == 1:
+                We take `steps` of DPM-Solver-1 (i.e. DDIM).
+            - If order == 2:
+                - Denote K = (steps // 2). We take K or (K + 1) intermediate time steps for sampling.
+                - If steps % 2 == 0, we use K steps of DPM-Solver-2.
+                - If steps % 2 == 1, we use K steps of DPM-Solver-2 and 1 step of DPM-Solver-1.
+            - If order == 3:
+                - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                - If steps % 3 == 0, we use (K - 2) steps of DPM-Solver-3, and 1 step of DPM-Solver-2 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 1, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-1.
+                - If steps % 3 == 2, we use (K - 1) steps of DPM-Solver-3 and 1 step of DPM-Solver-2.
+
+        ============================================
+        Args:
+            order: A `int`. The max order for the solver (2 or 3).
+            steps: A `int`. The total number of function evaluations (NFE).
+            skip_type: A `str`. The type for the spacing of the time steps. We support three types:
+                - 'logSNR': uniform logSNR for the time steps.
+                - 'time_uniform': uniform time for the time steps. (**Recommended for high-resolutional data**.)
+                - 'time_quadratic': quadratic time for the time steps. (Used in DDIM for low-resolutional data.)
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            device: A torch device.
+        Returns:
+            orders: A list of the solver order of each step.
+        """
+        if order == 3:
+            K = steps // 3 + 1
+            if steps % 3 == 0:
+                orders = [3,] * (
+                    K - 2
+                ) + [2, 1]
+            elif steps % 3 == 1:
+                orders = [3,] * (
+                    K - 1
+                ) + [1]
+            else:
+                orders = [3,] * (
+                    K - 1
+                ) + [2]
+        elif order == 2:
+            if steps % 2 == 0:
+                K = steps // 2
+                orders = [
+                    2,
+                ] * K
+            else:
+                K = steps // 2 + 1
+                orders = [2,] * (
+                    K - 1
+                ) + [1]
+        elif order == 1:
+            K = 1
+            orders = [
+                1,
+            ] * steps
+        else:
+            raise ValueError("'order' must be '1' or '2' or '3'.")
+        if skip_type == "logSNR":
+            # To reproduce the results in DPM-Solver paper
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
+        else:
+            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[
+                torch.cumsum(
+                    torch.tensor(
+                        [
+                            0,
+                        ]
+                        + orders
+                    ),
+                    0,
+                ).to(device)
+            ]
+        return timesteps_outer, orders
+
+    def denoise_to_zero_fn(self, x, s):
+        """
+        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
+        """
+        return self.data_prediction_fn(x, s)
+
+    def dpm_solver_first_update(self, x, s, t, model_s=None, return_intermediate=False):
+        """
+        DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        dims = x.dim()
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        log_alpha_s, log_alpha_t = ns.marginal_log_mean_coeff(
+            s
+        ), ns.marginal_log_mean_coeff(t)
+        sigma_s, sigma_t = ns.marginal_std(s), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = sigma_t / sigma_s * x - alpha_t * phi_1 * model_s
+            if return_intermediate:
+                return x_t, {"model_s": model_s}
+            else:
+                return x_t
+        else:
+            phi_1 = torch.expm1(h)
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_t = torch.exp(log_alpha_t - log_alpha_s) * x - (sigma_t * phi_1) * model_s
+            if return_intermediate:
+                return x_t, {"model_s": model_s}
+            else:
+                return x_t
+
+    def singlestep_dpm_solver_second_update(
+        self,
+        x,
+        s,
+        t,
+        r1=0.5,
+        model_s=None,
+        return_intermediate=False,
+        solver_type="dpmsolver",
+    ):
+        """
+        Singlestep solver DPM-Solver-2 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            r1: A `float`. The hyperparameter of the second-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s` and `s1` (the intermediate time).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ["dpmsolver", "taylor"]:
+            raise ValueError(
+                "'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(
+                    solver_type
+                )
+            )
+        if r1 is None:
+            r1 = 0.5
+        ns = self.noise_schedule
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        log_alpha_s, log_alpha_s1, log_alpha_t = (
+            ns.marginal_log_mean_coeff(s),
+            ns.marginal_log_mean_coeff(s1),
+            ns.marginal_log_mean_coeff(t),
+        )
+        sigma_s, sigma_s1, sigma_t = (
+            ns.marginal_std(s),
+            ns.marginal_std(s1),
+            ns.marginal_std(t),
+        )
+        alpha_s1, alpha_t = torch.exp(log_alpha_s1), torch.exp(log_alpha_t)
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_11 = torch.expm1(-r1 * h)
+            phi_1 = torch.expm1(-h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (sigma_s1 / sigma_s) * x - (alpha_s1 * phi_11) * model_s
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    - (0.5 / r1) * (alpha_t * phi_1) * (model_s1 - model_s)
+                )
+            elif solver_type == "taylor":
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1.0 / r1) * (alpha_t * (phi_1 / h + 1.0)) * (model_s1 - model_s)
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_1 = torch.expm1(h)
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            x_s1 = (
+                torch.exp(log_alpha_s1 - log_alpha_s) * x
+                - (sigma_s1 * phi_11) * model_s
+            )
+            model_s1 = self.model_fn(x_s1, s1)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    torch.exp(log_alpha_t - log_alpha_s) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (0.5 / r1) * (sigma_t * phi_1) * (model_s1 - model_s)
+                )
+            elif solver_type == "taylor":
+                x_t = (
+                    torch.exp(log_alpha_t - log_alpha_s) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1.0 / r1) * (sigma_t * (phi_1 / h - 1.0)) * (model_s1 - model_s)
+                )
+        if return_intermediate:
+            return x_t, {"model_s": model_s, "model_s1": model_s1}
+        else:
+            return x_t
+
+    def singlestep_dpm_solver_third_update(
+        self,
+        x,
+        s,
+        t,
+        r1=1.0 / 3.0,
+        r2=2.0 / 3.0,
+        model_s=None,
+        model_s1=None,
+        return_intermediate=False,
+        solver_type="dpmsolver",
+    ):
+        """
+        Singlestep solver DPM-Solver-3 from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            r1: A `float`. The hyperparameter of the third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+            model_s: A pytorch tensor. The model function evaluated at time `s`.
+                If `model_s` is None, we evaluate the model by `x` and `s`; otherwise we directly use it.
+            model_s1: A pytorch tensor. The model function evaluated at time `s1` (the intermediate time given by `r1`).
+                If `model_s1` is None, we evaluate the model at `s1`; otherwise we directly use it.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ["dpmsolver", "taylor"]:
+            raise ValueError(
+                "'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(
+                    solver_type
+                )
+            )
+        if r1 is None:
+            r1 = 1.0 / 3.0
+        if r2 is None:
+            r2 = 2.0 / 3.0
+        ns = self.noise_schedule
+        lambda_s, lambda_t = ns.marginal_lambda(s), ns.marginal_lambda(t)
+        h = lambda_t - lambda_s
+        lambda_s1 = lambda_s + r1 * h
+        lambda_s2 = lambda_s + r2 * h
+        s1 = ns.inverse_lambda(lambda_s1)
+        s2 = ns.inverse_lambda(lambda_s2)
+        log_alpha_s, log_alpha_s1, log_alpha_s2, log_alpha_t = (
+            ns.marginal_log_mean_coeff(s),
+            ns.marginal_log_mean_coeff(s1),
+            ns.marginal_log_mean_coeff(s2),
+            ns.marginal_log_mean_coeff(t),
+        )
+        sigma_s, sigma_s1, sigma_s2, sigma_t = (
+            ns.marginal_std(s),
+            ns.marginal_std(s1),
+            ns.marginal_std(s2),
+            ns.marginal_std(t),
+        )
+        alpha_s1, alpha_s2, alpha_t = (
+            torch.exp(log_alpha_s1),
+            torch.exp(log_alpha_s2),
+            torch.exp(log_alpha_t),
+        )
+
+        if self.algorithm_type == "dpmsolver++":
+            phi_11 = torch.expm1(-r1 * h)
+            phi_12 = torch.expm1(-r2 * h)
+            phi_1 = torch.expm1(-h)
+            phi_22 = torch.expm1(-r2 * h) / (r2 * h) + 1.0
+            phi_2 = phi_1 / h + 1.0
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (sigma_s1 / sigma_s) * x - (alpha_s1 * phi_11) * model_s
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                (sigma_s2 / sigma_s) * x
+                - (alpha_s2 * phi_12) * model_s
+                + r2 / r1 * (alpha_s2 * phi_22) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1.0 / r2) * (alpha_t * phi_2) * (model_s2 - model_s)
+                )
+            elif solver_type == "taylor":
+                D1_0 = (1.0 / r1) * (model_s1 - model_s)
+                D1_1 = (1.0 / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2.0 * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (alpha_t * phi_2) * D1
+                    - (alpha_t * phi_3) * D2
+                )
+        else:
+            phi_11 = torch.expm1(r1 * h)
+            phi_12 = torch.expm1(r2 * h)
+            phi_1 = torch.expm1(h)
+            phi_22 = torch.expm1(r2 * h) / (r2 * h) - 1.0
+            phi_2 = phi_1 / h - 1.0
+            phi_3 = phi_2 / h - 0.5
+
+            if model_s is None:
+                model_s = self.model_fn(x, s)
+            if model_s1 is None:
+                x_s1 = (torch.exp(log_alpha_s1 - log_alpha_s)) * x - (
+                    sigma_s1 * phi_11
+                ) * model_s
+                model_s1 = self.model_fn(x_s1, s1)
+            x_s2 = (
+                (torch.exp(log_alpha_s2 - log_alpha_s)) * x
+                - (sigma_s2 * phi_12) * model_s
+                - r2 / r1 * (sigma_s2 * phi_22) * (model_s1 - model_s)
+            )
+            model_s2 = self.model_fn(x_s2, s2)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_s)) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1.0 / r2) * (sigma_t * phi_2) * (model_s2 - model_s)
+                )
+            elif solver_type == "taylor":
+                D1_0 = (1.0 / r1) * (model_s1 - model_s)
+                D1_1 = (1.0 / r2) * (model_s2 - model_s)
+                D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
+                D2 = 2.0 * (D1_1 - D1_0) / (r2 - r1)
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_s)) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (sigma_t * phi_2) * D1
+                    - (sigma_t * phi_3) * D2
+                )
+
+        if return_intermediate:
+            return x_t, {"model_s": model_s, "model_s1": model_s1, "model_s2": model_s2}
+        else:
+            return x_t
+
+    def multistep_dpm_solver_second_update(
+        self, x, model_prev_list, t_prev_list, t, solver_type="dpmsolver"
+    ):
+        """
+        Multistep solver DPM-Solver-2 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if solver_type not in ["dpmsolver", "taylor"]:
+            raise ValueError(
+                "'solver_type' must be either 'dpmsolver' or 'taylor', got {}".format(
+                    solver_type
+                )
+            )
+        ns = self.noise_schedule
+        model_prev_1, model_prev_0 = model_prev_list[-2], model_prev_list[-1]
+        t_prev_1, t_prev_0 = t_prev_list[-2], t_prev_list[-1]
+        lambda_prev_1, lambda_prev_0, lambda_t = (
+            ns.marginal_lambda(t_prev_1),
+            ns.marginal_lambda(t_prev_0),
+            ns.marginal_lambda(t),
+        )
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(
+            t_prev_0
+        ), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0 = h_0 / h
+        D1_0 = (1.0 / r0) * (model_prev_0 - model_prev_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    (sigma_t / sigma_prev_0) * x
+                    - (alpha_t * phi_1) * model_prev_0
+                    - 0.5 * (alpha_t * phi_1) * D1_0
+                )
+            elif solver_type == "taylor":
+                x_t = (
+                    (sigma_t / sigma_prev_0) * x
+                    - (alpha_t * phi_1) * model_prev_0
+                    + (alpha_t * (phi_1 / h + 1.0)) * D1_0
+                )
+        else:
+            phi_1 = torch.expm1(h)
+            if solver_type == "dpmsolver":
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                    - (sigma_t * phi_1) * model_prev_0
+                    - 0.5 * (sigma_t * phi_1) * D1_0
+                )
+            elif solver_type == "taylor":
+                x_t = (
+                    (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                    - (sigma_t * phi_1) * model_prev_0
+                    - (sigma_t * (phi_1 / h - 1.0)) * D1_0
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_update(
+        self, x, model_prev_list, t_prev_list, t, solver_type="dpmsolver"
+    ):
+        """
+        Multistep solver DPM-Solver-3 from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        ns = self.noise_schedule
+        model_prev_2, model_prev_1, model_prev_0 = model_prev_list
+        t_prev_2, t_prev_1, t_prev_0 = t_prev_list
+        lambda_prev_2, lambda_prev_1, lambda_prev_0, lambda_t = (
+            ns.marginal_lambda(t_prev_2),
+            ns.marginal_lambda(t_prev_1),
+            ns.marginal_lambda(t_prev_0),
+            ns.marginal_lambda(t),
+        )
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(
+            t_prev_0
+        ), ns.marginal_log_mean_coeff(t)
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h_1 = lambda_prev_1 - lambda_prev_2
+        h_0 = lambda_prev_0 - lambda_prev_1
+        h = lambda_t - lambda_prev_0
+        r0, r1 = h_0 / h, h_1 / h
+        D1_0 = (1.0 / r0) * (model_prev_0 - model_prev_1)
+        D1_1 = (1.0 / r1) * (model_prev_1 - model_prev_2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            phi_2 = phi_1 / h + 1.0
+            phi_3 = phi_2 / h - 0.5
+            x_t = (
+                (sigma_t / sigma_prev_0) * x
+                - (alpha_t * phi_1) * model_prev_0
+                + (alpha_t * phi_2) * D1
+                - (alpha_t * phi_3) * D2
+            )
+        else:
+            phi_1 = torch.expm1(h)
+            phi_2 = phi_1 / h - 1.0
+            phi_3 = phi_2 / h - 0.5
+            x_t = (
+                (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                - (sigma_t * phi_1) * model_prev_0
+                - (sigma_t * phi_2) * D1
+                - (sigma_t * phi_3) * D2
+            )
+        return x_t
+
+    def singlestep_dpm_solver_update(
+        self,
+        x,
+        s,
+        t,
+        order,
+        return_intermediate=False,
+        solver_type="dpmsolver",
+        r1=None,
+        r2=None,
+    ):
+        """
+        Singlestep DPM-Solver with the order `order` from time `s` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            s: A pytorch tensor. The starting time, with the shape (1,).
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            return_intermediate: A `bool`. If true, also return the model value at time `s`, `s1` and `s2` (the intermediate times).
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+            r1: A `float`. The hyperparameter of the second-order or third-order solver.
+            r2: A `float`. The hyperparameter of the third-order solver.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(
+                x, s, t, return_intermediate=return_intermediate
+            )
+        elif order == 2:
+            return self.singlestep_dpm_solver_second_update(
+                x,
+                s,
+                t,
+                return_intermediate=return_intermediate,
+                solver_type=solver_type,
+                r1=r1,
+            )
+        elif order == 3:
+            return self.singlestep_dpm_solver_third_update(
+                x,
+                s,
+                t,
+                return_intermediate=return_intermediate,
+                solver_type=solver_type,
+                r1=r1,
+                r2=r2,
+            )
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def multistep_dpm_solver_update(
+        self, x, model_prev_list, t_prev_list, t, order, solver_type="dpmsolver"
+    ):
+        """
+        Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `s`.
+            model_prev_list: A list of pytorch tensor. The previous computed model values.
+            t_prev_list: A list of pytorch tensor. The previous times, each time has the shape (1,)
+            t: A pytorch tensor. The ending time, with the shape (1,).
+            order: A `int`. The order of DPM-Solver. We only support order == 1 or 2 or 3.
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_t: A pytorch tensor. The approximated solution at time `t`.
+        """
+        if order == 1:
+            return self.dpm_solver_first_update(
+                x, t_prev_list[-1], t, model_s=model_prev_list[-1]
+            )
+        elif order == 2:
+            return self.multistep_dpm_solver_second_update(
+                x, model_prev_list, t_prev_list, t, solver_type=solver_type
+            )
+        elif order == 3:
+            return self.multistep_dpm_solver_third_update(
+                x, model_prev_list, t_prev_list, t, solver_type=solver_type
+            )
+        else:
+            raise ValueError("Solver order must be 1 or 2 or 3, got {}".format(order))
+
+    def dpm_solver_adaptive(
+        self,
+        x,
+        order,
+        t_T,
+        t_0,
+        h_init=0.05,
+        atol=0.0078,
+        rtol=0.05,
+        theta=0.9,
+        t_err=1e-5,
+        solver_type="dpmsolver",
+    ):
+        """
+        The adaptive step size solver based on singlestep DPM-Solver.
+
+        Args:
+            x: A pytorch tensor. The initial value at time `t_T`.
+            order: A `int`. The (higher) order of the solver. We only support order == 2 or 3.
+            t_T: A `float`. The starting time of the sampling (default is T).
+            t_0: A `float`. The ending time of the sampling (default is epsilon).
+            h_init: A `float`. The initial step size (for logSNR).
+            atol: A `float`. The absolute tolerance of the solver. For image data, the default setting is 0.0078, followed [1].
+            rtol: A `float`. The relative tolerance of the solver. The default setting is 0.05.
+            theta: A `float`. The safety hyperparameter for adapting the step size. The default setting is 0.9, followed [1].
+            t_err: A `float`. The tolerance for the time. We solve the diffusion ODE until the absolute error between the
+                current time and `t_0` is less than `t_err`. The default setting is 1e-5.
+            solver_type: either 'dpmsolver' or 'taylor'. The type for the high-order solvers.
+                The type slightly impacts the performance. We recommend to use 'dpmsolver' type.
+        Returns:
+            x_0: A pytorch tensor. The approximated solution at time `t_0`.
+
+        [1] A. Jolicoeur-Martineau, K. Li, R. Piché-Taillefer, T. Kachman, and I. Mitliagkas, "Gotta go fast when generating data with score-based models," arXiv preprint arXiv:2105.14080, 2021.
+        """
+        ns = self.noise_schedule
+        s = t_T * torch.ones((1,)).to(x)
+        lambda_s = ns.marginal_lambda(s)
+        lambda_0 = ns.marginal_lambda(t_0 * torch.ones_like(s).to(x))
+        h = h_init * torch.ones_like(s).to(x)
+        x_prev = x
+        nfe = 0
+        if order == 2:
+            r1 = 0.5
+
+            def lower_update(x, s, t):
+                return self.dpm_solver_first_update(x, s, t, return_intermediate=True)
+
+            def higher_update(x, s, t, **kwargs):
+                return self.singlestep_dpm_solver_second_update(
+                    x, s, t, r1=r1, solver_type=solver_type, **kwargs
+                )
+
+        elif order == 3:
+            r1, r2 = 1.0 / 3.0, 2.0 / 3.0
+
+            def lower_update(x, s, t):
+                return self.singlestep_dpm_solver_second_update(
+                    x, s, t, r1=r1, return_intermediate=True, solver_type=solver_type
+                )
+
+            def higher_update(x, s, t, **kwargs):
+                return self.singlestep_dpm_solver_third_update(
+                    x, s, t, r1=r1, r2=r2, solver_type=solver_type, **kwargs
+                )
+
+        else:
+            raise ValueError(
+                "For adaptive step size solver, order must be 2 or 3, got {}".format(
+                    order
+                )
+            )
+        while torch.abs((s - t_0)).mean() > t_err:
+            t = ns.inverse_lambda(lambda_s + h)
+            x_lower, lower_noise_kwargs = lower_update(x, s, t)
+            x_higher = higher_update(x, s, t, **lower_noise_kwargs)
+            delta = torch.max(
+                torch.ones_like(x).to(x) * atol,
+                rtol * torch.max(torch.abs(x_lower), torch.abs(x_prev)),
+            )
+
+            def norm_fn(v):
+                return torch.sqrt(
+                    torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True)
+                )
+
+            E = norm_fn((x_higher - x_lower) / delta).max()
+            if torch.all(E <= 1.0):
+                x = x_higher
+                s = t
+                x_prev = x_lower
+                lambda_s = ns.marginal_lambda(s)
+            h = torch.min(
+                theta * h * torch.float_power(E, -1.0 / order).float(),
+                lambda_0 - lambda_s,
+            )
+            nfe += order
+        print("adaptive solver nfe", nfe)
+        return x
+
+    def add_noise(self, x, t, noise=None):
+        """
+        Compute the noised input xt = alpha_t * x + sigma_t * noise.
+
+        Args:
+            x: A `torch.Tensor` with shape `(batch_size, *shape)`.
+            t: A `torch.Tensor` with shape `(t_size,)`.
+        Returns:
+            xt with shape `(t_size, batch_size, *shape)`.
+        """
+        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(
+            t
+        ), self.noise_schedule.marginal_std(t)
+        if noise is None:
+            noise = torch.randn((t.shape[0], *x.shape), device=x.device)
+        x = x.reshape((-1, *x.shape))
+        xt = expand_dims(alpha_t, x.dim()) * x + expand_dims(sigma_t, x.dim()) * noise
+        if t.shape[0] == 1:
+            return xt.squeeze(0)
+        else:
+            return xt
+
+    def inverse(
+        self,
+        x,
+        steps=20,
+        t_start=None,
+        t_end=None,
+        order=2,
+        skip_type="time_uniform",
+        method="multistep",
+        lower_order_final=True,
+        denoise_to_zero=False,
+        solver_type="dpmsolver",
+        atol=0.0078,
+        rtol=0.05,
+        return_intermediate=False,
+    ):
+        """
+        Inverse the sample `x` from time `t_start` to `t_end` by DPM-Solver.
+        For discrete-time DPMs, we use `t_start=1/N`, where `N` is the total time steps during training.
+        """
+        t_0 = 1.0 / self.noise_schedule.total_N if t_start is None else t_start
+        t_T = self.noise_schedule.T if t_end is None else t_end
+        assert (
+            t_0 > 0 and t_T > 0
+        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
+        return self.sample(
+            x,
+            steps=steps,
+            t_start=t_0,
+            t_end=t_T,
+            order=order,
+            skip_type=skip_type,
+            method=method,
+            lower_order_final=lower_order_final,
+            denoise_to_zero=denoise_to_zero,
+            solver_type=solver_type,
+            atol=atol,
+            rtol=rtol,
+            return_intermediate=return_intermediate,
+        )
+
+    def sample(
+        self,
+        x,
+        steps=20,
+        t_start=None,
+        t_end=None,
+        order=2,
+        skip_type="time_uniform",
+        method="multistep",
+        lower_order_final=True,
+        denoise_to_zero=False,
+        solver_type="dpmsolver",
+        atol=0.0078,
+        rtol=0.05,
+        return_intermediate=False,
+    ):
+        """
+        Compute the sample at time `t_end` by DPM-Solver, given the initial `x` at time `t_start`.
+
+        =====================================================
+
+        We support the following algorithms for both noise prediction model and data prediction model:
+            - 'singlestep':
+                Singlestep DPM-Solver (i.e. "DPM-Solver-fast" in the paper), which combines different orders of singlestep DPM-Solver.
+                We combine all the singlestep solvers with order <= `order` to use up all the function evaluations (steps).
+                The total number of function evaluations (NFE) == `steps`.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    - If `order` == 1:
+                        - Denote K = steps. We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - Denote K = (steps // 2) + (steps % 2). We take K intermediate time steps for sampling.
+                        - If steps % 2 == 0, we use K steps of singlestep DPM-Solver-2.
+                        - If steps % 2 == 1, we use (K - 1) steps of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                    - If `order` == 3:
+                        - Denote K = (steps // 3 + 1). We take K intermediate time steps for sampling.
+                        - If steps % 3 == 0, we use (K - 2) steps of singlestep DPM-Solver-3, and 1 step of singlestep DPM-Solver-2 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 1, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of DPM-Solver-1.
+                        - If steps % 3 == 2, we use (K - 1) steps of singlestep DPM-Solver-3 and 1 step of singlestep DPM-Solver-2.
+            - 'multistep':
+                Multistep DPM-Solver with the order of `order`. The total number of function evaluations (NFE) == `steps`.
+                We initialize the first `order` values by lower order multistep solvers.
+                Given a fixed NFE == `steps`, the sampling procedure is:
+                    Denote K = steps.
+                    - If `order` == 1:
+                        - We use K steps of DPM-Solver-1 (i.e. DDIM).
+                    - If `order` == 2:
+                        - We firstly use 1 step of DPM-Solver-1, then use (K - 1) step of multistep DPM-Solver-2.
+                    - If `order` == 3:
+                        - We firstly use 1 step of DPM-Solver-1, then 1 step of multistep DPM-Solver-2, then (K - 2) step of multistep DPM-Solver-3.
+            - 'singlestep_fixed':
+                Fixed order singlestep DPM-Solver (i.e. DPM-Solver-1 or singlestep DPM-Solver-2 or singlestep DPM-Solver-3).
+                We use singlestep DPM-Solver-`order` for `order`=1 or 2 or 3, with total [`steps` // `order`] * `order` NFE.
+            - 'adaptive':
+                Adaptive step size DPM-Solver (i.e. "DPM-Solver-12" and "DPM-Solver-23" in the paper).
+                We ignore `steps` and use adaptive step size DPM-Solver with a higher order of `order`.
+                You can adjust the absolute tolerance `atol` and the relative tolerance `rtol` to balance the computatation costs
+                (NFE) and the sample quality.
+                    - If `order` == 2, we use DPM-Solver-12 which combines DPM-Solver-1 and singlestep DPM-Solver-2.
+                    - If `order` == 3, we use DPM-Solver-23 which combines singlestep DPM-Solver-2 and singlestep DPM-Solver-3.
+
+        =====================================================
+
+        Some advices for choosing the algorithm:
+            - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
+                Use singlestep DPM-Solver or DPM-Solver++ ("DPM-Solver-fast" in the paper) with `order = 3`.
+                e.g., DPM-Solver:
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+                            skip_type='time_uniform', method='singlestep')
+                e.g., DPM-Solver++:
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=3,
+                            skip_type='time_uniform', method='singlestep')
+            - For **guided sampling with large guidance scale** by DPMs:
+                Use multistep DPM-Solver with `algorithm_type="dpmsolver++"` and `order = 2`.
+                e.g.
+                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+                    >>> x_sample = dpm_solver.sample(x, steps=steps, t_start=t_start, t_end=t_end, order=2,
+                            skip_type='time_uniform', method='multistep')
+
+        We support three types of `skip_type`:
+            - 'logSNR': uniform logSNR for the time steps. **Recommended for low-resolutional images**
+            - 'time_uniform': uniform time for the time steps. **Recommended for high-resolutional images**.
+            - 'time_quadratic': quadratic time for the time steps.
+
+        =====================================================
+        Args:
+            x: A pytorch tensor. The initial value at time `t_start`
+                e.g. if `t_start` == T, then `x` is a sample from the standard normal distribution.
+            steps: A `int`. The total number of function evaluations (NFE).
+            t_start: A `float`. The starting time of the sampling.
+                If `T` is None, we use self.noise_schedule.T (default is 1.0).
+            t_end: A `float`. The ending time of the sampling.
+                If `t_end` is None, we use 1. / self.noise_schedule.total_N.
+                e.g. if total_N == 1000, we have `t_end` == 1e-3.
+                For discrete-time DPMs:
+                    - We recommend `t_end` == 1. / self.noise_schedule.total_N.
+                For continuous-time DPMs:
+                    - We recommend `t_end` == 1e-3 when `steps` <= 15; and `t_end` == 1e-4 when `steps` > 15.
+            order: A `int`. The order of DPM-Solver.
+            skip_type: A `str`. The type for the spacing of the time steps. 'time_uniform' or 'logSNR' or 'time_quadratic'.
+            method: A `str`. The method for sampling. 'singlestep' or 'multistep' or 'singlestep_fixed' or 'adaptive'.
+            denoise_to_zero: A `bool`. Whether to denoise to time 0 at the final step.
+                Default is `False`. If `denoise_to_zero` is `True`, the total NFE is (`steps` + 1).
+
+                This trick is firstly proposed by DDPM (https://arxiv.org/abs/2006.11239) and
+                score_sde (https://arxiv.org/abs/2011.13456). Such trick can improve the FID
+                for diffusion models sampling by diffusion SDEs for low-resolutional images
+                (such as CIFAR-10). However, we observed that such trick does not matter for
+                high-resolutional images. As it needs an additional NFE, we do not recommend
+                it for high-resolutional images.
+            lower_order_final: A `bool`. Whether to use lower order solvers at the final steps.
+                Only valid for `method=multistep` and `steps < 15`. We empirically find that
+                this trick is a key to stabilizing the sampling by DPM-Solver with very few steps
+                (especially for steps <= 10). So we recommend to set it to be `True`.
+            solver_type: A `str`. The taylor expansion type for the solver. `dpmsolver` or `taylor`. We recommend `dpmsolver`.
+            atol: A `float`. The absolute tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+            rtol: A `float`. The relative tolerance of the adaptive step size solver. Valid when `method` == 'adaptive'.
+            return_intermediate: A `bool`. Whether to save the xt at each step.
+                When set to `True`, method returns a tuple (x0, intermediates); when set to False, method returns only x0.
+        Returns:
+            x_end: A pytorch tensor. The approximated solution at time `t_end`.
+
+        """
+        t_0 = 1.0 / self.noise_schedule.total_N if t_end is None else t_end
+        t_T = self.noise_schedule.T if t_start is None else t_start
+        assert (
+            t_0 > 0 and t_T > 0
+        ), "Time range needs to be greater than 0. For discrete-time DPMs, it needs to be in [1 / N, 1], where N is the length of betas array"
+        if return_intermediate:
+            assert method in [
+                "multistep",
+                "singlestep",
+                "singlestep_fixed",
+            ], "Cannot use adaptive solver when saving intermediate values"
+        if self.correcting_xt_fn is not None:
+            assert method in [
+                "multistep",
+                "singlestep",
+                "singlestep_fixed",
+            ], "Cannot use adaptive solver when correcting_xt_fn is not None"
+        device = x.device
+        intermediates = []
+        with torch.no_grad():
+            if method == "adaptive":
+                x = self.dpm_solver_adaptive(
+                    x,
+                    order=order,
+                    t_T=t_T,
+                    t_0=t_0,
+                    atol=atol,
+                    rtol=rtol,
+                    solver_type=solver_type,
+                )
+            elif method == "multistep":
+                assert steps >= order
+                timesteps = self.get_time_steps(
+                    skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device
+                )
+                assert timesteps.shape[0] - 1 == steps
+                # Init the initial values.
+                step = 0
+                t = timesteps[step]
+                t_prev_list = [t]
+                model_prev_list = [self.model_fn(x, t)]
+                if self.correcting_xt_fn is not None:
+                    x = self.correcting_xt_fn(x, t, step)
+                if return_intermediate:
+                    intermediates.append(x)
+                # Init the first `order` values by lower order multistep DPM-Solver.
+                for step in range(1, order):
+                    t = timesteps[step]
+                    x = self.multistep_dpm_solver_update(
+                        x,
+                        model_prev_list,
+                        t_prev_list,
+                        t,
+                        step,
+                        solver_type=solver_type,
+                    )
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+                    t_prev_list.append(t)
+                    model_prev_list.append(self.model_fn(x, t))
+                # Compute the remaining values by `order`-th order multistep DPM-Solver.
+                for step in range(order, steps + 1):
+                    t = timesteps[step]
+                    # We only use lower order for steps < 10
+                    if lower_order_final and steps < 10:
+                        step_order = min(order, steps + 1 - step)
+                    else:
+                        step_order = order
+                    x = self.multistep_dpm_solver_update(
+                        x,
+                        model_prev_list,
+                        t_prev_list,
+                        t,
+                        step_order,
+                        solver_type=solver_type,
+                    )
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+                    for i in range(order - 1):
+                        t_prev_list[i] = t_prev_list[i + 1]
+                        model_prev_list[i] = model_prev_list[i + 1]
+                    t_prev_list[-1] = t
+                    # We do not need to evaluate the final model value.
+                    if step < steps:
+                        model_prev_list[-1] = self.model_fn(x, t)
+            elif method in ["singlestep", "singlestep_fixed"]:
+                if method == "singlestep":
+                    (
+                        timesteps_outer,
+                        orders,
+                    ) = self.get_orders_and_timesteps_for_singlestep_solver(
+                        steps=steps,
+                        order=order,
+                        skip_type=skip_type,
+                        t_T=t_T,
+                        t_0=t_0,
+                        device=device,
+                    )
+                elif method == "singlestep_fixed":
+                    K = steps // order
+                    orders = [
+                        order,
+                    ] * K
+                    timesteps_outer = self.get_time_steps(
+                        skip_type=skip_type, t_T=t_T, t_0=t_0, N=K, device=device
+                    )
+                for step, order in enumerate(orders):
+                    s, t = timesteps_outer[step], timesteps_outer[step + 1]
+                    timesteps_inner = self.get_time_steps(
+                        skip_type=skip_type,
+                        t_T=s.item(),
+                        t_0=t.item(),
+                        N=order,
+                        device=device,
+                    )
+                    lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
+                    h = lambda_inner[-1] - lambda_inner[0]
+                    r1 = None if order <= 1 else (lambda_inner[1] - lambda_inner[0]) / h
+                    r2 = None if order <= 2 else (lambda_inner[2] - lambda_inner[0]) / h
+                    x = self.singlestep_dpm_solver_update(
+                        x, s, t, order, solver_type=solver_type, r1=r1, r2=r2
+                    )
+                    if self.correcting_xt_fn is not None:
+                        x = self.correcting_xt_fn(x, t, step)
+                    if return_intermediate:
+                        intermediates.append(x)
+            else:
+                raise ValueError("Got wrong method {}".format(method))
+            if denoise_to_zero:
+                t = torch.ones((1,)).to(device) * t_0
+                x = self.denoise_to_zero_fn(x, t)
+                if self.correcting_xt_fn is not None:
+                    x = self.correcting_xt_fn(x, t, step + 1)
+                if return_intermediate:
+                    intermediates.append(x)
+        if return_intermediate:
+            return x, intermediates
+        else:
+            return x
+
+
+#############################################################
+# other utility functions
+#############################################################
+
+
+def interpolate_fn(x, xp, yp):
+    """
+    A piecewise linear function y = f(x), using xp and yp as keypoints.
+    We implement f(x) in a differentiable way (i.e. applicable for autograd).
+    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
+
+    Args:
+        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
+        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
+        yp: PyTorch tensor with shape [C, K].
+    Returns:
+        The function values f(x), with shape [N, C].
+    """
+    N, K = x.shape[0], xp.shape[1]
+    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
+    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
+    x_idx = torch.argmin(x_indices, dim=2)
+    cand_start_idx = x_idx - 1
+    start_idx = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(1, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K),
+            torch.tensor(K - 2, device=x.device),
+            cand_start_idx,
+        ),
+    )
+    end_idx = torch.where(
+        torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1
+    )
+    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
+    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
+    start_idx2 = torch.where(
+        torch.eq(x_idx, 0),
+        torch.tensor(0, device=x.device),
+        torch.where(
+            torch.eq(x_idx, K),
+            torch.tensor(K - 2, device=x.device),
+            cand_start_idx,
+        ),
+    )
+    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
+    start_y = torch.gather(
+        y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)
+    ).squeeze(2)
+    end_y = torch.gather(
+        y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)
+    ).squeeze(2)
+    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
+    return cand
+
+
+def expand_dims(v, dims):
+    """
+    Expand the tensor `v` to the dim `dims`.
+
+    Args:
+        `v`: a PyTorch tensor with shape [N].
+        `dim`: a `int`.
+    Returns:
+        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
+    """
+    return v[(...,) + (None,) * (dims - 1)]

tortoise/get_conditioning_latents.py

@@ -0,0 +1,51 @@
+# +
+import argparse
+import os
+
+import torch
+from api import TextToSpeech
+
+from tortoise.utils.audio import get_voices, load_required_audio
+
+"""
+Dumps the conditioning latents for the specified voice to disk. These are expressive latents which can be used for
+other ML models, or can be augmented manually and fed back into Tortoise to affect vocal qualities.
+"""
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--voice",
+        type=str,
+        help="Selects the voice to convert to conditioning latents",
+        default="pat2",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        help="Where to store outputs.",
+        default="../results/conditioning_latents",
+    )
+    parser.add_argument(
+        "--latent_averaging_mode",
+        type=int,
+        help="How to average voice latents, 0 for standard, 1 for per-sample, 2 for per-minichunk",
+        default=0,
+    )
+
+    args = parser.parse_args()
+    os.makedirs(args.output_path, exist_ok=True)
+
+    tts = TextToSpeech()
+    voices = get_voices()
+    print(list(voices.keys()))
+    selected_voices = args.voice.split(",")
+    for voice in selected_voices:
+        cond_paths = voices[voice]
+        conds = []
+        for cond_path in cond_paths:
+            c = load_required_audio(cond_path)
+            conds.append(c)
+        conditioning_latents = tts.get_conditioning_latents(
+            conds, latent_averaging_mode=args.latent_averaging_mode
+        )
+        torch.save(conditioning_latents, os.path.join(args.output_path, f"{voice}.pth"))

tortoise/inference.py

@@ -0,0 +1,179 @@
+import os
+import sys
+from random import randint
+from typing import List, Optional, Set, Union
+
+from tortoise.utils.audio import get_voices, load_audio, load_voices
+from tortoise.utils.text import split_and_recombine_text
+
+
+def get_all_voices(extra_voice_dirs_str: str = ""):
+    extra_voice_dirs = extra_voice_dirs_str.split(",") if extra_voice_dirs_str else []
+    return sorted(get_voices(extra_voice_dirs)), extra_voice_dirs
+
+
+def parse_voice_str(voice_str: str, all_voices: List[str]):
+    selected_voices = all_voices if voice_str == "all" else voice_str.split(",")
+    selected_voices = [v.split("&") if "&" in v else [v] for v in selected_voices]
+    for voices in selected_voices:
+        for v in voices:
+            if v != "random" and v not in all_voices:
+                raise ValueError(
+                    f"voice {v} not available, use --list-voices to see available voices."
+                )
+
+    return selected_voices
+
+
+def voice_loader(selected_voices: list, extra_voice_dirs: List[str]):
+    for voices in selected_voices:
+        yield voices, *load_voices(voices, extra_voice_dirs)
+
+
+def parse_multiarg_text(text: List[str]):
+    return (" ".join(text) if text else "".join(line for line in sys.stdin)).strip()
+
+
+def split_text(text: str, text_split: str):
+    if text_split:
+        desired_length, max_length = map(int, text_split.split(","))
+        if desired_length > max_length:
+            raise ValueError(
+                f"--text-split: desired_length ({desired_length}) must be <= max_length ({max_length})"
+            )
+        texts = split_and_recombine_text(text, desired_length, max_length)
+    else:
+        texts = split_and_recombine_text(text)
+    #
+    if not texts:
+        raise ValueError("no text provided")
+    return texts
+
+
+def validate_output_dir(output_dir: str, selected_voices: list, candidates: int):
+    if output_dir:
+        os.makedirs(output_dir, exist_ok=True)
+    else:
+        if len(selected_voices) > 1:
+            raise ValueError('cannot have multiple voices without --output-dir"')
+        if candidates > 1:
+            raise ValueError('cannot have multiple candidates without --output-dir"')
+    return output_dir
+
+
+def check_pydub(play: bool):
+    if play:
+        try:
+            import pydub
+            import pydub.playback
+
+            return pydub
+        except ImportError:
+            raise RuntimeError(
+                '--play requires pydub to be installed, which can be done with "pip install pydub"'
+            )
+
+
+def get_seed(seed: Optional[int]):
+    return randint(0, 2**32 - 1) if seed is None else seed
+
+
+from pathlib import Path
+from typing import Any, Callable
+
+import torch
+import torchaudio
+
+
+def run_and_save_tts(
+    call_tts,
+    text,
+    output_dir: Path,
+    return_deterministic_state,
+    return_filepaths=False,
+    voicefixer=True,
+):
+    output_dir.mkdir(exist_ok=True)
+    if return_deterministic_state:
+        gen, dbg = call_tts(text)
+        torch.save(dbg, output_dir / "dbg.pt")
+    else:
+        gen = call_tts(text)
+    #
+    if not isinstance(gen, list):
+        gen = [gen]
+    gen = [g.squeeze(0).cpu() for g in gen]
+    fps = []
+    for i, g in enumerate(gen):
+        fps.append(output_dir / f"{i}.wav")
+        save_gen_with_voicefix(g, fps[-1], squeeze=False, voicefixer=voicefixer)
+        # torchaudio.save(output_dir/f'{i}.wav', g, 24000)
+    return fps if return_filepaths else gen
+
+
+def infer_on_texts(
+    call_tts: Callable[[str], Any],
+    texts: List[str],
+    output_dir: Union[str, Path],
+    return_deterministic_state: bool,
+    lines_to_regen: Set[int],
+    logger=print,
+    return_filepaths=False,
+    voicefixer=True,
+):
+    audio_chunks = []
+    base_p = Path(output_dir)
+    base_p.mkdir(exist_ok=True)
+
+    for text_idx, text in enumerate(texts):
+        line_p = base_p / f"{text_idx}"
+        line_p.mkdir(exist_ok=True)
+        #
+        if text_idx not in lines_to_regen:
+            files = list(line_p.glob("*.wav"))
+            if files:
+                logger(f"loading existing audio fragments for [{text_idx}]")
+                audio_chunks.append([load_audio(str(f), 24000) for f in files])
+                continue
+            else:
+                logger(f"no existing audio fragment for [{text_idx}]")
+        #
+        logger(f"generating audio for text {text_idx}: {text}")
+        audio_chunks.append(
+            run_and_save_tts(
+                call_tts,
+                text,
+                line_p,
+                return_deterministic_state,
+                voicefixer=voicefixer,
+            )
+        )
+
+    fnames = []
+    results = []
+    for i in range(len(audio_chunks[0])):
+        resultant = torch.cat([c[i] for c in audio_chunks], dim=-1)
+        fnames.append(base_p / f"combined-{i}.wav")
+        save_gen_with_voicefix(
+            resultant, fnames[-1], squeeze=False, voicefixer=False
+        )  # do not run fix on combined!!
+        results.append(resultant)
+        # torchaudio.save(base_p/'combined.wav', resultant, 24000)
+    return fnames if return_filepaths else results
+
+
+from voicefixer import VoiceFixer
+
+vfixer = VoiceFixer()
+
+
+def save_gen_with_voicefix(g, fpath, squeeze=True, voicefixer=True):
+    torchaudio.save(fpath, g.squeeze(0).cpu() if squeeze else g, 24000, format="wav")
+    if voicefixer:
+        vfixer.restore(
+            input=fpath,
+            output=fpath,
+            cuda=True,
+            mode=0,
+            # your_vocoder_func = convert_mel_to_wav # TODO test if integration with unvinet improves things
+        )

tortoise/models/arch_util.py

@@ -0,0 +1,425 @@
+import functools
+import math
+import os
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+
+from tortoise.models.xtransformers import (
+    ContinuousTransformerWrapper,
+    RelativePositionBias,
+)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class GroupNorm32(nn.GroupNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+
+def normalization(channels):
+    """
+    Make a standard normalization layer.
+
+    :param channels: number of input channels.
+    :return: an nn.Module for normalization.
+    """
+    groups = 32
+    if channels <= 16:
+        groups = 8
+    elif channels <= 64:
+        groups = 16
+    while channels % groups != 0:
+        groups = int(groups / 2)
+    assert groups > 2
+    return GroupNorm32(groups, channels)
+
+
+class QKVAttentionLegacy(nn.Module):
+    """
+    A module which performs QKV attention. Matches legacy QKVAttention + input/output heads shaping
+    """
+
+    def __init__(self, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+
+    def forward(self, qkv, mask=None, rel_pos=None):
+        """
+        Apply QKV attention.
+
+        :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+        :return: an [N x (H * C) x T] tensor after attention.
+        """
+        bs, width, length = qkv.shape
+        assert width % (3 * self.n_heads) == 0
+        ch = width // (3 * self.n_heads)
+        q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+        scale = 1 / math.sqrt(math.sqrt(ch))
+        weight = torch.einsum(
+            "bct,bcs->bts", q * scale, k * scale
+        )  # More stable with f16 than dividing afterwards
+        if rel_pos is not None:
+            weight = rel_pos(
+                weight.reshape(bs, self.n_heads, weight.shape[-2], weight.shape[-1])
+            ).reshape(bs * self.n_heads, weight.shape[-2], weight.shape[-1])
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        if mask is not None:
+            # The proper way to do this is to mask before the softmax using -inf, but that doesn't work properly on CPUs.
+            mask = mask.repeat(self.n_heads, 1).unsqueeze(1)
+            weight = weight * mask
+        a = torch.einsum("bts,bcs->bct", weight, v)
+
+        return a.reshape(bs, -1, length)
+
+
+class AttentionBlock(nn.Module):
+    """
+    An attention block that allows spatial positions to attend to each other.
+
+    Originally ported from here, but adapted to the N-d case.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+    """
+
+    def __init__(
+        self,
+        channels,
+        num_heads=1,
+        num_head_channels=-1,
+        do_checkpoint=True,
+        relative_pos_embeddings=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.do_checkpoint = do_checkpoint
+        if num_head_channels == -1:
+            self.num_heads = num_heads
+        else:
+            assert (
+                channels % num_head_channels == 0
+            ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+            self.num_heads = channels // num_head_channels
+        self.norm = normalization(channels)
+        self.qkv = nn.Conv1d(channels, channels * 3, 1)
+        # split heads before split qkv
+        self.attention = QKVAttentionLegacy(self.num_heads)
+
+        self.proj_out = zero_module(nn.Conv1d(channels, channels, 1))
+        if relative_pos_embeddings:
+            self.relative_pos_embeddings = RelativePositionBias(
+                scale=(channels // self.num_heads) ** 0.5,
+                causal=False,
+                heads=num_heads,
+                num_buckets=32,
+                max_distance=64,
+            )
+        else:
+            self.relative_pos_embeddings = None
+
+    def forward(self, x, mask=None):
+        b, c, *spatial = x.shape
+        x = x.reshape(b, c, -1)
+        qkv = self.qkv(self.norm(x))
+        h = self.attention(qkv, mask, self.relative_pos_embeddings)
+        h = self.proj_out(h)
+        return (x + h).reshape(b, c, *spatial)
+
+
+class Upsample(nn.Module):
+    """
+    An upsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.factor = factor
+        if use_conv:
+            ksize = 5
+            pad = 2
+            self.conv = nn.Conv1d(self.channels, self.out_channels, ksize, padding=pad)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        x = F.interpolate(x, scale_factor=self.factor, mode="nearest")
+        if self.use_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    """
+    A downsampling layer with an optional convolution.
+
+    :param channels: channels in the inputs and outputs.
+    :param use_conv: a bool determining if a convolution is applied.
+    """
+
+    def __init__(self, channels, use_conv, out_channels=None, factor=4, ksize=5, pad=2):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+
+        stride = factor
+        if use_conv:
+            self.op = nn.Conv1d(
+                self.channels, self.out_channels, ksize, stride=stride, padding=pad
+            )
+        else:
+            assert self.channels == self.out_channels
+            self.op = nn.AvgPool1d(kernel_size=stride, stride=stride)
+
+    def forward(self, x):
+        assert x.shape[1] == self.channels
+        return self.op(x)
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        up=False,
+        down=False,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+        padding = 1 if kernel_size == 3 else 2
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False)
+            self.x_upd = Upsample(channels, False)
+        elif down:
+            self.h_upd = Downsample(channels, False)
+            self.x_upd = Downsample(channels, False)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                nn.Conv1d(
+                    self.out_channels, self.out_channels, kernel_size, padding=padding
+                )
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = nn.Conv1d(
+                channels, self.out_channels, kernel_size, padding=padding
+            )
+        else:
+            self.skip_connection = nn.Conv1d(channels, self.out_channels, 1)
+
+    def forward(self, x):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AudioMiniEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        base_channels=128,
+        depth=2,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        dropout=0,
+        downsample_factor=2,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.init = nn.Sequential(nn.Conv1d(spec_dim, base_channels, 3, padding=1))
+        ch = base_channels
+        res = []
+        for l in range(depth):
+            for r in range(resnet_blocks):
+                res.append(ResBlock(ch, dropout, kernel_size=kernel_size))
+            res.append(
+                Downsample(
+                    ch, use_conv=True, out_channels=ch * 2, factor=downsample_factor
+                )
+            )
+            ch *= 2
+        self.res = nn.Sequential(*res)
+        self.final = nn.Sequential(
+            normalization(ch), nn.SiLU(), nn.Conv1d(ch, embedding_dim, 1)
+        )
+        attn = []
+        for a in range(attn_blocks):
+            attn.append(
+                AttentionBlock(
+                    embedding_dim,
+                    num_attn_heads,
+                )
+            )
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.res(h)
+        h = self.final(h)
+        h = self.attn(h)
+        return h[:, :, 0]
+
+
+DEFAULT_MEL_NORM_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "../data/mel_norms.pth"
+)
+
+
+class TorchMelSpectrogram(nn.Module):
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=80,
+        mel_fmin=0,
+        mel_fmax=8000,
+        sampling_rate=22050,
+        normalize=False,
+        mel_norm_file=DEFAULT_MEL_NORM_FILE,
+    ):
+        super().__init__()
+        # These are the default tacotron values for the MEL spectrogram.
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.mel_fmin = mel_fmin
+        self.mel_fmax = mel_fmax
+        self.sampling_rate = sampling_rate
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(
+            n_fft=self.filter_length,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            power=2,
+            normalized=normalize,
+            sample_rate=self.sampling_rate,
+            f_min=self.mel_fmin,
+            f_max=self.mel_fmax,
+            n_mels=self.n_mel_channels,
+            norm="slaney",
+        )
+        self.mel_norm_file = mel_norm_file
+        if self.mel_norm_file is not None:
+            self.mel_norms = torch.load(self.mel_norm_file)
+        else:
+            self.mel_norms = None
+
+    def forward(self, inp):
+        if (
+            len(inp.shape) == 3
+        ):  # Automatically squeeze out the channels dimension if it is present (assuming mono-audio)
+            inp = inp.squeeze(1)
+        assert len(inp.shape) == 2
+        self.mel_stft = self.mel_stft.to(inp.device)
+        mel = self.mel_stft(inp)
+        # Perform dynamic range compression
+        mel = torch.log(torch.clamp(mel, min=1e-5))
+        if self.mel_norms is not None:
+            self.mel_norms = self.mel_norms.to(mel.device)
+            mel = mel / self.mel_norms.unsqueeze(0).unsqueeze(-1)
+        return mel
+
+
+class CheckpointedLayer(nn.Module):
+    """
+    Wraps a module. When forward() is called, passes kwargs that require_grad through torch.checkpoint() and bypasses
+    checkpoint for all other args.
+    """
+
+    def __init__(self, wrap):
+        super().__init__()
+        self.wrap = wrap
+
+    def forward(self, x, *args, **kwargs):
+        for k, v in kwargs.items():
+            assert not (
+                isinstance(v, torch.Tensor) and v.requires_grad
+            )  # This would screw up checkpointing.
+        partial = functools.partial(self.wrap, **kwargs)
+        return partial(x, *args)
+
+
+class CheckpointedXTransformerEncoder(nn.Module):
+    """
+    Wraps a ContinuousTransformerWrapper and applies CheckpointedLayer to each layer and permutes from channels-mid
+    to channels-last that XTransformer expects.
+    """
+
+    def __init__(
+        self,
+        needs_permute=True,
+        exit_permute=True,
+        checkpoint=True,
+        **xtransformer_kwargs,
+    ):
+        super().__init__()
+        self.transformer = ContinuousTransformerWrapper(**xtransformer_kwargs)
+        self.needs_permute = needs_permute
+        self.exit_permute = exit_permute
+
+        if not checkpoint:
+            return
+        for i in range(len(self.transformer.attn_layers.layers)):
+            n, b, r = self.transformer.attn_layers.layers[i]
+            self.transformer.attn_layers.layers[i] = nn.ModuleList(
+                [n, CheckpointedLayer(b), r]
+            )
+
+    def forward(self, x, **kwargs):
+        if self.needs_permute:
+            x = x.permute(0, 2, 1)
+        h = self.transformer(x, **kwargs)
+        if self.exit_permute:
+            h = h.permute(0, 2, 1)
+        return h

tortoise/models/autoregressive.py

@@ -0,0 +1,810 @@
+# AGPL: a notification must be added stating that changes have been made to that file.
+import functools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GPT2Config, GPT2PreTrainedModel, LogitsProcessorList
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+from tortoise.models.arch_util import AttentionBlock
+from tortoise.utils.typical_sampling import TypicalLogitsWarper
+
+
+def null_position_embeddings(range, dim):
+    return torch.zeros((range.shape[0], range.shape[1], dim), device=range.device)
+
+
+def _p(t):
+    return t and (len(t), len(t[0]), t[0][0].shape)  # kv_cache debug
+
+
+class ResBlock(nn.Module):
+    """
+    Basic residual convolutional block that uses GroupNorm.
+    """
+
+    def __init__(self, chan):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+            nn.GroupNorm(chan // 8, chan),
+            nn.ReLU(),
+            nn.Conv1d(chan, chan, kernel_size=3, padding=1),
+            nn.GroupNorm(chan // 8, chan),
+        )
+
+    def forward(self, x):
+        return F.relu(self.net(x) + x)
+
+
+class GPT2InferenceModel(GPT2PreTrainedModel):
+    def __init__(self, config, gpt, text_pos_emb, embeddings, norm, linear, kv_cache):
+        super().__init__(config)
+        self.transformer = gpt
+        self.text_pos_embedding = text_pos_emb
+        self.embeddings = embeddings
+        self.lm_head = nn.Sequential(norm, linear)
+        self.kv_cache = kv_cache
+
+    def store_mel_emb(self, mel_emb):
+        self.cached_mel_emb = mel_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)  # usually None
+        if not self.kv_cache:
+            past_key_values = None
+        # only last token for inputs_ids if past is defined in kwargs
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        past_key_values=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        assert self.cached_mel_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # Create embedding
+        mel_len = self.cached_mel_emb.shape[1]
+        if input_ids.shape[1] != 1:
+            text_inputs = input_ids[:, mel_len:]
+            text_emb = self.embeddings(text_inputs)
+            text_emb = text_emb + self.text_pos_embedding(text_emb)
+            if self.cached_mel_emb.shape[0] != text_emb.shape[0]:
+                mel_emb = self.cached_mel_emb.repeat_interleave(
+                    text_emb.shape[0] // self.cached_mel_emb.shape[0], 0
+                )
+            else:  # this outcome only occurs once per loop in most cases
+                mel_emb = self.cached_mel_emb
+            emb = torch.cat([mel_emb, text_emb], dim=1)
+        else:
+            emb = self.embeddings(input_ids)
+            emb = emb + self.text_pos_embedding.get_fixed_embedding(
+                attention_mask.shape[1] - mel_len, attention_mask.device
+            )
+
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+class ConditioningEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        attn_blocks=6,
+        num_attn_heads=4,
+        do_checkpointing=False,
+        mean=False,
+    ):
+        super().__init__()
+        attn = []
+        self.init = nn.Conv1d(spec_dim, embedding_dim, kernel_size=1)
+        for a in range(attn_blocks):
+            attn.append(AttentionBlock(embedding_dim, num_attn_heads))
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+        self.do_checkpointing = do_checkpointing
+        self.mean = mean
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.attn(h)
+        if self.mean:
+            return h.mean(dim=2)
+        else:
+            return h[:, :, 0]
+
+
+class LearnedPositionEmbeddings(nn.Module):
+    def __init__(self, seq_len, model_dim, init=0.02):
+        super().__init__()
+        self.emb = nn.Embedding(seq_len, model_dim)
+        # Initializing this way is standard for GPT-2
+        self.emb.weight.data.normal_(mean=0.0, std=init)
+
+    def forward(self, x):
+        sl = x.shape[1]
+        return self.emb(torch.arange(0, sl, device=x.device))
+
+    def get_fixed_embedding(self, ind, dev):
+        return self.emb(torch.arange(0, ind, device=dev))[ind - 1 : ind]
+
+
+def build_hf_gpt_transformer(
+    layers, model_dim, heads, max_mel_seq_len, max_text_seq_len, checkpointing
+):
+    """
+    GPT-2 implemented by the HuggingFace library.
+    """
+    from transformers import GPT2Config, GPT2Model
+
+    gpt_config = GPT2Config(
+        vocab_size=256,  # Unused.
+        n_positions=max_mel_seq_len + max_text_seq_len,
+        n_ctx=max_mel_seq_len + max_text_seq_len,
+        n_embd=model_dim,
+        n_layer=layers,
+        n_head=heads,
+        gradient_checkpointing=checkpointing,
+        use_cache=not checkpointing,
+    )
+    gpt = GPT2Model(gpt_config)
+    # Override the built in positional embeddings
+    del (
+        gpt.wpe
+    )  # TODO: figure out relevance in fixing exported model definition: Embedding(1012, 1024)
+    gpt.wpe = functools.partial(null_position_embeddings, dim=model_dim)
+    # Built-in token embeddings are unused.
+    del gpt.wte
+    return (
+        gpt,
+        LearnedPositionEmbeddings(max_mel_seq_len, model_dim),
+        LearnedPositionEmbeddings(max_text_seq_len, model_dim),
+        None,
+        None,
+    )
+
+
+class MelEncoder(nn.Module):
+    def __init__(self, channels, mel_channels=80, resblocks_per_reduction=2):
+        super().__init__()
+        self.channels = channels
+        self.encoder = nn.Sequential(
+            nn.Conv1d(mel_channels, channels // 4, kernel_size=3, padding=1),
+            nn.Sequential(
+                *[ResBlock(channels // 4) for _ in range(resblocks_per_reduction)]
+            ),
+            nn.Conv1d(channels // 4, channels // 2, kernel_size=3, stride=2, padding=1),
+            nn.GroupNorm(channels // 16, channels // 2),
+            nn.ReLU(),
+            nn.Sequential(
+                *[ResBlock(channels // 2) for _ in range(resblocks_per_reduction)]
+            ),
+            nn.Conv1d(channels // 2, channels, kernel_size=3, stride=2, padding=1),
+            nn.GroupNorm(channels // 8, channels),
+            nn.ReLU(),
+            nn.Sequential(
+                *[ResBlock(channels) for _ in range(resblocks_per_reduction)]
+            ),
+        )
+        self.reduction = 4
+
+    def forward(self, x):
+        for e in self.encoder:
+            x = e(x)
+        return x.permute(0, 2, 1)
+
+
+class UnifiedVoice(nn.Module):
+    def __init__(
+        self,
+        layers=8,
+        model_dim=512,
+        heads=8,
+        max_text_tokens=120,
+        max_mel_tokens=250,
+        max_conditioning_inputs=1,
+        mel_length_compression=1024,
+        number_text_tokens=256,
+        start_text_token=None,
+        number_mel_codes=8194,
+        start_mel_token=8192,
+        stop_mel_token=8193,
+        train_solo_embeddings=False,
+        use_mel_codes_as_input=True,
+        checkpointing=True,
+        types=1,
+    ):
+        """
+        Args:
+            layers: Number of layers in transformer stack.
+            model_dim: Operating dimensions of the transformer
+            heads: Number of transformer heads. Must be divisible by model_dim. Recommend model_dim//64
+            max_text_tokens: Maximum number of text tokens that will be encountered by model.
+            max_mel_tokens: Maximum number of MEL tokens that will be encountered by model.
+            max_conditioning_inputs: Maximum number of conditioning inputs provided to the model. If (1), conditioning input can be of format (b,80,s), otherwise (b,n,80,s).
+            mel_length_compression: The factor between <number_input_samples> and <mel_tokens>. Used to compute MEL code padding given wav input length.
+            number_text_tokens:
+            start_text_token:
+            stop_text_token:
+            number_mel_codes:
+            start_mel_token:
+            stop_mel_token:
+            train_solo_embeddings:
+            use_mel_codes_as_input:
+            checkpointing:
+        """
+        super().__init__()
+
+        self.number_text_tokens = number_text_tokens
+        self.start_text_token = (
+            number_text_tokens * types if start_text_token is None else start_text_token
+        )
+        self.stop_text_token = 0
+        self.number_mel_codes = number_mel_codes
+        self.start_mel_token = start_mel_token
+        self.stop_mel_token = stop_mel_token
+        self.layers = layers
+        self.heads = heads
+        self.max_mel_tokens = max_mel_tokens
+        self.max_text_tokens = max_text_tokens
+        self.model_dim = model_dim
+        self.max_conditioning_inputs = max_conditioning_inputs
+        self.mel_length_compression = mel_length_compression
+        self.conditioning_encoder = ConditioningEncoder(
+            80, model_dim, num_attn_heads=heads
+        )
+        self.text_embedding = nn.Embedding(
+            self.number_text_tokens * types + 1, model_dim
+        )
+        if use_mel_codes_as_input:
+            self.mel_embedding = nn.Embedding(self.number_mel_codes, model_dim)
+        else:
+            self.mel_embedding = MelEncoder(model_dim, resblocks_per_reduction=1)
+        (
+            self.gpt,
+            self.mel_pos_embedding,
+            self.text_pos_embedding,
+            self.mel_layer_pos_embedding,
+            self.text_layer_pos_embedding,
+        ) = build_hf_gpt_transformer(
+            layers,
+            model_dim,
+            heads,
+            self.max_mel_tokens + 2 + self.max_conditioning_inputs,
+            self.max_text_tokens + 2,
+            checkpointing,
+        )
+        if train_solo_embeddings:
+            self.mel_solo_embedding = nn.Parameter(
+                torch.randn(1, 1, model_dim) * 0.02, requires_grad=True
+            )
+            self.text_solo_embedding = nn.Parameter(
+                torch.randn(1, 1, model_dim) * 0.02, requires_grad=True
+            )
+        else:
+            self.mel_solo_embedding = 0
+            self.text_solo_embedding = 0
+
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.text_head = nn.Linear(model_dim, self.number_text_tokens * types + 1)
+        self.mel_head = nn.Linear(model_dim, self.number_mel_codes)
+
+        # Initialize the embeddings per the GPT-2 scheme
+        embeddings = [self.text_embedding]
+        if use_mel_codes_as_input:
+            embeddings.append(self.mel_embedding)
+        for module in embeddings:
+            module.weight.data.normal_(mean=0.0, std=0.02)
+
+    def post_init_gpt2_config(self, kv_cache=True):
+        seq_length = self.max_mel_tokens + self.max_text_tokens + 2
+        gpt_config = GPT2Config(
+            vocab_size=self.max_mel_tokens,
+            n_positions=seq_length,
+            n_ctx=seq_length,
+            n_embd=self.model_dim,
+            n_layer=self.layers,
+            n_head=self.heads,
+            gradient_checkpointing=False,
+            use_cache=True,
+        )
+        self.inference_model = GPT2InferenceModel(
+            gpt_config,
+            self.gpt,
+            self.mel_pos_embedding,
+            self.mel_embedding,
+            self.final_norm,
+            self.mel_head,
+            kv_cache=kv_cache,
+        )
+        # self.inference_model = PrunedGPT2InferenceModel(gpt_config, self.gpt, self.mel_pos_embedding, self.mel_embedding, self.final_norm, self.mel_head)
+        self.gpt.wte = self.mel_embedding
+
+    def build_aligned_inputs_and_targets(self, input, start_token, stop_token):
+        inp = F.pad(input, (1, 0), value=start_token)
+        tar = F.pad(input, (0, 1), value=stop_token)
+        return inp, tar
+
+    def set_mel_padding(self, mel_input_tokens, wav_lengths):
+        """
+        Given mel tokens that are derived from a padded audio clip and the actual lengths of each batch element in
+        that audio clip, reformats the tokens with STOP_MEL_TOKEN in place of the zero padding. This is required
+        preformatting to create a working TTS model.
+        """
+        # Set padding areas within MEL (currently it is coded with the MEL code for <zero>).
+        mel_lengths = torch.div(
+            wav_lengths, self.mel_length_compression, rounding_mode="trunc"
+        )
+        for b in range(len(mel_lengths)):
+            actual_end = (
+                mel_lengths[b] + 1
+            )  # Due to the convolutional nature of how these tokens are generated, it would be best if the model predicts a token past the actual last token.
+            if actual_end < mel_input_tokens.shape[-1]:
+                mel_input_tokens[b, actual_end:] = self.stop_mel_token
+        return mel_input_tokens
+
+    def get_logits(
+        self,
+        speech_conditioning_inputs,
+        first_inputs,
+        first_head,
+        second_inputs=None,
+        second_head=None,
+        get_attns=False,
+        return_latent=False,
+    ):
+        if second_inputs is not None:
+            emb = torch.cat(
+                [speech_conditioning_inputs, first_inputs, second_inputs], dim=1
+            )
+        else:
+            emb = torch.cat([speech_conditioning_inputs, first_inputs], dim=1)
+
+        gpt_out = self.gpt(
+            inputs_embeds=emb, return_dict=True, output_attentions=get_attns
+        )
+        if get_attns:
+            return gpt_out.attentions
+
+        enc = gpt_out.last_hidden_state[
+            :, 1:
+        ]  # The first logit is tied to the speech_conditioning_input
+        enc = self.final_norm(enc)
+
+        if return_latent:
+            return (
+                enc[
+                    :,
+                    speech_conditioning_inputs.shape[
+                        1
+                    ] : speech_conditioning_inputs.shape[1]
+                    + first_inputs.shape[1],
+                ],
+                enc[:, -second_inputs.shape[1] :],
+            )
+
+        first_logits = enc[:, : first_inputs.shape[1]]
+        first_logits = first_head(first_logits)
+        first_logits = first_logits.permute(0, 2, 1)
+        if second_inputs is not None:
+            second_logits = enc[:, -second_inputs.shape[1] :]
+            second_logits = second_head(second_logits)
+            second_logits = second_logits.permute(0, 2, 1)
+            return first_logits, second_logits
+        else:
+            return first_logits
+
+    def get_conditioning(self, speech_conditioning_input):
+        speech_conditioning_input = (
+            speech_conditioning_input.unsqueeze(1)
+            if len(speech_conditioning_input.shape) == 3
+            else speech_conditioning_input
+        )
+        conds = []
+        for j in range(speech_conditioning_input.shape[1]):
+            conds.append(self.conditioning_encoder(speech_conditioning_input[:, j]))
+        conds = torch.stack(conds, dim=1)
+        conds = conds.mean(dim=1)
+        return conds
+
+    def forward(
+        self,
+        speech_conditioning_latent,
+        text_inputs,
+        text_lengths,
+        mel_codes,
+        wav_lengths,
+        types=None,
+        text_first=True,
+        raw_mels=None,
+        return_attentions=False,
+        return_latent=False,
+        clip_inputs=True,
+    ):
+        """
+        Forward pass that uses both text and voice in either text conditioning mode or voice conditioning mode
+        (actuated by `text_first`).
+
+        speech_conditioning_input: MEL float tensor, (b,1024)
+        text_inputs: long tensor, (b,t)
+        text_lengths: long tensor, (b,)
+        mel_inputs:  long tensor, (b,m)
+        wav_lengths: long tensor, (b,)
+        raw_mels: MEL float tensor (b,80,s)
+
+        If return_attentions is specified, only logits are returned.
+        If return_latent is specified, loss & logits are not computed or returned. Only the predicted latents are returned.
+        If clip_inputs is True, the inputs will be clipped to the smallest input size across each input modality.
+        """
+        # Types are expressed by expanding the text embedding space.
+        if types is not None:
+            text_inputs = text_inputs * (1 + types).unsqueeze(-1)
+
+        if clip_inputs:
+            # This model will receive micro-batches with a ton of padding for both the text and MELs. Ameliorate this by
+            # chopping the inputs by the maximum actual length.
+            max_text_len = text_lengths.max()
+            text_inputs = text_inputs[:, :max_text_len]
+            max_mel_len = wav_lengths.max() // self.mel_length_compression
+            mel_codes = mel_codes[:, :max_mel_len]
+            if raw_mels is not None:
+                raw_mels = raw_mels[:, :, : max_mel_len * 4]
+        mel_codes = self.set_mel_padding(mel_codes, wav_lengths)
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        mel_codes = F.pad(mel_codes, (0, 1), value=self.stop_mel_token)
+
+        conds = speech_conditioning_latent.unsqueeze(1)
+        text_inputs, text_targets = self.build_aligned_inputs_and_targets(
+            text_inputs, self.start_text_token, self.stop_text_token
+        )
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(
+            text_inputs
+        )
+        mel_codes, mel_targets = self.build_aligned_inputs_and_targets(
+            mel_codes, self.start_mel_token, self.stop_mel_token
+        )
+        if raw_mels is not None:
+            mel_inp = F.pad(raw_mels, (0, 8))
+        else:
+            mel_inp = mel_codes
+        mel_emb = self.mel_embedding(mel_inp)
+        mel_emb = mel_emb + self.mel_pos_embedding(mel_codes)
+
+        if text_first:
+            text_logits, mel_logits = self.get_logits(
+                conds,
+                text_emb,
+                self.text_head,
+                mel_emb,
+                self.mel_head,
+                get_attns=return_attentions,
+                return_latent=return_latent,
+            )
+            if return_latent:
+                return mel_logits[
+                    :, :-2
+                ]  # Despite the name, these are not logits. Strip off the two tokens added by this forward pass.
+        else:
+            mel_logits, text_logits = self.get_logits(
+                conds,
+                mel_emb,
+                self.mel_head,
+                text_emb,
+                self.text_head,
+                get_attns=return_attentions,
+                return_latent=return_latent,
+            )
+            if return_latent:
+                return text_logits[
+                    :, :-2
+                ]  # Despite the name, these are not logits. Strip off the two tokens added by this forward pass.
+
+        if return_attentions:
+            return mel_logits
+        loss_text = F.cross_entropy(text_logits, text_targets.long())
+        loss_mel = F.cross_entropy(mel_logits, mel_targets.long())
+        return loss_text.mean(), loss_mel.mean(), mel_logits
+
+    def inference_speech(
+        self,
+        speech_conditioning_latent,
+        text_inputs,
+        input_tokens=None,
+        num_return_sequences=1,
+        max_generate_length=None,
+        typical_sampling=False,
+        typical_mass=0.9,
+        **hf_generate_kwargs
+    ):
+        text_inputs = F.pad(text_inputs, (0, 1), value=self.stop_text_token)
+        text_inputs, text_targets = self.build_aligned_inputs_and_targets(
+            text_inputs, self.start_text_token, self.stop_text_token
+        )
+        text_emb = self.text_embedding(text_inputs) + self.text_pos_embedding(
+            text_inputs
+        )
+
+        conds = speech_conditioning_latent.unsqueeze(1)
+        emb = torch.cat([conds, text_emb], dim=1)
+        self.inference_model.store_mel_emb(emb)
+
+        fake_inputs = torch.full(
+            (
+                emb.shape[0],
+                conds.shape[1] + emb.shape[1],
+            ),
+            fill_value=1,
+            dtype=torch.long,
+            device=text_inputs.device,
+        )
+        fake_inputs[:, -1] = self.start_mel_token
+        trunc_index = fake_inputs.shape[1]
+        if input_tokens is None:
+            inputs = fake_inputs
+        else:
+            assert (
+                num_return_sequences % input_tokens.shape[0] == 0
+            ), "The number of return sequences must be divisible by the number of input sequences"
+            fake_inputs = fake_inputs.repeat(num_return_sequences, 1)
+            input_tokens = input_tokens.repeat(
+                num_return_sequences // input_tokens.shape[0], 1
+            )
+            inputs = torch.cat([fake_inputs, input_tokens], dim=1)
+
+        logits_processor = (
+            LogitsProcessorList([TypicalLogitsWarper(mass=typical_mass)])
+            if typical_sampling
+            else LogitsProcessorList()
+        )  # TODO disable this
+        max_length = (
+            trunc_index + self.max_mel_tokens - 1
+            if max_generate_length is None
+            else trunc_index + max_generate_length
+        )
+        gen = self.inference_model.generate(
+            inputs,
+            bos_token_id=self.start_mel_token,
+            pad_token_id=self.stop_mel_token,
+            eos_token_id=self.stop_mel_token,
+            max_length=max_length,
+            logits_processor=logits_processor,
+            num_return_sequences=num_return_sequences,
+            **hf_generate_kwargs
+        )
+        return gen[:, trunc_index:]
+
+
+class PrunedGPT2InferenceModel(GPT2PreTrainedModel):
+    def __init__(self, config, gpt, text_pos_emb, embeddings, norm, linear):
+        super().__init__(config)
+        self.transformer = gpt
+        self.text_pos_embedding = text_pos_emb
+        self.embeddings = embeddings
+        self.lm_head = nn.Sequential(norm, linear)
+
+    def store_mel_emb(self, mel_emb):
+        self.cached_mel_emb = mel_emb
+
+    def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs):
+        token_type_ids = kwargs.get("token_type_ids", None)
+        # only last token for inputs_ids if past is defined in kwargs
+        print(past)
+        if past:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+            if token_type_ids is not None:
+                token_type_ids = token_type_ids[:, -1].unsqueeze(-1)
+
+        attention_mask = kwargs.get("attention_mask", None)
+        position_ids = kwargs.get("position_ids", None)
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            print(position_ids)
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            print(position_ids)
+            if past:
+                position_ids = position_ids[:, -1].unsqueeze(-1)
+        else:
+            position_ids = None
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past,
+            "use_cache": kwargs.get("use_cache"),
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "token_type_ids": token_type_ids,
+        }
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, **kwargs):
+        past_key_values = None
+        token_type_ids = None
+        head_mask = None
+        inputs_embeds = None
+        encoder_hidden_states = None
+        encoder_attention_mask = None
+        labels = None
+        use_cache = True
+        output_attentions = False
+        output_hidden_states = False
+        return_dict = True
+        #
+        assert self.cached_mel_emb is not None
+        assert inputs_embeds is None  # Not supported by this inference model.
+        assert labels is None  # Training not supported by this inference model.
+        # return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        """
+        print(attention_mask)
+        print(position_ids)
+        print(attention_mask.dtype)
+        print(position_ids.dtype)
+        """
+
+        """
+        attention_mask=tensor([[1, 1, 1,  ..., 1, 1, 1],
+            [1, 1, 1,  ..., 1, 1, 1],
+            [1, 1, 1,  ..., 1, 1, 1],
+            ...,
+            [1, 1, 1,  ..., 1, 1, 1],
+            [1, 1, 1,  ..., 1, 1, 1],
+            [1, 1, 1,  ..., 1, 1, 1]], device='cuda:0')
+        """
+
+        # Create embedding
+        mel_len = self.cached_mel_emb.shape[1]
+        text_inputs = input_ids[:, mel_len:]
+        text_emb = self.embeddings(text_inputs)
+        text_emb = text_emb + self.text_pos_embedding(text_emb)
+        mel_emb = self.cached_mel_emb.repeat_interleave(
+            text_emb.shape[0] // self.cached_mel_emb.shape[0], 0
+        )
+        emb = torch.cat([mel_emb, text_emb], dim=1)
+
+        transformer_outputs = self.transformer(
+            inputs_embeds=emb,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+
+        lm_logits = self.lm_head(hidden_states)
+
+        if not return_dict:
+            return (lm_logits,) + transformer_outputs[1:]
+        return CausalLMOutputWithCrossAttentions(
+            loss=None,
+            logits=lm_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+            cross_attentions=transformer_outputs.cross_attentions,
+        )
+
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """
+        This function is used to re-order the :obj:`past_key_values` cache if
+        :meth:`~transformers.PreTrainedModel.beam_search` or :meth:`~transformers.PreTrainedModel.beam_sample` is
+        called. This is required to match :obj:`past_key_values` with the correct beam_idx at every generation step.
+        """
+        return tuple(
+            tuple(
+                past_state.index_select(0, beam_idx.to(past_state.device))
+                for past_state in layer_past
+            )
+            for layer_past in past
+        )
+
+
+if __name__ == "__main__":
+    gpt = UnifiedVoice(
+        model_dim=256,
+        heads=4,
+        train_solo_embeddings=True,
+        use_mel_codes_as_input=True,
+        max_conditioning_inputs=4,
+    )
+    l = gpt(
+        torch.randn(2, 3, 80, 800),
+        torch.randint(high=120, size=(2, 120)),
+        torch.tensor([32, 120]),
+        torch.randint(high=8192, size=(2, 250)),
+        torch.tensor([250 * 256, 195 * 256]),
+    )
+    gpt.text_forward(
+        torch.randn(2, 80, 800),
+        torch.randint(high=50, size=(2, 80)),
+        torch.tensor([32, 80]),
+    )

tortoise/models/classifier.py

@@ -0,0 +1,166 @@
+import torch
+import torch.nn as nn
+
+from tortoise.models.arch_util import (
+    AttentionBlock,
+    Downsample,
+    Upsample,
+    normalization,
+    zero_module,
+)
+
+
+class ResBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        dropout,
+        out_channels=None,
+        use_conv=False,
+        use_scale_shift_norm=False,
+        dims=2,
+        up=False,
+        down=False,
+        kernel_size=3,
+        do_checkpoint=True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_scale_shift_norm = use_scale_shift_norm
+        self.do_checkpoint = do_checkpoint
+        padding = 1 if kernel_size == 3 else 2
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, kernel_size, padding=padding),
+        )
+
+        self.updown = up or down
+
+        if up:
+            self.h_upd = Upsample(channels, False, dims)
+            self.x_upd = Upsample(channels, False, dims)
+        elif down:
+            self.h_upd = Downsample(channels, False, dims)
+            self.x_upd = Downsample(channels, False, dims)
+        else:
+            self.h_upd = self.x_upd = nn.Identity()
+
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(
+                nn.Conv1d(
+                    self.out_channels, self.out_channels, kernel_size, padding=padding
+                )
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        elif use_conv:
+            self.skip_connection = nn.Conv1d(
+                dims, channels, self.out_channels, kernel_size, padding=padding
+            )
+        else:
+            self.skip_connection = nn.Conv1d(dims, channels, self.out_channels, 1)
+
+    def forward(self, x):
+        if self.updown:
+            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+            h = in_rest(x)
+            h = self.h_upd(h)
+            x = self.x_upd(x)
+            h = in_conv(h)
+        else:
+            h = self.in_layers(x)
+        h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class AudioMiniEncoder(nn.Module):
+    def __init__(
+        self,
+        spec_dim,
+        embedding_dim,
+        base_channels=128,
+        depth=2,
+        resnet_blocks=2,
+        attn_blocks=4,
+        num_attn_heads=4,
+        dropout=0,
+        downsample_factor=2,
+        kernel_size=3,
+    ):
+        super().__init__()
+        self.init = nn.Sequential(nn.Conv1d(spec_dim, base_channels, 3, padding=1))
+        ch = base_channels
+        res = []
+        self.layers = depth
+        for l in range(depth):
+            for r in range(resnet_blocks):
+                res.append(
+                    ResBlock(ch, dropout, do_checkpoint=False, kernel_size=kernel_size)
+                )
+            res.append(
+                Downsample(
+                    ch, use_conv=True, out_channels=ch * 2, factor=downsample_factor
+                )
+            )
+            ch *= 2
+        self.res = nn.Sequential(*res)
+        self.final = nn.Sequential(
+            normalization(ch), nn.SiLU(), nn.Conv1d(ch, embedding_dim, 1)
+        )
+        attn = []
+        for a in range(attn_blocks):
+            attn.append(
+                AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=False)
+            )
+        self.attn = nn.Sequential(*attn)
+        self.dim = embedding_dim
+
+    def forward(self, x):
+        h = self.init(x)
+        h = self.res(h)
+        h = self.final(h)
+        for blk in self.attn:
+            h = blk(h)
+        return h[:, :, 0]
+
+
+class AudioMiniEncoderWithClassifierHead(nn.Module):
+    def __init__(self, classes, distribute_zero_label=True, **kwargs):
+        super().__init__()
+        self.enc = AudioMiniEncoder(**kwargs)
+        self.head = nn.Linear(self.enc.dim, classes)
+        self.num_classes = classes
+        self.distribute_zero_label = distribute_zero_label
+
+    def forward(self, x, labels=None):
+        h = self.enc(x)
+        logits = self.head(h)
+        if labels is None:
+            return logits
+        else:
+            if self.distribute_zero_label:
+                oh_labels = nn.functional.one_hot(labels, num_classes=self.num_classes)
+                zeros_indices = (labels == 0).unsqueeze(-1)
+                # Distribute 20% of the probability mass on all classes when zero is specified, to compensate for dataset noise.
+                zero_extra_mass = torch.full_like(
+                    oh_labels,
+                    dtype=torch.float,
+                    fill_value=0.2 / (self.num_classes - 1),
+                )
+                zero_extra_mass[:, 0] = -0.2
+                zero_extra_mass = zero_extra_mass * zeros_indices
+                oh_labels = oh_labels + zero_extra_mass
+            else:
+                oh_labels = labels
+            loss = nn.functional.cross_entropy(logits, oh_labels)
+            return loss

tortoise/models/clvp.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+
+from tortoise.models.arch_util import CheckpointedXTransformerEncoder
+from tortoise.models.transformer import Transformer
+from tortoise.models.xtransformers import Encoder
+
+
+def exists(val):
+    return val is not None
+
+
+def masked_mean(t, mask, dim=1):
+    t = t.masked_fill(~mask[:, :, None], 0.0)
+    return t.sum(dim=1) / mask.sum(dim=1)[..., None]
+
+
+class CLVP(nn.Module):
+    """
+    CLIP model retrofitted for performing contrastive evaluation between tokenized audio data and the corresponding
+    transcribed text.
+
+    Originally from https://github.com/lucidrains/DALLE-pytorch/blob/main/dalle_pytorch/dalle_pytorch.py
+    """
+
+    def __init__(
+        self,
+        *,
+        dim_text=512,
+        dim_speech=512,
+        dim_latent=512,
+        num_text_tokens=256,
+        text_enc_depth=6,
+        text_seq_len=120,
+        text_heads=8,
+        num_speech_tokens=8192,
+        speech_enc_depth=6,
+        speech_heads=8,
+        speech_seq_len=250,
+        text_mask_percentage=0,
+        voice_mask_percentage=0,
+        wav_token_compression=1024,
+        use_xformers=False,
+    ):
+        super().__init__()
+        self.text_emb = nn.Embedding(num_text_tokens, dim_text)
+        self.to_text_latent = nn.Linear(dim_text, dim_latent, bias=False)
+
+        self.speech_emb = nn.Embedding(num_speech_tokens, dim_speech)
+        self.to_speech_latent = nn.Linear(dim_speech, dim_latent, bias=False)
+
+        if use_xformers:
+            self.text_transformer = CheckpointedXTransformerEncoder(
+                needs_permute=False,
+                exit_permute=False,
+                max_seq_len=-1,
+                attn_layers=Encoder(
+                    dim=dim_text,
+                    depth=text_enc_depth,
+                    heads=text_heads,
+                    ff_dropout=0.1,
+                    ff_mult=2,
+                    attn_dropout=0.1,
+                    use_rmsnorm=True,
+                    ff_glu=True,
+                    rotary_pos_emb=True,
+                ),
+            )
+            self.speech_transformer = CheckpointedXTransformerEncoder(
+                needs_permute=False,
+                exit_permute=False,
+                max_seq_len=-1,
+                attn_layers=Encoder(
+                    dim=dim_speech,
+                    depth=speech_enc_depth,
+                    heads=speech_heads,
+                    ff_dropout=0.1,
+                    ff_mult=2,
+                    attn_dropout=0.1,
+                    use_rmsnorm=True,
+                    ff_glu=True,
+                    rotary_pos_emb=True,
+                ),
+            )
+        else:
+            self.text_transformer = Transformer(
+                causal=False,
+                seq_len=text_seq_len,
+                dim=dim_text,
+                depth=text_enc_depth,
+                heads=text_heads,
+            )
+            self.speech_transformer = Transformer(
+                causal=False,
+                seq_len=speech_seq_len,
+                dim=dim_speech,
+                depth=speech_enc_depth,
+                heads=speech_heads,
+            )
+
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+        self.text_mask_percentage = text_mask_percentage
+        self.voice_mask_percentage = voice_mask_percentage
+        self.wav_token_compression = wav_token_compression
+        self.xformers = use_xformers
+        if not use_xformers:
+            self.text_pos_emb = nn.Embedding(text_seq_len, dim_text)
+            self.speech_pos_emb = nn.Embedding(num_speech_tokens, dim_speech)
+
+    def forward(self, text, speech_tokens, return_loss=False):
+        b, device = text.shape[0], text.device
+        if self.training:
+            text_mask = torch.rand_like(text.float()) > self.text_mask_percentage
+            voice_mask = (
+                torch.rand_like(speech_tokens.float()) > self.voice_mask_percentage
+            )
+        else:
+            text_mask = torch.ones_like(text.float()).bool()
+            voice_mask = torch.ones_like(speech_tokens.float()).bool()
+
+        text_emb = self.text_emb(text)
+        speech_emb = self.speech_emb(speech_tokens)
+
+        if not self.xformers:
+            text_emb += self.text_pos_emb(torch.arange(text.shape[1], device=device))
+            speech_emb += self.speech_pos_emb(
+                torch.arange(speech_emb.shape[1], device=device)
+            )
+
+        enc_text = self.text_transformer(text_emb, mask=text_mask)
+        enc_speech = self.speech_transformer(speech_emb, mask=voice_mask)
+
+        text_latents = masked_mean(enc_text, text_mask, dim=1)
+        speech_latents = masked_mean(enc_speech, voice_mask, dim=1)
+
+        text_latents = self.to_text_latent(text_latents)
+        speech_latents = self.to_speech_latent(speech_latents)
+
+        text_latents, speech_latents = map(
+            lambda t: F.normalize(t, p=2, dim=-1), (text_latents, speech_latents)
+        )
+
+        temp = self.temperature.exp()
+
+        if not return_loss:
+            sim = einsum("n d, n d -> n", text_latents, speech_latents) * temp
+            return sim
+
+        sim = einsum("i d, j d -> i j", text_latents, speech_latents) * temp
+        labels = torch.arange(b, device=device)
+        loss = (F.cross_entropy(sim, labels) + F.cross_entropy(sim.t(), labels)) / 2
+        return loss
+
+
+if __name__ == "__main__":
+    clip = CLVP(text_mask_percentage=0.2, voice_mask_percentage=0.2)
+    clip(
+        torch.randint(0, 256, (2, 120)),
+        torch.tensor([50, 100]),
+        torch.randint(0, 8192, (2, 250)),
+        torch.tensor([101, 102]),
+        return_loss=True,
+    )
+    nonloss = clip(
+        torch.randint(0, 256, (2, 120)),
+        torch.tensor([50, 100]),
+        torch.randint(0, 8192, (2, 250)),
+        torch.tensor([101, 102]),
+        return_loss=False,
+    )
+    print(nonloss.shape)

tortoise/models/cvvp.py

@@ -0,0 +1,156 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import einsum
+
+from tortoise.models.arch_util import AttentionBlock
+from tortoise.models.xtransformers import ContinuousTransformerWrapper, Encoder
+
+
+def exists(val):
+    return val is not None
+
+
+def masked_mean(t, mask):
+    t = t.masked_fill(~mask, 0.0)
+    return t.sum(dim=1) / mask.sum(dim=1)
+
+
+class CollapsingTransformer(nn.Module):
+    def __init__(
+        self,
+        model_dim,
+        output_dims,
+        heads,
+        dropout,
+        depth,
+        mask_percentage=0,
+        **encoder_kwargs
+    ):
+        super().__init__()
+        self.transformer = ContinuousTransformerWrapper(
+            max_seq_len=-1,
+            use_pos_emb=False,
+            attn_layers=Encoder(
+                dim=model_dim,
+                depth=depth,
+                heads=heads,
+                ff_dropout=dropout,
+                ff_mult=1,
+                attn_dropout=dropout,
+                use_rmsnorm=True,
+                ff_glu=True,
+                rotary_pos_emb=True,
+                **encoder_kwargs,
+            ),
+        )
+        self.pre_combiner = nn.Sequential(
+            nn.Conv1d(model_dim, output_dims, 1),
+            AttentionBlock(output_dims, num_heads=heads, do_checkpoint=False),
+            nn.Conv1d(output_dims, output_dims, 1),
+        )
+        self.mask_percentage = mask_percentage
+
+    def forward(self, x, **transformer_kwargs):
+        h = self.transformer(x, **transformer_kwargs)
+        h = h.permute(0, 2, 1)
+        h = self.pre_combiner(h).permute(0, 2, 1)
+        if self.training:
+            mask = torch.rand_like(h.float()) > self.mask_percentage
+        else:
+            mask = torch.ones_like(h.float()).bool()
+        return masked_mean(h, mask)
+
+
+class ConvFormatEmbedding(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.emb = nn.Embedding(*args, **kwargs)
+
+    def forward(self, x):
+        y = self.emb(x)
+        return y.permute(0, 2, 1)
+
+
+class CVVP(nn.Module):
+    def __init__(
+        self,
+        model_dim=512,
+        transformer_heads=8,
+        dropout=0.1,
+        conditioning_enc_depth=8,
+        cond_mask_percentage=0,
+        mel_channels=80,
+        mel_codes=None,
+        speech_enc_depth=8,
+        speech_mask_percentage=0,
+        latent_multiplier=1,
+    ):
+        super().__init__()
+        latent_dim = latent_multiplier * model_dim
+        self.temperature = nn.Parameter(torch.tensor(1.0))
+
+        self.cond_emb = nn.Sequential(
+            nn.Conv1d(mel_channels, model_dim // 2, kernel_size=5, stride=2, padding=2),
+            nn.Conv1d(model_dim // 2, model_dim, kernel_size=3, stride=2, padding=1),
+        )
+        self.conditioning_transformer = CollapsingTransformer(
+            model_dim,
+            model_dim,
+            transformer_heads,
+            dropout,
+            conditioning_enc_depth,
+            cond_mask_percentage,
+        )
+        self.to_conditioning_latent = nn.Linear(latent_dim, latent_dim, bias=False)
+
+        if mel_codes is None:
+            self.speech_emb = nn.Conv1d(
+                mel_channels, model_dim, kernel_size=5, padding=2
+            )
+        else:
+            self.speech_emb = ConvFormatEmbedding(mel_codes, model_dim)
+        self.speech_transformer = CollapsingTransformer(
+            model_dim,
+            latent_dim,
+            transformer_heads,
+            dropout,
+            speech_enc_depth,
+            speech_mask_percentage,
+        )
+        self.to_speech_latent = nn.Linear(latent_dim, latent_dim, bias=False)
+
+    def get_grad_norm_parameter_groups(self):
+        return {
+            "conditioning": list(self.conditioning_transformer.parameters()),
+            "speech": list(self.speech_transformer.parameters()),
+        }
+
+    def forward(self, mel_cond, mel_input, return_loss=False):
+        cond_emb = self.cond_emb(mel_cond).permute(0, 2, 1)
+        enc_cond = self.conditioning_transformer(cond_emb)
+        cond_latents = self.to_conditioning_latent(enc_cond)
+
+        speech_emb = self.speech_emb(mel_input).permute(0, 2, 1)
+        enc_speech = self.speech_transformer(speech_emb)
+        speech_latents = self.to_speech_latent(enc_speech)
+
+        cond_latents, speech_latents = map(
+            lambda t: F.normalize(t, p=2, dim=-1), (cond_latents, speech_latents)
+        )
+        temp = self.temperature.exp()
+
+        if not return_loss:
+            sim = einsum("n d, n d -> n", cond_latents, speech_latents) * temp
+            return sim
+
+        sim = einsum("i d, j d -> i j", cond_latents, speech_latents) * temp
+        labels = torch.arange(cond_latents.shape[0], device=mel_input.device)
+        loss = (F.cross_entropy(sim, labels) + F.cross_entropy(sim.t(), labels)) / 2
+
+        return loss
+
+
+if __name__ == "__main__":
+    clvp = CVVP()
+    clvp(torch.randn(2, 80, 100), torch.randn(2, 80, 95), return_loss=True)

tortoise/models/diffusion_decoder.py

@@ -0,0 +1,445 @@
+import math
+import random
+from abc import abstractmethod
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import autocast
+
+from tortoise.models.arch_util import AttentionBlock, normalization
+
+
+def is_latent(t):
+    return t.dtype == torch.float
+
+
+def is_sequence(t):
+    return t.dtype == torch.long
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=torch.float32)
+        / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class TimestepBlock(nn.Module):
+    @abstractmethod
+    def forward(self, x, emb):
+        """
+        Apply the module to `x` given `emb` timestep embeddings.
+        """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+    def forward(self, x, emb):
+        for layer in self:
+            if isinstance(layer, TimestepBlock):
+                x = layer(x, emb)
+            else:
+                x = layer(x)
+        return x
+
+
+class ResBlock(TimestepBlock):
+    def __init__(
+        self,
+        channels,
+        emb_channels,
+        dropout,
+        out_channels=None,
+        dims=2,
+        kernel_size=3,
+        efficient_config=True,
+        use_scale_shift_norm=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.emb_channels = emb_channels
+        self.dropout = dropout
+        self.out_channels = out_channels or channels
+        self.use_scale_shift_norm = use_scale_shift_norm
+        padding = {1: 0, 3: 1, 5: 2}[kernel_size]
+        eff_kernel = 1 if efficient_config else 3
+        eff_padding = 0 if efficient_config else 1
+
+        self.in_layers = nn.Sequential(
+            normalization(channels),
+            nn.SiLU(),
+            nn.Conv1d(channels, self.out_channels, eff_kernel, padding=eff_padding),
+        )
+
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(
+                emb_channels,
+                2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+            ),
+        )
+        self.out_layers = nn.Sequential(
+            normalization(self.out_channels),
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            nn.Conv1d(
+                self.out_channels, self.out_channels, kernel_size, padding=padding
+            ),
+        )
+
+        if self.out_channels == channels:
+            self.skip_connection = nn.Identity()
+        else:
+            self.skip_connection = nn.Conv1d(
+                channels, self.out_channels, eff_kernel, padding=eff_padding
+            )
+
+    def forward(self, x, emb):
+        h = self.in_layers(x)
+        emb_out = self.emb_layers(emb).type(h.dtype)
+        while len(emb_out.shape) < len(h.shape):
+            emb_out = emb_out[..., None]
+        if self.use_scale_shift_norm:
+            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+            scale, shift = torch.chunk(emb_out, 2, dim=1)
+            h = out_norm(h) * (1 + scale) + shift
+            h = out_rest(h)
+        else:
+            h = h + emb_out
+            h = self.out_layers(h)
+        return self.skip_connection(x) + h
+
+
+class DiffusionLayer(TimestepBlock):
+    def __init__(self, model_channels, dropout, num_heads):
+        super().__init__()
+        self.resblk = ResBlock(
+            model_channels,
+            model_channels,
+            dropout,
+            model_channels,
+            dims=1,
+            use_scale_shift_norm=True,
+        )
+        self.attn = AttentionBlock(
+            model_channels, num_heads, relative_pos_embeddings=True
+        )
+
+    def forward(self, x, time_emb):
+        y = self.resblk(x, time_emb)
+        return self.attn(y)
+
+
+class DiffusionTts(nn.Module):
+    def __init__(
+        self,
+        model_channels=512,
+        num_layers=8,
+        in_channels=100,
+        in_latent_channels=512,
+        in_tokens=8193,
+        out_channels=200,  # mean and variance
+        dropout=0,
+        use_fp16=False,
+        num_heads=16,
+        # Parameters for regularization.
+        layer_drop=0.1,
+        unconditioned_percentage=0.1,  # This implements a mechanism similar to what is used in classifier-free training.
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.dropout = dropout
+        self.num_heads = num_heads
+        self.unconditioned_percentage = unconditioned_percentage
+        self.enable_fp16 = use_fp16
+        self.layer_drop = layer_drop
+
+        self.inp_block = nn.Conv1d(in_channels, model_channels, 3, 1, 1)
+        self.time_embed = nn.Sequential(
+            nn.Linear(model_channels, model_channels),
+            nn.SiLU(),
+            nn.Linear(model_channels, model_channels),
+        )
+
+        # Either code_converter or latent_converter is used, depending on what type of conditioning data is fed.
+        # This model is meant to be able to be trained on both for efficiency purposes - it is far less computationally
+        # complex to generate tokens, while generating latents will normally mean propagating through a deep autoregressive
+        # transformer network.
+        self.code_embedding = nn.Embedding(in_tokens, model_channels)
+        self.code_converter = nn.Sequential(
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+        )
+        self.code_norm = normalization(model_channels)
+        self.latent_conditioner = nn.Sequential(
+            nn.Conv1d(in_latent_channels, model_channels, 3, padding=1),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+            AttentionBlock(model_channels, num_heads, relative_pos_embeddings=True),
+        )
+        self.contextual_embedder = nn.Sequential(
+            nn.Conv1d(in_channels, model_channels, 3, padding=1, stride=2),
+            nn.Conv1d(model_channels, model_channels * 2, 3, padding=1, stride=2),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+            AttentionBlock(
+                model_channels * 2,
+                num_heads,
+                relative_pos_embeddings=True,
+                do_checkpoint=False,
+            ),
+        )
+        self.unconditioned_embedding = nn.Parameter(torch.randn(1, model_channels, 1))
+        self.conditioning_timestep_integrator = TimestepEmbedSequential(
+            DiffusionLayer(model_channels, dropout, num_heads),
+            DiffusionLayer(model_channels, dropout, num_heads),
+            DiffusionLayer(model_channels, dropout, num_heads),
+        )
+
+        self.integrating_conv = nn.Conv1d(
+            model_channels * 2, model_channels, kernel_size=1
+        )
+        self.mel_head = nn.Conv1d(model_channels, in_channels, kernel_size=3, padding=1)
+
+        self.layers = nn.ModuleList(
+            [
+                DiffusionLayer(model_channels, dropout, num_heads)
+                for _ in range(num_layers)
+            ]
+            + [
+                ResBlock(
+                    model_channels,
+                    model_channels,
+                    dropout,
+                    dims=1,
+                    use_scale_shift_norm=True,
+                )
+                for _ in range(3)
+            ]
+        )
+
+        self.out = nn.Sequential(
+            normalization(model_channels),
+            nn.SiLU(),
+            nn.Conv1d(model_channels, out_channels, 3, padding=1),
+        )
+
+    def get_grad_norm_parameter_groups(self):
+        groups = {
+            "minicoder": list(self.contextual_embedder.parameters()),
+            "layers": list(self.layers.parameters()),
+            "code_converters": list(self.code_embedding.parameters())
+            + list(self.code_converter.parameters())
+            + list(self.latent_conditioner.parameters())
+            + list(self.latent_conditioner.parameters()),
+            "timestep_integrator": list(
+                self.conditioning_timestep_integrator.parameters()
+            )
+            + list(self.integrating_conv.parameters()),
+            "time_embed": list(self.time_embed.parameters()),
+        }
+        return groups
+
+    def get_conditioning(self, conditioning_input):
+        speech_conditioning_input = (
+            conditioning_input.unsqueeze(1)
+            if len(conditioning_input.shape) == 3
+            else conditioning_input
+        )
+        conds = []
+        for j in range(speech_conditioning_input.shape[1]):
+            conds.append(self.contextual_embedder(speech_conditioning_input[:, j]))
+        conds = torch.cat(conds, dim=-1)
+        conds = conds.mean(dim=-1)
+        return conds
+
+    def timestep_independent(
+        self,
+        aligned_conditioning,
+        conditioning_latent,
+        expected_seq_len,
+        return_code_pred,
+    ):
+        # Shuffle aligned_latent to BxCxS format
+        if is_latent(aligned_conditioning):
+            aligned_conditioning = aligned_conditioning.permute(0, 2, 1)
+
+        cond_scale, cond_shift = torch.chunk(conditioning_latent, 2, dim=1)
+        if is_latent(aligned_conditioning):
+            code_emb = self.latent_conditioner(aligned_conditioning)
+        else:
+            code_emb = self.code_embedding(aligned_conditioning).permute(0, 2, 1)
+            code_emb = self.code_converter(code_emb)
+        code_emb = self.code_norm(code_emb) * (
+            1 + cond_scale.unsqueeze(-1)
+        ) + cond_shift.unsqueeze(-1)
+
+        unconditioned_batches = torch.zeros(
+            (code_emb.shape[0], 1, 1), device=code_emb.device
+        )
+        # Mask out the conditioning branch for whole batch elements, implementing something similar to classifier-free guidance.
+        if self.training and self.unconditioned_percentage > 0:
+            unconditioned_batches = (
+                torch.rand((code_emb.shape[0], 1, 1), device=code_emb.device)
+                < self.unconditioned_percentage
+            )
+            code_emb = torch.where(
+                unconditioned_batches,
+                self.unconditioned_embedding.repeat(
+                    aligned_conditioning.shape[0], 1, 1
+                ),
+                code_emb,
+            )
+        expanded_code_emb = F.interpolate(
+            code_emb, size=expected_seq_len, mode="nearest"
+        )
+
+        if not return_code_pred:
+            return expanded_code_emb
+        else:
+            mel_pred = self.mel_head(expanded_code_emb)
+            # Multiply mel_pred by !unconditioned_branches, which drops the gradient on unconditioned branches. This is because we don't want that gradient being used to train parameters through the codes_embedder as it unbalances contributions to that network from the MSE loss.
+            mel_pred = mel_pred * unconditioned_batches.logical_not()
+            return expanded_code_emb, mel_pred
+
+    def forward(
+        self,
+        x,
+        timesteps,
+        aligned_conditioning=None,
+        conditioning_latent=None,
+        precomputed_aligned_embeddings=None,
+        conditioning_free=False,
+        return_code_pred=False,
+    ):
+        """
+        Apply the model to an input batch.
+
+        :param x: an [N x C x ...] Tensor of inputs.
+        :param timesteps: a 1-D batch of timesteps.
+        :param aligned_conditioning: an aligned latent or sequence of tokens providing useful data about the sample to be produced.
+        :param conditioning_latent: a pre-computed conditioning latent; see get_conditioning().
+        :param precomputed_aligned_embeddings: Embeddings returned from self.timestep_independent()
+        :param conditioning_free: When set, all conditioning inputs (including tokens and conditioning_input) will not be considered.
+        :return: an [N x C x ...] Tensor of outputs.
+        """
+        assert precomputed_aligned_embeddings is not None or (
+            aligned_conditioning is not None and conditioning_latent is not None
+        )
+        assert not (
+            return_code_pred and precomputed_aligned_embeddings is not None
+        )  # These two are mutually exclusive.
+
+        unused_params = []
+        if conditioning_free:
+            code_emb = self.unconditioned_embedding.repeat(x.shape[0], 1, x.shape[-1])
+            unused_params.extend(
+                list(self.code_converter.parameters())
+                + list(self.code_embedding.parameters())
+            )
+            unused_params.extend(list(self.latent_conditioner.parameters()))
+        else:
+            if precomputed_aligned_embeddings is not None:
+                code_emb = precomputed_aligned_embeddings
+            else:
+                code_emb, mel_pred = self.timestep_independent(
+                    aligned_conditioning, conditioning_latent, x.shape[-1], True
+                )
+                if is_latent(aligned_conditioning):
+                    unused_params.extend(
+                        list(self.code_converter.parameters())
+                        + list(self.code_embedding.parameters())
+                    )
+                else:
+                    unused_params.extend(list(self.latent_conditioner.parameters()))
+
+            unused_params.append(self.unconditioned_embedding)
+
+        time_emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+        code_emb = self.conditioning_timestep_integrator(code_emb, time_emb)
+        x = self.inp_block(x)
+        x = torch.cat([x, code_emb], dim=1)
+        x = self.integrating_conv(x)
+        for i, lyr in enumerate(self.layers):
+            # Do layer drop where applicable. Do not drop first and last layers.
+            if (
+                self.training
+                and self.layer_drop > 0
+                and i != 0
+                and i != (len(self.layers) - 1)
+                and random.random() < self.layer_drop
+            ):
+                unused_params.extend(list(lyr.parameters()))
+            else:
+                # First and last blocks will have autocast disabled for improved precision.
+                with autocast(x.device.type, enabled=self.enable_fp16 and i != 0):
+                    x = lyr(x, time_emb)
+
+        x = x.float()
+        out = self.out(x)
+
+        # Involve probabilistic or possibly unused parameters in loss so we don't get DDP errors.
+        extraneous_addition = 0
+        for p in unused_params:
+            extraneous_addition = extraneous_addition + p.mean()
+        out = out + extraneous_addition * 0
+
+        if return_code_pred:
+            return out, mel_pred
+        return out
+
+
+if __name__ == "__main__":
+    clip = torch.randn(2, 100, 400)
+    aligned_latent = torch.randn(2, 388, 512)
+    aligned_sequence = torch.randint(0, 8192, (2, 100))
+    cond = torch.randn(2, 100, 400)
+    ts = torch.LongTensor([600, 600])
+    model = DiffusionTts(512, layer_drop=0.3, unconditioned_percentage=0.5)
+    # Test with latent aligned conditioning
+    # o = model(clip, ts, aligned_latent, cond)
+    # Test with sequence aligned conditioning
+    o = model(clip, ts, aligned_sequence, cond)

tortoise/models/random_latent_generator.py

@@ -0,0 +1,56 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2**0.5):
+    if bias is not None:
+        rest_dim = [1] * (input.ndim - bias.ndim - 1)
+        return (
+            F.leaky_relu(
+                input + bias.view(1, bias.shape[0], *rest_dim),
+                negative_slope=negative_slope,
+            )
+            * scale
+        )
+    else:
+        return F.leaky_relu(input, negative_slope=0.2) * scale
+
+
+class EqualLinear(nn.Module):
+    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1):
+        super().__init__()
+        self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+        else:
+            self.bias = None
+        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+        self.lr_mul = lr_mul
+
+    def forward(self, input):
+        out = F.linear(input, self.weight * self.scale)
+        out = fused_leaky_relu(out, self.bias * self.lr_mul)
+        return out
+
+
+class RandomLatentConverter(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.layers = nn.Sequential(
+            *[EqualLinear(channels, channels, lr_mul=0.1) for _ in range(5)],
+            nn.Linear(channels, channels)
+        )
+        self.channels = channels
+
+    def forward(self, ref):
+        r = torch.randn(ref.shape[0], self.channels, device=ref.device)
+        y = self.layers(r)
+        return y
+
+
+if __name__ == "__main__":
+    model = RandomLatentConverter(512)
+    model(torch.randn(5, 512))

tortoise/models/transformer.py

@@ -0,0 +1,237 @@
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+# helpers
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def cast_tuple(val, depth=1):
+    if isinstance(val, list):
+        val = tuple(val)
+    return val if isinstance(val, tuple) else (val,) * depth
+
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+
+def stable_softmax(t, dim=-1, alpha=32**2):
+    t = t / alpha
+    t = t - torch.amax(t, dim=dim, keepdim=True).detach()
+    return (t * alpha).softmax(dim=dim)
+
+
+def route_args(router, args, depth):
+    routed_args = [(dict(), dict()) for _ in range(depth)]
+    matched_keys = [key for key in args.keys() if key in router]
+
+    for key in matched_keys:
+        val = args[key]
+        for depth, ((f_args, g_args), routes) in enumerate(
+            zip(routed_args, router[key])
+        ):
+            new_f_args, new_g_args = map(
+                lambda route: ({key: val} if route else {}), routes
+            )
+            routed_args[depth] = ({**f_args, **new_f_args}, {**g_args, **new_g_args})
+    return routed_args
+
+
+# classes
+class SequentialSequence(nn.Module):
+    def __init__(self, layers, args_route={}, layer_dropout=0.0):
+        super().__init__()
+        assert all(
+            len(route) == len(layers) for route in args_route.values()
+        ), "each argument route map must have the same depth as the number of sequential layers"
+        self.layers = layers
+        self.args_route = args_route
+        self.layer_dropout = layer_dropout
+
+    def forward(self, x, **kwargs):
+        args = route_args(self.args_route, kwargs, len(self.layers))
+        layers_and_args = list(zip(self.layers, args))
+
+        for (f, g), (f_args, g_args) in layers_and_args:
+            x = x + f(x, **f_args)
+            x = x + g(x, **g_args)
+        return x
+
+
+class DivideMax(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        maxes = x.amax(dim=self.dim, keepdim=True).detach()
+        return x / maxes
+
+
+# https://arxiv.org/abs/2103.17239
+class LayerScale(nn.Module):
+    def __init__(self, dim, depth, fn):
+        super().__init__()
+        if depth <= 18:
+            init_eps = 0.1
+        elif depth > 18 and depth <= 24:
+            init_eps = 1e-5
+        else:
+            init_eps = 1e-6
+
+        scale = torch.zeros(1, 1, dim).fill_(init_eps)
+        self.scale = nn.Parameter(scale)
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) * self.scale
+
+
+# layer norm
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn, sandwich=False):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.norm_out = nn.LayerNorm(dim) if sandwich else nn.Identity()
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        x = self.norm(x)
+        x = self.fn(x, **kwargs)
+        return self.norm_out(x)
+
+
+# feed forward
+
+
+class GEGLU(nn.Module):
+    def forward(self, x):
+        x, gates = x.chunk(2, dim=-1)
+        return x * F.gelu(gates)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dropout=0.0, mult=4.0):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * mult * 2),
+            GEGLU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * mult, dim),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# Attention
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, seq_len, causal=True, heads=8, dim_head=64, dropout=0.0):
+        super().__init__()
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.seq_len = seq_len
+        self.scale = dim_head**-0.5
+
+        self.causal = causal
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))
+
+    def forward(self, x, mask=None):
+        b, n, _, h, device = *x.shape, self.heads, x.device
+        softmax = torch.softmax
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), qkv)
+
+        q = q * self.scale
+
+        dots = torch.einsum("b h i d, b h j d -> b h i j", q, k)
+        mask_value = max_neg_value(dots)
+
+        if exists(mask):
+            mask = rearrange(mask, "b j -> b () () j")
+            dots.masked_fill_(~mask, mask_value)
+            del mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
+            dots.masked_fill_(mask, mask_value)
+
+        attn = softmax(dots, dim=-1)
+
+        out = torch.einsum("b h i j, b h j d -> b h i d", attn, v)
+        out = rearrange(out, "b h n d -> b n (h d)")
+        out = self.to_out(out)
+        return out
+
+
+# main transformer class
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        seq_len,
+        causal=True,
+        heads=8,
+        dim_head=64,
+        ff_mult=4,
+        attn_dropout=0.0,
+        ff_dropout=0.0,
+        sparse_attn=False,
+        sandwich_norm=False,
+    ):
+        super().__init__()
+        layers = nn.ModuleList([])
+        sparse_layer = cast_tuple(sparse_attn, depth)
+
+        for ind, sparse_attn in zip(range(depth), sparse_layer):
+            attn = Attention(
+                dim,
+                causal=causal,
+                seq_len=seq_len,
+                heads=heads,
+                dim_head=dim_head,
+                dropout=attn_dropout,
+            )
+
+            ff = FeedForward(dim, mult=ff_mult, dropout=ff_dropout)
+
+            layers.append(
+                nn.ModuleList(
+                    [
+                        LayerScale(
+                            dim, ind + 1, PreNorm(dim, attn, sandwich=sandwich_norm)
+                        ),
+                        LayerScale(
+                            dim, ind + 1, PreNorm(dim, ff, sandwich=sandwich_norm)
+                        ),
+                    ]
+                )
+            )
+
+        execute_type = SequentialSequence
+        route_attn = ((True, False),) * depth
+        attn_route_map = {"mask": route_attn}
+
+        self.layers = execute_type(layers, args_route=attn_route_map)
+
+    def forward(self, x, **kwargs):
+        return self.layers(x, **kwargs)

tortoise/models/utils.py

@@ -0,0 +1,80 @@
+import os
+try: import gdown
+except ImportError:
+    raise ImportError(
+        "Sorry, gdown is required in order to download the new BigVGAN vocoder.\n"
+        "Please install it with `pip install gdown` and try again."
+    )
+from urllib import request
+
+import progressbar
+
+D_STEM = "https://drive.google.com/uc?id="
+
+DEFAULT_MODELS_DIR = os.path.join(
+    os.path.expanduser("~"), ".cache", "tortoise", "models"
+)
+MODELS_DIR = os.environ.get("TORTOISE_MODELS_DIR", DEFAULT_MODELS_DIR)
+MODELS = {
+    "autoregressive.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/autoregressive.pth",
+    "classifier.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/classifier.pth",
+    "clvp2.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/clvp2.pth",
+    "cvvp.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/cvvp.pth",
+    "diffusion_decoder.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/diffusion_decoder.pth",
+    "vocoder.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/vocoder.pth",
+    "rlg_auto.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/rlg_auto.pth",
+    "rlg_diffuser.pth": "https://huggingface.co/jbetker/tortoise-tts-v2/resolve/main/.models/rlg_diffuser.pth",
+    # these links are from the nvidia gdrive
+    "bigvgan_base_24khz_100band_g.pth": "https://drive.google.com/uc?id=1_cKskUDuvxQJUEBwdgjAxKuDTUW6kPdY",
+    "bigvgan_24khz_100band_g.pth": "https://drive.google.com/uc?id=1wmP_mAs7d00KHVfVEl8B5Gb72Kzpcavp",
+}
+
+
+pbar = None
+def download_models(specific_models=None):
+    """
+    Call to download all the models that Tortoise uses.
+    """
+    os.makedirs(MODELS_DIR, exist_ok=True)
+
+    def show_progress(block_num, block_size, total_size):
+        global pbar
+        if pbar is None:
+            pbar = progressbar.ProgressBar(maxval=total_size)
+            pbar.start()
+
+        downloaded = block_num * block_size
+        if downloaded < total_size:
+            pbar.update(downloaded)
+        else:
+            pbar.finish()
+            pbar = None
+
+    for model_name, url in MODELS.items():
+        if specific_models is not None and model_name not in specific_models:
+            continue
+        model_path = os.path.join(MODELS_DIR, model_name)
+        if os.path.exists(model_path):
+            continue
+        print(f"Downloading {model_name} from {url}...")
+        if D_STEM in url:
+            gdown.download(url, model_path, quiet=False)
+        else:
+            request.urlretrieve(url, model_path, show_progress)
+        print("Done.")
+
+
+def get_model_path(model_name, models_dir=MODELS_DIR):
+    """
+    Get path to given model, download it if it doesn't exist.
+    """
+    if model_name not in MODELS:
+        raise ValueError(f"Model {model_name} not found in available models.")
+    model_path = os.path.join(models_dir, model_name)
+    if not os.path.exists(model_path) and models_dir == MODELS_DIR:
+        download_models([model_name])
+    return model_path
+
+if __name__ == "__main__":
+    download_models() # to download all models
+    

tortoise/models/vocoder.py

@@ -0,0 +1,440 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import json
+from enum import Enum
+from typing import Optional, Callable
+from dataclasses import dataclass
+try:
+    from BigVGAN.models import BigVGAN as BVGModel
+    from BigVGAN.env import AttrDict
+except ImportError:
+    raise ImportError(
+        "BigVGAN not installed, can't use BigVGAN vocoder\n"
+        "Please see the installation instructions on README."
+    )
+
+MAX_WAV_VALUE = 32768.0
+
+
+class KernelPredictor(torch.nn.Module):
+    """Kernel predictor for the location-variable convolutions"""
+
+    def __init__(
+        self,
+        cond_channels,
+        conv_in_channels,
+        conv_out_channels,
+        conv_layers,
+        conv_kernel_size=3,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+        kpnet_nonlinear_activation="LeakyReLU",
+        kpnet_nonlinear_activation_params={"negative_slope": 0.1},
+    ):
+        """
+        Args:
+            cond_channels (int): number of channel for the conditioning sequence,
+            conv_in_channels (int): number of channel for the input sequence,
+            conv_out_channels (int): number of channel for the output sequence,
+            conv_layers (int): number of layers
+        """
+        super().__init__()
+
+        self.conv_in_channels = conv_in_channels
+        self.conv_out_channels = conv_out_channels
+        self.conv_kernel_size = conv_kernel_size
+        self.conv_layers = conv_layers
+
+        kpnet_kernel_channels = (
+            conv_in_channels * conv_out_channels * conv_kernel_size * conv_layers
+        )  # l_w
+        kpnet_bias_channels = conv_out_channels * conv_layers  # l_b
+
+        self.input_conv = nn.Sequential(
+            nn.utils.weight_norm(
+                nn.Conv1d(cond_channels, kpnet_hidden_channels, 5, padding=2, bias=True)
+            ),
+            getattr(nn, kpnet_nonlinear_activation)(
+                **kpnet_nonlinear_activation_params
+            ),
+        )
+
+        self.residual_convs = nn.ModuleList()
+        padding = (kpnet_conv_size - 1) // 2
+        for _ in range(3):
+            self.residual_convs.append(
+                nn.Sequential(
+                    nn.Dropout(kpnet_dropout),
+                    nn.utils.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(
+                        **kpnet_nonlinear_activation_params
+                    ),
+                    nn.utils.weight_norm(
+                        nn.Conv1d(
+                            kpnet_hidden_channels,
+                            kpnet_hidden_channels,
+                            kpnet_conv_size,
+                            padding=padding,
+                            bias=True,
+                        )
+                    ),
+                    getattr(nn, kpnet_nonlinear_activation)(
+                        **kpnet_nonlinear_activation_params
+                    ),
+                )
+            )
+        self.kernel_conv = nn.utils.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_kernel_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+        self.bias_conv = nn.utils.weight_norm(
+            nn.Conv1d(
+                kpnet_hidden_channels,
+                kpnet_bias_channels,
+                kpnet_conv_size,
+                padding=padding,
+                bias=True,
+            )
+        )
+
+    def forward(self, c):
+        """
+        Args:
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+        """
+        batch, _, cond_length = c.shape
+        c = self.input_conv(c)
+        for residual_conv in self.residual_convs:
+            residual_conv.to(c.device)
+            c = c + residual_conv(c)
+        k = self.kernel_conv(c)
+        b = self.bias_conv(c)
+        kernels = k.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_in_channels,
+            self.conv_out_channels,
+            self.conv_kernel_size,
+            cond_length,
+        )
+        bias = b.contiguous().view(
+            batch,
+            self.conv_layers,
+            self.conv_out_channels,
+            cond_length,
+        )
+
+        return kernels, bias
+
+    def remove_weight_norm(self):
+        nn.utils.remove_weight_norm(self.input_conv[0])
+        nn.utils.remove_weight_norm(self.kernel_conv)
+        nn.utils.remove_weight_norm(self.bias_conv)
+        for block in self.residual_convs:
+            nn.utils.remove_weight_norm(block[1])
+            nn.utils.remove_weight_norm(block[3])
+
+
+class LVCBlock(torch.nn.Module):
+    """the location-variable convolutions"""
+
+    def __init__(
+        self,
+        in_channels,
+        cond_channels,
+        stride,
+        dilations=[1, 3, 9, 27],
+        lReLU_slope=0.2,
+        conv_kernel_size=3,
+        cond_hop_length=256,
+        kpnet_hidden_channels=64,
+        kpnet_conv_size=3,
+        kpnet_dropout=0.0,
+    ):
+        super().__init__()
+
+        self.cond_hop_length = cond_hop_length
+        self.conv_layers = len(dilations)
+        self.conv_kernel_size = conv_kernel_size
+
+        self.kernel_predictor = KernelPredictor(
+            cond_channels=cond_channels,
+            conv_in_channels=in_channels,
+            conv_out_channels=2 * in_channels,
+            conv_layers=len(dilations),
+            conv_kernel_size=conv_kernel_size,
+            kpnet_hidden_channels=kpnet_hidden_channels,
+            kpnet_conv_size=kpnet_conv_size,
+            kpnet_dropout=kpnet_dropout,
+            kpnet_nonlinear_activation_params={"negative_slope": lReLU_slope},
+        )
+
+        self.convt_pre = nn.Sequential(
+            nn.LeakyReLU(lReLU_slope),
+            nn.utils.weight_norm(
+                nn.ConvTranspose1d(
+                    in_channels,
+                    in_channels,
+                    2 * stride,
+                    stride=stride,
+                    padding=stride // 2 + stride % 2,
+                    output_padding=stride % 2,
+                )
+            ),
+        )
+
+        self.conv_blocks = nn.ModuleList()
+        for dilation in dilations:
+            self.conv_blocks.append(
+                nn.Sequential(
+                    nn.LeakyReLU(lReLU_slope),
+                    nn.utils.weight_norm(
+                        nn.Conv1d(
+                            in_channels,
+                            in_channels,
+                            conv_kernel_size,
+                            padding=dilation * (conv_kernel_size - 1) // 2,
+                            dilation=dilation,
+                        )
+                    ),
+                    nn.LeakyReLU(lReLU_slope),
+                )
+            )
+
+    def forward(self, x, c):
+        """forward propagation of the location-variable convolutions.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length)
+            c (Tensor): the conditioning sequence (batch, cond_channels, cond_length)
+
+        Returns:
+            Tensor: the output sequence (batch, in_channels, in_length)
+        """
+        _, in_channels, _ = x.shape  # (B, c_g, L')
+
+        x = self.convt_pre(x)  # (B, c_g, stride * L')
+        kernels, bias = self.kernel_predictor(c)
+
+        for i, conv in enumerate(self.conv_blocks):
+            output = conv(x)  # (B, c_g, stride * L')
+
+            k = kernels[:, i, :, :, :, :]  # (B, 2 * c_g, c_g, kernel_size, cond_length)
+            b = bias[:, i, :, :]  # (B, 2 * c_g, cond_length)
+
+            output = self.location_variable_convolution(
+                output, k, b, hop_size=self.cond_hop_length
+            )  # (B, 2 * c_g, stride * L'): LVC
+            x = x + torch.sigmoid(output[:, :in_channels, :]) * torch.tanh(
+                output[:, in_channels:, :]
+            )  # (B, c_g, stride * L'): GAU
+
+        return x
+
+    def location_variable_convolution(self, x, kernel, bias, dilation=1, hop_size=256):
+        """perform location-variable convolution operation on the input sequence (x) using the local convolution kernl.
+        Time: 414 μs ± 309 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each), test on NVIDIA V100.
+        Args:
+            x (Tensor): the input sequence (batch, in_channels, in_length).
+            kernel (Tensor): the local convolution kernel (batch, in_channel, out_channels, kernel_size, kernel_length)
+            bias (Tensor): the bias for the local convolution (batch, out_channels, kernel_length)
+            dilation (int): the dilation of convolution.
+            hop_size (int): the hop_size of the conditioning sequence.
+        Returns:
+            (Tensor): the output sequence after performing local convolution. (batch, out_channels, in_length).
+        """
+        batch, _, in_length = x.shape
+        batch, _, out_channels, kernel_size, kernel_length = kernel.shape
+        assert in_length == (
+            kernel_length * hop_size
+        ), "length of (x, kernel) is not matched"
+
+        padding = dilation * int((kernel_size - 1) / 2)
+        x = F.pad(
+            x, (padding, padding), "constant", 0
+        )  # (batch, in_channels, in_length + 2*padding)
+        x = x.unfold(
+            2, hop_size + 2 * padding, hop_size
+        )  # (batch, in_channels, kernel_length, hop_size + 2*padding)
+
+        if hop_size < dilation:
+            x = F.pad(x, (0, dilation), "constant", 0)
+        x = x.unfold(
+            3, dilation, dilation
+        )  # (batch, in_channels, kernel_length, (hop_size + 2*padding)/dilation, dilation)
+        x = x[:, :, :, :, :hop_size]
+        x = x.transpose(
+            3, 4
+        )  # (batch, in_channels, kernel_length, dilation, (hop_size + 2*padding)/dilation)
+        x = x.unfold(
+            4, kernel_size, 1
+        )  # (batch, in_channels, kernel_length, dilation, _, kernel_size)
+
+        o = torch.einsum("bildsk,biokl->bolsd", x, kernel)
+        o = o.to(memory_format=torch.channels_last_3d)
+        bias = bias.unsqueeze(-1).unsqueeze(-1).to(memory_format=torch.channels_last_3d)
+        o = o + bias
+        o = o.contiguous().view(batch, out_channels, -1)
+
+        return o
+
+    def remove_weight_norm(self):
+        self.kernel_predictor.remove_weight_norm()
+        nn.utils.remove_weight_norm(self.convt_pre[1])
+        for block in self.conv_blocks:
+            nn.utils.remove_weight_norm(block[1])
+
+
+class UnivNetGenerator(nn.Module):
+    """
+    UnivNet Generator
+
+    Originally from https://github.com/mindslab-ai/univnet/blob/master/model/generator.py.
+    """
+
+    def __init__(
+        self,
+        noise_dim=64,
+        channel_size=32,
+        dilations=[1, 3, 9, 27],
+        strides=[8, 8, 4],
+        lReLU_slope=0.2,
+        kpnet_conv_size=3,
+        # Below are MEL configurations options that this generator requires.
+        hop_length=256,
+        n_mel_channels=100,
+    ):
+        super(UnivNetGenerator, self).__init__()
+        self.mel_channel = n_mel_channels
+        self.noise_dim = noise_dim
+        self.hop_length = hop_length
+        channel_size = channel_size
+        kpnet_conv_size = kpnet_conv_size
+
+        self.res_stack = nn.ModuleList()
+        hop_length = 1
+        for stride in strides:
+            hop_length = stride * hop_length
+            self.res_stack.append(
+                LVCBlock(
+                    channel_size,
+                    n_mel_channels,
+                    stride=stride,
+                    dilations=dilations,
+                    lReLU_slope=lReLU_slope,
+                    cond_hop_length=hop_length,
+                    kpnet_conv_size=kpnet_conv_size,
+                )
+            )
+
+        self.conv_pre = nn.utils.weight_norm(
+            nn.Conv1d(noise_dim, channel_size, 7, padding=3, padding_mode="reflect")
+        )
+
+        self.conv_post = nn.Sequential(
+            nn.LeakyReLU(lReLU_slope),
+            nn.utils.weight_norm(
+                nn.Conv1d(channel_size, 1, 7, padding=3, padding_mode="reflect")
+            ),
+            nn.Tanh(),
+        )
+
+    def forward(self, c, z):
+        """
+        Args:
+            c (Tensor): the conditioning sequence of mel-spectrogram (batch, mel_channels, in_length)
+            z (Tensor): the noise sequence (batch, noise_dim, in_length)
+
+        """
+        z = self.conv_pre(z)  # (B, c_g, L)
+
+        for res_block in self.res_stack:
+            res_block.to(z.device)
+            z = res_block(z, c)  # (B, c_g, L * s_0 * ... * s_i)
+
+        z = self.conv_post(z)  # (B, 1, L * 256)
+
+        return z
+
+    def eval(self, inference=False):
+        super(UnivNetGenerator, self).eval()
+        # don't remove weight norm while validation in training loop
+        if inference:
+            self.remove_weight_norm()
+
+    def remove_weight_norm(self):
+        nn.utils.remove_weight_norm(self.conv_pre)
+
+        for layer in self.conv_post:
+            if len(layer.state_dict()) != 0:
+                nn.utils.remove_weight_norm(layer)
+
+        for res_block in self.res_stack:
+            res_block.remove_weight_norm()
+
+    def inference(self, c, z=None):
+        # pad input mel with zeros to cut artifact
+        # see https://github.com/seungwonpark/melgan/issues/8
+        zero = torch.full((c.shape[0], self.mel_channel, 10), -11.5129).to(c.device)
+        mel = torch.cat((c, zero), dim=2)
+
+        if z is None:
+            z = torch.randn(c.shape[0], self.noise_dim, mel.size(2)).to(mel.device)
+
+        audio = self.forward(mel, z)
+        audio = audio[:, :, : -(self.hop_length * 10)]
+        audio = audio.clamp(min=-1, max=1)
+        return audio
+
+from pathlib import Path
+STATIC_DIR = Path(__file__).parent.parent.parent/'static'
+assert STATIC_DIR.is_dir()
+def BVGWithConf(fname: str):
+    json_config = json.loads(
+        (STATIC_DIR/fname).read_text()
+    )
+    return lambda: BVGModel(AttrDict(json_config))
+
+@dataclass
+class VocType:
+    constructor: Callable[[], nn.Module]
+    model_path: str
+    subkey: Optional[str] = None
+    def optionally_index(self, model_dict):
+        if self.subkey is not None:
+            return model_dict[self.subkey]
+        return model_dict
+class VocConf(Enum):
+    Univnet = VocType(UnivNetGenerator, "vocoder.pth", 'model_g')
+    BigVGAN_Base = VocType(BVGWithConf("bigvgan_base_24khz_100band_config.json"), "bigvgan_base_24khz_100band_g.pth", 'generator')
+    BigVGAN = VocType(BVGWithConf("bigvgan_24khz_100band_config.json"), "bigvgan_24khz_100band_g.pth", 'generator')
+
+
+if __name__ == "__main__":
+    model = UnivNetGenerator()
+
+    c = torch.randn(3, 100, 10)
+    z = torch.randn(3, 64, 10)
+    print(c.shape)
+
+    y = model(c, z)
+    print(y.shape)
+    assert y.shape == torch.Size([3, 1, 2560])
+
+    pytorch_total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(pytorch_total_params)

tortoise/models/xtransformers.py

@@ -0,0 +1,1436 @@
+import math
+from collections import namedtuple
+from functools import partial
+from inspect import isfunction
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import einsum, nn
+
+DEFAULT_DIM_HEAD = 64
+
+Intermediates = namedtuple("Intermediates", ["pre_softmax_attn", "post_softmax_attn"])
+
+LayerIntermediates = namedtuple(
+    "Intermediates",
+    [
+        "hiddens",
+        "attn_intermediates",
+        "past_key_values",
+    ],
+)
+
+
+# helpers
+
+
+def exists(val):
+    return val is not None
+
+
+def default(val, d):
+    if exists(val):
+        return val
+    return d() if isfunction(d) else d
+
+
+def cast_tuple(val, depth):
+    return val if isinstance(val, tuple) else (val,) * depth
+
+
+class always:
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, *args, **kwargs):
+        return self.val
+
+
+class not_equals:
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, x, *args, **kwargs):
+        return x != self.val
+
+
+class equals:
+    def __init__(self, val):
+        self.val = val
+
+    def __call__(self, x, *args, **kwargs):
+        return x == self.val
+
+
+def max_neg_value(tensor):
+    return -torch.finfo(tensor.dtype).max
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+# init helpers
+
+
+def init_zero_(layer):
+    nn.init.constant_(layer.weight, 0.0)
+    if exists(layer.bias):
+        nn.init.constant_(layer.bias, 0.0)
+
+
+# keyword argument helpers
+
+
+def pick_and_pop(keys, d):
+    values = list(map(lambda key: d.pop(key), keys))
+    return dict(zip(keys, values))
+
+
+def group_dict_by_key(cond, d):
+    return_val = [dict(), dict()]
+    for key in d.keys():
+        match = bool(cond(key))
+        ind = int(not match)
+        return_val[ind][key] = d[key]
+    return (*return_val,)
+
+
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
+
+
+def group_by_key_prefix(prefix, d):
+    return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+
+def groupby_prefix_and_trim(prefix, d):
+    kwargs_with_prefix, kwargs = group_dict_by_key(
+        partial(string_begins_with, prefix), d
+    )
+    kwargs_without_prefix = dict(
+        map(lambda x: (x[0][len(prefix) :], x[1]), tuple(kwargs_with_prefix.items()))
+    )
+    return kwargs_without_prefix, kwargs
+
+
+# activations
+
+
+class ReluSquared(nn.Module):
+    def forward(self, x):
+        return F.relu(x) ** 2
+
+
+# positional embeddings
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len):
+        super().__init__()
+        self.scale = dim**-0.5
+        self.emb = nn.Embedding(max_seq_len, dim)
+
+    def forward(self, x):
+        n = torch.arange(x.shape[1], device=x.device)
+        pos_emb = self.emb(n)
+        pos_emb = rearrange(pos_emb, "n d -> () n d")
+        return pos_emb * self.scale
+
+
+class FixedPositionalEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    def forward(self, x, seq_dim=1, offset=0):
+        t = (
+            torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq)
+            + offset
+        )
+        sinusoid_inp = torch.einsum("i , j -> i j", t, self.inv_freq)
+        emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+        return rearrange(emb, "n d -> () n d")
+
+
+class RelativePositionBias(nn.Module):
+    def __init__(self, scale, causal=False, num_buckets=32, max_distance=128, heads=8):
+        super().__init__()
+        self.scale = scale
+        self.causal = causal
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.relative_attention_bias = nn.Embedding(num_buckets, heads)
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position, causal=True, num_buckets=32, max_distance=128
+    ):
+        ret = 0
+        n = -relative_position
+        if not causal:
+            num_buckets //= 2
+            ret += (n < 0).long() * num_buckets
+            n = torch.abs(n)
+        else:
+            n = torch.max(n, torch.zeros_like(n))
+
+        max_exact = num_buckets // 2
+        is_small = n < max_exact
+
+        val_if_large = (
+            max_exact
+            + (
+                torch.log(n.float() / max_exact)
+                / math.log(max_distance / max_exact)
+                * (num_buckets - max_exact)
+            ).long()
+        )
+        val_if_large = torch.min(
+            val_if_large, torch.full_like(val_if_large, num_buckets - 1)
+        )
+
+        ret += torch.where(is_small, n, val_if_large)
+        return ret
+
+    def forward(self, qk_dots):
+        i, j, device = *qk_dots.shape[-2:], qk_dots.device
+        q_pos = torch.arange(i, dtype=torch.long, device=device)
+        k_pos = torch.arange(j, dtype=torch.long, device=device)
+        rel_pos = k_pos[None, :] - q_pos[:, None]
+        rp_bucket = self._relative_position_bucket(
+            rel_pos,
+            causal=self.causal,
+            num_buckets=self.num_buckets,
+            max_distance=self.max_distance,
+        )
+        values = self.relative_attention_bias(rp_bucket)
+        bias = rearrange(values, "i j h -> () h i j")
+        return qk_dots + (bias * self.scale)
+
+
+class AlibiPositionalBias(nn.Module):
+    def __init__(self, heads, **kwargs):
+        super().__init__()
+        self.heads = heads
+        slopes = torch.Tensor(self._get_slopes(heads))
+        slopes = rearrange(slopes, "h -> () h () ()")
+        self.register_buffer("slopes", slopes, persistent=False)
+        self.register_buffer("bias", None, persistent=False)
+
+    @staticmethod
+    def _get_slopes(heads):
+        def get_slopes_power_of_2(n):
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            ratio = start
+            return [start * ratio**i for i in range(n)]
+
+        if math.log2(heads).is_integer():
+            return get_slopes_power_of_2(heads)
+
+        closest_power_of_2 = 2 ** math.floor(math.log2(heads))
+        return (
+            get_slopes_power_of_2(closest_power_of_2)
+            + get_slopes_power_of_2(2 * closest_power_of_2)[0::2][
+                : heads - closest_power_of_2
+            ]
+        )
+
+    def forward(self, qk_dots):
+        h, i, j, device = *qk_dots.shape[-3:], qk_dots.device
+
+        if exists(self.bias) and self.bias.shape[-1] >= j:
+            return qk_dots + self.bias[..., :j]
+
+        bias = torch.arange(j, device=device)
+        bias = rearrange(bias, "j -> () () () j")
+        bias = bias * self.slopes
+
+        num_heads_unalibied = h - bias.shape[1]
+        bias = F.pad(bias, (0, 0, 0, 0, 0, num_heads_unalibied))
+
+        self.register_buffer("bias", bias, persistent=False)
+        return qk_dots + self.bias
+
+
+class LearnedAlibiPositionalBias(AlibiPositionalBias):
+    def __init__(self, heads, bidirectional=False):
+        super().__init__(heads)
+        los_slopes = torch.log(self.slopes)
+        self.learned_logslopes = nn.Parameter(los_slopes)
+
+        self.bidirectional = bidirectional
+        if self.bidirectional:
+            self.learned_logslopes_future = nn.Parameter(los_slopes)
+
+    def forward(self, qk_dots):
+        h, i, j, device = *qk_dots.shape[-3:], qk_dots.device
+
+        def get_slopes(param):
+            return F.pad(param.exp(), (0, 0, 0, 0, 0, h - param.shape[1]))
+
+        if exists(self.bias) and self.bias.shape[-1] >= j:
+            bias = self.bias[..., :i, :j]
+        else:
+            i_arange = torch.arange(i, device=device)
+            j_arange = torch.arange(j, device=device)
+            bias = rearrange(j_arange, "j -> 1 1 1 j") - rearrange(
+                i_arange, "i -> 1 1 i 1"
+            )
+            self.register_buffer("bias", bias, persistent=False)
+
+        if self.bidirectional:
+            past_slopes = get_slopes(self.learned_logslopes)
+            future_slopes = get_slopes(self.learned_logslopes_future)
+            bias = torch.tril(bias * past_slopes) + torch.triu(bias * future_slopes)
+        else:
+            slopes = get_slopes(self.learned_logslopes)
+            bias = bias * slopes
+
+        return qk_dots + bias
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+    def forward(self, max_seq_len, device):
+        t = torch.arange(max_seq_len, device=device).type_as(self.inv_freq)
+        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return rearrange(emb, "n d -> () () n d")
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (j d) -> ... j d", j=2)
+    x1, x2 = x.unbind(dim=-2)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(t, freqs):
+    seq_len = t.shape[-2]
+    freqs = freqs[:, :, -seq_len:]
+    return (t * freqs.cos()) + (rotate_half(t) * freqs.sin())
+
+
+# norms
+
+
+class Scale(nn.Module):
+    def __init__(self, value, fn):
+        super().__init__()
+        self.value = value
+        self.fn = fn
+
+    def forward(self, x, **kwargs):
+        out = self.fn(x, **kwargs)
+
+        def scale_fn(t):
+            return t * self.value
+
+        if not isinstance(out, tuple):
+            return scale_fn(out)
+
+        return (scale_fn(out[0]), *out[1:])
+
+
+class Rezero(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+        self.g = nn.Parameter(torch.zeros(1))
+
+    def forward(self, x, **kwargs):
+        out = self.fn(x, **kwargs)
+
+        def rezero_fn(t):
+            return t * self.g
+
+        if not isinstance(out, tuple):
+            return rezero_fn(out)
+
+        return (rezero_fn(out[0]), *out[1:])
+
+
+class ScaleNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.scale = dim**-0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim**-0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSScaleShiftNorm(nn.Module):
+    def __init__(self, dim, eps=1e-8):
+        super().__init__()
+        self.scale = dim**-0.5
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(dim))
+        self.scale_shift_process = nn.Linear(dim * 2, dim * 2)
+
+    def forward(self, x, norm_scale_shift_inp):
+        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+        norm = x / norm.clamp(min=self.eps) * self.g
+
+        ss_emb = self.scale_shift_process(norm_scale_shift_inp)
+        scale, shift = torch.chunk(ss_emb, 2, dim=1)
+        h = norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return h
+
+
+# residual and residual gates
+
+
+class Residual(nn.Module):
+    def __init__(self, dim, scale_residual=False):
+        super().__init__()
+        self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
+
+    def forward(self, x, residual):
+        if exists(self.residual_scale):
+            residual = residual * self.residual_scale
+
+        return x + residual
+
+
+class GRUGating(nn.Module):
+    def __init__(self, dim, scale_residual=False):
+        super().__init__()
+        self.gru = nn.GRUCell(dim, dim)
+        self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
+
+    def forward(self, x, residual):
+        if exists(self.residual_scale):
+            residual = residual * self.residual_scale
+
+        gated_output = self.gru(
+            rearrange(x, "b n d -> (b n) d"), rearrange(residual, "b n d -> (b n) d")
+        )
+
+        return gated_output.reshape_as(x)
+
+
+# token shifting
+
+
+def shift(t, amount, mask=None):
+    if amount == 0:
+        return t
+
+    if exists(mask):
+        t = t.masked_fill(~mask[..., None], 0.0)
+
+    return F.pad(t, (0, 0, amount, -amount), value=0.0)
+
+
+class ShiftTokens(nn.Module):
+    def __init__(self, shifts, fn):
+        super().__init__()
+        self.fn = fn
+        self.shifts = tuple(shifts)
+
+    def forward(self, x, **kwargs):
+        mask = kwargs.get("mask", None)
+        shifts = self.shifts
+        segments = len(shifts)
+        feats_per_shift = x.shape[-1] // segments
+        splitted = x.split(feats_per_shift, dim=-1)
+        segments_to_shift, rest = splitted[:segments], splitted[segments:]
+        segments_to_shift = list(
+            map(lambda args: shift(*args, mask=mask), zip(segments_to_shift, shifts))
+        )
+        x = torch.cat((*segments_to_shift, *rest), dim=-1)
+        return self.fn(x, **kwargs)
+
+
+# feedforward
+
+
+class GLU(nn.Module):
+    def __init__(self, dim_in, dim_out, activation):
+        super().__init__()
+        self.act = activation
+        self.proj = nn.Linear(dim_in, dim_out * 2)
+
+    def forward(self, x):
+        x, gate = self.proj(x).chunk(2, dim=-1)
+        return x * self.act(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_out=None,
+        mult=4,
+        glu=False,
+        relu_squared=False,
+        post_act_ln=False,
+        dropout=0.0,
+        zero_init_output=False,
+    ):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = default(dim_out, dim)
+        activation = ReluSquared() if relu_squared else nn.GELU()
+
+        project_in = (
+            nn.Sequential(nn.Linear(dim, inner_dim), activation)
+            if not glu
+            else GLU(dim, inner_dim, activation)
+        )
+
+        self.net = nn.Sequential(
+            project_in,
+            nn.LayerNorm(inner_dim) if post_act_ln else nn.Identity(),
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out),
+        )
+
+        # init last linear layer to 0
+        if zero_init_output:
+            init_zero_(self.net[-1])
+
+    def forward(self, x):
+        return self.net(x)
+
+
+# attention.
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_head=DEFAULT_DIM_HEAD,
+        heads=8,
+        causal=False,
+        talking_heads=False,
+        head_scale=False,
+        collab_heads=False,
+        collab_compression=0.3,
+        sparse_topk=None,
+        use_entmax15=False,
+        num_mem_kv=0,
+        dropout=0.0,
+        on_attn=False,
+        gate_values=False,
+        zero_init_output=False,
+        max_attend_past=None,
+        qk_norm=False,
+        scale_init_value=None,
+        rel_pos_bias=False,
+        rel_pos_num_buckets=32,
+        rel_pos_max_distance=128,
+    ):
+        super().__init__()
+        self.scale = dim_head**-0.5
+
+        self.heads = heads
+        self.causal = causal
+        self.max_attend_past = max_attend_past
+
+        qk_dim = v_dim = dim_head * heads
+
+        # collaborative heads
+        self.collab_heads = collab_heads
+        if self.collab_heads:
+            qk_dim = int(collab_compression * qk_dim)
+            self.collab_mixing = nn.Parameter(torch.randn(heads, qk_dim))
+
+        self.to_q = nn.Linear(dim, qk_dim, bias=False)
+        self.to_k = nn.Linear(dim, qk_dim, bias=False)
+        self.to_v = nn.Linear(dim, v_dim, bias=False)
+
+        self.dropout = nn.Dropout(dropout)
+
+        # add GLU gating for aggregated values, from alphafold2
+        self.to_v_gate = None
+        if gate_values:
+            self.to_v_gate = nn.Linear(dim, v_dim)
+            nn.init.constant_(self.to_v_gate.weight, 0)
+            nn.init.constant_(self.to_v_gate.bias, 1)
+
+        # cosine sim attention
+        self.qk_norm = qk_norm
+        if qk_norm:
+            scale_init_value = default(
+                scale_init_value, -3
+            )  # if not provided, initialize as though it were sequence length of 1024
+            self.scale = nn.Parameter(torch.ones(1, heads, 1, 1) * scale_init_value)
+
+        # talking heads
+        self.talking_heads = talking_heads
+        if talking_heads:
+            self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+            self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+
+        # head scaling
+        self.head_scale = head_scale
+        if head_scale:
+            self.head_scale_params = nn.Parameter(torch.ones(1, heads, 1, 1))
+
+        # explicit topk sparse attention
+        self.sparse_topk = sparse_topk
+
+        # entmax
+        self.attn_fn = F.softmax
+
+        # add memory key / values
+        self.num_mem_kv = num_mem_kv
+        if num_mem_kv > 0:
+            self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+            self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+
+        # attention on attention
+        self.attn_on_attn = on_attn
+        self.to_out = (
+            nn.Sequential(nn.Linear(v_dim, dim * 2), nn.GLU())
+            if on_attn
+            else nn.Linear(v_dim, dim)
+        )
+
+        self.rel_pos_bias = rel_pos_bias
+        if rel_pos_bias:
+            assert (
+                rel_pos_num_buckets <= rel_pos_max_distance
+            ), "number of relative position buckets must be less than the relative position max distance"
+            self.rel_pos = RelativePositionBias(
+                scale=dim_head**0.5,
+                causal=causal,
+                heads=heads,
+                num_buckets=rel_pos_num_buckets,
+                max_distance=rel_pos_max_distance,
+            )
+
+        # init output projection 0
+        if zero_init_output:
+            init_zero_(self.to_out)
+
+    def forward(
+        self,
+        x,
+        context=None,
+        mask=None,
+        context_mask=None,
+        attn_mask=None,
+        sinusoidal_emb=None,
+        rotary_pos_emb=None,
+        prev_attn=None,
+        mem=None,
+        layer_past=None,
+    ):
+        (
+            b,
+            n,
+            _,
+            h,
+            talking_heads,
+            collab_heads,
+            head_scale,
+            scale,
+            device,
+            has_context,
+        ) = (
+            *x.shape,
+            self.heads,
+            self.talking_heads,
+            self.collab_heads,
+            self.head_scale,
+            self.scale,
+            x.device,
+            exists(context),
+        )
+        kv_input = default(context, x)
+
+        q_input = x
+        k_input = kv_input
+        v_input = kv_input
+
+        if exists(mem):
+            k_input = torch.cat((mem, k_input), dim=-2)
+            v_input = torch.cat((mem, v_input), dim=-2)
+
+        if exists(sinusoidal_emb):
+            # in shortformer, the query would start at a position offset depending on the past cached memory
+            offset = k_input.shape[-2] - q_input.shape[-2]
+            q_input = q_input + sinusoidal_emb(q_input, offset=offset)
+            k_input = k_input + sinusoidal_emb(k_input)
+
+        q = self.to_q(q_input)
+        k = self.to_k(k_input)
+        v = self.to_v(v_input)
+
+        if not collab_heads:
+            q, k, v = map(
+                lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)
+            )
+        else:
+            q = einsum("b i d, h d -> b h i d", q, self.collab_mixing)
+            k = rearrange(k, "b n d -> b () n d")
+            v = rearrange(v, "b n (h d) -> b h n d", h=h)
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            k = torch.cat([past_key, k], dim=-2)
+            v = torch.cat([past_value, v], dim=-2)
+        k_cache = k
+        v_cache = v
+
+        if exists(rotary_pos_emb) and not has_context:
+            l = rotary_pos_emb.shape[-1]
+            (ql, qr), (kl, kr), (vl, vr) = map(
+                lambda t: (t[..., :l], t[..., l:]), (q, k, v)
+            )
+            ql, kl, vl = map(
+                lambda t: apply_rotary_pos_emb(t, rotary_pos_emb), (ql, kl, vl)
+            )
+            q, k, v = map(
+                lambda t: torch.cat(t, dim=-1), ((ql, qr), (kl, kr), (vl, vr))
+            )
+
+        input_mask = None
+        if any(map(exists, (mask, context_mask))):
+            q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
+            k_mask = q_mask if not exists(context) else context_mask
+            k_mask = default(
+                k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool()
+            )
+            q_mask = rearrange(q_mask, "b i -> b () i ()")
+            k_mask = rearrange(k_mask, "b j -> b () () j")
+            input_mask = q_mask * k_mask
+
+        if self.num_mem_kv > 0:
+            mem_k, mem_v = map(
+                lambda t: repeat(t, "h n d -> b h n d", b=b), (self.mem_k, self.mem_v)
+            )
+            k = torch.cat((mem_k, k), dim=-2)
+            v = torch.cat((mem_v, v), dim=-2)
+            if exists(input_mask):
+                input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
+
+        if collab_heads:
+            k = k.expand(-1, h, -1, -1)
+
+        if self.qk_norm:
+            q, k = map(l2norm, (q, k))
+            scale = 1 / (self.scale.exp().clamp(min=1e-2))
+
+        dots = einsum("b h i d, b h j d -> b h i j", q, k) * scale
+        mask_value = max_neg_value(dots)
+
+        if exists(prev_attn):
+            dots = dots + prev_attn
+
+        pre_softmax_attn = dots.clone()
+
+        if talking_heads:
+            dots = einsum(
+                "b h i j, h k -> b k i j", dots, self.pre_softmax_proj
+            ).contiguous()
+
+        if self.rel_pos_bias:
+            dots = self.rel_pos(dots)
+
+        if exists(input_mask):
+            dots.masked_fill_(~input_mask, mask_value)
+            del input_mask
+
+        if exists(attn_mask):
+            assert (
+                2 <= attn_mask.ndim <= 4
+            ), "attention mask must have greater than 2 dimensions but less than or equal to 4"
+            if attn_mask.ndim == 2:
+                attn_mask = rearrange(attn_mask, "i j -> () () i j")
+            elif attn_mask.ndim == 3:
+                attn_mask = rearrange(attn_mask, "h i j -> () h i j")
+            dots.masked_fill_(~attn_mask, mask_value)
+
+        if exists(self.max_attend_past):
+            i, j = dots.shape[-2:]
+            range_q = torch.arange(j - i, j, device=device)
+            range_k = torch.arange(j, device=device)
+            dist = rearrange(range_q, "i -> () () i ()") - rearrange(
+                range_k, "j -> () () () j"
+            )
+            mask = dist > self.max_attend_past
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if self.causal:
+            i, j = dots.shape[-2:]
+            r = torch.arange(i, device=device)
+            mask = rearrange(r, "i -> () () i ()") < rearrange(r, "j -> () () () j")
+            mask = F.pad(mask, (j - i, 0), value=False)
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
+            top, _ = dots.topk(self.sparse_topk, dim=-1)
+            vk = top[..., -1].unsqueeze(-1).expand_as(dots)
+            mask = dots < vk
+            dots.masked_fill_(mask, mask_value)
+            del mask
+
+        attn = self.attn_fn(dots, dim=-1)
+        post_softmax_attn = attn.clone()
+
+        attn = self.dropout(attn)
+
+        if talking_heads:
+            attn = einsum(
+                "b h i j, h k -> b k i j", attn, self.post_softmax_proj
+            ).contiguous()
+
+        out = einsum("b h i j, b h j d -> b h i d", attn, v)
+
+        if head_scale:
+            out = out * self.head_scale_params
+
+        out = rearrange(out, "b h n d -> b n (h d)")
+
+        if exists(self.to_v_gate):
+            gates = self.to_v_gate(x)
+            out = out * gates.sigmoid()
+
+        intermediates = Intermediates(
+            pre_softmax_attn=pre_softmax_attn, post_softmax_attn=post_softmax_attn
+        )
+
+        return self.to_out(out), intermediates, k_cache, v_cache
+
+
+class AttentionLayers(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        heads=8,
+        causal=False,
+        cross_attend=False,
+        only_cross=False,
+        use_scalenorm=False,
+        use_rms_scaleshift_norm=False,
+        use_rmsnorm=False,
+        use_rezero=False,
+        alibi_pos_bias=False,
+        alibi_num_heads=None,
+        alibi_learned=False,
+        position_infused_attn=False,
+        rotary_pos_emb=False,
+        rotary_emb_dim=None,
+        custom_layers=None,
+        sandwich_coef=None,
+        par_ratio=None,
+        residual_attn=False,
+        cross_residual_attn=False,
+        macaron=False,
+        pre_norm=True,
+        gate_residual=False,
+        scale_residual=False,
+        shift_tokens=0,
+        sandwich_norm=False,
+        use_qk_norm_attn=False,
+        qk_norm_attn_seq_len=None,
+        zero_init_branch_output=False,
+        **kwargs,
+    ):
+        super().__init__()
+        ff_kwargs, kwargs = groupby_prefix_and_trim("ff_", kwargs)
+        attn_kwargs, _ = groupby_prefix_and_trim("attn_", kwargs)
+
+        dim_head = attn_kwargs.get("dim_head", DEFAULT_DIM_HEAD)
+
+        self.dim = dim
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+        self.causal = causal
+
+        rel_pos_bias = "rel_pos_bias" in attn_kwargs
+        self.has_pos_emb = position_infused_attn or rel_pos_bias or rotary_pos_emb
+        self.pia_pos_emb = (
+            FixedPositionalEmbedding(dim) if position_infused_attn else None
+        )
+
+        rotary_emb_dim = max(default(rotary_emb_dim, dim_head // 2), 32)
+        self.rotary_pos_emb = (
+            RotaryEmbedding(rotary_emb_dim) if rotary_pos_emb else None
+        )
+
+        assert not (
+            alibi_pos_bias and rel_pos_bias
+        ), "you can only choose Alibi positional bias or T5 relative positional bias, not both"
+
+        if alibi_pos_bias:
+            alibi_num_heads = default(alibi_num_heads, heads)
+            assert (
+                alibi_num_heads <= heads
+            ), "number of ALiBi heads must be less than the total number of heads"
+            alibi_pos_klass = (
+                LearnedAlibiPositionalBias
+                if alibi_learned or not causal
+                else AlibiPositionalBias
+            )
+            self.rel_pos = alibi_pos_klass(
+                heads=alibi_num_heads, bidirectional=not causal
+            )
+        else:
+            self.rel_pos = None
+
+        assert not (
+            not pre_norm and sandwich_norm
+        ), "sandwich norm cannot be used when not using prenorm"
+        self.pre_norm = pre_norm
+        self.sandwich_norm = sandwich_norm
+
+        self.residual_attn = residual_attn
+        self.cross_residual_attn = cross_residual_attn
+        self.cross_attend = cross_attend
+
+        norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
+        norm_class = RMSNorm if use_rmsnorm else norm_class
+        norm_class = RMSScaleShiftNorm if use_rms_scaleshift_norm else norm_class
+        norm_fn = partial(norm_class, dim)
+
+        norm_fn = nn.Identity if use_rezero else norm_fn
+        branch_fn = Rezero if use_rezero else None
+
+        if cross_attend and not only_cross:
+            default_block = ("a", "c", "f")
+        elif cross_attend and only_cross:
+            default_block = ("c", "f")
+        else:
+            default_block = ("a", "f")
+
+        if macaron:
+            default_block = ("f",) + default_block
+
+        # qk normalization
+
+        if use_qk_norm_attn:
+            attn_scale_init_value = (
+                -math.log(math.log2(qk_norm_attn_seq_len**2 - qk_norm_attn_seq_len))
+                if exists(qk_norm_attn_seq_len)
+                else None
+            )
+            attn_kwargs = {
+                **attn_kwargs,
+                "qk_norm": True,
+                "scale_init_value": attn_scale_init_value,
+            }
+
+        # zero init
+
+        if zero_init_branch_output:
+            attn_kwargs = {**attn_kwargs, "zero_init_output": True}
+            ff_kwargs = {**ff_kwargs, "zero_init_output": True}
+
+        # calculate layer block order
+
+        if exists(custom_layers):
+            layer_types = custom_layers
+        elif exists(par_ratio):
+            par_depth = depth * len(default_block)
+            assert 1 < par_ratio <= par_depth, "par ratio out of range"
+            default_block = tuple(filter(not_equals("f"), default_block))
+            par_attn = par_depth // par_ratio
+            depth_cut = (
+                par_depth * 2 // 3
+            )  # 2 / 3 attention layer cutoff suggested by PAR paper
+            par_width = (depth_cut + depth_cut // par_attn) // par_attn
+            assert (
+                len(default_block) <= par_width
+            ), "default block is too large for par_ratio"
+            par_block = default_block + ("f",) * (par_width - len(default_block))
+            par_head = par_block * par_attn
+            layer_types = par_head + ("f",) * (par_depth - len(par_head))
+        elif exists(sandwich_coef):
+            assert (
+                sandwich_coef > 0 and sandwich_coef <= depth
+            ), "sandwich coefficient should be less than the depth"
+            layer_types = (
+                ("a",) * sandwich_coef
+                + default_block * (depth - sandwich_coef)
+                + ("f",) * sandwich_coef
+            )
+        else:
+            layer_types = default_block * depth
+
+        self.layer_types = layer_types
+        self.num_attn_layers = len(list(filter(equals("a"), layer_types)))
+
+        # calculate token shifting
+
+        shift_tokens = cast_tuple(shift_tokens, len(layer_types))
+
+        # iterate and construct layers
+
+        for ind, (layer_type, layer_shift_tokens) in enumerate(
+            zip(self.layer_types, shift_tokens)
+        ):
+            is_last_layer = ind == (len(self.layer_types) - 1)
+
+            if layer_type == "a":
+                layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
+            elif layer_type == "c":
+                layer = Attention(dim, heads=heads, **attn_kwargs)
+            elif layer_type == "f":
+                layer = FeedForward(dim, **ff_kwargs)
+                layer = layer if not macaron else Scale(0.5, layer)
+            else:
+                raise Exception(f"invalid layer type {layer_type}")
+
+            if layer_shift_tokens > 0:
+                shift_range_upper = layer_shift_tokens + 1
+                shift_range_lower = -layer_shift_tokens if not causal else 0
+                layer = ShiftTokens(range(shift_range_lower, shift_range_upper), layer)
+
+            if exists(branch_fn):
+                layer = branch_fn(layer)
+
+            residual_fn = GRUGating if gate_residual else Residual
+            residual = residual_fn(dim, scale_residual=scale_residual)
+
+            layer_uses_qk_norm = use_qk_norm_attn and layer_type in ("a", "c")
+
+            pre_branch_norm = norm_fn() if pre_norm and not layer_uses_qk_norm else None
+            post_branch_norm = (
+                norm_fn() if sandwich_norm or layer_uses_qk_norm else None
+            )
+            post_main_norm = norm_fn() if not pre_norm and not is_last_layer else None
+
+            norms = nn.ModuleList([pre_branch_norm, post_branch_norm, post_main_norm])
+
+            self.layers.append(nn.ModuleList([norms, layer, residual]))
+
+    def forward(
+        self,
+        x,
+        context=None,
+        full_context=None,  # for passing a list of hidden states from an encoder
+        mask=None,
+        context_mask=None,
+        attn_mask=None,
+        mems=None,
+        return_hiddens=False,
+        norm_scale_shift_inp=None,
+        past_key_values=None,
+        expected_seq_len=None,
+    ):
+
+        assert not (
+            self.cross_attend ^ (exists(context) or exists(full_context))
+        ), "context must be passed in if cross_attend is set to True"
+        assert (
+            context is None or full_context is None
+        ), "only one of full_context or context can be provided"
+
+        hiddens = []
+        intermediates = []
+        prev_attn = None
+        prev_cross_attn = None
+
+        mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
+        norm_args = {}
+        if exists(norm_scale_shift_inp):
+            norm_args["norm_scale_shift_inp"] = norm_scale_shift_inp
+
+        rotary_pos_emb = None
+        if exists(self.rotary_pos_emb):
+            if not self.training and self.causal:
+                assert (
+                    expected_seq_len is not None
+                ), "To decode a transformer with rotary embeddings, you must specify an `expected_seq_len`"
+            elif expected_seq_len is None:
+                expected_seq_len = 0
+            seq_len = x.shape[1]
+            if past_key_values is not None:
+                seq_len += past_key_values[0][0].shape[-2]
+            max_rotary_emb_length = max(
+                list(map(lambda m: (m.shape[1] if exists(m) else 0) + seq_len, mems))
+                + [expected_seq_len]
+            )
+            rotary_pos_emb = self.rotary_pos_emb(max_rotary_emb_length, x.device)
+
+        present_key_values = []
+        cross_attn_count = 0
+        for ind, (layer_type, (norm, block, residual_fn)) in enumerate(
+            zip(self.layer_types, self.layers)
+        ):
+            if layer_type == "a":
+                layer_mem = mems.pop(0) if mems else None
+
+            residual = x
+
+            pre_branch_norm, post_branch_norm, post_main_norm = norm
+
+            if exists(pre_branch_norm):
+                x = pre_branch_norm(x, **norm_args)
+
+            if layer_type == "a" or layer_type == "c":
+                if past_key_values is not None:
+                    layer_kv = past_key_values.pop(0)
+                    layer_past = tuple(s.to(x.device) for s in layer_kv)
+                else:
+                    layer_past = None
+
+            if layer_type == "a":
+                out, inter, k, v = block(
+                    x,
+                    None,
+                    mask,
+                    None,
+                    attn_mask,
+                    self.pia_pos_emb,
+                    rotary_pos_emb,
+                    prev_attn,
+                    layer_mem,
+                    layer_past,
+                )
+            elif layer_type == "c":
+                if exists(full_context):
+                    out, inter, k, v = block(
+                        x,
+                        full_context[cross_attn_count],
+                        mask,
+                        context_mask,
+                        None,
+                        None,
+                        None,
+                        prev_attn,
+                        None,
+                        layer_past,
+                    )
+                else:
+                    out, inter, k, v = block(
+                        x,
+                        context,
+                        mask,
+                        context_mask,
+                        None,
+                        None,
+                        None,
+                        prev_attn,
+                        None,
+                        layer_past,
+                    )
+            elif layer_type == "f":
+                out = block(x)
+
+            if (
+                layer_type == "a"
+                or layer_type == "c"
+                and present_key_values is not None
+            ):
+                present_key_values.append((k.detach(), v.detach()))
+
+            if exists(post_branch_norm):
+                out = post_branch_norm(out, **norm_args)
+
+            x = residual_fn(out, residual)
+
+            if layer_type in ("a", "c"):
+                intermediates.append(inter)
+
+            if layer_type == "a" and self.residual_attn:
+                prev_attn = inter.pre_softmax_attn
+            elif layer_type == "c" and self.cross_residual_attn:
+                prev_cross_attn = inter.pre_softmax_attn
+
+            if exists(post_main_norm):
+                x = post_main_norm(x, **norm_args)
+
+            if layer_type == "c":
+                cross_attn_count += 1
+
+            if layer_type == "f":
+                hiddens.append(x)
+
+        if return_hiddens:
+            intermediates = LayerIntermediates(
+                hiddens=hiddens,
+                attn_intermediates=intermediates,
+                past_key_values=present_key_values,
+            )
+
+            return x, intermediates
+
+        return x
+
+
+class Encoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert "causal" not in kwargs, "cannot set causality on encoder"
+        super().__init__(causal=False, **kwargs)
+
+
+class Decoder(AttentionLayers):
+    def __init__(self, **kwargs):
+        assert "causal" not in kwargs, "cannot set causality on decoder"
+        super().__init__(causal=True, **kwargs)
+
+
+class CrossAttender(AttentionLayers):
+    def __init__(self, **kwargs):
+        super().__init__(cross_attend=True, only_cross=True, **kwargs)
+
+
+class ViTransformerWrapper(nn.Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        attn_layers,
+        num_classes=None,
+        dropout=0.0,
+        emb_dropout=0.0,
+    ):
+        super().__init__()
+        assert isinstance(attn_layers, Encoder), "attention layers must be an Encoder"
+        assert (
+            image_size % patch_size == 0
+        ), "image dimensions must be divisible by the patch size"
+        dim = attn_layers.dim
+        num_patches = (image_size // patch_size) ** 2
+        patch_dim = 3 * patch_size**2
+
+        self.patch_size = patch_size
+
+        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
+        self.patch_to_embedding = nn.Linear(patch_dim, dim)
+        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
+        self.dropout = nn.Dropout(emb_dropout)
+
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+        self.mlp_head = (
+            FeedForward(dim, dim_out=num_classes, dropout=dropout)
+            if exists(num_classes)
+            else None
+        )
+
+    def forward(self, img, return_embeddings=False):
+        p = self.patch_size
+
+        x = rearrange(img, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
+        x = self.patch_to_embedding(x)
+        b, n, _ = x.shape
+
+        cls_tokens = repeat(self.cls_token, "() n d -> b n d", b=b)
+        x = torch.cat((cls_tokens, x), dim=1)
+        x = x + self.pos_embedding[:, : (n + 1)]
+        x = self.dropout(x)
+
+        x = self.attn_layers(x)
+        x = self.norm(x)
+
+        if not exists(self.mlp_head) or return_embeddings:
+            return x
+
+        return self.mlp_head(x[:, 0])
+
+
+class TransformerWrapper(nn.Module):
+    def __init__(
+        self,
+        *,
+        num_tokens,
+        max_seq_len,
+        attn_layers,
+        emb_dim=None,
+        max_mem_len=0.0,
+        shift_mem_down=0,
+        emb_dropout=0.0,
+        num_memory_tokens=None,
+        tie_embedding=False,
+        use_pos_emb=True,
+    ):
+        super().__init__()
+        assert isinstance(
+            attn_layers, AttentionLayers
+        ), "attention layers must be one of Encoder or Decoder"
+
+        dim = attn_layers.dim
+        emb_dim = default(emb_dim, dim)
+
+        self.max_seq_len = max_seq_len
+        self.max_mem_len = max_mem_len
+        self.shift_mem_down = shift_mem_down
+
+        self.token_emb = nn.Embedding(num_tokens, emb_dim)
+        self.pos_emb = (
+            AbsolutePositionalEmbedding(emb_dim, max_seq_len)
+            if (use_pos_emb and not attn_layers.has_pos_emb)
+            else always(0)
+        )
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.init_()
+
+        self.to_logits = (
+            nn.Linear(dim, num_tokens)
+            if not tie_embedding
+            else lambda t: t @ self.token_emb.weight.t()
+        )
+
+        # memory tokens (like [cls]) from Memory Transformers paper
+        num_memory_tokens = default(num_memory_tokens, 0)
+        self.num_memory_tokens = num_memory_tokens
+        if num_memory_tokens > 0:
+            self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+    def init_(self):
+        nn.init.kaiming_normal_(self.token_emb.weight)
+
+    def forward(
+        self,
+        x,
+        return_embeddings=False,
+        mask=None,
+        return_hiddens=False,
+        return_attn=False,
+        mems=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
+        x = self.token_emb(x)
+        x = x + self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x = self.project_emb(x)
+
+        if num_mem > 0:
+            mem = repeat(self.memory_tokens, "n d -> b n d", b=b)
+            x = torch.cat((mem, x), dim=1)
+
+            # auto-handle masking after appending memory tokens
+            if exists(mask):
+                mask = F.pad(mask, (num_mem, 0), value=True)
+
+        if self.shift_mem_down and exists(mems):
+            mems_l, mems_r = mems[: self.shift_mem_down], mems[self.shift_mem_down :]
+            mems = [*mems_r, *mems_l]
+
+        x, intermediates = self.attn_layers(
+            x, mask=mask, mems=mems, return_hiddens=True, **kwargs
+        )
+        x = self.norm(x)
+
+        mem, x = x[:, :num_mem], x[:, num_mem:]
+
+        out = self.to_logits(x) if not return_embeddings else x
+
+        if return_hiddens:
+            hiddens = intermediates.hiddens
+            return out, hiddens
+
+        res = [out]
+        if return_attn:
+            attn_maps = list(
+                map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates)
+            )
+            res.append(attn_maps)
+        if use_cache:
+            res.append(intermediates.past_key_values)
+
+        if len(res) > 1:
+            return tuple(res)
+        return res[0]
+
+
+class ContinuousTransformerWrapper(nn.Module):
+    def __init__(
+        self,
+        *,
+        max_seq_len,
+        attn_layers,
+        dim_in=None,
+        dim_out=None,
+        emb_dim=None,
+        emb_dropout=0.0,
+        use_pos_emb=True,
+    ):
+        super().__init__()
+        assert isinstance(
+            attn_layers, AttentionLayers
+        ), "attention layers must be one of Encoder or Decoder"
+
+        dim = attn_layers.dim
+
+        self.max_seq_len = max_seq_len
+
+        self.pos_emb = (
+            AbsolutePositionalEmbedding(dim, max_seq_len)
+            if (use_pos_emb and not attn_layers.has_pos_emb)
+            else always(0)
+        )
+        self.emb_dropout = nn.Dropout(emb_dropout)
+
+        self.project_in = nn.Linear(dim_in, dim) if exists(dim_in) else nn.Identity()
+
+        self.attn_layers = attn_layers
+        self.norm = nn.LayerNorm(dim)
+
+        self.project_out = nn.Linear(dim, dim_out) if exists(dim_out) else nn.Identity()
+
+    def forward(
+        self,
+        x,
+        return_embeddings=False,
+        mask=None,
+        return_attn=False,
+        mems=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        b, n, _, device = *x.shape, x.device
+
+        x = self.project_in(x)
+        x = x + self.pos_emb(x)
+        x = self.emb_dropout(x)
+
+        x, intermediates = self.attn_layers(
+            x, mask=mask, mems=mems, return_hiddens=True, **kwargs
+        )
+        x = self.norm(x)
+
+        out = self.project_out(x) if not return_embeddings else x
+
+        res = [out]
+        if return_attn:
+            attn_maps = list(
+                map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates)
+            )
+            res.append(attn_maps)
+        if use_cache:
+            res.append(intermediates.past_key_values)
+
+        if len(res) > 1:
+            return tuple(res)
+        return res[0]

tortoise/utils/audio.py

@@ -0,0 +1,238 @@
+import os
+from glob import glob
+from typing import Dict, List
+
+import librosa
+import numpy as np
+import torch
+import torchaudio
+from scipy.io.wavfile import read
+
+from tortoise.utils.stft import STFT
+
+BUILTIN_VOICES_DIR = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "../voices"
+)
+
+
+def load_wav_to_torch(full_path):
+    sampling_rate, data = read(full_path)
+    if data.dtype == np.int32:
+        norm_fix = 2**31
+    elif data.dtype == np.int16:
+        norm_fix = 2**15
+    elif data.dtype == np.float16 or data.dtype == np.float32:
+        norm_fix = 1.0
+    else:
+        raise NotImplementedError(f"Provided data dtype not supported: {data.dtype}")
+    return (torch.FloatTensor(data.astype(np.float32)) / norm_fix, sampling_rate)
+
+
+def check_audio(audio, audiopath: str):
+    # Check some assumptions about audio range. This should be automatically fixed in load_wav_to_torch, but might not be in some edge cases, where we should squawk.
+    # '2' is arbitrarily chosen since it seems like audio will often "overdrive" the [-1,1] bounds.
+    if torch.any(audio > 2) or not torch.any(audio < 0):
+        print(f"Error with {audiopath}. Max={audio.max()} min={audio.min()}")
+    audio.clip_(-1, 1)
+
+
+def read_audio_file(audiopath: str):
+    if audiopath[-4:] == ".wav":
+        audio, lsr = load_wav_to_torch(audiopath)
+    elif audiopath[-4:] == ".mp3":
+        audio, lsr = librosa.load(audiopath, sr=None)
+        audio = torch.FloatTensor(audio)
+    else:
+        assert False, f"Unsupported audio format provided: {audiopath[-4:]}"
+
+    # Remove any channel data.
+    if len(audio.shape) > 1:
+        if audio.shape[0] < 5:
+            audio = audio[0]
+        else:
+            assert audio.shape[1] < 5
+            audio = audio[:, 0]
+
+    return audio, lsr
+
+
+def load_required_audio(audiopath: str):
+    audio, lsr = read_audio_file(audiopath)
+
+    audios = [
+        torchaudio.functional.resample(audio, lsr, sampling_rate)
+        for sampling_rate in (22050, 24000)
+    ]
+    for audio in audios:
+        check_audio(audio, audiopath)
+
+    return [audio.unsqueeze(0) for audio in audios]
+
+
+def load_audio(audiopath, sampling_rate):
+    audio, lsr = read_audio_file(audiopath)
+
+    if lsr != sampling_rate:
+        audio = torchaudio.functional.resample(audio, lsr, sampling_rate)
+    check_audio(audio, audiopath)
+
+    return audio.unsqueeze(0)
+
+
+TACOTRON_MEL_MAX = 2.3143386840820312
+TACOTRON_MEL_MIN = -11.512925148010254
+
+
+def denormalize_tacotron_mel(norm_mel):
+    return ((norm_mel + 1) / 2) * (
+        TACOTRON_MEL_MAX - TACOTRON_MEL_MIN
+    ) + TACOTRON_MEL_MIN
+
+
+def normalize_tacotron_mel(mel):
+    return 2 * ((mel - TACOTRON_MEL_MIN) / (TACOTRON_MEL_MAX - TACOTRON_MEL_MIN)) - 1
+
+
+def dynamic_range_compression(x, C=1, clip_val=1e-5):
+    """
+    PARAMS
+    ------
+    C: compression factor
+    """
+    return torch.log(torch.clamp(x, min=clip_val) * C)
+
+
+def dynamic_range_decompression(x, C=1):
+    """
+    PARAMS
+    ------
+    C: compression factor used to compress
+    """
+    return torch.exp(x) / C
+
+
+def get_voices(extra_voice_dirs: List[str] = []):
+    dirs = [BUILTIN_VOICES_DIR] + extra_voice_dirs
+    voices: Dict[str, List[str]] = {}
+    for d in dirs:
+        subs = os.listdir(d)
+        for sub in subs:
+            subj = os.path.join(d, sub)
+            if os.path.isdir(subj):
+                voices[sub] = (
+                    list(glob(f"{subj}/*.wav"))
+                    + list(glob(f"{subj}/*.mp3"))
+                    + list(glob(f"{subj}/*.pth"))
+                )
+    return voices
+
+
+def load_voice(voice: str, extra_voice_dirs: List[str] = []):
+    if voice == "random":
+        return None, None
+
+    voices = get_voices(extra_voice_dirs)
+    paths = voices[voice]
+    if len(paths) == 1 and paths[0].endswith(".pth"):
+        return None, torch.load(paths[0])
+    else:
+        conds = []
+        for cond_path in paths:
+            c = load_required_audio(cond_path)
+            conds.append(c)
+        return conds, None
+
+
+def load_voices(voices: List[str], extra_voice_dirs: List[str] = []):
+    latents = []
+    clips = []
+    for voice in voices:
+        if voice == "random":
+            if len(voices) > 1:
+                print(
+                    "Cannot combine a random voice with a non-random voice. Just using a random voice."
+                )
+            return None, None
+        clip, latent = load_voice(voice, extra_voice_dirs)
+        if latent is None:
+            assert (
+                len(latents) == 0
+            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
+            clips.extend(clip)
+        elif clip is None:
+            assert (
+                len(clips) == 0
+            ), "Can only combine raw audio voices or latent voices, not both. Do it yourself if you want this."
+            latents.append(latent)
+    if len(latents) == 0:
+        return clips, None
+    else:
+        latents_0 = torch.stack([l[0] for l in latents], dim=0).mean(dim=0)
+        latents_1 = torch.stack([l[1] for l in latents], dim=0).mean(dim=0)
+        latents = (latents_0, latents_1)
+        return None, latents
+
+
+class TacotronSTFT(torch.nn.Module):
+    def __init__(
+        self,
+        filter_length=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=80,
+        sampling_rate=22050,
+        mel_fmin=0.0,
+        mel_fmax=8000.0,
+    ):
+        super(TacotronSTFT, self).__init__()
+        self.n_mel_channels = n_mel_channels
+        self.sampling_rate = sampling_rate
+        self.stft_fn = STFT(filter_length, hop_length, win_length)
+        from librosa.filters import mel as librosa_mel_fn
+
+        mel_basis = librosa_mel_fn(
+            sr=sampling_rate,
+            n_fft=filter_length,
+            n_mels=n_mel_channels,
+            fmin=mel_fmin,
+            fmax=mel_fmax,
+        )
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+
+    def spectral_normalize(self, magnitudes):
+        output = dynamic_range_compression(magnitudes)
+        return output
+
+    def spectral_de_normalize(self, magnitudes):
+        output = dynamic_range_decompression(magnitudes)
+        return output
+
+    def mel_spectrogram(self, y):
+        """Computes mel-spectrograms from a batch of waves
+        PARAMS
+        ------
+        y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1]
+
+        RETURNS
+        -------
+        mel_output: torch.FloatTensor of shape (B, n_mel_channels, T)
+        """
+        assert torch.min(y.data) >= -10
+        assert torch.max(y.data) <= 10
+        y = torch.clip(y, min=-1, max=1)
+
+        magnitudes, phases = self.stft_fn.transform(y)
+        magnitudes = magnitudes.data
+        mel_output = torch.matmul(self.mel_basis, magnitudes)
+        mel_output = self.spectral_normalize(mel_output)
+        return mel_output
+
+
+def wav_to_univnet_mel(wav, do_normalization=False, device="cuda"):
+    stft = TacotronSTFT(1024, 256, 1024, 100, 24000, 0, 12000)
+    stft = stft.to(device)
+    mel = stft.mel_spectrogram(wav)
+    if do_normalization:
+        mel = normalize_tacotron_mel(mel)
+    return mel

tortoise/utils/diffusion.py

@@ -0,0 +1,1469 @@
+"""
+This is an almost carbon copy of gaussian_diffusion.py from OpenAI's ImprovedDiffusion repo, which itself:
+
+This code started out as a PyTorch port of Ho et al's diffusion models:
+https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py
+
+Docstrings have been added, as well as DDIM sampling and a new collection of beta schedules.
+"""
+# AGPL: a notification must be added stating that changes have been made to that file.
+
+import enum
+import math
+
+import numpy as np
+import torch
+import torch as th
+from k_diffusion.sampling import sample_dpmpp_2m, sample_euler_ancestral
+from tqdm import tqdm
+
+from tortoise.dpm_solver_pytorch import DPM_Solver, NoiseScheduleVP, model_wrapper
+
+K_DIFFUSION_SAMPLERS = {"k_euler_a": sample_euler_ancestral, "dpm++2m": sample_dpmpp_2m}
+SAMPLERS = ["dpm++2m", "p", "ddim"]
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = "previous_x"  # the model predicts x_{t-1}
+    START_X = "start_x"  # the model predicts x_0
+    EPSILON = "epsilon"  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = "learned"
+    FIXED_SMALL = "fixed_small"
+    FIXED_LARGE = "fixed_large"
+    LEARNED_RANGE = "learned_range"
+
+
+class LossType(enum.Enum):
+    MSE = "mse"  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        "rescaled_mse"  # use raw MSE loss (with RESCALED_KL when learning variances)
+    )
+    KL = "kl"  # use the variational lower-bound
+    RESCALED_KL = "rescaled_kl"  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here, and then adapted over time to further experimentation.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    :param model_mean_type: a ModelMeanType determining what the model outputs.
+    :param model_var_type: a ModelVarType determining how variance is output.
+    :param loss_type: a LossType determining the loss function to use.
+    :param rescale_timesteps: if True, pass floating point timesteps into the
+                              model so that they are always scaled like in the
+                              original paper (0 to 1000).
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type,
+        rescale_timesteps=False,  # this is generally False
+        conditioning_free=False,
+        conditioning_free_k=1,
+        ramp_conditioning_free=True,
+        sampler="ddim",
+    ):
+        self.sampler = sampler
+        self.model_mean_type = ModelMeanType(model_mean_type)
+        self.model_var_type = ModelVarType(model_var_type)
+        self.loss_type = LossType(loss_type)
+        self.rescale_timesteps = rescale_timesteps
+        self.conditioning_free = conditioning_free
+        self.conditioning_free_k = conditioning_free_k
+        self.ramp_conditioning_free = ramp_conditioning_free
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # log calculation clipped because the posterior variance is 0 at the
+        # beginning of the diffusion chain.
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev)
+            * np.sqrt(alphas)
+            / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        )
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(
+            self.log_one_minus_alphas_cumprod, t, x_start.shape
+        )
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+            * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        assert self.model_var_type == ModelVarType.LEARNED_RANGE
+        assert self.model_mean_type == ModelMeanType.EPSILON
+        assert denoised_fn is None
+        assert clip_denoised is True
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, self._scale_timesteps(t), **model_kwargs)
+        if self.conditioning_free:
+            model_output_no_conditioning = model(
+                x, self._scale_timesteps(t), conditioning_free=True, **model_kwargs
+            )
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, C * 2, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.conditioning_free:
+                model_output_no_conditioning, _ = th.split(
+                    model_output_no_conditioning, C, dim=1
+                )
+            if self.model_var_type == ModelVarType.LEARNED:
+                assert False
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(
+                    self.posterior_log_variance_clipped, t, x.shape
+                )
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            assert False
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        if self.conditioning_free:
+            if self.ramp_conditioning_free:
+                assert t.shape[0] == 1  # This should only be used in inference.
+                cfk = self.conditioning_free_k * (
+                    1 - self._scale_timesteps(t)[0].item() / self.num_timesteps
+                )
+            else:
+                cfk = self.conditioning_free_k
+            model_output = (1 + cfk) * model_output - cfk * model_output_no_conditioning
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                assert False
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            assert False
+            return x
+
+        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+            assert False
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
+            if self.model_mean_type == ModelMeanType.START_X:
+                assert False
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(
+                x_start=pred_xstart, x_t=x, t=t
+            )
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert (
+            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        )
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _scale_timesteps(self, t):
+        if self.rescale_timesteps:
+            return t.float() * (1000.0 / self.num_timesteps)
+        return t
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
+        new_mean = (
+            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        )
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
+            x, self._scale_timesteps(t), **model_kwargs
+        )
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(
+            x_start=out["pred_xstart"], x_t=x, t=t
+        )
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(
+                cond_fn, out, x, t, model_kwargs=model_kwargs
+            )
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def k_diffusion_sample_loop(
+        self,
+        k_sampler,
+        pbar,
+        model,
+        shape,
+        noise=None,  # all given
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        device=None,  # ALL UNUSED
+        model_kwargs=None,  # {'precomputed_aligned_embeddings': precomputed_embeddings},
+        progress=False,  # unused as well
+    ):
+        assert isinstance(model_kwargs, dict)
+        if device is None:
+            device = next(model.parameters()).device
+        s_in = noise.new_ones([noise.shape[0]])
+
+        def model_split(*args, **kwargs):
+            model_output = model(*args, **kwargs)
+            model_epsilon, model_var = th.split(
+                model_output, model_output.shape[1] // 2, dim=1
+            )
+            return model_epsilon, model_var
+
+        #
+        """
+        print(self.betas)
+        print(th.tensor(self.betas))
+        noise_schedule = NoiseScheduleVP(schedule='discrete', betas=th.tensor(self.betas))
+        """
+        noise_schedule = NoiseScheduleVP(
+            schedule="linear", continuous_beta_0=0.1 / 4, continuous_beta_1=20.0 / 4
+        )
+
+        def model_fn_prewrap(x, t, *args, **kwargs):
+            """
+                x_in = torch.cat([x] * 2)
+                t_in = torch.cat([t_continuous] * 2)
+                c_in = torch.cat([unconditional_condition, condition])
+                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
+            print(t)
+            print(self.timestep_map)
+            exit()
+            """
+            """
+            model_output = model(x, self._scale_timesteps(t*4000), **model_kwargs)
+            out = self.p_mean_variance(model, x, t*4000, model_kwargs=model_kwargs)
+            return out['pred_xstart']
+            """
+            x, _ = x.chunk(2)
+            t, _ = (t * 1000).chunk(2)
+            res = torch.cat(
+                [
+                    model_split(x, t, conditioning_free=True, **model_kwargs)[0],
+                    model_split(x, t, **model_kwargs)[0],
+                ]
+            )
+            pbar.update(1)
+            return res
+
+        model_fn = model_wrapper(
+            model_fn_prewrap,
+            noise_schedule,
+            model_type="noise",  # "noise" or "x_start" or "v" or "score"
+            model_kwargs=model_kwargs,
+            guidance_type="classifier-free",
+            condition=th.Tensor(1),
+            unconditional_condition=th.Tensor(1),
+            guidance_scale=self.conditioning_free_k,
+        )
+        """
+        model_fn = model_wrapper(
+            model_fn_prewrap,
+            noise_schedule,
+            model_type='x_start',
+            model_kwargs={}
+        )
+        #
+        dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver")
+        x_sample = dpm_solver.sample(
+            noise,
+            steps=20,
+            order=3,
+            skip_type="time_uniform",
+            method="singlestep",
+        )
+        """
+        dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+        x_sample = dpm_solver.sample(
+            noise,
+            steps=self.num_timesteps,
+            order=2,
+            skip_type="time_uniform",
+            method="multistep",
+        )
+        #'''
+        return x_sample
+
+        # HF DIFFUSION ATTEMPT
+        """
+        from .hf_diffusion import EulerAncestralDiscreteScheduler
+        Scheduler = EulerAncestralDiscreteScheduler()
+        Scheduler.set_timesteps(100)
+        for timestep in Scheduler.timesteps:
+            noise_input = Scheduler.scale_model_input(noise, timestep)
+            ts = s_in * timestep
+            model_output = model(noise_input, ts, **model_kwargs)
+            model_epsilon, _model_var = th.split(model_output, model_output.shape[1]//2, dim=1)
+            noise, _x0 = Scheduler.step(model_epsilon, timestep, noise)
+        return noise
+        """
+
+        # KARRAS DIFFUSION ATTEMPT
+        """
+        TRAINED_DIFFUSION_STEPS = 4000 # HARDCODED
+        ratio = TRAINED_DIFFUSION_STEPS/14.5
+        def call_model(*args, **kwargs):
+            model_output = model(*args, **kwargs)
+            model_output, model_var_values = th.split(model_output, model_output.shape[1]//2, dim=1)
+            return model_output
+        print(get_sigmas_karras(self.num_timesteps, sigma_min=0.0, sigma_max=4000, device=device))
+        exit()
+        sigmas = get_sigmas_karras(self.num_timesteps, sigma_min=0.03, sigma_max=14.5, device=device)
+        return k_sampler(call_model, noise, sigmas, extra_args=model_kwargs, disable=not progress)
+        '''
+        sigmas = get_sigmas_karras(self.num_timesteps, sigma_min=0.03, sigma_max=14.5, device=device)
+        step = 0 # LMAO
+        global_sigmas = None
+        #
+        def fakemodel(x, t, **model_kwargs):
+            print(t,global_sigmas*ratio)
+            return model(x, t, **model_kwargs)
+        def denoised(x, sigmas, **extra_args):
+            t = th.tensor([self.num_timesteps-step-1] * shape[0], device=device)
+            nonlocal global_sigmas
+            global_sigmas = sigmas
+            with th.no_grad():
+                out = self.p_sample(
+                    fakemodel,
+                    x,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                return out["sample"]
+        def callback(d):
+            nonlocal step
+            step += 1
+
+        return k_sampler(denoised, noise, sigmas, extra_args=model_kwargs, callback=callback, disable=not progress)
+        '''
+        """
+
+    def sample_loop(self, *args, **kwargs):
+        s = self.sampler
+        if s == "p":
+            return self.p_sample_loop(*args, **kwargs)
+        elif s == "ddim":
+            return self.ddim_sample_loop(*args, **kwargs)
+        elif s == "dpm++2m":
+            if self.conditioning_free is not True:
+                raise RuntimeError("cond_free must be true")
+            with tqdm(total=self.num_timesteps) as pbar:
+                return self.k_diffusion_sample_loop(
+                    K_DIFFUSION_SAMPLERS[s], pbar, *args, **kwargs
+                )
+        else:
+            raise RuntimeError("sampler not impl")
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        for i in tqdm(indices, disable=not progress):
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma**2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_next)
+            + th.sqrt(1 - alpha_bar_next) * eps
+        )
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices, disable=not progress)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+    ):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            # TODO: support multiple model outputs for this mode.
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_outputs = model(x_t, self._scale_timesteps(t), **model_kwargs)
+            if isinstance(model_outputs, tuple):
+                model_output = model_outputs[0]
+                terms["extra_outputs"] = model_outputs[1:]
+            else:
+                model_output = model_outputs
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+                target = self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[
+                    0
+                ]
+                x_start_pred = torch.zeros(x_start)  # Not supported.
+            elif self.model_mean_type == ModelMeanType.START_X:
+                target = x_start
+                x_start_pred = model_output
+            elif self.model_mean_type == ModelMeanType.EPSILON:
+                target = noise
+                x_start_pred = self._predict_xstart_from_eps(x_t, t, model_output)
+            else:
+                raise NotImplementedError(self.model_mean_type)
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            terms["x_start_predicted"] = x_start_pred
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def autoregressive_training_losses(
+        self,
+        model,
+        x_start,
+        t,
+        model_output_keys,
+        gd_out_key,
+        model_kwargs=None,
+        noise=None,
+    ):
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+        terms = {}
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            assert False  # not currently supported for this type of diffusion.
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_outputs = model(
+                x_t, x_start, self._scale_timesteps(t), **model_kwargs
+            )
+            terms.update({k: o for k, o in zip(model_output_keys, model_outputs)})
+            model_output = terms[gd_out_key]
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C, 2, *x_t.shape[2:])
+                model_output, model_var_values = (
+                    model_output[:, :, 0],
+                    model_output[:, :, 1],
+                )
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+                target = self.q_posterior_mean_variance(x_start=x_start, x_t=x_t, t=t)[
+                    0
+                ]
+                x_start_pred = torch.zeros(x_start)  # Not supported.
+            elif self.model_mean_type == ModelMeanType.START_X:
+                target = x_start
+                x_start_pred = model_output
+            elif self.model_mean_type == ModelMeanType.EPSILON:
+                target = noise
+                x_start_pred = self._predict_xstart_from_eps(x_t, t, model_output)
+            else:
+                raise NotImplementedError(self.model_mean_type)
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2)
+            terms["x_start_predicted"] = x_start_pred
+            if "vb" in terms:
+                terms["loss"] = terms["mse"] + terms["vb"]
+            else:
+                terms["loss"] = terms["mse"]
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+class SpacedDiffusion(GaussianDiffusion):
+    """
+    A diffusion process which can skip steps in a base diffusion process.
+
+    :param use_timesteps: a collection (sequence or set) of timesteps from the
+                          original diffusion process to retain.
+    :param kwargs: the kwargs to create the base diffusion process.
+    """
+
+    def __init__(self, use_timesteps, **kwargs):
+        self.use_timesteps = set(use_timesteps)
+        self.timestep_map = []
+        self.original_num_steps = len(kwargs["betas"])
+
+        base_diffusion = GaussianDiffusion(**kwargs)  # pylint: disable=missing-kwoa
+        last_alpha_cumprod = 1.0
+        new_betas = []
+        for i, alpha_cumprod in enumerate(base_diffusion.alphas_cumprod):
+            if i in self.use_timesteps:
+                new_betas.append(1 - alpha_cumprod / last_alpha_cumprod)
+                last_alpha_cumprod = alpha_cumprod
+                self.timestep_map.append(i)
+        kwargs["betas"] = np.array(new_betas)
+        super().__init__(**kwargs)
+
+    def p_mean_variance(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().p_mean_variance(self._wrap_model(model), *args, **kwargs)
+
+    def training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().training_losses(self._wrap_model(model), *args, **kwargs)
+
+    def autoregressive_training_losses(
+        self, model, *args, **kwargs
+    ):  # pylint: disable=signature-differs
+        return super().autoregressive_training_losses(
+            self._wrap_model(model, True), *args, **kwargs
+        )
+
+    def condition_mean(self, cond_fn, *args, **kwargs):
+        return super().condition_mean(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def condition_score(self, cond_fn, *args, **kwargs):
+        return super().condition_score(self._wrap_model(cond_fn), *args, **kwargs)
+
+    def _wrap_model(self, model, autoregressive=False):
+        if isinstance(model, _WrappedModel) or isinstance(
+            model, _WrappedAutoregressiveModel
+        ):
+            return model
+        mod = _WrappedAutoregressiveModel if autoregressive else _WrappedModel
+        return mod(
+            model, self.timestep_map, self.rescale_timesteps, self.original_num_steps
+        )
+
+    def _scale_timesteps(self, t):
+        # Scaling is done by the wrapped model.
+        return t
+
+
+def space_timesteps(num_timesteps, section_counts):
+    """
+    Create a list of timesteps to use from an original diffusion process,
+    given the number of timesteps we want to take from equally-sized portions
+    of the original process.
+
+    For example, if there's 300 timesteps and the section counts are [10,15,20]
+    then the first 100 timesteps are strided to be 10 timesteps, the second 100
+    are strided to be 15 timesteps, and the final 100 are strided to be 20.
+
+    If the stride is a string starting with "ddim", then the fixed striding
+    from the DDIM paper is used, and only one section is allowed.
+
+    :param num_timesteps: the number of diffusion steps in the original
+                          process to divide up.
+    :param section_counts: either a list of numbers, or a string containing
+                           comma-separated numbers, indicating the step count
+                           per section. As a special case, use "ddimN" where N
+                           is a number of steps to use the striding from the
+                           DDIM paper.
+    :return: a set of diffusion steps from the original process to use.
+    """
+    if isinstance(section_counts, str):
+        if section_counts.startswith("ddim"):
+            desired_count = int(section_counts[len("ddim") :])
+            for i in range(1, num_timesteps):
+                if len(range(0, num_timesteps, i)) == desired_count:
+                    return set(range(0, num_timesteps, i))
+            raise ValueError(
+                f"cannot create exactly {num_timesteps} steps with an integer stride"
+            )
+        section_counts = [int(x) for x in section_counts.split(",")]
+    size_per = num_timesteps // len(section_counts)
+    extra = num_timesteps % len(section_counts)
+    start_idx = 0
+    all_steps = []
+    for i, section_count in enumerate(section_counts):
+        size = size_per + (1 if i < extra else 0)
+        if size < section_count:
+            raise ValueError(
+                f"cannot divide section of {size} steps into {section_count}"
+            )
+        if section_count <= 1:
+            frac_stride = 1
+        else:
+            frac_stride = (size - 1) / (section_count - 1)
+        cur_idx = 0.0
+        taken_steps = []
+        for _ in range(section_count):
+            taken_steps.append(start_idx + round(cur_idx))
+            cur_idx += frac_stride
+        all_steps += taken_steps
+        start_idx += size
+    return set(all_steps)
+
+
+class _WrappedModel:
+    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        if self.rescale_timesteps:
+            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, new_ts, **kwargs)
+
+
+class _WrappedAutoregressiveModel:
+    def __init__(self, model, timestep_map, rescale_timesteps, original_num_steps):
+        self.model = model
+        self.timestep_map = timestep_map
+        self.rescale_timesteps = rescale_timesteps
+        self.original_num_steps = original_num_steps
+
+    def __call__(self, x, x0, ts, **kwargs):
+        map_tensor = th.tensor(self.timestep_map, device=ts.device, dtype=ts.dtype)
+        new_ts = map_tensor[ts]
+        if self.rescale_timesteps:
+            new_ts = new_ts.float() * (1000.0 / self.original_num_steps)
+        return self.model(x, x0, new_ts, **kwargs)
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)

tortoise/utils/stft.py

@@ -0,0 +1,215 @@
+"""
+BSD 3-Clause License
+
+Copyright (c) 2017, Prem Seetharaman
+All rights reserved.
+
+* Redistribution and use in source and binary forms, with or without
+  modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice,
+  this list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice, this
+  list of conditions and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from this
+  software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+"""
+
+import librosa.util as librosa_util
+import numpy as np
+import torch
+import torch.nn.functional as F
+from librosa.util import pad_center, tiny
+from scipy.signal import get_window
+from torch.autograd import Variable
+
+
+def window_sumsquare(
+    window,
+    n_frames,
+    hop_length=200,
+    win_length=800,
+    n_fft=800,
+    dtype=np.float32,
+    norm=None,
+):
+    """
+    # from librosa 0.6
+    Compute the sum-square envelope of a window function at a given hop length.
+
+    This is used to estimate modulation effects induced by windowing
+    observations in short-time fourier transforms.
+
+    Parameters
+    ----------
+    window : string, tuple, number, callable, or list-like
+        Window specification, as in `get_window`
+
+    n_frames : int > 0
+        The number of analysis frames
+
+    hop_length : int > 0
+        The number of samples to advance between frames
+
+    win_length : [optional]
+        The length of the window function.  By default, this matches `n_fft`.
+
+    n_fft : int > 0
+        The length of each analysis frame.
+
+    dtype : np.dtype
+        The data type of the output
+
+    Returns
+    -------
+    wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
+        The sum-squared envelope of the window function
+    """
+    if win_length is None:
+        win_length = n_fft
+
+    n = n_fft + hop_length * (n_frames - 1)
+    x = np.zeros(n, dtype=dtype)
+
+    # Compute the squared window at the desired length
+    win_sq = get_window(window, win_length, fftbins=True)
+    win_sq = librosa_util.normalize(win_sq, norm=norm) ** 2
+    win_sq = librosa_util.pad_center(win_sq, n_fft)
+
+    # Fill the envelope
+    for i in range(n_frames):
+        sample = i * hop_length
+        x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
+    return x
+
+
+class STFT(torch.nn.Module):
+    """adapted from Prem Seetharaman's https://github.com/pseeth/pytorch-stft"""
+
+    def __init__(
+        self, filter_length=800, hop_length=200, win_length=800, window="hann"
+    ):
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.window = window
+        self.forward_transform = None
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack(
+            [np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
+        )
+
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :]
+        )
+
+        if window is not None:
+            assert filter_length >= win_length
+            # get window and zero center pad it to filter_length
+            fft_window = get_window(window, win_length, fftbins=True)
+            fft_window = pad_center(fft_window, size=filter_length)
+            fft_window = torch.from_numpy(fft_window).float()
+
+            # window the bases
+            forward_basis *= fft_window
+            inverse_basis *= fft_window
+
+        self.register_buffer("forward_basis", forward_basis.float())
+        self.register_buffer("inverse_basis", inverse_basis.float())
+
+    def transform(self, input_data):
+        num_batches = input_data.size(0)
+        num_samples = input_data.size(1)
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (int(self.filter_length / 2), int(self.filter_length / 2), 0, 0),
+            mode="reflect",
+        )
+        input_data = input_data.squeeze(1)
+
+        forward_transform = F.conv1d(
+            input_data,
+            Variable(self.forward_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part**2 + imag_part**2)
+        phase = torch.autograd.Variable(torch.atan2(imag_part.data, real_part.data))
+
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        recombine_magnitude_phase = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
+        )
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            Variable(self.inverse_basis, requires_grad=False),
+            stride=self.hop_length,
+            padding=0,
+        )
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window,
+                magnitude.size(-1),
+                hop_length=self.hop_length,
+                win_length=self.win_length,
+                n_fft=self.filter_length,
+                dtype=np.float32,
+            )
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0]
+            )
+            window_sum = torch.autograd.Variable(
+                torch.from_numpy(window_sum), requires_grad=False
+            )
+            window_sum = window_sum.cuda() if magnitude.is_cuda else window_sum
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
+                approx_nonzero_indices
+            ]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[:, :, int(self.filter_length / 2) :]
+        inverse_transform = inverse_transform[:, :, : -int(self.filter_length / 2) :]
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction

tortoise/utils/text.py

@@ -0,0 +1,144 @@
+import re
+
+
+def split_and_recombine_text(text, desired_length=200, max_length=300):
+    """Split text it into chunks of a desired length trying to keep sentences intact."""
+    # normalize text, remove redundant whitespace and convert non-ascii quotes to ascii
+    text = re.sub(r"\n\n+", "\n", text)
+    text = re.sub(r"\s+", " ", text)
+    text = re.sub(r"[“”]", '"', text)
+
+    rv = []
+    in_quote = False
+    current = ""
+    split_pos = []
+    pos = -1
+    end_pos = len(text) - 1
+
+    def seek(delta):
+        nonlocal pos, in_quote, current
+        is_neg = delta < 0
+        for _ in range(abs(delta)):
+            if is_neg:
+                pos -= 1
+                current = current[:-1]
+            else:
+                pos += 1
+                current += text[pos]
+            if text[pos] == '"':
+                in_quote = not in_quote
+        return text[pos]
+
+    def peek(delta):
+        p = pos + delta
+        return text[p] if p < end_pos and p >= 0 else ""
+
+    def commit():
+        nonlocal rv, current, split_pos
+        rv.append(current)
+        current = ""
+        split_pos = []
+
+    while pos < end_pos:
+        c = seek(1)
+        # do we need to force a split?
+        if len(current) >= max_length:
+            if len(split_pos) > 0 and len(current) > (desired_length / 2):
+                # we have at least one sentence and we are over half the desired length, seek back to the last split
+                d = pos - split_pos[-1]
+                seek(-d)
+            else:
+                # no full sentences, seek back until we are not in the middle of a word and split there
+                while c not in "!?.\n " and pos > 0 and len(current) > desired_length:
+                    c = seek(-1)
+            commit()
+        # check for sentence boundaries
+        elif not in_quote and (c in "!?\n" or (c == "." and peek(1) in "\n ")):
+            # seek forward if we have consecutive boundary markers but still within the max length
+            while (
+                pos < len(text) - 1 and len(current) < max_length and peek(1) in "!?."
+            ):
+                c = seek(1)
+            split_pos.append(pos)
+            if len(current) >= desired_length:
+                commit()
+        # treat end of quote as a boundary if its followed by a space or newline
+        elif in_quote and peek(1) == '"' and peek(2) in "\n ":
+            seek(2)
+            split_pos.append(pos)
+    rv.append(current)
+
+    # clean up, remove lines with only whitespace or punctuation
+    rv = [s.strip() for s in rv]
+    rv = [s for s in rv if len(s) > 0 and not re.match(r"^[\s\.,;:!?]*$", s)]
+
+    return rv
+
+
+if __name__ == "__main__":
+    import os
+    import unittest
+
+    class Test(unittest.TestCase):
+        def test_split_and_recombine_text(self):
+            text = """
+            This is a sample sentence.
+            This is another sample sentence.
+            This is a longer sample sentence that should force a split inthemiddlebutinotinthislongword.
+            "Don't split my quote... please"
+            """
+            self.assertEqual(
+                split_and_recombine_text(text, desired_length=20, max_length=40),
+                [
+                    "This is a sample sentence.",
+                    "This is another sample sentence.",
+                    "This is a longer sample sentence that",
+                    "should force a split",
+                    "inthemiddlebutinotinthislongword.",
+                    '"Don\'t split my quote... please"',
+                ],
+            )
+
+        def test_split_and_recombine_text_2(self):
+            text = """
+            When you are really angry sometimes you use consecutive exclamation marks!!!!!! Is this a good thing to do?!?!?!
+            I don't know but we should handle this situation..........................
+            """
+            self.assertEqual(
+                split_and_recombine_text(text, desired_length=30, max_length=50),
+                [
+                    "When you are really angry sometimes you use",
+                    "consecutive exclamation marks!!!!!!",
+                    "Is this a good thing to do?!?!?!",
+                    "I don't know but we should handle this situation.",
+                ],
+            )
+
+        def test_split_and_recombine_text_3(self):
+            text_src = os.path.join(
+                os.path.dirname(__file__), "../data/riding_hood.txt"
+            )
+            with open(text_src, "r") as f:
+                text = f.read()
+            self.assertEqual(
+                split_and_recombine_text(text),
+                [
+                    "Once upon a time there lived in a certain village a little country girl, the prettiest creature who was ever seen. Her mother was excessively fond of her; and her grandmother doted on her still more. This good woman had a little red riding hood made for her.",
+                    'It suited the girl so extremely well that everybody called her Little Red Riding Hood. One day her mother, having made some cakes, said to her, "Go, my dear, and see how your grandmother is doing, for I hear she has been very ill. Take her a cake, and this little pot of butter."',
+                    "Little Red Riding Hood set out immediately to go to her grandmother, who lived in another village. As she was going through the wood, she met with a wolf, who had a very great mind to eat her up, but he dared not, because of some woodcutters working nearby in the forest.",
+                    'He asked her where she was going. The poor child, who did not know that it was dangerous to stay and talk to a wolf, said to him, "I am going to see my grandmother and carry her a cake and a little pot of butter from my mother." "Does she live far off?" said the wolf "Oh I say,"',
+                    'answered Little Red Riding Hood; "it is beyond that mill you see there, at the first house in the village." "Well," said the wolf, "and I\'ll go and see her too. I\'ll go this way and go you that, and we shall see who will be there first."',
+                    "The wolf ran as fast as he could, taking the shortest path, and the little girl took a roundabout way, entertaining herself by gathering nuts, running after butterflies, and gathering bouquets of little flowers.",
+                    'It was not long before the wolf arrived at the old woman\'s house. He knocked at the door: tap, tap. "Who\'s there?" "Your grandchild, Little Red Riding Hood," replied the wolf, counterfeiting her voice; "who has brought you a cake and a little pot of butter sent you by mother."',
+                    'The good grandmother, who was in bed, because she was somewhat ill, cried out, "Pull the bobbin, and the latch will go up."',
+                    "The wolf pulled the bobbin, and the door opened, and then he immediately fell upon the good woman and ate her up in a moment, for it been more than three days since he had eaten.",
+                    "He then shut the door and got into the grandmother's bed, expecting Little Red Riding Hood, who came some time afterwards and knocked at the door: tap, tap. \"Who's there?\"",
+                    'Little Red Riding Hood, hearing the big voice of the wolf, was at first afraid; but believing her grandmother had a cold and was hoarse, answered, "It is your grandchild Little Red Riding Hood, who has brought you a cake and a little pot of butter mother sends you."',
+                    'The wolf cried out to her, softening his voice as much as he could, "Pull the bobbin, and the latch will go up." Little Red Riding Hood pulled the bobbin, and the door opened.',
+                    'The wolf, seeing her come in, said to her, hiding himself under the bedclothes, "Put the cake and the little pot of butter upon the stool, and come get into bed with me." Little Red Riding Hood took off her clothes and got into bed.',
+                    'She was greatly amazed to see how her grandmother looked in her nightclothes, and said to her, "Grandmother, what big arms you have!" "All the better to hug you with, my dear." "Grandmother, what big legs you have!" "All the better to run with, my child." "Grandmother, what big ears you have!"',
+                    '"All the better to hear with, my child." "Grandmother, what big eyes you have!" "All the better to see with, my child." "Grandmother, what big teeth you have got!" "All the better to eat you up with." And, saying these words, this wicked wolf fell upon Little Red Riding Hood, and ate her all up.',
+                ],
+            )
+
+    unittest.main()

tortoise/utils/tokenizer.py

@@ -0,0 +1,201 @@
+import os
+import re
+
+import inflect
+import torch
+from tokenizers import Tokenizer
+
+# Regular expression matching whitespace:
+from unidecode import unidecode
+
+_whitespace_re = re.compile(r"\s+")
+
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [
+    (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+    for x in [
+        ("mrs", "misess"),
+        ("mr", "mister"),
+        ("dr", "doctor"),
+        ("st", "saint"),
+        ("co", "company"),
+        ("jr", "junior"),
+        ("maj", "major"),
+        ("gen", "general"),
+        ("drs", "doctors"),
+        ("rev", "reverend"),
+        ("lt", "lieutenant"),
+        ("hon", "honorable"),
+        ("sgt", "sergeant"),
+        ("capt", "captain"),
+        ("esq", "esquire"),
+        ("ltd", "limited"),
+        ("col", "colonel"),
+        ("ft", "fort"),
+    ]
+]
+
+
+def expand_abbreviations(text):
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text
+
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
+_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
+_pounds_re = re.compile(r"£([0-9\,]*[0-9]+)")
+_dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)")
+_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
+_number_re = re.compile(r"[0-9]+")
+
+
+def _remove_commas(m):
+    return m.group(1).replace(",", "")
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace(".", " point ")
+
+
+def _expand_dollars(m):
+    match = m.group(1)
+    parts = match.split(".")
+    if len(parts) > 2:
+        return match + " dollars"  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s, %s %s" % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = "dollar" if dollars == 1 else "dollars"
+        return "%s %s" % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = "cent" if cents == 1 else "cents"
+        return "%s %s" % (cents, cent_unit)
+    else:
+        return "zero dollars"
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return "two thousand"
+        elif num > 2000 and num < 2010:
+            return "two thousand " + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + " hundred"
+        else:
+            return _inflect.number_to_words(
+                num, andword="", zero="oh", group=2
+            ).replace(", ", " ")
+    else:
+        return _inflect.number_to_words(num, andword="")
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r"\1 pounds", text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text
+
+
+def expand_numbers(text):
+    return normalize_numbers(text)
+
+
+def lowercase(text):
+    return text.lower()
+
+
+def collapse_whitespace(text):
+    return re.sub(_whitespace_re, " ", text)
+
+
+def convert_to_ascii(text):
+    return unidecode(text)
+
+
+def basic_cleaners(text):
+    """Basic pipeline that lowercases and collapses whitespace without transliteration."""
+    text = lowercase(text)
+    text = collapse_whitespace(text)
+    return text
+
+
+def transliteration_cleaners(text):
+    """Pipeline for non-English text that transliterates to ASCII."""
+    text = convert_to_ascii(text)
+    text = lowercase(text)
+    text = collapse_whitespace(text)
+    return text
+
+
+def english_cleaners(text):
+    """Pipeline for English text, including number and abbreviation expansion."""
+    text = convert_to_ascii(text)
+    text = lowercase(text)
+    text = expand_numbers(text)
+    text = expand_abbreviations(text)
+    text = collapse_whitespace(text)
+    text = text.replace('"', "")
+    return text
+
+
+def lev_distance(s1, s2):
+    if len(s1) > len(s2):
+        s1, s2 = s2, s1
+
+    distances = range(len(s1) + 1)
+    for i2, c2 in enumerate(s2):
+        distances_ = [i2 + 1]
+        for i1, c1 in enumerate(s1):
+            if c1 == c2:
+                distances_.append(distances[i1])
+            else:
+                distances_.append(
+                    1 + min((distances[i1], distances[i1 + 1], distances_[-1]))
+                )
+        distances = distances_
+    return distances[-1]
+
+
+DEFAULT_VOCAB_FILE = os.path.join(
+    os.path.dirname(os.path.realpath(__file__)), "../data/tokenizer.json"
+)
+
+
+class VoiceBpeTokenizer:
+    def __init__(self, vocab_file=DEFAULT_VOCAB_FILE):
+        if vocab_file is not None:
+            self.tokenizer = Tokenizer.from_file(vocab_file)
+
+    def preprocess_text(self, txt):
+        txt = english_cleaners(txt)
+        return txt
+
+    def encode(self, txt):
+        txt = self.preprocess_text(txt)
+        txt = txt.replace(" ", "[SPACE]")
+        return self.tokenizer.encode(txt).ids
+
+    def decode(self, seq):
+        if isinstance(seq, torch.Tensor):
+            seq = seq.cpu().numpy()
+        txt = self.tokenizer.decode(seq, skip_special_tokens=False).replace(" ", "")
+        txt = txt.replace("[SPACE]", " ")
+        txt = txt.replace("[STOP]", "")
+        txt = txt.replace("[UNK]", "")
+        return txt

tortoise/utils/typical_sampling.py

@@ -0,0 +1,44 @@
+import torch
+from transformers import LogitsWarper
+
+
+class TypicalLogitsWarper(LogitsWarper):
+    def __init__(
+        self,
+        mass: float = 0.9,
+        filter_value: float = -float("Inf"),
+        min_tokens_to_keep: int = 1,
+    ):
+        self.filter_value = filter_value
+        self.mass = mass
+        self.min_tokens_to_keep = min_tokens_to_keep
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        # calculate entropy
+        normalized = torch.nn.functional.log_softmax(scores, dim=-1)
+        p = torch.exp(normalized)
+        ent = -(normalized * p).nansum(-1, keepdim=True)
+
+        # shift and sort
+        shifted_scores = torch.abs((-normalized) - ent)
+        sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
+        sorted_logits = scores.gather(-1, sorted_indices)
+        cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+
+        # Remove tokens with cumulative mass above the threshold
+        last_ind = (cumulative_probs < self.mass).sum(dim=1)
+        last_ind[last_ind < 0] = 0
+        sorted_indices_to_remove = sorted_scores > sorted_scores.gather(
+            1, last_ind.view(-1, 1)
+        )
+        if self.min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., : self.min_tokens_to_keep] = 0
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+
+        scores = scores.masked_fill(indices_to_remove, self.filter_value)
+        return scores

tortoise/utils/wav2vec_alignment.py

@@ -0,0 +1,164 @@
+import torch
+import torchaudio
+from transformers import Wav2Vec2CTCTokenizer, Wav2Vec2FeatureExtractor, Wav2Vec2ForCTC
+
+
+def max_alignment(s1, s2, skip_character="~", record=None):
+    """
+    A clever function that aligns s1 to s2 as best it can. Wherever a character from s1 is not found in s2, a '~' is
+    used to replace that character.
+
+    Finally got to use my DP skills!
+    """
+    if record is None:
+        record = {}
+    assert (
+        skip_character not in s1
+    ), f"Found the skip character {skip_character} in the provided string, {s1}"
+    if len(s1) == 0:
+        return ""
+    if len(s2) == 0:
+        return skip_character * len(s1)
+    if s1 == s2:
+        return s1
+    if s1[0] == s2[0]:
+        return s1[0] + max_alignment(s1[1:], s2[1:], skip_character, record)
+
+    take_s1_key = (len(s1), len(s2) - 1)
+    if take_s1_key in record:
+        take_s1, take_s1_score = record[take_s1_key]
+    else:
+        take_s1 = max_alignment(s1, s2[1:], skip_character, record)
+        take_s1_score = len(take_s1.replace(skip_character, ""))
+        record[take_s1_key] = (take_s1, take_s1_score)
+
+    take_s2_key = (len(s1) - 1, len(s2))
+    if take_s2_key in record:
+        take_s2, take_s2_score = record[take_s2_key]
+    else:
+        take_s2 = max_alignment(s1[1:], s2, skip_character, record)
+        take_s2_score = len(take_s2.replace(skip_character, ""))
+        record[take_s2_key] = (take_s2, take_s2_score)
+
+    return take_s1 if take_s1_score > take_s2_score else skip_character + take_s2
+
+
+class Wav2VecAlignment:
+    """
+    Uses wav2vec2 to perform audio<->text alignment.
+    """
+
+    def __init__(self, device="cuda"):
+        self.model = Wav2Vec2ForCTC.from_pretrained(
+            "jbetker/wav2vec2-large-robust-ft-libritts-voxpopuli"
+        ).cpu()
+        self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
+            "facebook/wav2vec2-large-960h"
+        )
+        self.tokenizer = Wav2Vec2CTCTokenizer.from_pretrained(
+            "jbetker/tacotron-symbols"
+        )
+        self.device = device
+
+    def align(self, audio, expected_text, audio_sample_rate=24000):
+        orig_len = audio.shape[-1]
+
+        with torch.no_grad():
+            self.model = self.model.to(self.device)
+            audio = audio.to(self.device)
+            audio = torchaudio.functional.resample(audio, audio_sample_rate, 16000)
+            clip_norm = (audio - audio.mean()) / torch.sqrt(audio.var() + 1e-7)
+            logits = self.model(clip_norm).logits
+            self.model = self.model.cpu()
+
+        logits = logits[0]
+        pred_string = self.tokenizer.decode(logits.argmax(-1).tolist())
+
+        fixed_expectation = max_alignment(expected_text.lower(), pred_string)
+        w2v_compression = orig_len // logits.shape[0]
+        expected_tokens = self.tokenizer.encode(fixed_expectation)
+        expected_chars = list(fixed_expectation)
+        if len(expected_tokens) == 1:
+            return [0]  # The alignment is simple; there is only one token.
+        expected_tokens.pop(0)  # The first token is a given.
+        expected_chars.pop(0)
+
+        alignments = [0]
+
+        def pop_till_you_win():
+            if len(expected_tokens) == 0:
+                return None
+            popped = expected_tokens.pop(0)
+            popped_char = expected_chars.pop(0)
+            while popped_char == "~":
+                alignments.append(-1)
+                if len(expected_tokens) == 0:
+                    return None
+                popped = expected_tokens.pop(0)
+                popped_char = expected_chars.pop(0)
+            return popped
+
+        next_expected_token = pop_till_you_win()
+        for i, logit in enumerate(logits):
+            top = logit.argmax()
+            if next_expected_token == top:
+                alignments.append(i * w2v_compression)
+                if len(expected_tokens) > 0:
+                    next_expected_token = pop_till_you_win()
+                else:
+                    break
+
+        pop_till_you_win()
+        if not (len(expected_tokens) == 0 and len(alignments) == len(expected_text)):
+            torch.save([audio, expected_text], "alignment_debug.pth")
+            assert False, (
+                "Something went wrong with the alignment algorithm. I've dumped a file, 'alignment_debug.pth' to"
+                "your current working directory. Please report this along with the file so it can get fixed."
+            )
+
+        # Now fix up alignments. Anything with -1 should be interpolated.
+        alignments.append(
+            orig_len
+        )  # This'll get removed but makes the algorithm below more readable.
+        for i in range(len(alignments)):
+            if alignments[i] == -1:
+                for j in range(i + 1, len(alignments)):
+                    if alignments[j] != -1:
+                        next_found_token = j
+                        break
+                for j in range(i, next_found_token):
+                    gap = alignments[next_found_token] - alignments[i - 1]
+                    alignments[j] = (j - i + 1) * gap // (
+                        next_found_token - i + 1
+                    ) + alignments[i - 1]
+
+        return alignments[:-1]
+
+    def redact(self, audio, expected_text, audio_sample_rate=24000):
+        if "[" not in expected_text:
+            return audio
+        splitted = expected_text.split("[")
+        fully_split = [splitted[0]]
+        for spl in splitted[1:]:
+            assert (
+                "]" in spl
+            ), 'Every "[" character must be paired with a "]" with no nesting.'
+            fully_split.extend(spl.split("]"))
+
+        # At this point, fully_split is a list of strings, with every other string being something that should be redacted.
+        non_redacted_intervals = []
+        last_point = 0
+        for i in range(len(fully_split)):
+            if i % 2 == 0:
+                end_interval = max(0, last_point + len(fully_split[i]) - 1)
+                non_redacted_intervals.append((last_point, end_interval))
+            last_point += len(fully_split[i])
+
+        bare_text = "".join(fully_split)
+        alignments = self.align(audio, bare_text, audio_sample_rate)
+
+        output_audio = []
+        for nri in non_redacted_intervals:
+            start, stop = nri
+            output_audio.append(audio[:, alignments[start] : alignments[stop]])
+        return torch.cat(output_audio, dim=-1)