update sovits to 4.1

.gitignore

@@ -0,0 +1,165 @@
+
+# Created by https://www.toptal.com/developers/gitignore/api/python
+# Edit at https://www.toptal.com/developers/gitignore?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+checkpoints/
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+pytestdebug.log
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+doc/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# End of https://www.toptal.com/developers/gitignore/api/python
+
+/shelf/
+/workspace.xml
+
+dataset
+dataset_raw
+raw
+results
+inference/chunks_temp.json
+logs
+hubert/checkpoint_best_legacy_500.pt
+configs/config.json
+filelists/test.txt
+filelists/train.txt
+filelists/val.txt
+.idea/
+.vscode/
+.idea/modules.xml
+.idea/so-vits-svc.iml
+.idea/vcs.xml
+.idea/inspectionProfiles/profiles_settings.xml
+.idea/inspectionProfiles/Project_Default.xml
+pretrain/
+.vscode/launch.json

.idea/.gitignore

@@ -1,3 +1,3 @@
-# 默认忽略的文件
-/shelf/
-/workspace.xml
+# 默认忽略的文件
+/shelf/
+/workspace.xml

.idea/misc.xml

@@ -1,4 +1,4 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="ProjectRootManager" version="2" project-jdk-name="sovits" project-jdk-type="Python SDK" />
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (so-vits-svc)" project-jdk-type="Python SDK" />
 </project>

.idea/modules.xml

@@ -1,8 +1,8 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="ProjectModuleManager">
-    <modules>
-      <module fileurl="file://$PROJECT_DIR$/.idea/sovits_app.iml" filepath="$PROJECT_DIR$/.idea/sovits_app.iml" />
-    </modules>
-  </component>
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/so-vits-svc.iml" filepath="$PROJECT_DIR$/.idea/so-vits-svc.iml" />
+    </modules>
+  </component>
 </project>

.idea/vcs.xml

@@ -1,6 +1,6 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project version="4">
-  <component name="VcsDirectoryMappings">
-    <mapping directory="" vcs="Git" />
-  </component>
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
 </project>

.ruff.toml

@@ -0,0 +1,4 @@
+select = ["E", "F", "I"]
+
+# Never enforce `E501` (line length violations).
+ignore = ["E501", "E741"]

G_20800.pth

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

LICENSE

@@ -0,0 +1,661 @@
+                    GNU AFFERO GENERAL PUBLIC LICENSE
+                       Version 3, 19 November 2007
+
+ 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/>.

README.md

@@ -5,7 +5,7 @@ colorFrom: green
 colorTo: indigo
 sdk: gradio
 sdk_version: 3.32.0
-app_file: app.py
+app_file: webUI2.py
 pinned: false
 license: apache-2.0
 ---

README_zh_CN.md

