Upload 51 files

.gitattributes

@@ -36,3 +36,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 examples/web_6f93090a-81f6-489e-bb35-1a2838b18c01.png filter=lfs diff=lfs merge=lfs -text
 examples/web_dfacd48d-d2c2-492f-b94c-41e6a34ea99f.png filter=lfs diff=lfs merge=lfs -text
 illustration.png filter=lfs diff=lfs merge=lfs -text
+demo/5895_34622_000026_000002.wav filter=lfs diff=lfs merge=lfs -text
+generated_tts/generated.wav filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Puyuan Peng
+
+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.

MODEL-LICENSE

@@ -0,0 +1,395 @@
+Attribution 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 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 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. 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.
+
+  j. 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.
+
+  k. 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; and
+
+            b. produce, reproduce, and Share Adapted Material.
+
+       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.
+
+
+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;
+
+  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.

config.py

@@ -0,0 +1,254 @@
+import argparse
+
+def int_or_str(value):
+    """Custom function to allow both int and str types."""
+    try:
+        return int(value)  # Try converting to integer
+    except ValueError:
+        return value  # If conversion fails, return as string
+
+def MyParser():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    # general training 
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--debug", type=int, default=0)
+    parser.add_argument("--multinodes", type=int, default=0)
+    parser.add_argument("--dist_url", default="env://", type=str)
+    parser.add_argument("--dist_backend", default="nccl", type=str)
+    parser.add_argument("--precision", type=str, default="float16", help="we might need float32 for NAR model")
+    parser.add_argument("--num_workers", type=int, default=8, help="per gpu")
+    parser.add_argument("--resume", action="store_true", default=False)
+    parser.add_argument("--tb_write_every_n_steps", type=int, default=100)
+    parser.add_argument("--print_every_n_steps", type=int, default=250)
+    parser.add_argument("--val_every_n_steps", type=int, default=500)
+    parser.add_argument("--inference_every_n_steps", type=int, default=3000, help="will only get to inference when model is saved, and therefore this needs to be multiple of val_every_n_steps")
+    parser.add_argument("--save_every_n_steps", type=int, default=10000000, help="save the model every n steps, will save the model as bundle_step$step.pth")
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--batch_size", type=int, default=100, help="this is the effective batch size per gpu, no matter whether using gradient_accumulation_steps")
+    parser.add_argument("--weight_decay", type=float, default=1e-2)
+    parser.add_argument("--warmup_fraction", type=float, default=0.1, help="use linear warmup, the proportion of the training steps that are used for warming up")
+    parser.add_argument("--num_epochs", type=int, default=10)
+    parser.add_argument("--num_steps", type=int, default=None, help="if not None, will ignore n_epochs and use num_steps as the total number of amount of training, can try e.g. 400000 i.e. 400k steps")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
+    parser.add_argument("--gradient_clip_val", type=float, default=1.0, help="the value for torch.nn.utils.clip_grad_norm_()")
+    parser.add_argument("--early_stop_step", type=int, default=3200, help="stop training after this many steps of non-improvement")
+    parser.add_argument("--early_stop_threshold", type=float, default=-1.0, help="early stop after the improvement is below this threshold for certain number of steps")
+
+
+    # path
+    parser.add_argument("--exp_dir", type=str, default='/saltpool0/scratch/pyp/VoiceEditor/', help="will be combined with dataset name")
+    parser.add_argument("--dataset", type=str, help="e.g. 'libritts', 'librilight', 'spotify', they are folder name in the data dir also")
+    parser.add_argument("--dataset_dir", type=str, help="need to be compatible with corresponding dataset py file")
+    parser.add_argument("--compact_folder_name", type=str, default=None, help="if not None, will use compact_combined_dataset.py, and this is the folder name of the compact dataset")
+    parser.add_argument("--inference_dataset_dir", type=str, default="/data/scratch/pyp/datasets/librilight/preprocessed", help="need to be compatible with corresponding dataset py file")
+
+    parser.add_argument("--training_stage", type=int, default=1, help="if 1, train VoiceEditor_one, if 2 train VoiceEditor_seven")
+    parser.add_argument("--local_wandb", type=int, default=0, help="if 1, will use local wandb, otherwise use the global one")
+    parser.add_argument("--wandb_entity", type=str, default="puyuanpeng", help="the entity (usually your username) for wandb")
+    # data
+    parser.add_argument("--librilight_ratio", type=float, default=1, help='the portion of lightlight compared to gigaspeech, 1 means equal, 2 means librilight data is twice as much as gigaspeech')
+    parser.add_argument("--plus_librilight_root", type=str, default=None, help="if not None, will combine gigaspeech and librilight, this is the root folder to librilight. Note that will need to merge the vocab.txt based on gigaspeech's, in order to be able to load a pretrained model")
+    parser.add_argument("--plus_librilight_phn_folder_name", type=str, default=None, help="if not None, will combine gigaspeech and librilight, this is the phoneme folder name of librilight")
+    parser.add_argument("--plus_librilight_encodec_folder_name", type=str, default=None, help="if not None, will combine gigaspeech and librilight, this is the encodec folder name of librilight")
+    parser.add_argument("--plus_librilight_manifest_name", type=str, default=None, help="if not None, will combine gigaspeech and librilight, this is the manifest folder name of librilight")
+    parser.add_argument("--skip_us", type=int, default=0, help="skip the giga utterances that contains 'j uː ɛ s' because of the tokenization issue")
+    parser.add_argument("--pseudo_epoch_size", type=int, default=37901, help="only use for Eden scheduler. 37901 is the epoch size in the default optim setting, this is probably too big")
+    parser.add_argument("--switch_order", type=int, default=0, help="this is only for hificodec, where we switch the order of 2 and 3nd codebook")
+    parser.add_argument("--phn_folder_name", type=str, default="phoneme", help="for libritts I also have arpa phns, in which case should be phonemes_arpa")
+    parser.add_argument("--encodec_folder_name", type=str, default="mimi_8cb", help="folder where encodec codes are stored")
+    parser.add_argument("--manifest_name", type=str, default="manifest_final", help="if using hificodec, it should be hificodec_menifest, if using encodec, it is the default")
+    parser.add_argument("--pad_x", type=int, default=1, help="whether or not always pad x to have text_max_length. select 1 to get the maximal memory consumption, but the actual case should be smaller, better to have it being 0")
+    parser.add_argument("--max_num_tokens", type=int, default=18750, help="max number of encodec tokens per gpu, this is only used when using dynamic batching, will ignore batch size. Note that batch size is the final effective batch size (sum of batch on each gpu), but max_num_tokens is per gpu")
+    parser.add_argument("--val_max_num_tokens", type=int, default=6000, help="FOR validation, this basically is for music-gen because of high mem consumption. max number of encodec tokens per gpu, this is only used when using dynamic batching, will ignore batch size. Note that batch size is the final effective batch size (sum of batch on each gpu), but max_num_tokens is per gpu")
+    parser.add_argument("--num_buckets", type=int, default=10)
+    parser.add_argument("--dynamic_batching", type=int, default=1)
+    parser.add_argument("--audio_max_length", type=float, default=120, help="in second, crop the audio is length is longer than this")
+    parser.add_argument("--audio_min_length", type=float, default=2, help="in second, drop the audio if length is shorter than this")
+    parser.add_argument("--text_max_length", type=int, default=1000, help='if too long, we crop')
+    parser.add_argument("--text_min_length", type=float, default=10, help="if too short, will drop")
+    parser.add_argument("--encodec_sr", type=float, default=50, help="for 24kHz mimi model, it produces 12.5 codes for 1 sec of audio")
+    parser.add_argument('--mask_len_min', type=int, default=20, help='Minimum mask length')
+    parser.add_argument('--mask_len_max', type=int, default=400, help='Maximum mask length')
+    parser.add_argument('--extra_mask_len_min', type=int, default=2, help='Minimum extra mask length')
+    parser.add_argument('--extra_mask_len_max', type=int, default=20, help='Maximum extra mask length')
+    parser.add_argument('--final_audio_token_len', type=int, default=772, help="this is only for stage 1 training, since we add eog, start_of_continue, and a random amount of extra mask, --audio_max_length won't be the final max length, the self.args.final_audio_token_len = self.args.audio_max_length*self.args.encodec_sr+self.args.extra_mask_len_max+2 ")
+
+    # model
+    parser.add_argument("--ttsonly", default=0, type=int, help="if 1, only train tts model, no CM3")
+    parser.add_argument("--load_existing_text_embedding", type=int, default=0, help="if 1, when load model and the text vocab doesn't match, will load the existing weights while the new weights will be initialized randomly")
+    parser.add_argument("--fly", type=int, default=0, help="if 1, encode chunked audio on the fly")
+    parser.add_argument("--encodec_ckpt", type=str, default="/data/scratch/pyp/exp_pyp/audiocraft/encodec/xps/6f79c6a8/checkpoint.th")
+    parser.add_argument("--downsample_rate", type=int, default=320, help="the downsample rate for the encodec model, 16000/320 = 50Hz")
+    parser.add_argument("--segtts_mask", type=int, default=0, help="if 1, use segtts_mask model, where we have a prefix and segment utterance into two and shifted separately for modeling, and use make use of mask:0, by insert two mask:0 in the middle of the two segments")
+    parser.add_argument("--segtts", type=int, default=0, help="if 1, use segtts model, where we have a prefix and segment utterance into two and shifted separately for modeling")
+    parser.add_argument("--edge", type=int, default=0, help="if 1, use edge prediction for the first codebook")
+    parser.add_argument("--duration_loss_weight", type=float, default=1.0, help="weight on the duration loss")
+    parser.add_argument("--drop_long", type=int, default=1, help="if this is true, will drop example whose encodec sequence or phone sequence is too long, rather than cropping as we did before, to avoid hellucination")
+    parser.add_argument("--eos", type=int, default=2051, help="this is to be used with reduced_eog, where we end the utterance with eos, and end the generated segment with eog, also when this is used, the n_special should be 4")
+    parser.add_argument("--reduced_eog", type=int, default=1, help="for the non-final segments, do not insert eog at the end, this could hopefully solve the early stopping issue when doing tts")
+
+    parser.add_argument("--valle_orig", type=int, default=0, help="the original valle model, trained for TTS")
+    parser.add_argument("--valle_max_prompt_len", type=float, default=6, help='in sec.')
+    # randomly choose a portion as tts examples during training
+    parser.add_argument("--tts_portion", type=float, default=0, help="randomly choose a portion of the training examples as tts examples, where no mask and rearrangement is used")
+
+    # put special tokens first to handle different vocab_size
+    parser.add_argument("--special_first", type=int, default=0, help="if 1, need to have special tokens to be the first few tokens, e.g. 0, 1, 2, which means we need to adjust the preprocessing and postprocessing of the encodec codes. note that we hard coded to have 3 special tokens")
+    parser.add_argument("--n_special", type=int, default=4, help="empty, eog, pad, eos")
+
+    # weight codebook differently
+    parser.add_argument("--codebook_weight", type=str, default=None, help="e.g. ['5','1','0.5','0.1']")
+
+    # args for MusicGen
+    parser.add_argument("--mask_span_weight", default=1.0, type=float, help="the weight on the tokens in masked span")
+    parser.add_argument("--unmask_span_weight", default=1.0, type=float, help="the weight on unmasked span")
+    parser.add_argument("--start_end_weight", default=None, type=str, help="weight the start x tokens and end x tokens differently, e.g. (10,2.0), means x == 10, weight==2.0")
+    # for now not consider the two weights above, only consider eog_weight, which is defined below somewhere, as the above two are not super principled
+
+    parser.add_argument("--musicgen", type=int, default=0, help="whether or not use this model, will also have an impact on the output shape of the dataset")
+    parser.add_argument("--enc_dec", default=0, type=int, help="use enc-dec architecture, text is from the enc, only for musicgen")
+    parser.add_argument("--dec", default=0, type=int, help="use dec only architecture, text is from the enc, only for musicgen. Exclusive with --enc_dec")
+    parser.add_argument("--empty_token", default=2048, type=int, help="indicating the no token at the position for the codebook")
+    # args for the optimizer and scheduler from Feiteng
+    # original setup for the 3 params are 5000 4 and 1000
+    # but that's because set_epoch is run on num_gradient_accumulation_step*step (with 4 being the accumulation step)
+    # so I scaled down them a little bit
+    # will try scaling them back if this doesn't work
+    parser.add_argument("--optimizer_name", type=str, default="AdamW", help="can also use ScaledAdam, in which case we'll also use the Eden scheduler")
+    parser.add_argument("--reduce_lr_start_step", type=int, default=3000, help='after which significantly reduce the lr. a param for the eden optimizer')
+    parser.add_argument("--reduce_lr_start_epoch", type=int, default=4)
+    parser.add_argument("--clipping_update_period", type=int, default=600)
+
+
+    # below are args for valle
+    # below are args for valle
+    parser.add_argument("--valle", type=int, default=0, help="if 1, use valle model (cm3)")
+    parser.add_argument("--decoder_dim", type=int, default=1024)
+    parser.add_argument("--norm_first", action="store_true", default=True)
+    parser.add_argument("--add_prenet", action="store_true", default=False)
+    parser.add_argument("--prefix_mode", type=int, default=5, help="this is for NAR, we only do 5, which is CM3")
+    parser.add_argument("--share_embedding", action="store_true", default=False)
+    parser.add_argument("--nar_scale_factor", type=float, default=1.0)
+    parser.add_argument("--prepend_bos", action="store_true", default=False)
+    parser.add_argument("--sync_nar", type=int, default=0, help="whether to choose the same NAR model to run for training_stage==2 across different process (this is only for DDP)")
+    # above are args for valle
+    # above are args for valle
+
+
+
+    # add parallel_pattern
+    parser.add_argument("--parallel_pattern", type=int, default=0, help="if 1, use parallel pattern, we also use LFSC codec")
+    parser.add_argument("--full_prediction", type=int, default=0, help='this is for ve1, if 1, use full autoregressive mask, and calculate loss over all tokens, except for mask_tokens')
+    parser.add_argument("--multicm3", type=int, default=0, help='cm3 model but allows multiple mask spans')
+    parser.add_argument("--max_mask_portion",type=float,default=0.7,help="should mask a utterance for more than this portion")
+    parser.add_argument("--max_n_spans", type=int, default=8, help='maximal number of spans, only use when using multicm3, this is used to decide number of mask_embedding, and max clamp value if use Poisson distribution, if use uniform distribution to sample number of spans if will be uniform(1,max_n_spans)')
+    parser.add_argument("--shuffle_mask_embedding", type=int, default=0, help="whether shuffle the mask embedding, so that mask:0 is not the most well trained, default is not shuffling. The default has it's benefit, as it make sure that mask:0 always appear the first")
+    parser.add_argument("--mask_sample_dist", type=str, default="uniform", help="uniform or poissonx, e.g. poisson1, meaning the parameter lambda is 1, it will most likely sample 1 masks")
+    parser.add_argument("--min_gap", type=int, default=10, help="after sampled starts, delete later one if it closer to the former start than the min_gap")
+
+    parser.add_argument('--cm3', type=int, default=0, help="use cm3 style for ve1, the input from dataloader is going to be just raw data, all masking and rearrangement will happen whin the model")
+
+    parser.add_argument('--sep_special_token', type=int, default=0, help="remove text/audio pad token, set audio_mask_token and start of continue to be separately learned embeddings. Therefore, for ve1 self.n_text_tokens == self.args.text_vocab_size, self.n_audio_tokens == self.args.audio_vocab_size + 2, for ve7, self.n_text_tokens == self.args.text_vocab_size, self.n_audio_tokens == self.args.audio_vocab_size")
+    parser.add_argument('--one_causal', type=int, default=0, help="whether model VE_one generation as autoregressive gen or non-autoregressive gen")
+    parser.add_argument('--n_codebooks', type=int, default=8)
+    parser.add_argument('--weight_sharing', type=int, default=0, help="sharing weights between VE_seven predict layer and embedding layer")
+    parser.add_argument('--text_vocab_size', type=int, default=86, help='Size of text vocabulary')
+    parser.add_argument('--text_pad_token', type=int, default=86, help='padding of the text tokens, not attended')
+    # parser.add_argument('--audio_vocab_size', type=int, default=1024, help='Size of audio vocabulary')
+    parser.add_argument('--audio_vocab_size', type=str, default='2048', help="Size of audio vocabulary, can be specified as '[128,512,1024,2048]'")
+    parser.add_argument('--audio_mask_token', type=int, default=1024, help='Audio mask token, this the the extra mask used in the masked region for AR, for NAR, the entire masked region will be filled with it')
+    parser.add_argument('--bog', type=int, default=1025, help='Begin of generation token')
+    parser.add_argument('--eog', type=int, default=2049, help='End of generation token')
+    parser.add_argument('--start_of_continue', type=int, default=1027, help='this token follows the masked region, proceeds the first unmasked token, to indicate that gt tokens starts')
+    parser.add_argument('--audio_pad_token', type=int, default=2050, help='padding of the encodec codes, not attended')
+    parser.add_argument('--d_model', type=int, default=1024, help='Model dimension')
+    parser.add_argument('--audio_embedding_dim', type=int, default=128, help='dimension for encodec continues embedding (before being quantized)')
+    parser.add_argument('--text_embedding_dropout', type=float, default=0.1, help='Dropout for text embedding')
+    parser.add_argument('--audio_embedding_dropout', type=float, default=0, help='Dropout for audio embedding')
+    parser.add_argument('--text_positional_embedding_dropout', type=float, default=0.1, help='Dropout for text positional embedding')
+    parser.add_argument('--audio_positional_embedding_dropout', type=float, default=0.1, help='Dropout for audio positional embedding')
+    parser.add_argument('--trm_dropout', type=float, default=0.1, help='Dropout for transformer')
+    parser.add_argument('--nhead', type=int, default=16, help='Number of attention heads')
+    parser.add_argument('--num_encoder_layers', type=int, default=12, help='Number of encoder layers')
+    parser.add_argument('--num_decoder_layers', type=int, default=12, help='Number of decoder layers')
+    parser.add_argument('--eog_weight', type=float, default=1.0, help='Weight for End of generation token')
+    parser.add_argument('--stage_one_load_encodec_embedding', type=str, default=None, help='Path to load encodec embedding for stage one. On our lab machine it is /saltpool0/scratch/pyp/VoiceEditor/encodec_embedding/24khz_8codebooks.pth, 8 is the n_codebooks')
+    parser.add_argument('--stage_two_load_encodec_embedding', type=str, default=None, help='Path to load encodec embedding for stage two。 On our lab machine it is /saltpool0/scratch/pyp/VoiceEditor/encodec_embedding/24khz_8codebooks.pth, 8 is the n_codebooks')
+    parser.add_argument('--stage_two_load_ve_one_embedding', type=str, default=None, help='Path to load VoiceEditor_one audio embedding for stage two')
+    parser.add_argument('--load_model_from', type=str, default=None, help='Path to load model from, this will be effective last, so will overwrite all previous load, including resume')
+    parser.add_argument('--load_model_from_ve1', type=str, default=None, help='Path to load ve1 model weights from, this will be effective last, designed for loading the encoder weights of the VE7 from a pretrained VE1')
+
+
+    ## below are args for the new long model
+    parser.add_argument("--target_time_stretch_prob", type=float, default=0, help="the probability of time stretching the target audio")
+    parser.add_argument("--target_time_stretch_bound", type=float, default=0.1, help="the bound of the time stretching target audio, e.g. 0.1 means the audio will be stretched by 0.9 to 1.1")
+    parser.add_argument("--time_stretch_prob", type=float, default=0, help="the probability of time stretching the audio")
+    parser.add_argument("--time_stretch_bound", type=float, default=0.3, help="the bound of the time stretching, e.g. 0.3 means the audio will be stretched by 0.7 to 1.3")
+    parser.add_argument("--no_loss_on_prefix", type=int, default=0, help="if 1, will not calculate loss on the prefix acoustic tokens")
+    parser.add_argument("--x_sep_token", type=int, default=None, help="if not None, will use this token in between prompt text and target generation text")
+    parser.add_argument("--y_sep_token", type=int, default=None, help="if not None, will use this token in between prompt codec tokens and target codec tokens")
+    parser.add_argument("--neighbor_prompt_prob", type=float, default=0, help="the probability of using the prompt from the neighbor")
+    parser.add_argument("--neighbor_folder_name", type=str, default='neighbors',help="folder where the neighbors of the current audio files are stored, each row contains three tab separated entries: neighbor_fn, neighbor_temporal_distance, neighbor_duration")
+    parser.add_argument("--alignment_folder_name", type=str, default='alignment', help="folder where the forced alignment of the current audio files are stored, in csv format, each row contains five comma separated entries: begin, end, label, type, speaker, the first row is header")
+    parser.add_argument("--ipa_alignment_folder_name", type=str, default='ipa_alignment', help="folder where the forced alignment of the current audio files are stored, in txt format, each row contains three tab separated entries: begin, end, ipa phn sequence, generated using data/ll60k_preprocessing/step7_ipa_alignment.py")
+    parser.add_argument("--max_prompt_len", type=float, default=30, help="in sec., maximal prompt length selected from some neighboring file")
+    parser.add_argument("--min_prompt_len", type=float, default=0.5, help="in sec., minimal prompt length selected from some neighboring file")
+    parser.add_argument("--neighbor_selection_method", type=str, default="maxdist_60", help="maxdist_60 means uniformly select a neighbor that's within 60 sec of the current audio file")
+    parser.add_argument("--num_trial", type=int, default=5, help="number of tries to select a neighbor")
+    parser.add_argument("--prompt_start_from_begining_prob", type=float, default=0.5, help="the probability of starting the prompt from the beginning of the neighbor")
+    parser.add_argument("--min_alignment_len", type=int, default=5, help="in number of words")
+    parser.add_argument("--audio_folder_name", type=str, default='audio', help="folder where the audio files are stored")
+    
+    # rope parameters
+    parser.add_argument("--decoder_regular_rope", type=int, default=0, help="if 1, will use regular rope for the decoder (note that we always use regular rope for encoder). ")
+    parser.add_argument("--progress_no_multiple", type=int, default=0, help="if 1, will not multiple the percentage progress by the length of the key, see apply_rotary_pos_emb in models/modules/activation.py, this applies to both rope and sinusoidal positional encoding. Note that progress scale is still applied, i.e. when we only apply progress scale, but not multiple, the scaling factor is constant for every sample, rather than sample dependent")
+    parser.add_argument("--add_eos_to_text", type=int, default=0, help="if not 0, use this number as eos and add to the end of text token, usually use the second to last token in the vocab size")
+    parser.add_argument("--add_bos_to_text", type=int, default=0, help="if not 0, use this number as bos and add to the begining of text token, usually use the third to last token in the vocab size")
+    parser.add_argument("--use_sinusoidal", type=int, default=0, help="if 1, will use sinusoidal positional encoding, otherwise use rope. BUT if rope_base is None, will use sinusoidal")
+    parser.add_argument("--sinusoidal_base", type=int, default=1e4, help="the base of the exponential function, default is 1e4")
+    parser.add_argument("--use_sinusoidal_progress", type=int, default=0, help="if 1, will use sinusoidal positional encoding for progress, otherwise use rope")
+    parser.add_argument("--rope_base", type=int, default=None, help="the base of the exponential function, default is 1e4, if None, will not use rope")
+    parser.add_argument("--multiple_key_length", type=int, default=0, help="if 1, during progress calculation, will multiple the precentage progress by the length of the key, otherwise multiple with length of query. see models/rope_playground.ipynb")
+    parser.add_argument("--progress_scale", type=float, default=1.0, help="scale the progress, the smaller the value, the bigger the diagonal in attention score, see models/rope_playground.ipynb")
+
+    # attention alignment loss
+    parser.add_argument("--attention_alignment_loss", type=float, default=0.0, help="the weight on the attention alignment loss, if 0, will not calculate the loss")
+    parser.add_argument("--alignment_loss_layer", type=str, default="['0-1', '2', '3']", help='the layers to calculate the alignment loss, e.g. ["0-1", "2", "3"]')
+    parser.add_argument("--alignment_loss_head", type=str, default="['0-1', '2', '3']", help='the attention heads to calculate the alignment loss, e.g. ["0-1", "2", "3"]')
+    parser.add_argument("--alignment_blank_logit", type=float, default=-1.0, help="the logit for the blank token added to the attention weights")
+
+    # inference parameters
+    parser.add_argument("--metrics", type=str, default="['spk_sim','wer','mcd','pitch','energy','pesq','utmos']")
+    parser.add_argument("--res_jsonl_root", type=str, default="/home/pyp/BoostedVoiceEditor/res")
+    parser.add_argument("--res_name", type=str, default="2jan25.jsonl")
+    parser.add_argument("--inference_seed", type=int, default=1)
+    parser.add_argument("--codec_audio_sr", type=int, default=16000)
+    parser.add_argument("--codec_sr", type=float, default=50)
+    parser.add_argument("--top_k", type=int, default=0)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--temperature", type=float, default=1)
+    parser.add_argument("--silence_tokens", type=list, default=[])
+    parser.add_argument("--kvcache", type=int, default=0)
+    parser.add_argument("--stop_repetition", type=int, default=3)
+    parser.add_argument("--sample_batch_size", type=int, default=1)
+    parser.add_argument("--inference_manifest_fns", type=str, default="['/home/pyp/BoostedVoiceEditor/manifests/debug.jsonl']")
+    parser.add_argument("--use_gt_duration", type=int, default=1)
+    parser.add_argument("--save_root", type=str, default="/data/scratch/pyp/exp_pyp/BoostedVoiceEditor/gens")
+    parser.add_argument("--encodec_signature", type=str, default="/data/scratch/pyp/exp_pyp/audiocraft/encodec/xps/6f79c6a8/checkpoint.th")  
+    parser.add_argument("--extra_cutoff", type=float, default=5, help="in rare cases where the model doesn't follow specified target duration (only happened in extrapolation cases), we will terminate generation once the extra duration exceeds this value")
+    parser.add_argument("--duration_margin", type=float, default=0.04, help="used along with extra_cutoff, when extra_cutoff is used (i.e. model doesn't follow specified target_duration), we terminate the generate, and cut the results to target_duration + duration_margin")
+    # add repeat_prompt and asr_model_name
+    parser.add_argument("--repeat_prompt", type=int_or_str, default=0, help="if 1, will repeat the prompt for each segment")
+    parser.add_argument("--asr_model_name", type=str, default="w2v2", help="the name of the asr model, if not None, will use the asr model to generate the prompt")
+
+    # depth transformer parameters
+    parser.add_argument("--depth_dec_num_layers", type=int, default=0)
+    parser.add_argument("--depth_dec_d_model", type=int, default=768)
+    parser.add_argument("--depth_dec_nhead", type=int, default=12)
+    parser.add_argument("--moshi_depth", type=int, default=0, help="if 1, will use the same parameterization as moshi, i.e. temporal trm output will gets added to every transformed token embedding")
+
+    parser.add_argument("--validation_sample_cap", type=int, default=None, help="cap the validation data to this number")
+    parser.add_argument("--no_libri_in_training", type=int, default=None, help="if 1, will not use librilight in training, only use in validation")
+    parser.add_argument("--uniform_weight_start_step", type=int, default=1e50, help="set all codebook weight to be uniform starting from this step")
+
+    return parser

copy_codebase.py

@@ -0,0 +1,56 @@
+
+import os
+import shutil
+import fnmatch
+
+def parse_gitignore(gitignore_path):
+    """Parse a .gitignore file and return a list of patterns."""
+    patterns = []
+    with open(gitignore_path, "r") as f:
+        for line in f:
+            # Ignore comments and blank lines
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+            # Handle wildcards and directory separators
+            patterns.append(line)
+    return patterns
+
+def file_matches_patterns(file_path, patterns):
+    """Check if a file matches any of the patterns in .gitignore."""
+    for pattern in patterns:
+        if fnmatch.fnmatch(file_path, pattern):
+            return True
+    return False
+
+def copy_codebase(src, dst, max_size_mb=5, gitignore_path=None):
+    """ Copy files from src to dst, skipping files larger than max_size_mb and matching .gitignore patterns. """
+    if gitignore_path and os.path.exists(gitignore_path):
+        patterns = parse_gitignore(gitignore_path)
+    else:
+        patterns = []
+    print("patterns to ignore: ", patterns)
+    os.makedirs(dst, exist_ok=True)
+    for root, dirs, files in os.walk(src):
+        for file in files:
+            file_path = os.path.join(root, file)
+            relative_path = os.path.relpath(file_path, src)
+            dst_path = os.path.join(dst, relative_path)
+            # ignore .git because of permission issues
+            if "/.git/" in file_path:
+                continue
+
+            # Check .gitignore patterns
+            if file_matches_patterns(file_path, patterns):
+                # print(f"Skipping {file_path} because it matches a pattern in .gitignore")
+                continue
+
+            # Check file size
+            if os.path.getsize(file_path) > max_size_mb * 1024 * 1024:
+                print(f"Skipping {file_path} because it's larger than {max_size_mb}MB")
+                continue
+
+
+            # Make sure the destination directory exists
+            os.makedirs(os.path.dirname(dst_path), exist_ok=True)
+            shutil.copy(file_path, dst_path)

data/combined_dataset.py

@@ -0,0 +1,466 @@
+import os
+import ffmpeg
+import torch
+import random
+import copy
+import logging
+import torch.distributed as dist
+import shutil
+import csv
+import torchaudio
+import glob
+import numpy as np
+from data.tokenizer import TextTokenizer, tokenize_text, AudioTokenizer
+def find_files(root_dir, endswith=".wav"):
+    files = []
+    # os.walk generates the file names in a directory tree
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        for filename in filenames:
+            # os.path.splitext splits the file name into a base and extension
+            # base, ext = os.path.splitext(filename)
+            if filename.lower().endswith(endswith):
+                # os.path.join combines one or more path names into a single path
+                full_path = os.path.join(dirpath, filename)
+                files.append(full_path)
+    return files
+
+class dataset(torch.utils.data.Dataset):
+    def __init__(self, args, split):
+        super().__init__()
+        self.args = args
+        self.args.target_time_stretch_prob = getattr(self.args, "target_time_stretch_prob", 0)
+        self.args.target_time_stretch_bound = getattr(self.args, "target_time_stretch_bound", 0.1)
+        self.split = split
+        
+        assert self.split in ['train', 'valid', 'test'], f"split should be one of ['train', 'valid', 'test'], but it's {split}"
+        
+        if "[" not in self.args.dataset_dir or "]" not in self.args.dataset_dir:
+            self.dataset_dir = f"['{self.args.dataset_dir}']"
+        else:
+            self.dataset_dir = copy.deepcopy(self.args.dataset_dir)
+        self.dataset_dir = eval(self.dataset_dir)
+        data = []
+        if "[" not in self.args.manifest_name or "]" not in self.args.manifest_name:
+            self.args.manifest_name = f"['{self.args.manifest_name}']"
+        else:
+            self.args.manifest_name = copy.deepcopy(self.args.manifest_name)
+        self.manifest_name = eval(self.args.manifest_name)
+        if len(self.manifest_name) != len(self.dataset_dir):
+            assert len(self.manifest_name) == 1, f"len(self.manifest_name) should be 1 or equal to len(self.dataset_dir), but it's {len(self.manifest_name)}"
+            self.manifest_name = self.manifest_name * len(self.dataset_dir)
+        for i_data, dataset_dir in enumerate(self.dataset_dir):
+            if getattr(self.args, "no_libri_in_training", None) != None and ("librilight" in dataset_dir) and self.split == "train":
+                if not dist.is_initialized() or dist.get_rank() == 0:
+                    logging.info(f"skipping librilight in training split")
+                continue
+            n_datapoints = 0
+            manifest_fn = os.path.join(dataset_dir, self.manifest_name[i_data], self.split+".txt")
+            if not os.path.isfile(manifest_fn):
+                all_manifest_fn = glob.glob(manifest_fn.replace(".txt", "_*=*.txt"))
+                if len(all_manifest_fn) == 0:
+                    logging.info(f"no manifest file found for {split} split in {dataset_dir}")
+                    continue
+                if self.args.debug:
+                    logging.info(f"debugging mode, only using the frist found manifest file: {all_manifest_fn[0]}")
+                    all_manifest_fn = all_manifest_fn[:1]
+                else:
+                    if dist.is_initialized() and dist.get_rank() == 0:
+                        logging.info(f"Combining found manifest files for {split}: {all_manifest_fn}")
+                for cur_manifest_fn in all_manifest_fn:
+                    with open(cur_manifest_fn, "r") as rf:
+                        tmp = [l.strip().split("\t") + [i_data] for l in rf.readlines()] # i_data is the index of the dataset
+                        n_datapoints += len(tmp)
+                        data += tmp
+            else:
+                with open(manifest_fn, "r") as rf:
+                    tmp = [l.strip().split("\t") + [i_data] for l in rf.readlines()]
+                    data += tmp
+                    n_datapoints += len(tmp)
+            if dist.is_initialized() and dist.get_rank() == 0:
+                logging.info(f"number of data points for {split} split in {dataset_dir}: {n_datapoints}")
+        assert len(data) > 0, f"no data found for {split} split"
+        lengths_list = [int(item[1]) for item in data] # use 1 because there might be more than 1 columns (for gigaspeech we have 3 columns: path, duration, selfsim)
+        self.data = []
+        self.lengths_list = []
+        total_duration = 0
+        for d, l in zip(data, lengths_list):
+            if l >= self.args.encodec_sr*self.args.audio_min_length:
+                if self.args.drop_long and l > self.args.encodec_sr*self.args.audio_max_length:
+                    continue
+                self.data.append(d)
+                self.lengths_list.append(l)
+                total_duration += l / self.args.encodec_sr / 3600
+        # logging.info(f"for now cut the dataset to only have 500 examples for debugging")
+        # self.data = self.data[:1000]
+        # self.lengths_list = self.lengths_list[:1000]
+        if dist.is_initialized() and dist.get_rank() == 0:
+            logging.info(f"TOTAL number of data points for {self.split} split: {len(self.lengths_list)}")
+            logging.info(f"TOTAL duration for {self.split} split: {total_duration:.1f} hours")
+        # phoneme vocabulary
+        phn_set = set()
+        for dataset_dir in self.dataset_dir:
+            vocab_fn = os.path.join(dataset_dir, "vocab.txt")
+            with open(vocab_fn, "r") as f:
+                temp = [l.strip().split("\t") for l in f.readlines() if len(l) != 0]
+                phn_set.update([item[-1] for item in temp])
+        self.phn2num = {item:i for i, item in enumerate(phn_set)}
+        assert self.args.text_vocab_size > len(self.phn2num), f"need self.args.text_vocab_size to be bigger than number of phns in vocab to handle OOD phn, but the former is {self.args.text_vocab_size} while the latter is {len(self.phn2num)}"
+
+        if (self.args.neighbor_prompt_prob > 0 and self.args.time_stretch_prob > 0) or self.args.target_time_stretch_prob > 0:
+            userdir = os.path.expanduser("~")
+            encodec_signature = getattr(self.args, "encodec_signature", os.path.join(userdir, "VoiceStar", "pretrained", "encodec_6f79c6a8.th"))
+            self.audio_tokenizer = AudioTokenizer(signature=encodec_signature, device=torch.device("cpu"), encode_only=True)
+            assert self.audio_tokenizer.sample_rate == self.args.codec_audio_sr, f"audio_tokenizer.sample_rate: {self.audio_tokenizer.sample_rate}, self.args.encodec_sr: {self.args.encodec_sr}"
+            if dist.is_initialized() and dist.get_rank() == 0:
+                logging.info(f"rank: {dist.get_rank()}, audio_tokenizer device: {self.audio_tokenizer._device}")
+    
+    def __len__(self):
+        return len(self.lengths_list)
+    
+    def _load_phn_enc(self, index):
+        item = self.data[index]
+        dataset_dir = self.dataset_dir[item[-1]]
+        pf = os.path.join(dataset_dir, self.args.phn_folder_name, item[0]+".txt")
+        ef = os.path.join(dataset_dir, self.args.encodec_folder_name, item[0]+".txt")
+        # with certain probability, we load the audio, and time stretch it, note that we should not hit self.args.audio_max_length
+        if "/librilight" in dataset_dir:
+            audio_ext = ".flac"
+        elif "/emilia" in dataset_dir:
+            audio_ext = ".mp3"
+        else:
+            raise NotImplementedError(f"dataset_dir: {dataset_dir}")
+        
+        audio_fn = os.path.join(dataset_dir, self.args.audio_folder_name, item[0].replace(".txt", "")+audio_ext)
+        speed_factor = random.uniform(-self.args.target_time_stretch_bound, self.args.target_time_stretch_bound) + 1
+        length_ok = (float(item[1]) / self.args.encodec_sr) / speed_factor < self.args.audio_max_length # NOTE to calculate the maximal duration after time stretching, we should be used as orig/(1-bound), rather than orig*(1+bound)
+        if self.args.target_time_stretch_prob > 0 and random.random() < self.args.target_time_stretch_prob and os.path.isfile(audio_fn) and length_ok: 
+            try:
+                with open(pf, "r") as p:
+                    phns = [l.strip() for l in p.readlines()]
+                    assert len(phns) == 1, phns
+                    all_phns = phns[0].split(" ")
+                    x = [self.phn2num[item] for item in all_phns if item in self.phn2num]
+            except:
+                logging.info(f"loading failed for {pf}, maybe files don't exist or are corrupted")
+                return [], [[]], dataset_dir, audio_ext
+            # time stretch
+            try:
+                process = (
+                    ffmpeg.input(audio_fn, ss=0, t=float(item[1]) / self.args.encodec_sr)
+                    .output('pipe:1', format='f32le', ac=1, ar=self.audio_tokenizer.sample_rate, filter='atempo={}'.format(speed_factor))
+                    .run_async(pipe_stdout=True, pipe_stderr=True)
+                )
+                # Read the processed audio from ffmpeg stdout
+                output, _ = process.communicate()
+
+                # Convert the output to a numpy array
+                output_np = np.frombuffer(output, dtype=np.float32).copy()
+
+                # Reshape the numpy array back to the expected shape (1, samples for mono)
+                waveform = torch.from_numpy(output_np)
+                waveform = waveform.unsqueeze(0).unsqueeze(0)
+                assert waveform.ndim == 3 and waveform.shape[0] == 1 and waveform.shape[1] == 1, waveform.shape
+                with torch.no_grad():
+                    encos = self.audio_tokenizer.encode(waveform.to(self.audio_tokenizer._device))
+                assert encos.shape[1] == self.args.n_codebooks, f"encos.shape: {encos.shape}"
+                encos = encos.cpu().squeeze(0).numpy().tolist() # [K, T]
+                if self.args.special_first:
+                    raise NotImplementedError
+                    # y = [[int(n)+self.args.n_special for n in l] for l in encos]
+                else:
+                    y = [[int(n) for n in l] for l in encos]
+                return x, y, dataset_dir, audio_ext
+            except Exception as e:
+                logging.info(f"failed with time stretch and codec encode for {audio_fn}")
+                logging.info(f"error: {e}")
+                pass
+
+        try:
+            with open(pf, "r") as p, open(ef, "r") as e:
+                phns = [l.strip() for l in p.readlines()]
+                assert len(phns) == 1, phns
+                all_phns = phns[0].split(" ")
+                x = [self.phn2num[item] for item in all_phns if item in self.phn2num] # we assume that OOD will not happen, because phn vocab is small
+                encos = [l.strip().split() for k, l in enumerate(e.readlines()) if k < self.args.n_codebooks]
+                
+                assert len(encos) == self.args.n_codebooks, ef
+
+                if self.args.special_first:
+                    raise NotImplementedError
+                    # y = [[int(n)+self.args.n_special for n in l] for l in encos]
+                else:
+                    y = [[int(n) for n in l] for l in encos]
+        except:
+            logging.info(f"loading failed for {pf} and {ef}, maybe files don't exist or are corrupted")
+            return [], [[]], dataset_dir, audio_ext
+
+        return x, y, dataset_dir, audio_ext
+
+    # this uses the output of step7_ipa_alignment.py
+    def find_neighbor(self, neighbors, y_len, dataset_dir, audio_ext):
+        neighbor = random.choice(neighbors)
+        neighbor_enc_fn = os.path.join(dataset_dir, self.args.encodec_folder_name, neighbor[0])
+        if not os.path.isfile(neighbor_enc_fn):
+            return None, None
+        neighbor_audio_path = os.path.join(dataset_dir, self.args.audio_folder_name, neighbor[0].replace(".txt", audio_ext))
+        if getattr(self.args, "time_stretch_prob", 0) > 0 and not os.path.isfile(neighbor_audio_path):
+            logging.info(f"audio file not found: {neighbor_audio_path}")
+            return None, None
+        if random.random() < getattr(self.args, "time_stretch_prob", 0):
+            time_stretch_flag = True
+            speed_factor = random.uniform(-self.args.time_stretch_bound, self.args.time_stretch_bound) + 1
+            duration_factor = 1 / speed_factor
+        else:
+            time_stretch_flag = False
+            duration_factor = 1
+
+        ####################### TODO for now always use the entire neighbor for emilia
+        ####################### TODO for now always use the entire neighbor for emilia
+        # if it's gigaspeech or emilia, we did not run MFA forced alignment, and therefore no ipa alignment, and will just use the entire neighbor as the prompt
+
+        if "/emilia" in dataset_dir:
+            # get neighbor duration
+            neighbor_dur = float(neighbor[2])
+            if neighbor_dur * duration_factor + y_len / self.args.encodec_sr > self.args.audio_max_length or neighbor_dur * duration_factor < self.args.min_prompt_len:
+                return None, None
+            try:
+                neighbor_pf = os.path.join(dataset_dir, self.args.phn_folder_name, neighbor[0])
+                with open(neighbor_pf, "r") as p:
+                    phns = [l.strip() for l in p.readlines()]
+                    assert len(phns) == 1, phns
+                    all_phns = phns[0].split(" ")
+                    phn_token = [self.phn2num[item] for item in all_phns if item in self.phn2num]
+            except:
+                logging.info(f"loading failed for {neighbor_pf}, maybe files don't exist")
+                return None, None
+            # if do not stretch the audio
+            if not time_stretch_flag:
+                with open(neighbor_enc_fn, "r") as f:
+                    neighbor_enc = [l.strip().split() for l in f.readlines()]
+                if len(neighbor_enc) != self.args.n_codebooks:
+                    return None, None
+                # if too long 
+                else:
+                    if self.args.special_first:
+                        raise NotImplementedError
+                        # neighbor_enc = [[int(n)+self.args.n_special for n in l] for l in neighbor_enc]
+                    else:
+                        neighbor_enc = [[int(n) for n in l] for l in neighbor_enc]
+                    
+                    return phn_token, neighbor_enc
+            else: # stretch the audio with ffmpeg-python
+                process = (
+                    ffmpeg.input(neighbor_audio_path, ss=0, t=neighbor_dur)
+                    .output('pipe:1', format='f32le', ac=1, ar=self.audio_tokenizer.sample_rate, filter='atempo={}'.format(speed_factor))
+                    .run_async(pipe_stdout=True, pipe_stderr=True)
+                    )
+                # Read the processed audio from ffmpeg stdout
+                output, _ = process.communicate()
+
+                # Convert the output to a numpy array
+                output_np = np.frombuffer(output, dtype=np.float32).copy()
+
+                # Reshape the numpy array back to the expected shape (1, samples for mono)
+                waveform = torch.from_numpy(output_np)
+                waveform = waveform.unsqueeze(0).unsqueeze(0)
+                assert waveform.ndim == 3 and waveform.shape[0] == 1 and waveform.shape[1] == 1, waveform.shape
+                with torch.no_grad():
+                    encos = self.audio_tokenizer.encode(waveform.to(self.audio_tokenizer._device))
+                assert encos.shape[1] == self.args.n_codebooks, f"encos.shape: {encos.shape}"
+                neighbor_enc = encos.cpu().squeeze(0).numpy().tolist() # [K, T]
+                return phn_token, neighbor_enc
+        ####################### TODO for now always use the entire neighbor for emilia
+        ####################### TODO for now always use the entire neighbor for emilia
+        ipa_alignment_fn = os.path.join(dataset_dir, self.args.ipa_alignment_folder_name, neighbor[0])
+        if not os.path.isfile(ipa_alignment_fn):
+            # print(f"file not found: {ipa_alignment_fn}", flush=True)
+            return None, None
+        with open(ipa_alignment_fn, "r") as f:
+            alignments = [l.strip().split("\t") for l in f.readlines()]
+        alignments = [[float(l[0]), float(l[1]), l[2]] for l in alignments if len(l) == 3]
+        alignments = [l for l in alignments if self.args.min_prompt_len < (l[1] - l[0]) * duration_factor < self.args.max_prompt_len]
+        if len(alignments) == 0:
+            # print(f"no valid alignment found for {ipa_alignment_fn}")
+            return None, None
+        idx = random.choice(range(len(alignments)))
+        while (alignments[idx][1] - alignments[idx][0]) * duration_factor + y_len / self.args.encodec_sr > self.args.audio_max_length:
+            idx -= 1
+            if idx < 0:
+                # print(f"too long combined with y_len {ipa_alignment_fn=}, and {y_len=}")
+                return None, None
+        if (alignments[idx][1] - alignments[idx][0]) * duration_factor < self.args.min_prompt_len:
+            return None, None
+        
+        start_time, end_time = alignments[idx][:2]
+        phn = alignments[idx][2].split(" ")
+        phn_token = [self.phn2num[item] for item in phn if item in self.phn2num]
+        if len(phn_token) == 0:
+            return None, None
+
+        if time_stretch_flag:
+            duration = end_time - start_time
+            process = (
+                ffmpeg.input(neighbor_audio_path, ss=start_time, t=duration)
+                .output('pipe:1', format='f32le', ac=1, ar=self.audio_tokenizer.sample_rate, filter='atempo={}'.format(speed_factor))
+                .run_async(pipe_stdout=True, pipe_stderr=True)
+            )
+            # Read the processed audio from ffmpeg stdout
+            output, _ = process.communicate()
+
+            # Convert the output to a numpy array
+            output_np = np.frombuffer(output, dtype=np.float32).copy()
+
+            # Reshape the numpy array back to the expected shape (1, samples for mono)
+            waveform = torch.from_numpy(output_np)
+            waveform = waveform.unsqueeze(0).unsqueeze(0)
+            assert waveform.ndim == 3 and waveform.shape[0] == 1 and waveform.shape[1] == 1, waveform.shape
+            try:
+                with torch.no_grad():
+                    encos = self.audio_tokenizer.encode(waveform.to(self.audio_tokenizer._device))
+            except:
+                logging.info(f"failed with time stretch for {neighbor_audio_path}, from {start_time} to {end_time} with duration factor {duration_factor}, which leads to {duration*duration_factor} seconds")
+                return None, None
+            assert encos.shape[1] == self.args.n_codebooks, f"encos.shape: {encos.shape}"
+            neighbor_enc = encos.cpu().squeeze(0).numpy().tolist() # [K, T]
+            return phn_token, neighbor_enc
+        else:
+            # get encodec codes from storage
+            with open(neighbor_enc_fn, "r") as f:
+                neighbor_enc = [l.strip().split() for l in f.readlines()]
+                if len(neighbor_enc) != self.args.n_codebooks:
+                    # print(f"wrong number of codebooks for {neighbor_enc_fn}")
+                    return None, None
+                else:
+                    # trim the encodec codes to the segment
+                    start_enc_frame = int(start_time * self.args.encodec_sr)
+                    end_enc_frame = int(end_time * self.args.encodec_sr)
+                    neighbor_enc = [l[start_enc_frame:end_enc_frame] for l in neighbor_enc]
+                    if len(neighbor_enc[0]) == 0:
+                        # print(f"no valid encodec codes found for {neighbor_enc_fn}")
+                        return None, None
+                    if self.args.special_first:
+                        raise NotImplementedError
+                    else:
+                        neighbor_enc = [[int(n) for n in l] for l in neighbor_enc]
+                    return phn_token, neighbor_enc
+
+    def __getitem__(self, index):
+        x, y, dataset_dir, audio_ext = self._load_phn_enc(index)
+        x_len, y_len = len(x), len(y[0])
+        extra_ret = {'x_sep_token_position': 0, 'y_sep_token_position': 0}
+        if x_len == 0 or y_len == 0:
+            ret = {
+                "x": None, 
+                "x_len": None, 
+                "y": None, 
+                "y_len": None,
+                }
+            ret.update(extra_ret)
+            return ret
+        while y_len < self.args.encodec_sr*self.args.audio_min_length:
+            assert not self.args.dynamic_batching
+            index = random.choice(range(len(self))) # regenerate an index
+            x, y, dataset_dir, audio_ext = self._load_phn_enc(index)
+            x_len, y_len = len(x), len(y[0])
+        
+        # if use neighbor prompt
+        x_neighbor, y_neighbor = None, None
+        use_neighbor_prob = random.random()
+        neighbor_fn = os.path.join(dataset_dir, self.args.neighbor_folder_name, self.data[index][0]+".txt")
+        if self.args.neighbor_prompt_prob > 0 and use_neighbor_prob < self.args.neighbor_prompt_prob and os.path.isfile(neighbor_fn): # it might not exist, just because we didn't find neighbor for this file (other than itself, which is common for emilia)
+            with open(neighbor_fn, "r") as f:
+                neighbors = [l.strip().split("\t") for l in f.readlines()]
+            # select neighbors
+            if "maxdist" in self.args.neighbor_selection_method:
+                maxdist = int(self.args.neighbor_selection_method.split("_")[-1])
+                # only keep neighbors with distance within maxdist
+                neighbors = [n for n in neighbors if float(n[1]) <= maxdist]
+            else:
+                raise NotImplementedError
+            x_neighbor, y_neighbor = None, None
+            if len(neighbors) > 0:
+                x_neighbor, y_neighbor = self.find_neighbor(neighbors, y_len, dataset_dir, audio_ext)
+                i_trial = 0
+                while x_neighbor is None and i_trial < self.args.num_trial and i_trial < len(neighbors):
+                    x_neighbor, y_neighbor = self.find_neighbor(neighbors, y_len, dataset_dir, audio_ext)
+                    i_trial += 1
+        
+        if x_neighbor != None:
+            if self.args.x_sep_token != None:
+                x = x_neighbor + [self.args.x_sep_token] + x
+            else:
+                x = x_neighbor + x
+            if self.args.y_sep_token != None:
+                y = [y_neighbor[i] + [self.args.y_sep_token] + y[i] for i in range(len(y))]
+            else:
+                y = [y_neighbor[i] + y[i] for i in range(len(y))]
+            extra_ret['y_sep_token_position'] = len(y_neighbor[0]) + 1 # if using y_sep_token, this is actually the position of the token right before the y_sep_token, but since y_sep_token is ignored in loss computation, it's fine that we use the position of the token right before it
+            extra_ret['x_sep_token_position'] = len(x_neighbor) + 1
+            x_len, y_len = len(x), len(y[0])
+
+
+        # consider adding eos to the end of the text
+        if self.args.add_eos_to_text != 0:
+            x.append(self.args.add_eos_to_text)
+            x_len += 1
+        if getattr(self.args, "add_bos_to_text", 0) != 0:
+            x = [self.args.add_bos_to_text] + x
+            x_len += 1
+        ### padding and cropping ###
+        ### padding and cropping ###
+        # adjust the length of encodec codes, pad to max_len or randomly crop
+        orig_y_len = copy.copy(y_len)
+        max_len = int(self.args.audio_max_length * self.args.encodec_sr)
+        if y_len > max_len + 10: # give it some margin for rounding error
+            raise RuntimeError(f"audio is too long, {y_len=}, {max_len=}")
+        else:
+            audio_start = 0
+            if not self.args.dynamic_batching:
+                pad = [0] * (max_len - y_len) if self.args.sep_special_token else [self.args.audio_pad_token] * (max_len - y_len)
+                for i in range(len(y)):
+                    y[i] = y[i] + pad
+        
+        if self.args.pad_x and x_len <= self.args.text_max_length:
+            pad = [0] * (self.args.text_max_length - x_len) if self.args.sep_special_token else [self.args.text_pad_token] * (self.args.text_max_length - x_len)
+            x = x + pad
+        
+        ret = {
+            "x": torch.LongTensor(x), 
+            "x_len": x_len, 
+            "y": torch.LongTensor(y), 
+            "y_len": y_len,
+            }
+        ret.update(extra_ret)
+
+        return ret
+            
+
+    def collate(self, batch):
+        # make sure keys in every batch is the same
+        for batch1, batch2 in zip(batch[:-1], batch[1:]):
+            assert set(batch1.keys()) == set(batch2.keys()), f"keys in batch1: {batch1.keys()} and keys in batch2: {batch2.keys()} are different"
+        out = {key:[] for key in batch[0]}
+        for item in batch:
+            if item['x'] == None: # deal with load failure
+                continue
+            for key, val in item.items():
+                out[key].append(val)
+        res = {}
+        if self.args.pad_x:
+            res["x"] = torch.stack(out["x"], dim=0)
+        else:
+            res["x"] = torch.nn.utils.rnn.pad_sequence(out["x"], batch_first=True, padding_value=self.args.text_pad_token)
+        res["x_lens"] = torch.LongTensor(out["x_len"])
+        if self.args.dynamic_batching:
+            res['y'] = torch.nn.utils.rnn.pad_sequence([item.transpose(1,0) for item in out['y']],padding_value=self.args.audio_pad_token)
+            res['y'] = res['y'].permute(1,2,0) # T B K -> B K T
+        else:
+            res['y'] = torch.stack(out['y'], dim=0)
+        res["y_lens"] = torch.LongTensor(out["y_len"])
+        res["text_padding_mask"] = torch.arange(res['x'][0].shape[-1]).unsqueeze(0) >= res['x_lens'].unsqueeze(1)
+        res["audio_padding_mask"] = torch.arange(res['y'][0].shape[-1]).unsqueeze(0) >= res['y_lens'].unsqueeze(1)
+        if "y_sep_token_position" in out:
+            res["y_sep_token_position"] = torch.LongTensor(out["y_sep_token_position"])
+        if "x_sep_token_position" in out:
+            res["x_sep_token_position"] = torch.LongTensor(out["x_sep_token_position"])
+        return res

data/emilia_preprocessing/delete_tar_files.sh

@@ -0,0 +1,42 @@
+#!/bin/bash
+
+# Define the root directory where the tar files are located
+root=${root:-/data/scratch/pyp/datasets/emilia/downloads} # Example: /data/scratch/pyp/datasets/emilia/downloads
+exist_log_file="file_log_debug.txt" # Log of files to delete
+delete_log="deleted_files.log" # Log of successfully deleted files
+error_log="delete_errors.log"  # Log of errors (e.g., missing files)
+
+# Clear previous logs
+> "$delete_log"
+> "$error_log"
+
+echo "Starting deletion of tar files listed in $exist_log_file..."
+
+# Loop through each line in exist_log_file
+while IFS=',' read -r filename size local_sha256 original_filename url; do
+    # Trim leading/trailing whitespace
+    original_filename=$(echo "$original_filename" | xargs)
+
+    # Construct the full path to the tar file
+    tar_file="${root}/${original_filename}"
+
+    # Check if the tar file exists
+    if [ -f "$tar_file" ]; then
+        # Attempt to delete the file
+        if rm -f "$tar_file"; then
+            echo "✅ Deleted: $tar_file"
+            echo "$tar_file" >> "$delete_log"
+        else
+            echo "❌ Failed to delete: $tar_file"
+            echo "$tar_file" >> "$error_log"
+        fi
+    else
+        # Log missing files
+        echo "❌ File not found: $tar_file"
+        echo "$tar_file" >> "$error_log"
+    fi
+done < "$exist_log_file"
+
+echo "Deletion process completed."
+echo "Deleted files are logged in $delete_log."
+echo "Errors (if any) are logged in $error_log."

data/emilia_preprocessing/encodec.py

@@ -0,0 +1,1554 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compression models or wrapper around existing models.
+Also defines the main interface that a model must follow to be usable as an audio tokenizer.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+import logging
+import math
+from pathlib import Path
+import typing as tp
+
+import numpy as np
+import torch
+from torch import nn
+from torch import einsum
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+import logging
+import warnings
+from einops import rearrange, repeat
+import omegaconf
+# import flashy
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                 'time_group_norm'])
+
+def dict_from_config(cfg: omegaconf.DictConfig) -> dict:
+    """Convenience function to map an omegaconf configuration to a dictionary.
+
+    Args:
+        cfg (omegaconf.DictConfig): Original configuration to map to dict.
+    Returns:
+        dict: Config as dictionary object.
+    """
+    dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)
+    assert isinstance(dct, dict)
+    return dct
+
+@dataclass
+class QuantizedResult:
+    x: torch.Tensor
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+class BaseQuantizer(nn.Module):
+    """Base class for quantizers.
+    """
+
+    def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:
+        """
+        Given input tensor x, returns first the quantized (or approximately quantized)
+        representation along with quantized codes, bandwidth, and any penalty term for the loss.
+        Finally, this returns a dict of metrics to update logging etc.
+        Frame rate must be passed so that the bandwidth is properly computed.
+        """
+        raise NotImplementedError()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth."""
+        raise NotImplementedError()
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        raise NotImplementedError()
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def num_codebooks(self):
+        """Number of active codebooks."""
+        raise NotImplementedError()
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise NotImplementedError()
+
+class CompressionModel(ABC, nn.Module):
+    """Base API for all compression model that aim at being used as audio tokenizers
+    with a language model.
+    """
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> QuantizedResult:
+        ...
+
+    @abstractmethod
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """See `EncodecModel.encode`."""
+        ...
+
+    @abstractmethod
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """See `EncodecModel.decode`."""
+        ...
+
+    @abstractmethod
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        ...
+
+    @property
+    @abstractmethod
+    def channels(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def frame_rate(self) -> float:
+        ...
+
+    @property
+    @abstractmethod
+    def sample_rate(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def cardinality(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def num_codebooks(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def total_codebooks(self) -> int:
+        ...
+
+    @abstractmethod
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        ...
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none'):
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs):
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                 padding_total: int = 0) -> int:
+    """See `pad_for_conv1d`."""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+    """Wrapper around ConvTranspose2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class StreamableConv1d(nn.Module):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, dilation: int = 1,
+                 groups: int = 1, bias: bool = True, causal: bool = False,
+                 norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 pad_mode: str = 'reflect'):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn("StreamableConv1d has been initialized with stride > 1 and dilation > 1"
+                          f" (kernel_size={kernel_size} stride={stride}, dilation={dilation}).")
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                               dilation=dilation, groups=groups, bias=bias, causal=causal,
+                               norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (kernel_size - 1) * dilation + 1  # effective kernel size with dilations
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        return self.conv(x)
+
+
+class StreamableConvTranspose1d(nn.Module):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, causal: bool = False,
+                 norm: str = 'none', trim_right_ratio: float = 1.,
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                          causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+class StreamableLSTM(nn.Module):
+    """LSTM without worrying about the hidden state, nor the layout of the data.
+    Expects input as convolutional layout.
+    """
+    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+
+
+class SEANetResnetBlock(nn.Module):
+    """Residual block from SEANet model.
+
+    Args:
+        dim (int): Dimension of the input/output.
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection.
+    """
+    def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],
+                 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,
+                 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'
+        act = getattr(nn, activation)
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                act(**activation_params),
+                StreamableConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,
+                                 norm=norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+        self.block = nn.Sequential(*block)
+        self.shortcut: nn.Module
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = StreamableConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,
+                                             causal=causal, pad_mode=pad_mode)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x)
+
+
+class SEANetEncoder(nn.Module):
+    """SEANet encoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order. We use the decoder order as some models may only employ the decoder.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the encoder, it corresponds to the N first blocks.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = 1
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(channels, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= i + 2 else norm
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      norm=block_norm, norm_params=norm_params,
+                                      activation=activation, activation_params=activation_params,
+                                      causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConv1d(mult * n_filters, mult * n_filters * 2,
+                                 kernel_size=ratio * 2, stride=ratio,
+                                 norm=block_norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+            mult *= 2
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        model += [
+            act(**activation_params),
+            StreamableConv1d(mult * n_filters, dimension, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(nn.Module):
+    """SEANet decoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple.
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the decoder, it corresponds to the N last blocks.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0, trim_right_ratio: float = 1.0):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(dimension, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) else norm
+            # Add upsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConvTranspose1d(mult * n_filters, mult * n_filters // 2,
+                                          kernel_size=ratio * 2, stride=ratio,
+                                          norm=block_norm, norm_kwargs=norm_params,
+                                          causal=causal, trim_right_ratio=trim_right_ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      activation=activation, activation_params=activation_params,
+                                      norm=block_norm, norm_params=norm_params, causal=causal,
+                                      pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params),
+            StreamableConv1d(n_filters, channels, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None:
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [
+                final_act(**final_activation_params)
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, z):
+        y = self.model(z)
+        return y
+
+
+def exists(val: tp.Optional[tp.Any]) -> bool:
+    return val is not None
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if exists(val) else d
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay: float):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters: int, num_iters: int = 10):
+    dim, dtype = samples.shape[-1], samples.dtype
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        diffs = rearrange(samples, "n d -> n () d") - rearrange(
+            means, "c d -> () c d"
+        )
+        dists = -(diffs ** 2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+def orthogonal_loss_fn(t):
+    # eq (2) from https://arxiv.org/abs/2112.00384
+    n = t.shape[0]
+    normed_codes = l2norm(t)
+    identity = torch.eye(n, device=t.device)
+    cosine_sim = einsum("i d, j d -> i j", normed_codes, normed_codes)
+    return ((cosine_sim - identity) ** 2).sum() / (n ** 2)
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
+            If set to true, run the k-means algorithm on the first training batch and use
+            the learned centroids as initialization.
+        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        kmeans_init: int = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.kmeans_iters = kmeans_iters
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed.clone())
+        self.cluster_size.data.copy_(cluster_size)
+        self.inited.data.copy_(torch.Tensor([True]))
+        # Make sure all buffers across workers are in sync after initialization
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def replace_(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes)
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(
+            x.pow(2).sum(1, keepdim=True)
+            - 2 * x @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+        embed_ind = dist.max(dim=-1).indices
+        return embed_ind
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        return embed_ind
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x):
+        raise NotImplementedError()
+        shape, dtype = x.shape, x.dtype
+        x = self.preprocess(x)
+        self.init_embed_(x)
+
+        embed_ind = self.quantize(x)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        quantize = self.dequantize(embed_ind)
+
+        if self.training:
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x)
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = x.t() @ embed_onehot
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+
+        return quantize, embed_ind
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int):
+        channels_last (bool): Channels are the last dimension in the input tensors.
+        commitment_weight (float): Weight for commitment loss.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider
+            for orthogonal regularization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        channels_last: bool = False,
+        commitment_weight: float = 1.,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        _codebook_dim: int = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim
+        self.project_in = (nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity())
+        self.project_out = (nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity())
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+
+        self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,
+                                           kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,
+                                           decay=decay, epsilon=epsilon,
+                                           threshold_ema_dead_code=threshold_ema_dead_code)
+        self.codebook_size = codebook_size
+
+        self.channels_last = channels_last
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    @property
+    def inited(self):
+        return self._codebook.inited
+
+    def _preprocess(self, x):
+        if not self.channels_last:
+            x = rearrange(x, "b d n -> b n d")
+        return x
+
+    def _postprocess(self, quantize):
+        if not self.channels_last:
+            quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def encode(self, x):
+        x = self._preprocess(x)
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+        return quantize
+
+    def forward(self, x):
+        device = x.device
+        x = self._preprocess(x)
+
+        x = self.project_in(x)
+        quantize, embed_ind = self._codebook(x)
+
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
+        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
+
+        if self.training:
+            if self.commitment_weight > 0:
+                commit_loss = F.mse_loss(quantize.detach(), x)
+                loss = loss + commit_loss * self.commitment_weight
+
+            if self.orthogonal_reg_weight > 0:
+                codebook = self.codebook
+
+                if self.orthogonal_reg_active_codes_only:
+                    # only calculate orthogonal loss for the activated codes for this batch
+                    unique_code_ids = torch.unique(embed_ind)
+                    codebook = codebook[unique_code_ids]
+
+                num_codes = codebook.shape[0]
+                if exists(self.orthogonal_reg_max_codes) and num_codes > self.orthogonal_reg_max_codes:
+                    rand_ids = torch.randperm(num_codes, device=device)[:self.orthogonal_reg_max_codes]
+                    codebook = codebook[rand_ids]
+
+                orthogonal_reg_loss = orthogonal_loss_fn(codebook)
+                loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
+
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+
+        return quantize, embed_ind, loss
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        codebook_size = kwargs.pop('codebook_size', None)
+        if codebook_size is None:
+            raise ValueError("codebook_size must be provided in kwargs")
+        if type(codebook_size) != list:
+            codebook_size = [codebook_size] * num_quantizers
+        self.layers = nn.ModuleList(
+            [VectorQuantization(codebook_size=cur_codebook_size, **kwargs) for _,cur_codebook_size in zip(range(num_quantizers), codebook_size)]
+        )
+            
+
+        # self.layers = nn.ModuleList(
+        #     [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        # )
+
+    def forward(self, x, n_q: tp.Optional[int] = None):
+        quantized_out = 0.0
+        residual = x
+
+        all_losses = []
+        all_indices = []
+
+        n_q = n_q or len(self.layers)
+
+        for i, layer in enumerate(self.layers[:n_q]):
+            quantized, indices, loss = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+            all_indices.append(indices)
+            all_losses.append(loss)
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        return quantized_out, out_indices, out_losses
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        for layer in self.layers[:n_q]:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            # the original code is below
+            # since quantize has the gradient of residual, according to line 321
+            #  quantize = x + (quantize - x).detach()
+            # the code below will make commitment loss to be 0 for all codebooks except for codebook1
+            # https://github.com/facebookresearch/encodec/issues/25
+            # therefore we change it
+
+            residual = residual - quantized
+            # residual = residual - quantized.detach()
+            # since commitment loss is averaged, the scale of the loss won't get change (not as said in the issue above)
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out
+
+
+class ResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer.
+
+    Args:
+        dimension (int): Dimension of the codebooks.
+        n_q (int): Number of residual vector quantizers used.
+        q_dropout (bool): Random quantizer drop out at train time.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider.
+            for orthogonal regularization.
+    """
+    def __init__(
+        self,
+        dimension: int = 256,
+        n_q: int = 8,
+        q_dropout: bool = False,
+        bins: tp.Union[int, tp.List[int]] = 1024,
+        decay: float = 0.99,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.max_n_q = n_q
+        self.n_q = n_q
+        self.q_dropout = q_dropout
+        self.dimension = dimension
+        self.bins = bins
+        self.decay = decay
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            kmeans_init=self.kmeans_init,
+            kmeans_iters=self.kmeans_iters,
+            threshold_ema_dead_code=self.threshold_ema_dead_code,
+            orthogonal_reg_weight=self.orthogonal_reg_weight,
+            orthogonal_reg_active_codes_only=self.orthogonal_reg_active_codes_only,
+            orthogonal_reg_max_codes=self.orthogonal_reg_max_codes,
+            channels_last=False
+        )
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        n_q = self.n_q
+        if self.training and self.q_dropout:
+            n_q = int(torch.randint(1, self.n_q + 1, (1,)).item())
+        if type(self.bins) == list:
+            bins = self.bins
+        else:
+            bins = [self.bins] * self.n_q
+        bw_per_q = [math.log2(bin) * frame_rate / 1000 for bin in bins]
+        bw = torch.tensor(sum(bw_per_q)).to(x)
+        quantized, codes, commit_loss = self.vq(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        n_q = self.n_q
+        codes = self.vq.encode(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        codes = codes.transpose(0, 1)
+        quantized = self.vq.decode(codes)
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    def set_num_codebooks(self, n: int):
+        assert n > 0 and n <= self.max_n_q
+        self.n_q = n
+
+class DummyQuantizer(BaseQuantizer):
+    """Fake quantizer that actually does not perform any quantization.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        q = x.unsqueeze(1)
+        return QuantizedResult(x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return x.unsqueeze(1)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return codes.squeeze(1)
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        return 1
+
+    @property
+    def num_codebooks(self):
+        """Total number of codebooks."""
+        return self.total_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise AttributeError("Cannot override the number of codebooks for the dummy quantizer")
+
+
+class EncodecModel(CompressionModel):
+    """Encodec model operating on the raw waveform.
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        decoder (nn.Module): Decoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 decoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        return self.quantizer.decode(codes)
+
+class EncodecModel_encode_only(CompressionModel):
+    """Encodec model operating on the raw waveform. Encode only, so no decoder
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        raise NotImplementedError("Decode is not supported for encode only model")
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        raise NotImplementedError("Decode is not supported for encode only model")
+        return self.quantizer.decode(codes)
+
+def get_quantizer(quantizer: str, cfg: omegaconf.DictConfig, dimension: int) -> BaseQuantizer:
+    klass = {
+        'no_quant': DummyQuantizer,
+        'rvq': ResidualVectorQuantizer
+    }[quantizer]
+    kwargs = dict_from_config(getattr(cfg, quantizer))
+    if quantizer != 'no_quant':
+        kwargs['dimension'] = dimension
+    return klass(**kwargs)
+
+def get_encodec_autoencoder(encoder_name: str, cfg: omegaconf.DictConfig):
+    if encoder_name == 'seanet':
+        kwargs = dict_from_config(getattr(cfg, 'seanet'))
+        encoder_override_kwargs = kwargs.pop('encoder')
+        decoder_override_kwargs = kwargs.pop('decoder')
+        encoder_kwargs = {**kwargs, **encoder_override_kwargs}
+        decoder_kwargs = {**kwargs, **decoder_override_kwargs}
+        encoder = SEANetEncoder(**encoder_kwargs)
+        decoder = SEANetDecoder(**decoder_kwargs)
+        return encoder, decoder
+    else:
+        raise KeyError(f"Unexpected compression model {cfg.compression_model}")
+
+
+def get_compression_model(ckpt_fn, encode_only=False, device="cpu") -> CompressionModel:
+    """Instantiate a compression model."""
+    if device == None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    state = torch.load(ckpt_fn, map_location='cpu')
+    cfg = state['xp.cfg']
+    cfg.device = str(device)
+    weights = state['best_state']['model']
+    assert cfg.compression_model == 'encodec', "Only Encodec model is supported for now."
+    if encode_only:
+        all_keys = list(weights.keys())
+        for key in all_keys:
+            if key.startswith('decoder'):
+                del weights[key]
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, _ = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel_encode_only(encoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+    else:
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, decoder = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel(encoder, decoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+if __name__ == "__main__":
+    import torchaudio
+    ckpt_fn = "/home/pyp/BoostedVoiceEditor/pretrained/encodec_6f79c6a8.th"
+    audio_in_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam.wav", "/home/pyp/BoostedVoiceEditor/demo/ray.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean.wav", "/home/pyp/BoostedVoiceEditor/demo/bible.wav", "/home/pyp/BoostedVoiceEditor/demo/miley.wav"]
+    audio_out_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/ray_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/bible_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/miley_encodecTest.wav"]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = get_compression_model(ckpt_fn, device=device)
+
+    for audio_in_fn, audio_out_fn in zip(audio_in_fns, audio_out_fns):
+        audio_in, sr = torchaudio.load(audio_in_fn)
+        if sr != model.sample_rate:
+            audio_in = torchaudio.transforms.Resample(sr, model.sample_rate)(audio_in)
+        if audio_in.shape[0] == 2:
+            audio_in = audio_in.mean(dim=0, keepdim=True)
+        audio_in = audio_in.unsqueeze(0)
+        audio_in = audio_in.to(torch.float32).to(device)
+        codes = model.encode(audio_in)[0]
+        audio_out = model.decode(codes)[0].cpu()
+        torchaudio.save(audio_out_fn, audio_out, model.sample_rate)

data/emilia_preprocessing/sha256hash.py

@@ -0,0 +1,14 @@
+import hashlib
+import sys
+
+def sha256_hash_file(filename):
+    sha256_hash = hashlib.sha256()
+    with open(filename, "rb") as file:
+        # Read and update hash string in chunks to handle large files
+        for byte_block in iter(lambda: file.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest()
+
+# Usage example
+filename = sys.argv[1]
+print(sha256_hash_file(filename))

data/emilia_preprocessing/step1_download.py

@@ -0,0 +1,9 @@
+# conda activate emilia
+from datasets import load_dataset
+import fire
+def main(root: str="/data/scratch/pyp/datasets/emilia"):
+    path = "EN/*.tar.gz*"
+    dataset = load_dataset("amphion/Emilia-Dataset", data_files={"en": path}, split="en", streaming=False, revision="fc71e07e8572f5f3be1dbd02ed3172a4d298f152", cache_dir=root)
+
+if __name__ == "__main__":
+    fire.Fire(main)

data/emilia_preprocessing/step2_log_tar_files.sh

@@ -0,0 +1,27 @@
+emilia_root=$1 # /data/scratch/pyp/datasets/emilia/downloads
+for file in ${emilia_root}/* ; do
+# Check if gzip compressed archive
+if file "$file" | grep -q 'gzip compressed data'; then
+    # Extract string of form 'was "EN_B00100.tar"'' from the output of the file command to keep EN_B00100
+    filename=$(file "$file" | grep -oP '(?<=was ")[^"]*' | sed 's/\.tar$//')
+    # Get the file size
+    size=$(du -sh "$file" | cut -f1)
+    original_filename=$(basename "$file")
+    # Get URL string from corresponding JSON file with same basename
+    json_file=$file.json
+    if [ -f "$json_file" ]; then
+        # url=$(jq -r '.url' "$json_file") # jq is not installed on the server
+        url=$(python3 -c "import sys, json; print(json.load(open('$json_file'))['url'])")
+    else
+        url="N/A"
+    fi
+    # Compute SHA256 hash of the file
+    hash=$(python sha256hash.py "$file")
+    echo $original_filename
+    # Write filename, size, hash, original filename, URL to output file
+    echo "$filename, $size, $hash, $original_filename, $url" >> file_log.txt
+fi
+done
+
+# Sort the output file by filename
+sort -o file_log.txt file_log.txt

data/emilia_preprocessing/step3_untar.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+
+# Define the root directory where the tar files are located
+root=$1 # /data/scratch/pyp/datasets/emilia/downloads
+save_root=$2 # /data/scratch/pyp/datasets/emilia/preprocessed/audio
+
+mkdir -p "${save_root}"
+
+# Input log files
+log_file="file_log.txt"       # Full log of files to process
+exist_log_file="file_log_debug.txt" # Log of already processed files
+failure_log="untar_failures.log" # Log file for untar failures
+
+# Clear previous failure log
+> "$failure_log"
+
+# Create an array of filenames already processed (from exist_log_file)
+if [ -f "$exist_log_file" ]; then
+    mapfile -t existing_files < "$exist_log_file"
+else
+    existing_files=()
+fi
+
+# Create a temporary filtered log of files to process
+filtered_log="filtered_file_log.txt"
+grep -v -F -f "$exist_log_file" "$log_file" > "$filtered_log"
+
+# Count total filtered files
+total_files=$(wc -l < "$filtered_log")
+echo "Found $total_files entries to process in $filtered_log."
+
+# Print the filtered files
+echo "Filtered files to process:"
+cat "$filtered_log"
+echo
+
+# Confirm before starting processing
+read -p "Do you want to proceed with the above files? (y/n): " confirm
+if [[ "$confirm" != "y" ]]; then
+    echo "Operation canceled."
+    rm -f "$filtered_log"
+    exit 1
+fi
+
+# Start time
+start_time=$(date +%s)
+
+# Counter for how many lines we've processed
+count=0
+
+# Process filtered log
+while IFS=',' read -r filename size local_sha256 original_filename url; do
+    count=$((count + 1))
+
+    # Trim leading/trailing whitespace
+    filename=$(echo "$filename" | xargs)
+    size=$(echo "$size" | xargs)
+    local_sha256=$(echo "$local_sha256" | xargs)
+    original_filename=$(echo "$original_filename" | xargs)
+    url=$(echo "$url" | xargs)
+
+    # Construct the full path to the tar file
+    tar_file="${root}/${original_filename}"
+
+    # Check if the tar file exists
+    if [ ! -f "$tar_file" ]; then
+        echo "❌ File not found: $tar_file"
+        echo "$filename, $size, $local_sha256, $original_filename, $url" >> "$failure_log"
+    else
+        # Try to untar the file
+        echo "[$count/$total_files] Untarring $tar_file..."
+        
+        if ! tar -xf "$tar_file" -C "${save_root}"; then
+            # If untar fails, log the failure
+            echo "❌ Failed to untar: $tar_file"
+            echo "$filename, $size, $local_sha256, $original_filename, $url" >> "$failure_log"
+        else
+            echo "✅ Successfully untarred: $tar_file"
+            # Append successfully untarred filename to exist_log_file
+            echo "$filename" >> "$exist_log_file"
+        fi
+    fi
+
+    # Calculate elapsed time, average time per file, and ETA
+    now=$(date +%s)
+    elapsed=$(( now - start_time ))  # total seconds since the start
+    if [ $count -gt 0 ]; then
+        avg_time=$(awk "BEGIN { printf \"%.2f\", $elapsed / $count }")
+        remain=$(( total_files - count ))
+        eta_seconds=$(awk "BEGIN { printf \"%.0f\", $avg_time * $remain }")
+        eta_formatted=$(date -ud "@${eta_seconds}" +'%H:%M:%S')
+        echo "Elapsed: ${elapsed}s | Avg/f: ${avg_time}s | Remaining: $remain files | ETA: ~${eta_formatted}"
+    fi
+
+done < "$filtered_log"
+
+# Clean up temporary filtered log
+rm -f "$filtered_log"
+
+# Summary
+echo "Untar operation completed. Check $failure_log for any failures."

data/emilia_preprocessing/step4_construct_manifest.py

@@ -0,0 +1,251 @@
+# construct manifest file for training, note that we only have one train split
+# also create neighbors folder for each sample, which is simply done through speaker label in the original manifest where each file has rows
+# path\tdistance\tduration
+# where distance is always 0 because we don't know the distance between the samples
+
+# waiting on Yushen Chen to provide data filtering approach
+import sys, copy
+import os, random, numpy as np, socket
+
+import json
+import tqdm
+from multiprocessing import Pool
+import glob, os
+from collections import defaultdict
+def write_jsonl(data, fn):
+    with open(fn, "w") as file:
+        for entry in data:
+            file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+def read_jsonl(file_path):
+    cur_data = []
+    with open(file_path, 'r', encoding='utf-8-sig') as file:
+        for line in file:
+            cur_data.append(json.loads(line.strip()))
+    return cur_data
+
+def repetition_found(text, length=2, tolerance=10):
+    pattern_count = defaultdict(int)
+    for i in range(len(text) - length + 1):
+        pattern = text[i : i + length]
+        pattern_count[pattern] += 1
+    for pattern, count in pattern_count.items():
+        if count > tolerance:
+            return True
+    return False
+
+
+
+out_en = {
+    "EN_B00013_S00913",
+    "EN_B00042_S00120",
+    "EN_B00055_S04111",
+    "EN_B00061_S00693",
+    "EN_B00061_S01494",
+    "EN_B00061_S03375",
+    "EN_B00059_S00092",
+    "EN_B00111_S04300",
+    "EN_B00100_S03759",
+    "EN_B00087_S03811",
+    "EN_B00059_S00950",
+    "EN_B00089_S00946",
+    "EN_B00078_S05127",
+    "EN_B00070_S04089",
+    "EN_B00074_S09659",
+    "EN_B00061_S06983",
+    "EN_B00061_S07060",
+    "EN_B00059_S08397",
+    "EN_B00082_S06192",
+    "EN_B00091_S01238",
+    "EN_B00089_S07349",
+    "EN_B00070_S04343",
+    "EN_B00061_S02400",
+    "EN_B00076_S01262",
+    "EN_B00068_S06467",
+    "EN_B00076_S02943",
+    "EN_B00064_S05954",
+    "EN_B00061_S05386",
+    "EN_B00066_S06544",
+    "EN_B00076_S06944",
+    "EN_B00072_S08620",
+    "EN_B00076_S07135",
+    "EN_B00076_S09127",
+    "EN_B00065_S00497",
+    "EN_B00059_S06227",
+    "EN_B00063_S02859",
+    "EN_B00075_S01547",
+    "EN_B00061_S08286",
+    "EN_B00079_S02901",
+    "EN_B00092_S03643",
+    "EN_B00096_S08653",
+    "EN_B00063_S04297",
+    "EN_B00063_S04614",
+    "EN_B00079_S04698",
+    "EN_B00104_S01666",
+    "EN_B00061_S09504",
+    "EN_B00061_S09694",
+    "EN_B00065_S05444",
+    "EN_B00063_S06860",
+    "EN_B00065_S05725",
+    "EN_B00069_S07628",
+    "EN_B00083_S03875",
+    "EN_B00071_S07665",
+    "EN_B00071_S07665",
+    "EN_B00062_S04187",
+    "EN_B00065_S09873",
+    "EN_B00065_S09922",
+    "EN_B00084_S02463",
+    "EN_B00067_S05066",
+    "EN_B00106_S08060",
+    "EN_B00073_S06399",
+    "EN_B00073_S09236",
+    "EN_B00087_S00432",
+    "EN_B00085_S05618",
+    "EN_B00064_S01262",
+    "EN_B00072_S01739",
+    "EN_B00059_S03913",
+    "EN_B00069_S04036",
+    "EN_B00067_S05623",
+    "EN_B00060_S05389",
+    "EN_B00060_S07290",
+    "EN_B00062_S08995",
+}
+en_filters = ["ا", "い", "て"]
+
+
+from multiprocessing import Pool
+
+def process_meta_item(item, root, sub_root, audio_folder, audio_ext, text_ext):
+    global filtered_duration, filtered_count, total_duration, total_count
+    # Data filtering following Yushen's approach
+    if (
+        item["wav"].split("/")[-1] in out_en
+        or any(t in item["text"] for t in en_filters)
+        or repetition_found(item["text"], length=4)
+    ):
+        return None, item["duration"], 1, 0, 0, (None, None)  # Return filtered results
+    
+    # Trim leading space from text if exists
+    if item["text"].startswith(" "):
+        item["text"] = item["text"][1:]
+    
+    # write text to text file
+    text_fn = os.path.join(root, sub_root, audio_folder, item["wav"].replace(audio_ext, text_ext))
+    os.makedirs(os.path.dirname(text_fn), exist_ok=True)
+    with open(text_fn, "w") as f:
+        f.write(item["text"])
+
+    # spk2info[item["speaker"]].append(item)
+    return (
+        f"{item['wav']}\t{item['duration']}\n",
+        0,
+        0,
+        item["duration"],
+        1,
+        (item['speaker'], item)
+    )  # Return processed results
+
+
+def parallel_process_meta(meta, root, sub_root, audio_folder, num_workers, audio_ext, text_ext):
+    with Pool(num_workers) as pool:
+        results = pool.starmap(
+            process_meta_item,
+            [(item, root, sub_root, audio_folder, audio_ext, text_ext) for item in meta],
+        )
+    
+    processed_items = []
+    spkitem = []
+    filtered_duration = 0
+    filtered_count = 0
+    total_duration = 0
+    total_count = 0
+    
+    for result in results:
+        if result[0]:  # If the item was processed
+            processed_items.append(result[0])
+        filtered_duration += result[1]
+        filtered_count += result[2]
+        total_duration += result[3]
+        total_count += result[4]
+        spkitem.append(result[5])
+    
+    return processed_items, filtered_duration, filtered_count, total_duration, total_count, spkitem
+
+
+def main(
+    root: str = "/data/scratch/pyp/datasets/emilia",
+    sub_root: str = "preprocessed",
+    audio_folder: str = "audio",
+    manifest_folder: str = "manifest_for_codec",
+    neighbors_folder: str = "neighbors",
+    audio_ext: str = ".mp3",
+    text_ext: str = ".txt",
+    num_workers: int = 8,  # Specify the number of workers
+):
+    # Find the segments that are untarred
+    all_fns = [
+        item
+        for item in glob.glob(f"{root}/{sub_root}/{audio_folder}/*")
+        if os.path.basename(item).startswith("EN_") and os.path.isdir(item)
+    ]
+    print(f"found {len(all_fns)} untarred segments")
+    print(f"{all_fns[:3]}")
+    
+    res = []
+    total_duration = 0
+    total_count = 0
+    filtered_duration = 0
+    filtered_count = 0
+
+    for fn in tqdm.tqdm(all_fns, desc="overall progress"):
+        spk2info = defaultdict(list)
+        metafn = os.path.join(root, "EN", os.path.basename(fn) + ".jsonl")
+        meta = read_jsonl(metafn)
+        
+        # Parallel process metadata
+        processed_items, fd, fc, td, tc, spkitem = parallel_process_meta(
+            meta, root, sub_root, audio_folder, num_workers, audio_ext, text_ext
+        )
+        
+        # Aggregate results
+        res.extend(processed_items)
+        filtered_duration += fd
+        filtered_count += fc
+        total_duration += td
+        total_count += tc
+
+        for spk, item in spkitem:
+            if spk:
+                spk2info[spk].append(item)
+        
+        # Save neighbor files
+        for spk in spk2info:
+            for item in spk2info[spk]:
+                neighbor_fn = os.path.join(
+                    root,
+                    sub_root,
+                    neighbors_folder,
+                    item["wav"].replace(audio_ext, text_ext),
+                )
+                os.makedirs(os.path.dirname(neighbor_fn), exist_ok=True)
+                tobe_write = [f"{neighbor_item['wav'].replace(audio_ext, text_ext)}\t0\t{neighbor_item['duration']}\n" for neighbor_item in spk2info[spk] if neighbor_item["wav"] != item["wav"]]
+                if tobe_write:
+                    with open(neighbor_fn, "w") as f:
+                        f.writelines(tobe_write)
+    
+    print(
+        f"total duration: {total_duration / 3600:.2f} hours, total count: {total_count}"
+    )
+    print(
+        f"filtered duration: {filtered_duration / 3600:.2f} hours, filtered count: {filtered_count}"
+    )
+    save_fn = os.path.join(root, sub_root, manifest_folder, "train.txt")
+    os.makedirs(os.path.dirname(save_fn), exist_ok=True)
+    with open(save_fn, "w") as f:
+        for item in res:
+            f.write(item)
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire(main)

data/emilia_preprocessing/step5_phonemize.py

@@ -0,0 +1,158 @@
+import sys, copy
+import os, random, numpy as np, socket
+
+import json
+import tqdm
+from multiprocessing import Pool
+import glob, os, fire
+from collections import defaultdict
+sys.path.insert(0, "../../")
+from data.tokenizer import TextTokenizer, tokenize_text
+
+def write_jsonl(data, fn):
+    with open(fn, "w") as file:
+        for entry in data:
+            file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+
+
+def read_jsonl(file_path):
+    cur_data = []
+    with open(file_path, 'r', encoding='utf-8-sig') as file:
+        for line in file:
+            cur_data.append(json.loads(line.strip()))
+    return cur_data
+
+
+def phonemize_and_save(text, fn, text_tokenizer):
+    """Phonemizes the text and saves the result to a file."""
+    phn = tokenize_text(text_tokenizer, text)
+    os.makedirs(os.path.dirname(fn), exist_ok=True)
+    with open(fn, "w") as f:
+        f.write(" ".join(phn))
+    return set(phn)
+
+
+def process_item(item, root, sub_root, audio_folder, phn_folder, audio_ext, text_ext, phn_ext, text_tokenizer):
+    """Worker function to process a single item."""
+    text_path = os.path.join(root, sub_root, audio_folder, item[0].replace(audio_ext, text_ext))
+    if not os.path.exists(text_path):
+        return {"missing_text": text_path, "success": False, "cur_phn_set": set()}
+
+    with open(text_path, "r") as f:
+        text = [line.strip() for line in f.readlines()]
+        text = " ".join(text)
+
+    phn_path = os.path.join(root, sub_root, phn_folder, item[0].replace(audio_ext, phn_ext))
+    cur_phn_set = phonemize_and_save(text, phn_path, text_tokenizer)
+    return {"missing_text": None, "success": True, "cur_phn_set": cur_phn_set}
+
+
+def process_item_star(args):
+    """Unpacks arguments for `process_item` to work with `imap`."""
+    return process_item(*args)
+
+def main(
+    root="/data/scratch/pyp/datasets/emilia",
+    sub_root="preprocessed",
+    manifest_folder="manifest_for_codec",
+    audio_folder="audio",
+    phn_folder="phoneme",
+    audio_ext=".mp3",
+    text_ext=".txt",
+    phn_ext=".txt",
+    num_workers=8,
+):
+    """Main function to process phoneme generation in parallel."""
+    # # Initialize the tokenizer
+    text_tokenizer = TextTokenizer()
+    all_fns = glob.glob(f"{root}/{sub_root}/{manifest_folder}/*.txt")
+    print(f"found {len(all_fns)} manifest files")
+    print(f"{all_fns[:3]=}")
+
+    data = []
+    for fn in all_fns:
+        with open(fn, "r") as f:
+            data += [line.strip().split("\t") for line in f]
+    
+    vocab = set()
+
+    ################## parallel processing ##################
+    ################## parallel processing ##################
+    ################## parallel processing ##################
+    # Prepare arguments for the worker function
+    # tasks = [
+    #     (
+    #         item,
+    #         root,
+    #         sub_root,
+    #         audio_folder,
+    #         phn_folder,
+    #         audio_ext,
+    #         text_ext,
+    #         phn_ext,
+    #         text_tokenizer,
+    #     )
+    #     for item in data
+    # ]
+
+    # # Parallel processing with progress monitoring
+    # results = []
+    # with Pool(num_workers) as pool:
+    #     for result in tqdm.tqdm(
+    #         pool.imap_unordered(process_item_star, tasks),
+    #         total=len(tasks),
+    #         desc="Processing items",
+    #     ):
+    #         results.append(result)
+    # # read all manifest endswith .txt
+    # missing_text = [result["missing_text"] for result in results if not result["success"]]
+    # for result in results:
+    #     if result['success']:
+    #         vocab.update(result['cur_phn_set'])
+    ################## parallel processing ##################
+    ################## parallel processing ##################
+    ################## parallel processing ##################
+
+    ################## sequential processing ##################
+    ################## sequential processing ##################
+    ################## sequential processing ##################
+    missing_text = []
+    for item in tqdm.tqdm(data):
+        text_path = os.path.join(root, sub_root, audio_folder, item[0].replace(audio_ext, text_ext))
+        if not os.path.exists(text_path):
+            missing_text.append(text_path)
+            continue
+        try:
+            with open(text_path, "r") as f:
+                text = [line.strip() for line in f.readlines()]
+                text = " ".join(text)
+        except:
+            print(f"Error reading {text_path}")
+            continue
+        cur_phn_set = phonemize_and_save(text, os.path.join(root, sub_root, phn_folder, item[0].replace(audio_ext, phn_ext)), text_tokenizer)
+        vocab.update(cur_phn_set)
+    ################## sequential processing ##################
+    ################## sequential processing ##################
+    ################## sequential processing ##################
+
+    # save the vocab
+    vocab = list(vocab)
+    # sort the vocab
+    vocab.sort()
+    with open(os.path.join(root, sub_root, "vocab.txt"), "w") as f:
+        f.write("\n".join(vocab))
+
+    # Collect missing text paths
+    print(f"Missing text files: {len(missing_text)}")
+    if missing_text:
+        print("Some missing files:", missing_text[:10])  # Print the first 10 missing files as an example
+    
+
+if __name__ == "__main__":
+    fire.Fire(main)
+
+
+    
+
+
+

data/emilia_preprocessing/step6_encodec_encode.py

@@ -0,0 +1,177 @@
+import argparse
+from email.policy import default
+def parse_args():
+    parser = argparse.ArgumentParser(description="encode the dataset using codec model")
+    parser.add_argument('--root', type=str, default="/data/scratch/pyp/datasets/emilia", help="Path to the directory")
+    parser.add_argument('--sub_root', type=str, default="preprocessed", help="sub directory")
+    parser.add_argument('--encodec_name', type=str, default="encodec_6f79c6a8.th", help="name of the codec model")
+    parser.add_argument('--n_workers', type=int, default=16, help="Number of parallel worker processes")
+    parser.add_argument('--batch_size', type=int, default=16, help="batch size for codec encoding, decrease it if OOM. This is the sum of batch size *over each gpu*, so increase it if you are using more gpus")
+    parser.add_argument('--audio_sr', type=int, default=16000, help='input audio sample rate')
+    parser.add_argument('--model_sr', type=int, default=16000, help='encodec input audio sample rate')
+    parser.add_argument('--downsample_rate', type=int, default=320, help='encodec downsample rate')
+    parser.add_argument('--model_code_sr', type=float, default=50, help='codec model code sample rate')
+    parser.add_argument('--len_cap', type=float, default=1000, help='will drop audios that are longer than this number')
+    parser.add_argument('--min_len', type=float, default=0.5, help='will drop audios that are shorter than this number')
+    parser.add_argument('--partition', type=str, default="1/1", help='split for parallel processing')
+    parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'])
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    import logging
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+
+    import os, sys
+    import numpy as np
+    import torch
+    import torchaudio
+    import tqdm
+    import time
+
+    args = parse_args()
+
+    def sort_by_audio_len(lens):
+        inds = np.argsort(lens).tolist()
+        
+        logging.info(f"longest: {lens[inds[-1]]/args.downsample_rate} encodec codes, {lens[inds[-1]]/args.model_sr:.2f} sec.")
+        logging.info(f"shortest: {lens[inds[0]]/args.downsample_rate} encodec codes, {lens[inds[0]]/args.model_sr:.2f} sec.")
+        logging.info(f"median: {lens[inds[len(inds)//2]]/args.downsample_rate} encodec codes, {lens[inds[len(inds)//2]]/args.model_sr:.2f} sec.")
+        logging.info(f"95 percentile longest: {lens[inds[int(len(inds)*0.95)]]/args.downsample_rate} encodec codes, {lens[inds[int(len(inds)*0.95)]]/args.model_sr:.2f} sec.")
+        return inds[::-1]
+
+    def write_array_to_txt_file(array, filename):
+        with open(filename, 'w') as f:
+            for a in array[:-1]:
+                f.write(' '.join(map(str, a))+'\n')
+            f.write(' '.join(map(str, array[-1])))
+
+    class mydataset(torch.utils.data.Dataset):
+        def __init__(self, split):
+            super().__init__()
+            self.split = split
+            self.audio_dir = audio_dir
+            manifest_fn = os.path.join(encodec_manifest_dir, split+".txt")
+            cur_sp = int(args.partition.split("/")[0])-1
+            total_sp = int(args.partition.split("/")[1])
+            with open(manifest_fn, "r") as rf:
+                self.data = [l.strip().split("\t") for l in rf.readlines()][cur_sp::total_sp]
+        def __len__(self):
+            return len(self.data)
+        def __getitem__(self, ind):
+            try:
+                afn = self.data[ind][0]
+                fn = os.path.join(self.audio_dir, afn)
+                audio, sr = torchaudio.load(fn)
+                if sr != args.model_sr:
+                    audio = torchaudio.transforms.Resample(sr, args.model_sr)(audio)
+                    sr = args.model_sr
+                assert sr == args.model_sr, sr
+            except Exception as e:
+                # logging.info(f"{e}")
+                return None, None, None
+            assert audio.ndim==2 and audio.shape[0] == 1, audio.shape
+            return audio.type(torch.float32).squeeze(0), audio.shape[-1], os.path.splitext(afn)[0]
+        def collate(self, batch):
+            lens, audios, segment_ids = [], [], []
+            for item in batch:
+                if item[0] != None:
+                    audios.append(item[0])
+                    lens.append(item[1])
+                    segment_ids.append(item[2])
+            return audios, lens, segment_ids
+    
+    # roots
+    sub_root = args.sub_root
+    encodec_manifest_dir = os.path.join(args.root, sub_root, "manifest_for_codec")
+    audio_dir = os.path.join(args.root, sub_root, "audio")
+    save_manifest_dir = os.path.join(args.root, sub_root,"manifest_final_encodec")
+    if args.encodec_name == "encodec_6f79c6a8.th":
+        save_codes_dir = os.path.join(args.root, sub_root,"encodec_4cb")
+    elif args.encodec_name == "encodec_8cb1024_giga.th":
+        save_codes_dir = os.path.join(args.root, sub_root,"encodec_8cb")
+
+    os.makedirs(save_manifest_dir, exist_ok=True)
+    os.makedirs(save_codes_dir, exist_ok=True)
+    
+    def import_encodec():
+        from encodec import get_compression_model
+        userdir = os.path.expanduser("~")
+        model = get_compression_model(os.path.join(userdir, "VoiceStar", f"pretrained/{args.encodec_name}"), encode_only=True, device="cuda")
+        model = torch.nn.DataParallel(model)
+        return model
+    model = import_encodec()
+    
+    # setup dataloader
+    mega_batch_size = 2048
+    batch_size = args.batch_size
+    
+    dataset = mydataset(args.split)
+    if len(dataset) == 0:
+        logging.info(f"no data found for split {args.split} partition {args.partition}")
+        sys.exit(0)
+    loader = torch.torch.utils.data.DataLoader(dataset, batch_size=mega_batch_size, shuffle=False, drop_last=False, num_workers=args.n_workers, collate_fn=dataset.collate)
+    split = args.split
+
+    skip = 0
+    logging.info(f"now processing split {split} partition {args.partition}...")
+    mega_n_steps = int(np.ceil(len(loader.dataset) / mega_batch_size))
+    # mega_n_steps = int(np.ceil(len(gs) / mega_batch_size))
+    logging.info(f"partition the split {split} into {mega_n_steps} parts, each has at most {mega_batch_size} samples")
+    mani_fn = os.path.join(save_manifest_dir, f"{split}_{args.partition.replace('/', '=')}.txt")
+    logging.info(f"manifest for split {split} partition {args.partition.replace('/', '=')}.txt will be saved at {mani_fn}")
+    with open(mani_fn, "w") as mani_wf:
+    # with open(mani_fn, "a") as mani_wf: # resume from where we failed
+        for m, mega_batch in enumerate(tqdm.tqdm(loader, mininterval=60, maxinterval=60)):
+
+            logging.info(f"====================================")
+            logging.info(f"====================================")
+            logging.info(f"now processing mega step {m+1}/{mega_n_steps}")
+
+            try:
+                lengths = np.array(mega_batch[1])
+                sorted_inds = sort_by_audio_len(lengths)
+                for j in range(len(sorted_inds))[::-1]:
+                    if lengths[sorted_inds[j]] < args.model_sr*args.min_len or lengths[sorted_inds[j]] > args.model_sr*args.len_cap: # skip samples that are too short (shorter than 0.2s), or too big (bigger than 80s)
+                        skip += 1
+                        del sorted_inds[j]
+                
+                n_steps = int(np.ceil(len(sorted_inds) / batch_size))
+                for n in tqdm.tqdm(range(n_steps), disable=True):
+                    inds_used = sorted_inds[n*batch_size:(n+1)*batch_size]
+                    wav_batch = [mega_batch[0][id] for id in inds_used]
+                    all_lens = [mega_batch[1][id] for id in inds_used]
+                    segment_id_batch = [mega_batch[2][id] for id in inds_used]
+                    padded_wav = torch.nn.utils.rnn.pad_sequence(wav_batch, batch_first=True).unsqueeze(1) # [B, T] -> [B, 1, T]
+                    # Extract discrete codes from EnCodec
+                    with torch.no_grad():
+                        if max(all_lens) > 300000 and len(all_lens) > 1: # if utterances are long, simply pass half of them at a time
+                            codes = []
+                            inwav = padded_wav.cuda()
+                            codes.append(model(inwav[:len(inwav)//2])[0].cpu())
+                            codes.append(model(inwav[len(inwav)//2:])[0].cpu())
+                            codes = torch.cat(codes, dim=0)
+                        else:
+                            encoded_frames = model(padded_wav.cuda()) 
+                            codes = encoded_frames[0].cpu() # [B, n_codebook, T]
+
+                    for i, length in enumerate(all_lens):
+                        save_fn = os.path.join(save_codes_dir, segment_id_batch[i]+".txt")
+                        actual_len = round(length / args.downsample_rate) # 320 is downsample rate for this model
+                        cur_code = codes[i].tolist() if type(codes) == list else codes[i, :, :actual_len].tolist()
+                        os.makedirs(os.path.dirname(save_fn), exist_ok=True)
+                        write_array_to_txt_file(cur_code, save_fn)
+
+                        mani_wf.write(f"{segment_id_batch[i]}\t{len(cur_code[0])}\n") # write to manifest file
+                        # if i == 10:
+                        #    raise
+            except Exception as e:
+                print(f'exception!! at {m+1}')
+                print(e)
+                continue
+
+            # break
+    logging.info(f"split {split} partition {args.partition} has {len(loader.dataset)} samples in total, skipped {skip} due to utterance being too long or too short")
+        # break

data/emilia_preprocessing/step6_encodec_encode_script.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate voicestar
+
+dir=${processed_dir:-/data/scratch/pyp/datasets/emilia}
+sub_root=${sub_root:-preprocessed}
+encodec_name=${encodec_name:-"encodec_6f79c6a8.th"}
+n_workers=${n_workers:-64}
+batch_size=${batch_size:-512}
+audio_sr=16000
+model_sr=16000
+downsample_rate=320
+model_code_sr=50
+len_cap=1000
+min_len=0.5
+partition=${partition:-"1/1"}
+split=${split:-"train"}
+
+python step6_encodec_encode.py --root $dir --sub_root ${sub_root} --encodec_name ${encodec_name} --n_workers $n_workers --batch_size $batch_size --audio_sr $audio_sr --model_sr $model_sr --downsample_rate $downsample_rate --model_code_sr $model_code_sr --len_cap $len_cap --min_len $min_len --partition $partition --split $split

data/encodec.py

@@ -0,0 +1,1554 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compression models or wrapper around existing models.
+Also defines the main interface that a model must follow to be usable as an audio tokenizer.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+import logging
+import math
+from pathlib import Path
+import typing as tp
+
+import numpy as np
+import torch
+from torch import nn
+from torch import einsum
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+import logging
+import warnings
+from einops import rearrange, repeat
+import omegaconf
+# import flashy
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                 'time_group_norm'])
+
+def dict_from_config(cfg: omegaconf.DictConfig) -> dict:
+    """Convenience function to map an omegaconf configuration to a dictionary.
+
+    Args:
+        cfg (omegaconf.DictConfig): Original configuration to map to dict.
+    Returns:
+        dict: Config as dictionary object.
+    """
+    dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)
+    assert isinstance(dct, dict)
+    return dct
+
+@dataclass
+class QuantizedResult:
+    x: torch.Tensor
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+class BaseQuantizer(nn.Module):
+    """Base class for quantizers.
+    """
+
+    def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:
+        """
+        Given input tensor x, returns first the quantized (or approximately quantized)
+        representation along with quantized codes, bandwidth, and any penalty term for the loss.
+        Finally, this returns a dict of metrics to update logging etc.
+        Frame rate must be passed so that the bandwidth is properly computed.
+        """
+        raise NotImplementedError()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth."""
+        raise NotImplementedError()
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        raise NotImplementedError()
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def num_codebooks(self):
+        """Number of active codebooks."""
+        raise NotImplementedError()
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise NotImplementedError()
+
+class CompressionModel(ABC, nn.Module):
+    """Base API for all compression model that aim at being used as audio tokenizers
+    with a language model.
+    """
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> QuantizedResult:
+        ...
+
+    @abstractmethod
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """See `EncodecModel.encode`."""
+        ...
+
+    @abstractmethod
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """See `EncodecModel.decode`."""
+        ...
+
+    @abstractmethod
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        ...
+
+    @property
+    @abstractmethod
+    def channels(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def frame_rate(self) -> float:
+        ...
+
+    @property
+    @abstractmethod
+    def sample_rate(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def cardinality(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def num_codebooks(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def total_codebooks(self) -> int:
+        ...
+
+    @abstractmethod
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        ...
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none'):
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs):
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                 padding_total: int = 0) -> int:
+    """See `pad_for_conv1d`."""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+    """Wrapper around ConvTranspose2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class StreamableConv1d(nn.Module):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, dilation: int = 1,
+                 groups: int = 1, bias: bool = True, causal: bool = False,
+                 norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 pad_mode: str = 'reflect'):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn("StreamableConv1d has been initialized with stride > 1 and dilation > 1"
+                          f" (kernel_size={kernel_size} stride={stride}, dilation={dilation}).")
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                               dilation=dilation, groups=groups, bias=bias, causal=causal,
+                               norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (kernel_size - 1) * dilation + 1  # effective kernel size with dilations
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        return self.conv(x)
+
+
+class StreamableConvTranspose1d(nn.Module):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, causal: bool = False,
+                 norm: str = 'none', trim_right_ratio: float = 1.,
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                          causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+class StreamableLSTM(nn.Module):
+    """LSTM without worrying about the hidden state, nor the layout of the data.
+    Expects input as convolutional layout.
+    """
+    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+
+
+class SEANetResnetBlock(nn.Module):
+    """Residual block from SEANet model.
+
+    Args:
+        dim (int): Dimension of the input/output.
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection.
+    """
+    def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],
+                 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,
+                 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'
+        act = getattr(nn, activation)
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                act(**activation_params),
+                StreamableConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,
+                                 norm=norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+        self.block = nn.Sequential(*block)
+        self.shortcut: nn.Module
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = StreamableConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,
+                                             causal=causal, pad_mode=pad_mode)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x)
+
+
+class SEANetEncoder(nn.Module):
+    """SEANet encoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order. We use the decoder order as some models may only employ the decoder.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the encoder, it corresponds to the N first blocks.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = 1
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(channels, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= i + 2 else norm
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      norm=block_norm, norm_params=norm_params,
+                                      activation=activation, activation_params=activation_params,
+                                      causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConv1d(mult * n_filters, mult * n_filters * 2,
+                                 kernel_size=ratio * 2, stride=ratio,
+                                 norm=block_norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+            mult *= 2
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        model += [
+            act(**activation_params),
+            StreamableConv1d(mult * n_filters, dimension, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(nn.Module):
+    """SEANet decoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple.
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the decoder, it corresponds to the N last blocks.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0, trim_right_ratio: float = 1.0):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(dimension, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) else norm
+            # Add upsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConvTranspose1d(mult * n_filters, mult * n_filters // 2,
+                                          kernel_size=ratio * 2, stride=ratio,
+                                          norm=block_norm, norm_kwargs=norm_params,
+                                          causal=causal, trim_right_ratio=trim_right_ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      activation=activation, activation_params=activation_params,
+                                      norm=block_norm, norm_params=norm_params, causal=causal,
+                                      pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params),
+            StreamableConv1d(n_filters, channels, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None:
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [
+                final_act(**final_activation_params)
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, z):
+        y = self.model(z)
+        return y
+
+
+def exists(val: tp.Optional[tp.Any]) -> bool:
+    return val is not None
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if exists(val) else d
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay: float):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters: int, num_iters: int = 10):
+    dim, dtype = samples.shape[-1], samples.dtype
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        diffs = rearrange(samples, "n d -> n () d") - rearrange(
+            means, "c d -> () c d"
+        )
+        dists = -(diffs ** 2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+def orthogonal_loss_fn(t):
+    # eq (2) from https://arxiv.org/abs/2112.00384
+    n = t.shape[0]
+    normed_codes = l2norm(t)
+    identity = torch.eye(n, device=t.device)
+    cosine_sim = einsum("i d, j d -> i j", normed_codes, normed_codes)
+    return ((cosine_sim - identity) ** 2).sum() / (n ** 2)
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
+            If set to true, run the k-means algorithm on the first training batch and use
+            the learned centroids as initialization.
+        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        kmeans_init: int = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.kmeans_iters = kmeans_iters
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed.clone())
+        self.cluster_size.data.copy_(cluster_size)
+        self.inited.data.copy_(torch.Tensor([True]))
+        # Make sure all buffers across workers are in sync after initialization
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def replace_(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes)
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(
+            x.pow(2).sum(1, keepdim=True)
+            - 2 * x @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+        embed_ind = dist.max(dim=-1).indices
+        return embed_ind
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        return embed_ind
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x):
+        raise NotImplementedError()
+        shape, dtype = x.shape, x.dtype
+        x = self.preprocess(x)
+        self.init_embed_(x)
+
+        embed_ind = self.quantize(x)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        quantize = self.dequantize(embed_ind)
+
+        if self.training:
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x)
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = x.t() @ embed_onehot
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+
+        return quantize, embed_ind
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int):
+        channels_last (bool): Channels are the last dimension in the input tensors.
+        commitment_weight (float): Weight for commitment loss.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider
+            for orthogonal regularization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        channels_last: bool = False,
+        commitment_weight: float = 1.,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        _codebook_dim: int = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim
+        self.project_in = (nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity())
+        self.project_out = (nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity())
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+
+        self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,
+                                           kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,
+                                           decay=decay, epsilon=epsilon,
+                                           threshold_ema_dead_code=threshold_ema_dead_code)
+        self.codebook_size = codebook_size
+
+        self.channels_last = channels_last
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    @property
+    def inited(self):
+        return self._codebook.inited
+
+    def _preprocess(self, x):
+        if not self.channels_last:
+            x = rearrange(x, "b d n -> b n d")
+        return x
+
+    def _postprocess(self, quantize):
+        if not self.channels_last:
+            quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def encode(self, x):
+        x = self._preprocess(x)
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+        return quantize
+
+    def forward(self, x):
+        device = x.device
+        x = self._preprocess(x)
+
+        x = self.project_in(x)
+        quantize, embed_ind = self._codebook(x)
+
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
+        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
+
+        if self.training:
+            if self.commitment_weight > 0:
+                commit_loss = F.mse_loss(quantize.detach(), x)
+                loss = loss + commit_loss * self.commitment_weight
+
+            if self.orthogonal_reg_weight > 0:
+                codebook = self.codebook
+
+                if self.orthogonal_reg_active_codes_only:
+                    # only calculate orthogonal loss for the activated codes for this batch
+                    unique_code_ids = torch.unique(embed_ind)
+                    codebook = codebook[unique_code_ids]
+
+                num_codes = codebook.shape[0]
+                if exists(self.orthogonal_reg_max_codes) and num_codes > self.orthogonal_reg_max_codes:
+                    rand_ids = torch.randperm(num_codes, device=device)[:self.orthogonal_reg_max_codes]
+                    codebook = codebook[rand_ids]
+
+                orthogonal_reg_loss = orthogonal_loss_fn(codebook)
+                loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
+
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+
+        return quantize, embed_ind, loss
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        codebook_size = kwargs.pop('codebook_size', None)
+        if codebook_size is None:
+            raise ValueError("codebook_size must be provided in kwargs")
+        if type(codebook_size) != list:
+            codebook_size = [codebook_size] * num_quantizers
+        self.layers = nn.ModuleList(
+            [VectorQuantization(codebook_size=cur_codebook_size, **kwargs) for _,cur_codebook_size in zip(range(num_quantizers), codebook_size)]
+        )
+            
+
+        # self.layers = nn.ModuleList(
+        #     [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        # )
+
+    def forward(self, x, n_q: tp.Optional[int] = None):
+        quantized_out = 0.0
+        residual = x
+
+        all_losses = []
+        all_indices = []
+
+        n_q = n_q or len(self.layers)
+
+        for i, layer in enumerate(self.layers[:n_q]):
+            quantized, indices, loss = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+            all_indices.append(indices)
+            all_losses.append(loss)
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        return quantized_out, out_indices, out_losses
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        for layer in self.layers[:n_q]:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            # the original code is below
+            # since quantize has the gradient of residual, according to line 321
+            #  quantize = x + (quantize - x).detach()
+            # the code below will make commitment loss to be 0 for all codebooks except for codebook1
+            # https://github.com/facebookresearch/encodec/issues/25
+            # therefore we change it
+
+            residual = residual - quantized
+            # residual = residual - quantized.detach()
+            # since commitment loss is averaged, the scale of the loss won't get change (not as said in the issue above)
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out
+
+
+class ResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer.
+
+    Args:
+        dimension (int): Dimension of the codebooks.
+        n_q (int): Number of residual vector quantizers used.
+        q_dropout (bool): Random quantizer drop out at train time.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider.
+            for orthogonal regularization.
+    """
+    def __init__(
+        self,
+        dimension: int = 256,
+        n_q: int = 8,
+        q_dropout: bool = False,
+        bins: tp.Union[int, tp.List[int]] = 1024,
+        decay: float = 0.99,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.max_n_q = n_q
+        self.n_q = n_q
+        self.q_dropout = q_dropout
+        self.dimension = dimension
+        self.bins = bins
+        self.decay = decay
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            kmeans_init=self.kmeans_init,
+            kmeans_iters=self.kmeans_iters,
+            threshold_ema_dead_code=self.threshold_ema_dead_code,
+            orthogonal_reg_weight=self.orthogonal_reg_weight,
+            orthogonal_reg_active_codes_only=self.orthogonal_reg_active_codes_only,
+            orthogonal_reg_max_codes=self.orthogonal_reg_max_codes,
+            channels_last=False
+        )
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        n_q = self.n_q
+        if self.training and self.q_dropout:
+            n_q = int(torch.randint(1, self.n_q + 1, (1,)).item())
+        if type(self.bins) == list:
+            bins = self.bins
+        else:
+            bins = [self.bins] * self.n_q
+        bw_per_q = [math.log2(bin) * frame_rate / 1000 for bin in bins]
+        bw = torch.tensor(sum(bw_per_q)).to(x)
+        quantized, codes, commit_loss = self.vq(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        n_q = self.n_q
+        codes = self.vq.encode(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        codes = codes.transpose(0, 1)
+        quantized = self.vq.decode(codes)
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    def set_num_codebooks(self, n: int):
+        assert n > 0 and n <= self.max_n_q
+        self.n_q = n
+
+class DummyQuantizer(BaseQuantizer):
+    """Fake quantizer that actually does not perform any quantization.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        q = x.unsqueeze(1)
+        return QuantizedResult(x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return x.unsqueeze(1)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return codes.squeeze(1)
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        return 1
+
+    @property
+    def num_codebooks(self):
+        """Total number of codebooks."""
+        return self.total_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise AttributeError("Cannot override the number of codebooks for the dummy quantizer")
+
+
+class EncodecModel(CompressionModel):
+    """Encodec model operating on the raw waveform.
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        decoder (nn.Module): Decoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 decoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        return self.quantizer.decode(codes)
+
+class EncodecModel_encode_only(CompressionModel):
+    """Encodec model operating on the raw waveform. Encode only, so no decoder
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        raise NotImplementedError("Decode is not supported for encode only model")
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        raise NotImplementedError("Decode is not supported for encode only model")
+        return self.quantizer.decode(codes)
+
+def get_quantizer(quantizer: str, cfg: omegaconf.DictConfig, dimension: int) -> BaseQuantizer:
+    klass = {
+        'no_quant': DummyQuantizer,
+        'rvq': ResidualVectorQuantizer
+    }[quantizer]
+    kwargs = dict_from_config(getattr(cfg, quantizer))
+    if quantizer != 'no_quant':
+        kwargs['dimension'] = dimension
+    return klass(**kwargs)
+
+def get_encodec_autoencoder(encoder_name: str, cfg: omegaconf.DictConfig):
+    if encoder_name == 'seanet':
+        kwargs = dict_from_config(getattr(cfg, 'seanet'))
+        encoder_override_kwargs = kwargs.pop('encoder')
+        decoder_override_kwargs = kwargs.pop('decoder')
+        encoder_kwargs = {**kwargs, **encoder_override_kwargs}
+        decoder_kwargs = {**kwargs, **decoder_override_kwargs}
+        encoder = SEANetEncoder(**encoder_kwargs)
+        decoder = SEANetDecoder(**decoder_kwargs)
+        return encoder, decoder
+    else:
+        raise KeyError(f"Unexpected compression model {cfg.compression_model}")
+
+
+def get_compression_model(ckpt_fn, encode_only=False, device="cpu") -> CompressionModel:
+    """Instantiate a compression model."""
+    if device == None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    state = torch.load(ckpt_fn, map_location='cpu')
+    cfg = state['xp.cfg']
+    cfg.device = str(device)
+    weights = state['best_state']['model']
+    assert cfg.compression_model == 'encodec', "Only Encodec model is supported for now."
+    if encode_only:
+        all_keys = list(weights.keys())
+        for key in all_keys:
+            if key.startswith('decoder'):
+                del weights[key]
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, _ = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel_encode_only(encoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+    else:
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, decoder = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel(encoder, decoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+if __name__ == "__main__":
+    import torchaudio
+    ckpt_fn = "/home/pyp/BoostedVoiceEditor/pretrained/encodec_6f79c6a8.th"
+    audio_in_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam.wav", "/home/pyp/BoostedVoiceEditor/demo/ray.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean.wav", "/home/pyp/BoostedVoiceEditor/demo/bible.wav", "/home/pyp/BoostedVoiceEditor/demo/miley.wav"]
+    audio_out_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/ray_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/bible_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/miley_encodecTest.wav"]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = get_compression_model(ckpt_fn, device=device)
+
+    for audio_in_fn, audio_out_fn in zip(audio_in_fns, audio_out_fns):
+        audio_in, sr = torchaudio.load(audio_in_fn)
+        if sr != model.sample_rate:
+            audio_in = torchaudio.transforms.Resample(sr, model.sample_rate)(audio_in)
+        if audio_in.shape[0] == 2:
+            audio_in = audio_in.mean(dim=0, keepdim=True)
+        audio_in = audio_in.unsqueeze(0)
+        audio_in = audio_in.to(torch.float32).to(device)
+        codes = model.encode(audio_in)[0]
+        audio_out = model.decode(codes)[0].cpu()
+        torchaudio.save(audio_out_fn, audio_out, model.sample_rate)

data/ll60k_preprocessing/config.yaml

@@ -0,0 +1,75 @@
+# WARNING: This is the base configuration file shared across ALL solvers in AudioCraft
+# Please don't update this file directly. Instead use distinct configuration files
+# to override the below configuration.
+defaults:
+  - _self_
+  - dset: default
+  - solver: default
+
+device: cuda
+dtype: float32
+autocast: false
+autocast_dtype: bfloat16
+seed: 2036
+show: false  # just show the model and its size and exit
+continue_from:  # continue from a given sig or path
+execute_only:  # can be set to generate/evaluate/valid to run that stage
+execute_inplace: false # don't enforce continue_from to be set
+                       # to enable inplace execution of the stage. This assume
+                       # that you know what you are doing and execute stage
+                       # preserving the original xp sig.
+benchmark_no_load: false  # if set to true, will repeat the same batch instead of loading them
+
+efficient_attention_backend: torch  # can be torch or xformers.
+num_threads: 1                      # called with torch.set_num_thread.
+mp_start_method: forkserver               # multiprocessing method (spawn, fork or fork_server).
+
+
+label:  # use this if you want twice the same exp, with a name.
+
+# logging parameters
+logging:
+  level: INFO
+  log_updates: 10
+  log_tensorboard: false
+  log_wandb: false
+tensorboard:
+  with_media_logging: false
+  name:  # optional name for the experiment
+  sub_dir:  # optional sub directory to store tensorboard data
+wandb:
+  with_media_logging: true
+  project:  # project name
+  name:  # optional name for the experiment
+  group:  # optional group
+
+# SLURM launcher configuration.
+slurm:
+  gpus: 4  # convenience parameter, number of GPUs to use.
+  mem_per_gpu: 40  # in GB, total mem is automatically scaled with `gpus`.
+  time: 3600
+  constraint:
+  partition:
+  comment:
+  setup: []
+  exclude: ''
+
+# dora parameters
+dora:
+  # Output folder for all artifacts of an experiment.
+  dir: /checkpoint/${oc.env:USER}/experiments/audiocraft/outputs
+  # The following entries will be ignored by dora when computing the unique XP signature.
+  # Note that slurm.* and dora.* are automatically ignored.
+  exclude: [
+    'device', 'wandb.*', 'tensorboard.*', 'logging.*',
+    'dataset.num_workers', 'eval.num_workers', 'special.*',
+    'metrics.visqol.bin', 'metrics.fad.bin',
+    'execute_only', 'execute_best', 'generate.every',
+    'optim.eager_sync', 'profiler.*', 'deadlock.*',
+    'efficient_attention_backend', 'num_threads', 'mp_start_method',
+  ]
+  use_rendezvous: false
+  # for grids, always run from a clean repo, allowing reliable runs and storing
+  # the exact commit. Your repo must be absolutely pristine clean.
+  # Local `dora run` are not impacted for easier debugging.
+  git_save: true

data/ll60k_preprocessing/encodec.py

@@ -0,0 +1,1554 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+"""Compression models or wrapper around existing models.
+Also defines the main interface that a model must follow to be usable as an audio tokenizer.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+import logging
+import math
+from pathlib import Path
+import typing as tp
+
+import numpy as np
+import torch
+from torch import nn
+from torch import einsum
+import torch.nn.functional as F
+from torch.nn.utils import spectral_norm, weight_norm
+
+import logging
+import warnings
+from einops import rearrange, repeat
+import omegaconf
+# import flashy
+
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                 'time_group_norm'])
+
+def dict_from_config(cfg: omegaconf.DictConfig) -> dict:
+    """Convenience function to map an omegaconf configuration to a dictionary.
+
+    Args:
+        cfg (omegaconf.DictConfig): Original configuration to map to dict.
+    Returns:
+        dict: Config as dictionary object.
+    """
+    dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)
+    assert isinstance(dct, dict)
+    return dct
+
+@dataclass
+class QuantizedResult:
+    x: torch.Tensor
+    codes: torch.Tensor
+    bandwidth: torch.Tensor  # bandwidth in kb/s used, per batch item.
+    penalty: tp.Optional[torch.Tensor] = None
+    metrics: dict = field(default_factory=dict)
+
+class BaseQuantizer(nn.Module):
+    """Base class for quantizers.
+    """
+
+    def forward(self, x: torch.Tensor, frame_rate: int) -> QuantizedResult:
+        """
+        Given input tensor x, returns first the quantized (or approximately quantized)
+        representation along with quantized codes, bandwidth, and any penalty term for the loss.
+        Finally, this returns a dict of metrics to update logging etc.
+        Frame rate must be passed so that the bandwidth is properly computed.
+        """
+        raise NotImplementedError()
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth."""
+        raise NotImplementedError()
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        raise NotImplementedError()
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        raise NotImplementedError()
+
+    @property
+    def num_codebooks(self):
+        """Number of active codebooks."""
+        raise NotImplementedError()
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise NotImplementedError()
+
+class CompressionModel(ABC, nn.Module):
+    """Base API for all compression model that aim at being used as audio tokenizers
+    with a language model.
+    """
+
+    @abstractmethod
+    def forward(self, x: torch.Tensor) -> QuantizedResult:
+        ...
+
+    @abstractmethod
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """See `EncodecModel.encode`."""
+        ...
+
+    @abstractmethod
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """See `EncodecModel.decode`."""
+        ...
+
+    @abstractmethod
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        ...
+
+    @property
+    @abstractmethod
+    def channels(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def frame_rate(self) -> float:
+        ...
+
+    @property
+    @abstractmethod
+    def sample_rate(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def cardinality(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def num_codebooks(self) -> int:
+        ...
+
+    @property
+    @abstractmethod
+    def total_codebooks(self) -> int:
+        ...
+
+    @abstractmethod
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        ...
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none'):
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return weight_norm(module)
+    elif norm == 'spectral_norm':
+        return spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs):
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                 padding_total: int = 0) -> int:
+    """See `pad_for_conv1d`."""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
+    """Pad for a convolution to make sure that the last window is full.
+    Extra padding is added at the end. This is required to ensure that we can rebuild
+    an output of the same length, as otherwise, even with padding, some time steps
+    might get removed.
+    For instance, with total padding = 4, kernel size = 4, stride = 2:
+        0 0 1 2 3 4 5 0 0   # (0s are padding)
+        1   2   3           # (output frames of a convolution, last 0 is never used)
+        0 0 1 2 3 4 5 0     # (output of tr. conv., but pos. 5 is going to get removed as padding)
+            1 2 3 4         # once you removed padding, we are missing one time step !
+    """
+    extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+    return F.pad(x, (0, extra_padding))
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'constant', value: float = 0.):
+    """Tiny wrapper around F.pad, just to allow for reflect padding on small input.
+    If this is the case, we insert extra 0 padding to the right before the reflection happen.
+    """
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConv2d(nn.Module):
+    """Wrapper around Conv2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal=False, norm=norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose2d(nn.Module):
+    """Wrapper around ConvTranspose2d and normalization applied to this conv
+    to provide a uniform interface across normalization approaches.
+    """
+    def __init__(self, *args, norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose2d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal=False, norm=norm, **norm_kwargs)
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class StreamableConv1d(nn.Module):
+    """Conv1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, dilation: int = 1,
+                 groups: int = 1, bias: bool = True, causal: bool = False,
+                 norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                 pad_mode: str = 'reflect'):
+        super().__init__()
+        # warn user on unusual setup between dilation and stride
+        if stride > 1 and dilation > 1:
+            warnings.warn("StreamableConv1d has been initialized with stride > 1 and dilation > 1"
+                          f" (kernel_size={kernel_size} stride={stride}, dilation={dilation}).")
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                               dilation=dilation, groups=groups, bias=bias, causal=causal,
+                               norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+
+    def forward(self, x):
+        B, C, T = x.shape
+        kernel_size = self.conv.conv.kernel_size[0]
+        stride = self.conv.conv.stride[0]
+        dilation = self.conv.conv.dilation[0]
+        kernel_size = (kernel_size - 1) * dilation + 1  # effective kernel size with dilations
+        padding_total = kernel_size - stride
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        if self.causal:
+            # Left padding for causal
+            x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        return self.conv(x)
+
+
+class StreamableConvTranspose1d(nn.Module):
+    """ConvTranspose1d with some builtin handling of asymmetric or causal padding
+    and normalization.
+    """
+    def __init__(self, in_channels: int, out_channels: int,
+                 kernel_size: int, stride: int = 1, causal: bool = False,
+                 norm: str = 'none', trim_right_ratio: float = 1.,
+                 norm_kwargs: tp.Dict[str, tp.Any] = {}):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                          causal=causal, norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+    def forward(self, x):
+        kernel_size = self.convtr.convtr.kernel_size[0]
+        stride = self.convtr.convtr.stride[0]
+        padding_total = kernel_size - stride
+
+        y = self.convtr(x)
+
+        # We will only trim fixed padding. Extra padding from `pad_for_conv1d` would be
+        # removed at the very end, when keeping only the right length for the output,
+        # as removing it here would require also passing the length at the matching layer
+        # in the encoder.
+        if self.causal:
+            # Trim the padding on the right according to the specified ratio
+            # if trim_right_ratio = 1.0, trim everything from right
+            padding_right = math.ceil(padding_total * self.trim_right_ratio)
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        else:
+            # Asymmetric padding required for odd strides
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            y = unpad1d(y, (padding_left, padding_right))
+        return y
+
+
+class StreamableLSTM(nn.Module):
+    """LSTM without worrying about the hidden state, nor the layout of the data.
+    Expects input as convolutional layout.
+    """
+    def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):
+        super().__init__()
+        self.skip = skip
+        self.lstm = nn.LSTM(dimension, dimension, num_layers)
+
+    def forward(self, x):
+        x = x.permute(2, 0, 1)
+        y, _ = self.lstm(x)
+        if self.skip:
+            y = y + x
+        y = y.permute(1, 2, 0)
+        return y
+
+
+class SEANetResnetBlock(nn.Module):
+    """Residual block from SEANet model.
+
+    Args:
+        dim (int): Dimension of the input/output.
+        kernel_sizes (list): List of kernel sizes for the convolutions.
+        dilations (list): List of dilations for the convolutions.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection.
+    """
+    def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],
+                 activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,
+                 pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):
+        super().__init__()
+        assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'
+        act = getattr(nn, activation)
+        hidden = dim // compress
+        block = []
+        for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
+            in_chs = dim if i == 0 else hidden
+            out_chs = dim if i == len(kernel_sizes) - 1 else hidden
+            block += [
+                act(**activation_params),
+                StreamableConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,
+                                 norm=norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+        self.block = nn.Sequential(*block)
+        self.shortcut: nn.Module
+        if true_skip:
+            self.shortcut = nn.Identity()
+        else:
+            self.shortcut = StreamableConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,
+                                             causal=causal, pad_mode=pad_mode)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x)
+
+
+class SEANetEncoder(nn.Module):
+    """SEANet encoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios. The encoder uses downsampling ratios instead of
+            upsampling ratios, hence it will use the ratios in the reverse order to the ones specified here
+            that must match the decoder order. We use the decoder order as some models may only employ the decoder.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the encoder, it corresponds to the N first blocks.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0):
+        super().__init__()
+        self.channels = channels
+        self.dimension = dimension
+        self.n_filters = n_filters
+        self.ratios = list(reversed(ratios))
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = 1
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(channels, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Downsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= i + 2 else norm
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      norm=block_norm, norm_params=norm_params,
+                                      activation=activation, activation_params=activation_params,
+                                      causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            # Add downsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConv1d(mult * n_filters, mult * n_filters * 2,
+                                 kernel_size=ratio * 2, stride=ratio,
+                                 norm=block_norm, norm_kwargs=norm_params,
+                                 causal=causal, pad_mode=pad_mode),
+            ]
+            mult *= 2
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        model += [
+            act(**activation_params),
+            StreamableConv1d(mult * n_filters, dimension, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class SEANetDecoder(nn.Module):
+    """SEANet decoder.
+
+    Args:
+        channels (int): Audio channels.
+        dimension (int): Intermediate representation dimension.
+        n_filters (int): Base width for the model.
+        n_residual_layers (int): nb of residual layers.
+        ratios (Sequence[int]): kernel size and stride ratios.
+        activation (str): Activation function.
+        activation_params (dict): Parameters to provide to the activation function.
+        final_activation (str): Final activation function after all convolutions.
+        final_activation_params (dict): Parameters to provide to the activation function.
+        norm (str): Normalization method.
+        norm_params (dict): Parameters to provide to the underlying normalization used along with the convolution.
+        kernel_size (int): Kernel size for the initial convolution.
+        last_kernel_size (int): Kernel size for the initial convolution.
+        residual_kernel_size (int): Kernel size for the residual layers.
+        dilation_base (int): How much to increase the dilation with each layer.
+        causal (bool): Whether to use fully causal convolution.
+        pad_mode (str): Padding mode for the convolutions.
+        true_skip (bool): Whether to use true skip connection or a simple.
+            (streamable) convolution as the skip connection in the residual network blocks.
+        compress (int): Reduced dimensionality in residual branches (from Demucs v3).
+        lstm (int): Number of LSTM layers at the end of the encoder.
+        disable_norm_outer_blocks (int): Number of blocks for which we don't apply norm.
+            For the decoder, it corresponds to the N last blocks.
+        trim_right_ratio (float): Ratio for trimming at the right of the transposed convolution under the causal setup.
+            If equal to 1.0, it means that all the trimming is done at the right.
+    """
+    def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 3,
+                 ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},
+                 final_activation: tp.Optional[str] = None, final_activation_params: tp.Optional[dict] = None,
+                 norm: str = 'none', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,
+                 last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,
+                 pad_mode: str = 'reflect', true_skip: bool = True, compress: int = 2, lstm: int = 0,
+                 disable_norm_outer_blocks: int = 0, trim_right_ratio: float = 1.0):
+        super().__init__()
+        self.dimension = dimension
+        self.channels = channels
+        self.n_filters = n_filters
+        self.ratios = ratios
+        del ratios
+        self.n_residual_layers = n_residual_layers
+        self.hop_length = np.prod(self.ratios)
+        self.n_blocks = len(self.ratios) + 2  # first and last conv + residual blocks
+        self.disable_norm_outer_blocks = disable_norm_outer_blocks
+        assert self.disable_norm_outer_blocks >= 0 and self.disable_norm_outer_blocks <= self.n_blocks, \
+            "Number of blocks for which to disable norm is invalid." \
+            "It should be lower or equal to the actual number of blocks in the network and greater or equal to 0."
+
+        act = getattr(nn, activation)
+        mult = int(2 ** len(self.ratios))
+        model: tp.List[nn.Module] = [
+            StreamableConv1d(dimension, mult * n_filters, kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks == self.n_blocks else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+
+        if lstm:
+            model += [StreamableLSTM(mult * n_filters, num_layers=lstm)]
+
+        # Upsample to raw audio scale
+        for i, ratio in enumerate(self.ratios):
+            block_norm = 'none' if self.disable_norm_outer_blocks >= self.n_blocks - (i + 1) else norm
+            # Add upsampling layers
+            model += [
+                act(**activation_params),
+                StreamableConvTranspose1d(mult * n_filters, mult * n_filters // 2,
+                                          kernel_size=ratio * 2, stride=ratio,
+                                          norm=block_norm, norm_kwargs=norm_params,
+                                          causal=causal, trim_right_ratio=trim_right_ratio),
+            ]
+            # Add residual layers
+            for j in range(n_residual_layers):
+                model += [
+                    SEANetResnetBlock(mult * n_filters // 2, kernel_sizes=[residual_kernel_size, 1],
+                                      dilations=[dilation_base ** j, 1],
+                                      activation=activation, activation_params=activation_params,
+                                      norm=block_norm, norm_params=norm_params, causal=causal,
+                                      pad_mode=pad_mode, compress=compress, true_skip=true_skip)]
+
+            mult //= 2
+
+        # Add final layers
+        model += [
+            act(**activation_params),
+            StreamableConv1d(n_filters, channels, last_kernel_size,
+                             norm='none' if self.disable_norm_outer_blocks >= 1 else norm,
+                             norm_kwargs=norm_params, causal=causal, pad_mode=pad_mode)
+        ]
+        # Add optional final activation to decoder (eg. tanh)
+        if final_activation is not None:
+            final_act = getattr(nn, final_activation)
+            final_activation_params = final_activation_params or {}
+            model += [
+                final_act(**final_activation_params)
+            ]
+        self.model = nn.Sequential(*model)
+
+    def forward(self, z):
+        y = self.model(z)
+        return y
+
+
+def exists(val: tp.Optional[tp.Any]) -> bool:
+    return val is not None
+
+
+def default(val: tp.Any, d: tp.Any) -> tp.Any:
+    return val if exists(val) else d
+
+
+def l2norm(t):
+    return F.normalize(t, p=2, dim=-1)
+
+
+def ema_inplace(moving_avg, new, decay: float):
+    moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))
+
+
+def laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):
+    return (x + epsilon) / (x.sum() + n_categories * epsilon)
+
+
+def uniform_init(*shape: int):
+    t = torch.empty(shape)
+    nn.init.kaiming_uniform_(t)
+    return t
+
+
+def sample_vectors(samples, num: int):
+    num_samples, device = samples.shape[0], samples.device
+
+    if num_samples >= num:
+        indices = torch.randperm(num_samples, device=device)[:num]
+    else:
+        indices = torch.randint(0, num_samples, (num,), device=device)
+
+    return samples[indices]
+
+
+def kmeans(samples, num_clusters: int, num_iters: int = 10):
+    dim, dtype = samples.shape[-1], samples.dtype
+
+    means = sample_vectors(samples, num_clusters)
+
+    for _ in range(num_iters):
+        diffs = rearrange(samples, "n d -> n () d") - rearrange(
+            means, "c d -> () c d"
+        )
+        dists = -(diffs ** 2).sum(dim=-1)
+
+        buckets = dists.max(dim=-1).indices
+        bins = torch.bincount(buckets, minlength=num_clusters)
+        zero_mask = bins == 0
+        bins_min_clamped = bins.masked_fill(zero_mask, 1)
+
+        new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)
+        new_means.scatter_add_(0, repeat(buckets, "n -> n d", d=dim), samples)
+        new_means = new_means / bins_min_clamped[..., None]
+
+        means = torch.where(zero_mask[..., None], means, new_means)
+
+    return means, bins
+
+
+def orthogonal_loss_fn(t):
+    # eq (2) from https://arxiv.org/abs/2112.00384
+    n = t.shape[0]
+    normed_codes = l2norm(t)
+    identity = torch.eye(n, device=t.device)
+    cosine_sim = einsum("i d, j d -> i j", normed_codes, normed_codes)
+    return ((cosine_sim - identity) ** 2).sum() / (n ** 2)
+
+
+class EuclideanCodebook(nn.Module):
+    """Codebook with Euclidean distance.
+
+    Args:
+        dim (int): Dimension.
+        codebook_size (int): Codebook size.
+        kmeans_init (bool): Whether to use k-means to initialize the codebooks.
+            If set to true, run the k-means algorithm on the first training batch and use
+            the learned centroids as initialization.
+        kmeans_iters (int): Number of iterations used for k-means algorithm at initialization.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        kmeans_init: int = False,
+        kmeans_iters: int = 10,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        threshold_ema_dead_code: int = 2,
+    ):
+        super().__init__()
+        self.decay = decay
+        init_fn: tp.Union[tp.Callable[..., torch.Tensor], tp.Any] = uniform_init if not kmeans_init else torch.zeros
+        embed = init_fn(codebook_size, dim)
+
+        self.codebook_size = codebook_size
+
+        self.kmeans_iters = kmeans_iters
+        self.epsilon = epsilon
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+
+        self.register_buffer("inited", torch.Tensor([not kmeans_init]))
+        self.register_buffer("cluster_size", torch.zeros(codebook_size))
+        self.register_buffer("embed", embed)
+        self.register_buffer("embed_avg", embed.clone())
+
+    @torch.jit.ignore
+    def init_embed_(self, data):
+        if self.inited:
+            return
+
+        embed, cluster_size = kmeans(data, self.codebook_size, self.kmeans_iters)
+        self.embed.data.copy_(embed)
+        self.embed_avg.data.copy_(embed.clone())
+        self.cluster_size.data.copy_(cluster_size)
+        self.inited.data.copy_(torch.Tensor([True]))
+        # Make sure all buffers across workers are in sync after initialization
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def replace_(self, samples, mask):
+        modified_codebook = torch.where(
+            mask[..., None], sample_vectors(samples, self.codebook_size), self.embed
+        )
+        self.embed.data.copy_(modified_codebook)
+
+    def expire_codes_(self, batch_samples):
+        if self.threshold_ema_dead_code == 0:
+            return
+
+        expired_codes = self.cluster_size < self.threshold_ema_dead_code
+        if not torch.any(expired_codes):
+            return
+
+        batch_samples = rearrange(batch_samples, "... d -> (...) d")
+        self.replace_(batch_samples, mask=expired_codes)
+        flashy.distrib.broadcast_tensors(self.buffers())
+
+    def preprocess(self, x):
+        x = rearrange(x, "... d -> (...) d")
+        return x
+
+    def quantize(self, x):
+        embed = self.embed.t()
+        dist = -(
+            x.pow(2).sum(1, keepdim=True)
+            - 2 * x @ embed
+            + embed.pow(2).sum(0, keepdim=True)
+        )
+        embed_ind = dist.max(dim=-1).indices
+        return embed_ind
+
+    def postprocess_emb(self, embed_ind, shape):
+        return embed_ind.view(*shape[:-1])
+
+    def dequantize(self, embed_ind):
+        quantize = F.embedding(embed_ind, self.embed)
+        return quantize
+
+    def encode(self, x):
+        shape = x.shape
+        # pre-process
+        x = self.preprocess(x)
+        # quantize
+        embed_ind = self.quantize(x)
+        # post-process
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        return embed_ind
+
+    def decode(self, embed_ind):
+        quantize = self.dequantize(embed_ind)
+        return quantize
+
+    def forward(self, x):
+        raise NotImplementedError()
+        shape, dtype = x.shape, x.dtype
+        x = self.preprocess(x)
+        self.init_embed_(x)
+
+        embed_ind = self.quantize(x)
+        embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)
+        embed_ind = self.postprocess_emb(embed_ind, shape)
+        quantize = self.dequantize(embed_ind)
+
+        if self.training:
+            # We do the expiry of code at that point as buffers are in sync
+            # and all the workers will take the same decision.
+            self.expire_codes_(x)
+            ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)
+            embed_sum = x.t() @ embed_onehot
+            ema_inplace(self.embed_avg, embed_sum.t(), self.decay)
+            cluster_size = (
+                laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)
+                * self.cluster_size.sum()
+            )
+            embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)
+            self.embed.data.copy_(embed_normalized)
+
+        return quantize, embed_ind
+
+
+class VectorQuantization(nn.Module):
+    """Vector quantization implementation.
+    Currently supports only euclidean distance.
+
+    Args:
+        dim (int): Dimension
+        codebook_size (int): Codebook size
+        codebook_dim (int): Codebook dimension. If not defined, uses the specified dimension in dim.
+        decay (float): Decay for exponential moving average over the codebooks.
+        epsilon (float): Epsilon value for numerical stability.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int):
+        channels_last (bool): Channels are the last dimension in the input tensors.
+        commitment_weight (float): Weight for commitment loss.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider
+            for orthogonal regularization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+    """
+    def __init__(
+        self,
+        dim: int,
+        codebook_size: int,
+        codebook_dim: tp.Optional[int] = None,
+        decay: float = 0.8,
+        epsilon: float = 1e-5,
+        kmeans_init: bool = False,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        channels_last: bool = False,
+        commitment_weight: float = 1.,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        _codebook_dim: int = default(codebook_dim, dim)
+
+        requires_projection = _codebook_dim != dim
+        self.project_in = (nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity())
+        self.project_out = (nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity())
+
+        self.epsilon = epsilon
+        self.commitment_weight = commitment_weight
+
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+
+        self._codebook = EuclideanCodebook(dim=_codebook_dim, codebook_size=codebook_size,
+                                           kmeans_init=kmeans_init, kmeans_iters=kmeans_iters,
+                                           decay=decay, epsilon=epsilon,
+                                           threshold_ema_dead_code=threshold_ema_dead_code)
+        self.codebook_size = codebook_size
+
+        self.channels_last = channels_last
+
+    @property
+    def codebook(self):
+        return self._codebook.embed
+
+    @property
+    def inited(self):
+        return self._codebook.inited
+
+    def _preprocess(self, x):
+        if not self.channels_last:
+            x = rearrange(x, "b d n -> b n d")
+        return x
+
+    def _postprocess(self, quantize):
+        if not self.channels_last:
+            quantize = rearrange(quantize, "b n d -> b d n")
+        return quantize
+
+    def encode(self, x):
+        x = self._preprocess(x)
+        x = self.project_in(x)
+        embed_in = self._codebook.encode(x)
+        return embed_in
+
+    def decode(self, embed_ind):
+        quantize = self._codebook.decode(embed_ind)
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+        return quantize
+
+    def forward(self, x):
+        device = x.device
+        x = self._preprocess(x)
+
+        x = self.project_in(x)
+        quantize, embed_ind = self._codebook(x)
+
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
+        loss = torch.tensor([0.0], device=device, requires_grad=self.training)
+
+        if self.training:
+            if self.commitment_weight > 0:
+                commit_loss = F.mse_loss(quantize.detach(), x)
+                loss = loss + commit_loss * self.commitment_weight
+
+            if self.orthogonal_reg_weight > 0:
+                codebook = self.codebook
+
+                if self.orthogonal_reg_active_codes_only:
+                    # only calculate orthogonal loss for the activated codes for this batch
+                    unique_code_ids = torch.unique(embed_ind)
+                    codebook = codebook[unique_code_ids]
+
+                num_codes = codebook.shape[0]
+                if exists(self.orthogonal_reg_max_codes) and num_codes > self.orthogonal_reg_max_codes:
+                    rand_ids = torch.randperm(num_codes, device=device)[:self.orthogonal_reg_max_codes]
+                    codebook = codebook[rand_ids]
+
+                orthogonal_reg_loss = orthogonal_loss_fn(codebook)
+                loss = loss + orthogonal_reg_loss * self.orthogonal_reg_weight
+
+        quantize = self.project_out(quantize)
+        quantize = self._postprocess(quantize)
+
+        return quantize, embed_ind, loss
+
+
+class ResidualVectorQuantization(nn.Module):
+    """Residual vector quantization implementation.
+
+    Follows Algorithm 1. in https://arxiv.org/pdf/2107.03312.pdf
+    """
+    def __init__(self, *, num_quantizers, **kwargs):
+        super().__init__()
+        codebook_size = kwargs.pop('codebook_size', None)
+        if codebook_size is None:
+            raise ValueError("codebook_size must be provided in kwargs")
+        if type(codebook_size) != list:
+            codebook_size = [codebook_size] * num_quantizers
+        self.layers = nn.ModuleList(
+            [VectorQuantization(codebook_size=cur_codebook_size, **kwargs) for _,cur_codebook_size in zip(range(num_quantizers), codebook_size)]
+        )
+            
+
+        # self.layers = nn.ModuleList(
+        #     [VectorQuantization(**kwargs) for _ in range(num_quantizers)]
+        # )
+
+    def forward(self, x, n_q: tp.Optional[int] = None):
+        quantized_out = 0.0
+        residual = x
+
+        all_losses = []
+        all_indices = []
+
+        n_q = n_q or len(self.layers)
+
+        for i, layer in enumerate(self.layers[:n_q]):
+            quantized, indices, loss = layer(residual)
+            residual = residual - quantized
+            quantized_out = quantized_out + quantized
+            all_indices.append(indices)
+            all_losses.append(loss)
+
+        out_losses, out_indices = map(torch.stack, (all_losses, all_indices))
+        return quantized_out, out_indices, out_losses
+
+    def encode(self, x: torch.Tensor, n_q: tp.Optional[int] = None) -> torch.Tensor:
+        residual = x
+        all_indices = []
+        n_q = n_q or len(self.layers)
+        for layer in self.layers[:n_q]:
+            indices = layer.encode(residual)
+            quantized = layer.decode(indices)
+            # the original code is below
+            # since quantize has the gradient of residual, according to line 321
+            #  quantize = x + (quantize - x).detach()
+            # the code below will make commitment loss to be 0 for all codebooks except for codebook1
+            # https://github.com/facebookresearch/encodec/issues/25
+            # therefore we change it
+
+            residual = residual - quantized
+            # residual = residual - quantized.detach()
+            # since commitment loss is averaged, the scale of the loss won't get change (not as said in the issue above)
+            all_indices.append(indices)
+        out_indices = torch.stack(all_indices)
+        return out_indices
+
+    def decode(self, q_indices: torch.Tensor) -> torch.Tensor:
+        quantized_out = torch.tensor(0.0, device=q_indices.device)
+        for i, indices in enumerate(q_indices):
+            layer = self.layers[i]
+            quantized = layer.decode(indices)
+            quantized_out = quantized_out + quantized
+        return quantized_out
+
+
+class ResidualVectorQuantizer(BaseQuantizer):
+    """Residual Vector Quantizer.
+
+    Args:
+        dimension (int): Dimension of the codebooks.
+        n_q (int): Number of residual vector quantizers used.
+        q_dropout (bool): Random quantizer drop out at train time.
+        bins (int): Codebook size.
+        decay (float): Decay for exponential moving average over the codebooks.
+        kmeans_init (bool): Whether to use kmeans to initialize the codebooks.
+        kmeans_iters (int): Number of iterations used for kmeans initialization.
+        threshold_ema_dead_code (int): Threshold for dead code expiration. Replace any codes
+            that have an exponential moving average cluster size less than the specified threshold with
+            randomly selected vector from the current batch.
+        orthogonal_reg_weight (float): Orthogonal regularization weights.
+        orthogonal_reg_active_codes_only (bool): Apply orthogonal regularization only on active codes.
+        orthogonal_reg_max_codes (optional int): Maximum number of codes to consider.
+            for orthogonal regularization.
+    """
+    def __init__(
+        self,
+        dimension: int = 256,
+        n_q: int = 8,
+        q_dropout: bool = False,
+        bins: tp.Union[int, tp.List[int]] = 1024,
+        decay: float = 0.99,
+        kmeans_init: bool = True,
+        kmeans_iters: int = 10,
+        threshold_ema_dead_code: int = 2,
+        orthogonal_reg_weight: float = 0.0,
+        orthogonal_reg_active_codes_only: bool = False,
+        orthogonal_reg_max_codes: tp.Optional[int] = None,
+    ):
+        super().__init__()
+        self.max_n_q = n_q
+        self.n_q = n_q
+        self.q_dropout = q_dropout
+        self.dimension = dimension
+        self.bins = bins
+        self.decay = decay
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.orthogonal_reg_weight = orthogonal_reg_weight
+        self.orthogonal_reg_active_codes_only = orthogonal_reg_active_codes_only
+        self.orthogonal_reg_max_codes = orthogonal_reg_max_codes
+        self.vq = ResidualVectorQuantization(
+            dim=self.dimension,
+            codebook_size=self.bins,
+            num_quantizers=self.n_q,
+            decay=self.decay,
+            kmeans_init=self.kmeans_init,
+            kmeans_iters=self.kmeans_iters,
+            threshold_ema_dead_code=self.threshold_ema_dead_code,
+            orthogonal_reg_weight=self.orthogonal_reg_weight,
+            orthogonal_reg_active_codes_only=self.orthogonal_reg_active_codes_only,
+            orthogonal_reg_max_codes=self.orthogonal_reg_max_codes,
+            channels_last=False
+        )
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        n_q = self.n_q
+        if self.training and self.q_dropout:
+            n_q = int(torch.randint(1, self.n_q + 1, (1,)).item())
+        if type(self.bins) == list:
+            bins = self.bins
+        else:
+            bins = [self.bins] * self.n_q
+        bw_per_q = [math.log2(bin) * frame_rate / 1000 for bin in bins]
+        bw = torch.tensor(sum(bw_per_q)).to(x)
+        quantized, codes, commit_loss = self.vq(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return QuantizedResult(quantized, codes, bw, penalty=torch.mean(commit_loss))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified frame rate at the given bandwidth.
+        The RVQ encode method sets the appropriate number of quantizer to use
+        and returns indices for each quantizer.
+        """
+        n_q = self.n_q
+        codes = self.vq.encode(x, n_q=n_q)
+        codes = codes.transpose(0, 1)
+        # codes is [B, K, T], with T frames, K nb of codebooks.
+        return codes
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation."""
+        # codes is [B, K, T], with T frames, K nb of codebooks, vq.decode expects [K, B, T].
+        codes = codes.transpose(0, 1)
+        quantized = self.vq.decode(codes)
+        return quantized
+
+    @property
+    def total_codebooks(self):
+        return self.max_n_q
+
+    @property
+    def num_codebooks(self):
+        return self.n_q
+
+    def set_num_codebooks(self, n: int):
+        assert n > 0 and n <= self.max_n_q
+        self.n_q = n
+
+class DummyQuantizer(BaseQuantizer):
+    """Fake quantizer that actually does not perform any quantization.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor, frame_rate: int):
+        q = x.unsqueeze(1)
+        return QuantizedResult(x, q, torch.tensor(q.numel() * 32 * frame_rate / 1000 / len(x)).to(x))
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """Encode a given input tensor with the specified sample rate at the given bandwidth.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return x.unsqueeze(1)
+
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """Decode the given codes to the quantized representation.
+        In the case of the DummyQuantizer, the codes are actually identical
+        to the input and resulting quantized representation as no quantization is done.
+        """
+        return codes.squeeze(1)
+
+    @property
+    def total_codebooks(self):
+        """Total number of codebooks."""
+        return 1
+
+    @property
+    def num_codebooks(self):
+        """Total number of codebooks."""
+        return self.total_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the number of active codebooks."""
+        raise AttributeError("Cannot override the number of codebooks for the dummy quantizer")
+
+
+class EncodecModel(CompressionModel):
+    """Encodec model operating on the raw waveform.
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        decoder (nn.Module): Decoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 decoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        return self.quantizer.decode(codes)
+
+class EncodecModel_encode_only(CompressionModel):
+    """Encodec model operating on the raw waveform. Encode only, so no decoder
+
+    Args:
+        encoder (nn.Module): Encoder network.
+        quantizer (BaseQuantizer): Quantizer network.
+        frame_rate (int): Frame rate for the latent representation.
+        sample_rate (int): Audio sample rate.
+        channels (int): Number of audio channels.
+        causal (bool): Whether to use a causal version of the model.
+        renormalize (bool): Whether to renormalize the audio before running the model.
+    """
+    # we need assignment to override the property in the abstract class,
+    # I couldn't find a better way...
+    frame_rate: float = 0
+    sample_rate: int = 0
+    channels: int = 0
+
+    def __init__(self,
+                 encoder: nn.Module,
+                 quantizer: BaseQuantizer,
+                 frame_rate: int,
+                 sample_rate: int,
+                 channels: int,
+                 causal: bool = False,
+                 renormalize: bool = False):
+        super().__init__()
+        self.encoder = encoder
+        self.quantizer = quantizer
+        self.frame_rate = frame_rate
+        self.sample_rate = sample_rate
+        self.channels = channels
+        self.renormalize = renormalize
+        self.causal = causal
+        if self.causal:
+            # we force disabling here to avoid handling linear overlap of segments
+            # as supported in original EnCodec codebase.
+            assert not self.renormalize, 'Causal model does not support renormalize'
+
+    @property
+    def total_codebooks(self):
+        """Total number of quantizer codebooks available."""
+        return self.quantizer.total_codebooks
+
+    @property
+    def num_codebooks(self):
+        """Active number of codebooks used by the quantizer."""
+        return self.quantizer.num_codebooks
+
+    def set_num_codebooks(self, n: int):
+        """Set the active number of codebooks used by the quantizer."""
+        self.quantizer.set_num_codebooks(n)
+
+    @property
+    def cardinality(self):
+        """Cardinality of each codebook."""
+        return self.quantizer.bins
+
+    def preprocess(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        scale: tp.Optional[torch.Tensor]
+        if self.renormalize:
+            mono = x.mean(dim=1, keepdim=True)
+            volume = mono.pow(2).mean(dim=2, keepdim=True).sqrt()
+            scale = 1e-8 + volume
+            x = x / scale
+            scale = scale.view(-1, 1)
+        else:
+            scale = None
+        return x, scale
+
+    def postprocess(self,
+                    x: torch.Tensor,
+                    scale: tp.Optional[torch.Tensor] = None) -> torch.Tensor:
+        if scale is not None:
+            assert self.renormalize
+            x = x * scale.view(-1, 1, 1)
+        return x
+
+    def forward(self, x: torch.Tensor, encode=False) -> QuantizedResult:
+        if encode:
+            return self.encode(x)
+        else:
+            raise NotImplementedError("model forward and training is not supported.")
+        assert x.dim() == 3
+        length = x.shape[-1]
+        x, scale = self.preprocess(x)
+
+        emb = self.encoder(x)
+        q_res = self.quantizer(emb, self.frame_rate)
+        out = self.decoder(q_res.x)
+
+        # remove extra padding added by the encoder and decoder
+        assert out.shape[-1] >= length, (out.shape[-1], length)
+        out = out[..., :length]
+
+        q_res.x = self.postprocess(out, scale)
+
+        return q_res
+
+    def encode(self, x: torch.Tensor) -> tp.Tuple[torch.Tensor, tp.Optional[torch.Tensor]]:
+        """Encode the given input tensor to quantized representation along with scale parameter.
+
+        Args:
+            x (torch.Tensor): Float tensor of shape [B, C, T]
+
+        Returns:
+            codes, scale (tuple of torch.Tensor, torch.Tensor): Tuple composed of:
+                codes a float tensor of shape [B, K, T] with K the number of codebooks used and T the timestep.
+                scale a float tensor containing the scale for audio renormalizealization.
+        """
+        assert x.dim() == 3
+        x, scale = self.preprocess(x)
+        emb = self.encoder(x)
+        codes = self.quantizer.encode(emb)
+        return codes, scale
+
+    def decode(self, codes: torch.Tensor, scale: tp.Optional[torch.Tensor] = None):
+        """Decode the given codes to a reconstructed representation, using the scale to perform
+        audio denormalization if needed.
+
+        Args:
+            codes (torch.Tensor): Int tensor of shape [B, K, T]
+            scale (torch.Tensor, optional): Float tensor containing the scale value.
+
+        Returns:
+            out (torch.Tensor): Float tensor of shape [B, C, T], the reconstructed audio.
+        """
+        raise NotImplementedError("Decode is not supported for encode only model")
+        emb = self.decode_latent(codes)
+        out = self.decoder(emb)
+        out = self.postprocess(out, scale)
+        # out contains extra padding added by the encoder and decoder
+        return out
+
+    def decode_latent(self, codes: torch.Tensor):
+        """Decode from the discrete codes to continuous latent space."""
+        raise NotImplementedError("Decode is not supported for encode only model")
+        return self.quantizer.decode(codes)
+
+def get_quantizer(quantizer: str, cfg: omegaconf.DictConfig, dimension: int) -> BaseQuantizer:
+    klass = {
+        'no_quant': DummyQuantizer,
+        'rvq': ResidualVectorQuantizer
+    }[quantizer]
+    kwargs = dict_from_config(getattr(cfg, quantizer))
+    if quantizer != 'no_quant':
+        kwargs['dimension'] = dimension
+    return klass(**kwargs)
+
+def get_encodec_autoencoder(encoder_name: str, cfg: omegaconf.DictConfig):
+    if encoder_name == 'seanet':
+        kwargs = dict_from_config(getattr(cfg, 'seanet'))
+        encoder_override_kwargs = kwargs.pop('encoder')
+        decoder_override_kwargs = kwargs.pop('decoder')
+        encoder_kwargs = {**kwargs, **encoder_override_kwargs}
+        decoder_kwargs = {**kwargs, **decoder_override_kwargs}
+        encoder = SEANetEncoder(**encoder_kwargs)
+        decoder = SEANetDecoder(**decoder_kwargs)
+        return encoder, decoder
+    else:
+        raise KeyError(f"Unexpected compression model {cfg.compression_model}")
+
+
+def get_compression_model(ckpt_fn, encode_only=False, device="cpu") -> CompressionModel:
+    """Instantiate a compression model."""
+    if device == None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    state = torch.load(ckpt_fn, map_location='cpu')
+    cfg = state['xp.cfg']
+    cfg.device = str(device)
+    weights = state['best_state']['model']
+    assert cfg.compression_model == 'encodec', "Only Encodec model is supported for now."
+    if encode_only:
+        all_keys = list(weights.keys())
+        for key in all_keys:
+            if key.startswith('decoder'):
+                del weights[key]
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, _ = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel_encode_only(encoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+    else:
+        kwargs = dict_from_config(getattr(cfg, 'encodec'))
+        encoder_name = kwargs.pop('autoencoder')
+        quantizer_name = kwargs.pop('quantizer')
+        encoder, decoder = get_encodec_autoencoder(encoder_name, cfg)
+        quantizer = get_quantizer(quantizer_name, cfg, encoder.dimension)
+        frame_rate = kwargs['sample_rate'] // encoder.hop_length
+        renormalize = kwargs.pop('renormalize', False)
+        # deprecated params
+        kwargs.pop('renorm', None)
+        compression_model = EncodecModel(encoder, decoder, quantizer,
+                            frame_rate=frame_rate, renormalize=renormalize, **kwargs).to(cfg.device)
+        assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
+        compression_model.load_state_dict(weights)
+        compression_model.eval()
+        return compression_model
+
+if __name__ == "__main__":
+    import torchaudio
+    ckpt_fn = "/home/pyp/BoostedVoiceEditor/pretrained/encodec_6f79c6a8.th"
+    audio_in_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam.wav", "/home/pyp/BoostedVoiceEditor/demo/ray.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean.wav", "/home/pyp/BoostedVoiceEditor/demo/bible.wav", "/home/pyp/BoostedVoiceEditor/demo/miley.wav"]
+    audio_out_fns = ["/home/pyp/BoostedVoiceEditor/demo/pam_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/ray_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/84_121550_000074_000000_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/caribbean_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/bible_encodecTest.wav", "/home/pyp/BoostedVoiceEditor/demo/miley_encodecTest.wav"]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = get_compression_model(ckpt_fn, device=device)
+
+    for audio_in_fn, audio_out_fn in zip(audio_in_fns, audio_out_fns):
+        audio_in, sr = torchaudio.load(audio_in_fn)
+        if sr != model.sample_rate:
+            audio_in = torchaudio.transforms.Resample(sr, model.sample_rate)(audio_in)
+        if audio_in.shape[0] == 2:
+            audio_in = audio_in.mean(dim=0, keepdim=True)
+        audio_in = audio_in.unsqueeze(0)
+        audio_in = audio_in.to(torch.float32).to(device)
+        codes = model.encode(audio_in)[0]
+        audio_out = model.decode(codes)[0].cpu()
+        torchaudio.save(audio_out_fn, audio_out, model.sample_rate)

data/ll60k_preprocessing/step1_download.sh

@@ -0,0 +1,42 @@
+# define where to store the the downloaded data
+dataroot=$DATAROOT
+mkdir -p $dataroot
+manifestroot=$dataroot/libriheavy
+mkdir -p $manifestroot
+audioroot=$dataroot/audio
+mkdir -p $audioroot
+
+# download libriheavy_long and libriheavy
+cd $manifestroot
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_dev.jsonl.gz?download=true -O libriheavy_cuts_dev.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_test_clean.jsonl.gz?download=true -O libriheavy_cuts_test_clean.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_test_other.jsonl.gz?download=true -O libriheavy_cuts_test_other.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_small.jsonl.gz?download=true -O libriheavy_cuts_small.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_medium.jsonl.gz?download=true -O libriheavy_cuts_medium.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy/resolve/main/libriheavy_cuts_large.jsonl.gz?download=true -O libriheavy_cuts_large.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy_long/resolve/main/libriheavy_cuts_small.jsonl.gz?download=true -O libriheavy_long_original_cuts_small.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy_long/resolve/main/libriheavy_cuts_medium.jsonl.gz?download=true -O libriheavy_long_original_cuts_medium.jsonl.gz
+wget https://huggingface.co/datasets/pkufool/libriheavy_long/resolve/main/libriheavy_cuts_large.jsonl.gz?download=true -O libriheavy_long_original_cuts_large.jsonl.gz
+
+# turn .jsonl.gz to .jsonl
+gunzip -k libriheavy_cuts_dev.jsonl.gz
+gunzip -k libriheavy_cuts_test_clean.jsonl.gz
+gunzip -k libriheavy_cuts_test_other.jsonl.gz
+gunzip -k libriheavy_cuts_small.jsonl.gz
+gunzip -k libriheavy_cuts_medium.jsonl.gz
+gunzip -k libriheavy_cuts_large.jsonl.gz
+gunzip -k libriheavy_long_original_cuts_small.jsonl.gz
+gunzip -k libriheavy_long_original_cuts_medium.jsonl.gz
+gunzip -k libriheavy_long_original_cuts_large.jsonl.gz
+
+# if librilight is already unzipped in origDATAROOT, then skip this step
+# download ll
+cd $audioroot
+wget https://dl.fbaipublicfiles.com/librilight/data/small.tar
+wget https://dl.fbaipublicfiles.com/librilight/data/medium.tar
+wget https://dl.fbaipublicfiles.com/librilight/data/large.tar
+
+# untar small, medium, large
+tar -xf small.tar
+tar -xf medium.tar
+tar -xf large.tar

data/ll60k_preprocessing/step2_resplit_long.py

@@ -0,0 +1,51 @@
+# find split, spk, books in libriheavy_cuts_dev.jsonl, libriheavy_cuts_test_clean.jsonl, libriheavy_cuts_test_other.jsonl
+# those would be in "id" field
+
+import sys
+import os, random, numpy as np, socket
+import json
+import tqdm
+def write_jsonl(data, fn):
+    with open(fn, "w") as file:
+        for entry in data:
+            file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+def read_jsonl(file_path):
+    cur_data = []
+    with open(file_path, 'r', encoding='utf-8-sig') as file:
+        for line in file:
+            cur_data.append(json.loads(line.strip()))
+    return cur_data
+import os
+dataroot=os.environ["DATAROOT"]
+manifestroot=os.path.join(dataroot, "libriheavy")
+tgt_names = ['libriheavy_cuts_dev.jsonl', 'libriheavy_cuts_test_clean.jsonl', 'libriheavy_cuts_test_other.jsonl']
+orig_names = ['libriheavy_long_original_cuts_small.jsonl', 'libriheavy_long_original_cuts_medium.jsonl', 'libriheavy_long_original_cuts_large.jsonl']
+
+id2split = {}
+data = read_jsonl(os.path.join(manifestroot, "libriheavy_cuts_dev.jsonl"))
+dev_ids = set(["/".join(item['id'].split("/")[:3]) for item in data])
+data = read_jsonl(os.path.join(manifestroot, "libriheavy_cuts_test_clean.jsonl"))
+test_clean_ids = set(["/".join(item['id'].split("/")[:3]) for item in data])
+data = read_jsonl(os.path.join(manifestroot, "libriheavy_cuts_test_other.jsonl"))
+test_other_ids = set(["/".join(item['id'].split("/")[:3]) for item in data])
+
+long_dev = []
+long_test_clean = []
+long_test_other = []
+for orig_name in orig_names:
+    keep = []
+    data = read_jsonl(os.path.join(manifestroot, orig_name))
+    for item in tqdm.tqdm(data):
+        if "/".join(item['id'].split("/")[:3]) in dev_ids:
+            long_dev.append(item)
+        elif "/".join(item['id'].split("/")[:3]) in test_clean_ids:
+            long_test_clean.append(item)
+        elif "/".join(item['id'].split("/")[:3]) in test_other_ids:
+            long_test_other.append(item)
+        else:
+            keep.append(item)
+    write_jsonl(keep, os.path.join(manifestroot, orig_name.replace("_original", "")))
+
+write_jsonl(long_dev, os.path.join(manifestroot, "libriheavy_long_cuts_dev.jsonl"))
+write_jsonl(long_test_clean, os.path.join(manifestroot, "libriheavy_long_cuts_test_clean.jsonl"))
+write_jsonl(long_test_other, os.path.join(manifestroot, "libriheavy_long_cuts_test_other.jsonl"))

data/ll60k_preprocessing/step3_seg_phn_manifest.py

@@ -0,0 +1,194 @@
+from importlib.resources import path
+import pathlib
+import soundfile as sf
+import numpy as np
+import json
+import multiprocessing
+import argparse
+import tqdm
+import gzip
+import time
+import os
+from tokenizer import TextTokenizer, tokenize_text
+import glob
+import sys
+import os, random, numpy as np, socket
+import json
+import tqdm
+def write_jsonl(data, fn):
+    with open(fn, "w") as file:
+        for entry in data:
+            file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+def read_jsonl(file_path):
+    cur_data = []
+    with open(file_path, 'r', encoding='utf-8-sig') as file:
+        for line in file:
+            cur_data.append(json.loads(line.strip()))
+    return cur_data
+def save_audio(seq, fn):
+    output = seq
+    os.makedirs(os.path.dirname(fn), exist_ok=True)
+    sf.write(fn, output, samplerate=16000)
+
+def save_text(text, fn):
+    os.makedirs(os.path.dirname(fn), exist_ok=True)
+    with open(fn, "w") as wwf:
+        wwf.writelines(text)
+
+def phonemize_and_save(text, fn):
+    phn = tokenize_text(text_tokenizer, text)
+    os.makedirs(os.path.dirname(fn), exist_ok=True)
+    with open(fn, "w") as f:
+        f.write(' '.join(phn))
+    return set(phn)
+
+def cut_sequence(task):
+    in_audio_fn, output_dir, metadata = task
+    if not os.path.isfile(in_audio_fn):
+        # print("missing: ", in_audio_fn)
+        return None
+    data, samplerate = sf.read(in_audio_fn)
+    assert len(data.shape) == 1
+    assert samplerate == 16000
+    all_phns = set()
+    for item in metadata:
+        out_fn = item['file_id']
+        out_audio_fn = os.path.join(output_dir, "audio", out_fn)
+        out_text_fn = os.path.join(output_dir, "audio", out_fn.replace(".flac", ".txt"))
+        out_phn_fn = os.path.join(output_dir, "phoneme", out_fn.replace(".flac", ".txt"))
+        save_audio(data[int(item['vad'][0]*samplerate):int(item['vad'][1]*samplerate)], out_audio_fn)
+        save_text(item['text'], out_text_fn)
+        phns = phonemize_and_save(item['text'], out_phn_fn)
+        all_phns.update(phns)
+    
+    return all_phns
+
+
+from collections import defaultdict
+# Function to create a defaultdict recursively
+def nested_defaultdict(levels, inner_type):
+    if levels <= 1:
+        return defaultdict(inner_type)
+    return defaultdict(lambda: nested_defaultdict(levels-1, inner_type))
+
+
+def open_mani(fn):
+    print("load segmentation and transcription metadata...")
+    stime = time.time()
+    data = []
+    with gzip.open(fn, 'rt', encoding='utf-8') as f:
+        for line in f:
+            data.append(json.loads(line))
+    print(f"loading done, took {time.time() - stime:.4f} seconds")
+    return data
+
+def cut(split,
+        audio_dir,
+        mani_dir,
+        output_dir,
+        n_process=32,
+        percent=0.5):
+    split2manifest = {
+            "train": [
+                "libriheavy_long_cuts_small.jsonl", 
+                "libriheavy_long_cuts_medium.jsonl", 
+                "libriheavy_long_cuts_large.jsonl",
+                "libriheavy_cuts_small.jsonl", 
+                "libriheavy_cuts_medium.jsonl", 
+                "libriheavy_cuts_large.jsonl",
+            ],
+            "valid": [
+                    "libriheavy_cuts_dev.jsonl",
+                    "libriheavy_long_cuts_dev.jsonl"
+            ],
+            "test": [
+                    "libriheavy_cuts_test_clean.jsonl",
+                    "libriheavy_cuts_test_other.jsonl",
+                    "libriheavy_long_cuts_test_clean.jsonl",
+                    "libriheavy_long_cuts_test_other.jsonl"
+            ]
+        }
+
+    print("organize data by recording_id (i.e. the original big .flac file name)...")
+    stime = time.time()
+    organized_data = nested_defaultdict(4, list)
+    manifest_fn = os.path.join(output_dir, "manifest_mimi", split+".txt")
+    os.makedirs(os.path.join(output_dir, "manifest_mimi"), exist_ok=True)
+    with open(manifest_fn, "w") as wf:
+        for mani_fn in split2manifest[split]:
+            # data = open_mani(os.path.join(mani_dir, mani_fn))
+            data = read_jsonl(os.path.join(mani_dir, mani_fn))
+            for item in data:
+                file_id = item['supervisions'][0]['id'] + '.flac'
+                recording_id = item['recording']['id'] + '.flac'
+                sizeSplit, spk, book, flac = recording_id.split("/") # e.g. 'medium/100/emerald_city_librivox_64kb_mp3/emeraldcity_01_baum_64kb'
+                if os.path.isfile(os.path.join(audio_dir, recording_id)):
+                    vad = (item['start'], item['start']+item['duration'])
+                    text = item['supervisions'][0]['custom']['texts'][0]
+                    file_id = file_id.replace(".flac", "") + f"_{vad[0]:.2f}_{vad[1]:.2f}.flac"
+                    organized_data[sizeSplit][spk][book][recording_id].append({"file_id": file_id, "vad":vad, "text": text})
+                    wf.writelines(f"{file_id}\t{item['duration']}\n")
+    
+    # #### take only a subet of tasks
+    tasks = [(os.path.join(audio_dir, recording_id), output_dir, organized_data[sizeSplit][spk][book][recording_id], spk) for sizeSplit in organized_data for spk in organized_data[sizeSplit] for book in organized_data[sizeSplit][spk] for recording_id in organized_data[sizeSplit][spk][book]]
+    ntasks = len(tasks)
+    spk2tasks = defaultdict(list)
+    for task in tasks:
+        spk2tasks[task[3]].append(task)
+    # randomly shuffle each task list for each speaker
+    for spk in spk2tasks:
+        random.shuffle(spk2tasks[spk])
+    # take only 20% of the tasks, uniformly sampled from each speaker
+    # randomly pick a speaker, and then randomly pick a task from that speaker
+    tasks = []
+    while len(tasks) < ntasks * percent:
+        spk = random.choice(list(spk2tasks.keys()))
+        if len(spk2tasks[spk]) == 0:
+            continue
+        tasks.append(spk2tasks[spk].pop()[:-1])
+    print(f"take only {percent*100:.2f}% of the tasks, {len(tasks)} out of {ntasks} tasks")
+    #### take only a subet of tasks
+
+    print(f"organizing done, took {time.time() - stime:.4f} seconds")
+    print(f"Launching {n_process} processes")
+    phn_vocab = set()
+    cnt = 0
+    with multiprocessing.Pool(processes=n_process) as pool:
+        for phns in tqdm.tqdm(pool.imap_unordered(cut_sequence, tasks), total=len(tasks)):
+            cnt += 1
+            if phns != None:
+                phn_vocab.update(phns)
+
+    # save phn vocabulary
+    if split == "train":
+        vocab_fn = os.path.join(output_dir, "vocab.txt")
+        with open(vocab_fn, "w") as f:
+            for i, phn in enumerate(list(phn_vocab)):
+                if i < len(list(phn_vocab)) - 1:
+                    f.write(f"{str(i)}\t{phn}\n")
+                else:
+                    f.write(f"{str(i)}\t{phn}")
+    
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Cut a dataset in small "
+                                     "sequences using VAD files")
+    parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'], help="train = libriheavy_cuts_{small,medium,large}.jsonl.gz, valid = libriheavy_cuts_dev_{clean,other}.jsonl.gz, test = libriheavy_cuts_test_{clean,other}.jsonl.gz")
+    parser.add_argument('--audio_dir', type=str, default="/data/scratch/pyp/datasets/librilight_example",
+                        help="Path to the audio directory")
+    parser.add_argument('--manifest_dir', type=str, default="/data/scratch/pyp/datasets/librilight/libriheavy", help="path to the transcription file's dir, can be downloaded https://huggingface.co/datasets/pkufool/libriheavy/tree/main/v0.1")
+    parser.add_argument('--output_dir', type=str, default="/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed",
+                        help="Path to the output directory")
+    parser.add_argument('--n_workers', type=int, default=16,
+                        help="Number of parallel worker processes")
+    parser.add_argument('--percent', type=float, default=0.5, help="take only this percent of the tasks, randomly sampled from each speaker")
+
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    pathlib.Path(args.output_dir).mkdir(exist_ok=True, parents=True)
+    text_tokenizer = TextTokenizer()
+    cut(args.split, args.audio_dir, args.manifest_dir, args.output_dir, args.n_workers, args.percent)

data/ll60k_preprocessing/step4_encodec_encode.py

@@ -0,0 +1,184 @@
+import argparse
+from email.policy import default
+def parse_args():
+    parser = argparse.ArgumentParser(description="encode the librilight dataset using codec model")
+    parser.add_argument('--dir', type=str, default="/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed", help="Path to the directory")
+    parser.add_argument('--sub_root', type=str, default="preprocessed", help="sub directory")
+    parser.add_argument('--encodec_name', type=str, default="encodec_6f79c6a8.th", help="name of the codec model")
+    parser.add_argument('--n_workers', type=int, default=16, help="Number of parallel worker processes")
+    parser.add_argument('--batch_size', type=int, default=16, help="batch size for codec encoding, decrease it if OOM. This is the sum of batch size *over each gpu*, so increase it if you are using more gpus")
+    parser.add_argument('--audio_sr', type=int, default=16000, help='input audio sample rate')
+    parser.add_argument('--model_sr', type=int, default=16000, help='encodec input audio sample rate')
+    parser.add_argument('--downsample_rate', type=int, default=320, help='encodec downsample rate')
+    parser.add_argument('--model_code_sr', type=float, default=50, help='codec model code sample rate')
+    parser.add_argument('--len_cap', type=float, default=1000, help='will drop audios that are longer than this number')
+    parser.add_argument('--min_len', type=float, default=0.5, help='will drop audios that are shorter than this number')
+    parser.add_argument('--partition', type=str, default="1/1", help='split for parallel processing')
+    parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'])
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    import logging
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+
+    import os, sys
+    import numpy as np
+    import torch
+    import torchaudio
+    import tqdm
+    import time
+
+    args = parse_args()
+
+    def sort_by_audio_len(lens):
+        inds = np.argsort(lens).tolist()
+        if len(inds) < 10:
+            return inds[::-1]
+        logging.info(f"longest: {lens[inds[-1]]/args.downsample_rate} encodec codes, {lens[inds[-1]]/args.model_sr:.2f} sec.")
+        logging.info(f"shortest: {lens[inds[0]]/args.downsample_rate} encodec codes, {lens[inds[0]]/args.model_sr:.2f} sec.")
+        logging.info(f"median: {lens[inds[len(inds)//2]]/args.downsample_rate} encodec codes, {lens[inds[len(inds)//2]]/args.model_sr:.2f} sec.")
+        logging.info(f"95 percentile longest: {lens[inds[int(len(inds)*0.95)]]/args.downsample_rate} encodec codes, {lens[inds[int(len(inds)*0.95)]]/args.model_sr:.2f} sec.")
+        return inds[::-1]
+
+    def write_array_to_txt_file(array, filename):
+        with open(filename, 'w') as f:
+            for a in array[:-1]:
+                f.write(' '.join(map(str, a))+'\n')
+            f.write(' '.join(map(str, array[-1])))
+
+    class mydataset(torch.utils.data.Dataset):
+        def __init__(self, split):
+            super().__init__()
+            self.split = split
+            self.audio_dir = audio_dir
+            manifest_fn = os.path.join(encodec_manifest_dir, split+".txt")
+            cur_sp = int(args.partition.split("/")[0])-1
+            total_sp = int(args.partition.split("/")[1])
+            with open(manifest_fn, "r") as rf:
+                self.data = [l.strip().split("\t") for l in rf.readlines()][cur_sp::total_sp]
+            self.data = [l for l in self.data if os.path.isfile(os.path.join(self.audio_dir, l[0]))]
+        def __len__(self):
+            return len(self.data)
+        def __getitem__(self, ind):
+            try:
+                afn = self.data[ind][0]
+                fn = os.path.join(self.audio_dir, afn)
+                audio, sr = torchaudio.load(fn)
+                if sr != args.model_sr:
+                    audio = torchaudio.transforms.Resample(sr, args.model_sr)(audio)
+                    sr = args.model_sr
+                assert sr == args.model_sr, sr
+            except Exception as e:
+                # logging.info(f"{e}")
+                return None, None, None
+            assert audio.ndim==2 and audio.shape[0] == 1, audio.shape
+            return audio.type(torch.float32).squeeze(0), audio.shape[-1], os.path.splitext(afn)[0]
+        def collate(self, batch):
+            lens, audios, segment_ids = [], [], []
+            for item in batch:
+                if item[0] != None:
+                    audios.append(item[0])
+                    lens.append(item[1])
+                    segment_ids.append(item[2])
+            return audios, lens, segment_ids
+
+    # roots
+    sub_root = args.sub_root
+    encodec_manifest_dir = os.path.join(args.dir, sub_root, "manifest_mimi")
+    audio_dir = os.path.join(args.dir, sub_root, "audio")
+    save_manifest_dir = os.path.join(args.dir, sub_root,"manifest_final_encodec")
+    if args.encodec_name == "encodec_6f79c6a8.th":
+        save_codes_dir = os.path.join(args.dir, sub_root,"encodec_4cb")
+    elif args.encodec_name == "encodec_8cb1024_giga.th":
+        save_codes_dir = os.path.join(args.dir, sub_root,"encodec_8cb")
+
+    os.makedirs(save_manifest_dir, exist_ok=True)
+    os.makedirs(save_codes_dir, exist_ok=True)
+    
+    # load the encodec model
+    def import_encodec():
+        from encodec import get_compression_model
+        userdir = os.path.expanduser("~")
+        model = get_compression_model(os.path.join(userdir, "VoiceStar", f"pretrained/{args.encodec_name}"), encode_only=True, device="cuda")
+        model = torch.nn.DataParallel(model)
+        return model
+    model = import_encodec()
+    
+    # setup dataloader
+    mega_batch_size = 1024
+    batch_size = args.batch_size
+    
+    dataset = mydataset(args.split)
+    if len(dataset) == 0:
+        logging.info(f"no data found for split {args.split} partition {args.partition}")
+        sys.exit(0)
+    loader = torch.torch.utils.data.DataLoader(dataset, batch_size=mega_batch_size, shuffle=False, drop_last=False, num_workers=args.n_workers, collate_fn=dataset.collate)
+    split = args.split
+
+    skip = 0
+    logging.info(f"now processing split {split} partition {args.partition}...")
+    mega_n_steps = int(np.ceil(len(loader.dataset) / mega_batch_size))
+    # mega_n_steps = int(np.ceil(len(gs) / mega_batch_size))
+    logging.info(f"partition the split {split} into {mega_n_steps} parts, each has at most {mega_batch_size} samples")
+    mani_fn = os.path.join(save_manifest_dir, f"{split}_{args.partition.replace('/', '=')}.txt")
+    logging.info(f"manifest for split {split} partition {args.partition.replace('/', '=')}.txt will be saved at {mani_fn}")
+    with open(mani_fn, "w") as mani_wf:
+    # with open(mani_fn, "a") as mani_wf: # resume from where we failed
+        for m, mega_batch in enumerate(tqdm.tqdm(loader)):
+
+            logging.info(f"====================================")
+            logging.info(f"====================================")
+            logging.info(f"now processing mega step {m+1}/{mega_n_steps}")
+            try:
+            # if True:
+                lengths = np.array(mega_batch[1])
+                if len(lengths) == 0: # the loader might not find any audio because step3 will write to manifest first, and then might selection a subset to cut and save audio
+                    continue
+                sorted_inds = sort_by_audio_len(lengths)
+                for j in range(len(sorted_inds))[::-1]:
+                    if lengths[sorted_inds[j]] < args.model_sr*args.min_len or lengths[sorted_inds[j]] > args.model_sr*args.len_cap: # skip samples that are too short (shorter than 0.2s), or too big (bigger than 80s)
+                        skip += 1
+                        del sorted_inds[j]
+                
+                n_steps = int(np.ceil(len(sorted_inds) / batch_size))
+                for n in tqdm.tqdm(range(n_steps), disable=True):
+                    inds_used = sorted_inds[n*batch_size:(n+1)*batch_size]
+                    while len(inds_used) < batch_size:
+                        inds_used += sorted_inds[:batch_size-len(inds_used)]
+                    wav_batch = [mega_batch[0][id] for id in inds_used]
+                    all_lens = [mega_batch[1][id] for id in inds_used]
+                    segment_id_batch = [mega_batch[2][id] for id in inds_used]
+                    padded_wav = torch.nn.utils.rnn.pad_sequence(wav_batch, batch_first=True).unsqueeze(1) # [B, T] -> [B, 1, T]
+                    # Extract discrete codes from EnCodec
+                    with torch.no_grad():
+                        if max(all_lens) > 300000 and len(all_lens) > 1: # if utterances are long, simply pass half of them at a time
+                            codes = []
+                            inwav = padded_wav.cuda()
+                            codes.append(model(inwav[:len(inwav)//2])[0].cpu())
+                            codes.append(model(inwav[len(inwav)//2:])[0].cpu())
+                            codes = torch.cat(codes, dim=0)
+                        else:
+                            encoded_frames = model(padded_wav.cuda()) 
+                            codes = encoded_frames[0].cpu() # [B, n_codebook, T]
+
+                    for i, length in enumerate(all_lens):
+                        save_fn = os.path.join(save_codes_dir, segment_id_batch[i]+".txt")
+                        actual_len = round(length / args.downsample_rate) # 320 is downsample rate for this model
+                        cur_code = codes[i].tolist() if type(codes) == list else codes[i, :, :actual_len].tolist()
+                        os.makedirs(os.path.dirname(save_fn), exist_ok=True)
+                        write_array_to_txt_file(cur_code, save_fn)
+
+                        mani_wf.write(f"{segment_id_batch[i]}\t{len(cur_code[0])}\n") # write to manifest file
+                        # if i == 10:
+                        #    raise
+            except Exception as e:
+                print(f'exception!! at {m+1}')
+                print(e)
+                continue
+
+            # break
+    logging.info(f"split {split} partition {args.partition} has {len(loader.dataset)} samples in total, skipped {skip} due to utterance being too long or too short")
+        # break

data/ll60k_preprocessing/step4_encodec_encode_script.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate voicestar
+
+dir=${dir:-/data/scratch/pyp/datasets/librilight}
+sub_root=${sub_root:-preprocessed}
+encodec_name=${encodec_name:-encodec_6f79c6a8.th} # or encodec_8cb1024_giga.th
+n_workers=${n_workers:-12}
+batch_size=${batch_size:-64}
+audio_sr=16000
+model_sr=16000
+downsample_rate=320
+model_code_sr=50
+len_cap=1000
+min_len=0.5
+partition=${partition:-"1/1"}
+split=${split:-"valid"}
+
+python step4_encodec_encode.py --dir $dir --sub_root ${sub_root} --encodec_name ${encodec_name} --n_workers $n_workers --batch_size $batch_size --audio_sr $audio_sr --model_sr $model_sr --downsample_rate $downsample_rate --model_code_sr $model_code_sr --len_cap $len_cap --min_len $min_len --partition $partition --split $split

data/ll60k_preprocessing/step5_find_nearest_neighbor.py

@@ -0,0 +1,157 @@
+# for each each audio segment, find the non-overlapping neighboring segments
+from importlib.resources import path
+import pathlib
+import soundfile as sf
+import numpy as np
+import json
+import multiprocessing
+import argparse
+import tqdm
+import gzip
+import time
+import os
+from tokenizer import TextTokenizer, tokenize_text
+import glob
+import sys
+import os, random, numpy as np, socket
+import json
+import tqdm
+import json
+import tqdm
+def write_jsonl(data, fn):
+    with open(fn, "w") as file:
+        for entry in data:
+            file.write(json.dumps(entry, ensure_ascii=False) + "\n")
+def read_jsonl(file_path):
+    cur_data = []
+    with open(file_path, 'r', encoding='utf-8-sig') as file:
+        for line in file:
+            cur_data.append(json.loads(line.strip()))
+    return cur_data
+from collections import defaultdict
+# Function to create a defaultdict recursively
+def nested_defaultdict(levels, inner_type):
+    if levels <= 1:
+        return defaultdict(inner_type)
+    return defaultdict(lambda: nested_defaultdict(levels-1, inner_type))
+
+def find_neighbor(args):
+    split2manifest = {
+            "train": [
+                "libriheavy_cuts_small.jsonl", 
+                "libriheavy_cuts_medium.jsonl", 
+                "libriheavy_cuts_large.jsonl",
+                "libriheavy_long_cuts_small.jsonl", 
+                "libriheavy_long_cuts_medium.jsonl", 
+                "libriheavy_long_cuts_large.jsonl"
+            ],
+            "valid": [
+                    "libriheavy_cuts_dev.jsonl",
+                    "libriheavy_long_cuts_dev.jsonl"
+            ],
+            "test": [
+                    "libriheavy_cuts_test_clean.jsonl",
+                    "libriheavy_cuts_test_other.jsonl",
+                    "libriheavy_long_cuts_test_clean.jsonl",
+                    "libriheavy_long_cuts_test_other.jsonl"
+            ]
+        }
+
+    stime = time.time()
+    organized_data = nested_defaultdict(4, list)
+    for mani_fn in split2manifest[args.split]:
+        # data = open_mani(os.path.join(mani_dir, mani_fn))
+        mani_full_fn = os.path.join(args.manifest_dir, mani_fn)
+        data = read_jsonl(mani_full_fn)
+        for item in data:
+            file_id = item['supervisions'][0]['id'] + '.flac'
+            recording_id = item['recording']['id'] + '.flac'
+            sizeSplit, spk, book, flac = recording_id.split("/") # e.g. 'medium/100/emerald_city_librivox_64kb_mp3/emeraldcity_01_baum_64kb'
+            if os.path.isfile(os.path.join(args.audio_dir, recording_id)):
+                vad = (item['start'], item['start']+item['duration'])
+                text = item['supervisions'][0]['custom']['texts'][0]
+                file_id = file_id.replace(".flac", "") + f"_{vad[0]:.2f}_{vad[1]:.2f}.flac"
+                organized_data[sizeSplit][spk][book][recording_id].append({"file_id": file_id, "vad":vad, "text": text})
+    
+    # # for each recording_id, find the non-overlapping neighboring segments based on vad
+    # for sizeSplit in organized_data:
+    #     for spk in organized_data[sizeSplit]:
+    #         for book in organized_data[sizeSplit][spk]:
+    #             for recording_id in organized_data[sizeSplit][spk][book]:
+    #                 segments = organized_data[sizeSplit][spk][book][recording_id]
+    #                 segments.sort(key=lambda x: x['vad'][0])
+    #                 for i in range(len(segments)):
+    #                     # for segment i, find the non-overlapping neighboring segments
+    #                     write_fn = os.path.join(args.output_dir, f"{segments[i]['file_id'].replace('.flac', '.txt')}")
+    #                     neighbors = []
+    #                     distance = []
+    #                     for j in range(len(segments)):
+    #                         if segments[i]['vad'][1] < segments[j]['vad'][0] or segments[i]['vad'][0] > segments[j]['vad'][0]:
+    #                             neighbors.append(segments[j]['file_id'].replace('.flac', '.txt'))
+    #                             distance.append(min(abs(segments[i]['vad'][1] - segments[j]['vad'][0]), abs(segments[i]['vad'][0] - segments[j]['vad'][1])))
+    #                     # order neighbors by distance
+    #                     neighbors_distance = [[x, dist] for dist, x in sorted(zip(distance, neighbors))]
+    #                     os.makedirs(os.path.dirname(write_fn), exist_ok=True)
+    #                     with open(write_fn, "w") as f:
+    #                         # note that there might be no neighbors, in which case the file is empty
+    #                         for neighbor, dist in neighbors_distance:
+    #                             f.write(f"{neighbor}\t{dist}\n")
+
+    # use multiprocessing.Pool for the above
+    segments = [organized_data[sizeSplit][spk][book][recording_id] for sizeSplit in organized_data for spk in organized_data[sizeSplit] for book in organized_data[sizeSplit][spk] for recording_id in organized_data[sizeSplit][spk][book]]
+    # only keep those that are exist
+    print(f"originally total {len(segments)} segments")
+    segments = [seg for seg in segments if os.path.isfile(os.path.join("/".join(args.output_dir.split("/")[:-1]),"audio", seg[0]['file_id']))]
+    print(f"after check existance, total {len(segments)} segments")
+    print(f"organizing took {(time.time()-stime)/60:.2f} minutes")
+    with multiprocessing.Pool(processes=args.n_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(find_neighbor_each, segments), total=len(segments)):
+            pass
+
+# audio_root = "/data/scratch/pyp/datasets/librilight/preprocessed/audio"
+def find_neighbor_each(segments):
+    # for each recording_id, find the non-overlapping neighboring segments based on vad
+    # only keep segments that have audio files
+    # actually only keep segments that have ipa_alignment files
+    segments = [seg for seg in segments if os.path.isfile(os.path.join("/".join(args.output_dir.split("/")[:-1]),"ipa_alignment", seg['file_id'].replace(".flac", ".txt")))]
+    if len(segments) <= 1:
+        return
+    for i in range(len(segments)):
+        # for segment i, find the non-overlapping neighboring segments
+        write_fn = os.path.join(args.output_dir, f"{segments[i]['file_id'].replace('.flac', '.txt')}")
+        neighbors = []
+        distance = []
+        for j in range(len(segments)):
+            if segments[i]['vad'][1] < segments[j]['vad'][0] or segments[i]['vad'][0] > segments[j]['vad'][0]:
+                neighbors.append(segments[j])
+                distance.append(min(abs(segments[i]['vad'][1] - segments[j]['vad'][0]), abs(segments[i]['vad'][0] - segments[j]['vad'][1])))
+        if len(neighbors) == 0:
+            continue
+        # order neighbors by distance
+        index = np.argsort(distance)
+        neighbors_distance = [[neighbors[ind], distance[ind]] for ind in index]
+        os.makedirs(os.path.dirname(write_fn), exist_ok=True)
+        with open(write_fn, "w") as f:
+            # note that there might be no neighbors, in which case the file is empty
+            for neighbor, dist in neighbors_distance:
+                f.write(f"{neighbor['file_id'].replace('.flac', '.txt')}\t{dist}\t{neighbor['vad'][1] - neighbor['vad'][0]}\n") # file_id, distance, duration
+
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Cut a dataset in small "
+                                     "sequences using VAD files")
+    parser.add_argument('--split', type=str, default='train', choices=['train', 'valid', 'test'], help="train = libriheavy_cuts_{small,medium,large}.jsonl.gz, valid = libriheavy_cuts_dev_{clean,other}.jsonl.gz, test = libriheavy_cuts_test_{clean,other}.jsonl.gz")
+    parser.add_argument('--audio_dir', type=str, default="/data/scratch/pyp/datasets/librilight_example",
+                        help="Path to the audio directory")
+    parser.add_argument('--manifest_dir', type=str, default="/data/scratch/pyp/datasets/librilight/libriheavy", help="path to the transcription file's dir, can be downloaded https://huggingface.co/datasets/pkufool/libriheavy/tree/main/v0.1")
+    parser.add_argument('--output_dir', type=str, default="/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed/neighbors",
+                        help="Path to the output directory")
+    parser.add_argument('--n_workers', type=int, default=16,
+                        help="Number of parallel worker processes")
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+    pathlib.Path(args.output_dir).mkdir(exist_ok=True, parents=True)
+    find_neighbor(args)

data/ll60k_preprocessing/step6_forced_alignment.py

@@ -0,0 +1,86 @@
+import os, sys
+import subprocess, tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+def align_folders(audio_root, subfolder, subsubfolder):
+    # Construct output folder path
+    file_root = os.path.dirname(audio_root)
+    out_folder = f"{file_root}/alignment/{subfolder}/{subsubfolder}"
+    
+    # Create the output directory
+    os.makedirs(out_folder, exist_ok=True)
+    
+    # Construct the MFA align command
+    command = [
+        "mfa", "align", "--single_speaker", "-j", "8", "--clean",
+        f"{audio_root}/{subfolder}/{subsubfolder}", "english_us_arpa", "english_us_arpa",
+        out_folder, "--beam", "50", "--retry_beam", "400", "--output_format", "csv"
+    ]
+    
+    # Run the command
+    subprocess.run(command, check=True)
+
+def main(file_root = "/data/scratch/pyp/datasets/librilight/librilight_example_preprocessed", max_parallel_jobs=10, max_spk=100, partition="1/10", n_workers=64):
+    # Find all subfolder/subsubfolder combinations
+    tasks = []
+    audio_root = os.path.join(file_root, "audio")
+    for subfolder in os.listdir(audio_root):
+        subfolder_path = os.path.join(audio_root, subfolder)
+        if os.path.isdir(subfolder_path):
+            for subsubfolder in os.listdir(subfolder_path):
+                subsubfolder_path = os.path.join(subfolder_path, subsubfolder)
+                if os.path.isdir(subsubfolder_path):
+                    tasks.append((audio_root, subfolder, subsubfolder))
+    speaker_folder_map = {}
+    for audio_root, subfolder, subsubfolder in tasks:
+        if os.path.join(audio_root, subfolder) not in speaker_folder_map:
+            speaker_folder_map[os.path.join(audio_root, subfolder)] = [os.path.join(audio_root, subfolder, subsubfolder)]
+        else:
+            speaker_folder_map[os.path.join(audio_root, subfolder)].append(os.path.join(audio_root, subfolder, subsubfolder))
+    speaker_folder_partitions = []
+    for audio_root_subfolder, speaker_folders in speaker_folder_map.items():
+        speaker_folder_partitions.extend([speaker_folders[i:i+max_spk] for i in range(0, len(speaker_folders), max_spk)])
+    s, e = partition.split("/")
+    s, e = int(s)-1, int(e)
+    cur_tasks = speaker_folder_partitions[s::e]
+    import secrets, string
+    import soundfile, glob
+    from joblib import Parallel, delayed
+    def delete_corrupted(fn):
+        try:
+            x = soundfile.read(fn)
+        except:
+            print(f"removing corrupted file: {fn}")
+            os.remove(fn)
+
+    for j, task in enumerate(tqdm.tqdm(cur_tasks)):
+        # get subfolder for the current task
+        subs = [item.split("/")[-2] for item in task]
+        # assert that all subs are the same
+        assert len(set(subs)) == 1, subs
+        sub = subs[0]
+        # randomly generate a foldername 
+        # generate a random character
+        # make softlink from item in task to temp folder
+        random_string = ''.join(secrets.choice(string.ascii_letters + string.digits) for i in range(10))
+        temp_folder = os.path.join(file_root, "softlink_audio", random_string)
+        os.makedirs(temp_folder, exist_ok=True)
+        out_folder = f"{file_root}/alignment/{sub}"
+        all_out_speaker_folders = [os.path.join(out_folder, os.path.basename(item)) for item in task]
+        if sum(os.path.isdir(curpath) for curpath in all_out_speaker_folders) == len(all_out_speaker_folders):
+            continue
+        # remove audio files that are corrupted
+        all_audio_files = [audiofile for item in task for audiofile in glob.glob(item+"/*/*.flac")]
+        Parallel(n_jobs=n_workers)(delayed(delete_corrupted)(audiofn) for audiofn in all_audio_files)
+        for item in task:
+            # make softlink from subsubfolder to a new folder in temp folder
+            os.symlink(item, os.path.join(temp_folder, os.path.basename(item)))
+        # run mfa on the linked folder, but save alignment to the correct folder
+        command = f"mfa align -j {n_workers} {temp_folder} english_us_arpa english_us_arpa {out_folder} --beam 50 --retry_beam 200 --output_format csv --quiet --use_mp --temporary_directory {temp_folder}_temp"
+        os.system(command)
+        # delete the temp_folder
+        os.system(f"rm -r {temp_folder}")
+
+if __name__ == "__main__":
+    import fire
+    fire.Fire(main)

data/ll60k_preprocessing/step6_forced_alignment.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+source ~/miniconda3/etc/profile.d/conda.sh
+conda activate voicecraft
+
+partition=$1
+file_root=$2
+max_spk=${max_spk:-100}
+n_workers=${n_workers:-64}
+python step6_forced_alignment.py \
+--partition $partition \
+--file_root $file_root \
+--max_spk $max_spk \
+--n_workers $n_workers

data/ll60k_preprocessing/step7_ipa_alignment.py

@@ -0,0 +1,114 @@
+# we have raw transcript at 
+# /data/scratch/pyp/datasets/librilight/preprocessed/audio
+# we have word and ARPA alignment at 
+# /data/scratch/pyp/datasets/librilight/preprocessed/alignment
+
+# we have manifest at /data/scratch/pyp/datasets/librilight/preprocessed/manifest_mimi
+# where each row is like large/10022/essayoncriticism_1505_librivox_64kb_mp3/essayoncriticism_01_pope_64kb_5_610.32_630.08.flac	19.76
+
+# we want to create IPA alignment from the raw transcript and word alignment, using phonemizer
+# save at /data/scratch/pyp/datasets/librilight/preprocessed/ipa_alignment
+
+# since ipa phonemized results are 1-to-1 with words (10 words might lead to a ipa sequence of 7 phonemes), we have to run phonemizer on each segment of the word sequence
+import os, string, csv, random, tqdm, glob
+from tokenizer import TextTokenizer, tokenize_text
+
+
+def remove_punctuation(input_string):
+    translator = str.maketrans('', '', string.punctuation)
+    return input_string.translate(translator)
+
+
+
+def create_alignment(fn, trans_dir, align_dir, audio_ext, trans_ext, arpa_ext, text_tokenizer, use_prob, ipa_alignment_fn, save=False, prompt_dur=30):
+    os.makedirs(os.path.dirname(ipa_alignment_fn), exist_ok=True)
+    trans_fn = os.path.join(trans_dir, fn.replace(audio_ext, trans_ext))
+    if not os.path.isfile(trans_fn):
+        return [], True
+    align_fn = os.path.join(align_dir, fn.replace(audio_ext, arpa_ext))
+    if not os.path.isfile(align_fn):
+        return [], True
+    # get raw transcript
+    with open(trans_fn, 'r') as f:
+        transcript = f.read().strip()
+    raw_word_list = transcript.split(" ")
+    # get word alignment
+    with open(align_fn, 'r') as f:
+        word_alignment = csv.reader(f)
+        word_alignment = [row for row in word_alignment if row[3]=='words']
+
+    ipa_alignment = []
+
+    for j, (item, raw_word) in enumerate(zip(word_alignment, raw_word_list)):
+        start, end, word = float(item[0]), float(item[1]), item[2]
+        if end > prompt_dur:
+            break
+        punc_re_raw_word = remove_punctuation(raw_word)
+        if not remove_punctuation(word).lower() == punc_re_raw_word.lower():
+            # print(f"word from alignment csv: {word}, word from txt: {raw_word}")
+            return ipa_alignment, True
+        if random.random() < use_prob:
+            cur_words = " ".join(raw_word_list[:j+1])
+            phn = tokenize_text(text_tokenizer, cur_words)
+            if len(phn) == 0:
+                continue
+            phn = " ".join(phn)
+            start = 0 # at this point, we always start from the beginning of the sentence
+            ipa_alignment.append([start, end, phn])
+    if save:
+        if ipa_alignment:
+            with open(ipa_alignment_fn, 'w') as f:
+                for item in ipa_alignment:
+                    f.write(f"{item[0]}\t{item[1]}\t{item[2]}\n")
+    else:
+        return ipa_alignment, False
+
+
+
+def main(
+    data_root: str = '/data/scratch/pyp/datasets/librilight/preprocessed',
+    audio_ext: str = '.flac',
+    arpa_ext: str = '.csv',
+    trans_ext: str = '.txt',
+    split: str = 'valid',
+    use_prob: float = 0.5,
+    max_dur: float = 30., # do not consider utterance longer than this
+    prompt_dur: float = 30., # do not consider prompt longer than this
+):
+    text_tokenizer = TextTokenizer()
+    trans_dir = f'{data_root}/audio'
+    align_dir = f'{data_root}/alignment'
+    manifest_fn = f"{data_root}/manifest_final_encodec/{split}*=*.txt"
+    manifest_fns = glob.glob(manifest_fn)
+    target_dir = f'{data_root}/ipa_alignment'
+    encodec_sr = 50
+    os.makedirs(target_dir, exist_ok=True)
+    manifest = []
+    for manifest_fn in manifest_fns:
+        with open(manifest_fn, 'r') as f:
+            temp = [l.strip().split("\t") for l in f.readlines()]
+            manifest += [l[0] + audio_ext for l in temp if float(l[1])/encodec_sr < max_dur]
+    # # sequential processing
+    n_flags = 0
+    zero_words = 0
+    for j, fn in enumerate(tqdm.tqdm(manifest)):
+        ipa_alignment_fn = os.path.join(target_dir, fn.replace(audio_ext, '.txt'))
+        ipa_alignment, flag = create_alignment(fn, trans_dir, align_dir, audio_ext, trans_ext, arpa_ext, text_tokenizer, use_prob, ipa_alignment_fn, prompt_dur=prompt_dur)
+        n_flags += flag
+        if not ipa_alignment:
+            zero_words += 1
+        # print(f"{n_flags} out of {j+1} utterances have mismatched words")
+        # print(f"{zero_words} out of {j+1} utterances have zero words")
+        if ipa_alignment:
+            with open(ipa_alignment_fn, 'w') as f:
+                for item in ipa_alignment:
+                    f.write(f"{item[0]}\t{item[1]}\t{item[2]}\n")
+    
+    # # # # do the above using joblib parallisim
+    # print(f"Processing {len(manifest)} utterances")
+    # from joblib import Parallel, delayed
+    # Parallel(n_jobs=32, verbose=2)(delayed(create_alignment)(fn, trans_dir, align_dir, audio_ext, trans_ext, arpa_ext, text_tokenizer, use_prob, os.path.join(target_dir, fn.replace(audio_ext, '.txt')), save=True) for fn in manifest)
+    
+if __name__ == "__main__":
+    import fire
+    fire.Fire(main)

data/ll60k_preprocessing/tokenizer.py

@@ -0,0 +1,460 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/data/tokenizer.py
+# Copyright    2023                            (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Pattern, Union
+
+import numpy as np
+import torch
+import torchaudio
+# from encodec import EncodecModel
+# from encodec.utils import convert_audio
+# from lhotse.features import FeatureExtractor
+# from lhotse.utils import Seconds, compute_num_frames
+from phonemizer.backend import EspeakBackend
+from phonemizer.backend.espeak.language_switch import LanguageSwitch
+from phonemizer.backend.espeak.words_mismatch import WordMismatch
+from phonemizer.punctuation import Punctuation
+from phonemizer.separator import Separator
+
+try:
+    from pypinyin import Style, pinyin
+    from pypinyin.style._utils import get_finals, get_initials
+except Exception:
+    pass
+
+
+class PypinyinBackend:
+    """PypinyinBackend for Chinese. Most codes is referenced from espnet.
+    There are two types pinyin or initials_finals, one is
+    just like "ni1 hao3", the other is like "n i1 h ao3".
+    """
+
+    def __init__(
+        self,
+        backend="initials_finals",
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+    ) -> None:
+        self.backend = backend
+        self.punctuation_marks = punctuation_marks
+
+    def phonemize(
+        self, text: List[str], separator: Separator, strip=True, njobs=1
+    ) -> List[str]:
+        assert isinstance(text, List)
+        phonemized = []
+        for _text in text:
+            _text = re.sub(" +", " ", _text.strip())
+            _text = _text.replace(" ", separator.word)
+            phones = []
+            if self.backend == "pypinyin":
+                for n, py in enumerate(
+                    pinyin(
+                        _text, style=Style.TONE3, neutral_tone_with_five=True
+                    )
+                ):
+                    if all([c in self.punctuation_marks for c in py[0]]):
+                        if len(phones):
+                            assert phones[-1] == separator.syllable
+                            phones.pop(-1)
+
+                        phones.extend(list(py[0]))
+                    else:
+                        phones.extend([py[0], separator.syllable])
+            elif self.backend == "pypinyin_initials_finals":
+                for n, py in enumerate(
+                    pinyin(
+                        _text, style=Style.TONE3, neutral_tone_with_five=True
+                    )
+                ):
+                    if all([c in self.punctuation_marks for c in py[0]]):
+                        if len(phones):
+                            assert phones[-1] == separator.syllable
+                            phones.pop(-1)
+                        phones.extend(list(py[0]))
+                    else:
+                        if py[0][-1].isalnum():
+                            initial = get_initials(py[0], strict=False)
+                            if py[0][-1].isdigit():
+                                final = (
+                                    get_finals(py[0][:-1], strict=False)
+                                    + py[0][-1]
+                                )
+                            else:
+                                final = get_finals(py[0], strict=False)
+                            phones.extend(
+                                [
+                                    initial,
+                                    separator.phone,
+                                    final,
+                                    separator.syllable,
+                                ]
+                            )
+                        else:
+                            assert ValueError
+            else:
+                raise NotImplementedError
+            phonemized.append(
+                "".join(phones).rstrip(f"{separator.word}{separator.syllable}")
+            )
+        return phonemized
+
+
+class TextTokenizer:
+    """Phonemize Text."""
+
+    def __init__(
+        self,
+        language="en-us",
+        backend="espeak",
+        separator=Separator(word="_", syllable="-", phone="|"),
+        preserve_punctuation=True,
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+        with_stress: bool = False,
+        tie: Union[bool, str] = False,
+        language_switch: LanguageSwitch = "keep-flags",
+        words_mismatch: WordMismatch = "ignore",
+    ) -> None:
+        if backend == "espeak":
+            phonemizer = EspeakBackend(
+                language,
+                punctuation_marks=punctuation_marks,
+                preserve_punctuation=preserve_punctuation,
+                with_stress=with_stress,
+                tie=tie,
+                language_switch=language_switch,
+                words_mismatch=words_mismatch,
+            )
+        elif backend in ["pypinyin", "pypinyin_initials_finals"]:
+            phonemizer = PypinyinBackend(
+                backend=backend,
+                punctuation_marks=punctuation_marks + separator.word,
+            )
+        else:
+            raise NotImplementedError(f"{backend}")
+
+        self.backend = phonemizer
+        self.separator = separator
+
+    def to_list(self, phonemized: str) -> List[str]:
+        fields = []
+        for word in phonemized.split(self.separator.word):
+            # "ɐ    m|iː|n?"    ɹ|ɪ|z|ɜː|v; h|ɪ|z.
+            pp = re.findall(r"\w+|[^\w\s]", word, re.UNICODE)
+            fields.extend(
+                [p for p in pp if p != self.separator.phone]
+                + [self.separator.word]
+            )
+        assert len("".join(fields[:-1])) == len(phonemized) - phonemized.count(
+            self.separator.phone
+        )
+        return fields[:-1]
+
+    def __call__(self, text, strip=True) -> List[List[str]]:
+        if isinstance(text, str):
+            text = [text]
+
+        phonemized = self.backend.phonemize(
+            text, separator=self.separator, strip=strip, njobs=1
+        )
+        return [self.to_list(p) for p in phonemized]
+
+
+def tokenize_text(tokenizer: TextTokenizer, text: str) -> List[str]:
+    phonemes = tokenizer([text.strip()])
+    return phonemes[0]  # k2symbols
+
+
+def remove_encodec_weight_norm(model):
+    from encodec.modules import SConv1d
+    from encodec.modules.seanet import SConvTranspose1d, SEANetResnetBlock
+    from torch.nn.utils import remove_weight_norm
+    encoder = model.encoder.model
+    for key in encoder._modules:
+        if isinstance(encoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(encoder._modules[key].shortcut.conv.conv)
+            block_modules = encoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(encoder._modules[key], SConv1d):
+            remove_weight_norm(encoder._modules[key].conv.conv)
+    decoder = model.decoder.model
+    for key in decoder._modules:
+        if isinstance(decoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(decoder._modules[key].shortcut.conv.conv)
+            block_modules = decoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(decoder._modules[key], SConvTranspose1d):
+            remove_weight_norm(decoder._modules[key].convtr.convtr)
+        elif isinstance(decoder._modules[key], SConv1d):
+            remove_weight_norm(decoder._modules[key].conv.conv)
+
+
+# class AudioTokenizer:
+#     """EnCodec audio."""
+
+#     def __init__(
+#         self,
+#         bandwidth, float=6.0,
+#         device: Any = None,
+#     ) -> None:
+#         # Instantiate a pretrained EnCodec model
+#         model = EncodecModel.encodec_model_24khz()
+#         model.set_target_bandwidth(bandwidth=bandwidth)
+#         remove_encodec_weight_norm(model)
+
+#         if not device:
+#             device = torch.device("cpu")
+#             if torch.cuda.is_available():
+#                 device = torch.device("cuda:0")
+
+#         self._device = device
+
+#         self.codec = model.to(device)
+#         self.sample_rate = model.sample_rate
+#         self.channels = model.channels
+
+#     @property
+#     def device(self):
+#         return self._device
+
+#     def encode(self, wav: torch.Tensor) -> torch.Tensor:
+#         return self.codec.encode(wav.to(self.device))
+
+#     def decode(self, frames: torch.Tensor) -> torch.Tensor:
+#         return self.codec.decode(frames)
+
+# class AudioTokenizer:
+#     """EnCodec audio."""
+
+#     def __init__(
+#         self,
+#         bandwidth: float=6.0,
+#         device: Any = None,
+#         hificodec=False,
+#         signature = None
+#     ) -> None:
+#         self.hificodec = hificodec
+#         self.customized = True if signature != None else False
+#         if hificodec:
+#             import sys
+#             sys.path.append("/home/pyp/AcademiCodec")
+#             from academicodec.models.hificodec.vqvae import VQVAE
+#             config_path = "/home/pyp/AcademiCodec/egs/HiFi-Codec-16k-320d/config_16k_320d.json"
+#             model_path = "/home/pyp/AcademiCodec/egs/HiFi-Codec-16k-320d/checkpoint/HiFi-Codec-16k-320d"
+#             self.sample_rate = 16000
+#             self.channels = 1
+#             model = VQVAE(config_path, model_path, with_encoder=True)
+#             model.generator.remove_weight_norm()
+#             model.encoder.remove_weight_norm()
+#             model.eval()
+#         else:
+#             if signature != None:
+#                 # use customized encodec model
+#                 # import sys
+#                 # sys.path.append("home/pyp/audiocraft")
+#                 from audiocraft.solvers import CompressionSolver
+#                 model_path = f'//sig/{signature}'
+#                 model = CompressionSolver.model_from_checkpoint(model_path)
+#                 self.sample_rate = model.sample_rate
+#                 self.channels = model.channels
+#             else:
+#                 # Instantiate a pretrained EnCodec model
+#                 model = EncodecModel.encodec_model_24khz()
+#                 model.set_target_bandwidth(bandwidth=bandwidth)
+#                 remove_encodec_weight_norm(model)
+#                 self.sample_rate = model.sample_rate
+#                 self.channels = model.channels
+
+#         if not device:
+#             device = torch.device("cpu")
+#             if torch.cuda.is_available():
+#                 device = torch.device("cuda:0")
+
+#         self._device = device
+
+#         self.codec = model.to(device)
+
+#     @property
+#     def device(self):
+#         return self._device
+
+#     def encode(self, wav: torch.Tensor) -> torch.Tensor:
+#         if self.hificodec:
+#             assert wav.ndim==3 and wav.shape[:2] == torch.Size((1,1)), wav.shape
+#             wav = wav.squeeze(0)
+#             codes = self.codec.encode(wav.to(self.device)) # [1,T,4]
+#             return [(codes.transpose(2,1),None)]
+#         elif self.customized:
+#             codes = self.codec.encode(wav.to(self.device))
+#             return [(codes[0], None)]
+#         return self.codec.encode(wav.to(self.device))
+
+#     def decode(self, frames: torch.Tensor) -> torch.Tensor:
+#         if self.hificodec:
+#             frames = frames[0][0] # [1,4,T]
+#             assert frames.shape[:2] == torch.Size((1,4))
+#             audio = self.codec(frames.transpose(2,1))
+#             assert audio.shape[0] == 1, audio.shape
+#             return audio
+#         elif self.customized:
+#             frames = frames[0][0] # [1,4,T]
+#             return self.codec.decode(frames)
+#         return self.codec.decode(frames)
+#         # try:
+#         #     return self.codec.decode(frames)
+#         # except:
+#         #     import logging
+#         #     logging.info(f"error when decoding frame of shape: {frames[0][0].shape}")
+#         #     self.codec.cpu()
+#         #     ret = self.codec.cpu().decode([(frames[0][0].cpu(),None)])[0].to(self._device)
+#         #     self.codec.to(self._device)
+#         #     return [ret]
+        
+# def tokenize_audio(tokenizer: AudioTokenizer, audio_path: str, offset = -1, num_frames=-1):
+#     # Load and pre-process the audio waveform
+#     if offset != -1 and num_frames!=-1:
+#         wav, sr = torchaudio.load(audio_path, frame_offset=offset, num_frames=num_frames)
+#     else:
+#         wav, sr = torchaudio.load(audio_path)
+#     wav = convert_audio(wav, sr, tokenizer.sample_rate, tokenizer.channels)
+#     wav = wav.unsqueeze(0)
+
+#     # Extract discrete codes from EnCodec
+#     with torch.no_grad():
+#         encoded_frames = tokenizer.encode(wav)
+#     return encoded_frames
+
+
+# @dataclass
+# class AudioTokenConfig:
+#     frame_shift: Seconds = 320.0 / 24000
+#     num_quantizers: int = 8
+
+#     def to_dict(self) -> Dict[str, Any]:
+#         return asdict(self)
+
+#     @staticmethod
+#     def from_dict(data: Dict[str, Any]) -> "AudioTokenConfig":
+#         return AudioTokenConfig(**data)
+
+
+# class AudioTokenExtractor(FeatureExtractor):
+#     name = "encodec"
+#     config_type = AudioTokenConfig
+
+#     def __init__(self, config: Optional[Any] = None):
+#         super(AudioTokenExtractor, self).__init__(config)
+#         self.tokenizer = AudioTokenizer()
+
+#     def extract(
+#         self, samples: Union[np.ndarray, torch.Tensor], sampling_rate: int
+#     ) -> np.ndarray:
+#         if not isinstance(samples, torch.Tensor):
+#             samples = torch.from_numpy(samples)
+#         if sampling_rate != self.tokenizer.sample_rate:
+#             samples = convert_audio(
+#                 samples,
+#                 sampling_rate,
+#                 self.tokenizer.sample_rate,
+#                 self.tokenizer.channels,
+#             )
+#         if len(samples.shape) == 2:
+#             samples = samples.unsqueeze(0)
+#         else:
+#             raise ValueError()
+
+#         device = self.tokenizer.device
+#         encoded_frames = self.tokenizer.encode(samples.detach().to(device))
+#         codes = encoded_frames[0][0]  # [B, n_q, T]
+#         if True:
+#             duration = round(samples.shape[-1] / sampling_rate, ndigits=12)
+#             expected_num_frames = compute_num_frames(
+#                 duration=duration,
+#                 frame_shift=self.frame_shift,
+#                 sampling_rate=sampling_rate,
+#             )
+#             assert abs(codes.shape[-1] - expected_num_frames) <= 1
+#             codes = codes[..., :expected_num_frames]
+#         return codes.cpu().squeeze(0).permute(1, 0).numpy()
+
+#     @property
+#     def frame_shift(self) -> Seconds:
+#         return self.config.frame_shift
+
+#     def feature_dim(self, sampling_rate: int) -> int:
+#         return self.config.num_quantizers
+
+#     def pad_tensor_list(self, tensor_list, device, padding_value=0):
+#         # 计算每个张量的长度
+#         lengths = [tensor.shape[0] for tensor in tensor_list]
+#         # 使用pad_sequence函数进行填充
+#         tensor_list = [torch.Tensor(t).to(device) for t in tensor_list]
+#         padded_tensor = torch.nn.utils.rnn.pad_sequence(
+#             tensor_list, batch_first=True, padding_value=padding_value
+#         )
+#         return padded_tensor, lengths
+
+#     def extract_batch(self, samples, sampling_rate, lengths) -> np.ndarray:
+#         samples = [wav.squeeze() for wav in samples]
+#         device = self.tokenizer.device
+#         samples, lengths = self.pad_tensor_list(samples, device)
+#         samples = samples.unsqueeze(1)
+
+#         if not isinstance(samples, torch.Tensor):
+#             samples = torch.from_numpy(samples)
+#         if len(samples.shape) != 3:
+#             raise ValueError()
+#         if sampling_rate != self.tokenizer.sample_rate:
+#             samples = [
+#                 convert_audio(
+#                     wav,
+#                     sampling_rate,
+#                     self.tokenizer.sample_rate,
+#                     self.tokenizer.channels,
+#                 )
+#                 for wav in samples
+#             ]
+#         # Extract discrete codes from EnCodec
+#         with torch.no_grad():
+#             encoded_frames = self.tokenizer.encode(samples.detach().to(device))
+#         encoded_frames = encoded_frames[0][0]  # [B, n_q, T]
+#         batch_codes = []
+#         for b, length in enumerate(lengths):
+#             codes = encoded_frames[b]
+#             duration = round(length / sampling_rate, ndigits=12)
+#             expected_num_frames = compute_num_frames(
+#                 duration=duration,
+#                 frame_shift=self.frame_shift,
+#                 sampling_rate=sampling_rate,
+#             )
+#             batch_codes.append(codes[..., :expected_num_frames])
+#         return [codes.cpu().permute(1, 0).numpy() for codes in batch_codes]
+
+
+if __name__ == "__main__":
+    model = EncodecModel.encodec_model_24khz()
+    model.set_target_bandwidth(6.0)
+    # model.cuda()
+    samples = torch.from_numpy(np.random.random([4, 1, 30000])).type(torch.float32)
+    codes_norm = model.encode(samples.cuda())
+    remove_encodec_weight_norm(model)
+    codes_raw = model.encode(samples.cuda())
+
+    assert torch.allclose(codes_raw[0][0], codes_norm[0][0])

data/tokenizer.py

@@ -0,0 +1,295 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/data/tokenizer.py
+# Copyright    2023                            (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Pattern, Union
+
+import numpy as np
+import torch
+import torchaudio
+# from encodec import EncodecModel
+# from encodec.utils import convert_audio
+# from lhotse.features import FeatureExtractor
+# from lhotse.utils import Seconds, compute_num_frames
+from phonemizer.backend import EspeakBackend
+from phonemizer.backend.espeak.language_switch import LanguageSwitch
+from phonemizer.backend.espeak.words_mismatch import WordMismatch
+from phonemizer.punctuation import Punctuation
+from phonemizer.separator import Separator
+
+
+try:
+    from pypinyin import Style, pinyin
+    from pypinyin.style._utils import get_finals, get_initials
+except Exception:
+    pass
+
+
+class PypinyinBackend:
+    """PypinyinBackend for Chinese. Most codes is referenced from espnet.
+    There are two types pinyin or initials_finals, one is
+    just like "ni1 hao3", the other is like "n i1 h ao3".
+    """
+
+    def __init__(
+        self,
+        backend="initials_finals",
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+    ) -> None:
+        self.backend = backend
+        self.punctuation_marks = punctuation_marks
+
+    def phonemize(
+        self, text: List[str], separator: Separator, strip=True, njobs=1
+    ) -> List[str]:
+        assert isinstance(text, List)
+        phonemized = []
+        for _text in text:
+            _text = re.sub(" +", " ", _text.strip())
+            _text = _text.replace(" ", separator.word)
+            phones = []
+            if self.backend == "pypinyin":
+                for n, py in enumerate(
+                    pinyin(
+                        _text, style=Style.TONE3, neutral_tone_with_five=True
+                    )
+                ):
+                    if all([c in self.punctuation_marks for c in py[0]]):
+                        if len(phones):
+                            assert phones[-1] == separator.syllable
+                            phones.pop(-1)
+
+                        phones.extend(list(py[0]))
+                    else:
+                        phones.extend([py[0], separator.syllable])
+            elif self.backend == "pypinyin_initials_finals":
+                for n, py in enumerate(
+                    pinyin(
+                        _text, style=Style.TONE3, neutral_tone_with_five=True
+                    )
+                ):
+                    if all([c in self.punctuation_marks for c in py[0]]):
+                        if len(phones):
+                            assert phones[-1] == separator.syllable
+                            phones.pop(-1)
+                        phones.extend(list(py[0]))
+                    else:
+                        if py[0][-1].isalnum():
+                            initial = get_initials(py[0], strict=False)
+                            if py[0][-1].isdigit():
+                                final = (
+                                    get_finals(py[0][:-1], strict=False)
+                                    + py[0][-1]
+                                )
+                            else:
+                                final = get_finals(py[0], strict=False)
+                            phones.extend(
+                                [
+                                    initial,
+                                    separator.phone,
+                                    final,
+                                    separator.syllable,
+                                ]
+                            )
+                        else:
+                            assert ValueError
+            else:
+                raise NotImplementedError
+            phonemized.append(
+                "".join(phones).rstrip(f"{separator.word}{separator.syllable}")
+            )
+        return phonemized
+
+
+class TextTokenizer:
+    """Phonemize Text."""
+
+    def __init__(
+        self,
+        language="en-us",
+        backend="espeak",
+        separator=Separator(word="_", syllable="-", phone="|"),
+        preserve_punctuation=True,
+        punctuation_marks: Union[str, Pattern] = Punctuation.default_marks(),
+        with_stress: bool = False,
+        tie: Union[bool, str] = False,
+        language_switch: LanguageSwitch = "keep-flags",
+        words_mismatch: WordMismatch = "ignore",
+    ) -> None:
+        if backend == "espeak":
+            phonemizer = EspeakBackend(
+                language,
+                punctuation_marks=punctuation_marks,
+                preserve_punctuation=preserve_punctuation,
+                with_stress=with_stress,
+                tie=tie,
+                language_switch=language_switch,
+                words_mismatch=words_mismatch,
+            )
+        elif backend in ["pypinyin", "pypinyin_initials_finals"]:
+            phonemizer = PypinyinBackend(
+                backend=backend,
+                punctuation_marks=punctuation_marks + separator.word,
+            )
+        else:
+            raise NotImplementedError(f"{backend}")
+
+        self.backend = phonemizer
+        self.separator = separator
+
+    def to_list(self, phonemized: str) -> List[str]:
+        fields = []
+        for word in phonemized.split(self.separator.word):
+            # "ɐ    m|iː|n?"    ɹ|ɪ|z|ɜː|v; h|ɪ|z.
+            pp = re.findall(r"\w+|[^\w\s]", word, re.UNICODE)
+            fields.extend(
+                [p for p in pp if p != self.separator.phone]
+                + [self.separator.word]
+            )
+        assert len("".join(fields[:-1])) == len(phonemized) - phonemized.count(
+            self.separator.phone
+        )
+        return fields[:-1]
+
+    def __call__(self, text, strip=True) -> List[List[str]]:
+        if isinstance(text, str):
+            text = [text]
+
+        phonemized = self.backend.phonemize(
+            text, separator=self.separator, strip=strip, njobs=1
+        )
+        return [self.to_list(p) for p in phonemized]
+
+
+def tokenize_text(tokenizer: TextTokenizer, text: str) -> List[str]:
+    phonemes = tokenizer([text.strip()])
+    return phonemes[0]  # k2symbols
+
+
+def remove_encodec_weight_norm(model):
+    from encodec.modules import SConv1d
+    from encodec.modules.seanet import SConvTranspose1d, SEANetResnetBlock
+    from torch.nn.utils import remove_weight_norm
+    encoder = model.encoder.model
+    for key in encoder._modules:
+        if isinstance(encoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(encoder._modules[key].shortcut.conv.conv)
+            block_modules = encoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(encoder._modules[key], SConv1d):
+            remove_weight_norm(encoder._modules[key].conv.conv)
+    decoder = model.decoder.model
+    for key in decoder._modules:
+        if isinstance(decoder._modules[key], SEANetResnetBlock):
+            remove_weight_norm(decoder._modules[key].shortcut.conv.conv)
+            block_modules = decoder._modules[key].block._modules
+            for skey in block_modules:
+                if isinstance(block_modules[skey], SConv1d):
+                    remove_weight_norm(block_modules[skey].conv.conv)
+        elif isinstance(decoder._modules[key], SConvTranspose1d):
+            remove_weight_norm(decoder._modules[key].convtr.convtr)
+        elif isinstance(decoder._modules[key], SConv1d):
+            remove_weight_norm(decoder._modules[key].conv.conv)
+
+
+class AudioTokenizer:
+    """mimi audio."""
+
+    def __init__(
+        self,
+        bandwidth: float=6.0,
+        device: Any = None,
+        hificodec=False,
+        signature = None,
+        encode_only = False
+    ) -> None:
+        self.signature = signature
+        from data.encodec import get_compression_model
+        model = get_compression_model(signature, encode_only=encode_only, device=device)
+        self.sample_rate = model.sample_rate
+        self.channels = model.channels
+
+        if not device:
+            device = torch.device("cpu")
+            if torch.cuda.is_available():
+                device = torch.device("cuda")
+
+        self._device = device
+
+        self.codec = model.to(device)
+
+    @property
+    def device(self):
+        return self._device
+
+    def encode(self, wav: torch.Tensor) -> torch.Tensor:
+        if self.signature != None:
+            if self.signature == "lfsc":
+                if wav.ndim==3:
+                    assert wav.shape[:2] == torch.Size((1,1)), wav.shape
+                    wav = wav.squeeze(0)
+                elif wav.ndim==2:
+                    assert wav.shape[0] == 1, wav.shape
+                else:
+                    raise ValueError(wav.shape)
+                audio_len = torch.tensor([wav.shape[1]]).to(self.device)
+                codes, encoded_len = self.codec.encode(audio=wav.to(self.device), audio_len=audio_len)
+                return codes[:, :, :encoded_len[0]]
+            else:
+                codes = self.codec.encode(wav.to(self.device))
+                return codes[0]
+        else:
+            assert wav.ndim==3 and wav.shape[:2] == torch.Size((1,1)), wav.shape
+            return self.codec.encode(wav.to(self.device))
+
+    def decode(self, frames: torch.Tensor) -> torch.Tensor:
+        if self.signature != None and self.signature == "lfsc":
+            encoded_len = torch.tensor([frames.shape[-1]]).to(self.device)
+            reconstructed_audio, decoded_len = self.codec.decode(tokens=frames, tokens_len=encoded_len)
+            return reconstructed_audio[:, :decoded_len[0]].unsqueeze(0)
+        else:
+            return self.codec.decode(frames)
+        
+def tokenize_audio(tokenizer: AudioTokenizer, audio_path: str, offset = -1, num_frames=-1):
+    # Load and pre-process the audio waveform
+    if offset != -1 and num_frames!=-1:
+        wav, sr = torchaudio.load(audio_path, frame_offset=offset, num_frames=num_frames)
+    else:
+        wav, sr = torchaudio.load(audio_path)
+    if sr != tokenizer.sample_rate:
+        wav = torchaudio.transforms.Resample(sr, tokenizer.sample_rate)(wav)
+        sr = tokenizer.sample_rate
+    if wav.shape[0] == 2:
+        wav = wav.mean(dim=0, keepdim=True)
+    wav = wav.unsqueeze(0)
+    # Extract discrete codes from mimi
+    with torch.no_grad():
+        encoded_frames = tokenizer.encode(wav)
+    return encoded_frames
+
+
+if __name__ == "__main__":
+    # tok = AudioTokenizer(signature="lfsc", device="cpu")
+    tok = AudioTokenizer(signature="/home/pyp/BoostedVoiceEditor/pretrained/encodec_6f79c6a8.th", device="cpu")
+    inaudio = "/home/pyp/BoostedVoiceEditor/demo/pam.wav"
+    encoded_frames = tokenize_audio(tok, inaudio)
+    print(encoded_frames.shape)
+    # decode it back
+    decoded_audio = tok.decode(encoded_frames)
+    torchaudio.save("/home/pyp/BoostedVoiceEditor/demo/pam_reconstructed_encodec_4cb_2nd.wav", decoded_audio[0], tok.sample_rate)
+

demo/5895_34622_000026_000002.wav

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

generated_tts/generated.wav

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

inference_commandline.py

@@ -0,0 +1,192 @@
+import os
+import torch
+import torchaudio
+import numpy as np
+import random
+import whisper
+import fire
+from argparse import Namespace
+
+from data.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+)
+
+from models import voice_star
+from inference_tts_utils import inference_one_sample
+
+############################################################
+# Utility Functions
+############################################################
+
+def seed_everything(seed=1):
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def estimate_duration(ref_audio_path, text):
+    """
+    Estimate duration based on seconds per character from the reference audio.
+    """
+    info = torchaudio.info(ref_audio_path)
+    audio_duration = info.num_frames / info.sample_rate
+    length_text = max(len(text), 1)
+    spc = audio_duration / length_text  # seconds per character
+    return len(text) * spc
+
+############################################################
+# Main Inference Function
+############################################################
+
+def run_inference(
+    reference_speech="./demo/5895_34622_000026_000002.wav",
+    target_text="I cannot believe that the same model can also do text to speech synthesis too! And you know what? this audio is 8 seconds long.",
+    # Model
+    model_name="VoiceStar_840M_30s", # or VoiceStar_840M_40s, the later model is trained on maximally 40s long speech
+    model_root="./pretrained",
+    # Additional optional
+    reference_text=None,  # if None => run whisper on reference_speech
+    target_duration=None, # if None => estimate from reference_speech and target_text
+    # Default hyperparameters from snippet
+    codec_audio_sr=16000, # do not change
+    codec_sr=50, # do not change
+    top_k=10, # try 10, 20, 30, 40
+    top_p=1, # do not change
+    min_p=1, # do not change
+    temperature=1,
+    silence_tokens=None, # do not change it
+    kvcache=1, # if OOM, set to 0
+    multi_trial=None, # do not change it
+    repeat_prompt=1, # increase this to improve speaker similarity, but it reference speech duration in total adding target duration is longer than maximal training duration, quality may drop
+    stop_repetition=3, # will not use it
+    sample_batch_size=1, # do not change
+    # Others
+    seed=1,
+    output_dir="./generated_tts",
+    # Some snippet-based defaults
+    cut_off_sec=100, # do not adjust this, we always use the entire reference speech. If you wish to change, also make sure to change the reference_transcript, so that it's only the trasnscript of the speech remained
+):
+    """
+    Inference script using Fire.
+
+    Example:
+        python inference_commandline.py \
+            --reference_speech "./demo/5895_34622_000026_000002.wav" \
+            --target_text "I cannot believe ... this audio is 10 seconds long." \
+            --reference_text "(optional) text to use as prefix" \
+            --target_duration (optional float) 
+    """
+
+    # Seed everything
+    seed_everything(seed)
+
+    # Load model, phn2num, and args
+    torch.serialization.add_safe_globals([Namespace])
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ckpt_fn = os.path.join(model_root, model_name+".pth")
+    if not os.path.exists(ckpt_fn):
+        # use wget to download
+        print(f"[Info] Downloading {model_name} checkpoint...")
+        os.system(f"wget https://huggingface.co/pyp1/VoiceStar/resolve/main/{model_name}.pth?download=true -O {ckpt_fn}")
+    bundle = torch.load(ckpt_fn, map_location=device, weights_only=True)
+    args = bundle["args"]
+    phn2num = bundle["phn2num"]
+    model = voice_star.VoiceStar(args)
+    model.load_state_dict(bundle["model"])
+    model.to(device)
+    model.eval()
+
+    # If reference_text not provided, use whisper large-v3-turbo
+    if reference_text is None:
+        print("[Info] No reference_text provided, transcribing reference_speech with Whisper.")
+        wh_model = whisper.load_model("large-v3-turbo")
+        result = wh_model.transcribe(reference_speech)
+        prefix_transcript = result["text"]
+        print(f"[Info] Whisper transcribed text: {prefix_transcript}")
+    else:
+        prefix_transcript = reference_text
+
+    # If target_duration not provided, estimate from reference speech + target_text
+    if target_duration is None:
+        target_generation_length = estimate_duration(reference_speech, target_text)
+        print(f"[Info] target_duration not provided, estimated as {target_generation_length:.2f} seconds. If not desired, please provide a target_duration.")
+    else:
+        target_generation_length = float(target_duration)
+
+    # signature from snippet
+    if args.n_codebooks == 4:
+        signature = "./pretrained/encodec_6f79c6a8.th"
+    elif args.n_codebooks == 8:
+        signature = "./pretrained/encodec_8cb1024_giga.th"
+    else:
+        # fallback, just use the 6-f79c6a8
+        signature = "./pretrained/encodec_6f79c6a8.th"
+
+    if silence_tokens is None:
+        # default from snippet
+        silence_tokens = []
+
+    if multi_trial is None:
+        # default from snippet
+        multi_trial = []
+
+    delay_pattern_increment = args.n_codebooks + 1  # from snippet
+
+    # We can compute prompt_end_frame if we want, from snippet
+    info = torchaudio.info(reference_speech)
+    prompt_end_frame = int(cut_off_sec * info.sample_rate)
+
+    # Prepare tokenizers
+    audio_tokenizer = AudioTokenizer(signature=signature)
+    text_tokenizer = TextTokenizer(backend="espeak")
+
+    # decode_config from snippet
+    decode_config = {
+        'top_k': top_k,
+        'top_p': top_p,
+        'min_p': min_p,
+        'temperature': temperature,
+        'stop_repetition': stop_repetition,
+        'kvcache': kvcache,
+        'codec_audio_sr': codec_audio_sr,
+        'codec_sr': codec_sr,
+        'silence_tokens': silence_tokens,
+        'sample_batch_size': sample_batch_size
+    }
+
+    # Run inference
+    print("[Info] Running TTS inference...")
+    concated_audio, gen_audio = inference_one_sample(
+        model, args, phn2num, text_tokenizer, audio_tokenizer,
+        reference_speech, target_text,
+        device, decode_config,
+        prompt_end_frame=prompt_end_frame,
+        target_generation_length=target_generation_length,
+        delay_pattern_increment=delay_pattern_increment,
+        prefix_transcript=prefix_transcript,
+        multi_trial=multi_trial,
+        repeat_prompt=repeat_prompt,
+    )
+
+    # The model returns a list of waveforms, pick the first
+    concated_audio, gen_audio = concated_audio[0].cpu(), gen_audio[0].cpu()
+
+    # Save the audio (just the generated portion, as the snippet does)
+    os.makedirs(output_dir, exist_ok=True)
+    out_filename = "generated.wav"
+    out_path = os.path.join(output_dir, out_filename)
+    torchaudio.save(out_path, gen_audio, codec_audio_sr)
+
+    print(f"[Success] Generated audio saved to {out_path}")
+
+
+def main():
+    fire.Fire(run_inference)
+
+if __name__ == "__main__":
+    main()

inference_gradio.py

@@ -0,0 +1,334 @@
+#!/usr/bin/env python3
+"""
+gradio_tts_app.py
+
+Run:
+    python gradio_tts_app.py
+
+Then open the printed local or public URL in your browser.
+"""
+
+import os
+import random
+import numpy as np
+import torch
+import torchaudio
+import whisper
+import gradio as gr
+from argparse import Namespace
+
+# ---------------------------------------------------------------------
+# The following imports assume your local project structure:
+#    data/tokenizer.py
+#    models/voice_star.py
+#    inference_tts_utils.py
+# Adjust if needed.
+# ---------------------------------------------------------------------
+from data.tokenizer import AudioTokenizer, TextTokenizer
+from models import voice_star
+from inference_tts_utils import inference_one_sample
+
+
+############################################################
+# Utility Functions
+############################################################
+
+def seed_everything(seed=1):
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+
+def estimate_duration(ref_audio_path, text):
+    """
+    Estimate duration based on seconds per character from the reference audio.
+    """
+    info = torchaudio.info(ref_audio_path)
+    audio_duration = info.num_frames / info.sample_rate
+    length_text = max(len(text), 1)
+    spc = audio_duration / length_text  # seconds per character
+    return len(text) * spc
+
+
+############################################################
+# Main Inference Function
+############################################################
+
+def run_inference(
+    # User-adjustable parameters (no "# do not change" in snippet)
+    reference_speech="./demo/5895_34622_000026_000002.wav",
+    target_text="VoiceStar is a very interesting model, it's duration controllable and can extrapolate",
+    model_name="VoiceStar_840M_40s",
+    model_root="./pretrained",
+    reference_text=None,    # optional
+    target_duration=None,   # optional
+    top_k=10,               # can try 10, 20, 30, 40
+    temperature=1,
+    kvcache=1,              # if OOM, set to 0
+    repeat_prompt=1,        # use higher to improve speaker similarity
+    stop_repetition=3,      # snippet says "will not use it" but not "do not change"
+    seed=1,
+    output_dir="./generated_tts",
+
+    # Non-adjustable parameters (based on snippet instructions)
+    codec_audio_sr=16000,   # do not change
+    codec_sr=50,            # do not change
+    top_p=1,                # do not change
+    min_p=1,                # do not change
+    silence_tokens=None,    # do not change it
+    multi_trial=None,       # do not change it
+    sample_batch_size=1,    # do not change
+    cut_off_sec=100,        # do not adjust
+):
+    """
+    Inference script for VoiceStar TTS.
+    """
+    # 1. Set seed
+    seed_everything(seed)
+
+    # 2. Load model checkpoint
+    torch.serialization.add_safe_globals([Namespace])
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    ckpt_fn = os.path.join(model_root, model_name + ".pth")
+    if not os.path.exists(ckpt_fn):
+        # use wget to download
+        print(f"[Info] Downloading {model_name} checkpoint...")
+        os.system(f"wget https://huggingface.co/pyp1/VoiceStar/resolve/main/{model_name}.pth?download=true -O {ckpt_fn}")
+    bundle = torch.load(ckpt_fn, map_location=device, weights_only=True)
+    args = bundle["args"]
+    phn2num = bundle["phn2num"]
+
+    model = voice_star.VoiceStar(args)
+    model.load_state_dict(bundle["model"])
+    model.to(device)
+    model.eval()
+
+    # 3. If reference_text not provided, transcribe reference speech with Whisper
+    if reference_text is None:
+        print("[Info] No reference_text provided. Transcribing reference_speech with Whisper (large-v3-turbo).")
+        wh_model = whisper.load_model("large-v3-turbo")
+        result = wh_model.transcribe(reference_speech)
+        prefix_transcript = result["text"]
+        print(f"[Info] Whisper transcribed text: {prefix_transcript}")
+    else:
+        prefix_transcript = reference_text
+
+    # 4. If target_duration not provided, estimate from reference speech + target_text
+    if target_duration is None:
+        target_generation_length = estimate_duration(reference_speech, target_text)
+        print(f"[Info] target_duration not provided, estimated as {target_generation_length:.2f}s. Provide --target_duration if needed.")
+    else:
+        target_generation_length = float(target_duration)
+
+    # 5. Prepare signature from snippet
+    if args.n_codebooks == 4:
+        signature = "./pretrained/encodec_6f79c6a8.th"
+    elif args.n_codebooks == 8:
+        signature = "./pretrained/encodec_8cb1024_giga.th"
+    else:
+        signature = "./pretrained/encodec_6f79c6a8.th"
+
+    if silence_tokens is None:
+        silence_tokens = []
+
+    if multi_trial is None:
+        multi_trial = []
+
+    delay_pattern_increment = args.n_codebooks + 1  # from snippet
+
+    info = torchaudio.info(reference_speech)
+    prompt_end_frame = int(cut_off_sec * info.sample_rate)
+
+    # 6. Tokenizers
+    audio_tokenizer = AudioTokenizer(signature=signature)
+    text_tokenizer = TextTokenizer(backend="espeak")
+
+    # 7. decode_config
+    decode_config = {
+        "top_k": top_k,
+        "top_p": top_p,
+        "min_p": min_p,
+        "temperature": temperature,
+        "stop_repetition": stop_repetition,
+        "kvcache": kvcache,
+        "codec_audio_sr": codec_audio_sr,
+        "codec_sr": codec_sr,
+        "silence_tokens": silence_tokens,
+        "sample_batch_size": sample_batch_size,
+    }
+
+    # 8. Run inference
+    print("[Info] Running TTS inference...")
+    concated_audio, gen_audio = inference_one_sample(
+        model, args, phn2num, text_tokenizer, audio_tokenizer,
+        reference_speech, target_text,
+        device, decode_config,
+        prompt_end_frame=prompt_end_frame,
+        target_generation_length=target_generation_length,
+        delay_pattern_increment=delay_pattern_increment,
+        prefix_transcript=prefix_transcript,
+        multi_trial=multi_trial,
+        repeat_prompt=repeat_prompt,
+    )
+
+    # The model returns a list of waveforms, pick the first
+    concated_audio, gen_audio = concated_audio[0].cpu(), gen_audio[0].cpu()
+
+    # 9. Save generated audio
+    os.makedirs(output_dir, exist_ok=True)
+    out_filename = "generated.wav"
+    out_path = os.path.join(output_dir, out_filename)
+    torchaudio.save(out_path, gen_audio, codec_audio_sr)
+
+    print(f"[Success] Generated audio saved to {out_path}")
+    return out_path  # Return the path for Gradio to load
+
+
+############################
+# Transcription function
+############################
+
+def transcribe_audio(reference_speech):
+    """
+    Transcribe uploaded reference audio with Whisper, return text.
+    If no file, return empty string.
+    """
+    if reference_speech is None:
+        return ""
+    audio_path = reference_speech  # Because type="filepath"
+
+    if not os.path.exists(audio_path):
+        return "File not found."
+
+    print("[Info] Transcribing with Whisper...")
+    model = whisper.load_model("medium")  # or "large-v2" etc.
+    result = model.transcribe(audio_path)
+    return result["text"]
+
+############################
+# Gradio UI
+############################
+
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("## VoiceStar TTS with Editable Reference Text")
+
+        with gr.Row():
+            reference_speech_input = gr.Audio(
+                label="Reference Speech",
+                type="filepath",
+                elem_id="ref_speech"
+            )
+            transcribe_button = gr.Button("Transcribe")
+
+        # The transcribed text appears here and can be edited
+        reference_text_box = gr.Textbox(
+            label="Reference Text (Editable)",
+            placeholder="Click 'Transcribe' to auto-fill from reference speech...",
+            lines=2
+        )
+
+        target_text_box = gr.Textbox(
+            label="Target Text",
+            value="VoiceStar is a very interesting model, it's duration controllable and can extrapolate to unseen duration.",
+            lines=3
+        )
+
+        model_name_box = gr.Textbox(
+            label="Model Name",
+            value="VoiceStar_840M_40s"
+        )
+
+        model_root_box = gr.Textbox(
+            label="Model Root Directory",
+            value="/data1/scratch/pyp/BoostedVoiceEditor/runs"
+        )
+
+        reference_duration_box = gr.Textbox(
+            label="Target Duration (Optional)",
+            placeholder="Leave empty for auto-estimate."
+        )
+
+        top_k_box = gr.Number(label="top_k", value=10)
+        temperature_box = gr.Number(label="temperature", value=1.0)
+        kvcache_box = gr.Number(label="kvcache (1 or 0)", value=1)
+        repeat_prompt_box = gr.Number(label="repeat_prompt", value=1)
+        stop_repetition_box = gr.Number(label="stop_repetition", value=3)
+        seed_box = gr.Number(label="Random Seed", value=1)
+        output_dir_box = gr.Textbox(label="Output Directory", value="./generated_tts")
+
+        generate_button = gr.Button("Generate TTS")
+        output_audio = gr.Audio(label="Generated Audio", type="filepath")
+
+        # 1) When user clicks "Transcribe", we call `transcribe_audio`
+        transcribe_button.click(
+            fn=transcribe_audio,
+            inputs=[reference_speech_input],
+            outputs=[reference_text_box],
+        )
+
+        # 2) The actual TTS generation function. 
+        def gradio_inference(
+            reference_speech,
+            reference_text,
+            target_text,
+            model_name,
+            model_root,
+            target_duration,
+            top_k,
+            temperature,
+            kvcache,
+            repeat_prompt,
+            stop_repetition,
+            seed,
+            output_dir
+        ):
+            # Convert any empty strings to None for optional fields
+            dur = float(target_duration) if target_duration else None
+
+            out_path = run_inference(
+                reference_speech=reference_speech,
+                reference_text=reference_text if reference_text else None,
+                target_text=target_text,
+                model_name=model_name,
+                model_root=model_root,
+                target_duration=dur,
+                top_k=int(top_k),
+                temperature=float(temperature),
+                kvcache=int(kvcache),
+                repeat_prompt=int(repeat_prompt),
+                stop_repetition=int(stop_repetition),
+                seed=int(seed),
+                output_dir=output_dir
+            )
+            return out_path
+
+        # 3) Link the "Generate TTS" button
+        generate_button.click(
+            fn=gradio_inference,
+            inputs=[
+                reference_speech_input,
+                reference_text_box,
+                target_text_box,
+                model_name_box,
+                model_root_box,
+                reference_duration_box,
+                top_k_box,
+                temperature_box,
+                kvcache_box,
+                repeat_prompt_box,
+                stop_repetition_box,
+                seed_box,
+                output_dir_box
+            ],
+            outputs=[output_audio],
+        )
+
+    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)
+
+if __name__ == "__main__":
+    main()

inference_tts_utils.py

@@ -0,0 +1,155 @@
+import argparse, pickle
+import logging
+import os, random
+import numpy as np
+import torch
+import torchaudio
+
+from data.tokenizer import (
+    AudioTokenizer,
+    TextTokenizer,
+    tokenize_audio,
+    tokenize_text
+)
+import argparse, time, tqdm
+
+
+# this script only works for the musicgen architecture
+def get_args():
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--manifest_fn", type=str, default="path/to/eval_metadata_file")
+    parser.add_argument("--audio_root", type=str, default="path/to/audio_folder")
+    parser.add_argument("--exp_dir", type=str, default="path/to/model_folder")
+    parser.add_argument("--seed", type=int, default=1)
+    parser.add_argument("--codec_audio_sr", type=int, default=16000, help='the sample rate of audio that the codec is trained for')
+    parser.add_argument("--codec_sr", type=int, default=50, help='the sample rate of the codec codes')
+    parser.add_argument("--top_k", type=int, default=0, help="sampling param")
+    parser.add_argument("--top_p", type=float, default=0.8, help="sampling param")
+    parser.add_argument("--temperature", type=float, default=1.0, help="sampling param")
+    parser.add_argument("--output_dir", type=str, default=None)
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--signature", type=str, default=None, help="path to the encodec model")
+    parser.add_argument("--crop_concat", type=int, default=0)
+    parser.add_argument("--stop_repetition", type=int, default=-1, help="used for inference, when the number of consecutive repetition of a token is bigger than this, stop it")
+    parser.add_argument("--kvcache", type=int, default=1, help='if true, use kv cache, which is 4-8x faster than without')
+    parser.add_argument("--sample_batch_size", type=int, default=1, help="batch size for sampling, NOTE that it's not running inference for several samples, but duplicate one input sample batch_size times, and during inference, we only return the shortest generation")
+    parser.add_argument("--silence_tokens", type=str, default="[1388,1898,131]", help="note that if you are not using the pretrained encodec 6f79c6a8, make sure you specified it yourself, rather than using the default")
+    return parser.parse_args()
+
+
+@torch.no_grad()
+def inference_one_sample(model, model_args, phn2num, text_tokenizer, audio_tokenizer, audio_fn, target_text, device, decode_config, prompt_end_frame, target_generation_length, delay_pattern_increment, prefix_transcript=None, quiet=False, repeat_prompt=0, multi_trial=[]):
+    # seq_len_thres = 500 # 10s, 26% of the data in seed tts
+    # encode audio
+    encoded_frames = tokenize_audio(audio_tokenizer, audio_fn, offset=0, num_frames=prompt_end_frame)
+    # if sequence length is shorter than seq_len_thres, repeat the audio
+    # if encoded_frames.shape[2] < seq_len_thres:
+    #     encoded_frames = torch.cat([encoded_frames, encoded_frames, encoded_frames], dim=2)
+    #     doubled = True
+    single_encoded_frames = encoded_frames
+
+    if isinstance(repeat_prompt, int) and repeat_prompt > 0:
+        cur_repeat_prompt = repeat_prompt
+        while cur_repeat_prompt > 0:
+            encoded_frames = torch.cat([encoded_frames, single_encoded_frames], dim=2)
+            cur_repeat_prompt -= 1
+    elif isinstance(repeat_prompt, str) and repeat_prompt.lower() == "max":
+        repeat_prompt = 0
+        while encoded_frames.shape[2] + decode_config['codec_sr'] * target_generation_length + delay_pattern_increment + single_encoded_frames.shape[2] < model_args.audio_max_length * decode_config['codec_sr']:
+            encoded_frames = torch.cat([encoded_frames, single_encoded_frames], dim=2)
+            repeat_prompt += 1
+    if getattr(model_args, "y_sep_token", None) != None:
+        encoded_frames = torch.cat([encoded_frames, torch.LongTensor([model_args.y_sep_token]*model_args.n_codebooks).unsqueeze(0).unsqueeze(2).to(encoded_frames.device)], dim=2)
+    # print(encoded_frames.shape)
+    original_audio = encoded_frames.transpose(2,1) # [1,T,K]
+    assert original_audio.ndim==3 and original_audio.shape[0] == 1 and original_audio.shape[2] == model_args.n_codebooks, original_audio.shape
+
+    # phonemize
+    if isinstance(target_text, list):
+        text_tokens = [phn2num[phn] for phn in target_text if phn in phn2num]
+    else:
+        text_tokens = [phn2num[phn] for phn in
+                tokenize_text(
+                    text_tokenizer, text=target_text.strip()
+                ) if phn in phn2num
+            ]
+    if getattr(model_args, "x_sep_token", None) != None:
+        assert prefix_transcript != None, "prefix_transcript must be provided if x_sep_token is not None"
+    if prefix_transcript is not None:
+        if isinstance(prefix_transcript, list):
+            prefix_tokens = [phn2num[phn] for phn in prefix_transcript if phn in phn2num]
+        else:
+            prefix_tokens = [phn2num[phn] for phn in
+                tokenize_text(
+                    text_tokenizer, text=prefix_transcript.strip()
+                ) if phn in phn2num
+            ]
+        # if doubled:
+        #     prefix_tokens = prefix_tokens + prefix_tokens + prefix_tokens
+        single_prefix_tokens = prefix_tokens
+        while repeat_prompt > 0:
+            prefix_tokens = prefix_tokens + single_prefix_tokens
+            repeat_prompt -= 1
+        if getattr(model_args, "x_sep_token", None) != None:
+            text_tokens = prefix_tokens + [getattr(model_args, "x_sep_token", None)] + text_tokens
+        else:
+            text_tokens = prefix_tokens + text_tokens
+    if getattr(model_args, "add_eos_to_text", 0) != 0:
+        text_tokens.append(model_args.add_eos_to_text)
+    if getattr(model_args, "add_bos_to_text", 0) != 0:
+        text_tokens = [model_args.add_bos_to_text] + text_tokens
+    text_tokens = torch.LongTensor(text_tokens).unsqueeze(0)
+    text_tokens_lens = torch.LongTensor([text_tokens.shape[-1]])
+
+    if not quiet:
+        logging.info(f"original audio length: {original_audio.shape[1]} codec frames, which is {original_audio.shape[1]/decode_config['codec_sr']:.2f} sec.")
+
+
+    if getattr(model_args, "parallel_pattern", 0) != 0:
+        tgt_y_lens = torch.LongTensor([int(original_audio.shape[1] + decode_config['codec_sr'] * target_generation_length + 2)]) # parallel pattern, therefore only add the empty_token (i.e. the sos token) and eos (i.e. 2 more tokens). Note that the delayed pattern between, both sos and eos is counted (sos is counted in the n_codebooks, eos is counted in the 1)
+    else:
+        tgt_y_lens = torch.LongTensor([int(original_audio.shape[1] + decode_config['codec_sr'] * target_generation_length + delay_pattern_increment)]) # delay pattern increment has accounted for the added eos
+
+    # forward
+    assert decode_config['sample_batch_size'] <= 1
+    stime = time.time()
+    assert multi_trial == [] 
+    if not quiet:
+        logging.info(f"running inference with batch size 1")
+    concat_frames, gen_frames = model.inference_tts(
+        text_tokens.to(device),
+        text_tokens_lens.to(device),
+        original_audio[...,:model_args.n_codebooks].to(device), # [1,T,8]
+        tgt_y_lens = tgt_y_lens.to(device),
+        top_k=decode_config['top_k'],
+        top_p=decode_config['top_p'],
+        min_p=decode_config['min_p'],
+        temperature=decode_config['temperature'],
+        stop_repetition=decode_config['stop_repetition'],
+        kvcache=decode_config['kvcache'],
+        silence_tokens=eval(decode_config['silence_tokens']) if type(decode_config['silence_tokens'])==str else decode_config['silence_tokens']
+    ) # output is [1,K,T]
+    if not quiet:
+        logging.info(f"inference on one sample take: {time.time() - stime:.4f} sec.")
+
+        logging.info(f"generated encoded_frames.shape: {gen_frames.shape}, which is {gen_frames.shape[-1]/decode_config['codec_sr']} sec.")
+    
+    # for timestamp, codes in enumerate(gen_frames[0].transpose(1,0)):
+    #     logging.info(f"{timestamp}: {codes.tolist()}")
+    # decode (both original and generated)
+    # concat_sample = audio_tokenizer.decode(
+    #     [(concat_frames, None)] # [1,T,8] -> [1,8,T]
+    # )
+    if getattr(model_args, "y_sep_token", None) != None:
+        concat_frames = torch.cat([concat_frames[:, :, :original_audio.shape[1]-1], concat_frames[:, :, original_audio.shape[1]:]], dim=2)
+    concat_sample = audio_tokenizer.decode(
+        concat_frames # [1,8,T]
+    )
+    gen_sample = audio_tokenizer.decode(
+        gen_frames
+    )
+    #Empty cuda cache between runs
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    # return
+    return concat_sample, gen_sample

main.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+import torch, os
+
+from tqdm import tqdm
+import pickle
+import argparse
+import logging, datetime
+import torch.distributed as dist
+from config import MyParser
+from steps import trainer
+from copy_codebase import copy_codebase
+
+def world_info_from_env():
+    local_rank = int(os.environ["LOCAL_RANK"])
+    global_rank = int(os.environ["RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    return local_rank, global_rank, world_size
+
+if __name__ == "__main__":
+    formatter = (
+        "%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d || %(message)s"
+    )
+    logging.basicConfig(format=formatter, level=logging.INFO)
+    
+    torch.cuda.empty_cache()
+    args = MyParser().parse_args()
+    exp_dir = Path(args.exp_dir)
+    exp_dir.mkdir(exist_ok=True, parents=True)
+    logging.info(f"exp_dir: {str(exp_dir)}")
+
+    if args.resume and (os.path.exists("%s/bundle.pth" % args.exp_dir) or os.path.exists("%s/bundle_prev.pth" % args.exp_dir)):
+        if not os.path.exists("%s/bundle.pth" % args.exp_dir):
+            os.system(f"cp {args.exp_dir}/bundle_prev.pth {args.exp_dir}/bundle.pth")
+        resume = args.resume
+        assert(bool(args.exp_dir))
+        with open("%s/args.pkl" % args.exp_dir, "rb") as f:
+            old_args = pickle.load(f)
+        new_args = vars(args)
+        old_args = vars(old_args)
+        for key in new_args:
+            if key not in old_args or old_args[key] != new_args[key]:
+                old_args[key] = new_args[key]
+        args = argparse.Namespace(**old_args)
+        args.resume = resume
+    else:
+        args.resume = False
+        with open("%s/args.pkl" % args.exp_dir, "wb") as f:
+            pickle.dump(args, f)
+    
+    # make timeout longer (for generation)
+    timeout = datetime.timedelta(seconds=7200)  # 60 minutes
+
+    if args.multinodes:
+        _local_rank, _, _ = world_info_from_env()
+        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, timeout=timeout)
+    else:
+        dist.init_process_group(backend='nccl', init_method='env://', timeout=timeout)
+
+    if args.local_wandb:
+        os.environ["WANDB_MODE"] = "offline"
+
+    rank = dist.get_rank()
+    if rank == 0:
+        logging.info(args)
+        logging.info(f"exp_dir: {str(exp_dir)}")
+    world_size = dist.get_world_size()
+
+    local_rank = int(_local_rank) if args.multinodes else rank
+    num_devices= torch.cuda.device_count()
+    logging.info(f"{local_rank=}, {rank=}, {world_size=}, {type(local_rank)=}, {type(rank)=}, {type(world_size)=}")
+    for device_idx in range(num_devices):
+        device_name = torch.cuda.get_device_name(device_idx)
+        logging.info(f"Device {device_idx}: {device_name}")
+
+    torch.cuda.set_device(local_rank)
+    if rank == 0:
+        user_dir = os.path.expanduser("~")
+        codebase_name = "VoiceStar"
+        now = datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
+        copy_codebase(os.path.join(user_dir, codebase_name), os.path.join(exp_dir, f"{codebase_name}_{now}"), max_size_mb=5, gitignore_path=os.path.join(user_dir, codebase_name, ".gitignore"))
+    my_trainer = trainer.Trainer(args, world_size, rank, local_rank)
+    my_trainer.train()

models/modules/activation.py

@@ -0,0 +1,781 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/activation.py
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.nn import Linear, Module
+from torch.nn import functional as F
+from torch.nn.init import constant_, xavier_normal_, xavier_uniform_
+from torch.nn.modules.linear import NonDynamicallyQuantizableLinear
+from torch.nn.parameter import Parameter
+import logging
+from typing import Callable, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from torch.types import _dtype as DType
+else:
+    # The JIT doesn't understand Union, nor torch.dtype here
+    DType = int
+
+def _canonical_mask(
+        mask: Optional[Tensor],
+        mask_name: str,
+        other_type: Optional[DType],
+        other_name: str,
+        target_type: DType,
+        check_other: bool = True,
+) -> Optional[Tensor]:
+
+    if mask is not None:
+        _mask_dtype = mask.dtype
+        _mask_is_float = torch.is_floating_point(mask)
+        if _mask_dtype != torch.bool and not _mask_is_float:
+            raise AssertionError(
+                f"only bool and floating types of {mask_name} are supported")
+        if check_other and other_type is not None:
+            if _mask_dtype != other_type:
+                warnings.warn(
+                    f"Support for mismatched {mask_name} and {other_name} "
+                    "is deprecated. Use same type for both instead."
+                )
+        if not _mask_is_float:
+            mask = (
+                torch.zeros_like(mask, dtype=target_type)
+                .masked_fill_(mask, float("-inf"))
+            )
+    return mask
+
+def _in_projection_packed(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    w: Tensor,
+    b: Optional[Tensor] = None,
+) -> List[Tensor]:
+    r"""
+    Performs the in-projection step of the attention operation, using packed weights.
+    Output is a triple containing projection tensors for query, key and value.
+
+    Args:
+        q, k, v: query, key and value tensors to be projected. For self-attention,
+            these are typically the same tensor; for encoder-decoder attention,
+            k and v are typically the same tensor. (We take advantage of these
+            identities for performance if they are present.) Regardless, q, k and v
+            must share a common embedding dimension; otherwise their shapes may vary.
+        w: projection weights for q, k and v, packed into a single tensor. Weights
+            are packed along dimension 0, in q, k, v order.
+        b: optional projection biases for q, k and v, packed into a single tensor
+            in q, k, v order.
+
+    Shape:
+        Inputs:
+        - q: :math:`(..., E)` where E is the embedding dimension
+        - k: :math:`(..., E)` where E is the embedding dimension
+        - v: :math:`(..., E)` where E is the embedding dimension
+        - w: :math:`(E * 3, E)` where E is the embedding dimension
+        - b: :math:`E * 3` where E is the embedding dimension
+
+        Output:
+        - in output list :math:`[q', k', v']`, each output tensor will have the
+            same shape as the corresponding input tensor.
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            proj = F.linear(q, w, b)
+            # reshape to 3, E and not E, 3 is deliberate for better memory coalescing and keeping same order as chunk()
+            proj = proj.unflatten(-1, (3, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return proj[0], proj[1], proj[2]
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            q_proj = F.linear(q, w_q, b_q)
+            kv_proj = F.linear(k, w_kv, b_kv)
+            # reshape to 2, E and not E, 2 is deliberate for better memory coalescing and keeping same order as chunk()
+            kv_proj = kv_proj.unflatten(-1, (2, E)).unsqueeze(0).transpose(0, -2).squeeze(-2).contiguous()
+            return (q_proj, kv_proj[0], kv_proj[1])
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
+
+def _none_or_dtype(input: Optional[Tensor]) -> Optional[DType]:
+    if input is None:
+        return None
+    elif isinstance(input, torch.Tensor):
+        return input.dtype
+    raise RuntimeError("input to _none_or_dtype() must be None or torch.Tensor")
+
+def rotate_half(x):
+    x1 = x[..., :x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat([-x2, x1], dim=-1)
+
+def apply_rotary_pos_emb(q, k, q_sinu=None, k_sinu=None, sinu=None, unsqueeze_dim=1, args=None, q_offset=0): 
+    if sinu is not None:
+        q_emb = q * sinu['cos'][:, q_offset:q_offset+q.shape[2]].unsqueeze(unsqueeze_dim) + rotate_half(q) * sinu['sin'][:, q_offset:q_offset+q.shape[2]].unsqueeze(unsqueeze_dim)
+        k_emb = k * sinu['cos'][:, :k.shape[2]].unsqueeze(unsqueeze_dim) + rotate_half(k) * sinu['sin'][:, :k.shape[2]].unsqueeze(unsqueeze_dim) 
+    if q_sinu is not None:
+        assert sinu is None, "sinu must be None"
+        q_emb = q * q_sinu['cos'][:, :, q_offset:q_offset+q.shape[2]] + rotate_half(q) * q_sinu['sin'][:, :, q_offset:q_offset+q.shape[2]]
+        k_emb = k * k_sinu['cos'][:, :, :k.shape[2]] + rotate_half(k) * k_sinu['sin'][:, :, :k.shape[2]]
+    # else:
+    #     assert freqs is not None, "freqs must be provided"
+    #     assert key_lens is not None, "key_lens must be provided"
+    #     assert query_lens is not None, "query_lens must be provided"
+    #     # key_multiple
+    #     assert key_lens.ndim==1, key_lens
+    #     assert query_lens.ndim==1, query_lens
+    #     q_lens_expanded = query_lens.unsqueeze(-1).unsqueeze(-1) # [B, 1, 1]
+    #     k_lens_expanded = key_lens.unsqueeze(-1).unsqueeze(-1) # [B, 1, 1]
+    #     query_ids_multiple = q_lens_expanded / (q_lens_expanded - 1)
+    #     key_ids_multiple = k_lens_expanded / (k_lens_expanded - 1)
+    #     # freqs.shape [1, x_len_max, d]
+    #     # torch.set_printoptions(edgeitems=200)
+    #     # logging.info(f"{freqs[:, :q.shape[2]]=}")
+    #     # logging.info(f"{query_ids_multiple=}")
+    #     # logging.info(f"{key_ids_multiple=}")
+    #     # print(f"q.shape: {q.shape}")
+    #     # print(f"q_offset: {q_offset}")
+    #     # print(f"k.shape: {k.shape}")
+    #     q_emb = freqs[:, q_offset:q_offset+q.shape[-2]] * query_ids_multiple # [B, q_len_max, d]
+    #     k_emb = freqs[:, :k.shape[2]] * key_ids_multiple # [B, k_len_max, d]
+    #     # logging.info(f"{q_emb[0, :, :5]=}")
+    #     # logging.info(f"{k_emb[0, :, :5]=}")
+    #     multiple = k_lens_expanded if multiple_key_length else q_lens_expanded
+    #     if progress_no_multiple:
+    #         multiple = 1
+    #     q_emb = q_emb / q_lens_expanded * multiple * progress_scale
+    #     k_emb = k_emb / k_lens_expanded * multiple * progress_scale
+    #     q_cos = q_emb.cos().unsqueeze(unsqueeze_dim) # [B, 1, q_len_max, d] # 1 is for nhead
+    #     q_sin = q_emb.sin().unsqueeze(unsqueeze_dim)
+    #     k_cos = k_emb.cos().unsqueeze(unsqueeze_dim)
+    #     k_sin = k_emb.sin().unsqueeze(unsqueeze_dim)
+    #     # # visualize rotary pos emb with dummy feature
+    #     # q_tmp = torch.ones_like(q)
+    #     # k_tmp = torch.ones_like(k)
+    #     # q_tmp_emb = q_tmp * q_cos + rotate_half(q_tmp) * q_sin
+    #     # k_tmp_emb = k_tmp * k_cos + rotate_half(k_tmp) * k_sin
+    #     # sims = q_tmp_emb @ k_tmp_emb.transpose(-2, -1)
+    #     # import matplotlib.pyplot as plt
+    #     # for i, sim in enumerate(sims):
+    #     #     plt.imshow(sim[0][:query_lens[i], :key_lens[i]].detach().cpu().numpy())
+    #     #     plt.savefig(f"sim{i}_head0.png")
+    #     #     plt.imshow(sim[5][:query_lens[i], :key_lens[i]].detach().cpu().numpy())
+    #     #     plt.savefig(f"sim{i}_head5.png")
+    #     q_emb = q * q_cos + rotate_half(q) * q_sin
+    #     k_emb = k * k_cos + rotate_half(k) * k_sin
+    #     # # visualize the real attention weights
+    #     # sims = q_emb @ k_emb.transpose(-2, -1)
+    #     # from datetime import datetime
+    #     # from matplotlib import pyplot as plt
+    #     # now = datetime.now()
+    #     # for i, sim in enumerate(sims):
+    #     #     for ihead, si in enumerate(sim):
+    #     #         if query_lens[i] == key_lens[i]:
+    #     #             continue
+    #     #         plt.imshow(si[:query_lens[i], :key_lens[i]].detach().cpu().numpy())
+    #     #         plt.savefig(f"sample{i}_head{ihead}_{now.strftime('%Y-%m-%d_%H-%M-%S')}.png")
+    return q_emb, k_emb
+
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces as described in the paper:
+    `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_.
+
+    Multi-Head Attention is defined as:
+
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+
+    where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`.
+
+    ``forward()`` will use a special optimized implementation if all of the following
+    conditions are met:
+
+    - self attention is being computed (i.e., ``query``, ``key``, and ``value`` are the same tensor. This
+      restriction will be loosened in the future.)
+    - Either autograd is disabled (using ``torch.inference_mode`` or ``torch.no_grad``) or no tensor argument ``requires_grad``
+    - training is disabled (using ``.eval()``)
+    - dropout is 0
+    - ``add_bias_kv`` is ``False``
+    - ``add_zero_attn`` is ``False``
+    - ``batch_first`` is ``True`` and the input is batched
+    - ``kdim`` and ``vdim`` are equal to ``embed_dim``
+    - at most one of ``key_padding_mask`` or ``attn_mask`` is passed
+    - if a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ is passed, neither ``key_padding_mask``
+      nor ``attn_mask`` is passed
+
+    If the optimized implementation is in use, a
+    `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_ can be passed for
+    ``query``/``key``/``value`` to represent padding more efficiently than using a
+    padding mask. In this case, a `NestedTensor <https://pytorch.org/docs/stable/nested.html>`_
+    will be returned, and an additional speedup proportional to the fraction of the input
+    that is padding can be expected.
+
+    Args:
+        embed_dim: Total dimension of the model.
+        num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split
+            across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``).
+        dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout).
+        bias: If specified, adds bias to input / output projection layers. Default: ``True``.
+        add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``.
+        add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1.
+            Default: ``False``.
+        kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``).
+        vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``).
+        batch_first: If ``True``, then the input and output tensors are provided
+            as (batch, seq, feature). Default: ``False`` (seq, batch, feature).
+
+    Examples::
+
+        >>> # xdoctest: +SKIP
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+
+    """
+    __constants__ = ["batch_first"]
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        linear1_cls=Linear,
+        linear2_cls=Linear,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = (
+            self.kdim == embed_dim and self.vdim == embed_dim
+        )
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+
+        if add_bias_kv:
+            self.bias_k = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+            self.bias_v = Parameter(
+                torch.empty((1, 1, embed_dim), **factory_kwargs)
+            )
+        else:
+            self.bias_k = self.bias_v = None
+
+        if linear1_cls == Linear:
+            if not self._qkv_same_embed_dim:
+                self.q_proj_weight = Parameter(
+                    torch.empty((embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.k_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.kdim), **factory_kwargs)
+                )
+                self.v_proj_weight = Parameter(
+                    torch.empty((embed_dim, self.vdim), **factory_kwargs)
+                )
+                self.register_parameter("in_proj_weight", None)
+            else:
+                # go down this route with music_gen
+                self.in_proj_weight = Parameter(
+                    torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)
+                )
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+            if bias: # True by default
+                self.in_proj_bias = Parameter(
+                    torch.empty(3 * embed_dim, **factory_kwargs)
+                )
+            else:
+                self.register_parameter("in_proj_bias", None)
+            self.out_proj = NonDynamicallyQuantizableLinear(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            self._reset_parameters()
+        else:
+            if not self._qkv_same_embed_dim:
+                raise NotImplementedError
+            else:
+                self.in_proj_linear = linear1_cls(
+                    embed_dim, 3 * embed_dim, bias=bias, **factory_kwargs
+                )
+                self.in_proj_weight = self.in_proj_linear.weight
+
+                self.register_parameter("q_proj_weight", None)
+                self.register_parameter("k_proj_weight", None)
+                self.register_parameter("v_proj_weight", None)
+
+                if bias:
+                    self.in_proj_bias = self.in_proj_linear.bias
+                else:
+                    self.register_parameter("in_proj_bias", None)
+
+            self.out_proj = linear2_cls(
+                embed_dim, embed_dim, bias=bias, **factory_kwargs
+            )
+
+            if self.bias_k is not None:
+                xavier_normal_(self.bias_k)
+            if self.bias_v is not None:
+                xavier_normal_(self.bias_v)
+
+        self.add_zero_attn = add_zero_attn
+
+    def _reset_parameters(self):
+        if self._qkv_same_embed_dim:
+            xavier_uniform_(self.in_proj_weight)
+        else:
+            xavier_uniform_(self.q_proj_weight)
+            xavier_uniform_(self.k_proj_weight)
+            xavier_uniform_(self.v_proj_weight)
+
+        if self.in_proj_bias is not None:
+            constant_(self.in_proj_bias, 0.0)
+            constant_(self.out_proj.bias, 0.0)
+
+        if self.bias_k is not None:
+            xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            xavier_normal_(self.bias_v)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if "_qkv_same_embed_dim" not in state:
+            state["_qkv_same_embed_dim"] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        key_padding_mask: Optional[Tensor] = None,
+        need_weights: bool = True,
+        attn_mask: Optional[Tensor] = None,
+        average_attn_weights: bool = True,
+        past: Optional[Tensor] = None,
+        q_sinu = None,
+        k_sinu = None,
+        sinu = None,
+        args = None,
+        q_offset = 0,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        r"""
+        Args:
+            query: Query embeddings of shape :math:`(L, E_q)` for unbatched input, :math:`(L, N, E_q)` when ``batch_first=False``
+                or :math:`(N, L, E_q)` when ``batch_first=True``, where :math:`L` is the target sequence length,
+                :math:`N` is the batch size, and :math:`E_q` is the query embedding dimension ``embed_dim``.
+                Queries are compared against key-value pairs to produce the output.
+                See "Attention Is All You Need" for more details.
+            key: Key embeddings of shape :math:`(S, E_k)` for unbatched input, :math:`(S, N, E_k)` when ``batch_first=False``
+                or :math:`(N, S, E_k)` when ``batch_first=True``, where :math:`S` is the source sequence length,
+                :math:`N` is the batch size, and :math:`E_k` is the key embedding dimension ``kdim``.
+                See "Attention Is All You Need" for more details.
+            value: Value embeddings of shape :math:`(S, E_v)` for unbatched input, :math:`(S, N, E_v)` when
+                ``batch_first=False`` or :math:`(N, S, E_v)` when ``batch_first=True``, where :math:`S` is the source
+                sequence length, :math:`N` is the batch size, and :math:`E_v` is the value embedding dimension ``vdim``.
+                See "Attention Is All You Need" for more details.
+            key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key``
+                to ignore for the purpose of attention (i.e. treat as "padding"). For unbatched `query`, shape should be :math:`(S)`.
+                Binary and byte masks are supported.
+                For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for
+                the purpose of attention. For a float mask, it will be directly added to the corresponding ``key`` value.
+            need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``.
+                Default: ``True``.
+            attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape
+                :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size,
+                :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be
+                broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch.
+                Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the
+                corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the
+                corresponding position is not allowed to attend. For a float mask, the mask values will be added to
+                the attention weight.
+            average_attn_weights: If true, indicates that the returned ``attn_weights`` should be averaged across
+                heads. Otherwise, ``attn_weights`` are provided separately per head. Note that this flag only has an
+                effect when ``need_weights=True``. Default: ``True`` (i.e. average weights across heads)
+            sinu: for direct original rope positional encoding
+            freqs: for progress monitoring with rope positional encoding
+            q_offset: for progress monitoring with rope positional encoding, during inference when kvcache is on, most of the time query is only of length 1, so we need to offset the query to get the correct progress
+
+        Outputs:
+            - **attn_output** - Attention outputs of shape :math:`(L, E)` when input is unbatched,
+              :math:`(L, N, E)` when ``batch_first=False`` or :math:`(N, L, E)` when ``batch_first=True``,
+              where :math:`L` is the target sequence length, :math:`N` is the batch size, and :math:`E` is the
+              embedding dimension ``embed_dim``.
+            - **attn_output_weights** - Only returned when ``need_weights=True``. If ``average_attn_weights=True``,
+              returns attention weights averaged across heads of shape :math:`(L, S)` when input is unbatched or
+              :math:`(N, L, S)`, where :math:`N` is the batch size, :math:`L` is the target sequence length, and
+              :math:`S` is the source sequence length. If ``average_attn_weights=False``, returns attention weights per
+              head of shape :math:`(\text{num\_heads}, L, S)` when input is unbatched or :math:`(N, \text{num\_heads}, L, S)`.
+
+            .. note::
+                `batch_first` argument is ignored for unbatched inputs.
+        """
+        is_batched = query.dim() == 3
+        if key_padding_mask is not None:
+            _kpm_dtype = key_padding_mask.dtype
+            if _kpm_dtype != torch.bool and not torch.is_floating_point(
+                key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        why_not_fast_path = ""
+        if not is_batched:
+            why_not_fast_path = f"input not batched; expected query.dim() of 3 but got {query.dim()}"
+        elif query is not key or key is not value:
+            # When lifting this restriction, don't forget to either
+            # enforce that the dtypes all match or test cases where
+            # they don't!
+            why_not_fast_path = "non-self attention was used (query, key, and value are not the same Tensor)"
+        elif (
+            self.in_proj_bias is not None
+            and query.dtype != self.in_proj_bias.dtype
+        ):
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_bias ({self.in_proj_bias.dtype}) don't match"
+        elif (
+            self.in_proj_weight is not None
+            and query.dtype != self.in_proj_weight.dtype
+        ):
+            # this case will fail anyway, but at least they'll get a useful error message.
+            why_not_fast_path = f"dtypes of query ({query.dtype}) and self.in_proj_weight ({self.in_proj_weight.dtype}) don't match"
+        elif self.training:
+            why_not_fast_path = "training is enabled"
+        elif not self.batch_first:
+            why_not_fast_path = "batch_first was not True"
+        elif self.bias_k is not None:
+            why_not_fast_path = "self.bias_k was not None"
+        elif self.bias_v is not None:
+            why_not_fast_path = "self.bias_v was not None"
+        elif self.dropout:
+            why_not_fast_path = f"dropout was {self.dropout}, required zero"
+        elif self.add_zero_attn:
+            why_not_fast_path = "add_zero_attn was enabled"
+        elif not self._qkv_same_embed_dim:
+            why_not_fast_path = "_qkv_same_embed_dim was not True"
+        elif attn_mask is not None:
+            why_not_fast_path = "attn_mask was not None"
+        elif query.is_nested and key_padding_mask is not None:
+            why_not_fast_path = (
+                "key_padding_mask is not supported with NestedTensor input"
+            )
+        elif self.num_heads % 2 == 1:
+            why_not_fast_path = "num_heads is odd"
+        elif torch.is_autocast_enabled():
+            why_not_fast_path = "autocast is enabled"
+
+        if not why_not_fast_path:
+            tensor_args = (
+                query,
+                key,
+                value,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.out_proj.weight,
+                self.out_proj.bias,
+            )
+            # We have to use list comprehensions below because TorchScript does not support
+            # generator expressions.
+            if torch.overrides.has_torch_function(tensor_args):
+                why_not_fast_path = "some Tensor argument has_torch_function"
+            elif not all(
+                [
+                    (x is None or x.is_cuda or "cpu" in str(x.device))
+                    for x in tensor_args
+                ]
+            ):
+                why_not_fast_path = (
+                    "some Tensor argument is neither CUDA nor CPU"
+                )
+            elif torch.is_grad_enabled() and any(
+                [x is not None and x.requires_grad for x in tensor_args]
+            ):
+                why_not_fast_path = (
+                    "grad is enabled and at least one of query or the "
+                    "input/output projection weights or biases requires_grad"
+                )
+            if not why_not_fast_path:
+                return torch._native_multi_head_attention(
+                    query,
+                    key,
+                    value,
+                    self.embed_dim,
+                    self.num_heads,
+                    self.in_proj_weight,
+                    self.in_proj_bias,
+                    self.out_proj.weight,
+                    self.out_proj.bias,
+                    key_padding_mask
+                    if key_padding_mask is not None
+                    else attn_mask,
+                    need_weights,
+                    average_attn_weights,
+                    1
+                    if key_padding_mask is not None
+                    else 0
+                    if attn_mask is not None
+                    else None,
+                )
+
+        any_nested = query.is_nested or key.is_nested or value.is_nested
+        assert not any_nested, (
+            "MultiheadAttention does not support NestedTensor outside of its fast path. "
+            + f"The fast path was not hit because {why_not_fast_path}"
+        )
+
+        if self.batch_first and is_batched:
+            # make sure that the transpose op does not affect the "is" property
+            if key is value:
+                if query is key:
+                    query = key = value = query.transpose(1, 0)
+                else:
+                    query, key = [x.transpose(1, 0) for x in (query, key)]
+                    value = key
+            else:
+                query, key, value = [
+                    x.transpose(1, 0) for x in (query, key, value)
+                ]
+
+        if not self._qkv_same_embed_dim:
+            attn_output, attn_output_weights = F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                self.in_proj_weight,
+                self.in_proj_bias,
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask,
+                need_weights=need_weights,
+                attn_mask=attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight,
+                k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight,
+                average_attn_weights=average_attn_weights,
+            )
+        else:
+            # music_gen should go down this route
+            # logging.info("using in_proj_weight")
+            # attn_output, attn_output_weights = F.multi_head_attention_forward(
+            #     query,
+            #     key,
+            #     value,
+            #     self.embed_dim,
+            #     self.num_heads,
+            #     self.in_proj_weight,
+            #     self.in_proj_bias,
+            #     self.bias_k,
+            #     self.bias_v,
+            #     self.add_zero_attn,
+            #     self.dropout,
+            #     self.out_proj.weight,
+            #     self.out_proj.bias,
+            #     training=self.training,
+            #     key_padding_mask=key_padding_mask,
+            #     need_weights=need_weights,
+            #     attn_mask=attn_mask,
+            #     average_attn_weights=average_attn_weights,
+            # )
+            # re-write the self.attention here, to get k, v cache
+            tgt_len, bsz, embed_dim = query.shape
+            src_len, _, _ = key.shape
+            num_heads = self.num_heads
+            key_padding_mask = _canonical_mask(
+                mask=key_padding_mask,
+                mask_name="key_padding_mask",
+                other_type=_none_or_dtype(attn_mask),
+                other_name="attn_mask",
+                target_type=query.dtype
+            )
+            attn_mask = _canonical_mask(
+                            mask=attn_mask,
+                            mask_name="attn_mask",
+                            other_type=None,
+                            other_name="",
+                            target_type=query.dtype,
+                            check_other=False,
+                            )
+            head_dim = self.embed_dim // self.num_heads
+            assert head_dim * self.num_heads == self.embed_dim, f"embed_dim {self.embed_dim} not divisible by num_heads {self.num_heads}"
+            assert key.shape == value.shape, f"key shape {key.shape} does not match value shape {value.shape}"
+            q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
+            # k_present, v_present = k, v
+            
+            #
+            # reshape q, k, v for multihead attention and make em batch first
+            #
+            
+
+            
+            q = q.view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+            k = k.view(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
+            v = v.view(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1) # (bsz * num_heads, src_len, head_dim)
+            src_len = k.size(1)
+            if past is not None and past.ndim > 2:
+                expected_src_len = src_len + past[0].shape[-2]
+            else:
+                expected_src_len = src_len
+
+
+            # ensure attn_mask's dim is 3
+            if attn_mask is not None:
+                if attn_mask.dim() == 2:
+                    correct_2d_size = (tgt_len, expected_src_len)
+                    if attn_mask.shape != correct_2d_size:
+                        raise RuntimeError(f"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.")
+                    attn_mask = attn_mask.unsqueeze(0)
+                elif attn_mask.dim() == 3:
+                    correct_3d_size = (bsz * num_heads, tgt_len, expected_src_len)
+                    if attn_mask.shape != correct_3d_size:
+                        raise RuntimeError(f"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.")
+                else:
+                    raise RuntimeError(f"attn_mask's dimension {attn_mask.dim()} is not supported")
+            
+            if key_padding_mask is not None:
+                assert key_padding_mask.shape == (bsz, expected_src_len), \
+                    f"expecting key_padding_mask shape of {(bsz, expected_src_len)}, but got {key_padding_mask.shape}"
+                key_padding_mask = key_padding_mask.view(bsz, 1, 1, expected_src_len).   \
+                    expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, expected_src_len)
+                if attn_mask is None:
+                    attn_mask = key_padding_mask
+                else:
+                    attn_mask = attn_mask + key_padding_mask
+            
+            if not self.training:
+                dropout_p = 0.0
+            else:
+                dropout_p = self.dropout
+
+            if need_weights:
+                raise NotImplementedError("need_weights not implemented for music_gen")
+                # B, Nt, E = q.shape
+                # q_scaled = q / math.sqrt(E)
+
+                # assert not (is_causal and attn_mask is None), "FIXME: is_causal not implemented for need_weights"
+
+                # if attn_mask is not None:
+                #     attn_output_weights = torch.baddbmm(attn_mask, q_scaled, k.transpose(-2, -1))
+                # else:
+                #     attn_output_weights = torch.bmm(q_scaled, k.transpose(-2, -1))
+                # attn_output_weights = softmax(attn_output_weights, dim=-1)
+                # if dropout_p > 0.0:
+                #     attn_output_weights = dropout(attn_output_weights, p=dropout_p)
+
+                # attn_output = torch.bmm(attn_output_weights, v)
+
+                # attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len * bsz, embed_dim)
+                # attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+                # attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+
+                # # optionally average attention weights over heads
+                # attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+                # if average_attn_weights:
+                #     attn_output_weights = attn_output_weights.mean(dim=1)
+
+                # if not is_batched:
+                #     # squeeze the output if input was unbatched
+                #     attn_output = attn_output.squeeze(1)
+                #     attn_output_weights = attn_output_weights.squeeze(0)
+                # return attn_output, attn_output_weights
+            else:
+                # attn_mask can be either (L,S) or (N*num_heads, L, S)
+                # if attn_mask's shape is (1, L, S) we need to unsqueeze to (1, 1, L, S)
+                # in order to match the input for SDPA of (N, num_heads, L, S)
+                if attn_mask is not None:
+                    if attn_mask.size(0) == 1 and attn_mask.dim() == 3:
+                        attn_mask = attn_mask.unsqueeze(0)
+                    else:
+                        attn_mask = attn_mask.view(bsz, num_heads, -1, expected_src_len)
+
+                q = q.view(bsz, num_heads, tgt_len, head_dim)
+                k = k.view(bsz, num_heads, src_len, head_dim)
+                v = v.view(bsz, num_heads, src_len, head_dim)
+                # logging.info(f"shape of past: {past.shape}")
+                if past is not None:
+                    present = torch.stack([k, v], dim=0) # (2, bsz, num_heads, src_len, head_dim)
+                    if past.ndim > 2: # this means we use kvcache, otherwise we just pass in a placeholder, but not actually using kvcache
+                        pk, pv = past
+                        k = torch.cat([pk, k], dim=-2)
+                        v = torch.cat([pv, v], dim=-2)
+                else:
+                    present = None
+                # when using kvcache, need to offset postion of q when applying rotary pos emb
+                # here we assume that this kvcache is only used in self-attention, and therefore k and q always have the same seq_len
+                # rope positional encoding
+                if sinu is not None:
+                    # direct rotary 
+                    # logging.info("perform rotary positional encoding")
+                    q, k = apply_rotary_pos_emb(q, k, sinu=sinu, args = args, q_offset=q_offset)
+                if q_sinu is not None:
+                    assert sinu is None, "sinu and q_sinu cannot be used together"
+                    assert k_sinu is not None, "k_sinu must be provided"
+                    q, k = apply_rotary_pos_emb(q, k, q_sinu=q_sinu, k_sinu=k_sinu, args = args, q_offset=q_offset)
+                
+                # if self.training and it's cross attention, will get attention_weights
+                if args != None and self.training and getattr(args, "attention_alignment_loss", 0) and not (query is key):
+                    attention_weights = q @ k.transpose(-1, -2)
+                else:
+                    attention_weights = None
+                attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask, dropout_p, is_causal=False)
+                attn_output = attn_output.permute(2, 0, 1, 3).contiguous().view(bsz * tgt_len, embed_dim)
+
+                attn_output = F.linear(attn_output, self.out_proj.weight, self.out_proj.bias)
+                attn_output = attn_output.view(tgt_len, bsz, attn_output.size(1))
+                if not is_batched:
+                    # squeeze the output if input was unbatched
+                    attn_output = attn_output.squeeze(1)
+                # if self.training:
+                #     return attn_output, None
+                # else:
+                #     return (attn_output, present), None
+
+        # harded coded, the code do not support returning attn weigths yet
+        attn_output_weights=None
+        if self.batch_first and is_batched:
+            if attention_weights != None:
+                return {"attn_output": attn_output.transpose(1, 0), "attention_weights": attention_weights}, present
+            return attn_output.transpose(1, 0), present
+        else:
+            if attention_weights != None:
+                return {"attn_output": attn_output, "attention_weights": attention_weights}, present
+            return attn_output, present
+

models/modules/embedding.py

@@ -0,0 +1,158 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/embedding.py
+# Copyright    2023                             (authors: Feiteng Li)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import logging
+import torch
+import torch.nn as nn
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        vocab_size: int,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        self.vocab_size = vocab_size
+        self.dim_model = dim_model
+
+        self.dropout = torch.nn.Dropout(p=dropout)
+        self.word_embeddings = nn.Embedding(self.vocab_size, self.dim_model)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.word_embeddings.weight
+
+    def embedding(self, index: int) -> torch.Tensor:
+        return self.word_embeddings.weight[index : index + 1]
+
+    def forward(self, x: torch.Tensor):
+        X = self.word_embeddings(x)
+        X = self.dropout(X)
+
+        return X
+
+
+class SinePositionalEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+    ):
+        super().__init__()
+        self.dim_model = dim_model
+        self.x_scale = math.sqrt(dim_model) if scale else 1.0
+        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+        self.reverse = False
+        self.pe = None
+        self.extend_pe(torch.tensor(0.0).expand(1, 4000))
+
+    def extend_pe(self, x):
+        """Reset the positional encodings."""
+        if self.pe is not None:
+            if self.pe.size(1) >= x.size(1):
+                if self.pe.dtype != x.dtype or self.pe.device != x.device:
+                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
+                return
+        pe = torch.zeros(x.size(1), self.dim_model)
+        if self.reverse:
+            position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(1)
+        else:
+            position = torch.arange(
+                0, x.size(1), dtype=torch.float32
+            ).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.dim_model, 2, dtype=torch.float32)
+            * -(math.log(10000.0) / self.dim_model)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.pe = pe.to(device=x.device, dtype=x.dtype).detach()
+
+    def forward(self, x: torch.Tensor, *args) -> torch.Tensor:
+        self.extend_pe(x)
+        output = x.unsqueeze(-1) if x.ndim == 2 else x
+        output = output * self.x_scale + self.alpha * self.pe[:, : x.size(1)]
+        return self.dropout(output)
+
+
+class SinePositionalEmbedding_progress(nn.Module):
+    def __init__(
+        self,
+        dim_model: int,
+        dropout: float = 0.0,
+        scale: bool = False,
+        alpha: bool = False,
+        args = None
+    ):
+        super().__init__()
+        self.args = args
+        self.dim_model = dim_model
+        self.x_scale = math.sqrt(dim_model) if scale else 1.0
+        self.alpha = nn.Parameter(torch.ones(1), requires_grad=alpha)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+        self.reverse = False
+        self.div_term = torch.exp(
+            torch.arange(0, self.dim_model, 2, dtype=torch.float32)
+            * -(math.log(args.sinusoidal_base) / self.dim_model)
+        ).unsqueeze(0).unsqueeze(0) # [1, 1, dim_model//2]
+        self.position = None
+        self.extend_position(torch.tensor(0.0).expand(1, 10000))
+        self.progress_scale = getattr(args, "progress_scale", 1.0)
+
+    def extend_position(self, x):
+        """Reset the positional encodings."""
+        if self.position is not None:
+            if self.div_term.dtype != x.dtype or self.div_term.device != x.device:
+                self.div_term = self.div_term.to(dtype=x.dtype, device=x.device)
+            if self.position.size(1) >= x.size(1):
+                if self.position.dtype != x.dtype or self.position.device != x.device:
+                    self.position = self.position.to(dtype=x.dtype, device=x.device)
+                return
+        if self.reverse:
+            self.position = torch.arange(
+                x.size(1) - 1, -1, -1.0, dtype=torch.float32
+            ).unsqueeze(0).unsqueeze(2).to(x)
+        else:
+            self.position = torch.arange(
+                0, x.size(1), dtype=torch.float32
+            ).unsqueeze(0).unsqueeze(2).to(x) # [1, seq_len, 1]
+
+    def forward(self, x: torch.Tensor, x_lens: torch.Tensor) -> torch.Tensor:
+        assert x.ndim == 3, x.shape
+        self.extend_position(x)
+        x_lens = x_lens.unsqueeze(1).unsqueeze(2) # [B, 1, 1]
+        multiple = x_lens / (x_lens - 1)
+        progress = self.position[:, :x.shape[1]] * multiple / x_lens * self.progress_scale
+        # torch.set_printoptions(edgeitems=100)
+        # for i in range(x_lens.shape[0]):
+        #     logging.info(f"{progress[i, :x_lens[i,0,0], 0]}")
+        invfreq = self.div_term * progress # might want to use a scale term here
+        pe = torch.zeros_like(x)
+        pe[..., 0::2] = torch.sin(invfreq)
+        pe[..., 1::2] = torch.cos(invfreq)
+        output = x * self.x_scale + self.alpha * pe
+        return self.dropout(output)

models/modules/sampling.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn.functional as F
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if top_k > 0:
+        top_k = min(
+            max(top_k, min_tokens_to_keep), logits.size(-1)
+        )  # Safety check
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(
+            F.softmax(sorted_logits, dim=-1), dim=-1
+        )
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+            ..., :-1
+        ].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+        # scatter sorted tensors to original indexing
+        indices_to_remove = sorted_indices_to_remove.scatter(
+            1, sorted_indices, sorted_indices_to_remove
+        )
+        logits[indices_to_remove] = filter_value
+    return logits
+    
+def topk_sampling(logits, top_k=10, top_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token

models/modules/scaling.py

@@ -0,0 +1,1406 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/modules/scaling.py
+# Copyright    2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import collections
+import logging
+import random
+import math
+from functools import reduce
+from itertools import repeat
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Embedding as ScaledEmbedding
+
+# from valle.utils import Transpose
+
+class Transpose(nn.Identity):
+    """(N, T, D) -> (N, D, T)"""
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return input.transpose(1, 2)
+    
+class ActivationBalancerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        scale_factor: Tensor,
+        sign_factor: Optional[Tensor],
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        if sign_factor is None:
+            ctx.save_for_backward(xgt0, scale_factor)
+        else:
+            ctx.save_for_backward(xgt0, scale_factor, sign_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        if len(ctx.saved_tensors) == 3:
+            xgt0, scale_factor, sign_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+                sign_factor = sign_factor.unsqueeze(-1)
+            factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        else:
+            xgt0, scale_factor = ctx.saved_tensors
+            for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+                scale_factor = scale_factor.unsqueeze(-1)
+            factor = scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+def _compute_scale_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_abs: float,
+    max_abs: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    x_abs_mean = torch.mean(x.abs(), dim=sum_dims).to(torch.float32)
+
+    if min_abs == 0.0:
+        below_threshold = 0.0
+    else:
+        # below_threshold is 0 if x_abs_mean > min_abs, can be at most max_factor if
+        # x_abs)_mean , min_abs.
+        below_threshold = (
+            (min_abs - x_abs_mean) * (gain_factor / min_abs)
+        ).clamp(min=0, max=max_factor)
+
+    above_threshold = ((x_abs_mean - max_abs) * (gain_factor / max_abs)).clamp(
+        min=0, max=max_factor
+    )
+
+    return below_threshold - above_threshold
+
+
+def _compute_sign_factor(
+    x: Tensor,
+    channel_dim: int,
+    min_positive: float,
+    max_positive: float,
+    gain_factor: float,
+    max_factor: float,
+) -> Tensor:
+    if channel_dim < 0:
+        channel_dim += x.ndim
+    sum_dims = [d for d in range(x.ndim) if d != channel_dim]
+    proportion_positive = torch.mean((x > 0).to(torch.float32), dim=sum_dims)
+    if min_positive == 0.0:
+        factor1 = 0.0
+    else:
+        # 0 if proportion_positive >= min_positive, else can be
+        # as large as max_factor.
+        factor1 = (
+            (min_positive - proportion_positive) * (gain_factor / min_positive)
+        ).clamp_(min=0, max=max_factor)
+
+    if max_positive == 1.0:
+        factor2 = 0.0
+    else:
+        # 0 if self.proportion_positive <= max_positive, else can be
+        # as large as -max_factor.
+        factor2 = (
+            (proportion_positive - max_positive)
+            * (gain_factor / (1.0 - max_positive))
+        ).clamp_(min=0, max=max_factor)
+    sign_factor = factor1 - factor2
+    # require min_positive != 0 or max_positive != 1:
+    assert not isinstance(sign_factor, float)
+    return sign_factor
+
+
+class ActivationScaleBalancerFunction(torch.autograd.Function):
+    """
+    This object is used in class ActivationBalancer when the user specified
+    min_positive=0, max_positive=1, so there are no constraints on the signs
+    of the activations and only the absolute value has a constraint.
+    """
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        sign_factor: Tensor,
+        scale_factor: Tensor,
+        channel_dim: int,
+    ) -> Tensor:
+        if channel_dim < 0:
+            channel_dim += x.ndim
+        ctx.channel_dim = channel_dim
+        xgt0 = x > 0
+        ctx.save_for_backward(xgt0, sign_factor, scale_factor)
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None]:
+        xgt0, sign_factor, scale_factor = ctx.saved_tensors
+        for _ in range(ctx.channel_dim, x_grad.ndim - 1):
+            sign_factor = sign_factor.unsqueeze(-1)
+            scale_factor = scale_factor.unsqueeze(-1)
+
+        factor = sign_factor + scale_factor * (xgt0.to(x_grad.dtype) - 0.5)
+        neg_delta_grad = x_grad.abs() * factor
+        return (
+            x_grad - neg_delta_grad,
+            None,
+            None,
+            None,
+        )
+
+
+class RandomClampFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        min: Optional[float],
+        max: Optional[float],
+        prob: float,
+        reflect: float,
+    ) -> Tensor:
+        x_clamped = torch.clamp(x, min=min, max=max)
+        mask = torch.rand_like(x) < prob
+        ans = torch.where(mask, x_clamped, x)
+        if x.requires_grad:
+            ctx.save_for_backward(ans == x)
+            ctx.reflect = reflect
+        if reflect != 0.0:
+            ans = ans * (1.0 + reflect) - (x * reflect)
+        return ans
+
+    @staticmethod
+    def backward(
+        ctx, ans_grad: Tensor
+    ) -> Tuple[Tensor, None, None, None, None]:
+        (is_same,) = ctx.saved_tensors
+        x_grad = ans_grad * is_same.to(ans_grad.dtype)
+        reflect = ctx.reflect
+        if reflect != 0.0:
+            x_grad = x_grad * (1.0 + reflect) - (ans_grad * reflect)
+        return x_grad, None, None, None, None
+
+
+def random_clamp(
+    x: Tensor,
+    min: Optional[float] = None,
+    max: Optional[float] = None,
+    prob: float = 0.5,
+    reflect: float = 0.0,
+):
+    return RandomClampFunction.apply(x, min, max, prob, reflect)
+
+
+def random_cast_to_half(x: Tensor, min_abs: float = 5.0e-06) -> Tensor:
+    """
+    A randomized way of casting a floating point value to half precision.
+    """
+    if x.dtype == torch.float16:
+        return x
+    x_abs = x.abs()
+    is_too_small = x_abs < min_abs
+    # for elements where is_too_small is true, random_val will contain +-min_abs with
+    # probability (x.abs() / min_abs), and 0.0 otherwise.  [so this preserves expectations,
+    # for those elements].
+    random_val = min_abs * x.sign() * (torch.rand_like(x) * min_abs < x_abs)
+    return torch.where(is_too_small, random_val, x).to(torch.float16)
+
+
+class RandomGradFunction(torch.autograd.Function):
+    """
+    Does nothing in forward pass; in backward pass, gets rid of very small grads using
+    randomized approach that preserves expectations (intended to reduce roundoff).
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, min_abs: float) -> Tensor:
+        ctx.min_abs = min_abs
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor) -> Tuple[Tensor, None]:
+        if ans_grad.dtype == torch.float16:
+            return (
+                random_cast_to_half(
+                    ans_grad.to(torch.float32), min_abs=ctx.min_abs
+                ),
+                None,
+            )
+        else:
+            return ans_grad, None
+
+
+class RandomGrad(torch.nn.Module):
+    """
+    Gets rid of very small gradients using an expectation-preserving method, intended to increase
+    accuracy of training when using amp (automatic mixed precision)
+    """
+
+    def __init__(self, min_abs: float = 5.0e-06):
+        super(RandomGrad, self).__init__()
+        self.min_abs = min_abs
+
+    def forward(self, x: Tensor):
+        if (
+            torch.jit.is_scripting()
+            or not self.training
+            or torch.jit.is_tracing()
+        ):
+            return x
+        else:
+            return RandomGradFunction.apply(x, self.min_abs)
+
+
+class SoftmaxFunction(torch.autograd.Function):
+    """
+    Tries to handle half-precision derivatives in a randomized way that should
+    be more accurate for training than the default behavior.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor, dim: int):
+        ans = x.softmax(dim=dim)
+        # if x dtype is float16, x.softmax() returns a float32 because
+        # (presumably) that op does not support float16, and autocast
+        # is enabled.
+        if torch.is_autocast_enabled():
+            ans = ans.to(torch.float16)
+        ctx.save_for_backward(ans)
+        ctx.x_dtype = x.dtype
+        ctx.dim = dim
+        return ans
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        (ans,) = ctx.saved_tensors
+        with torch.cuda.amp.autocast(enabled=False):
+            ans_grad = ans_grad.to(torch.float32)
+            ans = ans.to(torch.float32)
+            x_grad = ans_grad * ans
+            x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True)
+            return x_grad, None
+
+
+def softmax(x: Tensor, dim: int):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x.softmax(dim)
+
+    return SoftmaxFunction.apply(x, dim)
+
+
+class MaxEigLimiterFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        coeffs: Tensor,
+        direction: Tensor,
+        channel_dim: int,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.channel_dim = channel_dim
+        ctx.grad_scale = grad_scale
+        ctx.save_for_backward(x.detach(), coeffs.detach(), direction.detach())
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad, *args):
+        with torch.enable_grad():
+            (x_orig, coeffs, new_direction) = ctx.saved_tensors
+            x_orig.requires_grad = True
+            num_channels = x_orig.shape[ctx.channel_dim]
+            x = x_orig.transpose(ctx.channel_dim, -1).reshape(-1, num_channels)
+            new_direction.requires_grad = False
+            x = x - x.mean(dim=0)
+            x_var = (x ** 2).mean()
+            x_residual = x - coeffs * new_direction
+            x_residual_var = (x_residual ** 2).mean()
+            # `variance_proportion` is the proportion of the variance accounted for
+            # by the top eigen-direction.  This is to be minimized.
+            variance_proportion = (x_var - x_residual_var) / (x_var + 1.0e-20)
+            variance_proportion.backward()
+        x_orig_grad = x_orig.grad
+        x_extra_grad = (
+            x_orig.grad
+            * ctx.grad_scale
+            * x_grad.norm()
+            / (x_orig_grad.norm() + 1.0e-20)
+        )
+        return x_grad + x_extra_grad.detach(), None, None, None, None
+
+
+class BasicNorm(torch.nn.Module):
+    """
+    This is intended to be a simpler, and hopefully cheaper, replacement for
+    LayerNorm.  The observation this is based on, is that Transformer-type
+    networks, especially with pre-norm, sometimes seem to set one of the
+    feature dimensions to a large constant value (e.g. 50), which "defeats"
+    the LayerNorm because the output magnitude is then not strongly dependent
+    on the other (useful) features.  Presumably the weight and bias of the
+    LayerNorm are required to allow it to do this.
+
+    So the idea is to introduce this large constant value as an explicit
+    parameter, that takes the role of the "eps" in LayerNorm, so the network
+    doesn't have to do this trick.  We make the "eps" learnable.
+
+    Args:
+       num_channels: the number of channels, e.g. 512.
+      channel_dim: the axis/dimension corresponding to the channel,
+        interprted as an offset from the input's ndim if negative.
+        shis is NOT the num_channels; it should typically be one of
+        {-2, -1, 0, 1, 2, 3}.
+       eps: the initial "epsilon" that we add as ballast in:
+             scale = ((input_vec**2).mean() + epsilon)**-0.5
+          Note: our epsilon is actually large, but we keep the name
+          to indicate the connection with conventional LayerNorm.
+       learn_eps: if true, we learn epsilon; if false, we keep it
+         at the initial value.
+    eps_min: float
+    eps_max: float
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int = -1,  # CAUTION: see documentation.
+        eps: float = 0.25,
+        learn_eps: bool = True,
+        eps_min: float = -3.0,
+        eps_max: float = 3.0,
+    ) -> None:
+        super(BasicNorm, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        if learn_eps:
+            self.eps = nn.Parameter(torch.tensor(eps).log().detach())
+        else:
+            self.register_buffer("eps", torch.tensor(eps).log().detach())
+        self.eps_min = eps_min
+        self.eps_max = eps_max
+
+    def forward(self, x: Tensor) -> Tensor:
+        assert x.shape[self.channel_dim] == self.num_channels
+        eps = self.eps
+        if self.training and random.random() < 0.25:
+            # with probability 0.25, in training mode, clamp eps between the min
+            # and max; this will encourage it to learn parameters within the
+            # allowed range by making parameters that are outside the allowed
+            # range noisy.
+
+            # gradients to allow the parameter to get back into the allowed region if it happens to exit it.
+            eps = eps.clamp(min=self.eps_min, max=self.eps_max)
+        scales = (
+            torch.mean(x ** 2, dim=self.channel_dim, keepdim=True) + eps.exp()
+        ) ** -0.5
+        return x * scales
+
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+    """
+    Behaves like a constructor of a modified version of nn.Linear
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Linear(*args, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def ScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Conv1d:
+    """
+    Behaves like a constructor of a modified version of nn.Conv1d
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+
+
+def TransposeScaledConv1d(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+    )
+
+
+def ScaledConv1dTranspose(
+    *args,
+    initial_scale: float = 1.0,
+    kernel_size: int = 3,
+    padding: str = "same",
+    **kwargs,
+) -> nn.Sequential:
+    """
+    Transpose -> ScaledConv1d
+    """
+    return nn.Sequential(
+        ScaledConv1d(
+            *args,
+            initial_scale=initial_scale,
+            kernel_size=kernel_size,
+            padding=padding,
+            **kwargs,
+        ),
+        Transpose(),
+    )
+
+
+def TransposeConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> Conv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def Conv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    ScaledConv1d -> Transpose
+    """
+    return nn.Sequential(
+        nn.Conv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class SRLinear(nn.Linear):
+    """https://arxiv.org/abs/2303.06296
+    Stabilizing Transformer Training by Preventing Attention Entropy Collapse
+    """
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        self.register_buffer(
+            "u", nn.functional.normalize(torch.randn(in_features), dim=0)
+        )
+        with torch.no_grad():
+            sigma = self.get_sigma()
+        self.register_buffer("spectral_norm", sigma)
+        self.sigma = nn.Parameter(torch.ones(1))
+
+    def get_sigma(self):
+        with torch.no_grad():
+            u = self.u
+            v = self.weight.mv(u)
+            v = nn.functional.normalize(v, dim=0)
+            u = self.weight.T.mv(v)
+            u = nn.functional.normalize(u, dim=0)
+            self.u.data.copy_(u)
+        return torch.einsum("c,cd,d->", v, self.weight, u)
+
+    def get_weight(self):
+        sigma = self.get_sigma()
+        if self.training:
+            self.spectral_norm.data.copy_(sigma)
+        weight = (self.sigma / sigma) * self.weight
+        return weight
+
+    def forward(self, x):
+        return nn.functional.linear(x, self.get_weight(), self.bias)
+
+
+class SRConv1d(SRLinear):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        kernel_size,
+        stride: int = 1,
+        padding: str = "same",
+        bias: bool = True,
+        **kwargs,
+    ):
+        in_features = in_features * kernel_size
+        super().__init__(in_features, out_features, bias=bias, **kwargs)
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.padding = padding
+
+    def forward(self, x):
+        in_features = self.in_features // self.kernel_size
+        weight = self.get_weight().view(
+            self.out_features, in_features, self.kernel_size
+        )
+        return nn.functional.conv1d(
+            x, weight, bias=self.bias, stride=self.stride, padding=self.padding
+        )
+
+
+def TransposeSRConv1d(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    Transpose -> SRConv1d
+    """
+    return nn.Sequential(
+        Transpose(),
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+    )
+
+
+def SRConv1dTranspose(
+    *args, kernel_size: int = 3, padding: str = "same", **kwargs
+) -> nn.Sequential:
+    """
+    SRConv1d -> Transpose
+    """
+    return nn.Sequential(
+        SRConv1d(*args, kernel_size=kernel_size, padding=padding, **kwargs),
+        Transpose(),
+    )
+
+
+class ActivationBalancer(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to encourage, for
+    each channel, that it is positive at least a proportion `threshold` of the
+    time.  It does this by multiplying negative derivative values by up to
+    (1+max_factor), and positive derivative values by up to (1-max_factor),
+    interpolated from 1 at the threshold to those extremal values when none
+    of the inputs are positive.
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           min_positive: the minimum, per channel, of the proportion of the time
+               that (x > 0), below which we start to modify the derivatives.
+           max_positive: the maximum, per channel, of the proportion of the time
+               that (x > 0), above which we start to modify the derivatives.
+           max_factor: the maximum factor by which we modify the derivatives for
+              either the sign constraint or the magnitude constraint;
+              e.g. with max_factor=0.02, the the derivatives would be multiplied by
+              values in the range [0.98..1.02].
+           sign_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_positive and max_positive
+              are violated.
+           scale_gain_factor: determines the 'gain' with which we increase the
+              change in gradient once the constraints on min_abs and max_abs
+              are violated.
+           min_abs:  the minimum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+           max_abs:  the maximum average-absolute-value difference from the mean
+               value per channel, which we allow, before we start to modify
+               the derivatives to prevent this.
+          min_prob: determines the minimum probability with which we modify the
+             gradients for the {min,max}_positive and {min,max}_abs constraints,
+             on each forward().  This is done randomly to prevent all layers
+             from doing it at the same time.  Early in training we may use
+             higher probabilities than this; it will decay to this value.
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        min_positive: float = 0.05,
+        max_positive: float = 0.95,
+        max_factor: float = 0.04,
+        sign_gain_factor: float = 0.01,
+        scale_gain_factor: float = 0.02,
+        min_abs: float = 0.2,
+        max_abs: float = 100.0,
+        min_prob: float = 0.1,
+    ):
+        super(ActivationBalancer, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.min_positive = min_positive
+        self.max_positive = max_positive
+        self.max_factor = max_factor
+        self.min_abs = min_abs
+        self.max_abs = max_abs
+        self.min_prob = min_prob
+        self.sign_gain_factor = sign_gain_factor
+        self.scale_gain_factor = scale_gain_factor
+
+        # count measures how many times the forward() function has been called.
+        # We occasionally sync this to a tensor called `count`, that exists to
+        # make sure it is synced to disk when we load and save the model.
+        self.cpu_count = 0
+        self.register_buffer("count", torch.tensor(0, dtype=torch.int64))
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or not x.requires_grad
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        count = self.cpu_count
+        self.cpu_count += 1
+
+        if random.random() < 0.01:
+            # Occasionally sync self.cpu_count with self.count.
+            # count affects the decay of 'prob'.  don't do this on every iter,
+            # because syncing with the GPU is slow.
+            self.cpu_count = max(self.cpu_count, self.count.item())
+            self.count.fill_(self.cpu_count)
+
+        # the prob of doing some work exponentially decreases from 0.5 till it hits
+        # a floor at min_prob (==0.1, by default)
+        prob = max(self.min_prob, 0.5 ** (1 + (count / 4000.0)))
+
+        if random.random() < prob:
+            sign_gain_factor = 0.5
+            if self.min_positive != 0.0 or self.max_positive != 1.0:
+                sign_factor = _compute_sign_factor(
+                    x,
+                    self.channel_dim,
+                    self.min_positive,
+                    self.max_positive,
+                    gain_factor=self.sign_gain_factor / prob,
+                    max_factor=self.max_factor,
+                )
+            else:
+                sign_factor = None
+
+            scale_factor = _compute_scale_factor(
+                x.detach(),
+                self.channel_dim,
+                min_abs=self.min_abs,
+                max_abs=self.max_abs,
+                gain_factor=self.scale_gain_factor / prob,
+                max_factor=self.max_factor,
+            )
+            return ActivationBalancerFunction.apply(
+                x,
+                scale_factor,
+                sign_factor,
+                self.channel_dim,
+            )
+        else:
+            return _no_op(x)
+
+
+def penalize_abs_values_gt(x: Tensor, limit: float, penalty: float) -> Tensor:
+    """
+    Returns x unmodified, but in backprop will put a penalty for the excess of
+    the absolute values of elements of x over the limit "limit".  E.g. if
+    limit == 10.0, then if x has any values over 10 it will get a penalty.
+
+    Caution: the value of this penalty will be affected by grad scaling used
+    in automatic mixed precision training.  For this reasons we use this,
+    it shouldn't really matter, or may even be helpful; we just use this
+    to disallow really implausible values of scores to be given to softmax.
+    """
+    x_sign = x.sign()
+    over_limit = (x.abs() - limit) > 0
+    # The following is a memory efficient way to penalize the absolute values of
+    # x that's over the limit.  (The memory efficiency comes when you think
+    # about which items torch needs to cache for the autograd, and which ones it
+    # can throw away).  The numerical value of aux_loss as computed here will
+    # actually be larger than it should be, by limit * over_limit.sum(), but it
+    # has the same derivative as the real aux_loss which is penalty * (x.abs() -
+    # limit).relu().
+    aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x)
+    # note: we don't do sum() here on aux)_loss, but it's as if we had done
+    # sum() due to how with_loss() works.
+    x = with_loss(x, aux_loss)
+    # you must use x for something, or this will be ineffective.
+    return x
+
+
+def _diag(x: Tensor):  # like .diag(), but works for tensors with 3 dims.
+    if x.ndim == 2:
+        return x.diag()
+    else:
+        (batch, dim, dim) = x.shape
+        x = x.reshape(batch, dim * dim)
+        x = x[:, :: dim + 1]
+        assert x.shape == (batch, dim)
+        return x
+
+
+def _whitening_metric(x: Tensor, num_groups: int):
+    """
+    Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of
+    of the centered feature covariance are the same within each group's covariance matrix
+    and also between groups.
+    Args:
+        x: a Tensor of shape (*, num_channels)
+     num_groups:  the number of groups of channels, a number >=1 that divides num_channels
+    Returns:
+        Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and
+    greater than 1.0 otherwise.
+    """
+    assert x.dtype != torch.float16
+    x = x.reshape(-1, x.shape[-1])
+    (num_frames, num_channels) = x.shape
+    assert num_channels % num_groups == 0
+    channels_per_group = num_channels // num_groups
+    x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1)
+    # x now has shape (num_groups, num_frames, channels_per_group)
+    # subtract the mean so we use the centered, not uncentered, covariance.
+    # My experience has been that when we "mess with the gradients" like this,
+    # it's better not do anything that tries to move the mean around, because
+    # that can easily cause instability.
+    x = x - x.mean(dim=1, keepdim=True)
+    # x_covar: (num_groups, channels_per_group, channels_per_group)
+    x_covar = torch.matmul(x.transpose(1, 2), x)
+    x_covar_mean_diag = _diag(x_covar).mean()
+    # the following expression is what we'd get if we took the matrix product
+    # of each covariance and measured the mean of its trace, i.e.
+    # the same as _diag(torch.matmul(x_covar, x_covar)).mean().
+    x_covarsq_mean_diag = (x_covar ** 2).sum() / (
+        num_groups * channels_per_group
+    )
+    # this metric will be >= 1.0; the larger it is, the less 'white' the data was.
+    metric = x_covarsq_mean_diag / (x_covar_mean_diag ** 2 + 1.0e-20)
+    return metric
+
+
+class WhiteningPenaltyFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: Tensor,
+        num_groups: int,
+        whitening_limit: float,
+        grad_scale: float,
+    ) -> Tensor:
+        ctx.save_for_backward(x)
+        ctx.num_groups = num_groups
+        ctx.whitening_limit = whitening_limit
+        ctx.grad_scale = grad_scale
+        return x
+
+    @staticmethod
+    def backward(ctx, x_grad: Tensor):
+        (x_orig,) = ctx.saved_tensors
+        with torch.enable_grad():
+            with torch.cuda.amp.autocast(enabled=False):
+                x_detached = x_orig.to(torch.float32).detach()
+                x_detached.requires_grad = True
+
+                metric = _whitening_metric(x_detached, ctx.num_groups)
+
+                if random.random() < 0.005 or __name__ == "__main__":
+                    logging.info(
+                        f"Whitening: num_groups={ctx.num_groups}, num_channels={x_orig.shape[-1]}, "
+                        f"metric={metric.item():.2f} vs. limit={ctx.whitening_limit}"
+                    )
+
+                (metric - ctx.whitening_limit).relu().backward()
+                penalty_grad = x_detached.grad
+                scale = ctx.grad_scale * (
+                    x_grad.to(torch.float32).norm()
+                    / (penalty_grad.norm() + 1.0e-20)
+                )
+                penalty_grad = penalty_grad * scale
+        return x_grad + penalty_grad.to(x_grad.dtype), None, None, None
+
+
+class Whiten(nn.Module):
+    def __init__(
+        self,
+        num_groups: int,
+        whitening_limit: float,
+        prob: Union[float, Tuple[float, float]],
+        grad_scale: float,
+    ):
+        """
+        Args:
+          num_groups: the number of groups to divide the channel dim into before
+            whitening.  We will attempt to make the feature covariance
+            within each group, after mean subtraction, as "white" as possible,
+            while having the same trace across all groups.
+         whitening_limit: a value greater than 1.0, that dictates how much
+           freedom we have to violate the constraints.  1.0 would mean perfectly
+           white, with exactly the same trace across groups; larger values
+           give more freedom.  E.g. 2.0.
+         prob: the probability with which we apply the gradient modification
+           (also affects the grad scale).  May be supplied as a float,
+           or as a pair (min_prob, max_prob)
+
+          grad_scale: determines the scale on the gradient term from this object,
+            relative to the rest of the gradient on the attention weights.
+            E.g. 0.02 (you may want to use smaller values than this if prob is large)
+        """
+        super(Whiten, self).__init__()
+        assert num_groups >= 1
+        assert whitening_limit >= 1
+        assert grad_scale >= 0
+        self.num_groups = num_groups
+        self.whitening_limit = whitening_limit
+        if isinstance(prob, float):
+            assert 0 < prob <= 1
+            self.prob = prob
+        else:
+            (self.min_prob, self.max_prob) = prob
+            assert 0 < self.min_prob < self.max_prob <= 1
+            self.prob = self.max_prob
+
+        self.grad_scale = grad_scale
+
+    def forward(self, x: Tensor) -> Tensor:
+        """
+        In the forward pass, this function just returns the input unmodified.
+        In the backward pass, it will modify the gradients to ensure that the
+        distribution in each group has close to (lambda times I) as the covariance
+        after mean subtraction, with the same lambda across groups.
+        For whitening_limit > 1, there will be more freedom to violate this
+        constraint.
+
+        Args:
+           x: the input of shape (*, num_channels)
+
+        Returns:
+            x, unmodified.   You should make sure
+        you use the returned value, or the graph will be freed
+        and nothing will happen in backprop.
+        """
+        if (
+            not x.requires_grad
+            or random.random() > self.prob
+            or self.grad_scale == 0
+        ):
+            return _no_op(x)
+        else:
+            if hasattr(self, "min_prob") and random.random() < 0.25:
+                # occasionally switch between min_prob and max_prob, based on whether
+                # we are above or below the threshold.
+                if (
+                    _whitening_metric(x.to(torch.float32), self.num_groups)
+                    > self.whitening_limit
+                ):
+                    # there would be a change to the grad.
+                    self.prob = self.max_prob
+                else:
+                    self.prob = self.min_prob
+
+            return WhiteningPenaltyFunction.apply(
+                x, self.num_groups, self.whitening_limit, self.grad_scale
+            )
+
+
+class WithLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: Tensor, y: Tensor):
+        ctx.y_shape = y.shape
+        return x
+
+    @staticmethod
+    def backward(ctx, ans_grad: Tensor):
+        return ans_grad, torch.ones(
+            ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device
+        )
+
+
+def with_loss(x, y):
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    # returns x but adds y.sum() to the loss function.
+    return WithLoss.apply(x, y)
+
+
+def _no_op(x: Tensor) -> Tensor:
+    if torch.jit.is_scripting() or torch.jit.is_tracing():
+        return x
+    else:
+        # a no-op function that will have a node in the autograd graph,
+        # to avoid certain bugs relating to backward hooks
+        return x.chunk(1, dim=-1)[0]
+
+
+class Identity(torch.nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return _no_op(x)
+
+
+class MaxEig(torch.nn.Module):
+    """
+    Modifies the backpropped derivatives of a function to try to discourage
+    that any given direction in activation space accounts for more than
+    a specified proportion of the covariance (e.g. 0.2).
+
+
+    Args:
+           num_channels: the number of channels
+           channel_dim: the dimension/axis corresponding to the channel, e.g.
+               -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative.
+           max_var_per_eig:  the maximum proportion of the variance of the
+               features/channels, after mean subtraction, that can come from
+               any given eigenvalue.
+           min_prob: the minimum probability with which we apply this during any invocation
+               of forward(), assuming last time we applied the constraint it was
+               not active; supplied for speed.
+           scale: determines the scale with which we modify the gradients, relative
+               to the existing / unmodified gradients
+    """
+
+    def __init__(
+        self,
+        num_channels: int,
+        channel_dim: int,
+        max_var_per_eig: float = 0.2,
+        min_prob: float = 0.01,
+        scale: float = 0.01,
+    ):
+        super(MaxEig, self).__init__()
+        self.num_channels = num_channels
+        self.channel_dim = channel_dim
+        self.scale = scale
+        assert max_var_per_eig == 0.0 or max_var_per_eig > 1.0 / num_channels
+        self.max_var_per_eig = max_var_per_eig
+
+        # we figure out the dominant direction using the power method: starting with
+        # a random vector, keep multiplying by the covariance and renormalizing.
+        with torch.no_grad():
+            # arbitrary.. would use randn() but want to leave the rest of the model's
+            # random parameters unchanged for comparison
+            direction = torch.arange(num_channels).to(torch.float)
+            direction = direction / direction.norm()
+            self.register_buffer("max_eig_direction", direction)
+
+        self.min_prob = min_prob
+        # cur_prob is the current probability we'll use to apply the ActivationBalancer.
+        # We'll regress this towards prob, each tiem we try to apply it and it is not
+        # active.
+        self.cur_prob = 1.0
+
+    def forward(self, x: Tensor) -> Tensor:
+        if (
+            torch.jit.is_scripting()
+            or self.max_var_per_eig <= 0
+            or random.random() > self.cur_prob
+            or torch.jit.is_tracing()
+        ):
+            return _no_op(x)
+
+        with torch.cuda.amp.autocast(enabled=False):
+            eps = 1.0e-20
+            orig_x = x
+            x = x.to(torch.float32)
+            with torch.no_grad():
+                x = x.transpose(self.channel_dim, -1).reshape(
+                    -1, self.num_channels
+                )
+                x = x - x.mean(dim=0)
+                new_direction, coeffs = self._find_direction_coeffs(
+                    x, self.max_eig_direction
+                )
+                x_var = (x ** 2).mean()
+                x_residual = x - coeffs * new_direction
+                x_residual_var = (x_residual ** 2).mean()
+
+                # `variance_proportion` is the proportion of the variance accounted for
+                # by the top eigen-direction.
+                variance_proportion = (x_var - x_residual_var) / (
+                    x_var + 1.0e-20
+                )
+
+                # ensure new direction is nonzero even if x == 0, by including `direction`.
+                self._set_direction(
+                    0.1 * self.max_eig_direction + new_direction
+                )
+
+            if random.random() < 0.01 or __name__ == "__main__":
+                logging.info(
+                    f"variance_proportion = {variance_proportion.item()}, shape={tuple(orig_x.shape)}, cur_prob={self.cur_prob}"
+                )
+
+            if variance_proportion >= self.max_var_per_eig:
+                # The constraint is active.  Note, we should quite rarely
+                # reach here, only near the beginning of training if we are
+                # starting to diverge, should this constraint be active.
+                cur_prob = self.cur_prob
+                self.cur_prob = (
+                    1.0  # next time, do the update with probability 1.0.
+                )
+                return MaxEigLimiterFunction.apply(
+                    orig_x, coeffs, new_direction, self.channel_dim, self.scale
+                )
+            else:
+                # let self.cur_prob exponentially approach self.min_prob, as
+                # long as the constraint is inactive.
+                self.cur_prob = 0.75 * self.cur_prob + 0.25 * self.min_prob
+                return orig_x
+
+    def _set_direction(self, direction: Tensor):
+        """
+        Sets self.max_eig_direction to a normalized version of `direction`
+        """
+        direction = direction.detach()
+        direction = direction / direction.norm()
+        direction_sum = direction.sum().item()
+        if direction_sum - direction_sum == 0:  # no inf/nan
+            self.max_eig_direction[:] = direction
+        else:
+            logging.info(
+                f"Warning: sum of direction in MaxEig is {direction_sum}, "
+                "num_channels={self.num_channels}, channel_dim={self.channel_dim}"
+            )
+
+    def _find_direction_coeffs(
+        self, x: Tensor, prev_direction: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+            Figure out (an approximation to) the proportion of the variance of a set of
+            feature vectors that can be attributed to the top eigen-direction.
+            Args:
+             x: a Tensor of shape (num_frames, num_channels), with num_frames > 1.
+          prev_direction:  a Tensor of shape (num_channels,), that is our previous estimate
+                   of the top eigen-direction, or a random direction if this is the first
+                   iteration.  Does not have to be normalized, but should be nonzero.
+
+        Returns: (cur_direction, coeffs), where:
+             cur_direction: a Tensor of shape (num_channels,) that is the current
+                estimate of the top eigen-direction.
+             coeffs: a Tensor of shape (num_frames, 1) that minimizes, or
+                approximately minimizes, (x - coeffs * cur_direction).norm()
+        """
+        (num_frames, num_channels) = x.shape
+        assert num_channels > 1 and num_frames > 1
+        assert prev_direction.shape == (num_channels,)
+        # `coeffs` are the coefficients of `prev_direction` in x.
+        # actually represent the coeffs up to a constant positive factor.
+        coeffs = (x * prev_direction).sum(dim=1, keepdim=True) + 1.0e-10
+        cur_direction = (x * coeffs).sum(dim=0) / (
+            (coeffs ** 2).sum() + 1.0e-20
+        )
+        return cur_direction, coeffs
+
+
+class DoubleSwishFunction(torch.autograd.Function):
+    """
+      double_swish(x) = x * torch.sigmoid(x-1)
+    This is a definition, originally motivated by its close numerical
+    similarity to swish(swish(x)), where swish(x) =  x * sigmoid(x).
+
+    Memory-efficient derivative computation:
+     double_swish(x) = x * s, where s(x) = torch.sigmoid(x-1)
+     double_swish'(x) = d/dx double_swish(x) =  x * s'(x) + x' * s(x) = x * s'(x) + s(x).
+     Now, s'(x) = s(x) * (1-s(x)).
+     double_swish'(x) =  x * s'(x) + s(x).
+                      =  x * s(x) * (1-s(x)) + s(x).
+                     = double_swish(x) * (1-s(x)) + s(x)
+     ... so we just need to remember s(x) but not x itself.
+    """
+
+    @staticmethod
+    def forward(ctx, x: Tensor) -> Tensor:
+        requires_grad = x.requires_grad
+        x_dtype = x.dtype
+        if x.dtype == torch.float16:
+            x = x.to(torch.float32)
+
+        s = torch.sigmoid(x - 1.0)
+        y = x * s
+
+        if requires_grad:
+            deriv = y * (1 - s) + s
+            # notes on derivative of x * sigmoid(x - 1):
+            # https://www.wolframalpha.com/input?i=d%2Fdx+%28x+*+sigmoid%28x-1%29%29
+            # min \simeq -0.043638.  Take floor as -0.043637 so it's a lower bund
+            # max \simeq 1.1990.   Take ceil to be 1.2 so it's an upper bound.
+            # the combination of "+ torch.rand_like(deriv)" and casting to torch.uint8 (which
+            # floors), should be expectation-preserving.
+            floor = -0.043637
+            ceil = 1.2
+            d_scaled = (deriv - floor) * (
+                255.0 / (ceil - floor)
+            ) + torch.rand_like(deriv)
+            if __name__ == "__main__":
+                # for self-testing only.
+                assert d_scaled.min() >= 0.0
+                assert d_scaled.max() < 256.0
+            d_int = d_scaled.to(torch.uint8)
+            ctx.save_for_backward(d_int)
+        if x.dtype == torch.float16 or torch.is_autocast_enabled():
+            y = y.to(torch.float16)
+        return y
+
+    @staticmethod
+    def backward(ctx, y_grad: Tensor) -> Tensor:
+        (d,) = ctx.saved_tensors
+        # the same constants as used in forward pass.
+        floor = -0.043637
+        ceil = 1.2
+        d = d * ((ceil - floor) / 255.0) + floor
+        return y_grad * d
+
+
+class DoubleSwish(torch.nn.Module):
+    def forward(self, x: Tensor) -> Tensor:
+        """Return double-swish activation function which is an approximation to Swish(Swish(x)),
+        that we approximate closely with x * sigmoid(x-1).
+        """
+        if torch.jit.is_scripting() or torch.jit.is_tracing():
+            return x * torch.sigmoid(x - 1.0)
+        return DoubleSwishFunction.apply(x)
+
+
+def BalancedDoubleSwish(
+    d_model, channel_dim=-1, max_abs=10.0, min_prob=0.25
+) -> nn.Sequential:
+    """
+    ActivationBalancer -> DoubleSwish
+    """
+    balancer = ActivationBalancer(
+        d_model, channel_dim=channel_dim, max_abs=max_abs, min_prob=min_prob
+    )
+    return nn.Sequential(
+        balancer,
+        DoubleSwish(),
+    )
+
+
+def _test_max_eig():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = MaxEig(
+            num_channels, 1, 0.5, scale=0.1  # channel_dim  # max_var_per_eig
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad, atol=1.0e-02)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_whiten():
+    for proportion in [0.1, 0.5, 10.0]:
+        logging.info(f"_test_whiten(): proportion = {proportion}")
+        x = torch.randn(100, 128)
+        direction = torch.randn(128)
+        coeffs = torch.randn(100, 1)
+        x += proportion * direction * coeffs
+
+        x.requires_grad = True
+
+        num_channels = 128
+        m = Whiten(
+            1, 5.0, prob=1.0, grad_scale=0.1  # num_groups  # whitening_limit,
+        )  # grad_scale
+
+        for _ in range(4):
+            y = m(x)
+
+        y_grad = torch.randn_like(x)
+        y.backward(gradient=y_grad)
+
+        if proportion < 0.2:
+            assert torch.allclose(x.grad, y_grad)
+        elif proportion > 1.0:
+            assert not torch.allclose(x.grad, y_grad)
+
+
+def _test_activation_balancer_sign():
+    probs = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = 1.0 * (
+        (2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        probs.numel(),
+        channel_dim=0,
+        min_positive=0.05,
+        max_positive=0.95,
+        max_factor=0.2,
+        min_abs=0.0,
+    )
+
+    y_grad = torch.sign(torch.randn(probs.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_sign: x = ", x)
+    print("_test_activation_balancer_sign: y grad = ", y_grad)
+    print("_test_activation_balancer_sign: x grad = ", x.grad)
+
+
+def _test_activation_balancer_magnitude():
+    magnitudes = torch.arange(0, 1, 0.01)
+    N = 1000
+    x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(
+        -1
+    )
+    x = x.detach()
+    x.requires_grad = True
+    m = ActivationBalancer(
+        magnitudes.numel(),
+        channel_dim=0,
+        min_positive=0.0,
+        max_positive=1.0,
+        max_factor=0.2,
+        min_abs=0.2,
+        max_abs=0.8,
+        min_prob=1.0,
+    )
+
+    y_grad = torch.sign(torch.randn(magnitudes.numel(), N))
+
+    y = m(x)
+    y.backward(gradient=y_grad)
+    print("_test_activation_balancer_magnitude: x = ", x)
+    print("_test_activation_balancer_magnitude: y grad = ", y_grad)
+    print("_test_activation_balancer_magnitude: x grad = ", x.grad)
+
+
+def _test_basic_norm():
+    num_channels = 128
+    m = BasicNorm(num_channels=num_channels, channel_dim=1)
+
+    x = torch.randn(500, num_channels)
+
+    y = m(x)
+
+    assert y.shape == x.shape
+    x_rms = (x ** 2).mean().sqrt()
+    y_rms = (y ** 2).mean().sqrt()
+    print("x rms = ", x_rms)
+    print("y rms = ", y_rms)
+    assert y_rms < x_rms
+    assert y_rms > 0.5 * x_rms
+
+
+def _test_double_swish_deriv():
+    x = torch.randn(10, 12, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    m = DoubleSwish()
+
+    tol = (1.2 - (-0.043637)) / 255.0
+    torch.autograd.gradcheck(m, x, atol=tol)
+
+    # for self-test.
+    x = torch.randn(1000, 1000, dtype=torch.double) * 3.0
+    x.requires_grad = True
+    y = m(x)
+
+
+def _test_softmax():
+    a = torch.randn(2, 10, dtype=torch.float64)
+    b = a.clone()
+    a.requires_grad = True
+    b.requires_grad = True
+    a.softmax(dim=1)[:, 0].sum().backward()
+    print("a grad = ", a.grad)
+    softmax(b, dim=1)[:, 0].sum().backward()
+    print("b grad = ", b.grad)
+    assert torch.allclose(a.grad, b.grad)
+
+
+if __name__ == "__main__":
+    logging.getLogger().setLevel(logging.INFO)
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    _test_softmax()
+    _test_whiten()
+    _test_max_eig()
+    _test_activation_balancer_sign()
+    _test_activation_balancer_magnitude()
+    _test_basic_norm()
+    _test_double_swish_deriv()

models/modules/transformer.py

@@ -0,0 +1,1089 @@
+import copy, logging
+import numbers
+from functools import partial
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .activation import MultiheadAttention
+from .scaling import ActivationBalancer, BalancedDoubleSwish
+from .scaling import BasicNorm as _BasicNorm
+
+_shape_t = Union[int, List[int], torch.Size]
+
+
+class LayerNorm(nn.Module):
+    __constants__ = ["normalized_shape", "eps", "elementwise_affine"]
+    normalized_shape: Tuple[int, ...]
+    eps: float
+    elementwise_affine: bool
+
+    def __init__(
+        self,
+        normalized_shape: _shape_t,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(LayerNorm, self).__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            # mypy error: incompatible types in assignment
+            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
+        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+            self.bias = nn.Parameter(
+                torch.empty(self.normalized_shape, **factory_kwargs)
+            )
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        if self.elementwise_affine:
+            nn.init.ones_(self.weight)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                F.layer_norm(
+                    input,
+                    self.normalized_shape,
+                    self.weight,
+                    self.bias,
+                    self.eps,
+                ),
+                embedding,
+            )
+
+        assert embedding is None
+        return F.layer_norm(
+            input, self.normalized_shape, self.weight, self.bias, self.eps
+        )
+
+    def extra_repr(self) -> str:
+        return (
+            "{normalized_shape}, eps={eps}, "
+            "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+        )
+
+
+class AdaptiveLayerNorm(nn.Module):
+    r"""Adaptive Layer Normalization"""
+
+    def __init__(self, d_model, norm) -> None:
+        super(AdaptiveLayerNorm, self).__init__()
+        self.project_layer = nn.Linear(d_model, 2 * d_model)
+        self.norm = norm
+        self.d_model = d_model
+        self.eps = self.norm.eps
+
+    def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            weight, bias = torch.split(
+                self.project_layer(embedding),
+                split_size_or_sections=self.d_model,
+                dim=-1,
+            )
+            return (weight * self.norm(input) + bias, embedding)
+
+        weight, bias = torch.split(
+            self.project_layer(embedding),
+            split_size_or_sections=self.d_model,
+            dim=-1,
+        )
+        return weight * self.norm(input) + bias
+
+
+class BasicNorm(_BasicNorm):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BasicNorm, self).__init__(d_model, eps=eps)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return (
+                super(BasicNorm, self).forward(input),
+                embedding,
+            )
+
+        assert embedding is None
+        return super(BasicNorm, self).forward(input)
+
+
+class BalancedBasicNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ):
+        super(BalancedBasicNorm, self).__init__()
+        self.balancer = ActivationBalancer(
+            d_model,
+            channel_dim=-1,
+            min_positive=0.45,
+            max_positive=0.55,
+            max_abs=6.0,
+        )
+        self.norm = BasicNorm(d_model, eps, device=device, dtype=dtype)
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            input, embedding = input
+            return self.norm((self.balancer(input), embedding))
+
+        assert embedding is None
+        return self.norm(self.balancer(input))
+
+
+class IdentityNorm(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        eps: float = 1e-5,
+        device=None,
+        dtype=None,
+    ) -> None:
+        super(IdentityNorm, self).__init__()
+
+    def forward(self, input: Tensor, embedding: Any = None) -> Tensor:
+        if isinstance(input, tuple):
+            return input
+
+        assert embedding is None
+        return input
+
+
+class TransformerEncoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerEncoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            activation = BalancedDoubleSwish(d_model)
+
+        # # We can't test self.activation in forward() in TorchScript,
+        # # so stash some information about it instead.
+        # if activation is F.relu or isinstance(activation, torch.nn.ReLU):
+        #     self.activation_relu_or_gelu = 1
+        # elif activation is F.gelu or isinstance(activation, torch.nn.GELU):
+        #     self.activation_relu_or_gelu = 2
+        # else:
+        #     self.activation_relu_or_gelu = 0
+        self.activation = activation
+
+        norm1 = layer_norm_cls(d_model, eps=layer_norm_eps, **factory_kwargs)
+        if layer_norm_cls == IdentityNorm:
+            norm2 = BalancedBasicNorm(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+        else:
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+        if adaptive_layer_norm:
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+        else:
+            self.norm1 = norm1
+            self.norm2 = norm2
+
+    def __setstate__(self, state):
+        super(TransformerEncoderLayer, self).__setstate__(state)
+        if not hasattr(self, "activation"):
+            self.activation = F.relu
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        need_weights: Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            src: the sequence to the encoder layer (required).
+            src_mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if isinstance(src, dict):
+            sinu = src["sinu"]
+            pm_sinu = src["pm_sinu"]
+            src = src["input"]
+        else:
+            sinu = None
+            pm_sinu = None
+        x, stage_embedding = src, None
+        is_src_tuple = False
+        if isinstance(src, tuple):
+            x, stage_embedding = src
+            is_src_tuple = True
+        
+        if src_key_padding_mask is not None:
+            _skpm_dtype = src_key_padding_mask.dtype
+            if _skpm_dtype != torch.bool and not torch.is_floating_point(
+                src_key_padding_mask
+            ):
+                raise AssertionError(
+                    "only bool and floating types of key_padding_mask are supported"
+                )
+        if need_weights:
+            raise NotImplementedError
+            if self.norm_first:
+                out, attn = self._sa_block_attn(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask,
+                    past, sinu = sinu
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out, attn = self._sa_block_attn(x, src_mask, src_key_padding_mask, past, sinu = sinu)
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding,
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            assert not is_src_tuple
+                # return (x, stage_embedding)
+            return (x, attn)
+        else:
+            if self.norm_first:
+                out = self._sa_block(
+                    self.norm1(x, stage_embedding),
+                    src_mask,
+                    src_key_padding_mask, past, sinu = sinu, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['q']
+                )
+                out, present = out # present is the kvcache of the present timestep
+                x = x + out
+                x = x + self._ff_block(self.norm2(x, stage_embedding))
+            else:
+                out = self._sa_block(x, src_mask, src_key_padding_mask, sinu = sinu, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['q'])
+                out, present = out # present is the kvcache of the present timestep
+                x = self.norm1(
+                    x + out,
+                    stage_embedding, past
+                )
+                x = self.norm2(x + self._ff_block(x), stage_embedding)
+            
+            if is_src_tuple:
+                x = (x, stage_embedding)
+            if present != None:
+                x = [x, present]
+            return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+        sinu = None,
+        q_sinu = None,
+        k_sinu = None
+    ) -> Tensor:
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            past=past,
+            sinu = sinu,
+            q_sinu = q_sinu,
+            k_sinu = k_sinu
+        )
+        x, present = x
+        return self.dropout1(x), present
+
+    # self-attention block, also return attention weights
+    def _sa_block_attn(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        x, attn = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=True,
+            past=past
+        )
+        x, present = x
+        return (self.dropout1(x), present), attn
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout2(x)
+
+def pre_compute_sinusoidal(dim, base, max_len = 10000): # 4000 max length equivalent of mimi code is 320s, as mimi is 12.5hz
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+    position_ids_expanded = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) # [x_len_max, 1]
+    inv_freq_expanded = inv_freq.unsqueeze(0).float() # [1, d//2]
+    freqs = position_ids_expanded @ inv_freq_expanded # [x_len_max, d//2]
+    freqs = torch.cat((freqs, freqs), dim=-1).unsqueeze(0) # [1, x_len_max, d]
+    return {"sin": freqs.sin(), "cos": freqs.cos()}
+
+def pre_compute_freqs(dim, base, max_len = 10000): # 4000 max length equivalent of mimi code is 320s, as mimi is 12.5hz
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+    position_ids_expanded = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1) # [x_len_max, 1]
+    inv_freq_expanded = inv_freq.unsqueeze(0).float() # [1, d//2]
+    freqs = position_ids_expanded @ inv_freq_expanded # [x_len_max, d//2]
+    freqs = torch.cat((freqs, freqs), dim=-1).unsqueeze(0) # [1, x_len_max, d]
+    return freqs
+
+class TransformerEncoder(nn.Module):
+    r"""TransformerEncoder is a stack of N encoder layers. Users can build the
+    BERT(https://arxiv.org/abs/1810.04805) model with corresponding parameters.
+
+    Args:
+        encoder_layer: an instance of the TransformerEncoderLayer() class (required).
+        num_layers: the number of sub-encoder-layers in the encoder (required).
+        norm: the layer normalization component (optional).
+        enable_nested_tensor: if True, input will automatically convert to nested tensor
+            (and convert back on output). This will improve the overall performance of
+            TransformerEncoder when padding rate is high. Default: ``True`` (enabled).
+
+    Examples::
+        >>> encoder_layer = TransformerEncoderLayer(d_model=512, nhead=8)
+        >>> transformer_encoder = TransformerEncoder(encoder_layer, num_layers=6)
+        >>> src = torch.rand(10, 32, 512)
+        >>> out = transformer_encoder(src)
+    """
+    __constants__ = ["norm"]
+
+    def __init__(self, encoder_layer, num_layers, norm=None, rope_base=None, d_model=None, nhead=None, args=None):
+        super(TransformerEncoder, self).__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        if args != None:
+            self.progress_no_multiple = args.progress_no_multiple
+            self.progress_scale = args.progress_scale
+        else:
+            self.progress_no_multiple = False
+            self.progress_scale = 1
+
+        if rope_base is not None:
+            if self.progress_no_multiple:
+                self.pm_freqs = pre_compute_freqs(d_model//nhead, rope_base)
+                self.sinu = None
+            else:
+                self.sinu = pre_compute_sinusoidal(d_model/nhead, rope_base)
+                self.pm_freqs = None
+            # logging.info(f"get precomputed sinusoidal for {rope_base=}: {self.sinu=}")
+        else:
+            self.sinu = None
+            self.pm_freqs = None
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        return_layer_states: bool = False,
+        need_weights:Optional[bool] = False,
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the input through the encoder layers in turn.
+
+        Args:
+            src: the sequence to the encoder (required).
+            mask: the mask for the src sequence (optional).
+            src_key_padding_mask: the mask for the src keys per batch (optional).
+            return_layer_states: return layers' state (optional).
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if return_layer_states:
+            raise NotImplementedError
+            assert not need_weights
+            layer_states = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    past=past
+                )
+                layer_states.append(output[0])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_states, output
+        if need_weights:
+            raise NotImplementedError
+            assert not return_layer_states
+            layer_attn = []  # layers' output
+            output = src
+            for mod in self.layers:
+                output = mod(
+                    output,
+                    src_mask=mask,
+                    src_key_padding_mask=src_key_padding_mask,
+                    need_weights=True,
+                    past=past
+                )
+                layer_attn.append(output[1])
+
+            if self.norm is not None:
+                output = self.norm(output)
+
+            return layer_attn, output
+        
+        output = src
+        all_present = []
+        if self.sinu is not None:
+            # use rope
+            assert self.pm_freqs is None
+            for k, v in self.sinu.items():
+                self.sinu[k] = v.to(output.device)
+        if self.pm_freqs is not None:
+            assert self.sinu is None
+            self.pm_freqs = self.pm_freqs.to(output.device)
+            if src_key_padding_mask != None:
+                query_lens = (~src_key_padding_mask).int().sum(-1).to(output.device)
+            else:
+                query_lens = torch.tensor([output.shape[1]]*output.shape[0]).to(output.device)
+            assert query_lens.ndim==1, query_lens
+            q_lens_expanded = query_lens.unsqueeze(-1).unsqueeze(-1) # [B, 1, 1]
+            query_ids_multiple = q_lens_expanded / (q_lens_expanded - 1)
+            q_emb = self.pm_freqs * query_ids_multiple # [B, q_len_max, d]
+            q_emb = q_emb / q_lens_expanded * self.progress_scale
+            q_cos = q_emb.cos().unsqueeze(1) # [B, 1, q_len_max, d] # 1 is for nhead
+            q_sin = q_emb.sin().unsqueeze(1)
+            self.pm_sinu = {"q": {"cos": q_cos, "sin": q_sin}}
+        else:
+            self.pm_sinu = {"q": None}
+
+        output = {"input": output, "sinu": self.sinu, "pm_sinu": self.pm_sinu}
+        for n_layer, mod in enumerate(self.layers):
+            output = mod(
+                output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, past=None if past is None else past[n_layer]
+            )
+            if isinstance(output, list):
+                output, present = output
+                all_present.append(present)
+            if self.sinu is not None or self.pm_sinu is not None:
+                output = {"input": output, "sinu": self.sinu, "pm_sinu": self.pm_sinu}
+        if self.sinu is not None or self.pm_sinu is not None:
+            output = output["input"]
+        if self.norm is not None:
+            output = self.norm(output)
+        if all_present != []:
+            all_present = torch.stack(all_present, dim=0) # (num_layers, 2, batch_size, num_heads, seq_len, head_dim)
+            output = [output, all_present]
+        return output
+
+
+class TransformerDecoderLayer(nn.Module):
+    __constants__ = ["batch_first", "norm_first"]
+
+    def __init__(
+        self,
+        d_model: int,
+        nhead: int,
+        dim_feedforward: int = 2048,
+        dropout: float = 0.1,
+        activation: Union[str, Callable[[Tensor], Tensor]] = F.relu,
+        linear1_self_attention_cls: nn.Module = nn.Linear,
+        linear2_self_attention_cls: nn.Module = nn.Linear,
+        linear1_feedforward_cls: nn.Module = nn.Linear,
+        linear2_feedforward_cls: nn.Module = nn.Linear,
+        batch_first: bool = False,
+        norm_first: bool = False,
+        device=None,
+        dtype=None,
+        layer_norm_cls: nn.Module = LayerNorm,
+        layer_norm_eps: float = 1e-5,
+        adaptive_layer_norm=False,
+    ) -> None:
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super(TransformerDecoderLayer, self).__init__()
+        self.self_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        self.multihead_attn = MultiheadAttention(
+            d_model,
+            nhead,
+            dropout=dropout,
+            batch_first=batch_first,
+            linear1_cls=linear1_self_attention_cls,
+            linear2_cls=linear2_self_attention_cls,
+            **factory_kwargs,
+        )
+        # Implementation of Feedforward model
+        self.linear1 = linear1_feedforward_cls(
+            d_model, dim_feedforward, **factory_kwargs
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = linear2_feedforward_cls(
+            dim_feedforward, d_model, **factory_kwargs
+        )
+
+        self.norm_first = norm_first
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        # Legacy string support for activation function.
+        if isinstance(activation, str):
+            self.activation = _get_activation_fn(activation)
+        elif isinstance(activation, partial):
+            self.activation = activation(d_model)
+        elif activation == BalancedDoubleSwish:
+            self.activation = BalancedDoubleSwish(d_model)
+        else:
+            self.activation = activation
+
+        if adaptive_layer_norm:
+            norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            norm3 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+
+            self.norm1 = AdaptiveLayerNorm(d_model, norm1)
+            self.norm2 = AdaptiveLayerNorm(d_model, norm2)
+            self.norm3 = AdaptiveLayerNorm(d_model, norm3)
+        else:
+            self.norm1 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            self.norm2 = layer_norm_cls(
+                d_model, eps=layer_norm_eps, **factory_kwargs
+            )
+            if layer_norm_cls == IdentityNorm:
+                self.norm3 = BalancedBasicNorm(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+            else:
+                self.norm3 = layer_norm_cls(
+                    d_model, eps=layer_norm_eps, **factory_kwargs
+                )
+
+    def forward(
+        self,
+        tgt: Tensor,
+        memory: Tensor,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        tgt_is_causal: Optional[bool] = False, # for compatibility with the nn.TransformerDecoder, not used
+        memory_is_causal: Optional[bool] = False, # for compatibility with the nn.TransformerDecoder, not used
+        past: Optional[Tensor] = None,
+    ) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer.
+
+        Args:
+            tgt: the sequence to the decoder layer (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+            past: the previous kvcache of the decoder (optional). shape: (2, batch_size, num_heads, seq_len, head_dim)
+
+        Shape:
+            see the docs in Transformer class.
+        """
+        if isinstance(tgt, dict):
+            pm_sinu = tgt["pm_sinu"]
+            sinu = tgt["sinu"]
+            args = tgt["args"]
+            tgt = tgt["input"]
+        else:
+            pm_sinu = None
+            sinu = None
+            args = None
+        tgt_is_tuple = False
+        if isinstance(tgt, tuple):
+            x, stage_embedding = tgt
+            tgt_is_tuple = True
+        else:
+            x, stage_embedding = tgt, None
+        # logging.info(f"{tgt_key_padding_mask=}, {memory_key_padding_mask=}")
+        # logging.info(f"{tgt_key_padding_mask.shape=}, {memory_key_padding_mask.shape=}")
+        # logging.info(f"{query_lens=}, {key_lens=}")
+
+        # past stores the kvcache for self-attention, and it can also be used to infer q_offset
+        if past is not None and past.ndim > 2:
+            q_offset = past[0].shape[-2] # past is (2, batch_size, num_heads, seq_len, head_dim), 2 contains [k, v], these are for self-attn, therefore also reflect the length of q
+        else:
+            q_offset = 0
+
+
+        if self.norm_first:
+            temp = self._sa_block(
+                self.norm1(x, stage_embedding), tgt_mask, tgt_key_padding_mask, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['q'], sinu=sinu, args = args, past=past, q_offset=q_offset
+            )
+            present = temp[1]
+            x = x + temp[0]
+            cross_out = self._mha_block(
+                self.norm2(x, stage_embedding),
+                memory,
+                memory_mask,
+                memory_key_padding_mask, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['k'], sinu=sinu, args = args, q_offset=q_offset
+            )
+            if isinstance(cross_out, dict):
+                attention_weights = cross_out["attention_weights"]
+                cross_out = cross_out["x"]
+            else:
+                attention_weights = None
+            x = x + cross_out
+            x = x + self._ff_block(self.norm3(x, stage_embedding))
+        else:
+            temp = self._sa_block(x, tgt_mask, tgt_key_padding_mask, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['q'], sinu=sinu, args = args, past=past, q_offset=q_offset)
+            present = temp[1]
+            x = self.norm1(
+                x + temp[0],
+                stage_embedding,
+            )
+            cross_out = self._mha_block(
+                    x, memory, memory_mask, memory_key_padding_mask, q_sinu=pm_sinu['q'], k_sinu=pm_sinu['k'], sinu=sinu, args=args, q_offset=q_offset
+                )
+            if isinstance(cross_out, dict):
+                attention_weights = cross_out["attention_weights"]
+                cross_out = cross_out["x"]
+            else:
+                attention_weights = None 
+            x = self.norm2(
+                x
+                + cross_out,
+                stage_embedding,
+            )
+            x = self.norm3(x + self._ff_block(x), stage_embedding)
+        
+        if attention_weights is not None:
+            x = {"x": x, "attention_weights": attention_weights}
+        if tgt_is_tuple:
+            x = (x, stage_embedding)
+        if present != None:
+            x = [x, present]
+        return x
+
+    # self-attention block
+    def _sa_block(
+        self,
+        x: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        q_sinu=None,
+        k_sinu=None,
+        sinu = None, 
+        args = None,
+        past = None,
+        q_offset = 0
+    ) -> Tensor:
+        # if past is not None and past.ndim > 2:
+        #     print(f"self-attn, k len: {past[0].shape[-2] + x.shape[-2]}, q len: {x.shape[-2]} q_offset: {q_offset}")
+        # else:
+        #     print(f"self-attn, k len: {x.shape[-2]}, q len: {x.shape[-2]} q_offset: {q_offset}")
+        x = self.self_attn(
+            x,
+            x,
+            x,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            q_sinu = q_sinu,
+            k_sinu = k_sinu,
+            sinu = sinu,
+            past = past,
+            q_offset = q_offset
+        )
+        x, present = x
+        return self.dropout1(x), present
+
+    # multihead attention block
+    def _mha_block(
+        self,
+        x: Tensor,
+        mem: Tensor,
+        attn_mask: Optional[Tensor],
+        key_padding_mask: Optional[Tensor],
+        q_sinu = None,
+        k_sinu = None,
+        sinu = None, 
+        args = None,
+        q_offset = 0
+    ) -> Tensor:
+        # print(f"cross-attn, k len: {mem.shape[-2]}, q len: {x.shape[-2]} q_offset: {q_offset}")
+        x = self.multihead_attn(
+            x,
+            mem,
+            mem,
+            attn_mask=attn_mask,
+            key_padding_mask=key_padding_mask,
+            need_weights=False,
+            q_sinu = q_sinu,
+            k_sinu = k_sinu,
+            sinu = sinu,
+            args = args,
+            q_offset = q_offset
+        )
+        if len(x) == 2 and isinstance(x[0], dict) and "attention_weights" in x[0]:
+            x, present = x
+            attention_weights = x['attention_weights']
+            x = x['attn_output']
+            return {"x": self.dropout2(x), "attention_weights": attention_weights}
+        elif len(x) == 2:
+            x = x[0]
+        return self.dropout2(x)
+
+    # feed forward block
+    def _ff_block(self, x: Tensor) -> Tensor:
+        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        return self.dropout3(x)
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation: str) -> Callable[[Tensor], Tensor]:
+    if activation == "relu":
+        return F.relu
+    elif activation == "gelu":
+        return F.gelu
+
+    raise RuntimeError(
+        "activation should be relu/gelu, not {}".format(activation)
+    )
+def _generate_square_subsequent_mask(
+        sz: int,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    r"""Generate a square causal mask for the sequence.
+
+    The masked positions are filled with float('-inf'). Unmasked positions are filled with float(0.0).
+    """
+    if device is None:
+        device = torch.device('cpu')
+    if dtype is None:
+        dtype = torch.float32
+    return torch.triu(
+        torch.full((sz, sz), float('-inf'), dtype=dtype, device=device),
+        diagonal=1,
+    )
+def _get_seq_len(
+        src: Tensor,
+        batch_first: bool
+) -> Optional[int]:
+
+    if src.is_nested:
+        return None
+    else:
+        src_size = src.size()
+        if len(src_size) == 2:
+            # unbatched: S, E
+            return src_size[0]
+        else:
+            # batched: B, S, E if batch_first else S, B, E
+            seq_len_pos = 1 if batch_first else 0
+            return src_size[seq_len_pos]
+
+def _detect_is_causal_mask(
+        mask: Optional[Tensor],
+        is_causal: Optional[bool] = None,
+        size: Optional[int] = None,
+) -> bool:
+    """Return whether the given attention mask is causal.
+
+    Warning:
+    If ``is_causal`` is not ``None``, its value will be returned as is.  If a
+    user supplies an incorrect ``is_causal`` hint,
+
+    ``is_causal=False`` when the mask is in fact a causal attention.mask
+       may lead to reduced performance relative to what would be achievable
+       with ``is_causal=True``;
+    ``is_causal=True`` when the mask is in fact not a causal attention.mask
+       may lead to incorrect and unpredictable execution - in some scenarios,
+       a causal mask may be applied based on the hint, in other execution
+       scenarios the specified mask may be used.  The choice may not appear
+       to be deterministic, in that a number of factors like alignment,
+       hardware SKU, etc influence the decision whether to use a mask or
+       rely on the hint.
+    ``size`` if not None, check whether the mask is a causal mask of the provided size
+       Otherwise, checks for any causal mask.
+    """
+    # Prevent type refinement
+    make_causal = (is_causal is True)
+
+    if is_causal is None and mask is not None:
+        sz = size if size is not None else mask.size(-2)
+        causal_comparison = _generate_square_subsequent_mask(
+            sz, device=mask.device, dtype=mask.dtype)
+
+        # Do not use `torch.equal` so we handle batched masks by
+        # broadcasting the comparison.
+        if mask.size() == causal_comparison.size():
+            make_causal = bool((mask == causal_comparison).all())
+        else:
+            make_causal = False
+
+    return make_causal
+
+class TransformerDecoder(nn.Module):
+    r"""TransformerDecoder is a stack of N decoder layers.
+
+    Args:
+        decoder_layer: an instance of the TransformerDecoderLayer() class (required).
+        num_layers: the number of sub-decoder-layers in the decoder (required).
+        norm: the layer normalization component (optional).
+
+    Examples::
+        >>> decoder_layer = nn.TransformerDecoderLayer(d_model=512, nhead=8)
+        >>> transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=6)
+        >>> memory = torch.rand(10, 32, 512)
+        >>> tgt = torch.rand(20, 32, 512)
+        >>> out = transformer_decoder(tgt, memory)
+    """
+
+    __constants__ = ['norm']
+
+    def __init__(
+        self,
+        decoder_layer: "TransformerDecoderLayer",
+        num_layers: int,
+        norm: Optional[nn.Module] = None,
+        rope_base=None, 
+        d_model=None, 
+        nhead=None,
+        args=None
+    ) -> None:
+        super().__init__()
+        torch._C._log_api_usage_once(f"torch.nn.modules.{self.__class__.__name__}")
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.args = args
+        if getattr(self.args, 'decoder_regular_rope', False):
+            self.sinu = pre_compute_sinusoidal(d_model/nhead, rope_base)
+            self.pm_freqs = None
+        else:
+            self.sinu = None
+            if rope_base is not None:
+                self.pm_freqs = pre_compute_freqs(d_model/nhead, rope_base)
+                # logging.info(f"get precomputed freqs for {rope_base=}: {self.freqs=}")
+            else:
+                self.pm_freqs = None
+        self.progress_scale = getattr(self.args, "progress_scale", 1.0)
+
+    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None, tgt_is_causal: Optional[bool] = None,
+                memory_is_causal: bool = False, query_lens: Optional[Tensor] = None, key_lens: Optional[Tensor] = None, past: Optional[Tensor] = None) -> Tensor:
+        r"""Pass the inputs (and mask) through the decoder layer in turn.
+
+        Args:
+            tgt: the sequence to the decoder (required).
+            memory: the sequence from the last layer of the encoder (required).
+            tgt_mask: the mask for the tgt sequence (optional).
+            memory_mask: the mask for the memory sequence (optional).
+            tgt_key_padding_mask: the mask for the tgt keys per batch (optional).
+            memory_key_padding_mask: the mask for the memory keys per batch (optional).
+            tgt_is_causal: If specified, applies a causal mask as ``tgt mask``.
+                Default: ``None``; try to detect a causal mask.
+                Warning:
+                ``tgt_is_causal`` provides a hint that ``tgt_mask`` is
+                the causal mask. Providing incorrect hints can result in
+                incorrect execution, including forward and backward
+                compatibility.
+            memory_is_causal: If specified, applies a causal mask as
+                ``memory mask``.
+                Default: ``False``.
+                Warning:
+                ``memory_is_causal`` provides a hint that
+                ``memory_mask`` is the causal mask. Providing incorrect
+                hints can result in incorrect execution, including
+                forward and backward compatibility.
+
+        Shape:
+            see the docs in :class:`~torch.nn.Transformer`.
+        """
+        output = tgt
+
+        # seq_len = _get_seq_len(tgt, self.layers[0].self_attn.batch_first)
+        # tgt_is_causal = _detect_is_causal_mask(tgt_mask, tgt_is_causal, seq_len)
+        if self.sinu is not None:
+            assert self.pm_freqs is None
+            for key in self.sinu:
+                self.sinu[key] = self.sinu[key].to(output.device)
+        if self.pm_freqs is not None:
+            assert self.sinu is None
+            if not self.training and hasattr(self, "pm_sinu") and past is not None and past[0].ndim > 2: # inference mode, will use cached sinu for the same example
+                assert self.pm_sinu['q'] is not None and self.pm_sinu['k'] is not None
+                # check batch size, need to modify the batch size if we use multi_trial during inference
+                if self.pm_sinu['q']['cos'].shape[0] != tgt.shape[0]:
+                    if self.pm_sinu['q']['cos'].shape[0] > tgt.shape[0]:
+                        self.pm_sinu['q']['cos'] = self.pm_sinu['q']['cos'][:tgt.shape[0]]
+                        self.pm_sinu['q']['sin'] = self.pm_sinu['q']['sin'][:tgt.shape[0]]
+                        self.pm_sinu['k']['cos'] = self.pm_sinu['k']['cos'][:tgt.shape[0]]
+                        self.pm_sinu['k']['sin'] = self.pm_sinu['k']['sin'][:tgt.shape[0]]
+                    else:
+                        assert self.pm_sinu['q']['cos'].shape[0] == 1
+                        self.pm_sinu['q']['cos'] = self.pm_sinu['q']['cos'].repeat(tgt.shape[0], 1, 1, 1)
+                        self.pm_sinu['q']['sin'] = self.pm_sinu['q']['sin'].repeat(tgt.shape[0], 1, 1, 1)
+                        self.pm_sinu['k']['cos'] = self.pm_sinu['k']['cos'].repeat(tgt.shape[0], 1, 1, 1)
+                        self.pm_sinu['k']['sin'] = self.pm_sinu['k']['sin'].repeat(tgt.shape[0], 1, 1, 1)
+                pass
+            else:
+                self.pm_freqs = self.pm_freqs.to(output.device)
+                if query_lens is None:
+                    query_lens = (~tgt_key_padding_mask).int().sum(-1).to(tgt.device)
+                if key_lens is None:
+                    key_lens = (~memory_key_padding_mask).int().sum(-1).to(tgt.device)
+                assert key_lens.ndim==1, key_lens
+                assert query_lens.ndim==1, query_lens
+                q_lens_expanded = query_lens.unsqueeze(-1).unsqueeze(-1) # [B, 1, 1]
+                k_lens_expanded = key_lens.unsqueeze(-1).unsqueeze(-1) # [B, 1, 1]
+                query_ids_multiple = q_lens_expanded / (q_lens_expanded - 1)
+                key_ids_multiple = k_lens_expanded / (k_lens_expanded - 1)
+                q_emb = self.pm_freqs * query_ids_multiple # [B, q_len_max, d]
+                k_emb = self.pm_freqs * key_ids_multiple # [B, k_len_max, d]
+                q_emb = q_emb / q_lens_expanded * self.progress_scale
+                k_emb = k_emb / k_lens_expanded * self.progress_scale
+                q_cos = q_emb.cos().unsqueeze(1) # [B, 1, q_len_max, d] # 1 is for nhead
+                q_sin = q_emb.sin().unsqueeze(1)
+                k_cos = k_emb.cos().unsqueeze(1)
+                k_sin = k_emb.sin().unsqueeze(1)
+                self.pm_sinu = {"q": {"cos": q_cos, "sin": q_sin}, "k": {"cos": k_cos, "sin": k_sin}}
+        else:
+            self.pm_sinu = {"q": None, "k": None}
+
+        output = {"input": output, "pm_sinu": self.pm_sinu, "sinu": self.sinu, "args": self.args}
+        if past != None:
+            all_present = []
+        if self.training and getattr(self.args, "attention_alignment_loss", 0):
+            all_attn_weights = []
+        for i, mod in enumerate(self.layers):
+            output = mod(output, memory, tgt_mask=tgt_mask,
+                        memory_mask=memory_mask,
+                        tgt_key_padding_mask=tgt_key_padding_mask,
+                        memory_key_padding_mask=memory_key_padding_mask,
+                        past=past[i] if past != None else None
+                        #  tgt_is_causal=tgt_is_causal,
+                        #  memory_is_causal=memory_is_causal
+                        )
+            if past != None:
+                output, cur_present = output
+                all_present.append(cur_present)
+            if isinstance(output, dict):
+                current_attn_weights = output["attention_weights"]
+                all_attn_weights.append(current_attn_weights)
+                output = output["x"]
+            if self.sinu is not None or self.pm_sinu is not None:
+                output = {"input": output, "pm_sinu": self.pm_sinu, "sinu": self.sinu, "args": self.args}
+        if self.pm_sinu is not None or self.sinu is not None:
+            output = output["input"]
+        if self.norm is not None:
+            output = self.norm(output)
+        if self.training and getattr(self.args, "attention_alignment_loss", 0):
+            assert len(all_attn_weights) == self.num_layers, f"{len(all_attn_weights)=}, {self.num_layers=}"
+            output = {"output": output, "attention_weights": all_attn_weights}
+        if past != None:
+            all_present = torch.stack(all_present, dim=0)
+            output = [output, all_present]
+        else:
+            output = [output, None]
+        return output

models/modules/utils.py

@@ -0,0 +1,36 @@
+import torch
+
+
+def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
+    """
+    Args:
+      lengths:
+        A 1-D tensor containing sentence lengths.
+      max_len:
+        The length of masks.
+    Returns:
+      Return a 2-D bool tensor, where masked positions
+      are filled with `True` and non-masked positions are
+      filled with `False`.
+
+    >>> lengths = torch.tensor([1, 3, 2, 5])
+    >>> make_pad_mask(lengths)
+    tensor([[False,  True,  True,  True,  True],
+            [False, False, False,  True,  True],
+            [False, False,  True,  True,  True],
+            [False, False, False, False, False]])
+    """
+    assert lengths.ndim == 1, lengths.ndim
+    max_len = max(max_len, lengths.max())
+    n = lengths.size(0)
+    seq_range = torch.arange(0, max_len, device=lengths.device)
+    expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len)
+
+    return expaned_lengths >= lengths.unsqueeze(-1)
+
+def generate_partial_autoregressive_mask(sz, start, end):
+    mask = torch.zeros(sz, sz).bool()
+    mask[start:end, start:end] = torch.triu(torch.ones(end-start, end-start,dtype=torch.bool), diagonal=1)
+    mask[:start, start:end] = True
+    mask[end:, start:end] = True
+    return mask

models/modules/visualizer.py

@@ -0,0 +1,107 @@
+# cp from https://github.com/lifeiteng/vall-e/blob/main/valle/models/visualizer.py
+#!/usr/bin/env python3
+# Copyright    2023                           (authors: Feiteng Li)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Dict, List, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+
+def visualize(
+    predicts: Tuple[torch.Tensor],
+    batch: Dict[str, Union[List, torch.Tensor]],
+    output_dir: str,
+    limit: int = 4,
+) -> None:
+    text_tokens = batch["text_tokens"].to("cpu").detach().numpy()
+    text_tokens_lens = batch["text_tokens_lens"].to("cpu").detach().numpy()
+    audio_features = batch["audio_features"].to("cpu").detach().numpy()
+    audio_features_lens = (
+        batch["audio_features_lens"].to("cpu").detach().numpy()
+    )
+    assert text_tokens.ndim == 2
+
+    utt_ids, texts = batch["utt_id"], batch["text"]
+
+    encoder_outputs = predicts[0].to("cpu").type(torch.float32).detach().numpy()
+    decoder_outputs = predicts[1]
+    if isinstance(decoder_outputs, list):
+        decoder_outputs = decoder_outputs[-1]
+    decoder_outputs = (
+        decoder_outputs.to("cpu").type(torch.float32).detach().numpy()
+    )
+
+    vmin, vmax = 0, 1024  # Encodec
+    if decoder_outputs.dtype == np.float32:
+        vmin, vmax = -6, 0  # Fbank
+
+    num_figures = 3
+    for b, (utt_id, text) in enumerate(zip(utt_ids[:limit], texts[:limit])):
+        _ = plt.figure(figsize=(14, 8 * num_figures))
+
+        S = text_tokens_lens[b]
+        T = audio_features_lens[b]
+
+        # encoder
+        plt.subplot(num_figures, 1, 1)
+        plt.title(f"Text: {text}")
+        plt.imshow(
+            X=np.transpose(encoder_outputs[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=S - 0.4, linewidth=2, color="r")
+        plt.xlabel("Encoder Output")
+        plt.colorbar()
+
+        # decoder
+        plt.subplot(num_figures, 1, 2)
+        plt.imshow(
+            X=np.transpose(decoder_outputs[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+            vmin=vmin,
+            vmax=vmax,
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=T - 0.4, linewidth=2, color="r")
+        plt.xlabel("Decoder Output")
+        plt.colorbar()
+
+        # target
+        plt.subplot(num_figures, 1, 3)
+        plt.imshow(
+            X=np.transpose(audio_features[b]),
+            cmap=plt.get_cmap("jet"),
+            aspect="auto",
+            interpolation="nearest",
+            vmin=vmin,
+            vmax=vmax,
+        )
+        plt.gca().invert_yaxis()
+        plt.axvline(x=T - 0.4, linewidth=2, color="r")
+        plt.xlabel("Decoder Target")
+        plt.colorbar()
+
+        plt.savefig(f"{output_dir}/{utt_id}.png")
+        plt.close()

models/voice_star.py

@@ -0,0 +1,784 @@
+import random, os, copy
+from typing import Dict, Iterator, List, Tuple, Union
+import logging
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchmetrics.classification import MulticlassAccuracy
+import torch.distributed as dist
+
+from .modules.utils import make_pad_mask, generate_partial_autoregressive_mask
+
+from .modules.embedding import SinePositionalEmbedding, TokenEmbedding, SinePositionalEmbedding_progress
+from .modules.transformer import (
+    AdaptiveLayerNorm,
+    LayerNorm,
+    TransformerDecoderLayer,
+    TransformerDecoder,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+
+def top_k_top_p_filtering(
+    logits, top_k=0, top_p=1.0, min_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1
+):
+    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+    Args:
+        logits: logits distribution shape (batch size, vocabulary size)
+        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+        Make sure we keep at least min_tokens_to_keep per batch example in the output
+    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+    """
+    if min_p < 1.0:
+        probs = F.softmax(logits, dim=-1)
+        indices_to_remove = probs < min_p
+        if not torch.any(indices_to_remove.sum(-1) == logits.size(-1)):
+            logits[indices_to_remove] = filter_value
+            top_k = 0
+            top_p = 1.0
+        # else will use other types of sampling, or no filtering
+            
+    # If top_k is a single integer
+    if isinstance(top_k, int) and top_k > 0:
+        # Safety check to ensure we don't ask for more than available
+        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))
+
+        # Remove all tokens with a probability less than the last token of the top-k
+        threshold = torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
+        indices_to_remove = logits < threshold
+        logits[indices_to_remove] = filter_value
+
+    # If top_k is a list, assume it has the same length as M
+    elif isinstance(top_k, list):
+        # Ensure the length matches the first dimension
+        assert len(top_k) == logits.size(0), \
+            f"top_k list length ({len(top_k)}) must match logits.size(0) ({logits.size(0)})"
+
+        for i in range(logits.size(0)):
+            k_i = top_k[i]
+            if k_i > 0:
+                # Safety check
+                k_i = min(max(k_i, min_tokens_to_keep), logits.size(-1))
+                row_threshold = torch.topk(logits[i], k_i, dim=-1)[0][-1]
+                indices_to_remove_i = logits[i] < row_threshold
+                logits[i, indices_to_remove_i] = filter_value
+
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(
+            F.softmax(sorted_logits, dim=-1), dim=-1
+        )
+
+        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        if min_tokens_to_keep > 1:
+            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+            sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+            ..., :-1
+        ].clone()
+        sorted_indices_to_remove[..., 0] = 0
+
+    return logits
+
+
+def topk_sampling(logits, top_k=10, top_p=1.0, min_p=1.0, temperature=1.0):
+    # temperature: (`optional`) float
+    #     The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+    # top_k: (`optional`) int
+    #     The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+    # top_p: (`optional`) float
+    #     The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+    # Temperature (higher temperature => more likely to sample low probability tokens)
+    if temperature != 1.0:
+        logits = logits / temperature
+    # Top-p/top-k filtering
+    logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p, min_p=min_p)
+    # Sample
+    token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
+    return token
+
+
+
+class VoiceStar(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        assert self.args.enc_dec ^ self.args.dec, f"self.args.enc_dec: {self.args.enc_dec}, self.args.dec: {self.args.dec}"
+        if not getattr(self.args, "special_first", False):
+            self.args.special_first = 0
+        if not getattr(self.args, "n_special", False):
+            self.args.n_special = 3
+        self.args.eos = getattr(self.args, "eos", -1)
+        self.eog = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eog, dtype=torch.long), requires_grad=False) # [K 1]
+        if self.args.eos > 0:
+            assert self.args.eos != self.args.audio_pad_token and self.args.eos != self.args.empty_token, self.args.eos
+            self.eos = nn.Parameter(torch.full((self.args.n_codebooks, 1), self.args.eos, dtype=torch.long), requires_grad=False) # [K 1]
+        if type(self.args.audio_vocab_size) == str:
+            self.args.audio_vocab_size = eval(self.args.audio_vocab_size)
+        if type(self.args.audio_vocab_size) == list: # otherwise they are all lists
+            assert self.args.special_first
+            
+
+        self.n_text_tokens = self.args.text_vocab_size + 1
+        assert self.args.text_pad_token == self.args.text_vocab_size, f"self.args.text_vocab_size: {self.args.text_vocab_size}, self.args.text_pad_token: {self.args.text_pad_token}"
+        
+        if self.args.special_first and type(self.args.audio_vocab_size) == list:
+            self.n_audio_tokens = [tok + self.args.n_special for tok in self.args.audio_vocab_size] # special tokens: empty token, EOG token, audio pad token
+            assert self.args.empty_token == 0, self.args.empty_token
+            assert self.args.eog == 1, self.args.eog
+            assert self.args.audio_pad_token == 2, self.args.audio_pad_token
+        else:
+            self.n_audio_tokens = [self.args.audio_vocab_size + self.args.n_special] * self.args.n_codebooks # special tokens: empty token, EOG token, audio pad token
+            assert self.args.audio_vocab_size == self.args.empty_token, self.args.empty_token
+            assert self.args.eog == self.args.audio_vocab_size + 1, self.args.eog
+            assert self.args.audio_pad_token == self.args.audio_vocab_size + 2, self.args.audio_pad_token
+
+        self.text_embedding = TokenEmbedding(
+            dim_model=self.args.d_model,
+            vocab_size=self.n_text_tokens, 
+            dropout=self.args.text_embedding_dropout
+        )
+
+        self.audio_embedding = nn.ModuleList(
+            [
+                TokenEmbedding(
+                dim_model=self.args.audio_embedding_dim, 
+                vocab_size=self.n_audio_tokens[k], 
+                dropout=self.args.audio_embedding_dropout
+            ) for k in range(self.args.n_codebooks)
+            ]
+        )
+
+        rope_base = getattr(self.args, "rope_base", None)
+        use_sinusoidal = getattr(self.args, "use_sinusoidal", False)
+        use_sinusoidal_progress = getattr(self.args, "use_sinusoidal_progress", False)
+        logging.info(f"rope_base: {rope_base}, use_sinusoidal: {use_sinusoidal}")
+        if use_sinusoidal:
+            self.text_positional_embedding = SinePositionalEmbedding(
+                self.args.d_model,
+                dropout=self.args.text_positional_embedding_dropout,
+                scale=False,
+                alpha=True, # learnable scaler, scale the volume of positional embedding
+            )
+            self.audio_positional_embedding = SinePositionalEmbedding(
+                self.args.d_model,
+                dropout=self.args.audio_positional_embedding_dropout,
+                scale=False,
+                alpha=True, # learnable scaler, scale the volume of positional embedding
+            )
+        elif use_sinusoidal_progress:
+            self.text_positional_embedding = SinePositionalEmbedding_progress(
+                self.args.d_model,
+                dropout=self.args.text_positional_embedding_dropout,
+                scale=False,
+                alpha=True, # learnable scaler, scale the volume of positional embedding
+                args = self.args
+            )
+            self.audio_positional_embedding = SinePositionalEmbedding_progress(
+                self.args.d_model,
+                dropout=self.args.audio_positional_embedding_dropout,
+                scale=False,
+                alpha=True, # learnable scaler, scale the volume of positional embedding
+                args = self.args
+            )
+
+        else:
+            class NoOp:
+                def __init__(self):
+                    pass
+                def __call__(self, *args, **kwargs):
+                    return args[0]
+            self.text_positional_embedding = NoOp()
+            self.audio_positional_embedding = NoOp()
+
+        if self.args.enc_dec:
+            enc_layer = TransformerEncoderLayer(
+                d_model=self.args.d_model,
+                nhead=self.args.nhead,
+                dim_feedforward=self.args.d_model*4,
+                dropout=self.args.trm_dropout,
+                batch_first=True,
+                norm_first=True,
+                layer_norm_cls=LayerNorm
+            ) # use the pre-norm arch 
+
+            self.encoder = TransformerEncoder(
+                encoder_layer=enc_layer, 
+                num_layers=self.args.num_encoder_layers, 
+                norm=LayerNorm(self.args.d_model),
+                rope_base = self.args.rope_base,
+                d_model = self.args.d_model,
+                nhead = self.args.nhead,
+                args = self.args
+            ) # use the pre-norm arch
+
+            dec_layer = TransformerDecoderLayer(
+                d_model=self.args.d_model,
+                nhead=self.args.nhead,
+                dim_feedforward=self.args.d_model*4,
+                dropout=self.args.trm_dropout,
+                batch_first=True,
+                norm_first=True,
+                layer_norm_cls=LayerNorm
+            )
+
+            self.decoder = TransformerDecoder(
+                decoder_layer=dec_layer, 
+                num_layers=self.args.num_decoder_layers,
+                norm=LayerNorm(self.args.d_model),
+                rope_base = self.args.rope_base,
+                d_model = self.args.d_model,
+                nhead = self.args.nhead,
+                args = self.args
+            ) # NOTE: this one I use torch.nn native implementation, as it's not implemented in .modules
+
+        else:
+            dec_layer = TransformerEncoderLayer(
+                self.args.d_model,
+                self.args.nhead,
+                dim_feedforward=self.args.d_model * 4,
+                dropout=self.args.trm_dropout,
+                batch_first=True,
+                norm_first=True,
+                layer_norm_cls=LayerNorm
+            )
+            self.decoder = TransformerEncoder(
+            dec_layer,
+            num_layers=self.args.num_decoder_layers,
+            norm=LayerNorm(self.args.d_model),
+        )
+        
+        if type(self.args.audio_vocab_size) == int:
+            self.predict_layer = nn.ModuleList(
+                [
+                    nn.Sequential(nn.Linear(self.args.d_model, self.args.audio_vocab_size//2), nn.GELU(), nn.Linear(self.args.audio_vocab_size//2, self.n_audio_tokens[k])) for k in range(self.args.n_codebooks)
+                ]
+            )
+        else:
+            self.predict_layer = nn.ModuleList(
+                [
+                    nn.Sequential(nn.Linear(self.args.d_model, self.args.d_model//2), nn.GELU(), nn.Linear(self.args.d_model//2, self.n_audio_tokens[k])) for k in range(self.args.n_codebooks)
+                ]
+            )
+
+        self.accuracy_metrics = nn.ModuleList(
+            [MulticlassAccuracy(
+                self.n_audio_tokens[k],
+                top_k=10,
+                average="micro",
+                multidim_average="global",
+                ignore_index=None,
+            ) for k in range(self.args.n_codebooks)]
+        )
+
+        if self.args.eog_weight != 1:
+            raise NotImplementedError("now have different vocab_size for different codebooks, therefore currently don't support eog_weight")
+            self.class_weight = nn.Parameter(torch.ones(self.n_audio_tokens), requires_grad=False)
+            self.class_weight.data[self.args.eog] = self.args.eog_weight
+
+    def dec_forward(
+            self, 
+            x_input, 
+            x_lens,
+            x_attention_mask,
+            x_padding_mask,
+            y_input,
+            new_y_lens,
+            y_attention_mask,
+            y_padding_mask,
+            need_weights=False,
+            past=None,
+            last_3_tokens=False
+        ):
+            x_attn_mask = F.pad(
+                x_attention_mask,
+                (0, new_y_lens.max()),
+                value=True,
+            ) # x attn to all x, doesn't attn to any y, this follow figure 3 of the valle paper
+            y_attn_mask = F.pad(
+                y_attention_mask,
+                (x_lens.max(), 0), # y is padded at the front
+                value=False,
+            ) # y attn to all x, for y itself use lower triangle mask to ensure autoregressive
+            xy_attn_mask = torch.concat([x_attn_mask, y_attn_mask], dim=0)
+
+            # merge key padding and attention masks
+            bsz, src_len = x_input.shape[0], x_lens.max() + new_y_lens.max()
+            xy_padding_mask = torch.concat([x_padding_mask, y_padding_mask], dim=1)
+            _xy_padding_mask = (
+                xy_padding_mask.view(bsz, 1, 1, src_len)
+                .expand(-1, self.args.nhead, -1, -1)
+                .reshape(bsz * self.args.nhead, 1, src_len)
+            )
+            xy_attn_mask = xy_attn_mask.logical_or(_xy_padding_mask)
+
+            new_attn_mask = torch.zeros_like(xy_attn_mask)
+            new_attn_mask.masked_fill_(xy_attn_mask, float("-inf"))
+            xy_attn_mask = new_attn_mask
+
+            xy_input = torch.cat([x_input, y_input], dim=1)
+            if need_weights:
+                raise NotImplementedError("not implemented yet")
+                out, layer_attn_weights =  self.decoder((xy_input, None), mask=xy_attn_mask, need_weights=True)
+                return layer_attn_weights
+
+            if past == None: # do not use kvcache
+                out, _ =  self.decoder((xy_input, None), mask=xy_attn_mask)
+                return out[:, x_lens.max():], None
+            else: # use kvcache
+                if past.ndim > 3: # uses kvcache, only need to pass the last tokens, this doesn't work with multi-span speech editing yet
+                    if last_3_tokens:
+                        xy_input = xy_input[:, -3:]
+                        xy_attn_mask = xy_attn_mask[:, -3:]
+                    else:
+                        xy_input = xy_input[:, -1:]
+                        xy_attn_mask = xy_attn_mask[:, -1:]
+
+                out, present =  self.decoder((xy_input, None), mask=xy_attn_mask, past=past)
+                if isinstance(out, tuple): # get rid of stage_embedding
+                    out = out[0]
+
+                if out.shape[1] > x_lens.max(): # the first pass, not kvcache yet
+                    return out[:, x_lens.max():], present
+                else: # used kvcache
+                    return out, present
+
+    def enc_dec_forward(
+            self,
+            xa,
+            x_attention_mask,
+            x_padding_mask,
+            y_input,
+            new_y_lens,
+            y_attention_mask,
+            y_padding_mask,
+            tgt_y_lens=None,
+            need_weights=False,
+            past=None,
+            last_3_tokens=False
+        ):
+        assert not need_weights
+        if past != None and past.ndim > 3:
+            y_input = y_input[:, -1:]
+            y_attention_mask = y_attention_mask[-1:]
+        yhat, present = self.decoder(tgt=y_input, memory=xa, tgt_mask=y_attention_mask, tgt_key_padding_mask=y_padding_mask, memory_key_padding_mask=x_padding_mask, query_lens=tgt_y_lens, past=past)
+        return yhat, present
+
+    def forward(self, batch, calc_loss = False):
+        """
+        Args:
+          x:
+            A 2-D tensor of shape (N, S).
+          x_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (N, K, T).
+            where K is the number of codebooks
+          y_lens:
+            A 1-D tensor of shape (N,). It contains the number of tokens in `x`
+            before padding.
+        """
+        x, x_lens, y, y_lens = batch["x"], batch["x_lens"], batch["y"], batch["y_lens"]
+        if len(x) == 0:
+            return None
+        x = x[:, :x_lens.max()] # this deal with gradient accumulation, where x_lens.max() might not be longer than the length of the current slice of x
+        y = y[...,:y_lens.max()]
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3 and y.shape[1] == self.args.n_codebooks, y.shape
+        assert y_lens.ndim == 1, y_lens.shape
+        x_padding_mask = make_pad_mask(x_lens).to(x.device)
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x_padding_mask.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input, x_lens)
+        y_with_eos = [torch.cat([item[:, :y_lens[i]], self.eos], dim=-1) for i, item in enumerate(y)]
+        targets = y_with_eos
+        # apply delayed stacking on y
+        shifted_y = []
+        patterns = []
+        new_y_lens = []
+        if getattr(self, "empty_tokens", None) == None:
+            self.empty_tokens = torch.full((self.args.n_codebooks, self.args.n_codebooks), self.args.empty_token, dtype=torch.long).to(y.device) # [K, K]
+        for i in range(len(y)):
+            tmp = torch.cat([y_with_eos[i], self.empty_tokens], dim=-1) # [K, T+n_codebooks]
+            for ii in range(self.args.n_codebooks):
+                tmp[ii] = torch.roll(tmp[ii], shifts=ii+1, dims=0)
+            shifted_y.append(tmp.transpose(1,0)) # [K, T+n_codebooks] -> [T+n_codebooks, K]
+            new_y_lens.append(y_with_eos[i].shape[1] + self.empty_tokens.shape[1])
+            
+        new_y_lens = torch.LongTensor(new_y_lens).to(y.device)
+
+        cated_y = torch.nn.utils.rnn.pad_sequence(shifted_y, batch_first=False, padding_value=self.args.audio_pad_token)
+        assert cated_y.shape == torch.Size([max(new_y_lens), len(y), self.args.n_codebooks]), cated_y.shape
+        cated_y = cated_y.permute(2,0,1) # [T,B,K]->[K,T,B]
+        stacked_embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, T, B, D]
+        assert stacked_embedded_y.shape[0] == self.args.n_codebooks and stacked_embedded_y.shape[2] == len(y) and stacked_embedded_y.shape[-1] == self.args.d_model, stacked_embedded_y.shape
+        embedded_y = stacked_embedded_y.sum(dim=0) # [K,T,B,D]->[T,B,D]
+        embedded_y = embedded_y.transpose(1,0) # [T,B,D]->[B,T,D]
+        assert embedded_y.shape[1:] == torch.Size([max(new_y_lens), self.args.d_model]), embedded_y.shape
+        y_input = self.audio_positional_embedding(embedded_y, new_y_lens)
+        y_padding_mask = make_pad_mask(new_y_lens).to(y.device)
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y_padding_mask.device)
+        if self.args.dec:
+            y_out = self.dec_forward(
+                    x_input, 
+                    x_lens,
+                    x_attention_mask,
+                    x_padding_mask,
+                    y_input,
+                    new_y_lens,
+                    y_attention_mask,
+                    y_padding_mask
+                )
+        else:
+            xa = self.encoder(src=x_input, src_key_padding_mask=x_padding_mask)
+            y_out = self.enc_dec_forward(
+                        xa,
+                        x_attention_mask,
+                        x_padding_mask,
+                        y_input,
+                        new_y_lens,
+                        y_attention_mask,
+                        y_padding_mask
+                    )
+        y_out = y_out[0] # no kv-caching during training
+        assert y_out.shape == y_input.shape, f"y_out.shape: {y_out.shape}, y_input.shape: {y_input.shape}" # [B S D]
+        logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card]
+        assert logits.shape[1] == self.args.n_codebooks and logits.shape[3] == self.n_audio_tokens[0], logits.shape
+        logits_use = [logit[:, :new_y_lens[i]] for i, logit in enumerate(logits)] # each of shape [K, T, card]
+        logits_final = []
+        for i, logit in enumerate(logits_use):
+            logit_copy = logit.clone()
+            for ii in range(self.args.n_codebooks):
+                logit_copy[ii] = torch.roll(logit_copy[ii], shifts=-ii, dims=0)
+            logit = logit_copy[:, :-self.args.n_codebooks] # [K, T, card] -> [K, T-n_codebooks, card]
+            logits_final.append(logit)
+        if self.args.no_loss_on_prefix:
+            assert "y_sep_token_position" in batch, f"y_sep_token_position should be in batch, but it's not"
+            logit_temp = []
+            target_temp = []
+            for jj, (logit, target) in enumerate(zip(logits_final, targets)):
+                # TODO already taken into consideration in depth transformer
+                logit_temp.append(logit[:, batch['y_sep_token_position'][jj]:])
+                target_temp.append(target[:, batch['y_sep_token_position'][jj]:]) 
+            logits_final = logit_temp
+            targets = target_temp
+        logits = torch.cat(logits_final, dim=1) # [K, T1+T2+T3+..., card]
+        targets = torch.cat(targets, dim=1) # [K, T1+T2+T3+...]
+        
+        assert targets.shape[:2] == logits.shape[:2], f"{targets.shape}, {logits.shape}"
+        loss = []
+        ntokens = []
+        top10acc = []
+        for k, (logit, target) in enumerate(zip(logits, targets)): # even though the loss and top10acc is calculated in a loop (loop through n_codebooks), validation is still taking a lot of mem, need to optimize this a little more
+            loss.append(F.cross_entropy(logit, target, reduction='mean', weight=self.class_weight.data if self.args.eog_weight!=1 else None, ignore_index=self.args.y_sep_token if self.args.y_sep_token != None else -100)) # ignore audio sep token as it's unpredictable (like the random early stop bug happened in 2023)
+            # NOTE have to ignore the sep token in the loss calculation
+            top10acc.append(self.accuracy_metrics[k](logit.detach(), target))
+            ntokens.append(len(logit))
+        
+        all_ntokens = sum(ntokens)
+        if self.args.codebook_weight != None:
+            codebook_weight = eval(self.args.codebook_weight) if isinstance(self.args.codebook_weight, str) else self.args.codebook_weight
+        else:
+            codebook_weight = [1.] * self.args.n_codebooks
+        perplexity_by_codebook = [torch.exp(l) for l in loss]
+        loss = sum([l*nt*cw for l, nt, cw in zip(loss, ntokens, codebook_weight)])
+        
+        top10acc_by_codebook = [t10a*nt for t10a, nt in zip(top10acc, ntokens)]
+        top10acc = sum(top10acc_by_codebook)
+
+        ntokens = torch.tensor(all_ntokens).to(logits.device)
+        
+        ret = {
+            "loss": loss,
+            "perplexity_by_codebook": perplexity_by_codebook,
+            "top10acc": top10acc,
+            "top10acc_by_codebook": top10acc_by_codebook,
+            "effective_ntoken": ntokens,
+        }
+
+        return ret
+
+    def inference_tts(
+        self,
+        x: torch.Tensor,
+        x_lens: torch.Tensor,
+        y: torch.Tensor,
+        tgt_y_lens: torch.Tensor, # 
+        top_k: Union[int, list[int]]=-100,
+        top_p: float=1.0,
+        min_p: float=1.0,
+        temperature: float=1.0,
+        stop_repetition: int=3,
+        kvcache: int=1,
+        silence_tokens: list[int]=[],
+        multi_trial: list[int]=[],
+        *kargs
+    ) -> torch.Tensor:
+        """
+        different from inference_tts, this implementation uses kvcache, which should have significant speed up
+        Args:
+          x:
+            A 2-D tensor of shape (1, L).
+          x_lens:
+            A 1-D tensor of shape (1,). It contains the number of tokens in `x`
+            before padding.
+          y:
+            A 3-D tensor of shape (1, T, K).
+          tgt_y_lens:
+            *new arg* this specify the target length of y
+          top_k: (`optional`) int
+            The number of highest probability tokens to keep for top-k-filtering. Default to -100.
+          top_p: (`optional`) float
+            For Neucleus sampling
+          min_p: (`optional`) float
+            For min_p filtered sampling
+          temperature: (`optional`) float
+            The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+          multi_trial: (`optional`) list[int]
+            If not empty, it will be [n_trials, beam_size, trial_interval]
+            from the start and begining trial_interval, we duplicate the current sample by beam_size, 
+            at the end of every trial_interval, we choose the sample with the highest log likelihood to keep and throw away the rest
+        """
+        eog_inference = self.args.eos if self.args.eos>0 else self.args.eog
+        assert x.ndim == 2, x.shape
+        assert x_lens.ndim == 1, x_lens.shape
+        assert y.ndim == 3, y.shape
+        if self.args.special_first:
+            y = y + int(self.args.n_special)
+        y = y.transpose(2,1) # [1,T,K] -> [1,K,T]
+        assert y.shape[0] == 1 and y.shape[1] == self.args.n_codebooks, y.shape # there is no padding
+
+        # make x attention mask and x_input
+        x_attention_mask = torch.triu(torch.ones(x.shape[1], x.shape[1]), diagonal=1).bool().to(x.device)
+        # x_attention_mask = torch.zeros(x.shape[1], x.shape[1]).bool().to(x.device)
+        x_input = self.text_embedding(x)
+        x_input = self.text_positional_embedding(x_input, x_lens)
+
+        y_len = y.shape[2]
+        y_lens = torch.LongTensor([y_len]).to(y.device)
+
+        # rearrange y, we don't add eog to the end, this doesn't actually do anything in the tts scenario
+        rearranged_y = [[y[0]]]
+        assert rearranged_y[0][0].shape[0] == self.args.n_codebooks, rearranged_y[0][0].shape
+        
+        # # shift y to create the delayed pattern
+        if getattr(self, "empty_tokens", None) == None:
+            self.empty_tokens = torch.full((self.args.n_codebooks, self.args.n_codebooks), self.args.empty_token, dtype=torch.long).to(y.device) # [K, K]
+        temp = rearranged_y[0][0]
+        assert temp.ndim == 2 and temp.shape[0] == self.args.n_codebooks, temp.shape
+        temp = torch.cat([temp, self.empty_tokens], dim=-1) # [K, T+n_codebooks]
+        for ii in range(self.args.n_codebooks):
+            temp[ii] = torch.roll(temp[ii], shifts=ii+1, dims=0)
+        shifted_y = [[temp]]
+
+        # below is different from forward or inference
+        # where we cut this shifted part 
+        shifted_y[0][0] = shifted_y[0][0][:, :-(self.args.n_codebooks-1)]
+        assert not (shifted_y[0][0][self.args.n_codebooks:] == self.args.empty_token).any() and not (shifted_y[0][0][self.args.n_codebooks:] == self.args.eog).any(), shifted_y[0][0]
+        # next section in inference is insert mask at the intersection of each tensor in a sample, but we don't need to do that 
+        # next section is concate tensors of each sample to one tensor, which we also don't need
+        cated_y = shifted_y[0][0].unsqueeze(-1) #[K,S]->[K,S,B]
+        new_y_lens = torch.LongTensor([cated_y.shape[1]]).to(cated_y.device)
+        assert cated_y.shape == torch.Size((self.args.n_codebooks, cated_y.shape[1], 1))
+        assert not (cated_y == self.args.audio_pad_token).any(), cated_y
+
+        # replace tokens in y with the embeddings, add sum codebooks up
+        embedded_y = torch.stack([self.audio_embedding[k](cated_y[k]) for k in range(self.args.n_codebooks)], dim=0) # [K, S, B, D]
+        assert embedded_y.shape[0] == self.args.n_codebooks, embedded_y.shape
+        assert embedded_y.shape[-1] == self.args.d_model, embedded_y.shape
+        embedded_y = embedded_y.sum(dim=0) # [K,S,B,D]->[S,B,D]
+        embedded_y = embedded_y.transpose(1,0) # [S,B,D]->[B,S,D]
+        
+        # positional embedding
+        y_input = self.audio_positional_embedding(embedded_y, tgt_y_lens)
+
+        # make attention mask and padding mask
+        y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+
+        x_padding_mask = torch.full((1,x_lens[0]), False).to(x.device)
+        y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+
+        # entering the generation stage
+        # starting from line 708
+        codebook_eog = [False] * self.args.n_codebooks
+        generated = [] # doesn't contain any empty token, contain eog
+        cur_generated = []
+        # say 0 is empty, 4 is eog
+        # tensor([[ 1,  2,  3,  4,  0,  0],
+        #         [ 0,  1,  2,  3,  4,  0],
+        #         [ 0,  0,  1,  2,  3,  4]])
+        num_gen = []
+        cur_num_gen = 0
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+        # silence_tokens = [1388,1898,131] # [1388, 2045, 2041, 1996]
+        # silence_tokens = []
+        consec_silence_count = 0
+        prev_token = None
+        ##################### silence repetition handling #####################
+        ##################### silence repetition handling #####################
+
+        def sample_helper(n_eog, logits, codebook_eog, top_k, top_p, min_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen):
+            if n_eog == 0:
+                logits_adjust = logits
+                for jj in range(1,self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                if cur_num_gen <= self.args.encodec_sr // 5: # this shouldn't happen, but just in case the model stopped too early
+                    logits_adjust[0][eog_inference] = -10000
+                ##################### silence repetition handling #####################
+                if stop_repetition > 0 and prev_token in silence_tokens and consec_silence_count > stop_repetition:
+                    if logits_adjust[0, prev_token] < 0:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] * (consec_silence_count - (stop_repetition-1))
+                    else:
+                        logits_adjust[0, prev_token] = logits_adjust[0, prev_token] / (consec_silence_count - (stop_repetition-1))
+                ##################### silence repetition handling #####################
+                samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, min_p=min_p, temperature=temperature
+                    ) # [K, 1]
+                assert samples.shape == torch.Size((self.args.n_codebooks, 1)), f"samples.shape: {samples.shape}"
+                if cur_num_gen < self.args.n_codebooks-1:
+                    for jj in range(1, self.args.n_codebooks - cur_num_gen):
+                        samples[-jj, 0] = self.args.empty_token
+                
+                if (
+                    samples[0,0] == eog_inference or torch.argmax(logits[0], dim=-1) == eog_inference or y_input.shape[1] > x_lens[0] * (self.args.encodec_sr//4)
+                ) or self.args.rope_base is not None and not self.args.decoder_regular_rope and self.args.progress_no_multiple and cur_num_gen > (tgt_y_lens[0] + self.args.encodec_sr * getattr(self.args, "extra_cutoff", 5)): 
+                # last one condition in the first bracket means y is already too long, shouldn't happen, but put it here
+                # the second bracket means we are using progress-monitoring RoPE, but the model is generating excessively long sequence (5 seconds more than specified), in which case we terminate the generation
+                    samples[0,0] = eog_inference
+                    codebook_eog[0] = True
+                ##################### silence repetition handling #####################
+                if samples[0,0] in silence_tokens and samples[0,0] == prev_token:
+                    consec_silence_count += 1
+                else:
+                    consec_silence_count = 0
+                prev_token = samples[0,0]
+                ##################### silence repetition handling #####################
+                return samples, codebook_eog, prev_token, consec_silence_count
+            else:
+                assert sum(codebook_eog[i] for i in range(n_eog)) == n_eog, f"codebook_eog: {codebook_eog}, but n_eog: {n_eog}"
+                logits_adjust = logits
+                for jj in range(n_eog+1,self.args.n_codebooks):
+                    logits_adjust[jj][eog_inference] = -10000
+                    logits_adjust[jj][self.args.empty_token] = -10000
+                samples = topk_sampling(
+                        logits_adjust, top_k=top_k, top_p=top_p, min_p=min_p, temperature=temperature
+                    ) # [K, 1]
+                for jj in range(n_eog):
+                    samples[jj, 0] = self.args.empty_token
+                samples[n_eog, 0] = eog_inference
+                codebook_eog[n_eog] = True
+                return samples, codebook_eog, prev_token, consec_silence_count
+        
+        # prepare the cache placeholder
+        # n_layers, 2, bsz, num_heads, src_len, head_dim, 2 means [key, value]
+        past = torch.ones([self.args.num_decoder_layers, 2, x.shape[0]], device=x.device, dtype=torch.float32) if kvcache else None
+        if self.args.enc_dec:
+            xa = self.encoder(src=x_input, src_key_padding_mask=x_padding_mask)
+        while True:
+            if self.args.dec:
+                y_out, present = self.dec_forward(
+                                x_input, 
+                                x_lens,
+                                x_attention_mask,
+                                x_padding_mask,
+                                y_input,
+                                new_y_lens,
+                                y_attention_mask,
+                                y_padding_mask,
+                                past=past
+                            )
+            else:
+                y_out, present = self.enc_dec_forward(
+                                xa,
+                                x_attention_mask,
+                                x_padding_mask,
+                                y_input,
+                                new_y_lens,
+                                y_attention_mask,
+                                y_padding_mask,
+                                tgt_y_lens=tgt_y_lens,
+                                past=past
+                            )
+            if past != None:
+                past = torch.cat([past, present.to(past.dtype)], dim=-2) if past.ndim > 3 else present.to(past.dtype)
+
+
+            y_out = y_out[:, -1:] # only take the last token
+
+            logits = torch.stack([self.predict_layer[i](y_out) for i in range(self.args.n_codebooks)], dim=1) # [B K S card], B==S==1, so [1 K 1 card]
+            logits = logits.squeeze(0).squeeze(1) # [K card]
+            assert logits.shape == torch.Size((self.args.n_codebooks, self.n_audio_tokens[0])), f"{logits.shape}"
+
+            n_eog = sum(codebook_eog)
+            assert n_eog < self.args.n_codebooks
+            if self.args.eos > 0: # if we are using end-of-sentence token (which is used by default), eog shouldn't be used here, as there is no masked spans
+                for jj in range(self.args.n_codebooks):
+                    logits[jj][self.args.eog] = -10000.
+            
+            samples, codebook_eog, prev_token, consec_silence_count = sample_helper(n_eog, logits, codebook_eog, top_k, top_p, min_p, temperature, prev_token, consec_silence_count, stop_repetition, silence_tokens, cur_num_gen)
+            # samples.shape is [K,1]
+            # ge samples_emb
+            samples_emb = torch.stack([self.audio_embedding[k](samples[k]) for k in range(self.args.n_codebooks)], dim=0) # [K,1,D]
+            samples_emb = samples_emb.sum(dim=0,keepdim=True) # [1,1,D]
+
+            cur_num_gen += 1
+            cur_generated.append(samples.squeeze(-1)) # [K,1] -> [K]
+
+            if sum(codebook_eog) == self.args.n_codebooks: # generation for the current span is done
+                codebook_eog = [False] * self.args.n_codebooks
+                num_gen.append(cur_num_gen)
+                cur_num_gen = 0
+                generated.append(cur_generated)
+                cur_generated = []
+                break
+            else:
+                assert samples_emb.shape == torch.Size((1,1,self.args.d_model)), f"samples_emb.shape: {samples_emb.shape}"
+            
+            embedded_y = torch.cat([embedded_y, samples_emb], dim=1)
+            new_y_lens = torch.LongTensor([embedded_y.shape[1]]).to(y.device)
+            y_input = self.audio_positional_embedding(embedded_y, tgt_y_lens) # [B T D]
+            # make attention mask and padding mask
+            y_attention_mask = torch.triu(torch.ones(y_input.shape[1], y_input.shape[1]), diagonal=1).bool().to(y.device)
+            y_padding_mask = torch.full((1,new_y_lens[0]), False).to(y.device)
+        
+        assert len(generated) == 1, f"len(generated): {len(generated)}"
+
+        # revert the pattern
+        flatten_gen = []
+        for l, orig_span in enumerate(generated):
+            span = torch.stack(orig_span, dim=0) # [T, K]
+            span = span.transpose(1,0) # [K, T]
+            assert span.shape[0] == self.args.n_codebooks, span.shape
+            unshifted_span = []
+            for j, s in enumerate(span):
+                start_from = j
+                end_at = - (self.args.n_codebooks - start_from)
+                unshifted_span.append(s[start_from:end_at])
+            unshifted_span = torch.stack(unshifted_span, dim=0)
+
+            assert unshifted_span.shape[1] == num_gen[l] - self.args.n_codebooks, f"len(unshifted_spans[0]): {len(unshifted_span[0])}, num_gen[l]: {num_gen[l]}"
+
+            flatten_gen.append(unshifted_span)
+        assert len(flatten_gen) == 1, len(flatten_gen)
+        
+        # combine 
+        res = [y[0], flatten_gen[0]]
+        res = torch.cat(res, dim=1).unsqueeze(0) # [K, new_t] -> [1, K, new_T]
+        expected_y_len = y_len + sum([item - self.args.n_codebooks for item in num_gen])
+        assert res.shape == torch.Size((1, self.args.n_codebooks, expected_y_len)), f"res.shape: {res.shape}, expected_y_len: {expected_y_len}. y_len + sum([item - self.args.n_codebooks for item in num_gen]): {y_len} + {sum([item - self.args.n_codebooks for item in num_gen])}"
+
+        if self.args.special_first:
+            res = res - int(self.args.n_special)
+            flatten_gen = flatten_gen - int(self.args.n_special)
+        return res, flatten_gen[0].unsqueeze(0)

steps/optim.py

@@ -0,0 +1,1123 @@
+# Copyright      2022  Xiaomi Corp.        (authors: Daniel Povey)
+#
+# See ../LICENSE for clarification regarding multiple authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import contextlib
+import logging
+import random
+from collections import defaultdict
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.optim import Optimizer
+
+
+class BatchedOptimizer(Optimizer):
+    """
+    This class adds to class Optimizer the capability to optimize parameters in batches:
+    it will stack the parameters and their grads for you so the optimizer can work
+    on tensors with an extra leading dimension.  This is intended for speed with GPUs,
+    as it reduces the number of kernels launched in the optimizer.
+
+    Args:
+      params:
+    """
+
+    def __init__(self, params, defaults):
+        super(BatchedOptimizer, self).__init__(params, defaults)
+
+    @contextlib.contextmanager
+    def batched_params(self, param_group, group_params_names):
+        """
+        This function returns (technically, yields) a list of
+          of tuples (p, state), where
+        p is a `fake` parameter that is stacked (over axis 0) from real parameters
+        that share the same shape, and its gradient is also stacked;
+        `state` is the state corresponding to this batch of parameters
+        (it will be physically located in the "state" for one of the real
+        parameters, the last one that has any particular shape and dtype).
+
+        This function is decorated as a context manager so that it can
+        write parameters back to their "real" locations.
+
+        The idea is, instead of doing:
+        <code>
+          for p in group["params"]:
+             state = self.state[p]
+             ...
+        </code>
+        you can do:
+        <code>
+          with self.batched_params(group["params"]) as batches:
+             for p, state, p_names in batches:
+                 ...
+        </code>
+
+        Args:
+          group: a parameter group, which is a list of parameters; should be
+                one of self.param_groups.
+          group_params_names: name for each parameter in group,
+                which is List[str].
+        """
+        batches = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter
+        batches_names = defaultdict(
+            list
+        )  # `batches` maps from tuple (dtype_as_str,*shape) to list of str
+
+        assert len(param_group) == len(group_params_names), f"len(param_group): {len(param_group)}, len(group_params_names): {len(group_params_names)}"
+        for p, named_p in zip(param_group, group_params_names):
+            key = (str(p.dtype), *p.shape)
+            batches[key].append(p)
+            batches_names[key].append(named_p)
+
+        batches_names_keys = list(batches_names.keys())
+        sorted_idx = sorted(
+            range(len(batches_names)), key=lambda i: batches_names_keys[i]
+        )
+        batches_names = [
+            batches_names[batches_names_keys[idx]] for idx in sorted_idx
+        ]
+        batches = [batches[batches_names_keys[idx]] for idx in sorted_idx]
+
+        stacked_params_dict = dict()
+
+        # turn batches into a list, in deterministic order.
+        # tuples will contain tuples of (stacked_param, state, stacked_params_names),
+        # one for each batch in `batches`.
+        tuples = []
+
+        for batch, batch_names in zip(batches, batches_names):
+            p = batch[0]
+            # we arbitrarily store the state in the
+            # state corresponding to the 1st parameter in the
+            # group.  class Optimizer will take care of saving/loading state.
+            state = self.state[p]
+            p_stacked = torch.stack(batch)
+            grad = torch.stack(
+                [
+                    torch.zeros_like(p) if p.grad is None else p.grad
+                    for p in batch
+                ]
+            )
+            p_stacked.grad = grad
+            stacked_params_dict[key] = p_stacked
+            tuples.append((p_stacked, state, batch_names))
+
+        yield tuples  # <-- calling code will do the actual optimization here!
+
+        for ((stacked_params, _state, _names), batch) in zip(tuples, batches):
+            for i, p in enumerate(batch):  # batch is list of Parameter
+                p.copy_(stacked_params[i])
+
+
+class ScaledAdam(BatchedOptimizer):
+    """
+     Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update
+     proportional to the norm of that parameter; and also learn the scale of the parameter,
+     in log space, subject to upper and lower limits (as if we had factored each parameter as
+     param = underlying_param * log_scale.exp())
+
+
+     Args:
+          params:  The parameters or param_groups to optimize (like other Optimizer subclasses)
+              lr:  The learning rate.  We will typically use a learning rate schedule that starts
+                   at 0.03 and decreases over time, i.e. much higher than other common
+                   optimizers.
+     clipping_scale: (e.g. 2.0)
+                   A scale for gradient-clipping: if specified, the normalized gradients
+                   over the whole model will be clipped to have 2-norm equal to
+                   `clipping_scale` times the median 2-norm over the most recent period
+                   of `clipping_update_period` minibatches.  By "normalized gradients",
+                   we mean after multiplying by the rms parameter value for this tensor
+                   [for non-scalars]; this is appropriate because our update is scaled
+                   by this quantity.
+            betas: beta1,beta2 are momentum constants for regular momentum, and moving sum-sq grad.
+                   Must satisfy 0 < beta <= beta2 < 1.
+     scalar_lr_scale: A scaling factor on the learning rate, that we use to update the
+                   scale of each parameter tensor and scalar parameters of the mode..
+                   If each parameter were decomposed
+                   as p * p_scale.exp(), where (p**2).mean().sqrt() == 1.0, scalar_lr_scale
+                   would be a the scaling factor on the learning rate of p_scale.
+              eps:  A general-purpose epsilon to prevent division by zero
+    param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be >= this value)
+    param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of
+                   learning the scale on the parameters (we'll constrain the rms of each non-scalar
+                   parameter tensor to be <= this value)
+       scalar_max: Maximum absolute value for scalar parameters (applicable if your
+                   model has any parameters with numel() == 1).
+    size_update_period: The periodicity, in steps, with which we update the size (scale)
+                   of the parameter tensor.  This is provided to save a little time
+                   in the update.
+     clipping_update_period: if clipping_scale is specified, this is the period
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=3e-02,
+        clipping_scale=None,
+        betas=(0.9, 0.98),
+        scalar_lr_scale=0.1,
+        eps=1.0e-08,
+        param_min_rms=1.0e-05,
+        param_max_rms=3.0,
+        scalar_max=10.0,
+        size_update_period=4,
+        clipping_update_period=100,
+        parameters_names=None,
+        show_dominant_parameters=True,
+    ):
+
+        assert parameters_names is not None, (
+            "Please prepare parameters_names,"
+            "which is a List[List[str]]. Each List[str] is for a group"
+            "and each str is for a parameter"
+        )
+        defaults = dict(
+            lr=lr,
+            clipping_scale=clipping_scale,
+            betas=betas,
+            scalar_lr_scale=scalar_lr_scale,
+            eps=eps,
+            param_min_rms=param_min_rms,
+            param_max_rms=param_max_rms,
+            scalar_max=scalar_max,
+            size_update_period=size_update_period,
+            clipping_update_period=clipping_update_period,
+        )
+
+        super(ScaledAdam, self).__init__(params, defaults)
+        assert len(self.param_groups) == len(parameters_names)
+        self.parameters_names = parameters_names
+        self.show_dominant_parameters = show_dominant_parameters
+
+    def __setstate__(self, state):
+        super(ScaledAdam, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        batch = True
+
+        for group, group_params_names in zip(
+            self.param_groups, self.parameters_names
+        ):
+
+            with self.batched_params(
+                group["params"], group_params_names
+            ) as batches:
+
+                # batches is list of pairs (stacked_param, state).  stacked_param is like
+                # a regular parameter, and will have a .grad, but the 1st dim corresponds to
+                # a stacking dim, it is not a real dim.
+
+                if (
+                    len(batches[0][1]) == 0
+                ):  # if len(first state) == 0: not yet initialized
+                    clipping_scale = 1
+                else:
+                    clipping_scale = self._get_clipping_scale(group, batches)
+
+                for p, state, _ in batches:
+                    # Perform optimization step.
+                    # grad is not going to be None, we handled that when creating the batches.
+                    grad = p.grad
+                    if grad.is_sparse:
+                        raise RuntimeError(
+                            "ScaledAdam optimizer does not support sparse gradients"
+                        )
+                    # State initialization
+                    if len(state) == 0:
+                        self._init_state(group, p, state)
+
+                    self._step_one_batch(group, p, state, clipping_scale)
+
+        return loss
+
+    def _init_state(self, group: dict, p: Tensor, state: dict):
+        """
+        Initializes state dict for parameter 'p'.  Assumes that dim 0 of tensor p
+        is actually the batch dimension, corresponding to batched-together
+        parameters of a given shape.
+
+
+        Args:
+           group:   Dict to look up configuration values.
+               p: The parameter that we are initializing the state for
+           state: Dict from string to whatever state we are initializing
+        """
+        size_update_period = group["size_update_period"]
+
+        state["step"] = 0
+
+        kwargs = {"device": p.device, "dtype": p.dtype}
+
+        # 'delta' implements conventional momentum.  There are
+        # several different kinds of update going on, so rather than
+        # compute "exp_avg" like in Adam, we store and decay a
+        # parameter-change "delta", which combines all forms of
+        # update.  this is equivalent to how it's done in Adam,
+        # except for the first few steps.
+        state["delta"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format
+        )
+
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        numel = p.numel()
+
+        if numel > 1:
+            # "param_rms" just periodically records the scalar root-mean-square value of
+            # the parameter tensor.
+            # it has a shape like (batch_size, 1, 1, 1, 1)
+            param_rms = (
+                (p ** 2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()
+            )
+            state["param_rms"] = param_rms
+
+            state["scale_exp_avg_sq"] = torch.zeros_like(param_rms)
+            state["scale_grads"] = torch.zeros(
+                size_update_period, *param_rms.shape, **kwargs
+            )
+
+        # exp_avg_sq is the weighted sum of scaled gradients. as in Adam.
+        state["exp_avg_sq"] = torch.zeros_like(
+            p, memory_format=torch.preserve_format
+        )
+
+    def _get_clipping_scale(
+        self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]]
+    ) -> float:
+        """
+        Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will scale the gradients
+        by this amount before applying the rest of the update.
+
+        Args:
+           group: the parameter group, an item in self.param_groups
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+        """
+        assert len(tuples) >= 1
+        clipping_scale = group["clipping_scale"]
+        (first_p, first_state, _) = tuples[0]
+        step = first_state["step"]
+        if clipping_scale is None or step == 0:
+            # no clipping.  return early on step == 0 because the other
+            # parameters' state won't have been initialized yet.
+            return 1.0
+        clipping_update_period = group["clipping_update_period"]
+
+        tot_sumsq = torch.tensor(0.0, device=first_p.device)
+        for (p, state, param_names) in tuples:
+            grad = p.grad
+            if grad.is_sparse:
+                raise RuntimeError(
+                    "ScaledAdam optimizer does not support sparse gradients"
+                )
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                tot_sumsq += (
+                    grad ** 2
+                ).sum()  # sum() to change shape [1] to []
+            else:
+                tot_sumsq += ((grad * state["param_rms"]) ** 2).sum()
+
+        tot_norm = tot_sumsq.sqrt()
+        if "model_norms" not in first_state:
+            first_state["model_norms"] = torch.zeros(
+                clipping_update_period, device=p.device
+            )
+        first_state["model_norms"][step % clipping_update_period] = tot_norm
+
+        if step % clipping_update_period == 0:
+            # Print some stats.
+            # We don't reach here if step == 0 because we would have returned
+            # above.
+            sorted_norms = first_state["model_norms"].sort()[0].to("cpu")
+            quartiles = []
+            for n in range(0, 5):
+                index = min(
+                    clipping_update_period - 1,
+                    (clipping_update_period // 4) * n,
+                )
+                quartiles.append(sorted_norms[index].item())
+
+            median = quartiles[2]
+            threshold = clipping_scale * median
+            first_state["model_norm_threshold"] = threshold
+            percent_clipped = (
+                first_state["num_clipped"] * 100.0 / clipping_update_period
+                if "num_clipped" in first_state
+                else 0.0
+            )
+            first_state["num_clipped"] = 0
+            quartiles = " ".join(["%.3e" % x for x in quartiles])
+            logging.info(
+                f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, "
+                f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}"
+            )
+
+        if step < clipping_update_period:
+            return 1.0  # We have not yet estimated a norm to clip to.
+        else:
+            try:
+                model_norm_threshold = first_state["model_norm_threshold"]
+            except KeyError:
+                logging.info(
+                    "Warning: model_norm_threshold not in state: possibly "
+                    "you changed config when restarting, adding clipping_scale option?"
+                )
+                return 1.0
+            ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item())
+            if ans < 1.0:
+                first_state["num_clipped"] += 1
+            if ans < 0.1:
+                logging.warn(
+                    f"Scaling gradients by {ans}, model_norm_threshold={model_norm_threshold}"
+                )
+                if self.show_dominant_parameters:
+                    assert p.shape[0] == len(param_names)
+                    self._show_gradient_dominating_parameter(tuples, tot_sumsq)
+            return ans
+
+    def _show_gradient_dominating_parameter(
+        self, tuples: List[Tuple[Tensor, dict, List[str]]], tot_sumsq: Tensor
+    ):
+        """
+        Show information of parameter wihch dominanting tot_sumsq.
+
+        Args:
+           tuples: a list of tuples of (param, state, param_names)
+                where param is a batched set of parameters,
+                with a .grad (1st dim is batch dim)
+                and state is the state-dict where optimization parameters are kept.
+                param_names is a List[str] while each str is name for a parameter
+                in batched set of parameters "param".
+            tot_sumsq: sumsq of all parameters. Though it's could be calculated
+                from tuples, we still pass it to save some time.
+        """
+        all_sumsq_orig = {}
+        for (p, state, batch_param_names) in tuples:
+            # p is a stacked batch parameters.
+            batch_grad = p.grad
+            if p.numel() == p.shape[0]:  # a batch of scalars
+                batch_sumsq_orig = batch_grad ** 2
+                # Dummpy values used by following `zip` statement.
+                batch_rms_orig = torch.ones(p.shape[0])
+            else:
+                batch_rms_orig = state["param_rms"]
+                batch_sumsq_orig = ((batch_grad * batch_rms_orig) ** 2).sum(
+                    dim=list(range(1, batch_grad.ndim))
+                )
+
+            for name, sumsq_orig, rms, grad in zip(
+                batch_param_names, batch_sumsq_orig, batch_rms_orig, batch_grad
+            ):
+
+                proportion_orig = sumsq_orig / tot_sumsq
+                all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad)
+
+        assert torch.isclose(
+            sum([value[0] for value in all_sumsq_orig.values()]).cpu(),
+            torch.tensor(1.0),
+        )
+        sorted_by_proportion = {
+            k: v
+            for k, v in sorted(
+                all_sumsq_orig.items(),
+                key=lambda item: item[1][0],
+                reverse=True,
+            )
+        }
+        dominant_param_name = next(iter(sorted_by_proportion))
+        (
+            dominant_proportion,
+            dominant_sumsq,
+            dominant_rms,
+            dominant_grad,
+        ) = sorted_by_proportion[dominant_param_name]
+        logging.info(
+            f"Parameter Dominanting tot_sumsq {dominant_param_name}"
+            f" with proportion {dominant_proportion:.2f},"
+            f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)"
+            f"={dominant_sumsq:.3e},"
+            f" grad_sumsq = {(dominant_grad**2).sum():.3e},"
+            f" orig_rms_sq={(dominant_rms**2).item():.3e}"
+        )
+
+    def _step_one_batch(
+        self, group: dict, p: Tensor, state: dict, clipping_scale: float
+    ):
+        """
+        Do the step for one parameter, which is actually going to be a batch of
+        `real` parameters, with dim 0 as the batch dim.
+        Args:
+                  group:  dict to look up configuration values
+                    p: parameter to update (actually multiple parameters stacked together
+                       as a batch)
+                  state: state-dict for p, to look up the optimizer state
+        """
+        lr = group["lr"]
+        size_update_period = group["size_update_period"]
+        beta1 = group["betas"][0]
+
+        grad = p.grad
+        if clipping_scale != 1.0:
+            grad = grad * clipping_scale
+        step = state["step"]
+        delta = state["delta"]
+
+        delta.mul_(beta1)
+        batch_size = p.shape[0]
+        numel = p.numel() // batch_size
+        if numel > 1:
+            # Update the size/scale of p, and set param_rms
+            scale_grads = state["scale_grads"]
+            scale_grads[step % size_update_period] = (p * grad).sum(
+                dim=list(range(1, p.ndim)), keepdim=True
+            )
+            if step % size_update_period == size_update_period - 1:
+                param_rms = state["param_rms"]  # shape: (batch_size, 1, 1, ..)
+                param_rms.copy_(
+                    (p ** 2)
+                    .mean(dim=list(range(1, p.ndim)), keepdim=True)
+                    .sqrt()
+                )
+                if step > 0:
+                    # self._size_update() learns the overall scale on the
+                    # parameter, by shrinking or expanding it.
+                    self._size_update(group, scale_grads, p, state)
+
+        if numel == 1:
+            # For parameters with 1 element we just use regular Adam.
+            # Updates delta.
+            self._step_scalar(group, p, state)
+        else:
+            self._step(group, p, state)
+
+        state["step"] = step + 1
+
+    def _size_update(
+        self, group: dict, scale_grads: Tensor, p: Tensor, state: dict
+    ) -> None:
+        """
+               Called only where p.numel() > 1, this updates the scale of the parameter.
+               If we imagine: p =  underlying_param * scale.exp(), and we are doing
+               gradient descent on underlying param and on scale, this function does the update
+               on `scale`.
+
+               Args:
+              group: dict to look up configuration values
+        scale_grads: a tensor of shape (size_update_period, batch_size, 1, 1,...) containing
+                      grads w.r.t. the scales.
+                  p:  The parameter to update
+               state: The state-dict of p
+        """
+
+        param_rms = state["param_rms"]
+        beta1, beta2 = group["betas"]
+        size_lr = group["lr"] * group["scalar_lr_scale"]
+        param_min_rms = group["param_min_rms"]
+        param_max_rms = group["param_max_rms"]
+        eps = group["eps"]
+        step = state["step"]
+        batch_size = p.shape[0]
+
+        size_update_period = scale_grads.shape[0]
+        # correct beta2 for the size update period: we will have
+        # faster decay at this level.
+        beta2_corr = beta2 ** size_update_period
+
+        scale_exp_avg_sq = state[
+            "scale_exp_avg_sq"
+        ]  # shape: (batch_size, 1, 1, ..)
+        scale_exp_avg_sq.mul_(beta2_corr).add_(
+            (scale_grads ** 2).mean(
+                dim=0
+            ),  # mean over dim `size_update_period`
+            alpha=1 - beta2_corr,
+        )  # shape is (batch_size, 1, 1, ...)
+
+        # The 1st time we reach here is when size_step == 1.
+        size_step = (step + 1) // size_update_period
+        bias_correction2 = 1 - beta2_corr ** size_step
+        # we don't bother with bias_correction1; this will help prevent divergence
+        # at the start of training.
+
+        denom = scale_exp_avg_sq.sqrt() + eps
+
+        scale_step = (
+            -size_lr
+            * (bias_correction2 ** 0.5)
+            * scale_grads.sum(dim=0)
+            / denom
+        )
+
+        is_too_small = param_rms < param_min_rms
+        is_too_large = param_rms > param_max_rms
+
+        # when the param gets too small, just don't shrink it any further.
+        scale_step.masked_fill_(is_too_small, 0.0)
+        # when it gets too large, stop it from getting any larger.
+        scale_step.masked_fill_(is_too_large, -size_lr * size_update_period)
+        delta = state["delta"]
+        # the factor of (1-beta1) relates to momentum.
+        delta.add_(p * scale_step, alpha=(1 - beta1))
+
+    def _step(self, group: dict, p: Tensor, state: dict):
+        """
+        This function does the core update of self.step(), in the case where the members of
+        the batch have more than 1 element.
+
+        Args:
+            group: A dict which will be used to look up configuration values
+                p: The parameter to be updated
+             grad: The grad of p
+            state: The state-dict corresponding to parameter p
+
+        This function modifies p.
+        """
+        grad = p.grad
+        lr = group["lr"]
+        beta1, beta2 = group["betas"]
+        eps = group["eps"]
+        param_min_rms = group["param_min_rms"]
+        step = state["step"]
+
+        exp_avg_sq = state["exp_avg_sq"]
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1 - beta2))
+
+        this_step = state["step"] - (
+            state["zero_step"] if "zero_step" in state else 0
+        )
+        bias_correction2 = 1 - beta2 ** (this_step + 1)
+        if bias_correction2 < 0.99:
+            # note: not in-place.
+            exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2)
+
+        denom = exp_avg_sq.sqrt()
+        denom += eps
+        grad = grad / denom
+
+        alpha = -lr * (1 - beta1) * state["param_rms"].clamp(min=param_min_rms)
+
+        delta = state["delta"]
+        delta.add_(grad * alpha)
+        p.add_(delta)
+
+    def _step_scalar(self, group: dict, p: Tensor, state: dict):
+        """
+        A simplified form of the core update for scalar tensors, where we cannot get a good
+        estimate of the parameter rms.
+        """
+        beta1, beta2 = group["betas"]
+        scalar_max = group["scalar_max"]
+        eps = group["eps"]
+        lr = group["lr"] * group["scalar_lr_scale"]
+        grad = p.grad
+
+        exp_avg_sq = state["exp_avg_sq"]  # shape: (batch_size,)
+        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+
+        # bias_correction2 is like in Adam.  Don't bother with bias_correction1;
+        # slower update at the start will help stability anyway.
+        bias_correction2 = 1 - beta2 ** (state["step"] + 1)
+        denom = (exp_avg_sq / bias_correction2).sqrt() + eps
+
+        delta = state["delta"]
+        delta.add_(grad / denom, alpha=-lr * (1 - beta1))
+        p.clamp_(min=-scalar_max, max=scalar_max)
+        p.add_(delta)
+
+
+class LRScheduler(object):
+    """
+    Base-class for learning rate schedulers where the learning-rate depends on both the
+    batch and the epoch.
+    """
+
+    def __init__(self, optimizer: Optimizer, verbose: bool = False):
+        # Attach optimizer
+        if not isinstance(optimizer, Optimizer):
+            raise TypeError(
+                "{} is not an Optimizer".format(type(optimizer).__name__)
+            )
+        self.optimizer = optimizer
+        self.verbose = verbose
+
+        for group in optimizer.param_groups:
+            group.setdefault("base_lr", group["lr"])
+
+        self.base_lrs = [group["base_lr"] for group in optimizer.param_groups]
+
+        self.epoch = 0
+        self.batch = 0
+
+    def state_dict(self):
+        """Returns the state of the scheduler as a :class:`dict`.
+
+        It contains an entry for every variable in self.__dict__ which
+        is not the optimizer.
+        """
+        return {
+            "base_lrs": self.base_lrs,
+            "epoch": self.epoch,
+            "batch": self.batch,
+        }
+
+    def load_state_dict(self, state_dict):
+        """Loads the schedulers state.
+
+        Args:
+            state_dict (dict): scheduler state. Should be an object returned
+                from a call to :meth:`state_dict`.
+        """
+        self.__dict__.update(state_dict)
+
+    def get_last_lr(self) -> List[float]:
+        """Return last computed learning rate by current scheduler.  Will be a list of float."""
+        return self._last_lr
+
+    def get_lr(self):
+        # Compute list of learning rates from self.epoch and self.batch and
+        # self.base_lrs; this must be overloaded by the user.
+        # e.g. return [some_formula(self.batch, self.epoch, base_lr) for base_lr in self.base_lrs ]
+        raise NotImplementedError
+
+    def step_batch(self, batch: Optional[int] = None) -> None:
+        # Step the batch index, or just set it.  If `batch` is specified, it
+        # must be the batch index from the start of training, i.e. summed over
+        # all epochs.
+        # You can call this in any order; if you don't provide 'batch', it should
+        # of course be called once per batch.
+        if batch is not None:
+            self.batch = batch
+        else:
+            self.batch = self.batch + 1
+        self._set_lrs()
+
+    def step_epoch(self, epoch: Optional[int] = None):
+        # Step the epoch index, or just set it.  If you provide the 'epoch' arg,
+        # you should call this at the start of the epoch; if you don't provide the 'epoch'
+        # arg, you should call it at the end of the epoch.
+        if epoch is not None:
+            self.epoch = epoch
+        else:
+            self.epoch = self.epoch + 1
+        self._set_lrs()
+
+    def _set_lrs(self):
+        values = self.get_lr()
+        assert len(values) == len(self.optimizer.param_groups)
+
+        for i, data in enumerate(zip(self.optimizer.param_groups, values)):
+            param_group, lr = data
+            param_group["lr"] = lr
+            self.print_lr(self.verbose, i, lr)
+        self._last_lr = [group["lr"] for group in self.optimizer.param_groups]
+
+    def print_lr(self, is_verbose, group, lr):
+        """Display the current learning rate."""
+        if is_verbose:
+            logging.info(
+                f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate"
+                f" of group {group} to {lr:.4e}."
+            )
+
+
+class Eden(LRScheduler):
+    """
+    Eden scheduler.
+    The basic formula (before warmup) is:
+      lr = base_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 *
+                     (((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25)) * warmup
+    where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches
+    and then stays constant at 1.
+
+
+     E.g. suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam
+
+    Args:
+        optimizer: the optimizer to change the learning rates on
+        lr_batches: the number of batches after which we start significantly
+              decreasing the learning rate, suggest 5000.
+        lr_epochs: the number of epochs after which we start significantly
+              decreasing the learning rate, suggest 6 if you plan to do e.g.
+              20 to 40 epochs, but may need smaller number if dataset is huge
+              and you will do few epochs.
+    """
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        lr_batches: Union[int, float],
+        lr_epochs: Union[int, float],
+        warmup_batches: Union[int, float] = 500.0,
+        verbose: bool = False,
+    ):
+        super(Eden, self).__init__(optimizer, verbose)
+        self.lr_batches = lr_batches
+        self.lr_epochs = lr_epochs
+        self.warmup_batches = warmup_batches
+
+    def get_lr(self):
+        factor = (
+            (self.batch ** 2 + self.lr_batches ** 2) / self.lr_batches ** 2
+        ) ** -0.25 * (
+            ((self.epoch ** 2 + self.lr_epochs ** 2) / self.lr_epochs ** 2)
+            ** -0.25
+        )
+        warmup_factor = (
+            1.0
+            if self.batch >= self.warmup_batches
+            else 0.5 + 0.5 * (self.batch / self.warmup_batches)
+        )
+
+        return [x * factor * warmup_factor for x in self.base_lrs]
+
+
+def _test_eden():
+    m = torch.nn.Linear(100, 100)
+    optim = ScaledAdam(m.parameters(), lr=0.03)
+
+    scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True)
+
+    for epoch in range(10):
+        scheduler.step_epoch(epoch)  # sets epoch to `epoch`
+
+        for step in range(20):
+            x = torch.randn(200, 100).detach()
+            x.requires_grad = True
+            y = m(x)
+            dy = torch.randn(200, 100).detach()
+            f = (y * dy).sum()
+            f.backward()
+
+            optim.step()
+            scheduler.step_batch()
+            optim.zero_grad()
+
+    logging.info(f"last lr = {scheduler.get_last_lr()}")
+    logging.info(f"state dict = {scheduler.state_dict()}")
+
+
+# This is included mostly as a baseline for ScaledAdam.
+class Eve(Optimizer):
+    """
+    Implements Eve algorithm.  This is a modified version of AdamW with a special
+    way of setting the weight-decay / shrinkage-factor, which is designed to make the
+    rms of the parameters approach a particular target_rms (default: 0.1).  This is
+    for use with networks with 'scaled' versions of modules (see scaling.py), which
+    will be close to invariant to the absolute scale on the parameter matrix.
+
+    The original Adam algorithm was proposed in `Adam: A Method for Stochastic Optimization`_.
+    The AdamW variant was proposed in `Decoupled Weight Decay Regularization`_.
+    Eve is unpublished so far.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay coefficient (default: 3e-4;
+            this value means that the weight would decay significantly after
+            about 3k minibatches.  Is not multiplied by learning rate, but
+            is conditional on RMS-value of parameter being > target_rms.
+        target_rms (float, optional): target root-mean-square value of
+           parameters, if they fall below this we will stop applying weight decay.
+
+
+    .. _Adam: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _Decoupled Weight Decay Regularization:
+        https://arxiv.org/abs/1711.05101
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.9, 0.98),
+        eps=1e-8,
+        weight_decay=1e-3,
+        target_rms=0.1,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 0: {}".format(betas[0])
+            )
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(
+                "Invalid beta parameter at index 1: {}".format(betas[1])
+            )
+        if not 0 <= weight_decay <= 0.1:
+            raise ValueError(
+                "Invalid weight_decay value: {}".format(weight_decay)
+            )
+        if not 0 < target_rms <= 10.0:
+            raise ValueError("Invalid target_rms value: {}".format(target_rms))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            target_rms=target_rms,
+        )
+        super(Eve, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Eve, self).__setstate__(state)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                # Perform optimization step
+                grad = p.grad
+                if grad.is_sparse:
+                    raise RuntimeError(
+                        "AdamW does not support sparse gradients"
+                    )
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros_like(
+                        p, memory_format=torch.preserve_format
+                    )
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+                bias_correction1 = 1 - beta1 ** state["step"]
+                bias_correction2 = 1 - beta2 ** state["step"]
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
+                denom = (exp_avg_sq.sqrt() * (bias_correction2 ** -0.5)).add_(
+                    group["eps"]
+                )
+
+                step_size = group["lr"] / bias_correction1
+                target_rms = group["target_rms"]
+                weight_decay = group["weight_decay"]
+
+                if p.numel() > 1:
+                    # avoid applying this weight-decay on "scaling factors"
+                    # (which are scalar).
+                    is_above_target_rms = p.norm() > (
+                        target_rms * (p.numel() ** 0.5)
+                    )
+                    p.mul_(1 - (weight_decay * is_above_target_rms))
+
+                p.addcdiv_(exp_avg, denom, value=-step_size)
+
+                # if random.random() < 0.0005:
+                #     step = (exp_avg / denom) * step_size
+                #     logging.info(
+                #         f"Delta rms = {(step**2).mean().item()}, shape = {step.shape}"
+                #     )
+
+        return loss
+
+def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear:
+    """
+    Behaves like a constructor of a modified version of nn.Linear
+    that gives an easy way to set the default initial parameter scale.
+
+    Args:
+        Accepts the standard args and kwargs that nn.Linear accepts
+        e.g. in_features, out_features, bias=False.
+
+        initial_scale: you can override this if you want to increase
+           or decrease the initial magnitude of the module's output
+           (affects the initialization of weight_scale and bias_scale).
+           Another option, if you want to do something like this, is
+           to re-initialize the parameters.
+    """
+    ans = nn.Linear(*args, **kwargs)
+    with torch.no_grad():
+        ans.weight[:] *= initial_scale
+        if ans.bias is not None:
+            torch.nn.init.uniform_(
+                ans.bias, -0.1 * initial_scale, 0.1 * initial_scale
+            )
+    return ans
+def _test_scaled_adam(hidden_dim: int):
+    import timeit
+
+    E = 100
+    B = 4
+    T = 2
+    logging.info("in test_eve_cain")
+    # device = torch.device('cuda')
+    device = torch.device("cpu")
+    dtype = torch.float32
+
+    # these input_magnitudes and output_magnitudes are to test that
+    # Abel is working as we expect and is able to adjust scales of
+    # different dims differently.
+    input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+    output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp()
+
+    for iter in [1, 0]:
+        Linear = torch.nn.Linear if iter == 0 else ScaledLinear
+
+        m = torch.nn.Sequential(
+            Linear(E, hidden_dim),
+            torch.nn.PReLU(),
+            Linear(hidden_dim, hidden_dim),
+            torch.nn.PReLU(),
+            Linear(hidden_dim, E),
+        ).to(device)
+
+        train_pairs = [
+            (
+                100.0
+                * torch.randn(B, T, E, device=device, dtype=dtype)
+                * input_magnitudes,
+                torch.randn(B, T, E, device=device, dtype=dtype)
+                * output_magnitudes,
+            )
+            for _ in range(20)
+        ]
+
+        if iter == 0:
+            optim = Eve(m.parameters(), lr=0.003)
+        elif iter == 1:
+            optim = ScaledAdam(m.parameters(), lr=0.03, clipping_scale=2.0)
+        scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False)
+
+        start = timeit.default_timer()
+        avg_loss = 0.0
+        for epoch in range(180):
+            scheduler.step_epoch()
+            # if epoch == 100 and iter in [2,3]:
+            #    optim.reset_speedup()  # check it doesn't crash.
+
+            # if epoch == 130:
+            #    opts = diagnostics.TensorDiagnosticOptions(
+            #        2 ** 22
+            #    )  # allow 4 megabytes per sub-module
+            #    diagnostic = diagnostics.attach_diagnostics(m, opts)
+
+            for n, (x, y) in enumerate(train_pairs):
+                y_out = m(x)
+                loss = ((y_out - y) ** 2).mean() * 100.0
+                if epoch == 0 and n == 0:
+                    avg_loss = loss.item()
+                else:
+                    avg_loss = 0.98 * avg_loss + 0.02 * loss.item()
+                if n == 0 and epoch % 5 == 0:
+                    # norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item()
+                    # norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item()
+                    # norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item()
+                    # norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item()
+                    # scale1 = '%.2e' % (m[0].weight_scale.exp().item())
+                    # scale1b = '%.2e' % (m[0].bias_scale.exp().item())
+                    # scale2 = '%.2e' % (m[2].weight_scale.exp().item())
+                    # scale2b = '%.2e' % (m[2].bias_scale.exp().item())
+                    lr = scheduler.get_last_lr()[0]
+                    logging.info(
+                        f"Iter {iter}, epoch {epoch}, batch {n}, avg_loss {avg_loss:.4g}, lr={lr:.4e}"
+                    )  # , norms={norm1,norm1b,norm2,norm2b}") # scales={scale1,scale1b,scale2,scale2b}
+                loss.log().backward()
+                optim.step()
+                optim.zero_grad()
+                scheduler.step_batch()
+
+        # diagnostic.print_diagnostics()
+
+        stop = timeit.default_timer()
+        logging.info(f"Iter={iter}, Time taken: {stop - start}")
+
+        logging.info(f"last lr = {scheduler.get_last_lr()}")
+        # logging.info("state dict = ", scheduler.state_dict())
+        # logging.info("optim state_dict = ", optim.state_dict())
+        logging.info(f"input_magnitudes = {input_magnitudes}")
+        logging.info(f"output_magnitudes = {output_magnitudes}")
+
+
+if __name__ == "__main__":
+    torch.set_num_threads(1)
+    torch.set_num_interop_threads(1)
+    logging.getLogger().setLevel(logging.INFO)
+    import subprocess
+
+    s = subprocess.check_output(
+        "git status -uno .; git log -1; git diff HEAD .", shell=True
+    )
+    logging.info(s)
+    import sys
+
+    if len(sys.argv) > 1:
+        hidden_dim = int(sys.argv[1])
+    else:
+        hidden_dim = 200
+
+    _test_scaled_adam(hidden_dim)
+    _test_eden()

steps/trainer.py

@@ -0,0 +1,717 @@
+import time, sys, subprocess, json, re
+from pathlib import Path
+import os, random
+import torch
+import math, pickle
+from tqdm import tqdm
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
+import torch.nn as nn
+import torch.distributed as dist
+from torch.utils.data.sampler import Sampler
+import copy
+from torch.utils.tensorboard import SummaryWriter
+import numpy as np
+from torch.utils.data.distributed import DistributedSampler
+import logging
+# from data import librilight, gigaspeech, gigaspeech_waveform
+from data import combined_dataset
+from models import voice_star
+
+from .trainer_utils import DistributedDynamicBatchSampler, StatefulDistributedSampler, StatefulSampler, AverageMeter, print_model_info
+from .optim import ScaledAdam, Eden
+import run_gen
+import wandb, socket
+
+class Trainer:
+    
+    def __init__(self, args, world_size, rank, local_rank):
+        self.start_time = time.time()
+        self.args = args
+        if self.args.val_max_num_tokens == None:
+            self.args.val_max_num_tokens = self.args.max_num_tokens 
+        self.world_size, self.rank, self.local_rank = world_size, rank, local_rank
+        self.device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
+        if self.rank == 0:
+            self.writer = SummaryWriter(args.exp_dir)
+            self.wandb = wandb.init(project="voice_editor", name=args.exp_dir.split("/")[-1], config=args, dir=args.exp_dir, entity=self.args.wandb_entity)
+        self.seed_everything(seed=self.args.seed)
+        self.meters = self._setup_meters()
+
+        self.progress, self.total_progress = self._setup_progress()
+
+        self.model, self.trainables, self.optim_states, self.scheduler_states, self.phn2num = self._setup_models()
+
+        self.train_dataset_length, self.train_sampler, self.train_loader, self.valid_loader = self._setup_dataloader() # both are use DistributedSampler, train sampler is stateful
+        if self.args.num_steps != None:
+            self.total_step = self.args.num_steps
+            self.args.num_epochs = math.ceil(self.total_step / math.floor(self.train_dataset_length / self.args.batch_size)) if not self.args.dynamic_batching else None
+        else:
+            self.total_step = int(math.floor(self.train_dataset_length / self.args.batch_size))*self.args.num_epochs
+
+        self.optimizer, self.scheduler = self._setup_optimizer()
+        self.scaler = torch.cuda.amp.GradScaler()
+        self.model = torch.nn.parallel.DistributedDataParallel(self.model, device_ids=[self.local_rank], find_unused_parameters=False)
+        self.early_stop_accu_steps = 0
+        if self.rank == 0:
+            if self.args.dynamic_batching:
+                logging.info(f"max number of tokens per GPU in a training batch: {self.args.max_num_tokens}, max number of tokens per GPU in a inference batch: {self.args.val_max_num_tokens}")
+            else:
+                logging.info(f"batch size (per gpu): {self.args.batch_size}")
+        
+        self.args.inference_every_n_steps = getattr(self.args, "inference_every_n_steps", self.args.val_every_n_steps*5)
+        assert self.args.inference_every_n_steps > self.args.val_every_n_steps and self.args.inference_every_n_steps % self.args.val_every_n_steps == 0, "inference_every_n_steps should be divisible by val_every_n_steps, otherwise the code will not get a chance to run inference"
+
+    def train(self):
+        flag = True
+        skip_flag = False
+        data_start_time = time.time()
+        if self.progress['step'] >= self.total_step:
+            if self.rank == 0:
+                self.writer.close()
+                self.wandb.finish()
+            return
+        while flag:
+            self.train_sampler.set_epoch(self.progress['epoch'])
+            for i, batch in enumerate(self.train_loader):
+                if len(batch['y_lens']) < self.args.gradient_accumulation_steps:
+                    continue
+                data_end_time = time.time()
+                self.model.train()
+                if self.progress['step'] >= getattr(self.args, "uniform_weight_start_step", 1e50):
+                    if self.progress['step'] == getattr(self.args, "uniform_weight_start_step", 1e50) and self.rank == 0:
+                        logging.info("NOTE: start using uniform weight from step: {}".format(self.progress['step']))
+                    self.args.codebook_weight = [2.5,2,1.5,0.6]
+                    self.model.module.args.codebook_weight = [2.5,2,1.5,0.6]
+
+                if self.progress['step'] >= self.total_step:
+                    dist.barrier()
+                    flag = False
+                    self.validate_and_save()
+                    if self.rank == 0:
+                        self.writer.close()
+                        self.wandb.finish()
+                    break
+                if isinstance(self.scheduler, Eden):
+                    self.scheduler.step_epoch(self.progress['step']//self.args.pseudo_epoch_size + 1)
+                if self.args.optimizer_name == "ScaledAdam":
+                    cur_lr = self.scheduler.get_last_lr()[0]
+                else:
+                    lrs = [param_group['lr'] for param_group in self.optimizer.param_groups]
+                    assert lrs[0] == lrs[1]
+                    cur_lr = lrs[0]
+
+                if self.rank == 0 and self.progress['step'] % self.args.tb_write_every_n_steps == 0:
+                    self.writer.add_scalar("train/lr", cur_lr, self.progress['step'])
+                    self.wandb.log({"train/lr": cur_lr}, step=self.progress['step'])
+
+                all_inds = list(range(len(batch['y'])))
+                sum_losses = 0
+                sum_top10acc = 0
+                sum_ntoken = 0
+                sum_top10acc_cbi = [0 for _ in range(self.args.n_codebooks)]
+                # extra losses
+                sum_extra_losses = {}
+                # when using prompt-based training, it's likely that due to prompt, the total length gets much longer, which make effective batch size in each accumulation step much bigger and then lead to OOM.
+                # therefore we re-calculate graduent_accumulation_steps based on the effective batch size
+
+                if self.args.neighbor_prompt_prob > 0:
+                    effective_batch_size = self.args.max_num_tokens // self.args.gradient_accumulation_steps
+                    total_batch_size = sum(batch['y_lens']).item()
+                    cur_gradient_accumulation_steps = max(self.args.gradient_accumulation_steps, total_batch_size // effective_batch_size)
+                    gas = torch.tensor(cur_gradient_accumulation_steps, dtype=torch.int, device=self.local_rank)
+                    dist.all_reduce(gas, op=dist.ReduceOp.MAX)
+                    cur_gradient_accumulation_steps = gas.item()
+                    len_batch = torch.tensor(len(batch['y']), dtype=torch.int, device=self.local_rank)
+                    dist.all_reduce(len_batch, op=dist.ReduceOp.MIN)
+                    len_batch = len_batch.item()
+                    cur_gradient_accumulation_steps = min(cur_gradient_accumulation_steps, len_batch)
+                    # for those that cur_gradient_accumulation_steps * effective_batch_size < total_batch_size, we only use the first cur_gradient_accumulation_steps * effective_batch_size samples
+                    cur_len = 0
+                    final_all_inds = []
+                    pointer = 0
+                    while cur_len < self.args.max_num_tokens and pointer < len(all_inds):
+                        cur_len += batch['y_lens'][pointer]
+                        final_all_inds.append(all_inds[pointer])
+                        pointer += 1
+                    all_inds = final_all_inds
+                else:
+                    cur_gradient_accumulation_steps = self.args.gradient_accumulation_steps
+                
+
+                sum_losses_local = 0.0
+                sum_top10acc_local = 0.0
+                sum_entropy_loss_local = 0.0
+                sum_ctc_loss_local = 0.0
+                sum_ntoken_local = 0.0
+                sum_top10acc_cbi_local = [0.0 for _ in range(self.args.n_codebooks)]
+
+                global_nan_flag = 0
+                for j in range(cur_gradient_accumulation_steps):
+                    cur_ind = all_inds[j::cur_gradient_accumulation_steps]
+                    cur_batch = {key: batch[key][cur_ind] for key in batch}
+
+                    # Automatic casting
+                    if self.args.precision == "float16":
+                        precision_used = torch.float16
+                    elif self.args.precision in ["bf16", "bfloat16"]:
+                        precision_used = torch.bfloat16
+                    else:
+                        precision_used = torch.float32
+
+                    with torch.amp.autocast('cuda', dtype=precision_used):
+                        out = self.model(cur_batch, calc_loss=True)
+                    if out is None:
+                        continue
+
+                    if torch.isnan(out['loss']).any():
+                        local_nan_flag = torch.tensor(1, device=self.local_rank)
+                    else:
+                        local_nan_flag = torch.tensor(0, device=self.local_rank)
+
+                    # All ranks check if *any* rank got a NaN
+                    dist.all_reduce(local_nan_flag, op=dist.ReduceOp.SUM)
+                    global_nan_flag = local_nan_flag.item()
+                    if global_nan_flag > 0:
+                        # Now *all* ranks break at the same j
+                        logging.info(f"rank: {self.rank}. Loss at micro-batch {j} in step {self.progress['step']} was NaN on at least one rank; skipping.")
+                        break
+
+                    # Accumulate local values
+                    record_loss = out['loss'].detach()
+                    top10acc = out['top10acc'].detach()
+                    effective_ntoken = out['effective_ntoken'].detach()
+
+                    sum_losses_local += record_loss.item()
+                    sum_top10acc_local += top10acc.item()
+                    sum_ntoken_local += effective_ntoken.item()
+
+                    # Optional losses
+                    if 'entropy_loss' in out:
+                        sum_entropy_loss_local += out['entropy_loss'].detach().item()
+                    if 'ctc_loss' in out:
+                        sum_ctc_loss_local += out['ctc_loss'].detach().item()
+
+                    # Codebook accuracy
+                    if 'top10acc_by_codebook' in out:
+                        for cb in range(self.args.n_codebooks):
+                            sum_top10acc_cbi_local[cb] += out['top10acc_by_codebook'][cb].detach().item()
+
+                    # Backprop on this micro-batch
+                    if self.args.optimizer_name == "ScaledAdam":
+                        self.scaler.scale(out['loss']).backward()
+                    else:
+                        self.scaler.scale(out['loss'] / out['effective_ntoken']).backward()
+
+                if global_nan_flag > 0:
+                    # If *any* rank had NaN, skip this step
+                    logging.info(f"rank: {self.rank}. Loss at one micro-batch in step {self.progress['step']} was NaN on at least one rank; skipping.")
+                    self.progress['step'] += 1
+                    self.progress['cur_step'] += 1
+                    self.optimizer.zero_grad()
+                    continue
+
+                # Otherwise, do one big reduce for the summed metrics
+                metrics_tensor = torch.tensor([
+                    sum_losses_local,
+                    sum_top10acc_local,
+                    sum_entropy_loss_local,
+                    sum_ctc_loss_local,
+                    sum_ntoken_local
+                ], device=self.local_rank, dtype=torch.float32)
+
+                dist.all_reduce(metrics_tensor, op=dist.ReduceOp.SUM)
+
+                # Also reduce the codebook array in one shot if needed
+                codebook_tensor = torch.tensor(sum_top10acc_cbi_local, device=self.local_rank, dtype=torch.float32)
+                dist.all_reduce(codebook_tensor, op=dist.ReduceOp.SUM)
+
+                # Convert them back to Python scalars
+                sum_losses = metrics_tensor[0].item()
+                sum_top10acc = metrics_tensor[1].item()
+                sum_entropy_loss = metrics_tensor[2].item()
+                sum_ctc_loss = metrics_tensor[3].item()
+                sum_ntoken = metrics_tensor[4].item()
+
+                sum_top10acc_cbi = codebook_tensor.tolist()
+
+                if self.args.optimizer_name != "ScaledAdam":
+                    self.scaler.unscale_(self.optimizer)
+                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.gradient_clip_val)
+
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+                self.optimizer.zero_grad()
+
+                if self.args.optimizer_name == "ScaledAdam":
+                    self.scheduler.step_batch(self.progress['step'])
+                else:
+                    self.scheduler.step()
+                
+                # logging
+                if self.rank == 0:
+                    average_loss = sum_losses / sum_ntoken
+                    average_top10acc = sum_top10acc / sum_ntoken
+                    average_top10acc_cbi = [sum_top10acc_cbi[cb] / sum_ntoken * self.args.n_codebooks for cb in range(self.args.n_codebooks)]
+                    self.meters['train_loss'].update(average_loss, batch['x'].shape[0]*self.world_size)
+                    self.meters['train_top10acc'].update(average_top10acc, batch['x'].shape[0]*self.world_size)
+                    self.meters['train_top10acc'].update(average_top10acc, batch['x'].shape[0]*self.world_size)
+                    for cb in range(self.args.n_codebooks):
+                        self.meters[f'train_top10acc_cb{cb+1}'].update(average_top10acc_cbi[cb], batch['x'].shape[0]*self.world_size)
+                    self.meters['data_time'].update(data_end_time - data_start_time)
+                    self.meters['train_time'].update(time.time() - data_end_time)
+
+                    # log extra losses
+                    for key in sum_extra_losses:
+                        if "train_"+key not in self.meters:
+                            self.meters["train_"+key] = AverageMeter()
+                        self.meters["train_"+key].update(sum(sum_extra_losses[key])/len(sum_extra_losses[key]), batch['x'].shape[0]*self.world_size)
+
+                    if self.progress['step'] % self.args.tb_write_every_n_steps == 0:
+                        self.writer.add_scalar('train/loss', average_loss, self.progress['step'])
+                        self.writer.add_scalar('train/top10acc', average_top10acc, self.progress['step'])
+                        self.writer.add_scalar("train/ntokens", sum_ntoken, self.progress['step'])
+                        self.wandb.log({"train/loss": average_loss, "train/top10acc": average_top10acc, "train/ntokens": sum_ntoken, "train/data_time": data_end_time - data_start_time, "train/train_time": time.time() - data_end_time}, step=self.progress['step'])
+
+                        for cb in range(self.args.n_codebooks):
+                            self.writer.add_scalar(f'train/top10acc_cb{cb+1}', average_top10acc_cbi[cb], self.progress['step'])
+                            self.wandb.log({f'train/top10acc_cb{cb+1}': average_top10acc_cbi[cb]}, step=self.progress['step'])
+                        self.writer.add_scalar("train/data_time", data_end_time - data_start_time, self.progress['step'])
+                        self.writer.add_scalar("train/train_time", time.time() - data_end_time, self.progress['step'])
+                        # write extra losses
+                        for key in sum_extra_losses:
+                            self.writer.add_scalar(f"train/{key}", sum(sum_extra_losses[key])/len(sum_extra_losses[key]), self.progress['step'])
+                            self.wandb.log({f"train/{key}": sum(sum_extra_losses[key])/len(sum_extra_losses[key])}, step=self.progress['step'])
+                    # logging.info(f"ntoken: {sum_ntoken}")
+                        
+                    # logging
+                    if self.progress['step'] % self.args.print_every_n_steps == 0:
+                        log_out = {}
+                        log_out['cur_epoch'] = f"{self.progress['epoch']}/{self.args.num_epochs}" if self.args.num_epochs is not None else f"{self.progress['epoch']}"
+                        log_out['cur_step'] = f"{int(self.progress['cur_step']+1)}"
+                        log_out['total_step'] = f"{self.progress['step']}/{self.args.num_steps}"
+                        log_out['lr'] = f"{cur_lr:.7f}"
+                        log_out['ntokens'] = f"{sum_ntoken}"
+                        for key in self.meters:
+                            if self.meters[key].val != 0 or self.meters[key].avg != 0:
+                                log_out[key] = f"{self.meters[key].val:.4f} ({self.meters[key].avg:.4f})" if isinstance(self.meters[key].val, float) else f"{self.meters[key].val}"
+                        logging.info(log_out)
+                        if np.isnan(self.meters['train_loss'].avg):
+                            logging.warning("training diverged...")
+                            raise RuntimeError("training diverged...")
+
+                # save the model only
+                if self.progress['step'] % self.args.save_every_n_steps == 0:
+                    dist.barrier()
+                    if self.rank == 0:
+                        save_path = os.path.join(self.args.exp_dir,f"bundle_step{self.progress['step']}.pth")
+                        self.save_progress(name=f"step{self.progress['step']}")
+                        torch.save(
+                            {
+                                "model": self.model.module.state_dict(),
+                                "args": self.args,
+                                "phn2num": self.train_loader.dataset.phn2num,
+                                "optimizer":  self.optimizer.state_dict(),
+                                "scheduler": self.scheduler.state_dict(),
+                            },save_path
+                        )
+                        logging.info(f"save model, optimizer, scheduler and progress at {save_path} at global step {self.progress['step']}")
+                    dist.barrier()
+                
+                # validation and save models
+                if self.progress['step'] % self.args.val_every_n_steps == 0:
+                    dist.barrier()
+                    continue_training = self.validate_and_save()
+                    # broadcast continue_training to all processes, so that all processes gets into generation stage
+                    continue_training = torch.tensor(int(continue_training), dtype=torch.int, device=self.local_rank)
+                    dist.broadcast(continue_training, src=0)
+                    continue_training = bool(continue_training.item())
+                    dist.barrier() # need this to ensure all processes get to the next line?
+                    logging.info(f"rank: {self.rank}, continue_training: {continue_training}")
+                    if not continue_training: 
+                        if self.rank == 0:
+                            self.writer.close()
+                            self.wandb.finish()
+                        flag = False
+                        break
+
+                self.progress['step'] += 1
+                self.progress['cur_step'] += 1
+
+                data_start_time = time.time()
+            self.progress['epoch'] += 1
+            self.progress['cur_step'] = 0 # reset cur_step to be 0
+        dist.destroy_process_group()
+
+    def validate_and_save(self):
+        self.model.eval()
+        
+        score = self.validate(self.valid_loader)
+
+        if self.args.early_stop_threshold > 0:
+            if self.progress['best_score'] - score < self.args.early_stop_threshold:
+                self.early_stop_accu_steps += self.args.val_every_n_steps
+                if self.early_stop_accu_steps >= self.args.early_stop_step-1:
+                    logging.info(f"early stop based on self.args.early_stop_threshold: {self.args.early_stop_threshold}, and self.args.early_stop_step: {self.args.early_stop_step}")
+                    logging.info(f"best validation score at step: {self.progress['best_step']}, and the score is {self.progress['best_score']:.4f}")
+                    return False
+            else:
+                self.early_stop_accu_steps = 0
+
+        if self.rank == 0:
+            save_path = os.path.join(self.args.exp_dir,"bundle.pth")
+            if os.path.isfile(save_path):
+                os.system(f"mv {save_path} {save_path.replace('.pth', '_prev.pth')}")
+            torch.save(
+                {
+                    "model": self.model.module.state_dict(),
+                    "optimizer":  self.optimizer.state_dict(),
+                    "scheduler": self.scheduler.state_dict(),
+                    "args": self.args,
+                    "phn2num": self.train_loader.dataset.phn2num
+                    },save_path
+            )
+            self.save_progress()
+            logging.info(f"save models, indices, acc and other statistics at {save_path} and {self.args.exp_dir}/progress.pkl at global step {self.progress['step']}")
+            if (score < self.progress['best_score']):
+                self.progress['best_step'] = self.progress['step']
+                self.progress['best_score'] = score
+                save_path = os.path.join(self.args.exp_dir,"best_bundle.pth")
+                if os.path.isfile(save_path):
+                    os.system(f"mv {save_path} {save_path.replace('.pth', '_prev.pth')}")
+                torch.save(
+                    {
+                        "model": self.model.module.state_dict(),
+                        "optimizer":  self.optimizer.state_dict(),
+                        "scheduler": self.scheduler.state_dict(),
+                        "args": self.args,
+                        "phn2num": self.train_loader.dataset.phn2num
+                    },save_path
+                )
+                logging.info(f"save *best* models at {save_path} at global step {self.progress['step']}")
+        
+        # sync best score and best step, so that all processes early stop at the same time
+        best_score_tensor = torch.tensor(self.progress['best_score'], device=self.local_rank)
+        dist.broadcast(best_score_tensor, src=0)
+        self.progress['best_score'] = float(best_score_tensor.item())
+        best_step_tensor = torch.tensor(self.progress['best_step'], device=self.local_rank)
+        dist.broadcast(best_step_tensor, src=0)
+        self.progress['best_step'] = int(best_step_tensor.item())
+        dist.barrier()
+        return True
+
+    def validate(self, valid_loader=None, hide_progress=True):
+        if valid_loader == None:
+            valid_loader = self.valid_loader
+        self.model.eval()
+
+        start_val_time = time.time()
+        sum_losses = 0
+        sum_top10acc = 0
+        sum_ntoken = 0
+        sum_dur_loss = 0
+        sum_dur_acc = 0
+        sum_entropy_loss = 0
+        sum_ctc_loss = 0
+
+        sum_top10acc_cbi = [0 for _ in range(self.args.n_codebooks)]
+        mean_perplexity_cbi = [0 for _ in range(self.args.n_codebooks)]
+
+        with torch.no_grad():
+            for i, batch in enumerate(tqdm(valid_loader, disable=hide_progress)):
+                out = self.model(batch, calc_loss=True) # no reduction is applied to loss
+                sum_losses += out['loss']
+                sum_top10acc += out['top10acc']
+                sum_ntoken += out['effective_ntoken']
+                if "dur_loss" in out:
+                    sum_dur_loss += out['dur_loss']
+                    sum_dur_acc += out['dur_acc']
+                if "entropy_loss" in out:
+                    sum_entropy_loss += out['entropy_loss']
+                if "ctc_loss" in out:
+                    sum_ctc_loss += out['ctc_loss']
+                # logging.info(f"iter {i}::: {sum_losses}, {sum_top10acc}, {sum_ntoken}")
+                if 'top10acc_by_codebook' in out:
+                    for cb in range(self.args.n_codebooks):
+                        sum_top10acc_cbi[cb] += out['top10acc_by_codebook'][cb]
+                        
+                if 'perplexity_by_codebook' in out:
+                    for cb in range(self.args.n_codebooks):
+                        mean_perplexity_cbi[cb] += out['perplexity_by_codebook'][cb]
+                # if i > 10:
+                #     break
+
+        
+        dist.all_reduce(sum_losses, op=dist.ReduceOp.SUM)
+        dist.all_reduce(sum_top10acc, op=dist.ReduceOp.SUM)
+        dist.all_reduce(sum_ntoken, op=dist.ReduceOp.SUM)
+        if "dur_loss" in out:
+            dist.all_reduce(sum_dur_loss, op=dist.ReduceOp.SUM)
+            dist.all_reduce(sum_dur_acc, op=dist.ReduceOp.SUM)
+        if "entropy_loss" in out:
+            dist.all_reduce(sum_entropy_loss, op=dist.ReduceOp.SUM)
+        if "ctc_loss" in out:
+            dist.all_reduce(sum_ctc_loss, op=dist.ReduceOp.SUM)
+
+        if 'top10acc_by_codebook' in out:
+            for cb in range(self.args.n_codebooks):
+                dist.all_reduce(sum_top10acc_cbi[cb], op=dist.ReduceOp.SUM)
+        
+        if 'perplexity_by_codebook' in out:
+            for cb in range(self.args.n_codebooks):
+                dist.all_reduce(mean_perplexity_cbi[cb], op=dist.ReduceOp.SUM)
+        
+        val_loss = sum_losses / sum_ntoken
+        val_top10acc = sum_top10acc / sum_ntoken
+        
+        if self.rank == 0:
+            if "dur_loss" in out:
+                val_dur_loss = sum_dur_loss / sum_ntoken
+                val_dur_acc = sum_dur_acc / sum_ntoken
+                self.meters['val_dur_loss'].update(val_dur_loss)
+                logging.info(f"val dur_loss: {val_dur_loss:.5f}")
+                self.meters['val_dur_acc'].update(val_dur_acc)
+                logging.info(f"val dur_acc: {val_dur_acc:.5f}")
+                self.writer.add_scalar("val/dur_loss", val_dur_loss, self.progress['step'])
+                self.writer.add_scalar("val/dur_acc", val_dur_acc, self.progress['step'])
+                self.wandb.log({"val/dur_loss": val_dur_loss, "val/dur_acc": val_dur_acc}, step=self.progress['step'])
+            # logging
+            self.meters['val_loss'].update(val_loss)
+            logging.info(f"val loss: {val_loss:.5f}")
+            self.writer.add_scalar("val/loss", val_loss, self.progress['step'])
+            self.wandb.log({"val/loss": val_loss}, step=self.progress['step'])
+
+            self.meters['val_top10acc'].update(val_top10acc)
+            logging.info(f"val top10acc: {val_top10acc:.5f}")
+            self.writer.add_scalar("val/top10acc", val_top10acc, self.progress['step'])
+            self.wandb.log({"val/top10acc": val_top10acc}, step=self.progress['step'])
+            for cb in range(self.args.n_codebooks):
+                average_top10acc_cbi = sum_top10acc_cbi[cb] / sum_ntoken * self.args.n_codebooks
+                self.meters[f'val_top10acc_cb{cb+1}'].update(average_top10acc_cbi)
+                self.writer.add_scalar(f'val/top10acc_cb{cb+1}', average_top10acc_cbi, self.progress['step'])
+                self.wandb.log({f'val/top10acc_cb{cb+1}': average_top10acc_cbi}, step=self.progress['step'])
+
+                temp = mean_perplexity_cbi[cb]/len(valid_loader)
+                self.writer.add_scalar(f'val/perplexity_cb{cb+1}', temp, self.progress['step'])
+                self.wandb.log({f'val/perplexity_cb{cb+1}': temp}, step=self.progress['step'])
+                
+            average_perplexity = sum(mean_perplexity_cbi)/(self.args.n_codebooks*len(valid_loader))
+            self.wandb.log({"val/average_perplexity": average_perplexity}, step=self.progress['step'])
+            self.writer.add_scalar('val/average_perplexity', average_perplexity, self.progress['step'])
+            
+            # log entropy and ctc loss
+            if "entropy_loss" in out:
+                val_entropy_loss = sum_entropy_loss / ((i+1) * self.world_size)
+                self.meters['val_entropy_loss'].update(val_entropy_loss)
+                logging.info(f"val entropy_loss: {val_entropy_loss:.5f}")
+                self.writer.add_scalar("val/entropy_loss", val_entropy_loss, self.progress['step'])
+                self.wandb.log({"val/entropy_loss": val_entropy_loss}, step=self.progress['step'])
+            if "ctc_loss" in out:
+                val_ctc_loss = sum_ctc_loss / ((i+1) * self.world_size)
+                self.meters['val_ctc_loss'].update(val_ctc_loss)
+                logging.info(f"val ctc_loss: {val_ctc_loss:.5f}")
+                self.writer.add_scalar("val/ctc_loss", val_ctc_loss, self.progress['step'])
+                self.wandb.log({"val/ctc_loss": val_ctc_loss}, step=self.progress['step'])
+
+            logging.info(f"validation takes: {time.time() - start_val_time:.2f}s")
+            logging.info(f"Step [{self.progress['step']}/{self.total_step}]\t Time elapsed {(time.time() - self.start_time)/3600.:.2f}h, Val Loss: {val_loss:.4f}, Val Top10Acc: {val_top10acc:.4f}")
+
+        return val_loss.item()
+
+    def _setup_meters(self):
+        meters = {}
+        meter_names = ['train_loss', 'val_loss', 'train_top10acc', 'val_top10acc', 'data_time', 'train_time']
+        meter_names += ['train_dur_loss', 'train_dur_acc', 'val_dur_loss', 'val_dur_acc']
+        meter_names += ['val_perplexity']
+        meter_names += [f'train_top10acc_cb{cb+1}' for cb in range(self.args.n_codebooks)]
+        meter_names += [f'val_top10acc_cb{cb+1}' for cb in range(self.args.n_codebooks)]
+        meter_names += [f'val_perplexity_cb{cb+1}' for cb in range(self.args.n_codebooks)]
+        for name in meter_names:
+            meters[name] = AverageMeter()
+        return meters
+    def _setup_progress(self):
+        """
+        Need to customize it
+        """
+        progress = {}
+        progress['best_step'] = 1
+        progress['best_score'] = np.inf # this records loss value
+        progress['step'] = 1
+        progress['epoch'] = 1
+        progress['cur_step'] = 0 # step in the current epoch, for resuming the sampler
+        total_progress = []
+        # if self.args.resume or self.args.validate:
+        if self.args.resume:
+            progress_pkl = "%s/progress.pkl" % self.args.exp_dir
+            with open(progress_pkl, "rb") as f:
+                total_progress = pickle.load(f)
+                progress['best_step'], progress['best_score'], progress['step'], progress['epoch'], progress['cur_step'], _ = total_progress[-1]
+            if self.rank == 0:
+                logging.info("\nResume training from:")
+                logging.info("  epoch = %s" % progress['epoch'])
+                logging.info("  cur_step = %s" % progress['cur_step'])
+                logging.info("  step = %s" % progress['step'])
+                logging.info("  best_step = %s" % progress['best_step'])
+                logging.info("  best_score = %s" % progress['best_score'])
+        return progress, total_progress
+    
+    def save_progress(self, name=None):
+        self.total_progress.append([self.progress['best_step'], self.progress['best_score'], int(self.progress['step']+1), self.progress['epoch'], int(self.progress['cur_step']+1), time.time() - self.start_time])
+        if name is not None:
+            progress_fn = f"{self.args.exp_dir}/progress_{name}.pkl"
+        else:
+            progress_fn = f"{self.args.exp_dir}/progress.pkl"
+        with open(progress_fn, "wb") as f:
+            pickle.dump(self.total_progress, f)
+
+    def _setup_dataloader(self):
+        train_dataset, val_dataset = combined_dataset.dataset(self.args, 'train'), combined_dataset.dataset(self.args, 'valid') # need to change 'train' to 'valid' in actual training
+
+        if self.args.dynamic_batching:
+            train_sampler = DistributedDynamicBatchSampler(train_dataset, self.args, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True, lengths_list=train_dataset.lengths_list, verbose=True, epoch=0)
+            valid_sampler = DistributedDynamicBatchSampler(val_dataset, self.args, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True, lengths_list=val_dataset.lengths_list, verbose=True, epoch=0)
+        else:
+            train_sampler = StatefulDistributedSampler(train_dataset, self.args.batch_size//self.world_size, num_replicas=self.world_size, rank=self.rank, shuffle=True, seed=self.args.seed, drop_last=True)
+            valid_sampler = DistributedSampler(val_dataset, num_replicas=self.world_size, rank=self.rank, shuffle=False, seed=self.args.seed, drop_last=False)
+            
+        if self.progress['step'] > 1:
+            train_sampler.set_epoch_resume(self.progress['epoch'], self.progress['cur_step'])
+            assert self.phn2num != None
+
+        if self.phn2num != None:
+            train_dataset.phn2num = self.phn2num
+            val_dataset.phn2num = self.phn2num
+
+        if self.args.dynamic_batching:
+            train_loader = torch.utils.data.DataLoader(train_dataset, 
+                            batch_sampler=train_sampler, 
+                            num_workers=self.args.num_workers,
+                            collate_fn=train_dataset.collate, persistent_workers=True
+                            )
+            valid_loader = torch.utils.data.DataLoader(val_dataset, 
+                            batch_sampler=valid_sampler, 
+                            num_workers=self.args.num_workers,
+                            collate_fn=val_dataset.collate, persistent_workers=True
+                            )
+        else:
+            train_loader = torch.utils.data.DataLoader(train_dataset, 
+                            batch_size=self.args.batch_size, sampler=train_sampler, num_workers=self.args.num_workers,
+                            collate_fn=train_dataset.collate, persistent_workers=True
+                            )
+            valid_loader = torch.utils.data.DataLoader(val_dataset, 
+                            batch_size=self.args.batch_size, sampler=valid_sampler,
+                            num_workers=self.args.num_workers,
+                            collate_fn=val_dataset.collate, persistent_workers=True
+                            )
+        return len(train_dataset), train_sampler, train_loader, valid_loader
+        
+
+        
+    def _setup_models(self):
+        model = voice_star.VoiceStar(self.args)
+
+        if self.rank == 0:
+            logging.info(model)
+            logging.info("model parameters")
+            print_model_info(model)
+        
+        phn2num = None
+        optim_states = None
+        scheduler_states = None
+        if self.progress['step'] > 1:
+            bundle = torch.load(os.path.join(self.args.exp_dir, "bundle.pth"), map_location="cpu")
+            model.load_state_dict(bundle['model'])
+            optim_states = bundle['optimizer']
+            scheduler_states = bundle['scheduler']
+            phn2num = bundle['phn2num']
+            if self.rank == 0:
+                logging.info("loaded parameters and data indices from epoch %d, global step %d" % (self.progress['epoch'], self.progress['step']))
+            del bundle['model']
+
+        if self.args.load_model_from != None and self.progress['step'] <= 1:
+            logging.info(f"load weights from {self.args.load_model_from}")
+            sd = torch.load(self.args.load_model_from, map_location="cpu")
+            if hasattr(model, "carefully_load_state_dict"):
+                model.carefully_load_state_dict(sd['model'])
+            else:
+                model.load_state_dict(sd['model'])
+            phn2num = sd['phn2num']
+            del sd
+        
+
+        #### below operations is for getting params for optimizer, which is at wrapper level ###
+        if self.args.optimizer_name == "ScaledAdam":
+            trainables = [p for p in model.parameters() if p.requires_grad]
+        else:
+            no_decay = [".bias", ".audio_embeddings.weight", ".text_embeddings.weight", ".norm.weight", ".norm1.weight", ".norm2.weight"]
+            optimizer_grouped_parameters = [
+                {
+                    "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad],
+                    "weight_decay": self.args.weight_decay,
+                },
+                {
+                    "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad],
+                    "weight_decay": 0.0,
+                },
+            ]
+            if len(optimizer_grouped_parameters[1]['params']) == 0:
+                logging.info("there is no embedding weights, bias, and layernorm parameters in the model, which should be True, check model parameter names")
+                trainables = optimizer_grouped_parameters[0]
+            else:
+                trainables = optimizer_grouped_parameters
+        #### below operations is for getting params for optimizer, which is at wrapper level ###
+        model.to(self.device)
+
+        return model, trainables, optim_states, scheduler_states, phn2num
+
+    
+    def _setup_optimizer(self):
+        if self.args.optimizer_name == "ScaledAdam":
+            parameters_names = []
+            _model = self.model.module if isinstance(self.model, torch.nn.parallel.DistributedDataParallel) else self.model
+            parameters_names.append([n for n,p in self.model.named_parameters() if p.requires_grad])
+            optimizer = ScaledAdam(
+                self.trainables,
+                lr=self.args.lr,
+                betas=(0.9, 0.95),
+                clipping_scale=2.0,
+                parameters_names=parameters_names,
+                show_dominant_parameters=False,
+                clipping_update_period=self.args.clipping_update_period,
+            )
+            scheduler = Eden(optimizer, self.args.reduce_lr_start_step, self.args.reduce_lr_start_epoch, warmup_batches=self.total_step * self.args.warmup_fraction) # NOTE: if using ScaledAdam, we will use the Eden scheduler!
+
+        else:
+            optimizer = AdamW(self.trainables, lr=self.args.lr)
+            warmup_steps = self.total_step * self.args.warmup_fraction
+            def lr_lambda(current_step: int):
+                if current_step < warmup_steps:
+                    return float(current_step) / float(max(1, warmup_steps))
+                return max(
+                    0.0, float(self.total_step - current_step) / float(max(1, self.total_step - warmup_steps))
+                )
+
+            scheduler = LambdaLR(optimizer, lr_lambda, last_epoch=-1)
+            
+        # if resume
+        if self.progress['step'] > 1:
+            optimizer.load_state_dict(self.optim_states)
+            for state in optimizer.state.values():
+                for k, v in state.items():
+                    if isinstance(v, torch.Tensor):
+                        state[k] = v.cuda()
+            del self.optim_states
+
+            scheduler.load_state_dict(self.scheduler_states)
+
+        optimizer.zero_grad()
+        return optimizer, scheduler
+    
+    def seed_everything(self, seed=1):
+        os.environ['PYTHONHASHSEED'] = str(seed)
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True