upload files

RepCodec/.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+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/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__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/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/

RepCodec/LICENSE

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

RepCodec/README.md

@@ -0,0 +1,273 @@
+# RepCodec: A Speech Representation Codec for Speech Tokenization
+
+> [**RepCodec: A Speech Representation Codec for Speech Tokenization**](https://arxiv.org/abs/2309.00169)
+
+## Introduction
+
+**RepCodec** is a speech tokenization method for converting a speech waveform into a sequence of discrete semantic
+tokens.
+The main idea is to train a representation codec which learns a vector quantization codebook through reconstructing the
+input speech representations from speech encoders like HuBERT or data2vec.
+Extensive experiments show that RepCodec significantly outperforms the widely used k-means clustering approach in both
+speech understanding and generation.
+Also, RepCodec generalizes well across various speech encoders and languages.
+
+<img src="images/RepCodec.png" alt="se" width="1000" />
+
+## RepCodec Models
+
+| Feature Type                                                                                                          | Speech Data                                              | RepCodec Model                                                                                           |
+|-----------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------|----------------------------------------------------------------------------------------------------------|
+| [HuBERT base](https://github.com/facebookresearch/fairseq/tree/main/examples/hubert#pre-trained-and-fine-tuned-asr-models) layer 9     | [Librispeech](http://www.openslr.org/12) train-clean-100 | [hubert_base_l9](https://drive.google.com/file/d/1XD0HKl607FFjri2-VJT7lHQeSpxsCCFO/view?usp=sharing)     |
+| [HuBERT large](https://github.com/facebookresearch/fairseq/tree/main/examples/hubert#pre-trained-and-fine-tuned-asr-models) layer 18   | [Librispeech](http://www.openslr.org/12) train-clean-100 | [hubert_large_l18](https://drive.google.com/file/d/1mTbm5GeJ7gp_5L3QLP-JGXdf8RnRw5n6/view?usp=sharing)   |
+| [data2vec base](https://github.com/facebookresearch/fairseq/blob/main/examples/data2vec/README.md#speech-2) layer 6   | [Librispeech](http://www.openslr.org/12) train-clean-100 | [data2vec_base_l6](https://drive.google.com/file/d/1d8sf3Ko_fYM9zlaiwxK_4xusLRKV5EMd/view?usp=sharing)   |
+| [data2vec large](https://github.com/facebookresearch/fairseq/blob/main/examples/data2vec/README.md#speech-2) layer 18 | [Librispeech](http://www.openslr.org/12) train-clean-100 | [data2vec_large_l18](https://drive.google.com/file/d/1nuRIHaejT-uVi4cluftbT8o_JZqar5SU/view?usp=sharing) |
+| [Whisper medium](https://github.com/openai/whisper/tree/main#available-models-and-languages) layer 24                 | [Librispeech](http://www.openslr.org/12) train-clean-100 | [whisper_medium_l24](https://drive.google.com/file/d/1V6YJSA2V4iywXrecJAN0oqsa3aHowexZ/view?usp=sharing) |
+| [Whisper large-v2](https://github.com/openai/whisper/tree/main#available-models-and-languages) layer 32                                                                                         | [Librispeech](http://www.openslr.org/12) train-clean-100 | [whisper_large_l32](https://drive.google.com/file/d/1k_X7ZMPg8iOeDrIJe70v6CHfFygzufXC/view?usp=sharing)  |
+
+## Speech Tokenization Using Pre-Trained Models
+
+### Installation
+
+Please first install RepCodec by
+
+```
+git clone https://github.com/mct10/RepCodec.git
+cd RepCodec
+pip install .
+```
+
+We used Python 3.9.18 and PyTorch 1.12.1 to test the usage, but the code should be compatible with other recent Python
+and PyTorch versions.
+
+### Representation Preparation
+
+We adapt the `dump_hubert_feature.py` script
+from [fairseq](https://github.com/facebookresearch/fairseq/tree/main/examples/hubert/simple_kmeans#hubert-feature)
+to support dumping representations from **data2vec**, **HuBERT**, or **Whisper** encoders.
+
+If you use our script (`examples/dump_feature.py`), please also install the following packages:
+
+```
+pip install npy_append_array soundfile 
+```
+
+Additionally, if you want to dump representations from
+
+- **data2vec** or **HuBERT**: please
+  follow [fairseq's instruction](https://github.com/facebookresearch/fairseq#requirements-and-installation) to install
+  the latest fairseq.
+
+- **Whisper**: please follow [Whispers'instruction](https://github.com/openai/whisper/tree/main#setup) to install the
+  latest
+  Whisper.
+
+Then, you can follow the given examples to dump representations:
+
+```
+# Example 1: dump from HuBERT base layer 9 
+# (for data2vec, simply change "model_type" to data2vec and "ckpt_path" to the path of data2vec model)
+
+layer=9
+
+python3 examples/dump_feature.py \
+    --model_type hubert \
+    --tsv_path /path/to/tsv/file \
+    --ckpt_path /path/to/HuBERT/model  \
+    --layer ${layer} \
+    --feat_dir /dir/to/save/representations
+
+
+# Example 2: dump from Whisper medium layer 24
+
+layer=24
+
+python3 examples/dump_feature.py \
+    --model_type whisper \
+    --tsv_path /path/to/tsv/file \
+    --whisper_root /directory/to/save/whisper/model \
+    --whisper_name medium \
+    --layer ${layer} \
+    --feat_dir /dir/to/save/representations
+```
+
+Explanations about the args:
+
+- **model_type:** choose from `data2vec`, `hubert`, and `whisper`.
+
+- **tsv_path:** path of the tsv file.
+  Should have the format of
+
+```
+/dir/to/dataset
+path_of_utterance_1 number_of_frames
+path_of_utterance_2 number_of_frames
+```
+
+You can follow [this script](https://github.com/facebookresearch/fairseq/blob/main/examples/wav2vec/wav2vec_manifest.py)
+to generate the tsv file.
+
+For example, by running
+
+```
+python wav2vec_manifest.py \
+  /dir/to/LibriSpeech/dev-clean \
+  --dest /dir/to/manifest \
+  --ext flac \
+  --valid-percent 0
+```
+
+you can obtain the `dev-clean.tsv` in `/dir/to/manifest` for LibriSpeech. (By default, the output file name
+is `train.tsv`. Remember to rename the file.)
+
+It should be similar to:
+
+```
+/dir/to/LibriSpeech/dev-clean
+2277/149896/2277-149896-0026.flac	78720
+2277/149896/2277-149896-0005.flac	89600
+2277/149896/2277-149896-0033.flac	45520
+```
+
+- **ckpt_path**:
+  must provide for data2vec and HuBERT.
+  You need to download the model
+  from [data2vec website](https://github.com/facebookresearch/fairseq/blob/main/examples/data2vec/README.md#speech-2)
+  or [HuBERT website](https://github.com/facebookresearch/fairseq/tree/main/examples/hubert#pre-trained-and-fine-tuned-asr-models)
+  yourself.
+  `--ckpt_path` is the path of the data2vec/HuBERT model.
+- **whisper_root** and  **whisper_name**:
+  must provide **BOTH** `--whisper_root` and `--whisper_name` for Whisper.
+  If there is no corresponding model in `--whisper_root`, the script will download for you.
+
+- **layer**:
+  which Transformer encoder layer of the model should the representations be extracted from.
+  It is **1-based**.
+  For example, if layer=9, then the outputs from the 9<sup>th</sup> Transformer encoder layer are dumped.
+  Range: [1, number of Transformer encoder layers]
+
+- **feat_dir**: The output representations will be saved to `${feat_dir}/0_1.npy`
+  and `${feat_dir}/0_1.len`.
+
+For other useful functionalities (e.g., sharding), please check the argument list in `examples/dump_feature.py`.
+
+### Command Line Usage
+
+We expect to have `${feat_dir}/0_1.npy` and `${feat_dir}/0_1.len` in the provided
+directory `/dir/to/representaitons`.
+
+Also, the tsv file should be the **same** as the one used in [Representation Preparation](#representation-preparation).
+
+```
+repcodec /dir/to/representaitons \
+    --model /path/to/repcodec/model \
+    --tsv_path /path/to/tsv/file \
+    [--model_config_path /path/to/train/config] \
+    [--use_gpu] \
+    [--out_dir /path/to/output]
+```
+
+If you trained the model yourself following [Training New RepCodec Models](#training-new-repcodec-models), 
+please provide the training config file using `--model_config_path`.
+If you use the model we provide [here](#repcodec-models), then you do not have to provide that.
+
+This command will tokenize the representations and the output discrete tokens will be saved to `${out_dir}/tokens`.
+The tokens are in the same order as the provided tsv file.
+
+An example of the output file:
+
+```
+/dir/to/LibriSpeech/dev-clean
+2277/149896/2277-149896-0026.flac	696 696 198 198 198 498 ...
+2277/149896/2277-149896-0005.flac	696 696 198 198 198 907 ...
+2277/149896/2277-149896-0033.flac	696 696 198 198 198 696 ...
+```
+
+Under `examples/tokens`, we provide some token files as references. They are obtained from LibriSpeech dev-clean subset
+using the 6 types of representations and corresponding [RepCodec Models](#repcodec-models).
+Your results should be very similar to ours.
+
+### Python Usage
+
+```python
+import torch
+import yaml
+
+from repcodec.RepCodec import RepCodec
+
+# for feature types of HubERT base & data2vec base, please use repcodec_dim768.yaml;
+# for feature types of HuBERT large & data2vec large & Whisper medium, please use repcodec_dim1024.yaml;
+# for feature types of Whisper large-v2, please use repcodec_dim1280.yaml
+config = "repcodec/configs/repcodec_dim768.yaml"
+with open(config) as fp:
+    conf = yaml.load(fp, Loader=yaml.FullLoader)
+
+model = RepCodec(**conf)
+model.load_state_dict(torch.load("./hubert_base_l9.pkl", map_location="cpu")["model"]["repcodec"])
+model.quantizer.initial()
+model.eval()
+
+# input shape: (batch size, hidden dim, sequence length)
+random_features = torch.randn(size=(1, 768, 100))
+with torch.no_grad():
+    x = model.encoder(random_features)
+    z = model.projector(x)
+    _, idx = model.quantizer.codebook.forward_index(z.transpose(2, 1))
+    tokens = idx.cpu().data.numpy().tolist()[0]
+```
+
+## Training New RepCodec Models
+
+We use a config file to set up all the training configurations, e.g., data, model architecture,
+optimizer, scheduler. 
+We provide an example [here](./train_configs/ex_dim768_mse.yaml).
+
+Please first install required packages following [Installation](#installation) 
+and prepare the representations following [Representation Preparation](#representation-preparation). 
+
+The input data directory is expected to have the following structure
+```
+/dir/to/representations/
+  train_set_name/
+    0_1.npy
+    0_1.len
+  valid_set_name/
+    0_1.npy
+    0_1.len
+  test_set_name/
+    0_1.npy
+    0_1.len
+```
+
+The names of subsets should be the same as the fields in the config file.
+
+Then, you can run training by
+```
+python train.py \
+  -c /path/to/config/file \
+  --tag $tag \
+  --exp_root exp
+```
+
+`tag` is the name of the output folder.
+All outputs will be saved to `exp_root/tag/`.
+
+## Acknowledge
+
+Our implementation is based on [facebookresearch/AudioDec](https://github.com/facebookresearch/AudioDec).
+We thank them for open-sourcing their code!
+
+## Citation
+
+If you find our work useful, please cite the following article.
+
+```
+@misc{huang2023repcodec,
+      title={RepCodec: A Speech Representation Codec for Speech Tokenization}, 
+      author={Zhichao Huang and Chutong Meng and Tom Ko},
+      year={2023},
+      eprint={2309.00169},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS}
+}
+```

RepCodec/examples/data2vec_audio.py

@@ -0,0 +1,541 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq)
+
+# ref: https://github.com/facebookresearch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
+
+import logging
+import math
+from dataclasses import dataclass, field
+from typing import Optional
+
+from omegaconf import II
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from fairseq.modules import EMAModule, EMAModuleConfig
+from fairseq.data.data_utils import compute_mask_indices
+from fairseq.models import BaseFairseqModel, register_model
+from fairseq.models.wav2vec import (
+    ConvFeatureExtractionModel,
+    Wav2Vec2Config,
+    TransformerEncoder,
+)
+from fairseq.modules import (
+    GradMultiply,
+    LayerNorm,
+)
+from fairseq.utils import index_put
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Data2VecAudioConfig(Wav2Vec2Config):
+
+    loss_beta: float = field(
+        default=0, metadata={"help": "beta for smooth l1 loss. 0 means use l2 loss"}
+    )
+    loss_scale: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "scale the reconstruction loss by this constant. if None then scales by 1/sqrt(dim)"
+        },
+    )
+    average_top_k_layers: int = field(
+        default=8, metadata={"help": "how many layers to average"}
+    )
+
+    layer_norm_target_layer: bool = False
+    instance_norm_target_layer: bool = False
+    instance_norm_targets: bool = False
+    layer_norm_targets: bool = False
+    batch_norm_target_layer: bool = False
+    group_norm_target_layer: bool = False
+
+    ema_decay: float = field(default=0.999, metadata={"help": "initial ema decay rate"})
+    ema_end_decay: float = field(
+        default=0.9999, metadata={"help": "final ema decay rate"}
+    )
+
+    # when to finish annealing ema decay rate
+    ema_anneal_end_step: int = II("optimization.max_update")
+
+    ema_transformer_only: bool = field(
+        default=True,
+        metadata={"help": "whether to momentum update only the transformer"},
+    )
+    ema_layers_only: bool = field(
+        default=True,
+        metadata={"help": "whether to momentum update only the transformer layers"},
+    )
+
+    max_update: int = II("optimization.max_update")
+
+    min_target_var: float = field(
+        default=0.1, metadata={"help": "stop training if target var falls below this"}
+    )
+    min_pred_var: float = field(
+        default=0.01,
+        metadata={"help": "stop training if prediction var falls below this"},
+    )
+
+
+def get_annealed_rate(start, end, curr_step, total_steps):
+    r = end - start
+    pct_remaining = 1 - curr_step / total_steps
+    return end - r * pct_remaining
+
+
+@register_model("data2vec_audio", dataclass=Data2VecAudioConfig)
+class Data2VecAudioModel(BaseFairseqModel):
+    def __init__(self, cfg: Data2VecAudioConfig):
+        super().__init__()
+        self.cfg = cfg
+
+        feature_enc_layers = eval(cfg.conv_feature_layers)
+        self.extractor_embed = feature_enc_layers[-1][0]
+
+        self.ema = None
+        self.embed = cfg.encoder_embed_dim
+
+        self.average_top_k_layers = cfg.average_top_k_layers
+        self.loss_beta = cfg.loss_beta
+        self.loss_scale = cfg.loss_scale
+
+        self.feature_extractor = ConvFeatureExtractionModel(
+            conv_layers=feature_enc_layers,
+            dropout=0.0,
+            mode=cfg.extractor_mode,
+            conv_bias=cfg.conv_bias,
+        )
+
+        self.post_extract_proj = nn.Linear(self.extractor_embed, cfg.encoder_embed_dim)
+
+        self.mask_prob = cfg.mask_prob
+        self.mask_selection = cfg.mask_selection
+        self.mask_other = cfg.mask_other
+        self.mask_length = cfg.mask_length
+        self.no_mask_overlap = cfg.no_mask_overlap
+        self.mask_min_space = cfg.mask_min_space
+
+        self.mask_channel_prob = cfg.mask_channel_prob
+        self.mask_channel_before = cfg.mask_channel_before
+        self.mask_channel_selection = cfg.mask_channel_selection
+        self.mask_channel_other = cfg.mask_channel_other
+        self.mask_channel_length = cfg.mask_channel_length
+        self.no_mask_channel_overlap = cfg.no_mask_channel_overlap
+        self.mask_channel_min_space = cfg.mask_channel_min_space
+
+        self.dropout_input = nn.Dropout(cfg.dropout_input)
+        self.dropout_features = nn.Dropout(cfg.dropout_features)
+
+        self.feature_grad_mult = cfg.feature_grad_mult
+
+        self.mask_emb = nn.Parameter(
+            torch.FloatTensor(cfg.encoder_embed_dim).uniform_()
+        )
+
+        self.encoder = TransformerEncoder(cfg)
+        self.layer_norm = LayerNorm(self.extractor_embed)
+
+        self.final_proj = nn.Linear(self.embed, self.embed)
+
+        self.num_updates = 0
+
+    def make_ema_teacher(self):
+        ema_config = EMAModuleConfig(
+            ema_decay=self.cfg.ema_decay,
+            ema_fp32=True,
+        )
+        skip_keys = set()
+        if self.cfg.ema_layers_only:
+            self.cfg.ema_transformer_only = True
+            for k, _ in self.encoder.pos_conv.named_parameters():
+                skip_keys.add(f"pos_conv.{k}")
+
+        self.ema = EMAModule(
+            self.encoder if self.cfg.ema_transformer_only else self,
+            ema_config,
+            skip_keys=skip_keys,
+        )
+
+    def set_num_updates(self, num_updates):
+        super().set_num_updates(num_updates)
+
+        if self.ema is None and self.final_proj is not None:
+            logger.info(f"making ema teacher")
+            self.make_ema_teacher()
+        elif self.training and self.ema is not None:
+            if self.cfg.ema_decay != self.cfg.ema_end_decay:
+                if num_updates >= self.cfg.ema_anneal_end_step:
+                    decay = self.cfg.ema_end_decay
+                else:
+                    decay = get_annealed_rate(
+                        self.cfg.ema_decay,
+                        self.cfg.ema_end_decay,
+                        num_updates,
+                        self.cfg.ema_anneal_end_step,
+                    )
+                self.ema.set_decay(decay)
+            if self.ema.get_decay() < 1:
+                self.ema.step(self.encoder if self.cfg.ema_transformer_only else self)
+
+        self.num_updates = num_updates
+
+    def state_dict(self, destination=None, prefix="", keep_vars=False):
+        state = super().state_dict(destination, prefix, keep_vars)
+
+        if self.ema is not None:
+            state[prefix + "_ema"] = self.ema.fp32_params
+
+        return state
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        if self.ema is not None:
+            k = prefix + "_ema"
+            assert k in state_dict
+            self.ema.restore(state_dict[k], True)
+            del state_dict[k]
+        return super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
+    @classmethod
+    def build_model(cls, cfg: Data2VecAudioConfig, task=None):
+        """Build a new model instance."""
+
+        return cls(cfg)
+
+    def apply_mask(
+        self,
+        x,
+        padding_mask,
+        mask_indices=None,
+        mask_channel_indices=None,
+    ):
+        B, T, C = x.shape
+
+        if self.mask_channel_prob > 0 and self.mask_channel_before:
+            mask_channel_indices = compute_mask_indices(
+                (B, C),
+                None,
+                self.mask_channel_prob,
+                self.mask_channel_length,
+                self.mask_channel_selection,
+                self.mask_channel_other,
+                no_overlap=self.no_mask_channel_overlap,
+                min_space=self.mask_channel_min_space,
+            )
+            mask_channel_indices = (
+                torch.from_numpy(mask_channel_indices)
+                .to(x.device)
+                .unsqueeze(1)
+                .expand(-1, T, -1)
+            )
+            x[mask_channel_indices] = 0
+
+        if self.mask_prob > 0:
+            if mask_indices is None:
+                mask_indices = compute_mask_indices(
+                    (B, T),
+                    padding_mask,
+                    self.mask_prob,
+                    self.mask_length,
+                    self.mask_selection,
+                    self.mask_other,
+                    min_masks=1,
+                    no_overlap=self.no_mask_overlap,
+                    min_space=self.mask_min_space,
+                    require_same_masks=self.cfg.require_same_masks,
+                    mask_dropout=self.cfg.mask_dropout,
+                )
+                mask_indices = torch.from_numpy(mask_indices).to(x.device)
+            x = index_put(x, mask_indices, self.mask_emb)
+        else:
+            mask_indices = None
+
+        if self.mask_channel_prob > 0 and not self.mask_channel_before:
+            if mask_channel_indices is None:
+                mask_channel_indices = compute_mask_indices(
+                    (B, C),
+                    None,
+                    self.mask_channel_prob,
+                    self.mask_channel_length,
+                    self.mask_channel_selection,
+                    self.mask_channel_other,
+                    no_overlap=self.no_mask_channel_overlap,
+                    min_space=self.mask_channel_min_space,
+                )
+                mask_channel_indices = (
+                    torch.from_numpy(mask_channel_indices)
+                    .to(x.device)
+                    .unsqueeze(1)
+                    .expand(-1, T, -1)
+                )
+            x = index_put(x, mask_channel_indices, 0)
+
+        return x, mask_indices
+
+    def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+        """
+        Computes the output length of the convolutional layers
+        """
+
+        def _conv_out_length(input_length, kernel_size, stride):
+            return torch.floor((input_length - kernel_size) / stride + 1)
+
+        conv_cfg_list = eval(self.cfg.conv_feature_layers)
+
+        for i in range(len(conv_cfg_list)):
+            input_lengths = _conv_out_length(
+                input_lengths, conv_cfg_list[i][1], conv_cfg_list[i][2]
+            )
+
+        return input_lengths.to(torch.long)
+
+    def forward(
+        self,
+        source,
+        padding_mask=None,
+        mask=True,
+        features_only=False,
+        layer=None,
+        mask_indices=None,
+        mask_channel_indices=None,
+        padding_count=None,
+    ):
+        features = source
+
+        if self.feature_grad_mult > 0:
+            features = self.feature_extractor(features)
+            if self.feature_grad_mult != 1.0:
+                features = GradMultiply.apply(features, self.feature_grad_mult)
+        else:
+            with torch.no_grad():
+                features = self.feature_extractor(features)
+
+        features = features.transpose(1, 2)
+
+        features = self.layer_norm(features)
+
+        orig_padding_mask = padding_mask
+
+        if padding_mask is not None and padding_mask.any():
+            input_lengths = (1 - padding_mask.long()).sum(-1)
+            # apply conv formula to get real output_lengths
+            output_lengths = self._get_feat_extract_output_lengths(input_lengths)
+
+            padding_mask = torch.zeros(
+                features.shape[:2], dtype=features.dtype, device=features.device
+            )
+
+            # these two operations makes sure that all values
+            # before the output lengths indices are attended to
+            padding_mask[
+                (
+                    torch.arange(padding_mask.shape[0], device=padding_mask.device),
+                    output_lengths - 1,
+                )
+            ] = 1
+            padding_mask = (1 - padding_mask.flip([-1]).cumsum(-1).flip([-1])).bool()
+        else:
+            padding_mask = None
+
+        if self.post_extract_proj is not None:
+            features = self.post_extract_proj(features)
+
+        pre_encoder_features = None
+        if self.cfg.ema_transformer_only:
+            pre_encoder_features = features.clone()
+
+        features = self.dropout_input(features)
+
+        if mask:
+            x, mask_indices = self.apply_mask(
+                features,
+                padding_mask,
+                mask_indices=mask_indices,
+                mask_channel_indices=mask_channel_indices,
+            )
+        else:
+            x = features
+            mask_indices = None
+
+        x, layer_results = self.encoder(
+            x,
+            padding_mask=padding_mask,
+            layer=layer,
+        )
+
+        if features_only:
+            return {
+                "x": x,
+                "padding_mask": padding_mask,
+                "layer_results": layer_results,
+            }
+
+        result = {
+            "losses": {},
+        }
+
+        with torch.no_grad():
+            self.ema.model.eval()
+
+            if self.cfg.ema_transformer_only:
+                y, layer_results = self.ema.model.extract_features(
+                    pre_encoder_features,
+                    padding_mask=padding_mask,
+                    min_layer=self.cfg.encoder_layers - self.average_top_k_layers,
+                )
+                y = {
+                    "x": y,
+                    "padding_mask": padding_mask,
+                    "layer_results": layer_results,
+                }
+            else:
+                y = self.ema.model.extract_features(
+                    source=source,
+                    padding_mask=orig_padding_mask,
+                    mask=False,
+                )
+
+            target_layer_results = [l[2] for l in y["layer_results"]]
+
+            permuted = False
+            if self.cfg.instance_norm_target_layer or self.cfg.batch_norm_target_layer:
+                target_layer_results = [
+                    tl.permute(1, 2, 0) for tl in target_layer_results  # TBC -> BCT
+                ]
+                permuted = True
+
+            if self.cfg.batch_norm_target_layer:
+                target_layer_results = [
+                    F.batch_norm(
+                        tl.float(), running_mean=None, running_var=None, training=True
+                    )
+                    for tl in target_layer_results
+                ]
+
+            if self.cfg.instance_norm_target_layer:
+                target_layer_results = [
+                    F.instance_norm(tl.float()) for tl in target_layer_results
+                ]
+
+            if permuted:
+                target_layer_results = [
+                    tl.transpose(1, 2) for tl in target_layer_results  # BCT -> BTC
+                ]
+
+            if self.cfg.group_norm_target_layer:
+                target_layer_results = [
+                    F.layer_norm(tl.float(), tl.shape[-2:])
+                    for tl in target_layer_results
+                ]
+
+            if self.cfg.layer_norm_target_layer:
+                target_layer_results = [
+                    F.layer_norm(tl.float(), tl.shape[-1:])
+                    for tl in target_layer_results
+                ]
+
+            y = sum(target_layer_results) / len(target_layer_results)
+
+            if self.cfg.layer_norm_targets:
+                y = F.layer_norm(y.float(), y.shape[-1:])
+
+            if self.cfg.instance_norm_targets:
+                y = F.instance_norm(y.float().transpose(1, 2)).transpose(1, 2)
+
+            if not permuted:
+                y = y.transpose(0, 1)
+
+            y = y[mask_indices]
+
+        x = x[mask_indices]
+        x = self.final_proj(x)
+
+        sz = x.size(-1)
+
+        if self.loss_beta == 0:
+            loss = F.mse_loss(x.float(), y.float(), reduction="none").sum(dim=-1)
+        else:
+            loss = F.smooth_l1_loss(
+                x.float(), y.float(), reduction="none", beta=self.loss_beta
+            ).sum(dim=-1)
+
+        if self.loss_scale is not None:
+            scale = self.loss_scale
+        else:
+            scale = 1 / math.sqrt(sz)
+
+        result["losses"]["regression"] = loss.sum() * scale
+
+        if "sample_size" not in result:
+            result["sample_size"] = loss.numel()
+
+        with torch.no_grad():
+            result["target_var"] = self.compute_var(y)
+            result["pred_var"] = self.compute_var(x.float())
+
+        if self.num_updates > 5000 and result["target_var"] < self.cfg.min_target_var:
+            logger.error(
+                f"target var is {result['target_var'].item()} < {self.cfg.min_target_var}, exiting"
+            )
+            raise Exception(
+                f"target var is {result['target_var'].item()} < {self.cfg.min_target_var}, exiting"
+            )
+        if self.num_updates > 5000 and result["pred_var"] < self.cfg.min_pred_var:
+            logger.error(
+                f"pred var is {result['pred_var'].item()} < {self.cfg.min_pred_var}, exiting"
+            )
+            raise Exception(
+                f"pred var is {result['pred_var'].item()} < {self.cfg.min_pred_var}, exiting"
+            )
+
+        if self.ema is not None:
+            result["ema_decay"] = self.ema.get_decay() * 1000
+
+        return result
+
+    @staticmethod
+    def compute_var(y):
+        y = y.view(-1, y.size(-1))
+        if dist.is_initialized():
+            zc = torch.tensor(y.size(0)).cuda()
+            zs = y.sum(dim=0)
+            zss = (y ** 2).sum(dim=0)
+
+            dist.all_reduce(zc)
+            dist.all_reduce(zs)
+            dist.all_reduce(zss)
+
+            var = zss / (zc - 1) - (zs ** 2) / (zc * (zc - 1))
+            return torch.sqrt(var + 1e-6).mean()
+        else:
+            return torch.sqrt(y.var(dim=0) + 1e-6).mean()
+
+    def extract_features(
+        self, source, padding_mask, mask=False, layer=None
+    ):
+        res = self.forward(
+            source,
+            padding_mask,
+            mask=mask,
+            features_only=True,
+            layer=layer,
+        )
+        return res
+
+    def remove_pretraining_modules(self, last_layer=None):
+        self.final_proj = None
+        self.ema = None
+        if last_layer is not None:
+            self.encoder.layers = nn.ModuleList(
+                l for i, l in enumerate(self.encoder.layers) if i <= last_layer
+            )

RepCodec/examples/data2vec_feature_reader.py

@@ -0,0 +1,87 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq)
+
+import logging
+
+import torch
+import torch.nn.functional as F
+from fairseq import tasks
+from fairseq.checkpoint_utils import load_checkpoint_to_cpu
+from fairseq.data.audio.audio_utils import get_features_or_waveform
+from omegaconf import OmegaConf
+
+from data2vec_audio import Data2VecAudioModel
+
+logger = logging.getLogger("dump_feature")
+
+
+class Data2vecFeatureReader(object):
+    def __init__(self, ckpt_path: str, layer: int, device: str, max_chunk=1600000):
+        state = load_checkpoint_to_cpu(ckpt_path)
+        cfg = state["cfg"]
+        # load task
+        task = tasks.setup_task(cfg.task, from_checkpoint=True)
+        task.load_state_dict(state["task_state"])
+        # load model config
+        if "layer_type" not in cfg.model:
+            # fix a missing key
+            model_config = {k: v for k, v in cfg.model.items()}
+            model_config["layer_type"] = "transformer"
+            model_config = OmegaConf.create(model_config)
+        else:
+            model_config = cfg.model
+
+        # fix param name in the state
+        state["model"]["final_proj.weight"] = state["model"].pop("final_proj.0.weight")
+        state["model"]["final_proj.bias"] = state["model"].pop("final_proj.0.bias")
+        del state["model"]["_ema"]
+
+        # load model
+        model = Data2VecAudioModel.build_model(model_config)
+        model.load_state_dict(
+            state["model"], strict=True, model_cfg=model_config
+        )
+
+        self.device = device
+        logger.info(f"device = {self.device}")
+
+        self.model = model.eval().to(self.device)
+        self.task = task
+        self.layer = layer - 1  # make it 1-based
+        self.max_chunk = max_chunk
+        logger.info(f"TASK CONFIG:\n{self.task.cfg}")
+        logger.info(f" max_chunk = {self.max_chunk}")
+
+    def read_audio(self, path, ref_len=None):
+        wav = get_features_or_waveform(path, need_waveform=True, use_sample_rate=self.task.cfg.sample_rate)
+        if wav.ndim == 2:
+            wav = wav.mean(-1)
+        assert wav.ndim == 1, wav.ndim
+        if ref_len is not None and abs(ref_len - len(wav)) > 160:
+            logger.warning(f"ref {ref_len} != read {len(wav)} ({path})")
+        return wav
+
+    def get_feats(self, path, ref_len=None):
+        x = self.read_audio(path, ref_len=ref_len)
+        with torch.no_grad():
+            x = torch.from_numpy(x).float().to(self.device)
+            if self.task.cfg.normalize:
+                x = F.layer_norm(x, x.shape)
+            x = x.view(1, -1)
+
+            feat = []
+            for start in range(0, x.size(1), self.max_chunk):
+                x_chunk = x[:, start: start + self.max_chunk]
+                res = self.model.extract_features(
+                    source=x_chunk,
+                    padding_mask=None,
+                    mask=False,
+                    layer=self.layer,
+                )
+                feat_chunk = res["x"]
+                feat.append(feat_chunk)
+        return torch.cat(feat, 1).squeeze(0)

RepCodec/examples/dump_feature.py

@@ -0,0 +1,142 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq)
+
+import logging
+import os
+import sys
+
+from feature_utils import get_path_iterator, dump_feature
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    stream=sys.stdout,
+)
+logger = logging.getLogger("dump_feature")
+
+
+def main(
+        model_type: str,
+        tsv_path: str,
+        ckpt_path: str,
+        whisper_root: str,
+        whisper_name: str,
+        layer: int,
+        nshard: int,
+        rank: int,
+        feat_dir: str,
+        max_chunk: int,
+        use_cpu: bool = False
+):
+    device = "cpu" if use_cpu else "cuda"
+
+    # some checks
+    if model_type in ["hubert", "data2vec"]:
+        assert ckpt_path and os.path.exists(ckpt_path)
+    elif model_type in ["whisper"]:
+        assert whisper_name and whisper_root
+    else:
+        raise ValueError(f"Unsupported model type {model_type}")
+
+    reader = None
+    if model_type == "hubert":
+        from hubert_feature_reader import HubertFeatureReader
+        reader = HubertFeatureReader(ckpt_path, layer, device=device, max_chunk=max_chunk)
+    elif model_type == "data2vec":
+        from data2vec_feature_reader import Data2vecFeatureReader
+        reader = Data2vecFeatureReader(ckpt_path, layer, device=device, max_chunk=max_chunk)
+    elif model_type == "whisper":
+        from whisper_feature_reader import WhisperFeatureReader
+        reader = WhisperFeatureReader(whisper_root, whisper_name, layer, device=device)
+
+    assert reader is not None
+
+    generator, num = get_path_iterator(tsv_path, nshard, rank)
+    dump_feature(reader, generator, num, nshard, rank, feat_dir)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type",
+        required=True,
+        type=str,
+        choices=["data2vec", "hubert", "whisper"],
+        help="the type of the speech encoder."
+    )
+    parser.add_argument(
+        "--tsv_path",
+        required=True,
+        type=str,
+        help="the path to the tsv file."
+    )
+    parser.add_argument(
+        "--ckpt_path",
+        required=False,
+        type=str,
+        default=None,
+        help="path to the speech model. must provide for HuBERT and data2vec"
+    )
+    parser.add_argument(
+        "--whisper_root",
+        required=False,
+        type=str,
+        default=None,
+        help="root dir to download/store whisper model. must provide for whisper model."
+    )
+    parser.add_argument(
+        "--whisper_name",
+        required=False,
+        type=str,
+        default=None,
+        help="name of whisper model. e.g., large-v2. must provide for whisper model."
+    )
+    parser.add_argument(
+        "--layer",
+        required=True,
+        type=int,
+        help="which layer of the model. this is 1-based."
+    )
+    parser.add_argument(
+        "--feat_dir",
+        required=True,
+        type=str,
+        help="the output dir to save the representations."
+    )
+    parser.add_argument(
+        "--nshard",
+        required=False,
+        type=int,
+        default=1,
+        help="total number of shards."
+    )
+    parser.add_argument(
+        "--rank",
+        required=False,
+        type=int,
+        default=0,
+        help="shard id of this process."
+    )
+    parser.add_argument(
+        "--max_chunk",
+        type=int,
+        default=1600000,
+        help="max number of frames of each batch."
+    )
+    parser.add_argument(
+        "--use_cpu",
+        default=False,
+        action="store_true",
+        help="whether use cpu instead of gpu."
+    )
+    args = parser.parse_args()
+    logger.info(args)
+
+    main(**vars(args))

RepCodec/examples/feature_utils.py

@@ -0,0 +1,70 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq)
+
+# ref: https://github.com/facebookresearch/fairseq/blob/main/examples/hubert/simple_kmeans/feature_utils.py
+
+import logging
+import os
+import sys
+
+import tqdm
+from npy_append_array import NpyAppendArray
+
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    stream=sys.stdout,
+)
+logger = logging.getLogger("feature_utils")
+
+
+def get_shard_range(tot, nshard, rank):
+    assert rank < nshard and rank >= 0, f"invaid rank/nshard {rank}/{nshard}"
+    start = round(tot / nshard * rank)
+    end = round(tot / nshard * (rank + 1))
+    assert start < end, f"start={start}, end={end}"
+    logger.info(
+        f"rank {rank} of {nshard}, process {end-start} "
+        f"({start}-{end}) out of {tot}"
+    )
+    return start, end
+
+
+def get_path_iterator(tsv, nshard, rank):
+    with open(tsv, "r") as f:
+        root = f.readline().rstrip()
+        lines = [line.rstrip() for line in f]
+        start, end = get_shard_range(len(lines), nshard, rank)
+        lines = lines[start:end]
+        def iterate():
+            for line in lines:
+                subpath, nsample = line.split("\t")
+                yield f"{root}/{subpath}", int(nsample)
+    return iterate, len(lines)
+
+
+def dump_feature(reader, generator, num, nshard, rank, feat_dir):
+    iterator = generator()
+
+    feat_path = f"{feat_dir}/{rank}_{nshard}.npy"
+    leng_path = f"{feat_dir}/{rank}_{nshard}.len"
+
+    os.makedirs(feat_dir, exist_ok=True)
+    if os.path.exists(feat_path):
+        os.remove(feat_path)
+
+    feat_f = NpyAppendArray(feat_path)
+    with open(leng_path, "w") as leng_f:
+        for path, nsample in tqdm.tqdm(iterator, total=num):
+            feat = reader.get_feats(path, nsample)
+            feat_f.append(feat.cpu().numpy())
+            leng_f.write(f"{len(feat)}\n")
+    logger.info("finished successfully")
+
+

RepCodec/examples/hubert_feature_reader.py

@@ -0,0 +1,64 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq)
+
+import logging
+
+import fairseq
+import torch
+import torch.nn.functional as F
+
+from fairseq.data.audio.audio_utils import get_features_or_waveform
+
+logger = logging.getLogger("dump_feature")
+
+
+class HubertFeatureReader(object):
+    def __init__(self, ckpt_path: str, layer: int, device: str, max_chunk=1600000):
+        (
+            model,
+            cfg,
+            task,
+        ) = fairseq.checkpoint_utils.load_model_ensemble_and_task([ckpt_path])
+
+        self.device = device
+        logger.info(f"device = {self.device}")
+
+        self.model = model[0].eval().to(self.device)
+        self.task = task
+        self.layer = layer
+        self.max_chunk = max_chunk
+        logger.info(f"TASK CONFIG:\n{self.task.cfg}")
+        logger.info(f" max_chunk = {self.max_chunk}")
+
+    def read_audio(self, path, ref_len=None):
+        wav = get_features_or_waveform(path, need_waveform=True, use_sample_rate=self.task.cfg.sample_rate)
+        if wav.ndim == 2:
+            wav = wav.mean(-1)
+        assert wav.ndim == 1, wav.ndim
+        if ref_len is not None and abs(ref_len - len(wav)) > 160:
+            logger.warning(f"ref {ref_len} != read {len(wav)} ({path})")
+        return wav
+
+    def get_feats(self, path, ref_len=None):
+        x = self.read_audio(path, ref_len=ref_len)
+        with torch.no_grad():
+            x = torch.from_numpy(x).float().to(self.device)
+            if self.task.cfg.normalize:
+                x = F.layer_norm(x, x.shape)
+            x = x.view(1, -1)
+
+            feat = []
+            for start in range(0, x.size(1), self.max_chunk):
+                x_chunk = x[:, start: start + self.max_chunk]
+                feat_chunk, _ = self.model.extract_features(
+                    source=x_chunk,
+                    padding_mask=None,
+                    mask=False,
+                    output_layer=self.layer,
+                )
+                feat.append(feat_chunk)
+        return torch.cat(feat, 1).squeeze(0)

RepCodec/examples/whisper_feature_reader.py

@@ -0,0 +1,110 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq) and
+# Whisper (https://github.com/openai/whisper/)
+
+import io
+import logging
+import os
+from typing import Optional, Union
+
+import soundfile as sf
+import torch
+from whisper import _MODELS, _download, _ALIGNMENT_HEADS, available_models
+from whisper.audio import log_mel_spectrogram
+from whisper.model import ModelDimensions
+
+from whisper_model import Whisper_
+
+logger = logging.getLogger("dump_feature")
+
+
+def load_model(
+        name: str,
+        device: Optional[Union[str, torch.device]] = None,
+        download_root: str = None,
+        in_memory: bool = False,
+) -> Whisper_:
+    """
+    Reference: https://github.com/openai/whisper/blob/main/whisper/__init__.py#L97
+    But we will load a `Whisper_` model for feature extraction.
+
+    Parameters
+    ----------
+    name : str
+        one of the official model names listed by `whisper.available_models()`, or
+        path to a model checkpoint containing the model dimensions and the model state_dict.
+    device : Union[str, torch.device]
+        the PyTorch device to put the model into
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/whisper"
+    in_memory: bool
+        whether to preload the model weights into host memory
+
+    Returns
+    -------
+    model : Whisper
+        The Whisper ASR model instance
+    """
+
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if download_root is None:
+        default = os.path.join(os.path.expanduser("~"), ".cache")
+        download_root = os.path.join(os.getenv("XDG_CACHE_HOME", default), "whisper")
+
+    if name in _MODELS:
+        checkpoint_file = _download(_MODELS[name], download_root, in_memory)
+        alignment_heads = _ALIGNMENT_HEADS[name]
+    elif os.path.isfile(name):
+        checkpoint_file = open(name, "rb").read() if in_memory else name
+        alignment_heads = None
+    else:
+        raise RuntimeError(
+            f"Model {name} not found; available models = {available_models()}"
+        )
+
+    with (
+            io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")
+    ) as fp:
+        checkpoint = torch.load(fp, map_location=device)
+    del checkpoint_file
+
+    dims = ModelDimensions(**checkpoint["dims"])
+    model = Whisper_(dims)
+    model.load_state_dict(checkpoint["model_state_dict"])
+
+    if alignment_heads is not None:
+        model.set_alignment_heads(alignment_heads)
+
+    return model.to(device)
+
+
+class WhisperFeatureReader(object):
+    def __init__(self, root, ckpt, layer, device):
+        self.device = device
+        logger.info(f"device = {self.device}")
+
+        self.model: Whisper_ = load_model(name=ckpt, device=self.device, download_root=root).eval()
+        self.model.decoder = None  # to save some memory by deleting the decoder
+        self.layer = layer  # one-based
+
+    def read_audio(self, path, ref_len=None):
+        wav, sample_rate = sf.read(path)
+        assert sample_rate == 16000, sample_rate
+        if ref_len is not None and abs(ref_len - len(wav)) > 160:
+            logger.warning(f"ref {ref_len} != read {len(wav)} ({path})")
+        return wav
+
+    def get_feats(self, path, ref_len=None):
+        wav = self.read_audio(path, ref_len)
+        audio_length = len(wav)
+        with torch.no_grad():
+            mel = log_mel_spectrogram(torch.from_numpy(wav).float().to(self.device))
+            hidden = self.model.extract_features(mel.unsqueeze(0), target_layer=self.layer)
+            feature_length = audio_length // 320
+            hidden = hidden[0, :feature_length]
+        return hidden.contiguous()

RepCodec/examples/whisper_model.py

@@ -0,0 +1,58 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+# 
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on fairseq (https://github.com/facebookresearch/fairseq) and 
+# Whisper (https://github.com/openai/whisper/)
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from whisper.model import AudioEncoder, sinusoids, Whisper, ModelDimensions
+
+
+class AudioEncoder_(AudioEncoder):
+    def __init__(self, *args, **kwargs):
+        super(AudioEncoder_, self).__init__(*args, **kwargs)
+
+    def extract_feature(self, x: Tensor, target_layer: Optional[int] = None):
+        """
+        x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
+            the mel spectrogram of the audio
+        """
+        x = F.gelu(self.conv1(x))
+        x = F.gelu(self.conv2(x))
+        x = x.permute(0, 2, 1)
+
+        length_x = x.shape[1]
+        if length_x > self.positional_embedding.shape[0]:
+            self.register_buffer("positional_embedding", sinusoids(length_x, self.positional_embedding.shape[1]))
+            self.positional_embedding = self.positional_embedding.to(x.device)
+        x = (x + self.positional_embedding[:length_x, :]).to(x.dtype)
+
+        if target_layer is None:
+            target_layer = len(self.blocks)
+
+        for block in self.blocks[:target_layer]:
+            x = block(x)
+
+        return x
+
+
+class Whisper_(Whisper):
+    def __init__(self, dims: ModelDimensions):
+        super(Whisper_, self).__init__(dims)
+        # replace audio encoder with our audio encoder
+        self.encoder = AudioEncoder_(
+            self.dims.n_mels,
+            self.dims.n_audio_ctx,
+            self.dims.n_audio_state,
+            self.dims.n_audio_head,
+            self.dims.n_audio_layer,
+        )
+
+    def extract_features(self, mel: torch.Tensor, target_layer: Optional[int] = None):
+        return self.encoder.extract_feature(mel, target_layer)

RepCodec/repcodec/RepCodec.py

@@ -0,0 +1,84 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+from repcodec.modules.decoder import Decoder
+from repcodec.modules.encoder import Encoder
+from repcodec.modules.projector import Projector
+from repcodec.modules.quantizer import Quantizer
+
+
+class RepCodec(nn.Module):
+    def __init__(
+            self,
+            input_channels=768,
+            output_channels=768,
+            encode_channels=768,
+            decode_channels=768,
+            code_dim=768,
+            codebook_num=1,
+            codebook_size=1024,
+            bias=True,
+            enc_ratios=(1, 1),
+            dec_ratios=(1, 1),
+            enc_strides=(1, 1),
+            dec_strides=(1, 1),
+            enc_kernel_size=3,
+            dec_kernel_size=3,
+            enc_block_dilations=(1, 1),
+            enc_block_kernel_size=3,
+            dec_block_dilations=(1, 1),
+            dec_block_kernel_size=3
+    ):
+        super().__init__()
+
+        self.input_channels = input_channels
+
+        self.encoder = Encoder(
+            input_channels=input_channels,
+            encode_channels=encode_channels,
+            channel_ratios=enc_ratios,
+            strides=enc_strides,
+            kernel_size=enc_kernel_size,
+            bias=bias,
+            block_dilations=enc_block_dilations,
+            unit_kernel_size=enc_block_kernel_size
+        )
+
+        self.decoder = Decoder(
+            code_dim=code_dim,
+            output_channels=output_channels,
+            decode_channels=decode_channels,
+            channel_ratios=dec_ratios,
+            strides=dec_strides,
+            kernel_size=dec_kernel_size,
+            bias=bias,
+            block_dilations=dec_block_dilations,
+            unit_kernel_size=dec_block_kernel_size
+        )
+
+        self.projector = Projector(
+            input_channels=self.encoder.out_channels,
+            code_dim=code_dim,
+            kernel_size=3,
+            stride=1,
+            bias=False
+        )
+
+        self.quantizer = Quantizer(
+            code_dim=code_dim,
+            codebook_num=codebook_num,
+            codebook_size=codebook_size
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        z = self.projector(x)
+        zq, vqloss, perplexity = self.quantizer(z)
+        y = self.decoder(zq)
+        return y, zq, z, vqloss, perplexity

RepCodec/repcodec/configs/repcodec_dim1024.yaml

@@ -0,0 +1,18 @@
+input_channels: 1024
+output_channels: 1024
+encode_channels: 1024
+decode_channels: 1024
+code_dim: 1024
+codebook_num: 1
+codebook_size: 1024
+bias: true
+enc_ratios: [ 1, 1 ]
+dec_ratios: [ 1, 1 ]
+enc_strides: [ 1, 1 ]  # no downsampling
+dec_strides: [ 1, 1 ]
+enc_kernel_size: 3
+dec_kernel_size: 3
+enc_block_dilations: [ 1, 1 ]
+enc_block_kernel_size: 3
+dec_block_dilations: [ 1, 1 ]
+dec_block_kernel_size: 3

RepCodec/repcodec/configs/repcodec_dim1280.yaml

@@ -0,0 +1,18 @@
+input_channels: 1280
+output_channels: 1280
+encode_channels: 1280
+decode_channels: 1280
+code_dim: 1280
+codebook_num: 1
+codebook_size: 1024
+bias: true
+enc_ratios: [ 1, 1 ]
+dec_ratios: [ 1, 1 ]
+enc_strides: [ 1, 1 ]  # no downsampling
+dec_strides: [ 1, 1 ]
+enc_kernel_size: 3
+dec_kernel_size: 3
+enc_block_dilations: [ 1, 1 ]
+enc_block_kernel_size: 3
+dec_block_dilations: [ 1, 1 ]
+dec_block_kernel_size: 3

RepCodec/repcodec/configs/repcodec_dim768.yaml

@@ -0,0 +1,18 @@
+input_channels: 768
+output_channels: 768
+encode_channels: 768
+decode_channels: 768
+code_dim: 768
+codebook_num: 1
+codebook_size: 1024
+bias: true
+enc_ratios: [ 1, 1 ]
+dec_ratios: [ 1, 1 ]
+enc_strides: [ 1, 1 ]  # no downsampling
+dec_strides: [ 1, 1 ]
+enc_kernel_size: 3
+dec_kernel_size: 3
+enc_block_dilations: [ 1, 1 ]
+enc_block_kernel_size: 3
+dec_block_dilations: [ 1, 1 ]
+dec_block_kernel_size: 3

RepCodec/repcodec/layers/conv_layer.py

@@ -0,0 +1,95 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+
+class Conv1d1x1(nn.Conv1d):
+    """1x1 Conv1d."""
+
+    def __init__(self, in_channels, out_channels, bias=True):
+        super(Conv1d1x1, self).__init__(in_channels, out_channels, kernel_size=1, bias=bias)
+
+
+class Conv1d(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int,
+            stride: int = 1,
+            padding: int = -1,
+            dilation: int = 1,
+            groups: int = 1,
+            bias: bool = True
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        if padding < 0:
+            padding = (kernel_size - 1) // 2 * dilation
+        self.dilation = dilation
+        self.conv = nn.Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): Float tensor variable with the shape  (B, C, T).
+        Returns:
+            Tensor: Float tensor variable with the shape (B, C, T).
+        """
+        x = self.conv(x)
+        return x
+
+
+class ConvTranspose1d(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size: int,
+            stride: int,
+            padding=-1,
+            output_padding=-1,
+            groups=1,
+            bias=True,
+    ):
+        super().__init__()
+        if padding < 0:
+            padding = (stride + 1) // 2
+        if output_padding < 0:
+            output_padding = 1 if stride % 2 else 0
+        self.deconv = nn.ConvTranspose1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            groups=groups,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        """
+        Args:
+            x (Tensor): Float tensor variable with the shape  (B, C, T).
+        Returns:
+            Tensor: Float tensor variable with the shape (B, C', T').
+        """
+        x = self.deconv(x)
+        return x

RepCodec/repcodec/layers/vq_module.py

@@ -0,0 +1,155 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VectorQuantize(nn.Module):
+    """Vector quantization w/ exponential moving averages (EMA)"""
+
+    def __init__(
+            self,
+            dim: int,
+            codebook_size: int,
+            decay=0.8,
+            commitment=1.,
+            eps=1e-5,
+            n_embed=None,
+    ):
+        super().__init__()
+        n_embed = self.default(n_embed, codebook_size)
+
+        self.dim = dim
+        self.n_embed = n_embed
+        self.decay = decay
+        self.eps = eps
+        self.commitment = commitment
+
+        embed = torch.randn(dim, n_embed)
+        self.register_buffer('embed', embed)
+        self.register_buffer('cluster_size', torch.zeros(n_embed))
+        self.register_buffer('embed_avg', embed.clone())
+
+    @property
+    def codebook(self):
+        return self.embed.transpose(0, 1)
+
+    def exists(self, val):
+        return val is not None
+
+    def default(self, val, d):
+        return val if self.exists(val) else d
+
+    def ema_inplace(self, moving_avg, new, decay):
+        moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+    def laplace_smoothing(self, x, n_categories, eps=1e-5):
+        return (x + eps) / (x.sum() + n_categories * eps)
+
+    def forward(self, input):
+        dtype = input.dtype
+        flatten = input.reshape(-1, self.dim)
+        dist = (
+                flatten.pow(2).sum(1, keepdim=True)
+                - 2 * flatten @ self.embed
+                + self.embed.pow(2).sum(0, keepdim=True)
+        )
+        _, embed_ind = (-dist).max(1)
+        embed_onehot = F.one_hot(embed_ind, self.n_embed).type(dtype)
+        embed_ind = embed_ind.view(*input.shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed.transpose(0, 1))
+
+        if self.training:
+            self.ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = flatten.transpose(0, 1) @ embed_onehot
+            self.ema_inplace(self.embed_avg, embed_sum, self.decay)
+            cluster_size = self.laplace_smoothing(self.cluster_size, self.n_embed, self.eps) * self.cluster_size.sum()
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(0)
+            self.embed.data.copy_(embed_normalized)
+
+        loss = F.mse_loss(quantize.detach(), input) * self.commitment
+        quantize = input + (quantize - input).detach()
+
+        avg_probs = torch.mean(embed_onehot, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
+
+        return quantize, loss, perplexity
+
+    def forward_index(self, input):
+        dtype = input.dtype
+        flatten = input.reshape(-1, self.dim)
+        dist = (
+                flatten.pow(2).sum(1, keepdim=True)
+                - 2 * flatten @ self.embed
+                + self.embed.pow(2).sum(0, keepdim=True)
+        )
+        _, embed_ind = (-dist).max(1)
+        embed_onehot = F.one_hot(embed_ind, self.n_embed).type(dtype)
+        embed_ind = embed_ind.view(*input.shape[:-1])
+        quantize = F.embedding(embed_ind, self.embed.transpose(0, 1))
+        quantize = input + (quantize - input).detach()
+
+        return quantize, embed_ind
+
+
+class ResidualVQ(nn.Module):
+    """ Residual VQ following algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf """
+
+    def __init__(
+            self,
+            *,
+            num_quantizers,
+            **kwargs
+    ):
+        super().__init__()
+        self.layers = nn.ModuleList([VectorQuantize(**kwargs) for _ in range(num_quantizers)])
+
+    def forward(self, x):
+        quantized_out = 0.
+        residual = x
+        all_losses = []
+        all_perplexities = []
+        for layer in self.layers:
+            quantized, loss, perplexity = layer(residual)
+            # Issue: https://github.com/lucidrains/vector-quantize-pytorch/issues/33
+            # We found considering only the 1st layer VQ's graident results in better performance
+            # residual = residual - quantized.detach() # considering all layers' graidents
+            residual = residual - quantized  # considering only the first layer's graident
+            quantized_out = quantized_out + quantized
+            all_losses.append(loss)
+            all_perplexities.append(perplexity)
+        all_losses, all_perplexities = map(torch.stack, (all_losses, all_perplexities))
+        return quantized_out, all_losses, all_perplexities
+
+    def forward_index(self, x, flatten_idx=False):
+        quantized_out = 0.
+        residual = x
+        all_indices = []
+        for i, layer in enumerate(self.layers):
+            quantized, indices = layer.forward_index(residual)
+            # residual = residual - quantized.detach()
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+            if flatten_idx:
+                indices += (self.codebook_size * i)
+            all_indices.append(indices)
+        all_indices = torch.stack(all_indices)
+        return quantized_out, all_indices.squeeze(1)
+
+    def initial(self):
+        self.codebook = []
+        for layer in self.layers:
+            self.codebook.append(layer.codebook)
+        self.codebook_size = self.codebook[0].size(0)
+        self.codebook = torch.stack(self.codebook)
+        self.codebook = self.codebook.reshape(-1, self.codebook.size(-1))
+
+    def lookup(self, indices):
+        quantized_out = F.embedding(indices, self.codebook)  # Num x T x C
+        return torch.sum(quantized_out, dim=0, keepdim=True)

RepCodec/repcodec/modules/decoder.py

@@ -0,0 +1,109 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch
+import torch.nn as nn
+
+from RepCodec.repcodec.layers.conv_layer import Conv1d, ConvTranspose1d
+from RepCodec.repcodec.modules.residual_unit import ResidualUnit
+
+
+class DecoderBlock(nn.Module):
+    """ Decoder block (no up-sampling) """
+
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            stride: int,
+            dilations=(1, 1),
+            unit_kernel_size=3,
+            bias=True
+    ):
+        super().__init__()
+
+        if stride == 1:
+            self.conv = Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=3,  # fix kernel=3 when stride=1 for unchanged shape
+                stride=stride,
+                bias=bias,
+            )
+        else:
+            self.conv = ConvTranspose1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=(2 * stride),
+                stride=stride,
+                bias=bias,
+            )
+
+        self.res_units = torch.nn.ModuleList()
+        for idx, dilation in enumerate(dilations):
+            self.res_units += [
+                ResidualUnit(out_channels, out_channels,
+                             kernel_size=unit_kernel_size,
+                             dilation=dilation)
+            ]
+        self.num_res = len(self.res_units)
+
+    def forward(self, x):
+        x = self.conv(x)
+        for idx in range(self.num_res):
+            x = self.res_units[idx](x)
+        return x
+
+
+class Decoder(nn.Module):
+    def __init__(
+            self,
+            code_dim: int,
+            output_channels: int,
+            decode_channels: int,
+            channel_ratios=(1, 1),
+            strides=(1, 1),
+            kernel_size=3,
+            bias=True,
+            block_dilations=(1, 1),
+            unit_kernel_size=3,
+    ):
+        super().__init__()
+        assert len(channel_ratios) == len(strides)
+
+        self.conv1 = Conv1d(
+            in_channels=code_dim,
+            out_channels=int(decode_channels * channel_ratios[0]),
+            kernel_size=kernel_size,
+            stride=1,
+            bias=False
+        )
+
+        self.conv_blocks = torch.nn.ModuleList()
+        for idx, stride in enumerate(strides):
+            in_channels = int(decode_channels * channel_ratios[idx])
+            if idx < (len(channel_ratios) - 1):
+                out_channels = int(decode_channels * channel_ratios[idx + 1])
+            else:
+                out_channels = decode_channels
+            self.conv_blocks += [
+                DecoderBlock(
+                    in_channels, out_channels, stride,
+                    dilations=block_dilations, unit_kernel_size=unit_kernel_size,
+                    bias=bias
+                )
+            ]
+        self.num_blocks = len(self.conv_blocks)
+
+        self.conv2 = Conv1d(out_channels, output_channels, kernel_size, 1, bias=False)
+
+    def forward(self, z):
+        x = self.conv1(z)
+        for i in range(self.num_blocks):
+            x = self.conv_blocks[i](x)
+        x = self.conv2(x)
+        return x

RepCodec/repcodec/modules/encoder.py

@@ -0,0 +1,89 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch
+import torch.nn as nn
+
+from RepCodec.repcodec.layers.conv_layer import Conv1d
+from RepCodec.repcodec.modules.residual_unit import ResidualUnit
+
+
+class EncoderBlock(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            stride: int,
+            dilations=(1, 1),
+            unit_kernel_size=3,
+            bias=True
+    ):
+        super().__init__()
+        self.res_units = torch.nn.ModuleList()
+        for dilation in dilations:
+            self.res_units += [
+                ResidualUnit(in_channels, in_channels,
+                             kernel_size=unit_kernel_size,
+                             dilation=dilation)
+            ]
+        self.num_res = len(self.res_units)
+
+        self.conv = Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3 if stride == 1 else (2 * stride),  # special case: stride=1, do not use kernel=2
+            stride=stride,
+            bias=bias,
+        )
+
+    def forward(self, x):
+        for idx in range(self.num_res):
+            x = self.res_units[idx](x)
+        x = self.conv(x)
+        return x
+
+
+class Encoder(nn.Module):
+    def __init__(
+            self,
+            input_channels: int,
+            encode_channels: int,
+            channel_ratios=(1, 1),
+            strides=(1, 1),
+            kernel_size=3,
+            bias=True,
+            block_dilations=(1, 1),
+            unit_kernel_size=3
+    ):
+        super().__init__()
+        assert len(channel_ratios) == len(strides)
+
+        self.conv = Conv1d(
+            in_channels=input_channels,
+            out_channels=encode_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            bias=False
+        )
+        self.conv_blocks = torch.nn.ModuleList()
+        in_channels = encode_channels
+        for idx, stride in enumerate(strides):
+            out_channels = int(encode_channels * channel_ratios[idx])  # could be float
+            self.conv_blocks += [
+                EncoderBlock(in_channels, out_channels, stride,
+                             dilations=block_dilations, unit_kernel_size=unit_kernel_size,
+                             bias=bias)
+            ]
+            in_channels = out_channels
+        self.num_blocks = len(self.conv_blocks)
+        self.out_channels = out_channels
+
+    def forward(self, x):
+        x = self.conv(x)
+        for i in range(self.num_blocks):
+            x = self.conv_blocks[i](x)
+        return x

RepCodec/repcodec/modules/projector.py

@@ -0,0 +1,32 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+from repcodec.layers.conv_layer import Conv1d
+
+
+class Projector(nn.Module):
+    def __init__(
+            self,
+            input_channels: int,
+            code_dim: int,
+            kernel_size=3,
+            stride=1,
+            bias=False
+    ):
+        super().__init__()
+        self.project = Conv1d(
+            input_channels,
+            code_dim,
+            kernel_size=kernel_size,
+            stride=stride,
+            bias=bias
+        )
+
+    def forward(self, x):
+        return self.project(x)

RepCodec/repcodec/modules/quantizer.py

@@ -0,0 +1,46 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+from repcodec.layers.vq_module import ResidualVQ
+
+
+class Quantizer(nn.Module):
+    def __init__(
+            self,
+            code_dim: int,
+            codebook_num: int,
+            codebook_size: int,
+    ):
+        super().__init__()
+        self.codebook = ResidualVQ(
+            dim=code_dim,
+            num_quantizers=codebook_num,
+            codebook_size=codebook_size
+        )
+
+    def initial(self):
+        self.codebook.initial()
+
+    def forward(self, z):
+        zq, vqloss, perplexity = self.codebook(z.transpose(2, 1))
+        zq = zq.transpose(2, 1)
+        return zq, vqloss, perplexity
+
+    def inference(self, z):
+        zq, indices = self.codebook.forward_index(z.transpose(2, 1))
+        zq = zq.transpose(2, 1)
+        return zq, indices
+
+    def encode(self, z):
+        zq, indices = self.codebook.forward_index(z.transpose(2, 1), flatten_idx=True)
+        return zq, indices
+
+    def decode(self, indices):
+        z = self.codebook.lookup(indices)
+        return z

RepCodec/repcodec/modules/residual_unit.py

@@ -0,0 +1,39 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import torch.nn as nn
+
+from RepCodec.repcodec.layers.conv_layer import Conv1d, Conv1d1x1
+
+
+class ResidualUnit(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            out_channels: int,
+            kernel_size=3,
+            dilation=1,
+            bias=False,
+            nonlinear_activation="ELU",
+            nonlinear_activation_params={},
+    ):
+        super().__init__()
+        self.activation = getattr(nn, nonlinear_activation)(**nonlinear_activation_params)
+        self.conv1 = Conv1d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=1,
+            dilation=dilation,
+            bias=bias,
+        )
+        self.conv2 = Conv1d1x1(out_channels, out_channels, bias)
+
+    def forward(self, x):
+        y = self.conv1(self.activation(x))
+        y = self.conv2(self.activation(y))
+        return x + y

RepCodec/repcodec/tokenize.py

@@ -0,0 +1,212 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import os
+from pathlib import Path
+from typing import Tuple, List, Optional
+
+import numpy as np
+import torch
+import yaml
+from tqdm import tqdm
+
+from repcodec.RepCodec import RepCodec
+
+ALL_MODELS = {
+    "data2vec_base_l6": 768,
+    "data2vec_large_l18": 1024,
+    "hubert_base_l9": 768,
+    "hubert_large_l18": 1024,
+    "whisper_medium_l24": 1024,
+    "whisper_large_l32": 1280
+}
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "in_dir",
+        type=str,
+        help="directory of representations to be tokenized."
+    )
+    parser.add_argument(
+        "--model",
+        required=True,
+        type=str,
+        help="path of the RepCodec model."
+    )
+    parser.add_argument(
+        "--tsv_path",
+        required=True,
+        type=str,
+        help="path of the tsv file."
+    )
+    parser.add_argument(
+        "--model_config_path",
+        default=None,
+        type=str,
+        help="please provide this training config if you are using the model you trained yourself."
+    )
+    parser.add_argument(
+        "--n_shard",
+        required=False,
+        type=int,
+        default=1,
+        help="number of shards of representations."
+    )
+    parser.add_argument(
+        "--use_gpu",
+        default=False,
+        action="store_true",
+        help="whether use gpu for inference."
+    )
+    parser.add_argument(
+        "--batch_size",
+        default=1,
+        type=int,
+        help="number of utterances for each mini batch."
+    )
+    parser.add_argument(
+        "--out_dir",
+        type=str,
+        default=".",
+        help="the directory to save the output."
+    )
+    return parser.parse_args()
+
+
+def load_model(model_path: str, config_path: Optional[str] = None):
+    if config_path is None:
+        name = os.path.basename(model_path).strip(".pkl")
+        assert name in ALL_MODELS.keys(), f"Cannot find configs for {model_path}. " \
+                                          f"Please provide the config file you used for training."
+        config = os.path.join(os.path.dirname(__file__), "configs", f"repcodec_dim{ALL_MODELS[name]}.yaml")
+        with open(config) as fp:
+            conf = yaml.load(fp, Loader=yaml.FullLoader)
+    else:
+        with open(config_path) as fp:
+            conf = yaml.load(fp, Loader=yaml.FullLoader)["model_params"]
+
+    model = RepCodec(**conf)
+    model.load_state_dict(torch.load(model_path, map_location="cpu")["model"]["repcodec"])
+    model.quantizer.initial()
+    model.eval()
+    return model
+
+
+def load_shard(in_dir: Path, rank: int, n_shard: int) -> Tuple[np.ndarray, List[int]]:
+    feat_path = in_dir / f"{rank}_{n_shard}.npy"
+    len_path = in_dir / f"{rank}_{n_shard}.len"
+
+    with open(len_path) as fp:
+        lengths = [int(line.strip()) for line in fp]
+
+    return np.load(feat_path.as_posix(), mmap_mode="r"), lengths
+
+
+def pad_data(data: List[np.ndarray]) -> List[np.ndarray]:
+    max_len = max([d.shape[0] for d in data])
+    data = [
+        np.pad(d, [(0, max_len - d.shape[0]), (0, 0)], "constant", constant_values=0.0)
+        for d in data
+    ]
+    return data
+
+
+def make_batch_data(data: np.ndarray, shard_lengths: List[int], batch_size: int):
+    batch_data = []
+    batch_lens = []
+    offsets = np.cumsum([0] + shard_lengths)
+    assert len(data) == offsets[-1], f"{len(data)} {offsets[-1]}"
+
+    # from longest to shortest
+    for i in range(len(shard_lengths)):
+        if batch_size > len(batch_data):
+            batch_data.append(data[offsets[i]: offsets[i + 1]])
+            batch_lens.append(shard_lengths[i])
+        else:
+            yield {
+                "data": torch.tensor(np.stack(pad_data(batch_data)), dtype=torch.float),  # (bsz, seq len, hidden dim)
+                "lengths": batch_lens
+            }
+            batch_data = [data[offsets[i]: offsets[i + 1]]]
+            batch_lens = [shard_lengths[i]]
+    if len(batch_data) > 0:
+        yield {
+            "data": torch.tensor(np.stack(pad_data(batch_data)), dtype=torch.float),
+            "lengths": batch_lens
+        }
+
+
+def tokenize_batch(model: RepCodec, batch: dict, device: str) -> List[List[int]]:
+    with torch.no_grad():
+        data = batch["data"].transpose(1, 2).to(device)  # (bsz, hidden dim, seq len)
+        x = model.encoder(data)
+        z = model.projector(x)
+        _, idx = model.quantizer.codebook.forward_index(z.transpose(2, 1))
+
+    # when bsz=1: (1, seq len)
+    if idx.dim() == 2:
+        return idx.cpu().data.numpy().tolist()
+    # when bsz>1: (1, bsz, seq len)
+    tokens = idx.cpu().data.numpy().tolist()[0]
+    res = []
+    batch_lens = batch["lengths"]
+    for i in range(len(tokens)):
+        n_tokens = batch_lens[i]
+        res.append(tokens[i][:n_tokens])
+    return res
+
+
+def load_tsv(path: str):
+    with open(path) as fp:
+        root = fp.readline().strip()
+        names = []
+        for line in fp:
+            names.append(line.strip().split("\t")[0])
+    return root, names
+
+
+def cli():
+    args = parse_args()
+    device = "cuda" if args.use_gpu else "cpu"
+
+    model = load_model(model_path=args.model, config_path=args.model_config_path)
+    model.to(device)
+
+    in_dir = Path(args.in_dir)
+    n_shard = args.n_shard
+    batch_size = args.batch_size
+
+    root_dir, file_names = load_tsv(args.tsv_path)
+
+    output_dir = args.out_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    processed_cnt = 0
+    pbar = tqdm(total=len(file_names))
+    with open(os.path.join(output_dir, "tokens"), mode="w+") as fp:
+        fp.write(f"{root_dir}\n")
+
+        for rank in range(n_shard):
+            shard_data, shard_lengths = load_shard(in_dir, rank, n_shard)
+            for batch in make_batch_data(shard_data, shard_lengths, batch_size=batch_size):
+                batch_tokens = tokenize_batch(model, batch, device)
+
+                for tokens in batch_tokens:
+                    fp.write(f"{file_names[processed_cnt]}\t{' '.join(map(str, tokens))}\n")
+                    processed_cnt += 1
+
+                pbar.update(len(batch_tokens))
+    assert processed_cnt == len(file_names), f"# lines of tsv do not match # of representations!"
+
+    pbar.close()
+    print("Tokenize successfully!")
+
+
+if __name__ == '__main__':
+    cli()

RepCodec/setup.py

@@ -0,0 +1,31 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from setuptools import setup
+
+try:
+    with open("README.md") as fp:
+        long_description = fp.read()
+except Exception:
+    long_description = ""
+
+setup(
+    name="RepCodec",
+    version="v1.0.0",
+    description="A Speech Representation Codec for Speech Tokenization",
+    long_description=long_description,
+    url="https://github.com/mct10/RepCodec",
+    packages=["repcodec", "repcodec.modules", "repcodec.layers"],
+    package_data={
+        "repcodec": ["configs/*.yaml"]
+    },
+    install_requires=["numpy", "tqdm", "torch", "tensorboardX", "PyYAML"],
+    entry_points={
+        'console_scripts': [
+            "repcodec=repcodec.tokenize:cli"
+        ]
+    }
+)

RepCodec/train.py

@@ -0,0 +1,228 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import argparse
+import logging
+
+import os
+
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+)
+logger = logging.getLogger("repcodec_train")  # init logger before other modules
+
+import random
+
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader
+
+from dataloader import ReprDataset, ReprCollater
+from losses.repr_reconstruct_loss import ReprReconstructLoss
+from repcodec.RepCodec import RepCodec
+from trainer.autoencoder import Trainer
+
+
+class TrainMain:
+    def __init__(self, args):
+        # Fix seed and make backends deterministic
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if not torch.cuda.is_available():
+            self.device = torch.device('cpu')
+            logger.info(f"device: cpu")
+        else:
+            self.device = torch.device('cuda:0')  # only supports single gpu for now
+            logger.info(f"device: gpu")
+            torch.cuda.manual_seed_all(args.seed)
+            if args.disable_cudnn == "False":
+                torch.backends.cudnn.benchmark = True
+
+        # initialize config
+        with open(args.config, 'r') as f:
+            self.config = yaml.load(f, Loader=yaml.FullLoader)
+        self.config.update(vars(args))
+
+        # initialize model folder
+        expdir = os.path.join(args.exp_root, args.tag)
+        os.makedirs(expdir, exist_ok=True)
+        self.config["outdir"] = expdir
+
+        # save config
+        with open(os.path.join(expdir, "config.yml"), "w") as f:
+            yaml.dump(self.config, f, Dumper=yaml.Dumper)
+        for key, value in self.config.items():
+            logger.info(f"{key} = {value}")
+
+        # initialize attribute
+        self.resume: str = args.resume
+        self.data_loader = None
+        self.model = None
+        self.optimizer = None
+        self.scheduler = None
+        self.criterion = None
+        self.trainer = None
+
+        # initialize batch_length
+        self.batch_length: int = self.config['batch_length']
+        self.data_path: str = self.config['data']['path']
+
+    def initialize_data_loader(self):
+        train_set = self._build_dataset("train")
+        valid_set = self._build_dataset("valid")
+        collater = ReprCollater()
+
+        logger.info(f"The number of training files = {len(train_set)}.")
+        logger.info(f"The number of validation files = {len(valid_set)}.")
+        dataset = {"train": train_set, "dev": valid_set}
+        self._set_data_loader(dataset, collater)
+
+    def define_model_optimizer_scheduler(self):
+        # model arch
+        self.model = {
+            "repcodec": RepCodec(**self.config["model_params"]).to(self.device)
+        }
+        logger.info(f"Model Arch:\n{self.model['repcodec']}")
+
+        # opt
+        optimizer_class = getattr(
+            torch.optim,
+            self.config["model_optimizer_type"]
+        )
+        self.optimizer = {
+            "repcodec": optimizer_class(
+                self.model["repcodec"].parameters(),
+                **self.config["model_optimizer_params"]
+            )
+        }
+
+        # scheduler
+        scheduler_class = getattr(
+            torch.optim.lr_scheduler,
+            self.config.get("model_scheduler_type", "StepLR"),
+        )
+        self.scheduler = {
+            "repcodec": scheduler_class(
+                optimizer=self.optimizer["repcodec"],
+                **self.config["model_scheduler_params"]
+            )
+        }
+
+    def define_criterion(self):
+        self.criterion = {
+            "repr_reconstruct_loss": ReprReconstructLoss(
+                **self.config.get("repr_reconstruct_loss_params", {}),
+            ).to(self.device)
+        }
+
+    def define_trainer(self):
+        self.trainer = Trainer(
+            steps=0,
+            epochs=0,
+            data_loader=self.data_loader,
+            model=self.model,
+            criterion=self.criterion,
+            optimizer=self.optimizer,
+            scheduler=self.scheduler,
+            config=self.config,
+            device=self.device
+        )
+
+    def initialize_model(self):
+        initial = self.config.get("initial", "")
+        if os.path.exists(self.resume):  # resume from trained model
+            self.trainer.load_checkpoint(self.resume)
+            logger.info(f"Successfully resumed from {self.resume}.")
+        elif os.path.exists(initial):  # initial new model with the pre-trained model
+            self.trainer.load_checkpoint(initial, load_only_params=True)
+            logger.info(f"Successfully initialize parameters from {initial}.")
+        else:
+            logger.info("Train from scrach")
+
+    def run(self):
+        assert self.trainer is not None
+        self.trainer: Trainer
+        try:
+            logger.info(f"The current training step: {self.trainer.steps}")
+            self.trainer.train_max_steps = self.config["train_max_steps"]
+            if not self.trainer._check_train_finish():
+                self.trainer.run()
+        finally:
+            self.trainer.save_checkpoint(
+                os.path.join(self.config["outdir"], f"checkpoint-{self.trainer.steps}steps.pkl")
+            )
+            logger.info(f"Successfully saved checkpoint @ {self.trainer.steps}steps.")
+
+    def _build_dataset(
+            self, subset: str
+    ) -> ReprDataset:
+        data_dir = os.path.join(
+            self.data_path, self.config['data']['subset'][subset]
+        )
+        params = {
+            "data_dir": data_dir,
+            "batch_len": self.batch_length
+        }
+        return ReprDataset(**params)
+
+    def _set_data_loader(self, dataset, collater):
+        self.data_loader = {
+            "train": DataLoader(
+                dataset=dataset["train"],
+                shuffle=True,
+                collate_fn=collater,
+                batch_size=self.config["batch_size"],
+                num_workers=self.config["num_workers"],
+                pin_memory=self.config["pin_memory"],
+            ),
+            "dev": DataLoader(
+                dataset=dataset["dev"],
+                shuffle=False,
+                collate_fn=collater,
+                batch_size=self.config["batch_size"],
+                num_workers=0,
+                pin_memory=False,  # save some memory. set to True if you have enough memory.
+            ),
+        }
+
+
+def train():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c", "--config", type=str, required=True,
+        help="the path of config yaml file."
+    )
+    parser.add_argument(
+        "--tag", type=str, required=True,
+        help="the outputs will be saved to exp_root/tag/"
+    )
+    parser.add_argument(
+        "--exp_root", type=str, default="exp"
+    )
+    parser.add_argument(
+        "--resume", default="", type=str, nargs="?",
+        help='checkpoint file path to resume training. (default="")',
+    )
+    parser.add_argument("--seed", default=1337, type=int)
+    parser.add_argument("--disable_cudnn", choices=("True", "False"), default="False", help="Disable CUDNN")
+    args = parser.parse_args()
+
+    train_main = TrainMain(args)
+    train_main.initialize_data_loader()
+    train_main.define_model_optimizer_scheduler()
+    train_main.define_criterion()
+    train_main.define_trainer()
+    train_main.initialize_model()
+    train_main.run()
+
+
+if __name__ == '__main__':
+    train()

RepCodec/train_configs/ex_dim768_mse.yaml

@@ -0,0 +1,74 @@
+###########################################################
+#                   DATA SETTING                          #
+###########################################################
+data:
+    path: "/dir/to/representations/"
+    subset:
+        train: "train_set_name"
+        valid: "valid_set_name"
+        test:  "test_set_name"
+
+###########################################################
+#                   MODEL SETTING                         #
+###########################################################
+model_params:
+    input_channels: 768
+    output_channels: 768
+    encode_channels: 768
+    decode_channels: 768
+    code_dim: 768
+    codebook_num: 1
+    codebook_size: 1024
+    bias: true
+    enc_ratios: [1, 1]
+    dec_ratios: [1, 1]
+    enc_strides: [1, 1]  # no downsampling
+    dec_strides: [1, 1]
+    enc_kernel_size: 3
+    dec_kernel_size: 3
+    enc_block_dilations: [1, 1]
+    enc_block_kernel_size: 3
+    dec_block_dilations: [1, 1]
+    dec_block_kernel_size: 3
+
+###########################################################
+#                 METRIC LOSS SETTING                     #
+###########################################################
+repr_reconstruct_loss_params:
+    loss_type: l2
+
+###########################################################
+#                  LOSS WEIGHT SETTING                    #
+###########################################################
+lambda_vq_loss: 1.0      # Loss weight of vector quantize loss.
+lambda_repr_reconstruct_loss: 45.0
+
+###########################################################
+#                  DATA LOADER SETTING                    #
+###########################################################
+batch_size: 32              # Batch size.
+batch_length: 96            # Length of each audio in batch (training w/o adv).
+pin_memory: true            # Whether to pin memory in Pytorch DataLoader.
+num_workers: 4              # Number of workers in Pytorch DataLoader.
+
+###########################################################
+#             OPTIMIZER & SCHEDULER SETTING               #
+###########################################################
+model_optimizer_type: Adam
+model_optimizer_params:
+    lr: 1.0e-4
+    betas: [0.5, 0.9]
+    weight_decay: 0.0
+model_scheduler_type: StepLR
+model_scheduler_params:
+    step_size: 200000      # Model's scheduler step size.
+    gamma: 1.0
+grad_norm: -1
+
+###########################################################
+#                    INTERVAL SETTING                     #
+###########################################################
+train_max_steps: 200000            # Number of training steps. (w/o adv)
+save_interval_steps: 20000         # Interval steps to save checkpoint.
+eval_interval_steps: 2000          # Interval steps to evaluate the network.
+log_interval_steps: 100            # Interval steps to record the training log.

RepCodec/trainer/autoencoder.py

@@ -0,0 +1,287 @@
+# Copyright (c) ByteDance, Inc. and its affiliates.
+# Copyright (c) Chutong Meng
+#
+# This source code is licensed under the CC BY-NC license found in the
+# LICENSE file in the root directory of this source tree.
+# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
+
+import logging
+import os
+from collections import defaultdict
+
+import torch
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+logger = logging.getLogger("repcodec_train")
+
+
+class Trainer:
+    def __init__(
+            self,
+            steps: int,
+            epochs: int,
+            data_loader: dict,
+            model: dict,
+            criterion: dict,
+            optimizer: dict,
+            scheduler: dict,
+            config: dict,
+            device=torch.device("cpu"),
+    ):
+        self.steps = steps
+        self.epochs = epochs
+        self.data_loader = data_loader
+        self.model = model
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.config = config
+        self.device = device
+        self.writer = SummaryWriter(config["outdir"])
+        self.total_train_loss = defaultdict(float)
+        self.total_eval_loss = defaultdict(float)
+        self.train_max_steps = config.get("train_max_steps", 0)
+
+    def _train_step(self, batch):
+        """Single step of training."""
+        mode = "train"
+        x = batch
+        x = x.to(self.device)
+
+        codec_loss = 0.0
+        y_, zq, z, vqloss, perplexity = self.model["repcodec"](x)
+        self._perplexity(perplexity, mode=mode)
+        codec_loss += self._vq_loss(vqloss, mode=mode)
+        codec_loss += self._metric_loss(y_, x, mode=mode)
+
+        self._record_loss("codec_loss", codec_loss, mode=mode)
+        self._update_repcodec(codec_loss)
+
+        self.steps += 1
+        self.tqdm.update(1)
+        self._check_train_finish()
+
+    @torch.no_grad()
+    def _eval_step(self, batch):
+        """Single step of evaluation."""
+        mode = "eval"
+        x = batch
+        x = x.to(self.device)
+
+        codec_loss = 0.0
+        y_, zq, z, vqloss, perplexity = self.model["repcodec"](x)
+        self._perplexity(perplexity, mode=mode)
+        codec_loss += self._vq_loss(vqloss, mode=mode)
+        codec_loss += self._metric_loss(y_, x, mode=mode)
+
+        self._record_loss("codec_loss", codec_loss, mode=mode)
+
+    def run(self):
+        """Run training."""
+        self.finish_train = False
+        self.tqdm = tqdm(
+            initial=self.steps, total=self.train_max_steps, desc="[train]"
+        )
+        while True:
+            self._train_epoch()
+
+            # check whether training is finished
+            if self.finish_train:
+                break
+
+        self.tqdm.close()
+        logger.info("Finished training.")
+
+    def save_checkpoint(self, checkpoint_path: str):
+        state_dict = {
+            "model": {
+                "repcodec": self.model["repcodec"].state_dict()
+            },
+            "optimizer": {
+                "repcodec": self.optimizer["repcodec"].state_dict(),
+            },
+            "scheduler": {
+                "repcodec": self.scheduler["repcodec"].state_dict(),
+            },
+            "steps": self.steps,
+            "epochs": self.epochs,
+        }
+
+        if not os.path.exists(os.path.dirname(checkpoint_path)):
+            os.makedirs(os.path.dirname(checkpoint_path))
+        torch.save(state_dict, checkpoint_path)
+
+    def load_checkpoint(
+            self,
+            checkpoint_path: str,
+            strict: bool = True,
+            load_only_params: bool = False
+    ):
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        self.model["repcodec"].load_state_dict(
+            state_dict["model"]["repcodec"], strict=strict
+        )
+
+        if not load_only_params:
+            self.steps = state_dict["steps"]
+            self.epochs = state_dict["epochs"]
+            self.optimizer["repcodec"].load_state_dict(
+                state_dict["optimizer"]["repcodec"]
+            )
+            self.scheduler["repcodec"].load_state_dict(
+                state_dict["scheduler"]["repcodec"]
+            )
+
+    def _train_epoch(self):
+        """One epoch of training."""
+        for train_steps_per_epoch, batch in enumerate(self.data_loader["train"], 1):
+            # train one step
+            self._train_step(batch)
+
+            # check interval
+            self._check_log_interval()
+            self._check_eval_interval()
+            self._check_save_interval()
+
+            # check whether training is finished
+            if self.finish_train:
+                return
+
+        # update
+        self.epochs += 1
+        self.train_steps_per_epoch = train_steps_per_epoch
+        if train_steps_per_epoch > 200:
+            logger.info(
+                f"(Steps: {self.steps}) Finished {self.epochs} epoch training "
+                f"({self.train_steps_per_epoch} steps per epoch)."
+            )
+
+    def _eval_epoch(self):
+        """One epoch of evaluation."""
+        logger.info(f"(Steps: {self.steps}) Start evaluation.")
+        # change mode
+        for key in self.model.keys():
+            self.model[key].eval()
+
+        # calculate loss for each batch
+        for eval_steps_per_epoch, batch in enumerate(
+                tqdm(self.data_loader["dev"], desc="[eval]"), 1
+        ):
+            # eval one step
+            self._eval_step(batch)
+
+        logger.info(
+            f"(Steps: {self.steps}) Finished evaluation "
+            f"({eval_steps_per_epoch} steps per epoch)."
+        )
+
+        # average loss
+        for key in self.total_eval_loss.keys():
+            self.total_eval_loss[key] /= eval_steps_per_epoch
+            logger.info(
+                f"(Steps: {self.steps}) {key} = {self.total_eval_loss[key]:.4f}."
+            )
+
+        # record
+        self._write_to_tensorboard(self.total_eval_loss)
+
+        # reset
+        self.total_eval_loss = defaultdict(float)
+
+        # restore mode
+        for key in self.model.keys():
+            self.model[key].train()
+
+    def _metric_loss(self, predict_y, natural_y, mode='train'):
+        """Metric losses."""
+        metric_loss = 0.0
+
+        repr_reconstruct_loss = self.criterion["repr_reconstruct_loss"](predict_y, natural_y)
+        repr_reconstruct_loss *= self.config["lambda_repr_reconstruct_loss"]
+        self._record_loss("reconstruct_loss", repr_reconstruct_loss, mode=mode)
+        metric_loss += repr_reconstruct_loss
+
+        return metric_loss
+
+    def _update_repcodec(self, repr_loss):
+        """Update generator."""
+        self.optimizer["repcodec"].zero_grad()
+        repr_loss.backward()
+        if self.config["grad_norm"] > 0:
+            torch.nn.utils.clip_grad_norm_(
+                self.model["repcodec"].parameters(),
+                self.config["grad_norm"],
+            )
+        self.optimizer["repcodec"].step()
+        self.scheduler["repcodec"].step()
+
+    def _record_loss(self, name: str, loss, mode='train'):
+        """Record loss."""
+        if torch.is_tensor(loss):
+            loss = loss.item()
+
+        if mode == 'train':
+            self.total_train_loss[f"train/{name}"] += loss
+        elif mode == 'eval':
+            self.total_eval_loss[f"eval/{name}"] += loss
+        else:
+            raise NotImplementedError(f"Mode ({mode}) is not supported!")
+
+    def _write_to_tensorboard(self, loss):
+        """Write to tensorboard."""
+        for key, value in loss.items():
+            self.writer.add_scalar(key, value, self.steps)
+
+    def _check_save_interval(self):
+        if self.steps and (self.steps % self.config["save_interval_steps"] == 0):
+            self.save_checkpoint(
+                os.path.join(self.config["outdir"], f"checkpoint-{self.steps}steps.pkl")
+            )
+            logger.info(f"Successfully saved checkpoint @ {self.steps} steps.")
+
+    def _check_eval_interval(self):
+        if self.steps % self.config["eval_interval_steps"] == 0:
+            self._eval_epoch()
+
+    def _check_log_interval(self):
+        if self.steps % self.config["log_interval_steps"] == 0:
+            for key in self.total_train_loss.keys():
+                self.total_train_loss[key] /= self.config["log_interval_steps"]
+                logger.info(
+                    f"(Steps: {self.steps}) {key} = {self.total_train_loss[key]:.4f}."
+                )
+            self._write_to_tensorboard(self.total_train_loss)
+
+            # reset
+            self.total_train_loss = defaultdict(float)
+
+    def _check_train_finish(self):
+        if self.steps >= self.train_max_steps:
+            self.finish_train = True
+        else:
+            self.finish_train = False
+        return self.finish_train
+
+    def _perplexity(self, perplexity, label=None, mode='train'):
+        if label:
+            name = f"{mode}/ppl_{label}"
+        else:
+            name = f"{mode}/ppl"
+        if torch.numel(perplexity) > 1:
+            perplexity = perplexity.tolist()
+            for idx, ppl in enumerate(perplexity):
+                self._record_loss(f"{name}_{idx}", ppl, mode=mode)
+        else:
+            self._record_loss(name, perplexity, mode=mode)
+
+    def _vq_loss(self, vqloss, label=None, mode='train'):
+        if label:
+            name = f"{mode}/vqloss_{label}"
+        else:
+            name = f"{mode}/vqloss"
+        vqloss = torch.sum(vqloss)
+        vqloss *= self.config["lambda_vq_loss"]
+        self._record_loss(name, vqloss, mode=mode)
+
+        return vqloss

decoders/config.yaml

@@ -0,0 +1,15 @@
+backbone:
+  class_path: vocos.models.VocosBackbone
+  init_args:
+    input_channels: 1024
+    dim: 512
+    intermediate_dim: 1536
+    num_layers: 8
+
+head:
+  class_path: vocos.heads.ISTFTHead
+  init_args:
+    dim: 512
+    n_fft: 3528
+    hop_length: 882
+    padding: same

decoders/decoder_131000.pth

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

decoders/decoder_151000.pth

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

descriptaudiocodec/dac/__init__.py

@@ -0,0 +1,16 @@
+__version__ = "1.0.0"
+
+# preserved here for legacy reasons
+__model_version__ = "latest"
+
+import audiotools
+
+audiotools.ml.BaseModel.INTERN += ["dac.**"]
+audiotools.ml.BaseModel.EXTERN += ["einops"]
+
+
+from . import nn
+from . import model
+from . import utils
+from .model import DAC
+from .model import DACFile

descriptaudiocodec/dac/__main__.py

@@ -0,0 +1,36 @@
+import sys
+
+import argbind
+
+from dac.utils import download
+from dac.utils.decode import decode
+from dac.utils.encode import encode
+
+STAGES = ["encode", "decode", "download"]
+
+
+def run(stage: str):
+    """Run stages.
+
+    Parameters
+    ----------
+    stage : str
+        Stage to run
+    """
+    if stage not in STAGES:
+        raise ValueError(f"Unknown command: {stage}. Allowed commands are {STAGES}")
+    stage_fn = globals()[stage]
+
+    if stage == "download":
+        stage_fn()
+        return
+
+    stage_fn()
+
+
+if __name__ == "__main__":
+    group = sys.argv.pop(1)
+    args = argbind.parse_args(group=group)
+
+    with argbind.scope(args):
+        run(group)

descriptaudiocodec/dac/compare/encodec.py

@@ -0,0 +1,54 @@
+import torch
+from audiotools import AudioSignal
+from audiotools.ml import BaseModel
+from encodec import EncodecModel
+
+
+class Encodec(BaseModel):
+    def __init__(self, sample_rate: int = 24000, bandwidth: float = 24.0):
+        super().__init__()
+
+        if sample_rate == 24000:
+            self.model = EncodecModel.encodec_model_24khz()
+        else:
+            self.model = EncodecModel.encodec_model_48khz()
+        self.model.set_target_bandwidth(bandwidth)
+        self.sample_rate = 44100
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = 44100,
+        n_quantizers: int = None,
+    ):
+        signal = AudioSignal(audio_data, sample_rate)
+        signal.resample(self.model.sample_rate)
+        recons = self.model(signal.audio_data)
+        recons = AudioSignal(recons, self.model.sample_rate)
+        recons.resample(sample_rate)
+        return {"audio": recons.audio_data}
+
+
+if __name__ == "__main__":
+    import numpy as np
+    from functools import partial
+
+    model = Encodec()
+
+    for n, m in model.named_modules():
+        o = m.extra_repr()
+        p = sum([np.prod(p.size()) for p in m.parameters()])
+        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
+        setattr(m, "extra_repr", partial(fn, o=o, p=p))
+    print(model)
+    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
+
+    length = 88200 * 2
+    x = torch.randn(1, 1, length).to(model.device)
+    x.requires_grad_(True)
+    x.retain_grad()
+
+    # Make a forward pass
+    out = model(x)["audio"]
+
+    print(x.shape, out.shape)

descriptaudiocodec/dac/model/__init__.py

@@ -0,0 +1,4 @@
+from .base import CodecMixin
+from .base import DACFile
+from .dac import DAC
+from .discriminator import Discriminator