@@ -0,0 +1,532 @@
+# SoftVC VITS Singing Voice Conversion
+
+[**English**](./README.md) | [**中文简体**](./README_zh_CN.md)
+
+#### ✨ 带有F0曲线编辑器，角色混合时间轴编辑器的推理端 (Onnx模型的用途) : [MoeVoiceStudio](https://github.com/NaruseMioShirakana/MoeVoiceStudio)
+
+#### ✨ 改善了交互的一个分支推荐：[34j/so-vits-svc-fork](https://github.com/34j/so-vits-svc-fork)
+
+#### ✨ 支持实时转换的一个客户端：[w-okada/voice-changer](https://github.com/w-okada/voice-changer)
+
+**本项目与Vits有着根本上的不同。Vits是TTS，本项目是SVC。本项目无法实现TTS，Vits也无法实现SVC，这两个项目的模型是完全不通用的。**
+
+## 重要通知
+
+这个项目是为了让开发者最喜欢的动画角色唱歌而开发的，任何涉及真人的东西都与开发者的意图背道而驰。
+
+## 声明
+
+本项目为开源、离线的项目，SvcDevelopTeam的所有成员与本项目的所有开发者以及维护者（以下简称贡献者）对本项目没有控制力。本项目的贡献者从未向任何组织或个人提供包括但不限于数据集提取、数据集加工、算力支持、训练支持、推理等一切形式的帮助；本项目的贡献者不知晓也无法知晓使用者使用该项目的用途。故一切基于本项目训练的AI模型和合成的音频都与本项目贡献者无关。一切由此造成的问题由使用者自行承担。
+
+此项目完全离线运行，不能收集任何用户信息或获取用户输入数据。因此，这个项目的贡献者不知道所有的用户输入和模型，因此不负责任何用户输入。
+
+本项目只是一个框架项目，本身并没有语音合成的功能，所有的功能都需要用户自己训练模型。同时，这个项目没有任何模型，任何二次分发的项目都与这个项目的贡献者无关。
+
+## 📏 使用规约
+
+# Warning：请自行解决数据集授权问题，禁止使用非授权数据集进行训练！任何由于使用非授权数据集进行训练造成的问题，需自行承担全部责任和后果！与仓库、仓库维护者、svc develop team 无关！
+
+1. 本项目是基于学术交流目的建立，仅供交流与学习使用，并非为生产环境准备。
+2. 任何发布到视频平台的基于 sovits 制作的视频，都必须要在简介明确指明用于变声器转换的输入源歌声、音频，例如：使用他人发布的视频 / 音频，通过分离的人声作为输入源进行转换的，必须要给出明确的原视频、音乐链接；若使用是自己的人声，或是使用其他歌声合成引擎合成的声音作为输入源进行转换的，也必须在简介加以说明。
+3. 由输入源造成的侵权问题需自行承担全部责任和一切后果。使用其他商用歌声合成软件作为输入源时，请确保遵守该软件的使用条例，注意，许多歌声合成引擎使用条例中明确指明不可用于输入源进行转换！
+4. 禁止使用该项目从事违法行为与宗教、政治等活动，该项目维护者坚决抵制上述行为，不同意此条则禁止使用该项目。
+5. 继续使用视为已同意本仓库 README 所述相关条例，本仓库 README 已进行劝导义务，不对后续可能存在问题负责。
+6. 如果将此项目用于任何其他企划，请提前联系并告知本仓库作者，十分感谢。
+
+## 📝 模型简介
+
+歌声音色转换模型，通过SoftVC内容编码器提取源音频语音特征，与F0同时输入VITS替换原本的文本输入达到歌声转换的效果。同时，更换声码器为 [NSF HiFiGAN](https://github.com/openvpi/DiffSinger/tree/refactor/modules/nsf_hifigan)解决断音问题。
+
+### 🆕 4.1-Stable 版本更新内容
+
++ 特征输入更换为 [Content Vec](https://github.com/auspicious3000/contentvec) 的第12层Transformer输出，并兼容4.0分支
++ 更新浅层扩散，可以使用浅层扩散模型提升音质
++ 增加whisper语音编码器的支持
++ 增加静态/动态声线融合
++ 增加响度嵌入
++ 增加特征检索,来自于[RVC](https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI)
+
+### 🆕 关于兼容4.0模型的问题
+
++ 可通过修改4.0模型的config.json对4.0的模型进行支持，需要在config.json的model字段中添加speech_encoder字段，具体见下
+
+```
+  "model": {
+    .........
+    "ssl_dim": 256,
+    "n_speakers": 200,
+    "speech_encoder":"vec256l9"
+  }
+```
+
+### 🆕 关于浅扩散
+![Diagram](shadowdiffusion.png)
+
+## 💬 关于 Python 版本问题
+
+在进行测试后，我们认为`Python 3.8.9`能够稳定地运行该项目
+
+## 📥 预先下载的模型文件
+
+#### **必须项**
+
+**以下编码器需要选择一个使用**
+
+##### **1. 若使用contentvec作为声音编码器（推荐）**
+
+`vec768l12`与`vec256l9` 需要该编码器
+
++ contentvec ：[checkpoint_best_legacy_500.pt](https://ibm.box.com/s/z1wgl1stco8ffooyatzdwsqn2psd9lrr)
+  + 放在`pretrain`目录下
+
+或者下载下面的ContentVec，大小只有199MB，但效果相同:
++ contentvec ：[hubert_base.pt](https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/hubert_base.pt)
+  + 将文件名改为`checkpoint_best_legacy_500.pt`后，放在`pretrain`目录下
+
+```shell
+# contentvec
+wget -P pretrain/ http://obs.cstcloud.cn/share/obs/sankagenkeshi/checkpoint_best_legacy_500.pt
+# 也可手动下载放在pretrain目录
+```
+
+##### **2. 若使用hubertsoft作为声音编码器**
++ soft vc hubert：[hubert-soft-0d54a1f4.pt](https://github.com/bshall/hubert/releases/download/v0.1/hubert-soft-0d54a1f4.pt)
+  + 放在`pretrain`目录下
+
+##### **3. 若使用Whisper-ppg作为声音编码器**
++ 下载模型 [medium.pt](https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt), 该模型适配`whisper-ppg`
++ 下载模型 [large-v2.pt](https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt), 该模型适配`whisper-ppg-large`
+  + 放在`pretrain`目录下
+ 
+##### **4. 若使用cnhubertlarge作为声音编码器**
++ 下载模型 [chinese-hubert-large-fairseq-ckpt.pt](https://huggingface.co/TencentGameMate/chinese-hubert-large/resolve/main/chinese-hubert-large-fairseq-ckpt.pt)
+  + 放在`pretrain`目录下
+
+##### **5. 若使用dphubert作为声音编码器**
++ 下载模型 [DPHuBERT-sp0.75.pth](https://huggingface.co/pyf98/DPHuBERT/resolve/main/DPHuBERT-sp0.75.pth)
+  + 放在`pretrain`目录下
+
+##### **6. 若使用WavLM作为声音编码器**
++ 下载模型 [WavLM-Base+.pt](https://valle.blob.core.windows.net/share/wavlm/WavLM-Base+.pt?sv=2020-08-04&st=2023-03-01T07%3A51%3A05Z&se=2033-03-02T07%3A51%3A00Z&sr=c&sp=rl&sig=QJXmSJG9DbMKf48UDIU1MfzIro8HQOf3sqlNXiflY1I%3D), 该模型适配`wavlmbase+`
+  + 放在`pretrain`目录下
+
+##### **7. 若使用OnnxHubert/ContentVec作为声音编码器**
++ 下载模型 [MoeSS-SUBModel](https://huggingface.co/NaruseMioShirakana/MoeSS-SUBModel/tree/main)
+  + 放在`pretrain`目录下
+
+#### **编码器列表**
+- "vec768l12"
+- "vec256l9"
+- "vec256l9-onnx"
+- "vec256l12-onnx"
+- "vec768l9-onnx"
+- "vec768l12-onnx"
+- "hubertsoft-onnx"
+- "hubertsoft"
+- "whisper-ppg"
+- "cnhubertlarge"
+- "dphubert"
+- "whisper-ppg-large"
+- "wavlmbase+"
+
+#### **可选项(强烈建议使用)**
+
++ 预训练底模文件： `G_0.pth` `D_0.pth`
+  + 放在`logs/44k`目录下
+
++ 扩散模型预训练底模文件： `model_0.pt `
+  + 放在`logs/44k/diffusion`目录下
+
+从svc-develop-team(待定)或任何其他地方获取Sovits底模
+
+扩散模型引用了[Diffusion-SVC](https://github.com/CNChTu/Diffusion-SVC)的Diffusion Model，底模与[Diffusion-SVC](https://github.com/CNChTu/Diffusion-SVC)的扩散模型底模通用，可以去[Diffusion-SVC](https://github.com/CNChTu/Diffusion-SVC)获取扩散模型的底模
+
+虽然底模一般不会引起什么版权问题，但还是请注意一下，比如事先询问作者，又或者作者在模型描述中明确写明了可行的用途
+
+#### **可选项(根据情况选择)**
+
+如果使用`NSF-HIFIGAN增强器`或`浅层扩散`的话，需要下载预训练的NSF-HIFIGAN模型，如果不需要可以不下载
+
++ 预训练的NSF-HIFIGAN声码器 ：[nsf_hifigan_20221211.zip](https://github.com/openvpi/vocoders/releases/download/nsf-hifigan-v1/nsf_hifigan_20221211.zip)
+  + 解压后，将四个文件放在`pretrain/nsf_hifigan`目录下
+
+```shell
+# nsf_hifigan
+wget -P pretrain/ https://github.com/openvpi/vocoders/releases/download/nsf-hifigan-v1/nsf_hifigan_20221211.zip
+unzip -od pretrain/nsf_hifigan pretrain/nsf_hifigan_20221211.zip
+# 也可手动下载放在pretrain/nsf_hifigan目录
+# 地址：https://github.com/openvpi/vocoders/releases/tag/nsf-hifigan-v1
+```
+
+## 📊 数据集准备
+
+仅需要以以下文件结构将数据集放入dataset_raw目录即可
+
+```
+dataset_raw
+├───speaker0
+│   ├───xxx1-xxx1.wav
+│   ├───...
+│   └───Lxx-0xx8.wav
+└───speaker1
+    ├───xx2-0xxx2.wav
+    ├───...
+    └───xxx7-xxx007.wav
+```
+
+可以自定义说话人名称
+
+```
+dataset_raw
+└───suijiSUI
+    ├───1.wav
+    ├───...
+    └───25788785-20221210-200143-856_01_(Vocals)_0_0.wav
+```
+
+## 🛠️ 数据预处理
+
+### 0. 音频切片
+
+将音频切片至`5s - 15s`, 稍微长点也无伤大雅，实在太长可能会导致训练中途甚至预处理就爆显存
+
+可以使用[audio-slicer-GUI](https://github.com/flutydeer/audio-slicer)、[audio-slicer-CLI](https://github.com/openvpi/audio-slicer)
+
+一般情况下只需调整其中的`Minimum Interval`，普通陈述素材通常保持默认即可，歌唱素材可以调整至`100`甚至`50`
+
+切完之后手动删除过长过短的音频
+
+**如果你使用Whisper-ppg声音编码器进行训练，所有的切片长度必须小于30s**
+
+### 1. 重采样至44100Hz单声道
+
+```shell
+python resample.py
+```
+
+#### 注意
+
+虽然本项目拥有重采样、转换单声道与响度匹配的脚本resample.py，但是默认的响度匹配是匹配到0db。这可能会造成音质的受损。而python的响度匹配包pyloudnorm无法对电平进行压限，这会导致爆音。所以建议可以考虑使用专业声音处理软件如`adobe audition`等软件做响度匹配处理。若已经使用其他软件做响度匹配，可以在运行上述命令时添加`--skip_loudnorm`跳过响度匹配步骤。如：
+
+```shell
+python resample.py --skip_loudnorm
+```
+
+### 2. 自动划分训练集、验证集，以及自动生成配置文件
+
+```shell
+python preprocess_flist_config.py --speech_encoder vec768l12
+```
+
+speech_encoder拥有以下选择
+
+```
+vec768l12
+vec256l9
+hubertsoft
+whisper-ppg
+whisper-ppg-large
+cnhubertlarge
+dphubert
+wavlmbase+
+```
+
+如果省略speech_encoder参数，默认值为vec768l12
+
+**使用响度嵌入**
+
+若使用响度嵌入，需要增加`--vol_aug`参数，比如：
+
+```shell
+python preprocess_flist_config.py --speech_encoder vec768l12 --vol_aug
+```
+
+使用后训练出的模型将匹配到输入源响度，否则为训练集响度。
+
+#### 此时可以在生成的config.json与diffusion.yaml修改部分参数
+
+##### config.json
+
+* `keep_ckpts`：训练时保留最后几个模型，`0`为保留所有，默认只保留最后`3`个
+
+* `all_in_mem`：加载所有数据集到内存中，某些平台的硬盘IO过于低下、同时内存容量 **远大于** 数据集体积时可以启用
+
+* `batch_size`：单次训练加载到GPU的数据量，调整到低于显存容量的大小即可
+
+* `vocoder_name` : 选择一种声码器，默认为`nsf-hifigan`.
+
+##### diffusion.yaml
+
+* `cache_all_data`：加载所有数据集到内存中，某些平台的硬盘IO过于低下、同时内存容量 **远大于** 数据集体积时可以启用
+
+* `duration`：训练时音频切片时长，可根据显存大小调整，**注意，该值必须小于训练集内音频的最短时间！**
+
+* `batch_size`：单次训练加载到GPU的数据量，调整到低于显存容量的大小即可
+
+* `timesteps` : 扩散模型总步数，默认为1000.
+
+* `k_step_max` : 训练时可仅训练`k_step_max`步扩散以节约训练时间，注意，该值必须小于`timesteps`，0为训练整个扩散模型，**注意，如果不训练整个扩散模型将无法使用仅扩散模型推理!**
+  
+##### **声码器列表**
+
+```
+nsf-hifigan
+nsf-snake-hifigan
+```
+
+### 3. 生成hubert与f0
+
+```shell
+python preprocess_hubert_f0.py --f0_predictor dio
+```
+
+f0_predictor拥有四个选择
+
+```
+crepe
+dio
+pm
+harvest
+```
+
+如果训练集过于嘈杂，请使用crepe处理f0
+
+如果省略f0_predictor参数，默认值为dio
+
+尚若需要浅扩散功能（可选），需要增加--use_diff参数，比如
+
+```shell
+python preprocess_hubert_f0.py --f0_predictor dio --use_diff
+```
+
+执行完以上步骤后 dataset 目录便是预处理完成的数据，可以删除 dataset_raw 文件夹了
+
+## 🏋️‍ 训练
+
+### 主模型训练
+
+```shell
+python train.py -c configs/config.json -m 44k
+```
+
+### 扩散模型（可选）
+
+尚若需要浅扩散功能，需要训练扩散模型，扩散模型训练方法为:
+
+```shell
+python train_diff.py -c configs/diffusion.yaml
+```
+
+模型训练结束后，模型文件保存在`logs/44k`目录下，扩散模型在`logs/44k/diffusion`下
+
+## 🤖 推理
+
+使用 [inference_main.py](inference_main.py)
+
+```shell
+# 例
+python inference_main.py -m "logs/44k/G_30400.pth" -c "configs/config.json" -n "君の知らない物語-src.wav" -t 0 -s "nen"
+```
+
+必填项部分：
++ `-m` | `--model_path`：模型路径
++ `-c` | `--config_path`：配置文件路径
++ `-n` | `--clean_names`：wav 文件名列表，放在 raw 文件夹下
++ `-t` | `--trans`：音高调整，支持正负（半音）
++ `-s` | `--spk_list`：合成目标说话人名称
++ `-cl` | `--clip`：音频强制切片，默认0为自动切片，单位为秒/s
+
+可选项部分：部分具体见下一节
++ `-lg` | `--linear_gradient`：两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒
++ `-f0p` | `--f0_predictor`：选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意：crepe为原F0使用均值滤波器)
++ `-a` | `--auto_predict_f0`：语音转换自动预测音高，转换歌声时不要打开这个会严重跑调
++ `-cm` | `--cluster_model_path`：聚类模型或特征检索索引路径，留空则自动设为各方案模型的默认路径，如果没有训练聚类或特征检索则随便填
++ `-cr` | `--cluster_infer_ratio`：聚类方案或特征检索占比，范围0-1，若没有训练聚类模型或特征检索则默认0即可
++ `-eh` | `--enhance`：是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭
++ `-shd` | `--shallow_diffusion`：是否使用浅层扩散，使用后可解决一部分电音问题，默认关闭，该选项打开时，NSF_HIFIGAN增强器将会被禁止
++ `-usm` | `--use_spk_mix`：是否使用角色融合/动态声线融合
++ `-lea` | `--loudness_envelope_adjustment`：输入源响度包络替换输出响度包络融合比例，越靠近1越使用输出响度包络
++ `-fr` | `--feature_retrieval`：是否使用特征检索，如果使用聚类模型将被禁用，且cm与cr参数将会变成特征检索的索引路径与混合比例
+
+浅扩散设置：
++ `-dm` | `--diffusion_model_path`：扩散模型路径
++ `-dc` | `--diffusion_config_path`：扩散模型配置文件路径
++ `-ks` | `--k_step`：扩散步数，越大越接近扩散模型的结果，默认100
++ `-od` | `--only_diffusion`：纯扩散模式，该模式不会加载sovits模型，以扩散模型推理
++ `-se` | `--second_encoding`：二次编码，浅扩散前会对原始音频进行二次编码，玄学选项，有时候效果好，有时候效果差
+
+### 注意！
+
+如果使用`whisper-ppg` 声音编码器进行推理，需要将`--clip`设置为25，`-lg`设置为1。否则将无法正常推理。
+
+## 🤔 可选项
+
+如果前面的效果已经满意，或者没看明白下面在讲啥，那后面的内容都可以忽略，不影响模型使用(这些可选项影响比较小，可能在某些特定数据上有点效果，但大部分情况似乎都感知不太明显)
+
+### 自动f0预测
+
+4.0模型训练过程会训练一个f0预测器，对于语音转换可以开启自动音高预测，如果效果不好也可以使用手动的，但转换歌声时请不要启用此功能！！！会严重跑调！！
++ 在inference_main中设置auto_predict_f0为true即可
+
+### 聚类音色泄漏控制
+
+介绍：聚类方案可以减小音色泄漏，使得模型训练出来更像目标的音色（但其实不是特别明显），但是单纯的聚类方案会降低模型的咬字（会口齿不清）（这个很明显），本模型采用了融合的方式，可以线性控制聚类方案与非聚类方案的占比，也就是可以手动在"像目标音色" 和 "咬字清晰" 之间调整比例，找到合适的折中点
+
+使用聚类前面的已有步骤不用进行任何的变动，只需要额外训练一个聚类模型，虽然效果比较有限，但训练成本也比较低
+
++ 训练过程：
+  + 使用cpu性能较好的机器训练，据我的经验在腾讯云6核cpu训练每个speaker需要约4分钟即可完成训练
+  + 执行`python cluster/train_cluster.py`，模型的输出会在`logs/44k/kmeans_10000.pt`
+  + 聚类模型目前可以使用gpu进行训练，执行`python cluster/train_cluster.py --gpu`
++ 推理过程：
+  + `inference_main.py`中指定`cluster_model_path` 为模型输出文件, 留空则默认为`logs/44k/kmeans_10000.pt`
+  + `inference_main.py`中指定`cluster_infer_ratio`，`0`为完全不使用聚类，`1`为只使用聚类，通常设置`0.5`即可
+
+### 特征检索
+
+介绍：跟聚类方案一样可以减小音色泄漏，咬字比聚类稍好，但会降低推理速度，采用了融合的方式，可以线性控制特征检索与非特征检索的占比，
+
++ 训练过程：
+  首先需要在生成hubert与f0后执行：
+
+```shell
+python train_index.py -c configs/config.json
+```
+
+模型的输出会在`logs/44k/feature_and_index.pkl`
+
++ 推理过程：
+  + 需要首先指定`--feature_retrieval`，此时聚类方案会自动切换到特征检索方案
+  + `inference_main.py`中指定`cluster_model_path` 为模型输出文件, 留空则默认为`logs/44k/feature_and_index.pkl`
+  + `inference_main.py`中指定`cluster_infer_ratio`，`0`为完全不使用特征检索，`1`为只使用特征检索，通常设置`0.5`即可
+
+### [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/svc-develop-team/so-vits-svc/blob/4.1-Stable/sovits4_for_colab.ipynb) [sovits4_for_colab.ipynb](https://colab.research.google.com/github/svc-develop-team/so-vits-svc/blob/4.1-Stable/sovits4_for_colab.ipynb)
+
+## 🗜️ 模型压缩
+
+生成的模型含有继续训练所需的信息。如果确认不再训练，可以移除模型中此部分信息，得到约 1/3 大小的最终模型。
+
+使用 [compress_model.py](compress_model.py)
+
+```shell
+# 例
+python compress_model.py -c="configs/config.json" -i="logs/44k/G_30400.pth" -o="logs/44k/release.pth"
+```
+
+## 👨‍🔧 声线混合
+
+### 静态声线混合
+
+**参考`webUI.py`文件中，小工具/实验室特性的静态声线融合。**
+
+介绍:该功能可以将多个声音模型合成为一个声音模型(多个模型参数的凸组合或线性组合)，从而制造出现实中不存在的声线 
+**注意：**
+
+1. 该功能仅支持单说话人的模型
+2. 如果强行使用多说话人模型，需要保证多个模型的说话人数量相同，这样可以混合同一个SpaekerID下的声音
+3. 保证所有待混合模型的config.json中的model字段是相同的
+4. 输出的混合模型可以使用待合成模型的任意一个config.json，但聚类模型将不能使用
+5. 批量上传模型的时候最好把模型放到一个文件夹选中后一起上传
+6. 混合比例调整建议大小在0-100之间，也可以调为其他数字，但在线性组合模式下会出现未知的效果
+7. 混合完毕后，文件将会保存在项目根目录中，文件名为output.pth
+8. 凸组合模式会将混合比例执行Softmax使混合比例相加为1，而线性组合模式不会
+
+### 动态声线混合
+
+**参考`spkmix.py`文件中关于动态声线混合的介绍**
+
+角色混合轨道 编写规则：
+
+角色ID : \[\[起始时间1, 终止时间1, 起始数值1, 起始数值1], [起始时间2, 终止时间2, 起始数值2, 起始数值2]]
+
+起始时间和前一个的终止时间必须相同，第一个起始时间必须为0，最后一个终止时间必须为1 （时间的范围为0-1）
+
+全部角色必须填写，不使用的角色填\[\[0., 1., 0., 0.]]即可
+
+融合数值可以随便填，在指定的时间段内从起始数值线性变化为终止数值，内部会自动确保线性组合为1（凸组合条件），可以放心使用
+
+推理的时候使用`--use_spk_mix`参数即可启用动态声线混合
+
+## 📤 Onnx导出
+
+使用 [onnx_export.py](onnx_export.py)
+
++ 新建文件夹：`checkpoints` 并打开
++ 在`checkpoints`文件夹中新建一个文件夹作为项目文件夹，文件夹名为你的项目名称，比如`aziplayer`
++ 将你的模型更名为`model.pth`，配置文件更名为`config.json`，并放置到刚才创建的`aziplayer`文件夹下
++ 将 [onnx_export.py](onnx_export.py) 中`path = "NyaruTaffy"` 的 `"NyaruTaffy"` 修改为你的项目名称，`path = "aziplayer" (onnx_export_speaker_mix，为支持角色混合的onnx导出)`
++ 运行 [onnx_export.py](onnx_export.py) 
++ 等待执行完毕，在你的项目文件夹下会生成一个`model.onnx`，即为导出的模型
+
+注意：Hubert Onnx模型请使用MoeSS提供的模型，目前无法自行导出（fairseq中Hubert有不少onnx不支持的算子和涉及到常量的东西，在导出时会报错或者导出的模型输入输出shape和结果都有问题）
+
+
+## 📎 引用及论文
+
+| URL | 名称 | 标题 | 源码 |
+| --- | ----------- | ----- | --------------------- |
+|[2106.06103](https://arxiv.org/abs/2106.06103) | VITS (Synthesizer)| Conditional Variational Autoencoder with Adversarial Learning for End-to-End Text-to-Speech | [jaywalnut310/vits](https://github.com/jaywalnut310/vits) |
+|[2111.02392](https://arxiv.org/abs/2111.02392) | SoftVC (Speech Encoder)| A Comparison of Discrete and Soft Speech Units for Improved Voice Conversion | [bshall/hubert](https://github.com/bshall/hubert) |
+|[2204.09224](https://arxiv.org/abs/2204.09224) | ContentVec (Speech Encoder)| ContentVec: An Improved Self-Supervised Speech Representation by Disentangling Speakers | [auspicious3000/contentvec](https://github.com/auspicious3000/contentvec) |
+|[2212.04356](https://arxiv.org/abs/2212.04356) | Whisper (Speech Encoder) | Robust Speech Recognition via Large-Scale Weak Supervision | [openai/whisper](https://github.com/openai/whisper) |
+|[2110.13900](https://arxiv.org/abs/2110.13900) | WavLM (Speech Encoder) | WavLM: Large-Scale Self-Supervised Pre-Training for Full Stack Speech Processing | [microsoft/unilm/wavlm](https://github.com/microsoft/unilm/tree/master/wavlm) |
+|[2305.17651](https://arxiv.org/abs/2305.17651) | DPHubert (Speech Encoder) | DPHuBERT: Joint Distillation and Pruning of Self-Supervised Speech Models | [pyf98/DPHuBERT](https://github.com/pyf98/DPHuBERT) |
+|[DOI:10.21437/Interspeech.2017-68](http://dx.doi.org/10.21437/Interspeech.2017-68) | Harvest (F0 Predictor) | Harvest: A high-performance fundamental frequency estimator from speech signals | [mmorise/World/harvest](https://github.com/mmorise/World/blob/master/src/harvest.cpp) |
+|[aes35-000039](https://www.aes.org/e-lib/online/browse.cfm?elib=15165) | Dio (F0 Predictor) | Fast and reliable F0 estimation method based on the period extraction of vocal fold vibration of singing voice and speech | [mmorise/World/dio](https://github.com/mmorise/World/blob/master/src/dio.cpp) |
+|[8461329](https://ieeexplore.ieee.org/document/8461329) | Crepe (F0 Predictor) | Crepe: A Convolutional Representation for Pitch Estimation | [maxrmorrison/torchcrepe](https://github.com/maxrmorrison/torchcrepe) |
+|[DOI:10.1016/j.wocn.2018.07.001](https://doi.org/10.1016/j.wocn.2018.07.001) | Parselmouth (F0 Predictor) | Introducing Parselmouth: A Python interface to Praat | [YannickJadoul/Parselmouth](https://github.com/YannickJadoul/Parselmouth) |
+|[2010.05646](https://arxiv.org/abs/2010.05646) | HIFIGAN (Vocoder) | HiFi-GAN: Generative Adversarial Networks for Efficient and High Fidelity Speech Synthesis | [jik876/hifi-gan](https://github.com/jik876/hifi-gan) |
+|[1810.11946](https://arxiv.org/abs/1810.11946.pdf) | NSF (Vocoder) | Neural source-filter-based waveform model for statistical parametric speech synthesis | [openvpi/DiffSinger/modules/nsf_hifigan](https://github.com/openvpi/DiffSinger/tree/refactor/modules/nsf_hifigan)
+|[2006.08195](https://arxiv.org/abs/2006.08195) | Snake (Vocoder) | Neural Networks Fail to Learn Periodic Functions and How to Fix It | [EdwardDixon/snake](https://github.com/EdwardDixon/snake)
+|[2105.02446v3](https://arxiv.org/abs/2105.02446v3) | Shallow Diffusion (PostProcessing)| DiffSinger: Singing Voice Synthesis via Shallow Diffusion Mechanism | [CNChTu/Diffusion-SVC](https://github.com/CNChTu/Diffusion-SVC) |
+|[K-means](https://citeseerx.ist.psu.edu/viewdoc/download;jsessionid=01D65490BADCC216F350D06F84D721AD?doi=10.1.1.308.8619&rep=rep1&type=pdf) | Feature K-means Clustering (PreProcessing)| Some methods for classification and analysis of multivariate observations | 本代码库 |
+| | Feature TopK Retrieval (PreProcessing)| Retrieval based Voice Conversion | [RVC-Project/Retrieval-based-Voice-Conversion-WebUI](https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI) |
+
+## ☀️ 旧贡献者
+
+因为某些原因原作者进行了删库处理，本仓库重建之初由于组织成员疏忽直接重新上传了所有文件导致以前的contributors全部木大，现在在README里重新添加一个旧贡献者列表
+
+*某些成员已根据其个人意愿不将其列出*
+
+<table>
+  <tr>
+    <td align="center"><a href="https://github.com/MistEO"><img src="https://avatars.githubusercontent.com/u/18511905?v=4" width="100px;" alt=""/><br /><sub><b>MistEO</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/XiaoMiku01"><img src="https://avatars.githubusercontent.com/u/54094119?v=4" width="100px;" alt=""/><br /><sub><b>XiaoMiku01</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/ForsakenRei"><img src="https://avatars.githubusercontent.com/u/23041178?v=4" width="100px;" alt=""/><br /><sub><b>しぐれ</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/TomoGaSukunai"><img src="https://avatars.githubusercontent.com/u/25863522?v=4" width="100px;" alt=""/><br /><sub><b>TomoGaSukunai</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/Plachtaa"><img src="https://avatars.githubusercontent.com/u/112609742?v=4" width="100px;" alt=""/><br /><sub><b>Plachtaa</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/zdxiaoda"><img src="https://avatars.githubusercontent.com/u/45501959?v=4" width="100px;" alt=""/><br /><sub><b>zd小达</b></sub></a><br /></td>
+    <td align="center"><a href="https://github.com/Archivoice"><img src="https://avatars.githubusercontent.com/u/107520869?v=4" width="100px;" alt=""/><br /><sub><b>凍聲響世</b></sub></a><br /></td>
+  </tr>
+</table>
+
+## 📚 一些法律条例参考
+
+#### 任何国家，地区，组织和个人使用此项目必须遵守以下法律
+
+#### 《民法典》
+
+##### 第一千零一十九条
+
+任何组织或者个人不得以丑化、污损，或者利用信息技术手段伪造等方式侵害他人的肖像权。未经肖像权人同意，不得制作、使用、公开肖像权人的肖像，但是法律另有规定的除外。未经肖像权人同意，肖像作品权利人不得以发表、复制、发行、出租、展览等方式使用或者公开肖像权人的肖像。对自然人声音的保护，参照适用肖像权保护的有关规定。
+
+##### 第一千零二十四条
+
+【名誉权】民事主体享有名誉权。任何组织或者个人不得以侮辱、诽谤等方式侵害他人的名誉权。
+
+##### 第一千零二十七条
+
+【作品侵害名誉权】行为人发表的文学、艺术作品以真人真事或者特定人为描述对象，含有侮辱、诽谤内容，侵害他人名誉权的，受害人有权依法请求该行为人承担民事责任。行为人发表的文学、艺术作品不以特定人为描述对象，仅其中的情节与该特定人的情况相似的，不承担民事责任。
+
+#### 《[中华人民共和国宪法](http://www.gov.cn/guoqing/2018-03/22/content_5276318.htm)》
+
+#### 《[中华人民共和国刑法](http://gongbao.court.gov.cn/Details/f8e30d0689b23f57bfc782d21035c3.html?sw=中华人民共和国刑法)》
+
+#### 《[中华人民共和国民法典](http://gongbao.court.gov.cn/Details/51eb6750b8361f79be8f90d09bc202.html)》
+
+#### 《[中华人民共和国合同法](http://www.npc.gov.cn/zgrdw/npc/lfzt/rlyw/2016-07/01/content_1992739.htm)》
+
+## 💪 感谢所有的贡献者
+<a href="https://github.com/svc-develop-team/so-vits-svc/graphs/contributors" target="_blank">
+  <img src="https://contrib.rocks/image?repo=svc-develop-team/so-vits-svc" />
+</a>

cluster/__init__.py

@@ -1,7 +1,7 @@
-import numpy as np
 import torch
 from sklearn.cluster import KMeans
 
+
 def get_cluster_model(ckpt_path):
     checkpoint = torch.load(ckpt_path)
     kmeans_dict = {}

cluster/kmeans.py

@@ -1,201 +1,204 @@
-import math,pdb
-import torch,pynvml
-from torch.nn.functional import normalize
-from time import time
-import numpy as np
-# device=torch.device("cuda:0")
-def _kpp(data: torch.Tensor, k: int, sample_size: int = -1):
-    """ Picks k points in the data based on the kmeans++ method.
-
-    Parameters
-    ----------
-    data : torch.Tensor
-        Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D
-        data, rank 2 multidimensional data, in which case one
-        row is one observation.
-    k : int
-        Number of samples to generate.
-    sample_size : int
-        sample data to avoid memory overflow during calculation
-
-    Returns
-    -------
-    init : ndarray
-        A 'k' by 'N' containing the initial centroids.
-
-    References
-    ----------
-    .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of
-       careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium
-       on Discrete Algorithms, 2007.
-    .. [2] scipy/cluster/vq.py: _kpp
-    """
-    batch_size=data.shape[0]
-    if batch_size>sample_size:
-        data = data[torch.randint(0, batch_size,[sample_size], device=data.device)]
-    dims = data.shape[1] if len(data.shape) > 1 else 1
-    init = torch.zeros((k, dims)).to(data.device)
-    r = torch.distributions.uniform.Uniform(0, 1)
-    for i in range(k):
-        if i == 0:
-            init[i, :] = data[torch.randint(data.shape[0], [1])]
-        else:
-            D2 = torch.cdist(init[:i, :][None, :], data[None, :], p=2)[0].amin(dim=0)
-            probs = D2 / torch.sum(D2)
-            cumprobs = torch.cumsum(probs, dim=0)
-            init[i, :] = data[torch.searchsorted(cumprobs, r.sample([1]).to(data.device))]
-    return init
-class KMeansGPU:
-  '''
-  Kmeans clustering algorithm implemented with PyTorch
-
-  Parameters:
-    n_clusters: int, 
-      Number of clusters
-
-    max_iter: int, default: 100
-      Maximum number of iterations
-
-    tol: float, default: 0.0001
-      Tolerance
-    
-    verbose: int, default: 0
-      Verbosity
-
-    mode: {'euclidean', 'cosine'}, default: 'euclidean'
-      Type of distance measure
-      
-    init_method: {'random', 'point', '++'}
-      Type of initialization
-
-    minibatch: {None, int}, default: None
-      Batch size of MinibatchKmeans algorithm
-      if None perform full KMeans algorithm
-      
-  Attributes:
-    centroids: torch.Tensor, shape: [n_clusters, n_features]
-      cluster centroids
-  '''
-  def __init__(self, n_clusters, max_iter=200, tol=1e-4, verbose=0, mode="euclidean",device=torch.device("cuda:0")):
-    self.n_clusters = n_clusters
-    self.max_iter = max_iter
-    self.tol = tol
-    self.verbose = verbose
-    self.mode = mode
-    self.device=device
-    pynvml.nvmlInit()
-    gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device.index)
-    info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
-    self.minibatch=int(33e6/self.n_clusters*info.free/ 1024 / 1024 / 1024)
-    print("free_mem/GB:",info.free/ 1024 / 1024 / 1024,"minibatch:",self.minibatch)
-    
-  @staticmethod
-  def cos_sim(a, b):
-    """
-      Compute cosine similarity of 2 sets of vectors
-
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    return normalize(a, dim=-1) @ normalize(b, dim=-1).transpose(-2, -1)
-
-  @staticmethod
-  def euc_sim(a, b):
-    """
-      Compute euclidean similarity of 2 sets of vectors
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    return 2 * a @ b.transpose(-2, -1) -(a**2).sum(dim=1)[..., :, None] - (b**2).sum(dim=1)[..., None, :]
-
-  def max_sim(self, a, b):
-    """
-      Compute maximum similarity (or minimum distance) of each vector
-      in a with all of the vectors in b
-      Parameters:
-      a: torch.Tensor, shape: [m, n_features]
-      b: torch.Tensor, shape: [n, n_features]
-    """
-    if self.mode == 'cosine':
-      sim_func = self.cos_sim
-    elif self.mode == 'euclidean':
-      sim_func = self.euc_sim
-    sim = sim_func(a, b)
-    max_sim_v, max_sim_i = sim.max(dim=-1)
-    return max_sim_v, max_sim_i
-
-  def fit_predict(self, X):
-    """
-      Combination of fit() and predict() methods.
-      This is faster than calling fit() and predict() seperately.
-      Parameters:
-      X: torch.Tensor, shape: [n_samples, n_features]
-      centroids: {torch.Tensor, None}, default: None
-        if given, centroids will be initialized with given tensor
-        if None, centroids will be randomly chosen from X
-      Return:
-      labels: torch.Tensor, shape: [n_samples]
-
-            mini_=33kk/k*remain
-            mini=min(mini_,fea_shape)
-            offset=log2(k/1000)*1.5
-            kpp_all=min(mini_*10/offset,fea_shape)
-            kpp_sample=min(mini_/12/offset,fea_shape)
-    """
-    assert isinstance(X, torch.Tensor), "input must be torch.Tensor"
-    assert X.dtype in [torch.half, torch.float, torch.double], "input must be floating point"
-    assert X.ndim == 2, "input must be a 2d tensor with shape: [n_samples, n_features] "
-    # print("verbose:%s"%self.verbose)
-
-    offset = np.power(1.5,np.log(self.n_clusters / 1000))/np.log(2)
-    with torch.no_grad():
-      batch_size= X.shape[0]
-      # print(self.minibatch, int(self.minibatch * 10 / offset), batch_size)
-      start_time = time()
-      if (self.minibatch*10//offset< batch_size):
-        x = X[torch.randint(0, batch_size,[int(self.minibatch*10/offset)])].to(self.device)
-      else:
-        x = X.to(self.device)
-      # print(x.device)
-      self.centroids = _kpp(x, self.n_clusters, min(int(self.minibatch/12/offset),batch_size))
-      del x
-      torch.cuda.empty_cache()
-      # self.centroids = self.centroids.to(self.device)
-      num_points_in_clusters = torch.ones(self.n_clusters, device=self.device, dtype=X.dtype)#全1
-      closest = None#[3098036]#int64
-      if(self.minibatch>=batch_size//2 and self.minibatch<batch_size):
-        X = X[torch.randint(0, batch_size,[self.minibatch])].to(self.device)
-      elif(self.minibatch>=batch_size):
-        X=X.to(self.device)
-      for i in range(self.max_iter):
-        iter_time = time()
-        if self.minibatch<batch_size//2:#可用minibatch数太小，每次都得从内存倒腾到显存
-          x = X[torch.randint(0, batch_size, [self.minibatch])].to(self.device)
-        else:#否则直接全部缓存
-          x = X
-
-        closest = self.max_sim(a=x, b=self.centroids)[1].to(torch.int16)#[3098036]#int64#0~999
-        matched_clusters, counts = closest.unique(return_counts=True)#int64#1k
-        expanded_closest = closest[None].expand(self.n_clusters, -1)#[1000, 3098036]#int16#0~999
-        mask = (expanded_closest==torch.arange(self.n_clusters, device=self.device)[:, None]).to(X.dtype)#==后者是int64*1000
-        c_grad = mask @ x / mask.sum(-1)[..., :, None]
-        c_grad[c_grad!=c_grad] = 0 # remove NaNs
-        error = (c_grad - self.centroids).pow(2).sum()
-        if self.minibatch is not None:
-          lr = 1/num_points_in_clusters[:,None] * 0.9 + 0.1
-        else:
-          lr = 1
-        matched_clusters=matched_clusters.long()
-        num_points_in_clusters[matched_clusters] += counts#IndexError: tensors used as indices must be long, byte or bool tensors
-        self.centroids = self.centroids * (1-lr) + c_grad * lr
-        if self.verbose >= 2:
-          print('iter:', i, 'error:', error.item(), 'time spent:', round(time()-iter_time, 4))
-        if error <= self.tol:
-          break
-
-      if self.verbose >= 1:
-        print(f'used {i+1} iterations ({round(time()-start_time, 4)}s) to cluster {batch_size} items into {self.n_clusters} clusters')
-    return closest
+from time import time
+
+import numpy as np
+import pynvml
+import torch
+from torch.nn.functional import normalize
+
+
+# device=torch.device("cuda:0")
+def _kpp(data: torch.Tensor, k: int, sample_size: int = -1):
+    """ Picks k points in the data based on the kmeans++ method.
+
+    Parameters
+    ----------
+    data : torch.Tensor
+        Expect a rank 1 or 2 array. Rank 1 is assumed to describe 1-D
+        data, rank 2 multidimensional data, in which case one
+        row is one observation.
+    k : int
+        Number of samples to generate.
+    sample_size : int
+        sample data to avoid memory overflow during calculation
+
+    Returns
+    -------
+    init : ndarray
+        A 'k' by 'N' containing the initial centroids.
+
+    References
+    ----------
+    .. [1] D. Arthur and S. Vassilvitskii, "k-means++: the advantages of
+       careful seeding", Proceedings of the Eighteenth Annual ACM-SIAM Symposium
+       on Discrete Algorithms, 2007.
+    .. [2] scipy/cluster/vq.py: _kpp
+    """
+    batch_size=data.shape[0]
+    if batch_size>sample_size:
+        data = data[torch.randint(0, batch_size,[sample_size], device=data.device)]
+    dims = data.shape[1] if len(data.shape) > 1 else 1
+    init = torch.zeros((k, dims)).to(data.device)
+    r = torch.distributions.uniform.Uniform(0, 1)
+    for i in range(k):
+        if i == 0:
+            init[i, :] = data[torch.randint(data.shape[0], [1])]
+        else:
+            D2 = torch.cdist(init[:i, :][None, :], data[None, :], p=2)[0].amin(dim=0)
+            probs = D2 / torch.sum(D2)
+            cumprobs = torch.cumsum(probs, dim=0)
+            init[i, :] = data[torch.searchsorted(cumprobs, r.sample([1]).to(data.device))]
+    return init
+class KMeansGPU:
+  '''
+  Kmeans clustering algorithm implemented with PyTorch
+
+  Parameters:
+    n_clusters: int, 
+      Number of clusters
+
+    max_iter: int, default: 100
+      Maximum number of iterations
+
+    tol: float, default: 0.0001
+      Tolerance
+    
+    verbose: int, default: 0
+      Verbosity
+
+    mode: {'euclidean', 'cosine'}, default: 'euclidean'
+      Type of distance measure
+      
+    init_method: {'random', 'point', '++'}
+      Type of initialization
+
+    minibatch: {None, int}, default: None
+      Batch size of MinibatchKmeans algorithm
+      if None perform full KMeans algorithm
+      
+  Attributes:
+    centroids: torch.Tensor, shape: [n_clusters, n_features]
+      cluster centroids
+  '''
+  def __init__(self, n_clusters, max_iter=200, tol=1e-4, verbose=0, mode="euclidean",device=torch.device("cuda:0")):
+    self.n_clusters = n_clusters
+    self.max_iter = max_iter
+    self.tol = tol
+    self.verbose = verbose
+    self.mode = mode
+    self.device=device
+    pynvml.nvmlInit()
+    gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device.index)
+    info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
+    self.minibatch=int(33e6/self.n_clusters*info.free/ 1024 / 1024 / 1024)
+    print("free_mem/GB:",info.free/ 1024 / 1024 / 1024,"minibatch:",self.minibatch)
+    
+  @staticmethod
+  def cos_sim(a, b):
+    """
+      Compute cosine similarity of 2 sets of vectors
+
+      Parameters:
+      a: torch.Tensor, shape: [m, n_features]
+
+      b: torch.Tensor, shape: [n, n_features]
+    """
+    return normalize(a, dim=-1) @ normalize(b, dim=-1).transpose(-2, -1)
+
+  @staticmethod
+  def euc_sim(a, b):
+    """
+      Compute euclidean similarity of 2 sets of vectors
+      Parameters:
+      a: torch.Tensor, shape: [m, n_features]
+      b: torch.Tensor, shape: [n, n_features]
+    """
+    return 2 * a @ b.transpose(-2, -1) -(a**2).sum(dim=1)[..., :, None] - (b**2).sum(dim=1)[..., None, :]
+
+  def max_sim(self, a, b):
+    """
+      Compute maximum similarity (or minimum distance) of each vector
+      in a with all of the vectors in b
+      Parameters:
+      a: torch.Tensor, shape: [m, n_features]
+      b: torch.Tensor, shape: [n, n_features]
+    """
+    if self.mode == 'cosine':
+      sim_func = self.cos_sim
+    elif self.mode == 'euclidean':
+      sim_func = self.euc_sim
+    sim = sim_func(a, b)
+    max_sim_v, max_sim_i = sim.max(dim=-1)
+    return max_sim_v, max_sim_i
+
+  def fit_predict(self, X):
+    """
+      Combination of fit() and predict() methods.
+      This is faster than calling fit() and predict() seperately.
+      Parameters:
+      X: torch.Tensor, shape: [n_samples, n_features]
+      centroids: {torch.Tensor, None}, default: None
+        if given, centroids will be initialized with given tensor
+        if None, centroids will be randomly chosen from X
+      Return:
+      labels: torch.Tensor, shape: [n_samples]
+
+            mini_=33kk/k*remain
+            mini=min(mini_,fea_shape)
+            offset=log2(k/1000)*1.5
+            kpp_all=min(mini_*10/offset,fea_shape)
+            kpp_sample=min(mini_/12/offset,fea_shape)
+    """
+    assert isinstance(X, torch.Tensor), "input must be torch.Tensor"
+    assert X.dtype in [torch.half, torch.float, torch.double], "input must be floating point"
+    assert X.ndim == 2, "input must be a 2d tensor with shape: [n_samples, n_features] "
+    # print("verbose:%s"%self.verbose)
+
+    offset = np.power(1.5,np.log(self.n_clusters / 1000))/np.log(2)
+    with torch.no_grad():
+      batch_size= X.shape[0]
+      # print(self.minibatch, int(self.minibatch * 10 / offset), batch_size)
+      start_time = time()
+      if (self.minibatch*10//offset< batch_size):
+        x = X[torch.randint(0, batch_size,[int(self.minibatch*10/offset)])].to(self.device)
+      else:
+        x = X.to(self.device)
+      # print(x.device)
+      self.centroids = _kpp(x, self.n_clusters, min(int(self.minibatch/12/offset),batch_size))
+      del x
+      torch.cuda.empty_cache()
+      # self.centroids = self.centroids.to(self.device)
+      num_points_in_clusters = torch.ones(self.n_clusters, device=self.device, dtype=X.dtype)#全1
+      closest = None#[3098036]#int64
+      if(self.minibatch>=batch_size//2 and self.minibatch<batch_size):
+        X = X[torch.randint(0, batch_size,[self.minibatch])].to(self.device)
+      elif(self.minibatch>=batch_size):
+        X=X.to(self.device)
+      for i in range(self.max_iter):
+        iter_time = time()
+        if self.minibatch<batch_size//2:#可用minibatch数太小，每次都得从内存倒腾到显存
+          x = X[torch.randint(0, batch_size, [self.minibatch])].to(self.device)
+        else:#否则直接全部缓存
+          x = X
+
+        closest = self.max_sim(a=x, b=self.centroids)[1].to(torch.int16)#[3098036]#int64#0~999
+        matched_clusters, counts = closest.unique(return_counts=True)#int64#1k
+        expanded_closest = closest[None].expand(self.n_clusters, -1)#[1000, 3098036]#int16#0~999
+        mask = (expanded_closest==torch.arange(self.n_clusters, device=self.device)[:, None]).to(X.dtype)#==后者是int64*1000
+        c_grad = mask @ x / mask.sum(-1)[..., :, None]
+        c_grad[c_grad!=c_grad] = 0 # remove NaNs
+        error = (c_grad - self.centroids).pow(2).sum()
+        if self.minibatch is not None:
+          lr = 1/num_points_in_clusters[:,None] * 0.9 + 0.1
+        else:
+          lr = 1
+        matched_clusters=matched_clusters.long()
+        num_points_in_clusters[matched_clusters] += counts#IndexError: tensors used as indices must be long, byte or bool tensors
+        self.centroids = self.centroids * (1-lr) + c_grad * lr
+        if self.verbose >= 2:
+          print('iter:', i, 'error:', error.item(), 'time spent:', round(time()-iter_time, 4))
+        if error <= self.tol:
+          break
+
+      if self.verbose >= 1:
+        print(f'used {i+1} iterations ({round(time()-start_time, 4)}s) to cluster {batch_size} items into {self.n_clusters} clusters')
+    return closest

cluster/train_cluster.py

@@ -1,19 +1,17 @@
-import time,pdb
-import tqdm
-from time import time as ttime
+import argparse
+import logging
 import os
+import time
 from pathlib import Path
-import logging
-import argparse
-from kmeans import KMeansGPU
-import torch
+
 import numpy as np
-from sklearn.cluster import KMeans,MiniBatchKMeans
+import torch
+import tqdm
+from kmeans import KMeansGPU
+from sklearn.cluster import KMeans, MiniBatchKMeans
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-from time import time as ttime
-import pynvml,torch
 
 def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False,use_gpu=False):#gpu_minibatch真拉，虽然库支持但是也不考虑
     logger.info(f"Loading features from {in_dir}")
@@ -29,7 +27,7 @@ def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False,use_gpu=
     features = features.astype(np.float32)
     logger.info(f"Clustering features of shape: {features.shape}")
     t = time.time()
-    if(use_gpu==False):
+    if(use_gpu is False):
         if use_minibatch:
             kmeans = MiniBatchKMeans(n_clusters=n_clusters,verbose=verbose, batch_size=4096, max_iter=80).fit(features)
         else:
@@ -37,14 +35,14 @@ def train_cluster(in_dir, n_clusters, use_minibatch=True, verbose=False,use_gpu=
     else:
             kmeans = KMeansGPU(n_clusters=n_clusters, mode='euclidean', verbose=2 if verbose else 0,max_iter=500,tol=1e-2)#
             features=torch.from_numpy(features)#.to(device)
-            labels = kmeans.fit_predict(features)#
+            kmeans.fit_predict(features)#
 
     print(time.time()-t, "s")
 
     x = {
-            "n_features_in_": kmeans.n_features_in_ if use_gpu==False else features.shape[1],
-            "_n_threads": kmeans._n_threads if use_gpu==False else 4,
-            "cluster_centers_": kmeans.cluster_centers_ if use_gpu==False else kmeans.centroids.cpu().numpy(),
+            "n_features_in_": kmeans.n_features_in_ if use_gpu is False else features.shape[1],
+            "_n_threads": kmeans._n_threads if use_gpu is False else 4,
+            "cluster_centers_": kmeans.cluster_centers_ if use_gpu is False else kmeans.centroids.cpu().numpy(),
     }
     print("end")
 

compress_model.py

@@ -0,0 +1,71 @@
+from collections import OrderedDict
+
+import torch
+
+import utils
+from models import SynthesizerTrn
+
+
+def copyStateDict(state_dict):
+    if list(state_dict.keys())[0].startswith('module'):
+        start_idx = 1
+    else:
+        start_idx = 0
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        name = ','.join(k.split('.')[start_idx:])
+        new_state_dict[name] = v
+    return new_state_dict
+
+
+def removeOptimizer(config: str, input_model: str, ishalf: bool, output_model: str):
+    hps = utils.get_hparams_from_file(config)
+
+    net_g = SynthesizerTrn(hps.data.filter_length // 2 + 1,
+                           hps.train.segment_size // hps.data.hop_length,
+                           **hps.model)
+
+    optim_g = torch.optim.AdamW(net_g.parameters(),
+                                hps.train.learning_rate,
+                                betas=hps.train.betas,
+                                eps=hps.train.eps)
+
+    state_dict_g = torch.load(input_model, map_location="cpu")
+    new_dict_g = copyStateDict(state_dict_g)
+    keys = []
+    for k, v in new_dict_g['model'].items():
+        keys.append(k)
+    
+    new_dict_g = {k: new_dict_g['model'][k].half() for k in keys} if ishalf else {k: new_dict_g['model'][k] for k in keys}
+
+    torch.save(
+        {
+            'model': new_dict_g,
+            'iteration': 0,
+            'optimizer': optim_g.state_dict(),
+            'learning_rate': 0.0001
+        }, output_model)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-c",
+                        "--config",
+                        type=str,
+                        default='configs/config.json')
+    parser.add_argument("-i", "--input", type=str)
+    parser.add_argument("-o", "--output", type=str, default=None)
+    parser.add_argument('-hf', '--half', action='store_true', default=False, help='Save as FP16')
+    
+    args = parser.parse_args()
+
+    output = args.output
+
+    if output is None:
+        import os.path
+        filename, ext = os.path.splitext(args.input)
+        half = "_half" if args.half else ""
+        output = filename + "_release" + half + ext
+
+    removeOptimizer(args.config, args.input, args.half, output)

configs/diffusion.yaml

@@ -0,0 +1,51 @@
+data:
+  block_size: 512
+  cnhubertsoft_gate: 10
+  duration: 2
+  encoder: vec768l12
+  encoder_hop_size: 320
+  encoder_out_channels: 768
+  encoder_sample_rate: 16000
+  extensions:
+  - wav
+  sampling_rate: 44100
+  training_files: filelists/train.txt
+  unit_interpolate_mode: nearest
+  validation_files: filelists/val.txt
+device: cuda
+env:
+  expdir: logs/44k/diffusion
+  gpu_id: 0
+infer:
+  method: dpm-solver++
+  speedup: 10
+model:
+  k_step_max: 0
+  n_chans: 512
+  n_hidden: 256
+  n_layers: 20
+  n_spk: 1
+  timesteps: 1000
+  type: Diffusion
+  use_pitch_aug: true
+spk:
+  ATRI: 0
+train:
+  amp_dtype: fp32
+  batch_size: 48
+  cache_all_data: true
+  cache_device: cpu
+  cache_fp16: true
+  decay_step: 100000
+  epochs: 100000
+  gamma: 0.5
+  interval_force_save: 10000
+  interval_log: 10
+  interval_val: 2000
+  lr: 0.0002
+  num_workers: 2
+  save_opt: false
+  weight_decay: 0
+vocoder:
+  ckpt: pretrain/nsf_hifigan/model
+  type: nsf-hifigan

config.json → configs_template/config_template.json

@@ -24,7 +24,7 @@
     "port": "8001",
     "keep_ckpts": 3,
     "all_in_mem": false,
-    "vol_aug": false
+    "vol_aug":false
   },
   "data": {
     "training_files": "filelists/train.txt",
@@ -37,7 +37,7 @@
     "n_mel_channels": 80,
     "mel_fmin": 0.0,
     "mel_fmax": 22050,
-    "unit_interpolate_mode": "nearest"
+    "unit_interpolate_mode":"nearest"
   },
   "model": {
     "inter_channels": 192,
@@ -48,58 +48,30 @@
     "kernel_size": 3,
     "p_dropout": 0.1,
     "resblock": "1",
-    "resblock_kernel_sizes": [
-      3,
-      7,
-      11
-    ],
-    "resblock_dilation_sizes": [
-      [
-        1,
-        3,
-        5
-      ],
-      [
-        1,
-        3,
-        5
-      ],
-      [
-        1,
-        3,
-        5
-      ]
-    ],
-    "upsample_rates": [
-      8,
-      8,
-      2,
-      2,
-      2
-    ],
+    "resblock_kernel_sizes": [3,7,11],
+    "resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
+    "upsample_rates": [ 8, 8, 2, 2, 2],
     "upsample_initial_channel": 512,
-    "upsample_kernel_sizes": [
-      16,
-      16,
-      4,
-      4,
-      4
-    ],
+    "upsample_kernel_sizes": [16,16, 4, 4, 4],
     "n_layers_q": 3,
     "n_flow_layer": 4,
     "use_spectral_norm": false,
     "gin_channels": 768,
     "ssl_dim": 768,
-    "n_speakers": 1,
-    "vocoder_name": "nsf-hifigan",
-    "speech_encoder": "vec768l12",
-    "speaker_embedding": false,
-    "vol_embedding": false,
-    "use_depthwise_conv": false,
+    "n_speakers": 200,
+    "vocoder_name":"nsf-hifigan",
+    "speech_encoder":"vec768l12",
+    "speaker_embedding":false,
+    "vol_embedding":false,
+    "use_depthwise_conv":false,
     "flow_share_parameter": false,
     "use_automatic_f0_prediction": true
   },
   "spk": {
-    "ATRI": 0
+    "nyaru": 0,
+    "huiyu": 1,
+    "nen": 2,
+    "paimon": 3,
+    "yunhao": 4
   }
 }

configs_template/diffusion_template.yaml

@@ -0,0 +1,51 @@
+data:
+  sampling_rate: 44100
+  block_size: 512 # Equal to hop_length
+  duration: 2 # Audio duration during training, must be less than the duration of the shortest audio clip
+  encoder: 'vec768l12' # 'hubertsoft', 'vec256l9', 'vec768l12'
+  cnhubertsoft_gate: 10
+  encoder_sample_rate: 16000
+  encoder_hop_size: 320
+  encoder_out_channels: 768 # 256 if using 'hubertsoft'
+  training_files: "filelists/train.txt"
+  validation_files: "filelists/val.txt"
+  extensions: # List of extension included in the data collection
+    - wav
+  unit_interpolate_mode: "nearest"
+model:
+  type: 'Diffusion'
+  n_layers: 20
+  n_chans: 512
+  n_hidden: 256
+  use_pitch_aug: true
+  timesteps : 1000
+  k_step_max: 0 # must <= timesteps, If it is 0, train all
+  n_spk: 1 # max number of different speakers
+device: cuda
+vocoder:
+  type: 'nsf-hifigan'
+  ckpt: 'pretrain/nsf_hifigan/model'
+infer:
+  speedup: 10
+  method: 'dpm-solver++' # 'pndm' or 'dpm-solver' or 'ddim' or 'unipc' or 'dpm-solver++'
+env:
+  expdir: logs/44k/diffusion
+  gpu_id: 0
+train:
+  num_workers: 2 # If your cpu and gpu are both very strong, set to 0 may be faster!
+  amp_dtype: fp32 # fp32, fp16 or bf16 (fp16 or bf16 may be faster if it is supported by your gpu)
+  batch_size: 48
+  cache_all_data: true # Save Internal-Memory or Graphics-Memory if it is false, but may be slow
+  cache_device: 'cpu' # Set to 'cuda' to cache the data into the Graphics-Memory, fastest speed for strong gpu
+  cache_fp16: true
+  epochs: 100000
+  interval_log: 10
+  interval_val: 2000
+  interval_force_save: 10000
+  lr: 0.0002
+  decay_step: 100000
+  gamma: 0.5
+  weight_decay: 0
+  save_opt: false
+spk:
+  'nyaru': 0

data_utils.py

@@ -0,0 +1,184 @@
+import os
+import random
+
+import numpy as np
+import torch
+import torch.utils.data
+
+import utils
+from modules.mel_processing import spectrogram_torch
+from utils import load_filepaths_and_text, load_wav_to_torch
+
+# import h5py
+
+
+"""Multi speaker version"""
+
+
+class TextAudioSpeakerLoader(torch.utils.data.Dataset):
+    """
+        1) loads audio, speaker_id, text pairs
+        2) normalizes text and converts them to sequences of integers
+        3) computes spectrograms from audio files.
+    """
+
+    def __init__(self, audiopaths, hparams, all_in_mem: bool = False, vol_aug: bool = True):
+        self.audiopaths = load_filepaths_and_text(audiopaths)
+        self.hparams = hparams
+        self.max_wav_value = hparams.data.max_wav_value
+        self.sampling_rate = hparams.data.sampling_rate
+        self.filter_length = hparams.data.filter_length
+        self.hop_length = hparams.data.hop_length
+        self.win_length = hparams.data.win_length
+        self.unit_interpolate_mode = hparams.data.unit_interpolate_mode
+        self.sampling_rate = hparams.data.sampling_rate
+        self.use_sr = hparams.train.use_sr
+        self.spec_len = hparams.train.max_speclen
+        self.spk_map = hparams.spk
+        self.vol_emb = hparams.model.vol_embedding
+        self.vol_aug = hparams.train.vol_aug and vol_aug
+        random.seed(1234)
+        random.shuffle(self.audiopaths)
+        
+        self.all_in_mem = all_in_mem
+        if self.all_in_mem:
+            self.cache = [self.get_audio(p[0]) for p in self.audiopaths]
+
+    def get_audio(self, filename):
+        filename = filename.replace("\\", "/")
+        audio, sampling_rate = load_wav_to_torch(filename)
+        if sampling_rate != self.sampling_rate:
+            raise ValueError("{} SR doesn't match target {} SR".format(
+                sampling_rate, self.sampling_rate))
+        audio_norm = audio / self.max_wav_value
+        audio_norm = audio_norm.unsqueeze(0)
+        spec_filename = filename.replace(".wav", ".spec.pt")
+
+        # Ideally, all data generated after Mar 25 should have .spec.pt
+        if os.path.exists(spec_filename):
+            spec = torch.load(spec_filename)
+        else:
+            spec = spectrogram_torch(audio_norm, self.filter_length,
+                                     self.sampling_rate, self.hop_length, self.win_length,
+                                     center=False)
+            spec = torch.squeeze(spec, 0)
+            torch.save(spec, spec_filename)
+
+        spk = filename.split("/")[-2]
+        spk = torch.LongTensor([self.spk_map[spk]])
+
+        f0, uv = np.load(filename + ".f0.npy",allow_pickle=True)
+        
+        f0 = torch.FloatTensor(np.array(f0,dtype=float))
+        uv = torch.FloatTensor(np.array(uv,dtype=float))
+
+        c = torch.load(filename+ ".soft.pt")
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[0], mode=self.unit_interpolate_mode)
+        if self.vol_emb:
+            volume_path = filename + ".vol.npy"
+            volume = np.load(volume_path)
+            volume = torch.from_numpy(volume).float()
+        else:
+            volume = None
+
+        lmin = min(c.size(-1), spec.size(-1))
+        assert abs(c.size(-1) - spec.size(-1)) < 3, (c.size(-1), spec.size(-1), f0.shape, filename)
+        assert abs(audio_norm.shape[1]-lmin * self.hop_length) < 3 * self.hop_length
+        spec, c, f0, uv = spec[:, :lmin], c[:, :lmin], f0[:lmin], uv[:lmin]
+        audio_norm = audio_norm[:, :lmin * self.hop_length]
+        if volume is not None:
+            volume = volume[:lmin]
+        return c, f0, spec, audio_norm, spk, uv, volume
+
+    def random_slice(self, c, f0, spec, audio_norm, spk, uv, volume):
+        # if spec.shape[1] < 30:
+        #     print("skip too short audio:", filename)
+        #     return None
+
+        if random.choice([True, False]) and self.vol_aug and volume is not None:
+            max_amp = float(torch.max(torch.abs(audio_norm))) + 1e-5
+            max_shift = min(1, np.log10(1/max_amp))
+            log10_vol_shift = random.uniform(-1, max_shift)
+            audio_norm = audio_norm * (10 ** log10_vol_shift)
+            volume = volume * (10 ** log10_vol_shift)
+            spec = spectrogram_torch(audio_norm,
+            self.hparams.data.filter_length,
+            self.hparams.data.sampling_rate,
+            self.hparams.data.hop_length,
+            self.hparams.data.win_length,
+            center=False)[0]
+
+        if spec.shape[1] > 800:
+            start = random.randint(0, spec.shape[1]-800)
+            end = start + 790
+            spec, c, f0, uv = spec[:, start:end], c[:, start:end], f0[start:end], uv[start:end]
+            audio_norm = audio_norm[:, start * self.hop_length : end * self.hop_length]
+            if volume is not None:
+                volume = volume[start:end]
+        return c, f0, spec, audio_norm, spk, uv,volume
+
+    def __getitem__(self, index):
+        if self.all_in_mem:
+            return self.random_slice(*self.cache[index])
+        else:
+            return self.random_slice(*self.get_audio(self.audiopaths[index][0]))
+
+    def __len__(self):
+        return len(self.audiopaths)
+
+
+class TextAudioCollate:
+
+    def __call__(self, batch):
+        batch = [b for b in batch if b is not None]
+
+        input_lengths, ids_sorted_decreasing = torch.sort(
+            torch.LongTensor([x[0].shape[1] for x in batch]),
+            dim=0, descending=True)
+
+        max_c_len = max([x[0].size(1) for x in batch])
+        max_wav_len = max([x[3].size(1) for x in batch])
+
+        lengths = torch.LongTensor(len(batch))
+
+        c_padded = torch.FloatTensor(len(batch), batch[0][0].shape[0], max_c_len)
+        f0_padded = torch.FloatTensor(len(batch), max_c_len)
+        spec_padded = torch.FloatTensor(len(batch), batch[0][2].shape[0], max_c_len)
+        wav_padded = torch.FloatTensor(len(batch), 1, max_wav_len)
+        spkids = torch.LongTensor(len(batch), 1)
+        uv_padded = torch.FloatTensor(len(batch), max_c_len)
+        volume_padded = torch.FloatTensor(len(batch), max_c_len)
+
+        c_padded.zero_()
+        spec_padded.zero_()
+        f0_padded.zero_()
+        wav_padded.zero_()
+        uv_padded.zero_()
+        volume_padded.zero_()
+
+        for i in range(len(ids_sorted_decreasing)):
+            row = batch[ids_sorted_decreasing[i]]
+
+            c = row[0]
+            c_padded[i, :, :c.size(1)] = c
+            lengths[i] = c.size(1)
+
+            f0 = row[1]
+            f0_padded[i, :f0.size(0)] = f0
+
+            spec = row[2]
+            spec_padded[i, :, :spec.size(1)] = spec
+
+            wav = row[3]
+            wav_padded[i, :, :wav.size(1)] = wav
+
+            spkids[i, 0] = row[4]
+
+            uv = row[5]
+            uv_padded[i, :uv.size(0)] = uv
+            volume = row[6]
+            if volume is not None:
+                volume_padded[i, :volume.size(0)] = volume
+            else :
+                volume_padded = None
+        return c_padded, f0_padded, spec_padded, wav_padded, spkids, lengths, uv_padded, volume_padded

diffusion/data_loaders.py

@@ -1,13 +1,14 @@
 import os
 import random
-import re
-import numpy as np
+
 import librosa
+import numpy as np
 import torch
-import random
-from utils import repeat_expand_2d
-from tqdm import tqdm
 from torch.utils.data import Dataset
+from tqdm import tqdm
+
+from utils import repeat_expand_2d
+
 
 def traverse_dir(
         root_dir,
@@ -63,6 +64,7 @@ def get_data_loaders(args, whole_audio=False):
         spk=args.spk,
         device=args.train.cache_device,
         fp16=args.train.cache_fp16,
+        unit_interpolate_mode = args.data.unit_interpolate_mode,
         use_aug=True)
     loader_train = torch.utils.data.DataLoader(
         data_train ,
@@ -81,6 +83,7 @@ def get_data_loaders(args, whole_audio=False):
         whole_audio=True,
         spk=args.spk,
         extensions=args.data.extensions,
+        unit_interpolate_mode = args.data.unit_interpolate_mode,
         n_spk=args.model.n_spk)
     loader_valid = torch.utils.data.DataLoader(
         data_valid,
@@ -107,6 +110,7 @@ class AudioDataset(Dataset):
         device='cpu',
         fp16=False,
         use_aug=False,
+        unit_interpolate_mode = 'left'
     ):
         super().__init__()
         
@@ -118,6 +122,7 @@ class AudioDataset(Dataset):
         self.use_aug = use_aug
         self.data_buffer={}
         self.pitch_aug_dict = {}
+        self.unit_interpolate_mode = unit_interpolate_mode
         # np.load(os.path.join(self.path_root, 'pitch_aug_dict.npy'), allow_pickle=True).item()
         if load_all_data:
             print('Load all the data filelists:', filelists)
@@ -126,7 +131,6 @@ class AudioDataset(Dataset):
         with open(filelists,"r") as f:
             self.paths = f.read().splitlines()
         for name_ext in tqdm(self.paths, total=len(self.paths)):
-            name = os.path.splitext(name_ext)[0]
             path_audio = name_ext
             duration = librosa.get_duration(filename = path_audio, sr = self.sample_rate)
             
@@ -171,7 +175,7 @@ class AudioDataset(Dataset):
                 path_units = name_ext + ".soft.pt"
                 units = torch.load(path_units).to(device)
                 units = units[0]  
-                units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
+                units = repeat_expand_2d(units,f0.size(0),unit_interpolate_mode).transpose(0,1)
                 
                 if fp16:
                     mel = mel.half()
@@ -263,7 +267,7 @@ class AudioDataset(Dataset):
             path_units = name_ext + ".soft.pt"
             units = torch.load(path_units)
             units = units[0]  
-            units = repeat_expand_2d(units,f0.size(0)).transpose(0,1)
+            units = repeat_expand_2d(units,f0.size(0),self.unit_interpolate_mode).transpose(0,1)
             
         units = units[start_frame : start_frame + units_frame_len]
 

diffusion/diffusion.py

@@ -1,10 +1,10 @@
 from collections import deque
 from functools import partial
 from inspect import isfunction
-import torch.nn.functional as F
-import librosa.sequence
+
 import numpy as np
 import torch
+import torch.nn.functional as F
 from torch import nn
 from tqdm import tqdm
 
@@ -26,8 +26,10 @@ def extract(a, t, x_shape):
 
 
 def noise_like(shape, device, repeat=False):
-    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    noise = lambda: torch.randn(shape, device=device)
+    def repeat_noise():
+        return torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    def noise():
+        return torch.randn(shape, device=device)
     return repeat_noise() if repeat else noise()
 
 
@@ -67,6 +69,7 @@ class GaussianDiffusion(nn.Module):
                 max_beta=0.02,
                 spec_min=-12, 
                 spec_max=2):
+        
         super().__init__()
         self.denoise_fn = denoise_fn
         self.out_dims = out_dims
@@ -78,7 +81,7 @@ class GaussianDiffusion(nn.Module):
 
         timesteps, = betas.shape
         self.num_timesteps = int(timesteps)
-        self.k_step = k_step
+        self.k_step = k_step if k_step>0 and k_step<timesteps else timesteps
 
         self.noise_list = deque(maxlen=4)
 
@@ -139,6 +142,18 @@ class GaussianDiffusion(nn.Module):
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
         return model_mean, posterior_variance, posterior_log_variance
 
+    @torch.no_grad()
+    def p_sample_ddim(self, x, t, interval, cond):
+        """
+        Use the DDIM method from
+        """
+        a_t = extract(self.alphas_cumprod, t, x.shape)
+        a_prev = extract(self.alphas_cumprod, torch.max(t - interval, torch.zeros_like(t)), x.shape)
+
+        noise_pred = self.denoise_fn(x, t, cond=cond)
+        x_prev = a_prev.sqrt() * (x / a_t.sqrt() + (((1 - a_prev) / a_prev).sqrt()-((1 - a_t) / a_t).sqrt()) * noise_pred)
+        return x_prev
+
     @torch.no_grad()
     def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
         b, *_, device = *x.shape, x.device
@@ -239,8 +254,12 @@ class GaussianDiffusion(nn.Module):
                 x = self.q_sample(x_start=norm_spec, t=torch.tensor([t - 1], device=device).long())
                         
             if method is not None and infer_speedup > 1:
-                if method == 'dpm-solver':
-                    from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
+                if method == 'dpm-solver' or method == 'dpm-solver++':
+                    from .dpm_solver_pytorch import (
+                        DPM_Solver,
+                        NoiseScheduleVP,
+                        model_wrapper,
+                    )
                     # 1. Define the noise schedule.
                     noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
 
@@ -267,17 +286,20 @@ class GaussianDiffusion(nn.Module):
                     # (We recommend singlestep DPM-Solver for unconditional sampling)
                     # You can adjust the `steps` to balance the computation
                     # costs and the sample quality.
-                    dpm_solver = DPM_Solver(model_fn, noise_schedule)
-
+                    if method == 'dpm-solver':
+                        dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver")
+                    elif method == 'dpm-solver++':
+                        dpm_solver = DPM_Solver(model_fn, noise_schedule, algorithm_type="dpmsolver++")
+                        
                     steps = t // infer_speedup
                     if use_tqdm:
                         self.bar = tqdm(desc="sample time step", total=steps)
                     x = dpm_solver.sample(
                         x,
                         steps=steps,
-                        order=3,
+                        order=2,
                         skip_type="time_uniform",
-                        method="singlestep",
+                        method="multistep",
                     )
                     if use_tqdm:
                         self.bar.close()
@@ -298,6 +320,63 @@ class GaussianDiffusion(nn.Module):
                                 x, torch.full((b,), i, device=device, dtype=torch.long),
                                 infer_speedup, cond=cond
                             )
+                elif method == 'ddim':
+                    if use_tqdm:
+                        for i in tqdm(
+                                reversed(range(0, t, infer_speedup)), desc='sample time step',
+                                total=t // infer_speedup,
+                        ):
+                            x = self.p_sample_ddim(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                    else:
+                        for i in reversed(range(0, t, infer_speedup)):
+                            x = self.p_sample_ddim(
+                                x, torch.full((b,), i, device=device, dtype=torch.long),
+                                infer_speedup, cond=cond
+                            )
+                elif method == 'unipc':
+                    from .uni_pc import NoiseScheduleVP, UniPC, model_wrapper
+                    # 1. Define the noise schedule.
+                    noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
+
+                    # 2. Convert your discrete-time `model` to the continuous-time
+                    # noise prediction model. Here is an example for a diffusion model
+                    # `model` with the noise prediction type ("noise") .
+                    def my_wrapper(fn):
+                        def wrapped(x, t, **kwargs):
+                            ret = fn(x, t, **kwargs)
+                            if use_tqdm:
+                                self.bar.update(1)
+                            return ret
+
+                        return wrapped
+
+                    model_fn = model_wrapper(
+                        my_wrapper(self.denoise_fn),
+                        noise_schedule,
+                        model_type="noise",  # or "x_start" or "v" or "score"
+                        model_kwargs={"cond": cond}
+                    )
+
+                    # 3. Define uni_pc and sample by multistep UniPC.
+                    # You can adjust the `steps` to balance the computation
+                    # costs and the sample quality.
+                    uni_pc = UniPC(model_fn, noise_schedule, variant='bh2')
+
+                    steps = t // infer_speedup
+                    if use_tqdm:
+                        self.bar = tqdm(desc="sample time step", total=steps)
+                    x = uni_pc.sample(
+                        x,
+                        steps=steps,
+                        order=2,
+                        skip_type="time_uniform",
+                        method="multistep",
+                    )
+                    if use_tqdm:
+                        self.bar.close()
                 else:
                     raise NotImplementedError(method)
             else:

diffusion/diffusion_onnx.py

@@ -1,15 +1,14 @@
+import math
 from collections import deque
 from functools import partial
 from inspect import isfunction
-import torch.nn.functional as F
-import librosa.sequence
+
 import numpy as np
-from torch.nn import Conv1d
-from torch.nn import Mish
 import torch
+import torch.nn.functional as F
 from torch import nn
+from torch.nn import Conv1d, Mish
 from tqdm import tqdm
-import math
 
 
 def exists(x):
@@ -27,8 +26,10 @@ def extract(a, t):
 
 
 def noise_like(shape, device, repeat=False):
-    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
-    noise = lambda: torch.randn(shape, device=device)
+    def repeat_noise():
+        return torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    def noise():
+        return torch.randn(shape, device=device)
     return repeat_noise() if repeat else noise()
 
 
@@ -389,7 +390,11 @@ class GaussianDiffusion(nn.Module):
                         
             if method is not None and infer_speedup > 1:
                 if method == 'dpm-solver':
-                    from .dpm_solver_pytorch import NoiseScheduleVP, model_wrapper, DPM_Solver
+                    from .dpm_solver_pytorch import (
+                        DPM_Solver,
+                        NoiseScheduleVP,
+                        model_wrapper,
+                    )
                     # 1. Define the noise schedule.
                     noise_schedule = NoiseScheduleVP(schedule='discrete', betas=self.betas[:t])
 
@@ -576,9 +581,6 @@ class GaussianDiffusion(nn.Module):
         plms_noise_stage = torch.tensor(0, dtype=torch.long, device=device)
         noise_list = torch.zeros((0, 1, 1, self.mel_bins, n_frames), device=device)
 
-        ot = step_range[0]
-        ot_1 = torch.full((1,), ot, device=device, dtype=torch.long)
-
         for t in step_range:
             t_1 = torch.full((1,), t, device=device, dtype=torch.long)
             noise_pred = self.denoise_fn(x, t_1, cond)

diffusion/dpm_solver_pytorch.py

@@ -1,5 +1,3 @@
-import math
-
 import torch
 
 
@@ -11,7 +9,8 @@ class NoiseScheduleVP:
             alphas_cumprod=None,
             continuous_beta_0=0.1,
             continuous_beta_1=20.,
-    ):
+            dtype=torch.float32,
+        ):
         """Create a wrapper class for the forward SDE (VP type).
 
         ***
@@ -46,7 +45,7 @@ class NoiseScheduleVP:
                 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(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
+            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
@@ -59,21 +58,19 @@ class NoiseScheduleVP:
 
         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:
+            We support the linear VPSDE for the continuous time setting. The hyperparameters for the noise
+            schedule are the default settings in Yang Song's ScoreSDE:
 
             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.
+                    'linear' for continuous-time DPMs.
         Returns:
             A wrapper object of the forward SDE (VP type).
         
@@ -92,10 +89,8 @@ class NoiseScheduleVP:
 
         """
 
-        if schedule not in ['discrete', 'linear', 'cosine']:
-            raise ValueError(
-                "Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(
-                    schedule))
+        if schedule not in ['discrete', 'linear']:
+            raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear'".format(schedule))
 
         self.schedule = schedule
         if schedule == 'discrete':
@@ -104,40 +99,37 @@ class NoiseScheduleVP:
             else:
                 assert alphas_cumprod is not None
                 log_alphas = 0.5 * torch.log(alphas_cumprod)
-            self.total_N = len(log_alphas)
             self.T = 1.
-            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
-            self.log_alpha_array = log_alphas.reshape((1, -1,))
+            self.log_alpha_array = self.numerical_clip_alpha(log_alphas).reshape((1, -1,)).to(dtype=dtype)
+            self.total_N = self.log_alpha_array.shape[1]
+            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1)).to(dtype=dtype)
         else:
+            self.T = 1.
             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.
-            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (
-                        1. + self.cosine_s) / math.pi - self.cosine_s
-            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
-            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.
+
+    def numerical_clip_alpha(self, log_alphas, clipped_lambda=-5.1):
+        """
+        For some beta schedules such as cosine schedule, the log-SNR has numerical isssues. 
+        We clip the log-SNR near t=T within -5.1 to ensure the stability.
+        Such a trick is very useful for diffusion models with the cosine schedule, such as i-DDPM, guided-diffusion and GLIDE.
+        """
+        log_sigmas = 0.5 * torch.log(1. - torch.exp(2. * log_alphas))
+        lambs = log_alphas - log_sigmas  
+        idx = torch.searchsorted(torch.flip(lambs, [0]), clipped_lambda)
+        if idx > 0:
+            log_alphas = log_alphas[:-idx]
+        return log_alphas
 
     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))
+            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':
-            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
-            log_alpha_t = log_alpha_fn(t) - self.cosine_log_alpha_0
-            return log_alpha_t
 
     def marginal_alpha(self, t):
         """
@@ -165,32 +157,25 @@ class NoiseScheduleVP:
         """
         if self.schedule == 'linear':
             tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
-            Delta = self.beta_0 ** 2 + tmp
+            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. * 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]))
+            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. * lamb, torch.zeros((1,)).to(lamb))
-            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (
-                        1. + 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.,
-        classifier_fn=None,
-        classifier_kwargs={},
+    model,
+    noise_schedule,
+    model_type="noise",
+    model_kwargs={},
+    guidance_type="uncond",
+    condition=None,
+    unconditional_condition=None,
+    guidance_scale=1.,
+    classifier_fn=None,
+    classifier_kwargs={},
 ):
     """Create a wrapper function for the noise prediction model.
 
@@ -293,8 +278,6 @@ def model_wrapper(
             return t_continuous
 
     def noise_pred_fn(x, t_continuous, cond=None):
-        if t_continuous.reshape((-1,)).shape[0] == 1:
-            t_continuous = t_continuous.expand((x.shape[0]))
         t_input = get_model_input_time(t_continuous)
         if cond is None:
             output = model(x, t_input, **model_kwargs)
@@ -304,16 +287,13 @@ def model_wrapper(
             return output
         elif model_type == "x_start":
             alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
+            return (x - expand_dims(alpha_t, x.dim()) * output) / expand_dims(sigma_t, x.dim())
         elif model_type == "v":
             alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
+            return expand_dims(alpha_t, x.dim()) * output + expand_dims(sigma_t, x.dim()) * x
         elif model_type == "score":
             sigma_t = noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return -expand_dims(sigma_t, dims) * output
+            return -expand_dims(sigma_t, x.dim()) * output
 
     def cond_grad_fn(x, t_input):
         """
@@ -328,8 +308,6 @@ def model_wrapper(
         """
         The noise predicition model function that is used for DPM-Solver.
         """
-        if t_continuous.reshape((-1,)).shape[0] == 1:
-            t_continuous = t_continuous.expand((x.shape[0]))
         if guidance_type == "uncond":
             return noise_pred_fn(x, t_continuous)
         elif guidance_type == "classifier":
@@ -338,7 +316,7 @@ def model_wrapper(
             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 * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
+            return noise - guidance_scale * expand_dims(sigma_t, x.dim()) * cond_grad
         elif guidance_type == "classifier-free":
             if guidance_scale == 1. or unconditional_condition is None:
                 return noise_pred_fn(x, t_continuous, cond=condition)
@@ -349,20 +327,34 @@ def model_wrapper(
                 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"]
+    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, predict_x0=False, thresholding=False, max_val=1.):
+    def __init__(
+        self,
+        model_fn,
+        noise_schedule,
+        algorithm_type="dpmsolver++",
+        correcting_x0_fn=None,
+        correcting_xt_fn=None,
+        thresholding_max_val=1.,
+        dynamic_thresholding_ratio=0.995,
+    ):
         """Construct a DPM-Solver. 
 
-        We support both the noise prediction model ("predicting epsilon") and the data prediction model ("predicting x0").
-        If `predict_x0` is False, we use the solver for the noise prediction model (DPM-Solver).
-        If `predict_x0` is True, we use the solver for the data prediction model (DPM-Solver++).
-            In such case, we further support the "dynamic thresholding" in [1] when `thresholding` is True.
-            The "dynamic thresholding" can greatly improve the sample quality for pixel-space DPMs with large guidance scales.
+        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]):
@@ -370,18 +362,65 @@ class DPM_Solver:
                 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.
-            predict_x0: A `bool`. If true, use the data prediction model; else, use the noise prediction model.
-            thresholding: A `bool`. Valid when `predict_x0` is True. Whether to use the "dynamic thresholding" in [1].
-            max_val: A `float`. Valid when both `predict_x0` and `thresholding` are True. The max value for 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.
+            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 = model_fn
+        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
         self.noise_schedule = noise_schedule
-        self.predict_x0 = predict_x0
-        self.thresholding = thresholding
-        self.max_val = max_val
+        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):
         """
@@ -391,24 +430,20 @@ class DPM_Solver:
 
     def data_prediction_fn(self, x, t):
         """
-        Return the data prediction model (with thresholding).
+        Return the data prediction model (with corrector).
         """
         noise = self.noise_prediction_fn(x, t)
-        dims = x.dim()
         alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
-        x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
-        if self.thresholding:
-            p = 0.995  # A hyperparameter in the paper of "Imagen" [1].
-            s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
-            s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
-            x0 = torch.clamp(x0, -s, s) / s
+        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.predict_x0:
+        if self.algorithm_type == "dpmsolver++":
             return self.data_prediction_fn(x, t)
         else:
             return self.noise_prediction_fn(x, t)
@@ -437,11 +472,10 @@ class DPM_Solver:
             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. / t_order), t_0 ** (1. / t_order), N + 1).pow(t_order).to(device)
+            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / 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))
+            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):
         """
@@ -478,32 +512,31 @@ class DPM_Solver:
         if order == 3:
             K = steps // 3 + 1
             if steps % 3 == 0:
-                orders = [3, ] * (K - 2) + [2, 1]
+                orders = [3,] * (K - 2) + [2, 1]
             elif steps % 3 == 1:
-                orders = [3, ] * (K - 1) + [1]
+                orders = [3,] * (K - 1) + [1]
             else:
-                orders = [3, ] * (K - 1) + [2]
+                orders = [3,] * (K - 1) + [2]
         elif order == 2:
             if steps % 2 == 0:
                 K = steps // 2
-                orders = [2, ] * K
+                orders = [2,] * K
             else:
                 K = steps // 2 + 1
-                orders = [2, ] * (K - 1) + [1]
+                orders = [2,] * (K - 1) + [1]
         elif order == 1:
             K = 1
-            orders = [1, ] * steps
+            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), dim=0).to(device)]
+            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_fn(self, x, s):
+    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. 
         """
@@ -515,8 +548,8 @@ class DPM_Solver:
 
         Args:
             x: A pytorch tensor. The initial value at time `s`.
-            s: A pytorch tensor. The starting time, with the shape (x.shape[0],).
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            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`.
@@ -524,20 +557,19 @@ class DPM_Solver:
             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.predict_x0:
+        if self.algorithm_type == "dpmsolver++":
             phi_1 = torch.expm1(-h)
             if model_s is None:
                 model_s = self.model_fn(x, s)
             x_t = (
-                    expand_dims(sigma_t / sigma_s, dims) * x
-                    - expand_dims(alpha_t * phi_1, dims) * model_s
+                sigma_t / sigma_s * x
+                - alpha_t * phi_1 * model_s
             )
             if return_intermediate:
                 return x_t, {'model_s': model_s}
@@ -548,70 +580,66 @@ class DPM_Solver:
             if model_s is None:
                 model_s = self.model_fn(x, s)
             x_t = (
-                    expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
-                    - expand_dims(sigma_t * phi_1, dims) * model_s
+                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='dpm_solver'):
+    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 (x.shape[0],).
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            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 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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 ['dpm_solver', 'taylor']:
-            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        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
-        dims = x.dim()
         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)
+        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.predict_x0:
+        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 = (
-                    expand_dims(sigma_s1 / sigma_s, dims) * x
-                    - expand_dims(alpha_s1 * phi_11, dims) * model_s
+                (sigma_s1 / sigma_s) * x
+                - (alpha_s1 * phi_11) * model_s
             )
             model_s1 = self.model_fn(x_s1, s1)
-            if solver_type == 'dpm_solver':
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(sigma_t / sigma_s, dims) * x
-                        - expand_dims(alpha_t * phi_1, dims) * model_s
-                        - (0.5 / r1) * expand_dims(alpha_t * phi_1, dims) * (model_s1 - model_s)
+                    (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 = (
-                        expand_dims(sigma_t / sigma_s, dims) * x
-                        - expand_dims(alpha_t * phi_1, dims) * model_s
-                        + (1. / r1) * expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * (
-                                    model_s1 - model_s)
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1. / r1) * (alpha_t * (phi_1 / h + 1.)) * (model_s1 - model_s)
                 )
         else:
             phi_11 = torch.expm1(r1 * h)
@@ -620,36 +648,35 @@ class DPM_Solver:
             if model_s is None:
                 model_s = self.model_fn(x, s)
             x_s1 = (
-                    expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
-                    - expand_dims(sigma_s1 * phi_11, dims) * model_s
+                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 == 'dpm_solver':
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
-                        - expand_dims(sigma_t * phi_1, dims) * model_s
-                        - (0.5 / r1) * expand_dims(sigma_t * phi_1, dims) * (model_s1 - model_s)
+                    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 = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
-                        - expand_dims(sigma_t * phi_1, dims) * model_s
-                        - (1. / r1) * expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * (model_s1 - model_s)
+                    torch.exp(log_alpha_t - log_alpha_s) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1. / r1) * (sigma_t * (phi_1 / h - 1.)) * (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. / 3., r2=2. / 3., model_s=None, model_s1=None,
-                                           return_intermediate=False, solver_type='dpm_solver'):
+    def singlestep_dpm_solver_third_update(self, x, s, t, r1=1./3., r2=2./3., 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 (x.shape[0],).
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            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`.
@@ -657,32 +684,29 @@ class DPM_Solver:
             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 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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 ['dpm_solver', 'taylor']:
-            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        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. / 3.
         if r2 is None:
             r2 = 2. / 3.
         ns = self.noise_schedule
-        dims = x.dim()
         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)
+        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.predict_x0:
+        if self.algorithm_type == "dpmsolver++":
             phi_11 = torch.expm1(-r1 * h)
             phi_12 = torch.expm1(-r2 * h)
             phi_1 = torch.expm1(-h)
@@ -694,21 +718,21 @@ class DPM_Solver:
                 model_s = self.model_fn(x, s)
             if model_s1 is None:
                 x_s1 = (
-                        expand_dims(sigma_s1 / sigma_s, dims) * x
-                        - expand_dims(alpha_s1 * phi_11, dims) * model_s
+                    (sigma_s1 / sigma_s) * x
+                    - (alpha_s1 * phi_11) * model_s
                 )
                 model_s1 = self.model_fn(x_s1, s1)
             x_s2 = (
-                    expand_dims(sigma_s2 / sigma_s, dims) * x
-                    - expand_dims(alpha_s2 * phi_12, dims) * model_s
-                    + r2 / r1 * expand_dims(alpha_s2 * phi_22, dims) * (model_s1 - model_s)
+                (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 == 'dpm_solver':
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(sigma_t / sigma_s, dims) * x
-                        - expand_dims(alpha_t * phi_1, dims) * model_s
-                        + (1. / r2) * expand_dims(alpha_t * phi_2, dims) * (model_s2 - model_s)
+                    (sigma_t / sigma_s) * x
+                    - (alpha_t * phi_1) * model_s
+                    + (1. / r2) * (alpha_t * phi_2) * (model_s2 - model_s)
                 )
             elif solver_type == 'taylor':
                 D1_0 = (1. / r1) * (model_s1 - model_s)
@@ -716,10 +740,10 @@ class DPM_Solver:
                 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
                 D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
                 x_t = (
-                        expand_dims(sigma_t / sigma_s, dims) * x
-                        - expand_dims(alpha_t * phi_1, dims) * model_s
-                        + expand_dims(alpha_t * phi_2, dims) * D1
-                        - expand_dims(alpha_t * phi_3, dims) * D2
+                    (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)
@@ -733,21 +757,21 @@ class DPM_Solver:
                 model_s = self.model_fn(x, s)
             if model_s1 is None:
                 x_s1 = (
-                        expand_dims(torch.exp(log_alpha_s1 - log_alpha_s), dims) * x
-                        - expand_dims(sigma_s1 * phi_11, dims) * model_s
+                    (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 = (
-                    expand_dims(torch.exp(log_alpha_s2 - log_alpha_s), dims) * x
-                    - expand_dims(sigma_s2 * phi_12, dims) * model_s
-                    - r2 / r1 * expand_dims(sigma_s2 * phi_22, dims) * (model_s1 - model_s)
+                (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 == 'dpm_solver':
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
-                        - expand_dims(sigma_t * phi_1, dims) * model_s
-                        - (1. / r2) * expand_dims(sigma_t * phi_2, dims) * (model_s2 - model_s)
+                    (torch.exp(log_alpha_t - log_alpha_s)) * x
+                    - (sigma_t * phi_1) * model_s
+                    - (1. / r2) * (sigma_t * phi_2) * (model_s2 - model_s)
                 )
             elif solver_type == 'taylor':
                 D1_0 = (1. / r1) * (model_s1 - model_s)
@@ -755,10 +779,10 @@ class DPM_Solver:
                 D1 = (r2 * D1_0 - r1 * D1_1) / (r2 - r1)
                 D2 = 2. * (D1_1 - D1_0) / (r2 - r1)
                 x_t = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_s), dims) * x
-                        - expand_dims(sigma_t * phi_1, dims) * model_s
-                        - expand_dims(sigma_t * phi_2, dims) * D1
-                        - expand_dims(sigma_t * phi_3, dims) * D2
+                    (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:
@@ -766,28 +790,26 @@ class DPM_Solver:
         else:
             return x_t
 
-    def multistep_dpm_solver_second_update(self, x, model_prev_list, t_prev_list, t, solver_type="dpm_solver"):
+    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 (x.shape[0],)
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
-            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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 ['dpm_solver', 'taylor']:
-            raise ValueError("'solver_type' must be either 'dpm_solver' or 'taylor', got {}".format(solver_type))
+        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
-        dims = x.dim()
-        model_prev_1, model_prev_0 = model_prev_list
-        t_prev_1, t_prev_0 = t_prev_list
-        lambda_prev_1, lambda_prev_0, lambda_t = ns.marginal_lambda(t_prev_1), ns.marginal_lambda(
-            t_prev_0), ns.marginal_lambda(t)
+        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)
@@ -795,55 +817,55 @@ class DPM_Solver:
         h_0 = lambda_prev_0 - lambda_prev_1
         h = lambda_t - lambda_prev_0
         r0 = h_0 / h
-        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
-        if self.predict_x0:
-            if solver_type == 'dpm_solver':
+        D1_0 = (1. / r0) * (model_prev_0 - model_prev_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(sigma_t / sigma_prev_0, dims) * x
-                        - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
-                        - 0.5 * expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * D1_0
+                    (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 = (
-                        expand_dims(sigma_t / sigma_prev_0, dims) * x
-                        - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
-                        + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1_0
+                    (sigma_t / sigma_prev_0) * x
+                    - (alpha_t * phi_1) * model_prev_0
+                    + (alpha_t * (phi_1 / h + 1.)) * D1_0
                 )
         else:
-            if solver_type == 'dpm_solver':
+            phi_1 = torch.expm1(h)
+            if solver_type == 'dpmsolver':
                 x_t = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
-                        - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
-                        - 0.5 * expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * D1_0
+                    (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 = (
-                        expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
-                        - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
-                        - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1_0
+                    (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                    - (sigma_t * phi_1) * model_prev_0
+                    - (sigma_t * (phi_1 / h - 1.)) * D1_0
                 )
         return x_t
 
-    def multistep_dpm_solver_third_update(self, x, model_prev_list, t_prev_list, t, solver_type='dpm_solver'):
+    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 (x.shape[0],)
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
-            solver_type: either 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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
-        dims = x.dim()
         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)
+        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)
@@ -852,39 +874,44 @@ class DPM_Solver:
         h_0 = lambda_prev_0 - lambda_prev_1
         h = lambda_t - lambda_prev_0
         r0, r1 = h_0 / h, h_1 / h
-        D1_0 = expand_dims(1. / r0, dims) * (model_prev_0 - model_prev_1)
-        D1_1 = expand_dims(1. / r1, dims) * (model_prev_1 - model_prev_2)
-        D1 = D1_0 + expand_dims(r0 / (r0 + r1), dims) * (D1_0 - D1_1)
-        D2 = expand_dims(1. / (r0 + r1), dims) * (D1_0 - D1_1)
-        if self.predict_x0:
+        D1_0 = (1. / r0) * (model_prev_0 - model_prev_1)
+        D1_1 = (1. / r1) * (model_prev_1 - model_prev_2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1. / (r0 + r1)) * (D1_0 - D1_1)
+        if self.algorithm_type == "dpmsolver++":
+            phi_1 = torch.expm1(-h)
+            phi_2 = phi_1 / h + 1.
+            phi_3 = phi_2 / h - 0.5
             x_t = (
-                    expand_dims(sigma_t / sigma_prev_0, dims) * x
-                    - expand_dims(alpha_t * (torch.exp(-h) - 1.), dims) * model_prev_0
-                    + expand_dims(alpha_t * ((torch.exp(-h) - 1.) / h + 1.), dims) * D1
-                    - expand_dims(alpha_t * ((torch.exp(-h) - 1. + h) / h ** 2 - 0.5), dims) * D2
+                (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.
+            phi_3 = phi_2 / h - 0.5
             x_t = (
-                    expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
-                    - expand_dims(sigma_t * (torch.exp(h) - 1.), dims) * model_prev_0
-                    - expand_dims(sigma_t * ((torch.exp(h) - 1.) / h - 1.), dims) * D1
-                    - expand_dims(sigma_t * ((torch.exp(h) - 1. - h) / h ** 2 - 0.5), dims) * D2
+                (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='dpm_solver', r1=None,
-                                     r2=None):
+    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 (x.shape[0],).
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            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 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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:
@@ -893,26 +920,24 @@ class DPM_Solver:
         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)
+            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)
+            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='dpm_solver'):
+    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 (x.shape[0],)
-            t: A pytorch tensor. The ending time, with the shape (x.shape[0],).
+            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 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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`.
         """
@@ -925,8 +950,7 @@ class DPM_Solver:
         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='dpm_solver'):
+    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.
 
@@ -941,15 +965,15 @@ class DPM_Solver:
             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 'dpm_solver' or 'taylor'. The type for the high-order solvers.
-                The type slightly impacts the performance. We recommend to use 'dpm_solver' type.
+            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((x.shape[0],)).to(x)
+        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)
@@ -957,18 +981,16 @@ class DPM_Solver:
         nfe = 0
         if order == 2:
             r1 = 0.5
-            lower_update = lambda x, s, t: self.dpm_solver_first_update(x, s, t, return_intermediate=True)
-            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
-                                                                                               solver_type=solver_type,
-                                                                                               **kwargs)
+            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. / 3., 2. / 3.
-            lower_update = lambda x, s, t: self.singlestep_dpm_solver_second_update(x, s, t, r1=r1,
-                                                                                    return_intermediate=True,
-                                                                                    solver_type=solver_type)
-            higher_update = lambda x, s, t, **kwargs: self.singlestep_dpm_solver_third_update(x, s, t, r1=r1, r2=r2,
-                                                                                              solver_type=solver_type,
-                                                                                              **kwargs)
+            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:
@@ -976,7 +998,8 @@ class DPM_Solver:
             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)))
-            norm_fn = lambda v: torch.sqrt(torch.square(v.reshape((v.shape[0], -1))).mean(dim=-1, keepdim=True))
+            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.):
                 x = x_higher
@@ -988,10 +1011,45 @@ class DPM_Solver:
         print('adaptive solver nfe', nfe)
         return x
 
-    def sample(self, x, steps=20, t_start=None, t_end=None, order=3, skip_type='time_uniform',
-               method='singlestep', denoise=False, solver_type='dpm_solver', atol=0.0078,
-               rtol=0.05,
-               ):
+    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. / 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`.
 
@@ -1040,15 +1098,19 @@ class DPM_Solver:
 
         Some advices for choosing the algorithm:
             - For **unconditional sampling** or **guided sampling with small guidance scale** by DPMs:
-                Use singlestep DPM-Solver ("DPM-Solver-fast" in the paper) with `order = 3`.
-                e.g.
-                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=False)
+                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 `predict_x0 = True` and `order = 2`.
+                Use multistep DPM-Solver with `algorithm_type="dpmsolver++"` and `order = 2`.
                 e.g.
-                    >>> dpm_solver = DPM_Solver(model_fn, noise_schedule, predict_x0=True)
+                    >>> 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')
 
@@ -1074,72 +1136,116 @@ class DPM_Solver:
             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: A `bool`. Whether to denoise at the final step. Default is False.
-                If `denoise` is True, the total NFE is (`steps` + 1).
-            solver_type: A `str`. The taylor expansion type for the solver. `dpm_solver` or `taylor`. We recommend `dpm_solver`.
+            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. / 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
-        if method == 'adaptive':
-            with torch.no_grad():
-                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
-            with torch.no_grad():
-                vec_t = timesteps[0].expand((x.shape[0]))
-                model_prev_list = [self.model_fn(x, vec_t)]
-                t_prev_list = [vec_t]
+        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 init_order in range(1, order):
-                    vec_t = timesteps[init_order].expand(x.shape[0])
-                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, init_order,
-                                                         solver_type=solver_type)
-                    model_prev_list.append(self.model_fn(x, vec_t))
-                    t_prev_list.append(vec_t)
+                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):
-                    vec_t = timesteps[step].expand(x.shape[0])
-                    x = self.multistep_dpm_solver_update(x, model_prev_list, t_prev_list, vec_t, order,
-                                                         solver_type=solver_type)
+                    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] = vec_t
+                    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, vec_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 i, order in enumerate(orders):
-                t_T_inner, t_0_inner = timesteps_outer[i], timesteps_outer[i + 1]
-                timesteps_inner = self.get_time_steps(skip_type=skip_type, t_T=t_T_inner.item(), t_0=t_0_inner.item(),
-                                                      N=order, device=device)
-                lambda_inner = self.noise_schedule.marginal_lambda(timesteps_inner)
-                vec_s, vec_t = t_T_inner.repeat(x.shape[0]), t_0_inner.repeat(x.shape[0])
-                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, vec_s, vec_t, order, solver_type=solver_type, r1=r1, r2=r2)
-        if denoise:
-            x = self.denoise_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
-        return x
+                        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
+
 
 
 #############################################################
@@ -1198,4 +1304,4 @@ def expand_dims(v, dims):
     Returns:
         a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
     """
-    return v[(...,) + (None,) * (dims - 1)]
+    return v[(...,) + (None,)*(dims - 1)]

diffusion/how to export onnx.md

@@ -1,4 +1,4 @@
-- Open [onnx_export](onnx_export.py)
-- project_name = "dddsp" change "project_name" to your project name
-- model_path = f'{project_name}/model_500000.pt' change "model_path" to your model path
+- Open [onnx_export](onnx_export.py)
+- project_name = "dddsp" change "project_name" to your project name
+- model_path = f'{project_name}/model_500000.pt' change "model_path" to your model path
 - Run

diffusion/infer_gt_mel.py

@@ -1,6 +1,6 @@
-import numpy as np
 import torch
 import torch.nn.functional as F
+
 from diffusion.unit2mel import load_model_vocoder
 
 

diffusion/logger/saver.py

@@ -2,16 +2,16 @@
 author: wayn391@mastertones
 '''
 
+import datetime
 import os
-import json
 import time
-import yaml
-import datetime
-import torch
+
 import matplotlib.pyplot as plt
-from . import utils
+import torch
+import yaml
 from torch.utils.tensorboard import SummaryWriter
 
+
 class Saver(object):
     def __init__(
             self, 
@@ -125,12 +125,7 @@ class Saver(object):
             torch.save({
                 'global_step': self.global_step,
                 'model': model.state_dict()}, path_pt)
-            
-        # to json
-        if to_json:
-            path_json = os.path.join(
-                self.expdir , name+'.json')
-            utils.to_json(path_params, path_json)
+        
     
     def delete_model(self, name='model', postfix=''):
         # path

diffusion/logger/utils.py

@@ -1,8 +1,9 @@
-import os
-import yaml
 import json
-import pickle
+import os
+
 import torch
+import yaml
+
 
 def traverse_dir(
         root_dir,
@@ -121,6 +122,6 @@ def load_model(
         ckpt = torch.load(path_pt, map_location=torch.device(device))
         global_step = ckpt['global_step']
         model.load_state_dict(ckpt['model'], strict=False)
-        if ckpt.get('optimizer') != None:
+        if ckpt.get("optimizer") is not None:
             optimizer.load_state_dict(ckpt['optimizer'])
     return global_step, model, optimizer

diffusion/onnx_export.py

@@ -1,226 +1,235 @@
-from diffusion_onnx import GaussianDiffusion
-import os
-import yaml
-import torch
-import torch.nn as nn
-import numpy as np
-from wavenet import WaveNet
-import torch.nn.functional as F
-import diffusion
-
-class DotDict(dict):
-    def __getattr__(*args):         
-        val = dict.get(*args)         
-        return DotDict(val) if type(val) is dict else val   
-
-    __setattr__ = dict.__setitem__    
-    __delattr__ = dict.__delitem__
-
-    
-def load_model_vocoder(
-        model_path,
-        device='cpu'):
-    config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
-    with open(config_file, "r") as config:
-        args = yaml.safe_load(config)
-    args = DotDict(args)
-    
-    # load model
-    model = Unit2Mel(
-                args.data.encoder_out_channels, 
-                args.model.n_spk,
-                args.model.use_pitch_aug,
-                128,
-                args.model.n_layers,
-                args.model.n_chans,
-                args.model.n_hidden)
-    
-    print(' [Loading] ' + model_path)
-    ckpt = torch.load(model_path, map_location=torch.device(device))
-    model.to(device)
-    model.load_state_dict(ckpt['model'])
-    model.eval()
-    return model, args
-
-
-class Unit2Mel(nn.Module):
-    def __init__(
-            self,
-            input_channel,
-            n_spk,
-            use_pitch_aug=False,
-            out_dims=128,
-            n_layers=20, 
-            n_chans=384, 
-            n_hidden=256):
-        super().__init__()
-        self.unit_embed = nn.Linear(input_channel, n_hidden)
-        self.f0_embed = nn.Linear(1, n_hidden)
-        self.volume_embed = nn.Linear(1, n_hidden)
-        if use_pitch_aug:
-            self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
-        else:
-            self.aug_shift_embed = None
-        self.n_spk = n_spk
-        if n_spk is not None and n_spk > 1:
-            self.spk_embed = nn.Embedding(n_spk, n_hidden)
-            
-        # diffusion
-        self.decoder = GaussianDiffusion(out_dims, n_layers, n_chans, n_hidden)
-        self.hidden_size = n_hidden
-        self.speaker_map = torch.zeros((self.n_spk,1,1,n_hidden))
-    
-        
-
-    def forward(self, units, mel2ph, f0, volume, g = None):
-        
-        '''
-        input: 
-            B x n_frames x n_unit
-        return: 
-            dict of B x n_frames x feat
-        '''
-
-        decoder_inp = F.pad(units, [0, 0, 1, 0])
-        mel2ph_ = mel2ph.unsqueeze(2).repeat([1, 1, units.shape[-1]])
-        units = torch.gather(decoder_inp, 1, mel2ph_)  # [B, T, H]
-
-        x = self.unit_embed(units) + self.f0_embed((1 + f0.unsqueeze(-1) / 700).log()) + self.volume_embed(volume.unsqueeze(-1))
-
-        if self.n_spk is not None and self.n_spk > 1:   # [N, S]  *  [S, B, 1, H]
-            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
-            g = g * self.speaker_map  # [N, S, B, 1, H]
-            g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
-            g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
-            x = x.transpose(1, 2) + g
-            return x
-        else:
-            return x.transpose(1, 2)
-        
-
-    def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
-                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
-        
-        '''
-        input: 
-            B x n_frames x n_unit
-        return: 
-            dict of B x n_frames x feat
-        '''
-        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
-        if self.n_spk is not None and self.n_spk > 1:
-            if spk_mix_dict is not None:
-                spk_embed_mix = torch.zeros((1,1,self.hidden_size))
-                for k, v in spk_mix_dict.items():
-                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
-                    spk_embeddd = self.spk_embed(spk_id_torch)
-                    self.speaker_map[k] = spk_embeddd
-                    spk_embed_mix = spk_embed_mix + v * spk_embeddd
-                x = x + spk_embed_mix
-            else:
-                x = x + self.spk_embed(spk_id - 1)
-        self.speaker_map = self.speaker_map.unsqueeze(0)
-        self.speaker_map = self.speaker_map.detach()
-        return x.transpose(1, 2)
-
-    def OnnxExport(self, project_name=None, init_noise=None, export_encoder=True, export_denoise=True, export_pred=True, export_after=True):
-        hubert_hidden_size = 768
-        n_frames = 100
-        hubert = torch.randn((1, n_frames, hubert_hidden_size))
-        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
-        f0 = torch.randn((1, n_frames))
-        volume = torch.randn((1, n_frames))
-        spk_mix = []
-        spks = {}
-        if self.n_spk is not None and self.n_spk > 1:
-            for i in range(self.n_spk):
-                spk_mix.append(1.0/float(self.n_spk))
-                spks.update({i:1.0/float(self.n_spk)})
-        spk_mix = torch.tensor(spk_mix)
-        spk_mix = spk_mix.repeat(n_frames, 1)
-        orgouttt = self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
-        outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
-        if export_encoder:
-            torch.onnx.export(
-                self,
-                (hubert, mel2ph, f0, volume, spk_mix),
-                f"{project_name}_encoder.onnx",
-                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
-                output_names=["mel_pred"],
-                dynamic_axes={
-                    "hubert": [1],
-                    "f0": [1],
-                    "volume": [1],
-                    "mel2ph": [1],
-                    "spk_mix": [0],
-                },
-                opset_version=16
-            )
-        
-        self.decoder.OnnxExport(project_name, init_noise=init_noise, export_denoise=export_denoise, export_pred=export_pred, export_after=export_after)
-
-    def ExportOnnx(self, project_name=None):
-        hubert_hidden_size = 768
-        n_frames = 100
-        hubert = torch.randn((1, n_frames, hubert_hidden_size))
-        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
-        f0 = torch.randn((1, n_frames))
-        volume = torch.randn((1, n_frames))
-        spk_mix = []
-        spks = {}
-        if self.n_spk is not None and self.n_spk > 1:
-            for i in range(self.n_spk):
-                spk_mix.append(1.0/float(self.n_spk))
-                spks.update({i:1.0/float(self.n_spk)})
-        spk_mix = torch.tensor(spk_mix)
-        orgouttt = self.orgforward(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
-        outtt = self.forward(hubert, mel2ph, f0, volume, spk_mix)
-
-        torch.onnx.export(
-                self,
-                (hubert, mel2ph, f0, volume, spk_mix),
-                f"{project_name}_encoder.onnx",
-                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
-                output_names=["mel_pred"],
-                dynamic_axes={
-                    "hubert": [1],
-                    "f0": [1],
-                    "volume": [1],
-                    "mel2ph": [1]
-                },
-                opset_version=16
-            )
-
-        condition = torch.randn(1,self.decoder.n_hidden,n_frames)
-        noise = torch.randn((1, 1, self.decoder.mel_bins, condition.shape[2]), dtype=torch.float32)
-        pndm_speedup = torch.LongTensor([100])
-        K_steps = torch.LongTensor([1000])
-        self.decoder = torch.jit.script(self.decoder)
-        self.decoder(condition, noise, pndm_speedup, K_steps)
-
-        torch.onnx.export(
-                self.decoder,
-                (condition, noise, pndm_speedup, K_steps),
-                f"{project_name}_diffusion.onnx",
-                input_names=["condition", "noise", "pndm_speedup", "K_steps"],
-                output_names=["mel"],
-                dynamic_axes={
-                    "condition": [2],
-                    "noise": [3],
-                },
-                opset_version=16
-            )
-
-
-if __name__ == "__main__":
-    project_name = "dddsp"
-    model_path = f'{project_name}/model_500000.pt'
-
-    model, _ = load_model_vocoder(model_path)
-
-    # 分开Diffusion导出（需要使用MoeSS/MoeVoiceStudio或者自己编写Pndm/Dpm采样）
-    model.OnnxExport(project_name, export_encoder=True, export_denoise=True, export_pred=True, export_after=True)
-
-    # 合并Diffusion导出（Encoder和Diffusion分开，直接将Encoder的结果和初始噪声输入Diffusion即可）
-    # model.ExportOnnx(project_name)
-
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import yaml
+from diffusion_onnx import GaussianDiffusion
+
+
+class DotDict(dict):
+    def __getattr__(*args):         
+        val = dict.get(*args)         
+        return DotDict(val) if type(val) is dict else val   
+
+    __setattr__ = dict.__setitem__    
+    __delattr__ = dict.__delitem__
+
+    
+def load_model_vocoder(
+        model_path,
+        device='cpu'):
+    config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
+    with open(config_file, "r") as config:
+        args = yaml.safe_load(config)
+    args = DotDict(args)
+    
+    # load model
+    model = Unit2Mel(
+                args.data.encoder_out_channels, 
+                args.model.n_spk,
+                args.model.use_pitch_aug,
+                128,
+                args.model.n_layers,
+                args.model.n_chans,
+                args.model.n_hidden,
+                args.model.timesteps,
+                args.model.k_step_max)
+    
+    print(' [Loading] ' + model_path)
+    ckpt = torch.load(model_path, map_location=torch.device(device))
+    model.to(device)
+    model.load_state_dict(ckpt['model'])
+    model.eval()
+    return model, args
+
+
+class Unit2Mel(nn.Module):
+    def __init__(
+            self,
+            input_channel,
+            n_spk,
+            use_pitch_aug=False,
+            out_dims=128,
+            n_layers=20, 
+            n_chans=384, 
+            n_hidden=256,
+            timesteps=1000,
+            k_step_max=1000):
+        super().__init__()
+
+        self.unit_embed = nn.Linear(input_channel, n_hidden)
+        self.f0_embed = nn.Linear(1, n_hidden)
+        self.volume_embed = nn.Linear(1, n_hidden)
+        if use_pitch_aug:
+            self.aug_shift_embed = nn.Linear(1, n_hidden, bias=False)
+        else:
+            self.aug_shift_embed = None
+        self.n_spk = n_spk
+        if n_spk is not None and n_spk > 1:
+            self.spk_embed = nn.Embedding(n_spk, n_hidden)
+
+        self.timesteps = timesteps if timesteps is not None else 1000
+        self.k_step_max = k_step_max if k_step_max is not None and k_step_max>0 and k_step_max<self.timesteps else self.timesteps
+
+
+        # diffusion
+        self.decoder = GaussianDiffusion(out_dims, n_layers, n_chans, n_hidden,self.timesteps,self.k_step_max)
+        self.hidden_size = n_hidden
+        self.speaker_map = torch.zeros((self.n_spk,1,1,n_hidden))
+    
+        
+
+    def forward(self, units, mel2ph, f0, volume, g = None):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+
+        decoder_inp = F.pad(units, [0, 0, 1, 0])
+        mel2ph_ = mel2ph.unsqueeze(2).repeat([1, 1, units.shape[-1]])
+        units = torch.gather(decoder_inp, 1, mel2ph_)  # [B, T, H]
+
+        x = self.unit_embed(units) + self.f0_embed((1 + f0.unsqueeze(-1) / 700).log()) + self.volume_embed(volume.unsqueeze(-1))
+
+        if self.n_spk is not None and self.n_spk > 1:   # [N, S]  *  [S, B, 1, H]
+            g = g.reshape((g.shape[0], g.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+            g = g * self.speaker_map  # [N, S, B, 1, H]
+            g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
+            g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
+            x = x.transpose(1, 2) + g
+            return x
+        else:
+            return x.transpose(1, 2)
+        
+
+    def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
+                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
+        if self.n_spk is not None and self.n_spk > 1:
+            if spk_mix_dict is not None:
+                spk_embed_mix = torch.zeros((1,1,self.hidden_size))
+                for k, v in spk_mix_dict.items():
+                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
+                    spk_embeddd = self.spk_embed(spk_id_torch)
+                    self.speaker_map[k] = spk_embeddd
+                    spk_embed_mix = spk_embed_mix + v * spk_embeddd
+                x = x + spk_embed_mix
+            else:
+                x = x + self.spk_embed(spk_id - 1)
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+        self.speaker_map = self.speaker_map.detach()
+        return x.transpose(1, 2)
+
+    def OnnxExport(self, project_name=None, init_noise=None, export_encoder=True, export_denoise=True, export_pred=True, export_after=True):
+        hubert_hidden_size = 768
+        n_frames = 100
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spk_mix = []
+        spks = {}
+        if self.n_spk is not None and self.n_spk > 1:
+            for i in range(self.n_spk):
+                spk_mix.append(1.0/float(self.n_spk))
+                spks.update({i:1.0/float(self.n_spk)})
+        spk_mix = torch.tensor(spk_mix)
+        spk_mix = spk_mix.repeat(n_frames, 1)
+        self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
+        self.forward(hubert, mel2ph, f0, volume, spk_mix)
+        if export_encoder:
+            torch.onnx.export(
+                self,
+                (hubert, mel2ph, f0, volume, spk_mix),
+                f"{project_name}_encoder.onnx",
+                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
+                output_names=["mel_pred"],
+                dynamic_axes={
+                    "hubert": [1],
+                    "f0": [1],
+                    "volume": [1],
+                    "mel2ph": [1],
+                    "spk_mix": [0],
+                },
+                opset_version=16
+            )
+        
+        self.decoder.OnnxExport(project_name, init_noise=init_noise, export_denoise=export_denoise, export_pred=export_pred, export_after=export_after)
+
+    def ExportOnnx(self, project_name=None):
+        hubert_hidden_size = 768
+        n_frames = 100
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        mel2ph = torch.arange(end=n_frames).unsqueeze(0).long()
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spk_mix = []
+        spks = {}
+        if self.n_spk is not None and self.n_spk > 1:
+            for i in range(self.n_spk):
+                spk_mix.append(1.0/float(self.n_spk))
+                spks.update({i:1.0/float(self.n_spk)})
+        spk_mix = torch.tensor(spk_mix)
+        self.orgforward(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
+        self.forward(hubert, mel2ph, f0, volume, spk_mix)
+
+        torch.onnx.export(
+                self,
+                (hubert, mel2ph, f0, volume, spk_mix),
+                f"{project_name}_encoder.onnx",
+                input_names=["hubert", "mel2ph", "f0", "volume", "spk_mix"],
+                output_names=["mel_pred"],
+                dynamic_axes={
+                    "hubert": [1],
+                    "f0": [1],
+                    "volume": [1],
+                    "mel2ph": [1]
+                },
+                opset_version=16
+            )
+
+        condition = torch.randn(1,self.decoder.n_hidden,n_frames)
+        noise = torch.randn((1, 1, self.decoder.mel_bins, condition.shape[2]), dtype=torch.float32)
+        pndm_speedup = torch.LongTensor([100])
+        K_steps = torch.LongTensor([1000])
+        self.decoder = torch.jit.script(self.decoder)
+        self.decoder(condition, noise, pndm_speedup, K_steps)
+
+        torch.onnx.export(
+                self.decoder,
+                (condition, noise, pndm_speedup, K_steps),
+                f"{project_name}_diffusion.onnx",
+                input_names=["condition", "noise", "pndm_speedup", "K_steps"],
+                output_names=["mel"],
+                dynamic_axes={
+                    "condition": [2],
+                    "noise": [3],
+                },
+                opset_version=16
+            )
+
+
+if __name__ == "__main__":
+    project_name = "dddsp"
+    model_path = f'{project_name}/model_500000.pt'
+
+    model, _ = load_model_vocoder(model_path)
+
+    # 分开Diffusion导出（需要使用MoeSS/MoeVoiceStudio或者自己编写Pndm/Dpm采样）
+    model.OnnxExport(project_name, export_encoder=True, export_denoise=True, export_pred=True, export_after=True)
+
+    # 合并Diffusion导出（Encoder和Diffusion分开，直接将Encoder的结果和初始噪声输入Diffusion即可）
+    # model.ExportOnnx(project_name)
+

diffusion/solver.py

@@ -1,13 +1,15 @@
-import os
 import time
+
+import librosa
 import numpy as np
 import torch
-import librosa
-from diffusion.logger.saver import Saver
-from diffusion.logger import utils
 from torch import autocast
 from torch.cuda.amp import GradScaler
 
+from diffusion.logger import utils
+from diffusion.logger.saver import Saver
+
+
 def test(args, model, vocoder, loader_test, saver):
     print(' [*] testing...')
     model.eval()
@@ -40,10 +42,12 @@ def test(args, model, vocoder, loader_test, saver):
                     data['f0'], 
                     data['volume'], 
                     data['spk_id'],
-                    gt_spec=None,
+                    gt_spec=None if model.k_step_max == model.timesteps else data['mel'],
                     infer=True, 
                     infer_speedup=args.infer.speedup, 
-                    method=args.infer.method)
+                    method=args.infer.method,
+                    k_step=model.k_step_max
+                    )
             signal = vocoder.infer(mel, data['f0'])
             ed_time = time.time()
                         
@@ -62,7 +66,8 @@ def test(args, model, vocoder, loader_test, saver):
                     data['volume'], 
                     data['spk_id'], 
                     gt_spec=data['mel'],
-                    infer=False)
+                    infer=False,
+                    k_step=model.k_step_max)
                 test_loss += loss.item()
             
             # log mel
@@ -121,11 +126,11 @@ def train(args, initial_global_step, model, optimizer, scheduler, vocoder, loade
             # forward
             if dtype == torch.float32:
                 loss = model(data['units'].float(), data['f0'], data['volume'], data['spk_id'], 
-                                aug_shift = data['aug_shift'], gt_spec=data['mel'].float(), infer=False)
+                                aug_shift = data['aug_shift'], gt_spec=data['mel'].float(), infer=False, k_step=model.k_step_max)
             else:
                 with autocast(device_type=args.device, dtype=dtype):
                     loss = model(data['units'], data['f0'], data['volume'], data['spk_id'], 
-                                    aug_shift = data['aug_shift'], gt_spec=data['mel'], infer=False)
+                                    aug_shift = data['aug_shift'], gt_spec=data['mel'], infer=False, k_step=model.k_step_max)
             
             # handle nan loss
             if torch.isnan(loss):
@@ -171,25 +176,41 @@ def train(args, initial_global_step, model, optimizer, scheduler, vocoder, loade
                 optimizer_save = optimizer if args.train.save_opt else None
                 
                 # save latest
+                import os
+                import re
+                file_name_list = os.listdir("./logs/44k/diffusion/")
+                for i in range(len(file_name_list)):
+                    if re.search(".pt", file_name_list[i]):
+                        os.remove("./logs/44k/diffusion/" + file_name_list[i])
+                pwd = os.getcwd()
+
                 saver.save_model(model, optimizer_save, postfix=f'{saver.global_step}')
                 last_val_step = saver.global_step - args.train.interval_val
                 if last_val_step % args.train.interval_force_save != 0:
                     saver.delete_model(postfix=f'{last_val_step}')
                 
-                # run testing set
-                test_loss = test(args, model, vocoder, loader_test, saver)
-                
-                # log loss
-                saver.log_info(
-                    ' --- <validation> --- \nloss: {:.3f}. '.format(
-                        test_loss,
-                    )
-                )
-                
-                saver.log_value({
-                    'validation/loss': test_loss
-                })
-                
+                # # run testing set
+                # test_loss = test(args, model, vocoder, loader_test, saver)
+                #
+                # # log loss
+                # saver.log_info(
+                #     ' --- <validation> --- \nloss: {:.3f}. '.format(
+                #         test_loss,
+                #     )
+                # )
+                #
+                # saver.log_value({
+                #     'validation/loss': test_loss
+                # })
+
+                if os.path.exists("/content/so-vits-svc/TMP"):
+                    os.chdir("/content/so-vits-svc/TMP")
+                    os.system("git pull")
+                    os.system("git add .")
+                    os.system('git commit -m "update"')
+                    os.system("git push")
+                os.chdir(pwd)
+
                 model.train()
 
                           

diffusion/uni_pc.py

@@ -0,0 +1,733 @@
+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.,
+            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.
+            self.t_array = torch.linspace(0., 1., 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.
+            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
+            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
+            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.
+
+    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. - torch.exp(2. * 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. - torch.exp(2. * 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. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * 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. * 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. * 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.,
+    classifier_fn=None,
+    classifier_kwargs={},
+):
+    """Create a wrapper function for the noise prediction model.
+    """
+
+    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. / noise_schedule.total_N) * noise_schedule.total_N
+        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. 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"]
+    assert guidance_type in ["uncond", "classifier", "classifier-free"]
+    return model_fn
+
+
+class UniPC:
+    def __init__(
+        self,
+        model_fn,
+        noise_schedule,
+        algorithm_type="data_prediction",
+        correcting_x0_fn=None,
+        correcting_xt_fn=None,
+        thresholding_max_val=1.,
+        dynamic_thresholding_ratio=0.995,
+        variant='bh1'
+    ):
+        """Construct a UniPC. 
+
+        We support both data_prediction and noise_prediction.
+        """
+        self.model = lambda x, t: model_fn(x, t.expand((x.shape[0])))
+        self.noise_schedule = noise_schedule
+        assert algorithm_type in ["data_prediction", "noise_prediction"]
+        
+        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
+        
+        self.variant = variant
+        self.predict_x0 = algorithm_type == "data_prediction"
+
+    def dynamic_thresholding_fn(self, x0, t=None):
+        """
+        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)
+        return x0
+
+    def model_fn(self, x, t):
+        """
+        Convert the model to the noise prediction model or the data prediction model. 
+        """
+        if self.predict_x0:
+            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.
+        """
+        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. / t_order), t_0**(1. / 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.
+        """
+        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 = steps
+            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 multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
+        if len(t.shape) == 0:
+            t = t.view(-1)
+        if 'bh' in self.variant:
+            return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
+        else:
+            assert self.variant == 'vary_coeff'
+            return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
+
+    def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
+        #print(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
+        ns = self.noise_schedule
+        assert order <= len(model_prev_list)
+
+        # first compute rks
+        t_prev_0 = t_prev_list[-1]
+        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
+        lambda_t = ns.marginal_lambda(t)
+        model_prev_0 = model_prev_list[-1]
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        log_alpha_t = ns.marginal_log_mean_coeff(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h = lambda_t - lambda_prev_0
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            t_prev_i = t_prev_list[-(i + 1)]
+            model_prev_i = model_prev_list[-(i + 1)]
+            lambda_prev_i = ns.marginal_lambda(t_prev_i)
+            rk = (lambda_prev_i - lambda_prev_0) / h
+            rks.append(rk)
+            D1s.append((model_prev_i - model_prev_0) / rk)
+
+        rks.append(1.)
+        rks = torch.tensor(rks, device=x.device)
+
+        K = len(rks)
+        # build C matrix
+        C = []
+
+        col = torch.ones_like(rks)
+        for k in range(1, K + 1):
+            C.append(col)
+            col = col * rks / (k + 1) 
+        C = torch.stack(C, dim=1)
+
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1) # (B, K)
+            C_inv_p = torch.linalg.inv(C[:-1, :-1])
+            A_p = C_inv_p
+
+        if use_corrector:
+            #print('using corrector')
+            C_inv = torch.linalg.inv(C)
+            A_c = C_inv
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh)
+        h_phi_ks = []
+        factorial_k = 1
+        h_phi_k = h_phi_1
+        for k in range(1, K + 2):
+            h_phi_ks.append(h_phi_k)
+            h_phi_k = h_phi_k / hh - 1 / factorial_k
+            factorial_k *= (k + 1)
+
+        model_t = None
+        if self.predict_x0:
+            x_t_ = (
+                sigma_t / sigma_prev_0 * x
+                - alpha_t * h_phi_1 * model_prev_0
+            )
+            # now predictor
+            x_t = x_t_
+            if len(D1s) > 0:
+                # compute the residuals for predictor
+                for k in range(K - 1):
+                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
+            # now corrector
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_
+                k = 0
+                for k in range(K - 1):
+                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
+                x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
+        else:
+            log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+            x_t_ = (
+                (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
+                - (sigma_t * h_phi_1) * model_prev_0
+            )
+            # now predictor
+            x_t = x_t_
+            if len(D1s) > 0:
+                # compute the residuals for predictor
+                for k in range(K - 1):
+                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
+            # now corrector
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_
+                k = 0
+                for k in range(K - 1):
+                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
+                x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
+        return x_t, model_t
+
+    def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
+        #print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
+        ns = self.noise_schedule
+        assert order <= len(model_prev_list)
+
+        # first compute rks
+        t_prev_0 = t_prev_list[-1]
+        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
+        lambda_t = ns.marginal_lambda(t)
+        model_prev_0 = model_prev_list[-1]
+        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
+        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
+        alpha_t = torch.exp(log_alpha_t)
+
+        h = lambda_t - lambda_prev_0
+
+        rks = []
+        D1s = []
+        for i in range(1, order):
+            t_prev_i = t_prev_list[-(i + 1)]
+            model_prev_i = model_prev_list[-(i + 1)]
+            lambda_prev_i = ns.marginal_lambda(t_prev_i)
+            rk = (lambda_prev_i - lambda_prev_0) / h
+            rks.append(rk)
+            D1s.append((model_prev_i - model_prev_0) / rk)
+
+        rks.append(1.)
+        rks = torch.tensor(rks, device=x.device)
+
+        R = []
+        b = []
+
+        hh = -h if self.predict_x0 else h
+        h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+        h_phi_k = h_phi_1 / hh - 1
+
+        factorial_i = 1
+
+        if self.variant == 'bh1':
+            B_h = hh
+        elif self.variant == 'bh2':
+            B_h = torch.expm1(hh)
+        else:
+            raise NotImplementedError()
+            
+        for i in range(1, order + 1):
+            R.append(torch.pow(rks, i - 1))
+            b.append(h_phi_k * factorial_i / B_h)
+            factorial_i *= (i + 1)
+            h_phi_k = h_phi_k / hh - 1 / factorial_i 
+
+        R = torch.stack(R)
+        b = torch.cat(b)
+
+        # now predictor
+        use_predictor = len(D1s) > 0 and x_t is None
+        if len(D1s) > 0:
+            D1s = torch.stack(D1s, dim=1) # (B, K)
+            if x_t is None:
+                # for order 2, we use a simplified version
+                if order == 2:
+                    rhos_p = torch.tensor([0.5], device=b.device)
+                else:
+                    rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
+        else:
+            D1s = None
+
+        if use_corrector:
+            #print('using corrector')
+            # for order 1, we use a simplified version
+            if order == 1:
+                rhos_c = torch.tensor([0.5], device=b.device)
+            else:
+                rhos_c = torch.linalg.solve(R, b)
+
+        model_t = None
+        if self.predict_x0:
+            x_t_ = (
+                sigma_t / sigma_prev_0 * x
+                - alpha_t * h_phi_1 * model_prev_0
+            )
+
+            if x_t is None:
+                if use_predictor:
+                    pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
+                else:
+                    pred_res = 0
+                x_t = x_t_ - alpha_t * B_h * pred_res
+
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                if D1s is not None:
+                    corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
+                else:
+                    corr_res = 0
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        else:
+            x_t_ = (
+                torch.exp(log_alpha_t - log_alpha_prev_0) * x
+                - sigma_t * h_phi_1 * model_prev_0
+            )
+            if x_t is None:
+                if use_predictor:
+                    pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
+                else:
+                    pred_res = 0
+                x_t = x_t_ - sigma_t * B_h * pred_res
+
+            if use_corrector:
+                model_t = self.model_fn(x_t, t)
+                if D1s is not None:
+                    corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
+                else:
+                    corr_res = 0
+                D1_t = (model_t - model_prev_0)
+                x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+        return x_t, model_t
+
+    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, atol=0.0078, rtol=0.05, return_intermediate=False,
+    ):
+        """
+        Compute the sample at time `t_end` by UniPC, given the initial `x` at time `t_start`.
+        """
+        t_0 = 1. / 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 == '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 UniPC.
+                for step in range(1, order):
+                    t = timesteps[step]
+                    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, t, step, use_corrector=True)
+                    if model_x is None:
+                        model_x = 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)
+                    t_prev_list.append(t)
+                    model_prev_list.append(model_x)
+                    
+                # Compute the remaining values by `order`-th order multistep DPM-Solver.
+                for step in range(order, steps + 1):
+                    t = timesteps[step]
+                    if lower_order_final:
+                        step_order = min(order, steps + 1 - step)
+                    else:
+                        step_order = order
+                    if step == steps:
+                        #print('do not run corrector at the last step')
+                        use_corrector = False
+                    else:
+                        use_corrector = True
+                    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, t, step_order, use_corrector=use_corrector)
+                    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:
+                        if model_x is None:
+                            model_x = self.model_fn(x, t)
+                        model_prev_list[-1] = model_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)]

diffusion/unit2mel.py

@@ -1,11 +1,14 @@
 import os
-import yaml
+
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
+import yaml
+
 from .diffusion import GaussianDiffusion
-from .wavenet import WaveNet
 from .vocoder import Vocoder
+from .wavenet import WaveNet
+
 
 class DotDict(dict):
     def __getattr__(*args):         
@@ -21,9 +24,11 @@ def load_model_vocoder(
         device='cpu',
         config_path = None
         ):
-    if config_path is None: config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
-    else: config_file = config_path
-    
+    if config_path is None:
+        config_file = os.path.join(os.path.split(model_path)[0], 'config.yaml')
+    else:
+        config_file = config_path
+
     with open(config_file, "r") as config:
         args = yaml.safe_load(config)
     args = DotDict(args)
@@ -39,13 +44,17 @@ def load_model_vocoder(
                 vocoder.dimension,
                 args.model.n_layers,
                 args.model.n_chans,
-                args.model.n_hidden)
+                args.model.n_hidden,
+                args.model.timesteps,
+                args.model.k_step_max
+                )
     
     print(' [Loading] ' + model_path)
     ckpt = torch.load(model_path, map_location=torch.device(device))
     model.to(device)
     model.load_state_dict(ckpt['model'])
     model.eval()
+    print(f'Loaded diffusion model, sampler is {args.infer.method}, speedup: {args.infer.speedup} ')
     return model, vocoder, args
 
 
@@ -58,7 +67,10 @@ class Unit2Mel(nn.Module):
             out_dims=128,
             n_layers=20, 
             n_chans=384, 
-            n_hidden=256):
+            n_hidden=256,
+            timesteps=1000,
+            k_step_max=1000
+            ):
         super().__init__()
         self.unit_embed = nn.Linear(input_channel, n_hidden)
         self.f0_embed = nn.Linear(1, n_hidden)
@@ -70,9 +82,51 @@ class Unit2Mel(nn.Module):
         self.n_spk = n_spk
         if n_spk is not None and n_spk > 1:
             self.spk_embed = nn.Embedding(n_spk, n_hidden)
-            
+        
+        self.timesteps = timesteps if timesteps is not None else 1000
+        self.k_step_max = k_step_max if k_step_max is not None and k_step_max>0 and k_step_max<self.timesteps else self.timesteps
+
+        self.n_hidden = n_hidden
         # diffusion
-        self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden), out_dims=out_dims)
+        self.decoder = GaussianDiffusion(WaveNet(out_dims, n_layers, n_chans, n_hidden),timesteps=self.timesteps,k_step=self.k_step_max, out_dims=out_dims)
+        self.input_channel = input_channel
+    
+    def init_spkembed(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
+                gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
+        
+        '''
+        input: 
+            B x n_frames x n_unit
+        return: 
+            dict of B x n_frames x feat
+        '''
+        x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
+        if self.n_spk is not None and self.n_spk > 1:
+            if spk_mix_dict is not None:
+                spk_embed_mix = torch.zeros((1,1,self.hidden_size))
+                for k, v in spk_mix_dict.items():
+                    spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
+                    spk_embeddd = self.spk_embed(spk_id_torch)
+                    self.speaker_map[k] = spk_embeddd
+                    spk_embed_mix = spk_embed_mix + v * spk_embeddd
+                x = x + spk_embed_mix
+            else:
+                x = x + self.spk_embed(spk_id - 1)
+        self.speaker_map = self.speaker_map.unsqueeze(0)
+        self.speaker_map = self.speaker_map.detach()
+        return x.transpose(1, 2)
+
+    def init_spkmix(self, n_spk):
+        self.speaker_map = torch.zeros((n_spk,1,1,self.n_hidden))
+        hubert_hidden_size = self.input_channel
+        n_frames = 10
+        hubert = torch.randn((1, n_frames, hubert_hidden_size))
+        f0 = torch.randn((1, n_frames))
+        volume = torch.randn((1, n_frames))
+        spks = {}
+        for i in range(n_spk):
+            spks.update({i:1.0/float(self.n_spk)})
+        self.init_spkembed(hubert, f0.unsqueeze(-1), volume.unsqueeze(-1), spk_mix_dict=spks)
 
     def forward(self, units, f0, volume, spk_id = None, spk_mix_dict = None, aug_shift = None,
                 gt_spec=None, infer=True, infer_speedup=10, method='dpm-solver', k_step=300, use_tqdm=True):
@@ -84,6 +138,12 @@ class Unit2Mel(nn.Module):
             dict of B x n_frames x feat
         '''
 
+        if not self.training and gt_spec is not None and k_step>self.k_step_max:
+            raise Exception("The shallow diffusion k_step is greater than the maximum diffusion k_step(k_step_max)!")
+
+        if not self.training and gt_spec is None and self.k_step_max!=self.timesteps:
+            raise Exception("This model can only be used for shallow diffusion and can not infer alone!")
+
         x = self.unit_embed(units) + self.f0_embed((1+ f0 / 700).log()) + self.volume_embed(volume)
         if self.n_spk is not None and self.n_spk > 1:
             if spk_mix_dict is not None:
@@ -91,7 +151,14 @@ class Unit2Mel(nn.Module):
                     spk_id_torch = torch.LongTensor(np.array([[k]])).to(units.device)
                     x = x + v * self.spk_embed(spk_id_torch)
             else:
-                x = x + self.spk_embed(spk_id)
+                if spk_id.shape[1] > 1:
+                    g = spk_id.reshape((spk_id.shape[0], spk_id.shape[1], 1, 1, 1))  # [N, S, B, 1, 1]
+                    g = g * self.speaker_map  # [N, S, B, 1, H]
+                    g = torch.sum(g, dim=1) # [N, 1, B, 1, H]
+                    g = g.transpose(0, -1).transpose(0, -2).squeeze(0) # [B, H, N]
+                    x = x + g
+                else:
+                    x = x + self.spk_embed(spk_id)
         if self.aug_shift_embed is not None and aug_shift is not None:
             x = x + self.aug_shift_embed(aug_shift / 5) 
         x = self.decoder(x, gt_spec=gt_spec, infer=infer, infer_speedup=infer_speedup, method=method, k_step=k_step, use_tqdm=use_tqdm)

diffusion/vocoder.py

@@ -1,9 +1,10 @@
 import torch
-from vdecoder.nsf_hifigan.nvSTFT import STFT
-from vdecoder.nsf_hifigan.models import load_model,load_config
 from torchaudio.transforms import Resample
 
-    
+from vdecoder.nsf_hifigan.models import load_config, load_model
+from vdecoder.nsf_hifigan.nvSTFT import STFT
+
+
 class Vocoder:
     def __init__(self, vocoder_type, vocoder_ckpt, device = None):
         if device is None:

edgetts/tts.py

@@ -0,0 +1,47 @@
+import asyncio
+import random
+import sys
+
+import edge_tts
+from edge_tts import VoicesManager
+from langdetect import DetectorFactory, detect
+
+DetectorFactory.seed = 0
+
+TEXT = sys.argv[1]
+LANG = detect(TEXT) if sys.argv[2] == "Auto" else sys.argv[2]
+RATE = sys.argv[3]
+VOLUME = sys.argv[4]
+GENDER = sys.argv[5] if len(sys.argv) == 6 else None
+OUTPUT_FILE = "tts.wav"
+
+print("Running TTS...")
+print(f"Text: {TEXT}, Language: {LANG}, Gender: {GENDER}, Rate: {RATE}, Volume: {VOLUME}")
+
+async def _main() -> None:
+    voices = await VoicesManager.create()
+    if GENDER is not None:
+        # From "zh-cn" to "zh-CN" etc.
+        if LANG == "zh-cn" or LANG == "zh-tw":
+            LOCALE = LANG[:-2] + LANG[-2:].upper()
+            voice = voices.find(Gender=GENDER, Locale=LOCALE)
+        else:
+            voice = voices.find(Gender=GENDER, Language=LANG)
+        VOICE = random.choice(voice)["Name"]
+        print(f"Using random {LANG} voice: {VOICE}")
+    else:
+        VOICE = LANG
+        
+    communicate = edge_tts.Communicate(text = TEXT, voice = VOICE, rate = RATE, volume = VOLUME)
+    await communicate.save(OUTPUT_FILE)
+
+if __name__ == "__main__":
+    if sys.platform.startswith("win"):
+        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
+        asyncio.run(_main())
+    else:
+        loop = asyncio.get_event_loop_policy().get_event_loop()
+        try:
+            loop.run_until_complete(_main())
+        finally:
+            loop.close()

edgetts/tts_voices.py

@@ -0,0 +1,306 @@
+#List of Supported Voices for edge_TTS
+SUPPORTED_VOICES = {
+    'zh-CN-XiaoxiaoNeural': 'zh-CN',
+    'zh-CN-XiaoyiNeural': 'zh-CN',
+    'zh-CN-YunjianNeural': 'zh-CN',
+    'zh-CN-YunxiNeural': 'zh-CN',
+    'zh-CN-YunxiaNeural': 'zh-CN',
+    'zh-CN-YunyangNeural': 'zh-CN',
+    'zh-HK-HiuGaaiNeural': 'zh-HK',
+    'zh-HK-HiuMaanNeural': 'zh-HK',
+    'zh-HK-WanLungNeural': 'zh-HK',
+    'zh-TW-HsiaoChenNeural': 'zh-TW',
+    'zh-TW-YunJheNeural': 'zh-TW',
+    'zh-TW-HsiaoYuNeural': 'zh-TW',
+    'af-ZA-AdriNeural': 'af-ZA',
+    'af-ZA-WillemNeural': 'af-ZA',
+    'am-ET-AmehaNeural': 'am-ET',
+    'am-ET-MekdesNeural': 'am-ET',
+    'ar-AE-FatimaNeural': 'ar-AE',
+    'ar-AE-HamdanNeural': 'ar-AE',
+    'ar-BH-AliNeural': 'ar-BH',
+    'ar-BH-LailaNeural': 'ar-BH',
+    'ar-DZ-AminaNeural': 'ar-DZ',
+    'ar-DZ-IsmaelNeural': 'ar-DZ',
+    'ar-EG-SalmaNeural': 'ar-EG',
+    'ar-EG-ShakirNeural': 'ar-EG',
+    'ar-IQ-BasselNeural': 'ar-IQ',
+    'ar-IQ-RanaNeural': 'ar-IQ',
+    'ar-JO-SanaNeural': 'ar-JO',
+    'ar-JO-TaimNeural': 'ar-JO',
+    'ar-KW-FahedNeural': 'ar-KW',
+    'ar-KW-NouraNeural': 'ar-KW',
+    'ar-LB-LaylaNeural': 'ar-LB',
+    'ar-LB-RamiNeural': 'ar-LB',
+    'ar-LY-ImanNeural': 'ar-LY',
+    'ar-LY-OmarNeural': 'ar-LY',
+    'ar-MA-JamalNeural': 'ar-MA',
+    'ar-MA-MounaNeural': 'ar-MA',
+    'ar-OM-AbdullahNeural': 'ar-OM',
+    'ar-OM-AyshaNeural': 'ar-OM',
+    'ar-QA-AmalNeural': 'ar-QA',
+    'ar-QA-MoazNeural': 'ar-QA',
+    'ar-SA-HamedNeural': 'ar-SA',
+    'ar-SA-ZariyahNeural': 'ar-SA',
+    'ar-SY-AmanyNeural': 'ar-SY',
+    'ar-SY-LaithNeural': 'ar-SY',
+    'ar-TN-HediNeural': 'ar-TN',
+    'ar-TN-ReemNeural': 'ar-TN',
+    'ar-YE-MaryamNeural': 'ar-YE',
+    'ar-YE-SalehNeural': 'ar-YE',
+    'az-AZ-BabekNeural': 'az-AZ',
+    'az-AZ-BanuNeural': 'az-AZ',
+    'bg-BG-BorislavNeural': 'bg-BG',
+    'bg-BG-KalinaNeural': 'bg-BG',
+    'bn-BD-NabanitaNeural': 'bn-BD',
+    'bn-BD-PradeepNeural': 'bn-BD',
+    'bn-IN-BashkarNeural': 'bn-IN',
+    'bn-IN-TanishaaNeural': 'bn-IN',
+    'bs-BA-GoranNeural': 'bs-BA',
+    'bs-BA-VesnaNeural': 'bs-BA',
+    'ca-ES-EnricNeural': 'ca-ES',
+    'ca-ES-JoanaNeural': 'ca-ES',
+    'cs-CZ-AntoninNeural': 'cs-CZ',
+    'cs-CZ-VlastaNeural': 'cs-CZ',
+    'cy-GB-AledNeural': 'cy-GB',
+    'cy-GB-NiaNeural': 'cy-GB',
+    'da-DK-ChristelNeural': 'da-DK',
+    'da-DK-JeppeNeural': 'da-DK',
+    'de-AT-IngridNeural': 'de-AT',
+    'de-AT-JonasNeural': 'de-AT',
+    'de-CH-JanNeural': 'de-CH',
+    'de-CH-LeniNeural': 'de-CH',
+    'de-DE-AmalaNeural': 'de-DE',
+    'de-DE-ConradNeural': 'de-DE',
+    'de-DE-KatjaNeural': 'de-DE',
+    'de-DE-KillianNeural': 'de-DE',
+    'el-GR-AthinaNeural': 'el-GR',
+    'el-GR-NestorasNeural': 'el-GR',
+    'en-AU-NatashaNeural': 'en-AU',
+    'en-AU-WilliamNeural': 'en-AU',
+    'en-CA-ClaraNeural': 'en-CA',
+    'en-CA-LiamNeural': 'en-CA',
+    'en-GB-LibbyNeural': 'en-GB',
+    'en-GB-MaisieNeural': 'en-GB',
+    'en-GB-RyanNeural': 'en-GB',
+    'en-GB-SoniaNeural': 'en-GB',
+    'en-GB-ThomasNeural': 'en-GB',
+    'en-HK-SamNeural': 'en-HK',
+    'en-HK-YanNeural': 'en-HK',
+    'en-IE-ConnorNeural': 'en-IE',
+    'en-IE-EmilyNeural': 'en-IE',
+    'en-IN-NeerjaNeural': 'en-IN',
+    'en-IN-PrabhatNeural': 'en-IN',
+    'en-KE-AsiliaNeural': 'en-KE',
+    'en-KE-ChilembaNeural': 'en-KE',
+    'en-NG-AbeoNeural': 'en-NG',
+    'en-NG-EzinneNeural': 'en-NG',
+    'en-NZ-MitchellNeural': 'en-NZ',
+    'en-NZ-MollyNeural': 'en-NZ',
+    'en-PH-JamesNeural': 'en-PH',
+    'en-PH-RosaNeural': 'en-PH',
+    'en-SG-LunaNeural': 'en-SG',
+    'en-SG-WayneNeural': 'en-SG',
+    'en-TZ-ElimuNeural': 'en-TZ',
+    'en-TZ-ImaniNeural': 'en-TZ',
+    'en-US-AnaNeural': 'en-US',
+    'en-US-AriaNeural': 'en-US',
+    'en-US-ChristopherNeural': 'en-US',
+    'en-US-EricNeural': 'en-US',
+    'en-US-GuyNeural': 'en-US',
+    'en-US-JennyNeural': 'en-US',
+    'en-US-MichelleNeural': 'en-US',
+    'en-ZA-LeahNeural': 'en-ZA',
+    'en-ZA-LukeNeural': 'en-ZA',
+    'es-AR-ElenaNeural': 'es-AR',
+    'es-AR-TomasNeural': 'es-AR',
+    'es-BO-MarceloNeural': 'es-BO',
+    'es-BO-SofiaNeural': 'es-BO',
+    'es-CL-CatalinaNeural': 'es-CL',
+    'es-CL-LorenzoNeural': 'es-CL',
+    'es-CO-GonzaloNeural': 'es-CO',
+    'es-CO-SalomeNeural': 'es-CO',
+    'es-CR-JuanNeural': 'es-CR',
+    'es-CR-MariaNeural': 'es-CR',
+    'es-CU-BelkysNeural': 'es-CU',
+    'es-CU-ManuelNeural': 'es-CU',
+    'es-DO-EmilioNeural': 'es-DO',
+    'es-DO-RamonaNeural': 'es-DO',
+    'es-EC-AndreaNeural': 'es-EC',
+    'es-EC-LuisNeural': 'es-EC',
+    'es-ES-AlvaroNeural': 'es-ES',
+    'es-ES-ElviraNeural': 'es-ES',
+    'es-ES-ManuelEsCUNeural': 'es-ES',
+    'es-GQ-JavierNeural': 'es-GQ',
+    'es-GQ-TeresaNeural': 'es-GQ',
+    'es-GT-AndresNeural': 'es-GT',
+    'es-GT-MartaNeural': 'es-GT',
+    'es-HN-CarlosNeural': 'es-HN',
+    'es-HN-KarlaNeural': 'es-HN',
+    'es-MX-DaliaNeural': 'es-MX',
+    'es-MX-JorgeNeural': 'es-MX',
+    'es-MX-LorenzoEsCLNeural': 'es-MX',
+    'es-NI-FedericoNeural': 'es-NI',
+    'es-NI-YolandaNeural': 'es-NI',
+    'es-PA-MargaritaNeural': 'es-PA',
+    'es-PA-RobertoNeural': 'es-PA',
+    'es-PE-AlexNeural': 'es-PE',
+    'es-PE-CamilaNeural': 'es-PE',
+    'es-PR-KarinaNeural': 'es-PR',
+    'es-PR-VictorNeural': 'es-PR',
+    'es-PY-MarioNeural': 'es-PY',
+    'es-PY-TaniaNeural': 'es-PY',
+    'es-SV-LorenaNeural': 'es-SV',
+    'es-SV-RodrigoNeural': 'es-SV',
+    'es-US-AlonsoNeural': 'es-US',
+    'es-US-PalomaNeural': 'es-US',
+    'es-UY-MateoNeural': 'es-UY',
+    'es-UY-ValentinaNeural': 'es-UY',
+    'es-VE-PaolaNeural': 'es-VE',
+    'es-VE-SebastianNeural': 'es-VE',
+    'et-EE-AnuNeural': 'et-EE',
+    'et-EE-KertNeural': 'et-EE',
+    'fa-IR-DilaraNeural': 'fa-IR',
+    'fa-IR-FaridNeural': 'fa-IR',
+    'fi-FI-HarriNeural': 'fi-FI',
+    'fi-FI-NooraNeural': 'fi-FI',
+    'fil-PH-AngeloNeural': 'fil-PH',
+    'fil-PH-BlessicaNeural': 'fil-PH',
+    'fr-BE-CharlineNeural': 'fr-BE',
+    'fr-BE-GerardNeural': 'fr-BE',
+    'fr-CA-AntoineNeural': 'fr-CA',
+    'fr-CA-JeanNeural': 'fr-CA',
+    'fr-CA-SylvieNeural': 'fr-CA',
+    'fr-CH-ArianeNeural': 'fr-CH',
+    'fr-CH-FabriceNeural': 'fr-CH',
+    'fr-FR-DeniseNeural': 'fr-FR',
+    'fr-FR-EloiseNeural': 'fr-FR',
+    'fr-FR-HenriNeural': 'fr-FR',
+    'ga-IE-ColmNeural': 'ga-IE',
+    'ga-IE-OrlaNeural': 'ga-IE',
+    'gl-ES-RoiNeural': 'gl-ES',
+    'gl-ES-SabelaNeural': 'gl-ES',
+    'gu-IN-DhwaniNeural': 'gu-IN',
+    'gu-IN-NiranjanNeural': 'gu-IN',
+    'he-IL-AvriNeural': 'he-IL',
+    'he-IL-HilaNeural': 'he-IL',
+    'hi-IN-MadhurNeural': 'hi-IN',
+    'hi-IN-SwaraNeural': 'hi-IN',
+    'hr-HR-GabrijelaNeural': 'hr-HR',
+    'hr-HR-SreckoNeural': 'hr-HR',
+    'hu-HU-NoemiNeural': 'hu-HU',
+    'hu-HU-TamasNeural': 'hu-HU',
+    'id-ID-ArdiNeural': 'id-ID',
+    'id-ID-GadisNeural': 'id-ID',
+    'is-IS-GudrunNeural': 'is-IS',
+    'is-IS-GunnarNeural': 'is-IS',
+    'it-IT-DiegoNeural': 'it-IT',
+    'it-IT-ElsaNeural': 'it-IT',
+    'it-IT-IsabellaNeural': 'it-IT',
+    'ja-JP-KeitaNeural': 'ja-JP',
+    'ja-JP-NanamiNeural': 'ja-JP',
+    'jv-ID-DimasNeural': 'jv-ID',
+    'jv-ID-SitiNeural': 'jv-ID',
+    'ka-GE-EkaNeural': 'ka-GE',
+    'ka-GE-GiorgiNeural': 'ka-GE',
+    'kk-KZ-AigulNeural': 'kk-KZ',
+    'kk-KZ-DauletNeural': 'kk-KZ',
+    'km-KH-PisethNeural': 'km-KH',
+    'km-KH-SreymomNeural': 'km-KH',
+    'kn-IN-GaganNeural': 'kn-IN',
+    'kn-IN-SapnaNeural': 'kn-IN',
+    'ko-KR-InJoonNeural': 'ko-KR',
+    'ko-KR-SunHiNeural': 'ko-KR',
+    'lo-LA-ChanthavongNeural': 'lo-LA',
+    'lo-LA-KeomanyNeural': 'lo-LA',
+    'lt-LT-LeonasNeural': 'lt-LT',
+    'lt-LT-OnaNeural': 'lt-LT',
+    'lv-LV-EveritaNeural': 'lv-LV',
+    'lv-LV-NilsNeural': 'lv-LV',
+    'mk-MK-AleksandarNeural': 'mk-MK',
+    'mk-MK-MarijaNeural': 'mk-MK',
+    'ml-IN-MidhunNeural': 'ml-IN',
+    'ml-IN-SobhanaNeural': 'ml-IN',
+    'mn-MN-BataaNeural': 'mn-MN',
+    'mn-MN-YesuiNeural': 'mn-MN',
+    'mr-IN-AarohiNeural': 'mr-IN',
+    'mr-IN-ManoharNeural': 'mr-IN',
+    'ms-MY-OsmanNeural': 'ms-MY',
+    'ms-MY-YasminNeural': 'ms-MY',
+    'mt-MT-GraceNeural': 'mt-MT',
+    'mt-MT-JosephNeural': 'mt-MT',
+    'my-MM-NilarNeural': 'my-MM',
+    'my-MM-ThihaNeural': 'my-MM',
+    'nb-NO-FinnNeural': 'nb-NO',
+    'nb-NO-PernilleNeural': 'nb-NO',
+    'ne-NP-HemkalaNeural': 'ne-NP',
+    'ne-NP-SagarNeural': 'ne-NP',
+    'nl-BE-ArnaudNeural': 'nl-BE',
+    'nl-BE-DenaNeural': 'nl-BE',
+    'nl-NL-ColetteNeural': 'nl-NL',
+    'nl-NL-FennaNeural': 'nl-NL',
+    'nl-NL-MaartenNeural': 'nl-NL',
+    'pl-PL-MarekNeural': 'pl-PL',
+    'pl-PL-ZofiaNeural': 'pl-PL',
+    'ps-AF-GulNawazNeural': 'ps-AF',
+    'ps-AF-LatifaNeural': 'ps-AF',
+    'pt-BR-AntonioNeural': 'pt-BR',
+    'pt-BR-FranciscaNeural': 'pt-BR',
+    'pt-PT-DuarteNeural': 'pt-PT',
+    'pt-PT-RaquelNeural': 'pt-PT',
+    'ro-RO-AlinaNeural': 'ro-RO',
+    'ro-RO-EmilNeural': 'ro-RO',
+    'ru-RU-DmitryNeural': 'ru-RU',
+    'ru-RU-SvetlanaNeural': 'ru-RU',
+    'si-LK-SameeraNeural': 'si-LK',
+    'si-LK-ThiliniNeural': 'si-LK',
+    'sk-SK-LukasNeural': 'sk-SK',
+    'sk-SK-ViktoriaNeural': 'sk-SK',
+    'sl-SI-PetraNeural': 'sl-SI',
+    'sl-SI-RokNeural': 'sl-SI',
+    'so-SO-MuuseNeural': 'so-SO',
+    'so-SO-UbaxNeural': 'so-SO',
+    'sq-AL-AnilaNeural': 'sq-AL',
+    'sq-AL-IlirNeural': 'sq-AL',
+    'sr-RS-NicholasNeural': 'sr-RS',
+    'sr-RS-SophieNeural': 'sr-RS',
+    'su-ID-JajangNeural': 'su-ID',
+    'su-ID-TutiNeural': 'su-ID',
+    'sv-SE-MattiasNeural': 'sv-SE',
+    'sv-SE-SofieNeural': 'sv-SE',
+    'sw-KE-RafikiNeural': 'sw-KE',
+    'sw-KE-ZuriNeural': 'sw-KE',
+    'sw-TZ-DaudiNeural': 'sw-TZ',
+    'sw-TZ-RehemaNeural': 'sw-TZ',
+    'ta-IN-PallaviNeural': 'ta-IN',
+    'ta-IN-ValluvarNeural': 'ta-IN',
+    'ta-LK-KumarNeural': 'ta-LK',
+    'ta-LK-SaranyaNeural': 'ta-LK',
+    'ta-MY-KaniNeural': 'ta-MY',
+    'ta-MY-SuryaNeural': 'ta-MY',
+    'ta-SG-AnbuNeural': 'ta-SG',
+    'ta-SG-VenbaNeural': 'ta-SG',
+    'te-IN-MohanNeural': 'te-IN',
+    'te-IN-ShrutiNeural': 'te-IN',
+    'th-TH-NiwatNeural': 'th-TH',
+    'th-TH-PremwadeeNeural': 'th-TH',
+    'tr-TR-AhmetNeural': 'tr-TR',
+    'tr-TR-EmelNeural': 'tr-TR',
+    'uk-UA-OstapNeural': 'uk-UA',
+    'uk-UA-PolinaNeural': 'uk-UA',
+    'ur-IN-GulNeural': 'ur-IN',
+    'ur-IN-SalmanNeural': 'ur-IN',
+    'ur-PK-AsadNeural': 'ur-PK',
+    'ur-PK-UzmaNeural': 'ur-PK',
+    'uz-UZ-MadinaNeural': 'uz-UZ',
+    'uz-UZ-SardorNeural': 'uz-UZ',
+    'vi-VN-HoaiMyNeural': 'vi-VN',
+    'vi-VN-NamMinhNeural': 'vi-VN',
+    'zu-ZA-ThandoNeural': 'zu-ZA',
+    'zu-ZA-ThembaNeural': 'zu-ZA',
+}
+
+SUPPORTED_LANGUAGES = [
+    "Auto",
+    *SUPPORTED_VOICES.keys()
+]

flask_api.py

@@ -0,0 +1,60 @@
+import io
+import logging
+
+import soundfile
+import torch
+import torchaudio
+from flask import Flask, request, send_file
+from flask_cors import CORS
+
+from inference.infer_tool import RealTimeVC, Svc
+
+app = Flask(__name__)
+
+CORS(app)
+
+logging.getLogger('numba').setLevel(logging.WARNING)
+
+
+@app.route("/voiceChangeModel", methods=["POST"])
+def voice_change_model():
+    request_form = request.form
+    wave_file = request.files.get("sample", None)
+    # 变调信息
+    f_pitch_change = float(request_form.get("fPitchChange", 0))
+    # DAW所需的采样率
+    daw_sample = int(float(request_form.get("sampleRate", 0)))
+    speaker_id = int(float(request_form.get("sSpeakId", 0)))
+    # http获得wav文件并转换
+    input_wav_path = io.BytesIO(wave_file.read())
+
+    # 模型推理
+    if raw_infer:
+        # out_audio, out_sr = svc_model.infer(speaker_id, f_pitch_change, input_wav_path)
+        out_audio, out_sr = svc_model.infer(speaker_id, f_pitch_change, input_wav_path, cluster_infer_ratio=0,
+                                            auto_predict_f0=False, noice_scale=0.4, f0_filter=False)
+        tar_audio = torchaudio.functional.resample(out_audio, svc_model.target_sample, daw_sample)
+    else:
+        out_audio = svc.process(svc_model, speaker_id, f_pitch_change, input_wav_path, cluster_infer_ratio=0,
+                                auto_predict_f0=False, noice_scale=0.4, f0_filter=False)
+        tar_audio = torchaudio.functional.resample(torch.from_numpy(out_audio), svc_model.target_sample, daw_sample)
+    # 返回音频
+    out_wav_path = io.BytesIO()
+    soundfile.write(out_wav_path, tar_audio.cpu().numpy(), daw_sample, format="wav")
+    out_wav_path.seek(0)
+    return send_file(out_wav_path, download_name="temp.wav", as_attachment=True)
+
+
+if __name__ == '__main__':
+    # 启用则为直接切片合成，False为交叉淡化方式
+    # vst插件调整0.3-0.5s切片时间可以降低延迟，直接切片方法会有连接处爆音、交叉淡化会有轻微重叠声音
+    # 自行选择能接受的方法，或将vst最大切片时间调整为1s，此处设为Ture，延迟大音质稳定一些
+    raw_infer = True
+    # 每个模型和config是唯一对应的
+    model_name = "logs/32k/G_174000-Copy1.pth"
+    config_name = "configs/config.json"
+    cluster_model_path = "logs/44k/kmeans_10000.pt"
+    svc_model = Svc(model_name, config_name, cluster_model_path=cluster_model_path)
+    svc = RealTimeVC()
+    # 此处与vst插件对应，不建议更改
+    app.run(port=6842, host="0.0.0.0", debug=False, threaded=False)

flask_api_full_song.py

@@ -0,0 +1,55 @@
+import io
+
+import numpy as np
+import soundfile
+from flask import Flask, request, send_file
+
+from inference import infer_tool, slicer
+
+app = Flask(__name__)
+
+
+@app.route("/wav2wav", methods=["POST"])
+def wav2wav():
+    request_form = request.form
+    audio_path = request_form.get("audio_path", None)  # wav文件地址
+    tran = int(float(request_form.get("tran", 0)))  # 音调
+    spk = request_form.get("spk", 0)  # 说话人(id或者name都可以,具体看你的config)
+    wav_format = request_form.get("wav_format", 'wav')  # 范围文件格式
+    infer_tool.format_wav(audio_path)
+    chunks = slicer.cut(audio_path, db_thresh=-40)
+    audio_data, audio_sr = slicer.chunks2audio(audio_path, chunks)
+
+    audio = []
+    for (slice_tag, data) in audio_data:
+        print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
+
+        length = int(np.ceil(len(data) / audio_sr * svc_model.target_sample))
+        if slice_tag:
+            print('jump empty segment')
+            _audio = np.zeros(length)
+        else:
+            # padd
+            pad_len = int(audio_sr * 0.5)
+            data = np.concatenate([np.zeros([pad_len]), data, np.zeros([pad_len])])
+            raw_path = io.BytesIO()
+            soundfile.write(raw_path, data, audio_sr, format="wav")
+            raw_path.seek(0)
+            out_audio, out_sr = svc_model.infer(spk, tran, raw_path)
+            svc_model.clear_empty()
+            _audio = out_audio.cpu().numpy()
+            pad_len = int(svc_model.target_sample * 0.5)
+            _audio = _audio[pad_len:-pad_len]
+
+        audio.extend(list(infer_tool.pad_array(_audio, length)))
+    out_wav_path = io.BytesIO()
+    soundfile.write(out_wav_path, audio, svc_model.target_sample, format=wav_format)
+    out_wav_path.seek(0)
+    return send_file(out_wav_path, download_name=f"temp.{wav_format}", as_attachment=True)
+
+
+if __name__ == '__main__':
+    model_name = "logs/44k/G_60000.pth"  # 模型地址
+    config_name = "configs/config.json"  # config地址
+    svc_model = infer_tool.Svc(model_name, config_name)
+    app.run(port=1145, host="0.0.0.0", debug=False, threaded=False)

inference/infer_tool.py

@@ -1,15 +1,16 @@
+import gc
 import hashlib
 import io
 import json
 import logging
 import os
+import pickle
 import time
 from pathlib import Path
-from inference import slicer
-import gc
 
 import librosa
 import numpy as np
+
 # import onnxruntime
 import soundfile
 import torch
@@ -17,10 +18,9 @@ import torchaudio
 
 import cluster
 import utils
-from models import SynthesizerTrn
-
 from diffusion.unit2mel import load_model_vocoder
-import yaml
+from inference import slicer
+from models import SynthesizerTrn
 
 logging.getLogger('matplotlib').setLevel(logging.WARNING)
 
@@ -122,26 +122,27 @@ class Svc(object):
                  diffusion_config_path="configs/diffusion.yaml",
                  shallow_diffusion = False,
                  only_diffusion = False,
+                 spk_mix_enable = False,
+                 feature_retrieval = False
                  ):
         self.net_g_path = net_g_path
         self.only_diffusion = only_diffusion
         self.shallow_diffusion = shallow_diffusion
+        self.feature_retrieval = feature_retrieval
         if device is None:
             self.dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-            # self.dev = torch.device("cpu")
         else:
             self.dev = torch.device(device)
         self.net_g_ms = None
         if not self.only_diffusion:
-            self.hps_ms = utils.get_hparams_from_file(config_path)
+            self.hps_ms = utils.get_hparams_from_file(config_path,True)
             self.target_sample = self.hps_ms.data.sampling_rate
             self.hop_size = self.hps_ms.data.hop_length
             self.spk2id = self.hps_ms.spk
-            try:
-                self.speech_encoder = self.hps_ms.model.speech_encoder
-            except Exception as e:
-                self.speech_encoder = 'vec768l12'
-
+            self.unit_interpolate_mode = self.hps_ms.data.unit_interpolate_mode if self.hps_ms.data.unit_interpolate_mode is not None else 'left'
+            self.vol_embedding = self.hps_ms.model.vol_embedding if self.hps_ms.model.vol_embedding is not None else False
+            self.speech_encoder = self.hps_ms.model.speech_encoder if self.hps_ms.model.speech_encoder is not None else 'vec768l12'
+ 
         self.nsf_hifigan_enhance = nsf_hifigan_enhance
         if self.shallow_diffusion or self.only_diffusion:
             if os.path.exists(diffusion_model_path) and os.path.exists(diffusion_model_path):
@@ -151,13 +152,16 @@ class Svc(object):
                     self.hop_size = self.diffusion_args.data.block_size
                     self.spk2id = self.diffusion_args.spk
                     self.speech_encoder = self.diffusion_args.data.encoder
+                    self.unit_interpolate_mode = self.diffusion_args.data.unit_interpolate_mode if self.diffusion_args.data.unit_interpolate_mode is not None else 'left'
+                if spk_mix_enable:
+                    self.diffusion_model.init_spkmix(len(self.spk2id))
             else:
                 print("No diffusion model or config found. Shallow diffusion mode will False")
                 self.shallow_diffusion = self.only_diffusion = False
                 
         # load hubert and model
         if not self.only_diffusion:
-            self.load_model()
+            self.load_model(spk_mix_enable)
             self.hubert_model = utils.get_speech_encoder(self.speech_encoder,device=self.dev)
             self.volume_extractor = utils.Volume_Extractor(self.hop_size)
         else:
@@ -165,13 +169,23 @@ class Svc(object):
             self.volume_extractor = utils.Volume_Extractor(self.diffusion_args.data.block_size)
             
         if os.path.exists(cluster_model_path):
-            self.cluster_model = cluster.get_cluster_model(cluster_model_path)
-        if self.shallow_diffusion : self.nsf_hifigan_enhance = False
+            if self.feature_retrieval:
+                with open(cluster_model_path,"rb") as f:
+                    self.cluster_model = pickle.load(f)
+                self.big_npy = None
+                self.now_spk_id = -1
+            else:
+                self.cluster_model = cluster.get_cluster_model(cluster_model_path)
+        else:
+            self.feature_retrieval=False
+
+        if self.shallow_diffusion :
+            self.nsf_hifigan_enhance = False
         if self.nsf_hifigan_enhance:
             from modules.enhancer import Enhancer
             self.enhancer = Enhancer('nsf-hifigan', 'pretrain/nsf_hifigan/model',device=self.dev)
             
-    def load_model(self):
+    def load_model(self, spk_mix_enable=False):
         # get model configuration
         self.net_g_ms = SynthesizerTrn(
             self.hps_ms.data.filter_length // 2 + 1,
@@ -182,8 +196,8 @@ class Svc(object):
             _ = self.net_g_ms.half().eval().to(self.dev)
         else:
             _ = self.net_g_ms.eval().to(self.dev)
-
-
+        if spk_mix_enable:
+            self.net_g_ms.EnableCharacterMix(len(self.spk2id), self.dev)
 
     def get_unit_f0(self, wav, tran, cluster_infer_ratio, speaker, f0_filter ,f0_predictor,cr_threshold=0.05):
 
@@ -202,12 +216,33 @@ class Svc(object):
         wav16k = librosa.resample(wav, orig_sr=self.target_sample, target_sr=16000)
         wav16k = torch.from_numpy(wav16k).to(self.dev)
         c = self.hubert_model.encoder(wav16k)
-        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1])
+        c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1],self.unit_interpolate_mode)
 
         if cluster_infer_ratio !=0:
-            cluster_c = cluster.get_cluster_center_result(self.cluster_model, c.cpu().numpy().T, speaker).T
-            cluster_c = torch.FloatTensor(cluster_c).to(self.dev)
-            c = cluster_infer_ratio * cluster_c + (1 - cluster_infer_ratio) * c
+            if self.feature_retrieval:
+                speaker_id = self.spk2id.get(speaker)
+                if speaker_id is None:
+                    raise RuntimeError("The name you entered is not in the speaker list!")
+                if not speaker_id and type(speaker) is int:
+                    if len(self.spk2id.__dict__) >= speaker:
+                        speaker_id = speaker
+                feature_index = self.cluster_model[speaker_id]
+                feat_np = c.transpose(0,1).cpu().numpy()
+                if self.big_npy is None or self.now_spk_id != speaker_id:
+                   self.big_npy = feature_index.reconstruct_n(0, feature_index.ntotal)
+                   self.now_spk_id = speaker_id
+                print("starting feature retrieval...")
+                score, ix = feature_index.search(feat_np, k=8)
+                weight = np.square(1 / score)
+                weight /= weight.sum(axis=1, keepdims=True)
+                npy = np.sum(self.big_npy[ix] * np.expand_dims(weight, axis=2), axis=1)
+                c = cluster_infer_ratio * npy + (1 - cluster_infer_ratio) * feat_np
+                c = torch.FloatTensor(c).to(self.dev).transpose(0,1)
+                print("end feature retrieval...")
+            else:
+                cluster_c = cluster.get_cluster_center_result(self.cluster_model, c.cpu().numpy().T, speaker).T
+                cluster_c = torch.FloatTensor(cluster_c).to(self.dev)
+                c = cluster_infer_ratio * cluster_c + (1 - cluster_infer_ratio) * c
 
         c = c.unsqueeze(0)
         return c, f0, uv
@@ -220,30 +255,47 @@ class Svc(object):
               f0_predictor='pm',
               enhancer_adaptive_key = 0,
               cr_threshold = 0.05,
-              k_step = 100
+              k_step = 100,
+              frame = 0,
+              spk_mix = False,
+              second_encoding = False,
+              loudness_envelope_adjustment = 1
               ):
-
-        speaker_id = self.spk2id.get(speaker)
-        if not speaker_id and type(speaker) is int:
-            if len(self.spk2id.__dict__) >= speaker:
-                speaker_id = speaker
-        sid = torch.LongTensor([int(speaker_id)]).to(self.dev).unsqueeze(0)
         wav, sr = librosa.load(raw_path, sr=self.target_sample)
-        c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+        if spk_mix:
+            c, f0, uv = self.get_unit_f0(wav, tran, 0, None, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+            n_frames = f0.size(1)
+            sid = speaker[:, frame:frame+n_frames].transpose(0,1)
+        else:
+            speaker_id = self.spk2id.get(speaker)
+            if not speaker_id and type(speaker) is int:
+                if len(self.spk2id.__dict__) >= speaker:
+                    speaker_id = speaker
+            if speaker_id is None:
+                raise RuntimeError("The name you entered is not in the speaker list!")
+            sid = torch.LongTensor([int(speaker_id)]).to(self.dev).unsqueeze(0)
+            c, f0, uv = self.get_unit_f0(wav, tran, cluster_infer_ratio, speaker, f0_filter,f0_predictor,cr_threshold=cr_threshold)
+            n_frames = f0.size(1)
         if "half" in self.net_g_path and torch.cuda.is_available():
             c = c.half()
         with torch.no_grad():
             start = time.time()
+            vol = None
             if not self.only_diffusion:
-                audio,f0 = self.net_g_ms.infer(c, f0=f0, g=sid, uv=uv, predict_f0=auto_predict_f0, noice_scale=noice_scale)
+                vol = self.volume_extractor.extract(torch.FloatTensor(wav).to(self.dev)[None,:])[None,:].to(self.dev) if self.vol_embedding else None
+                audio,f0 = self.net_g_ms.infer(c, f0=f0, g=sid, uv=uv, predict_f0=auto_predict_f0, noice_scale=noice_scale,vol=vol)
                 audio = audio[0,0].data.float()
-                if self.shallow_diffusion:
-                    audio_mel = self.vocoder.extract(audio[None,:],self.target_sample)
+                audio_mel = self.vocoder.extract(audio[None,:],self.target_sample) if self.shallow_diffusion else None
             else:
                 audio = torch.FloatTensor(wav).to(self.dev)
                 audio_mel = None
             if self.only_diffusion or self.shallow_diffusion:
-                vol = self.volume_extractor.extract(audio[None,:])[None,:,None].to(self.dev)
+                vol = self.volume_extractor.extract(audio[None,:])[None,:,None].to(self.dev) if vol is None else vol[:,:,None]
+                if self.shallow_diffusion and second_encoding:
+                    audio16k = librosa.resample(audio.detach().cpu().numpy(), orig_sr=self.target_sample, target_sr=16000)
+                    audio16k = torch.from_numpy(audio16k).to(self.dev)
+                    c = self.hubert_model.encoder(audio16k)
+                    c = utils.repeat_expand_2d(c.squeeze(0), f0.shape[1],self.unit_interpolate_mode)
                 f0 = f0[:,:,None]
                 c = c.transpose(-1,-2)
                 audio_mel = self.diffusion_model(
@@ -265,9 +317,11 @@ class Svc(object):
                                     f0[:,:,None], 
                                     self.hps_ms.data.hop_length, 
                                     adaptive_key = enhancer_adaptive_key)
+            if loudness_envelope_adjustment != 1:
+                audio = utils.change_rms(wav,self.target_sample,audio,self.target_sample,loudness_envelope_adjustment)
             use_time = time.time() - start
             print("vits use time:{}".format(use_time))
-        return audio, audio.shape[-1]
+        return audio, audio.shape[-1], n_frames
 
     def clear_empty(self):
         # clean up vram
@@ -298,8 +352,15 @@ class Svc(object):
                         f0_predictor='pm',
                         enhancer_adaptive_key = 0,
                         cr_threshold = 0.05,
-                        k_step = 100
+                        k_step = 100,
+                        use_spk_mix = False,
+                        second_encoding = False,
+                        loudness_envelope_adjustment = 1
                         ):
+        if use_spk_mix:
+            if len(self.spk2id) == 1:
+                spk = self.spk2id.keys()[0]
+                use_spk_mix = False
         wav_path = Path(raw_audio_path).with_suffix('.wav')
         chunks = slicer.cut(wav_path, db_thresh=slice_db)
         audio_data, audio_sr = slicer.chunks2audio(wav_path, chunks)
@@ -309,7 +370,62 @@ class Svc(object):
         lg_size_c_l = (lg_size-lg_size_r)//2
         lg_size_c_r = lg_size-lg_size_r-lg_size_c_l
         lg = np.linspace(0,1,lg_size_r) if lg_size!=0 else 0
-        
+
+        if use_spk_mix:
+            assert len(self.spk2id) == len(spk)
+            audio_length = 0
+            for (slice_tag, data) in audio_data:
+                aud_length = int(np.ceil(len(data) / audio_sr * self.target_sample))
+                if slice_tag:
+                    audio_length += aud_length // self.hop_size
+                    continue
+                if per_size != 0:
+                    datas = split_list_by_n(data, per_size,lg_size)
+                else:
+                    datas = [data]
+                for k,dat in enumerate(datas):
+                    pad_len = int(audio_sr * pad_seconds)
+                    per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample))
+                    a_length = per_length + 2 * pad_len
+                    audio_length += a_length // self.hop_size
+            audio_length += len(audio_data)
+            spk_mix_tensor = torch.zeros(size=(len(spk), audio_length)).to(self.dev)
+            for i in range(len(spk)):
+                last_end = None
+                for mix in spk[i]:
+                    if mix[3]<0. or mix[2]<0.:
+                        raise RuntimeError("mix value must higer Than zero!")
+                    begin = int(audio_length * mix[0])
+                    end = int(audio_length * mix[1])
+                    length = end - begin
+                    if length<=0:                        
+                        raise RuntimeError("begin Must lower Than end!")
+                    step = (mix[3] - mix[2])/length
+                    if last_end is not None:
+                        if last_end != begin:
+                            raise RuntimeError("[i]EndTime Must Equal [i+1]BeginTime!")
+                    last_end = end
+                    if step == 0.:
+                        spk_mix_data = torch.zeros(length).to(self.dev) + mix[2]
+                    else:
+                        spk_mix_data = torch.arange(mix[2],mix[3],step).to(self.dev)
+                    if(len(spk_mix_data)<length):
+                        num_pad = length - len(spk_mix_data)
+                        spk_mix_data = torch.nn.functional.pad(spk_mix_data, [0, num_pad], mode="reflect").to(self.dev)
+                    spk_mix_tensor[i][begin:end] = spk_mix_data[:length]
+
+            spk_mix_ten = torch.sum(spk_mix_tensor,dim=0).unsqueeze(0).to(self.dev)
+            # spk_mix_tensor[0][spk_mix_ten<0.001] = 1.0
+            for i, x in enumerate(spk_mix_ten[0]):
+                if x == 0.0:
+                    spk_mix_ten[0][i] = 1.0
+                    spk_mix_tensor[:,i] = 1.0 / len(spk)
+            spk_mix_tensor = spk_mix_tensor / spk_mix_ten
+            if not ((torch.sum(spk_mix_tensor,dim=0) - 1.)<0.0001).all():
+                raise RuntimeError("sum(spk_mix_tensor) not equal 1")
+            spk = spk_mix_tensor
+
+        global_frame = 0
         audio = []
         for (slice_tag, data) in audio_data:
             print(f'#=====segment start, {round(len(data) / audio_sr, 3)}s======')
@@ -319,6 +435,7 @@ class Svc(object):
                 print('jump empty segment')
                 _audio = np.zeros(length)
                 audio.extend(list(pad_array(_audio, length)))
+                global_frame += length // self.hop_size
                 continue
             if per_size != 0:
                 datas = split_list_by_n(data, per_size,lg_size)
@@ -326,22 +443,28 @@ class Svc(object):
                 datas = [data]
             for k,dat in enumerate(datas):
                 per_length = int(np.ceil(len(dat) / audio_sr * self.target_sample)) if clip_seconds!=0 else length
-                if clip_seconds!=0: print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
+                if clip_seconds!=0: 
+                    print(f'###=====segment clip start, {round(len(dat) / audio_sr, 3)}s======')
                 # padd
                 pad_len = int(audio_sr * pad_seconds)
                 dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
                 raw_path = io.BytesIO()
                 soundfile.write(raw_path, dat, audio_sr, format="wav")
                 raw_path.seek(0)
-                out_audio, out_sr = self.infer(spk, tran, raw_path,
+                out_audio, out_sr, out_frame = self.infer(spk, tran, raw_path,
                                                     cluster_infer_ratio=cluster_infer_ratio,
                                                     auto_predict_f0=auto_predict_f0,
                                                     noice_scale=noice_scale,
                                                     f0_predictor = f0_predictor,
                                                     enhancer_adaptive_key = enhancer_adaptive_key,
                                                     cr_threshold = cr_threshold,
-                                                    k_step = k_step
+                                                    k_step = k_step,
+                                                    frame = global_frame,
+                                                    spk_mix = use_spk_mix,
+                                                    second_encoding = second_encoding,
+                                                    loudness_envelope_adjustment = loudness_envelope_adjustment
                                                     )
+                global_frame += out_frame
                 _audio = out_audio.cpu().numpy()
                 pad_len = int(self.target_sample * pad_seconds)
                 _audio = _audio[pad_len:-pad_len]
@@ -404,4 +527,4 @@ class RealTimeVC:
             self.last_chunk = audio[-self.pre_len:]
             self.last_o = audio
             return ret[self.chunk_len:2 * self.chunk_len]
-            
+            

inference/infer_tool_grad.py

@@ -1,26 +1,21 @@
-import hashlib
-import json
+import io
 import logging
 import os
-import time
-from pathlib import Path
-import io
+
 import librosa
-import maad
 import numpy as np
-from inference import slicer
 import parselmouth
 import soundfile
 import torch
 import torchaudio
 
-# from hubert import hubert_model
 import utils
+from inference import slicer
 from models import SynthesizerTrn
+
 logging.getLogger('numba').setLevel(logging.WARNING)
 logging.getLogger('matplotlib').setLevel(logging.WARNING)
 
-
 def resize2d_f0(x, target_len):
     source = np.array(x)
     source[source < 0.001] = np.nan
@@ -29,8 +24,7 @@ def resize2d_f0(x, target_len):
     res = np.nan_to_num(target)
     return res
 
-
-def get_f0(x, p_len, f0_up_key=0):
+def get_f0(x, p_len,f0_up_key=0):
 
     time_step = 160 / 16000 * 1000
     f0_min = 50
@@ -42,21 +36,18 @@ def get_f0(x, p_len, f0_up_key=0):
         time_step=time_step / 1000, voicing_threshold=0.6,
         pitch_floor=f0_min, pitch_ceiling=f0_max).selected_array['frequency']
 
-    pad_size = (p_len - len(f0) + 1) // 2
-    if(pad_size > 0 or p_len - len(f0) - pad_size > 0):
-        f0 = np.pad(
-            f0, [[pad_size, p_len - len(f0) - pad_size]], mode='constant')
+    pad_size=(p_len - len(f0) + 1) // 2
+    if(pad_size>0 or p_len - len(f0) - pad_size>0):
+        f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
 
     f0 *= pow(2, f0_up_key / 12)
     f0_mel = 1127 * np.log(1 + f0 / 700)
-    f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * \
-        254 / (f0_mel_max - f0_mel_min) + 1
+    f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * 254 / (f0_mel_max - f0_mel_min) + 1
     f0_mel[f0_mel <= 1] = 1
     f0_mel[f0_mel > 255] = 255
     f0_coarse = np.rint(f0_mel).astype(np.int)
     return f0_coarse, f0
 
-
 def clean_pitch(input_pitch):
     num_nan = np.sum(input_pitch == 1)
     if num_nan / len(input_pitch) > 0.9:
@@ -89,8 +80,7 @@ def mkdir(paths: list):
 
 class VitsSvc(object):
     def __init__(self):
-        self.device = torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.SVCVITS = None
         self.hps = None
         self.speakers = None
@@ -99,18 +89,16 @@ class VitsSvc(object):
     def set_device(self, device):
         self.device = torch.device(device)
         self.hubert_soft.to(self.device)
-        if self.SVCVITS != None:
+        if self.SVCVITS is not None:
             self.SVCVITS.to(self.device)
 
     def loadCheckpoint(self, path):
-        self.hps = utils.get_hparams_from_file(
-            f"checkpoints/{path}/config.json")
+        self.hps = utils.get_hparams_from_file(f"checkpoints/{path}/config.json")
         self.SVCVITS = SynthesizerTrn(
             self.hps.data.filter_length // 2 + 1,
             self.hps.train.segment_size // self.hps.data.hop_length,
             **self.hps.model)
-        _ = utils.load_checkpoint(
-            f"checkpoints/{path}/model.pth", self.SVCVITS, None)
+        _ = utils.load_checkpoint(f"checkpoints/{path}/model.pth", self.SVCVITS, None)
         _ = self.SVCVITS.eval().to(self.device)
         self.speakers = self.hps.spk
 
@@ -120,6 +108,7 @@ class VitsSvc(object):
             units = self.hubert_soft.units(source)
             return units
 
+
     def get_unit_pitch(self, in_path, tran):
         source, sr = torchaudio.load(in_path)
         source = torchaudio.functional.resample(source, sr, 16000)
@@ -137,27 +126,23 @@ class VitsSvc(object):
         stn_tst = torch.FloatTensor(soft)
         with torch.no_grad():
             x_tst = stn_tst.unsqueeze(0).to(self.device)
-            x_tst = torch.repeat_interleave(
-                x_tst, repeats=2, dim=1).transpose(1, 2)
-            audio, _ = self.SVCVITS.infer(x_tst, f0=f0, g=sid)[
-                0, 0].data.float()
+            x_tst = torch.repeat_interleave(x_tst, repeats=2, dim=1).transpose(1, 2)
+            audio,_ = self.SVCVITS.infer(x_tst, f0=f0, g=sid)[0,0].data.float()
         return audio, audio.shape[-1]
 
-    def inference(self, srcaudio, chara, tran, slice_db):
+    def inference(self,srcaudio,chara,tran,slice_db):
         sampling_rate, audio = srcaudio
         audio = (audio / np.iinfo(audio.dtype).max).astype(np.float32)
         if len(audio.shape) > 1:
             audio = librosa.to_mono(audio.transpose(1, 0))
         if sampling_rate != 16000:
-            audio = librosa.resample(
-                audio, orig_sr=sampling_rate, target_sr=16000)
+            audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
         soundfile.write("tmpwav.wav", audio, 16000, format="wav")
         chunks = slicer.cut("tmpwav.wav", db_thresh=slice_db)
         audio_data, audio_sr = slicer.chunks2audio("tmpwav.wav", chunks)
         audio = []
         for (slice_tag, data) in audio_data:
-            length = int(np.ceil(len(data) / audio_sr *
-                                 self.hps.data.sampling_rate))
+            length = int(np.ceil(len(data) / audio_sr * self.hps.data.sampling_rate))
             raw_path = io.BytesIO()
             soundfile.write(raw_path, data, audio_sr, format="wav")
             raw_path.seek(0)
@@ -168,4 +153,4 @@ class VitsSvc(object):
                 _audio = out_audio.cpu().numpy()
             audio.extend(list(_audio))
         audio = (np.array(audio) * 32768.0).astype('int16')
-        return (self.hps.data.sampling_rate, audio)
+        return (self.hps.data.sampling_rate,audio)

inference_main.py

@@ -0,0 +1,155 @@
+import logging
+
+import soundfile
+
+from inference import infer_tool
+from inference.infer_tool import Svc
+from spkmix import spk_mix_map
+
+logging.getLogger('numba').setLevel(logging.WARNING)
+chunks_dict = infer_tool.read_temp("inference/chunks_temp.json")
+
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser(description='sovits4 inference')
+
+    # 一定要设置的部分
+    parser.add_argument('-m', '--model_path', type=str, default="logs/44k/G_37600.pth", help='模型路径')
+    parser.add_argument('-c', '--config_path', type=str, default="logs/44k/config.json", help='配置文件路径')
+    parser.add_argument('-cl', '--clip', type=float, default=0, help='音频强制切片，默认0为自动切片，单位为秒/s')
+    parser.add_argument('-n', '--clean_names', type=str, nargs='+', default=["君の知らない物語-src.wav"], help='wav文件名列表，放在raw文件夹下')
+    parser.add_argument('-t', '--trans', type=int, nargs='+', default=[0], help='音高调整，支持正负（半音）')
+    parser.add_argument('-s', '--spk_list', type=str, nargs='+', default=['buyizi'], help='合成目标说话人名称')
+    
+    # 可选项部分
+    parser.add_argument('-a', '--auto_predict_f0', action='store_true', default=False, help='语音转换自动预测音高，转换歌声时不要打开这个会严重跑调')
+    parser.add_argument('-cm', '--cluster_model_path', type=str, default="", help='聚类模型或特征检索索引路径，留空则自动设为各方案模型的默认路径，如果没有训练聚类或特征检索则随便填')
+    parser.add_argument('-cr', '--cluster_infer_ratio', type=float, default=0, help='聚类方案或特征检索占比，范围0-1，若没有训练聚类模型或特征检索则默认0即可')
+    parser.add_argument('-lg', '--linear_gradient', type=float, default=0, help='两段音频切片的交叉淡入长度，如果强制切片后出现人声不连贯可调整该数值，如果连贯建议采用默认值0，单位为秒')
+    parser.add_argument('-f0p', '--f0_predictor', type=str, default="pm", help='选择F0预测器,可选择crepe,pm,dio,harvest,默认为pm(注意：crepe为原F0使用均值滤波器)')
+    parser.add_argument('-eh', '--enhance', action='store_true', default=False, help='是否使用NSF_HIFIGAN增强器,该选项对部分训练集少的模型有一定的音质增强效果，但是对训练好的模型有反面效果，默认关闭')
+    parser.add_argument('-shd', '--shallow_diffusion', action='store_true', default=False, help='是否使用浅层扩散，使用后可解决一部分电音问题，默认关闭，该选项打开时，NSF_HIFIGAN增强器将会被禁止')
+    parser.add_argument('-usm', '--use_spk_mix', action='store_true', default=False, help='是否使用角色融合')
+    parser.add_argument('-lea', '--loudness_envelope_adjustment', type=float, default=1, help='输入源响度包络替换输出响度包络融合比例，越靠近1越使用输出响度包络')
+    parser.add_argument('-fr', '--feature_retrieval', action='store_true', default=False, help='是否使用特征检索，如果使用聚类模型将被禁用，且cm与cr参数将会变成特征检索的索引路径与混合比例')
+
+    # 浅扩散设置
+    parser.add_argument('-dm', '--diffusion_model_path', type=str, default="logs/44k/diffusion/model_0.pt", help='扩散模型路径')
+    parser.add_argument('-dc', '--diffusion_config_path', type=str, default="logs/44k/diffusion/config.yaml", help='扩散模型配置文件路径')
+    parser.add_argument('-ks', '--k_step', type=int, default=100, help='扩散步数，越大越接近扩散模型的结果，默认100')
+    parser.add_argument('-se', '--second_encoding', action='store_true', default=False, help='二次编码，浅扩散前会对原始音频进行二次编码，玄学选项，有时候效果好，有时候效果差')
+    parser.add_argument('-od', '--only_diffusion', action='store_true', default=False, help='纯扩散模式，该模式不会加载sovits模型，以扩散模型推理')
+    
+
+    # 不用动的部分
+    parser.add_argument('-sd', '--slice_db', type=int, default=-40, help='默认-40，嘈杂的音频可以-30，干声保留呼吸可以-50')
+    parser.add_argument('-d', '--device', type=str, default=None, help='推理设备，None则为自动选择cpu和gpu')
+    parser.add_argument('-ns', '--noice_scale', type=float, default=0.4, help='噪音级别，会影响咬字和音质，较为玄学')
+    parser.add_argument('-p', '--pad_seconds', type=float, default=0.5, help='推理音频pad秒数，由于未知原因开头结尾会有异响，pad一小段静音段后就不会出现')
+    parser.add_argument('-wf', '--wav_format', type=str, default='flac', help='音频输出格式')
+    parser.add_argument('-lgr', '--linear_gradient_retain', type=float, default=0.75, help='自动音频切片后，需要舍弃每段切片的头尾。该参数设置交叉长度保留的比例，范围0-1,左开右闭')
+    parser.add_argument('-eak', '--enhancer_adaptive_key', type=int, default=0, help='使增强器适应更高的音域(单位为半音数)|默认为0')
+    parser.add_argument('-ft', '--f0_filter_threshold', type=float, default=0.05,help='F0过滤阈值，只有使用crepe时有效. 数值范围从0-1. 降低该值可减少跑调概率，但会增加哑音')
+
+
+    args = parser.parse_args()
+
+    clean_names = args.clean_names
+    trans = args.trans
+    spk_list = args.spk_list
+    slice_db = args.slice_db
+    wav_format = args.wav_format
+    auto_predict_f0 = args.auto_predict_f0
+    cluster_infer_ratio = args.cluster_infer_ratio
+    noice_scale = args.noice_scale
+    pad_seconds = args.pad_seconds
+    clip = args.clip
+    lg = args.linear_gradient
+    lgr = args.linear_gradient_retain
+    f0p = args.f0_predictor
+    enhance = args.enhance
+    enhancer_adaptive_key = args.enhancer_adaptive_key
+    cr_threshold = args.f0_filter_threshold
+    diffusion_model_path = args.diffusion_model_path
+    diffusion_config_path = args.diffusion_config_path
+    k_step = args.k_step
+    only_diffusion = args.only_diffusion
+    shallow_diffusion = args.shallow_diffusion
+    use_spk_mix = args.use_spk_mix
+    second_encoding = args.second_encoding
+    loudness_envelope_adjustment = args.loudness_envelope_adjustment
+
+    if cluster_infer_ratio != 0:
+        if args.cluster_model_path == "":
+            if args.feature_retrieval:  # 若指定了占比但没有指定模型路径，则按是否使用特征检索分配默认的模型路径
+                args.cluster_model_path = "logs/44k/feature_and_index.pkl"
+            else:
+                args.cluster_model_path = "logs/44k/kmeans_10000.pt"
+    else:  # 若未指定占比，则无论是否指定模型路径，都将其置空以避免之后的模型加载
+        args.cluster_model_path = ""
+
+    svc_model = Svc(args.model_path,
+                    args.config_path,
+                    args.device,
+                    args.cluster_model_path,
+                    enhance,
+                    diffusion_model_path,
+                    diffusion_config_path,
+                    shallow_diffusion,
+                    only_diffusion,
+                    use_spk_mix,
+                    args.feature_retrieval)
+    
+    infer_tool.mkdir(["raw", "results"])
+    
+    if len(spk_mix_map)<=1:
+        use_spk_mix = False
+    if use_spk_mix:
+        spk_list = [spk_mix_map]
+    
+    infer_tool.fill_a_to_b(trans, clean_names)
+    for clean_name, tran in zip(clean_names, trans):
+        raw_audio_path = f"raw/{clean_name}"
+        if "." not in raw_audio_path:
+            raw_audio_path += ".wav"
+        infer_tool.format_wav(raw_audio_path)
+        for spk in spk_list:
+            kwarg = {
+                "raw_audio_path" : raw_audio_path,
+                "spk" : spk,
+                "tran" : tran,
+                "slice_db" : slice_db,
+                "cluster_infer_ratio" : cluster_infer_ratio,
+                "auto_predict_f0" : auto_predict_f0,
+                "noice_scale" : noice_scale,
+                "pad_seconds" : pad_seconds,
+                "clip_seconds" : clip,
+                "lg_num": lg,
+                "lgr_num" : lgr,
+                "f0_predictor" : f0p,
+                "enhancer_adaptive_key" : enhancer_adaptive_key,
+                "cr_threshold" : cr_threshold,
+                "k_step":k_step,
+                "use_spk_mix":use_spk_mix,
+                "second_encoding":second_encoding,
+                "loudness_envelope_adjustment":loudness_envelope_adjustment
+            }
+            audio = svc_model.slice_inference(**kwarg)
+            key = "auto" if auto_predict_f0 else f"{tran}key"
+            cluster_name = "" if cluster_infer_ratio == 0 else f"_{cluster_infer_ratio}"
+            isdiffusion = "sovits"
+            if shallow_diffusion :
+                isdiffusion = "sovdiff"
+            if only_diffusion :
+                isdiffusion = "diff"
+            if use_spk_mix:
+                spk = "spk_mix"
+            res_path = f'results/{clean_name}_{key}_{spk}{cluster_name}_{isdiffusion}_{f0p}.{wav_format}'
+            soundfile.write(res_path, audio, svc_model.target_sample, format=wav_format)
+            svc_model.clear_empty()
+            
+if __name__ == '__main__':
+    main()

modules/DSConv.py

@@ -0,0 +1,76 @@
+import torch.nn as nn
+from torch.nn.utils import remove_weight_norm, weight_norm
+
+
+class Depthwise_Separable_Conv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = 0,
+        dilation = 1,
+        bias = True,
+        padding_mode = 'zeros',  # TODO: refine this type
+        device=None,
+        dtype=None
+    ):
+      super().__init__()
+      self.depth_conv = nn.Conv1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
+      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
+    
+    def forward(self, input):
+      return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = weight_norm(self.point_conv, name = 'weight')
+
+    def remove_weight_norm(self):
+      self.depth_conv = remove_weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = remove_weight_norm(self.point_conv, name = 'weight')
+
+class Depthwise_Separable_TransposeConv1D(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = 0, 
+        output_padding = 0,
+        bias = True,
+        dilation = 1,
+        padding_mode = 'zeros',  # TODO: refine this type
+        device=None,
+        dtype=None
+    ):
+      super().__init__()
+      self.depth_conv = nn.ConvTranspose1d(in_channels=in_channels, out_channels=in_channels, kernel_size=kernel_size, groups=in_channels,stride = stride,output_padding=output_padding,padding=padding,dilation=dilation,bias=bias,padding_mode=padding_mode,device=device,dtype=dtype)
+      self.point_conv = nn.Conv1d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias, device=device,dtype=dtype)
+    
+    def forward(self, input):
+      return self.point_conv(self.depth_conv(input))
+
+    def weight_norm(self):
+      self.depth_conv = weight_norm(self.depth_conv, name = 'weight')
+      self.point_conv = weight_norm(self.point_conv, name = 'weight')
+
+    def remove_weight_norm(self):
+      remove_weight_norm(self.depth_conv, name = 'weight')
+      remove_weight_norm(self.point_conv, name = 'weight')
+
+
+def weight_norm_modules(module, name = 'weight', dim = 0):
+    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
+      module.weight_norm()
+      return module
+    else:
+      return weight_norm(module,name,dim)
+
+def remove_weight_norm_modules(module, name = 'weight'):
+    if isinstance(module,Depthwise_Separable_Conv1D) or isinstance(module,Depthwise_Separable_TransposeConv1D):
+      module.remove_weight_norm()
+    else:
+      remove_weight_norm(module,name)

modules/F0Predictor/CrepeF0Predictor.py

@@ -1,7 +1,9 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-from modules.F0Predictor.crepe import CrepePitchExtractor
 import torch
 
+from modules.F0Predictor.crepe import CrepePitchExtractor
+from modules.F0Predictor.F0Predictor import F0Predictor
+
+
 class CrepeF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,device=None,sampling_rate=44100,threshold=0.05,model="full"):
         self.F0Creper = CrepePitchExtractor(hop_length=hop_length,f0_min=f0_min,f0_max=f0_max,device=device,threshold=threshold,model=model)

modules/F0Predictor/DioF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import pyworld
 import numpy as np
+import pyworld
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class DioF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -13,39 +15,25 @@ class DioF0Predictor(F0Predictor):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
-    
-        return ip_data[:,0], vuv_vector[:,0]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
+
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
 
     def resize_f0(self,x, target_len):
         source = np.array(x)

modules/F0Predictor/HarvestF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import pyworld
 import numpy as np
+import pyworld
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class HarvestF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -13,40 +15,25 @@ class HarvestF0Predictor(F0Predictor):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
-    
-        return ip_data[:,0], vuv_vector[:,0]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
 
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
     def resize_f0(self,x, target_len):
         source = np.array(x)
         source[source<0.001] = np.nan

modules/F0Predictor/PMF0Predictor.py

@@ -1,6 +1,8 @@
-from modules.F0Predictor.F0Predictor import F0Predictor
-import parselmouth
 import numpy as np
+import parselmouth
+
+from modules.F0Predictor.F0Predictor import F0Predictor
+
 
 class PMF0Predictor(F0Predictor):
     def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
@@ -14,39 +16,26 @@ class PMF0Predictor(F0Predictor):
         '''
         对F0进行插值处理
         '''
+        vuv_vector = np.zeros_like(f0, dtype=np.float32)
+        vuv_vector[f0 > 0.0] = 1.0
+        vuv_vector[f0 <= 0.0] = 0.0
     
-        data = np.reshape(f0, (f0.size, 1))
-    
-        vuv_vector = np.zeros((data.size, 1), dtype=np.float32)
-        vuv_vector[data > 0.0] = 1.0
-        vuv_vector[data <= 0.0] = 0.0
-    
-        ip_data = data
-    
-        frame_number = data.size
-        last_value = 0.0
-        for i in range(frame_number):
-            if data[i] <= 0.0:
-                j = i + 1
-                for j in range(i + 1, frame_number):
-                    if data[j] > 0.0:
-                        break
-                if j < frame_number - 1:
-                    if last_value > 0.0:
-                        step = (data[j] - data[i - 1]) / float(j - i)
-                        for k in range(i, j):
-                            ip_data[k] = data[i - 1] + step * (k - i + 1)
-                    else:
-                        for k in range(i, j):
-                            ip_data[k] = data[j]
-                else:
-                    for k in range(i, frame_number):
-                        ip_data[k] = last_value
-            else:
-                ip_data[i] = data[i] #这里可能存在一个没有必要的拷贝
-                last_value = data[i]
+        nzindex = np.nonzero(f0)[0]
+        data = f0[nzindex]
+        nzindex = nzindex.astype(np.float32)
+        time_org = self.hop_length / self.sampling_rate * nzindex
+        time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
+
+        if data.shape[0] <= 0:
+            return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
+
+        if data.shape[0] == 1:
+            return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
+
+        f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
+        
+        return f0,vuv_vector
     
-        return ip_data[:,0], vuv_vector[:,0]
 
     def compute_f0(self,wav,p_len=None):
         x = wav

modules/F0Predictor/crepe.py

@@ -1,14 +1,14 @@
-from typing import Optional,Union
+from typing import Optional, Union
+
 try:
     from typing import Literal
-except Exception as e:
+except Exception:
     from typing_extensions import Literal
 import numpy as np
 import torch
 import torchcrepe
 from torch import nn
 from torch.nn import functional as F
-import scipy
 
 #from:https://github.com/fishaudio/fish-diffusion
 
@@ -97,19 +97,19 @@ class BasePitchExtractor:
         f0 = torch.index_select(f0, dim=0, index=nzindex).cpu().numpy()
         time_org = self.hop_length / sampling_rate * nzindex.cpu().numpy()
         time_frame = np.arange(pad_to) * self.hop_length / sampling_rate
+        
+        vuv_vector = F.interpolate(vuv_vector[None,None,:],size=pad_to)[0][0]
 
         if f0.shape[0] <= 0:
-            return torch.zeros(pad_to, dtype=torch.float, device=x.device),torch.zeros(pad_to, dtype=torch.float, device=x.device)
-
+            return torch.zeros(pad_to, dtype=torch.float, device=x.device),vuv_vector.cpu().numpy()
         if f0.shape[0] == 1:
-            return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],torch.ones(pad_to, dtype=torch.float, device=x.device)
+            return torch.ones(pad_to, dtype=torch.float, device=x.device) * f0[0],vuv_vector.cpu().numpy()
     
         # 大概可以用 torch 重写?
         f0 = np.interp(time_frame, time_org, f0, left=f0[0], right=f0[-1])
-        vuv_vector = vuv_vector.cpu().numpy()
-        vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
+        #vuv_vector = np.ceil(scipy.ndimage.zoom(vuv_vector,pad_to/len(vuv_vector),order = 0))
         
-        return f0,vuv_vector
+        return f0,vuv_vector.cpu().numpy()
 
 
 class MaskedAvgPool1d(nn.Module):
@@ -323,7 +323,7 @@ class CrepePitchExtractor(BasePitchExtractor):
         else:
             pd = torchcrepe.filter.median(pd, 3)
 
-        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, 512)
+        pd = torchcrepe.threshold.Silence(-60.0)(pd, x, sampling_rate, self.hop_length)
         f0 = torchcrepe.threshold.At(self.threshold)(f0, pd)
         
         if self.use_fast_filters:
@@ -334,7 +334,7 @@ class CrepePitchExtractor(BasePitchExtractor):
         f0 = torch.where(torch.isnan(f0), torch.full_like(f0, 0), f0)[0]
 
         if torch.all(f0 == 0):
-            rtn = f0.cpu().numpy() if pad_to==None else np.zeros(pad_to)
+            rtn = f0.cpu().numpy() if pad_to is None else np.zeros(pad_to)
             return rtn,rtn
         
         return self.post_process(x, sampling_rate, f0, pad_to)

modules/attentions.py

@@ -1,12 +1,10 @@
-import copy
 import math
-import numpy as np
+
 import torch
 from torch import nn
 from torch.nn import functional as F
 
 import modules.commons as commons
-import modules.modules as modules
 from modules.modules import LayerNorm
 
 
@@ -243,7 +241,7 @@ class MultiHeadAttention(nn.Module):
     return ret
 
   def _get_relative_embeddings(self, relative_embeddings, length):
-    max_relative_position = 2 * self.window_size + 1
+    2 * self.window_size + 1
     # Pad first before slice to avoid using cond ops.
     pad_length = max(length - (self.window_size + 1), 0)
     slice_start_position = max((self.window_size + 1) - length, 0)

modules/commons.py

@@ -1,9 +1,9 @@
 import math
-import numpy as np
+
 import torch
-from torch import nn
 from torch.nn import functional as F
 
+
 def slice_pitch_segments(x, ids_str, segment_size=4):
   ret = torch.zeros_like(x[:, :segment_size])
   for i in range(x.size(0)):
@@ -24,10 +24,12 @@ def rand_slice_segments_with_pitch(x, pitch, x_lengths=None, segment_size=4):
 
 def init_weights(m, mean=0.0, std=0.01):
   classname = m.__class__.__name__
-  if classname.find("Conv") != -1:
+  if "Depthwise_Separable" in classname:
+    m.depth_conv.weight.data.normal_(mean, std)
+    m.point_conv.weight.data.normal_(mean, std) 
+  elif classname.find("Conv") != -1:
     m.weight.data.normal_(mean, std)
 
-
 def get_padding(kernel_size, dilation=1):
   return int((kernel_size*dilation - dilation)/2)
 
@@ -134,12 +136,6 @@ def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels):
   return acts
 
 
-def convert_pad_shape(pad_shape):
-  l = pad_shape[::-1]
-  pad_shape = [item for sublist in l for item in sublist]
-  return pad_shape
-
-
 def shift_1d(x):
   x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
   return x
@@ -157,7 +153,6 @@ def generate_path(duration, mask):
   duration: [b, 1, t_x]
   mask: [b, 1, t_y, t_x]
   """
-  device = duration.device
   
   b, _, t_y, t_x = mask.shape
   cum_duration = torch.cumsum(duration, -1)

modules/enhancer.py

@@ -1,10 +1,12 @@
 import numpy as np
 import torch
 import torch.nn.functional as F
-from vdecoder.nsf_hifigan.nvSTFT import STFT
-from vdecoder.nsf_hifigan.models import load_model
 from torchaudio.transforms import Resample
 
+from vdecoder.nsf_hifigan.models import load_model
+from vdecoder.nsf_hifigan.nvSTFT import STFT
+
+
 class Enhancer:
     def __init__(self, enhancer_type, enhancer_ckpt, device=None):
         if device is None:

modules/losses.py

@@ -1,7 +1,4 @@
-import torch 
-from torch.nn import functional as F
-
-import modules.commons as commons
+import torch
 
 
 def feature_loss(fmap_r, fmap_g):