Release Inflect-Nano-v1

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
+tiny_tts/text/cmudict_cache.pickle filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,201 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction,
+and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by
+the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all
+other entities that control, are controlled by, or are under common
+control with that entity. For the purposes of this definition,
+"control" means (i) the power, direct or indirect, to cause the
+direction or management of such entity, whether by contract or
+otherwise, or (ii) ownership of fifty percent (50%) or more of the
+outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity
+exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications,
+including but not limited to software source code, documentation
+source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical
+transformation or translation of a Source form, including but
+not limited to compiled object code, generated documentation,
+and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or
+Object form, made available under the License, as indicated by a
+copyright notice that is included in or attached to the work
+(an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object
+form, that is based on (or derived from) the Work and for which the
+editorial revisions, annotations, elaborations, or other modifications
+represent, as a whole, an original work of authorship. For the purposes
+of this License, Derivative Works shall not include works that remain
+separable from, or merely link (or bind by name) to the interfaces of,
+the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including
+the original version of the Work and any modifications or additions
+to that Work or Derivative Works thereof, that is intentionally
+submitted to Licensor for inclusion in the Work by the copyright owner
+or by an individual or Legal Entity authorized to submit on behalf of
+the copyright owner. For the purposes of this definition, "submitted"
+means any form of electronic, verbal, or written communication sent
+to the Licensor or its representatives, including but not limited to
+communication on electronic mailing lists, source code control systems,
+and issue tracking systems that are managed by, or on behalf of, the
+Licensor for the purpose of discussing and improving the Work, but
+excluding communication that is conspicuously marked or otherwise
+designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity
+on behalf of whom a Contribution has been received by Licensor and
+subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+copyright license to reproduce, prepare Derivative Works of,
+publicly display, publicly perform, sublicense, and distribute the
+Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+this License, each Contributor hereby grants to You a perpetual,
+worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+(except as stated in this section) patent license to make, have made,
+use, offer to sell, sell, import, and otherwise transfer the Work,
+where such license applies only to those patent claims licensable
+by such Contributor that are necessarily infringed by their
+Contribution(s) alone or by combination of their Contribution(s)
+with the Work to which such Contribution(s) was submitted. If You
+institute patent litigation against any entity (including a
+cross-claim or counterclaim in a lawsuit) alleging that the Work
+or a Contribution incorporated within the Work constitutes direct
+or contributory patent infringement, then any patent licenses
+granted to You under this License for that Work shall terminate
+as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+Work or Derivative Works thereof in any medium, with or without
+modifications, and in Source or Object form, provided that You
+meet the following conditions:
+
+(a) You must give any other recipients of the Work or
+Derivative Works a copy of this License; and
+
+(b) You must cause any modified files to carry prominent notices
+stating that You changed the files; and
+
+(c) You must retain, in the Source form of any Derivative Works
+that You distribute, all copyright, patent, trademark, and
+attribution notices from the Source form of the Work,
+excluding those notices that do not pertain to any part of
+the Derivative Works; and
+
+(d) If the Work includes a "NOTICE" text file as part of its
+distribution, then any Derivative Works that You distribute must
+include a readable copy of the attribution notices contained
+within such NOTICE file, excluding those notices that do not
+pertain to any part of the Derivative Works, in at least one
+of the following places: within a NOTICE text file distributed
+as part of the Derivative Works; within the Source form or
+documentation, if provided along with the Derivative Works; or,
+within a display generated by the Derivative Works, if and
+wherever such third-party notices normally appear. The contents
+of the NOTICE file are for informational purposes only and
+do not modify the License. You may add Your own attribution
+notices within Derivative Works that You distribute, alongside
+or as an addendum to the NOTICE text from the Work, provided
+that such additional attribution notices cannot be construed
+as modifying the License.
+
+You may add Your own copyright statement to Your modifications and
+may provide additional or different license terms and conditions
+for use, reproduction, or distribution of Your modifications, or
+for any such Derivative Works as a whole, provided Your use,
+reproduction, and distribution of the Work otherwise complies with
+the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+any Contribution intentionally submitted for inclusion in the Work
+by You to the Licensor shall be under the terms and conditions of
+this License, without any additional terms or conditions.
+Notwithstanding the above, nothing herein shall supersede or modify
+the terms of any separate license agreement you may have executed
+with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+names, trademarks, service marks, or product names of the Licensor,
+except as required for reasonable and customary use in describing the
+origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+agreed to in writing, Licensor provides the Work (and each
+Contributor provides its Contributions) on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+implied, including, without limitation, any warranties or conditions
+of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+PARTICULAR PURPOSE. You are solely responsible for determining the
+appropriateness of using or redistributing the Work and assume any
+risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+whether in tort (including negligence), contract, or otherwise,
+unless required by applicable law (such as deliberate and grossly
+negligent acts) or agreed to in writing, shall any Contributor be
+liable to You for damages, including any direct, indirect, special,
+incidental, or consequential damages of any character arising as a
+result of this License or out of the use or inability to use the
+Work (including but not limited to damages for loss of goodwill,
+work stoppage, computer failure or malfunction, or any and all
+other commercial damages or losses), even if such Contributor
+has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+the Work or Derivative Works thereof, You may choose to offer,
+and charge a fee for, acceptance of support, warranty, indemnity,
+or other liability obligations and/or rights consistent with this
+License. However, in accepting such obligations, You may act only
+on Your own behalf and on Your sole responsibility, not on behalf
+of any other Contributor, and only if You agree to indemnify,
+defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason
+of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+To apply the Apache License to your work, attach the following
+boilerplate notice, with the fields enclosed by brackets "{}"
+replaced with your own identifying information. (Don't include
+the brackets.) The text should be enclosed in the appropriate
+comment syntax for the file format. We also recommend that a
+file or class name and description of purpose be included on the
+same "printed page" as the copyright notice for easier
+identification within third-party archives.
+
+Copyright {yyyy} {name of copyright owner}
+
+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.

README.md

@@ -0,0 +1,251 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- text-to-speech
+- tts
+- speech-synthesis
+- pytorch
+- tiny-tts
+- experimental
+pipeline_tag: text-to-speech
+library_name: pytorch
+---
+
+# Inflect-Nano-v1
+
+Inflect-Nano-v1 is an experimental ultra-small English TTS stack. It is built to test how far a sub-5M-parameter text-to-speech system can be pushed with a compact non-autoregressive acoustic model and a small neural vocoder.
+
+This is **not** a production-quality or SOTA TTS model. It is a research/demo release: small, local, and runnable, but still audibly limited.
+
+## Quick Facts
+
+| Item | Value |
+|---|---:|
+| Total inference parameters | **4.632M** |
+| Acoustic model | **3.465M** |
+| Vocoder generator | **1.167M** |
+| Language | English |
+| Voice | single Mark-style synthetic male voice |
+| Sample rate | 24 kHz |
+| Acoustic output | 80-bin mel spectrogram |
+| Vocoder | custom Snake-activation HiFi-GAN-style generator |
+| Training source | synthetic Qwen3-TTS Mark-style teacher data |
+
+## Audio Examples
+
+These are unseen/OOD stress prompts, not hand-picked training rows.
+
+| Prompt | Audio |
+|---|---|
+| Wait, are you actually being for real now? I can't believe it! | <audio controls src="examples/example_01.wav"></audio> |
+| Sophia sent me 43 pictures of her doing stuff... interesting. | <audio controls src="examples/example_02.wav"></audio> |
+| Please say chrysanthemum, thoroughly, proprietary, and rural without rushing through the middle syllables. | <audio controls src="examples/example_03.wav"></audio> |
+| No, seriously, did Jordan leave the receipt in Albuquerque, or did Priya move it to Worcester? | <audio controls src="examples/example_04.wav"></audio> |
+| The Wi-Fi password is Q7-Delta-9921, but please do not say the dash like a minus sign. | <audio controls src="examples/example_05.wav"></audio> |
+| I appreciate the honesty, but that explanation sounded weirdly dramatic for a Tuesday morning. | <audio controls src="examples/example_06.wav"></audio> |
+| Could you whisper the first part, then brighten up when you say, 'we finally solved it'? | <audio controls src="examples/example_07.wav"></audio> |
+| The dermatologist, the anesthesiologist, and the statistician all disagreed about February. | <audio controls src="examples/example_08.wav"></audio> |
+
+## Install
+
+```bash
+git clone https://huggingface.co/owensong/Inflect-Nano-v1
+cd Inflect-Nano-v1
+pip install -r requirements.txt
+```
+
+The text frontend uses TinyTTS-style English G2P and may download `bert-base-uncased` tokenizer files on first run.
+
+## Usage
+
+```bash
+python inference.py \
+  --text "Wait, are you actually being for real now? I can't believe it!" \
+  --out sample.wav
+```
+
+CPU example:
+
+```bash
+python inference.py \
+  --device cpu \
+  --text "Please say chrysanthemum, thoroughly, proprietary, and rural clearly." \
+  --out sample_cpu.wav
+```
+
+Optional controls:
+
+```bash
+python inference.py \
+  --text "No, seriously, did Jordan leave the receipt in Albuquerque?" \
+  --length-scale 1.03 \
+  --pitch-scale 1.00 \
+  --energy-scale 1.00 \
+  --out sample_controlled.wav
+```
+
+Gradio demo:
+
+```bash
+python app.py
+```
+
+## Architecture
+
+Inflect-Nano-v1 is a two-part TTS stack:
+
+```text
+text
+-> TinyTTS-style normalization + G2P phoneme/tone/lang IDs
+-> compact FastSpeech-style acoustic model
+-> 80-bin mel spectrogram
+-> Snake V2Mid HiFi-GAN-style vocoder
+-> 24 kHz waveform
+```
+
+### Acoustic Model
+
+The acoustic model is a small non-autoregressive FastSpeech-style network. It predicts duration, energy, brightness, and pitch, then expands token states into frame states and decodes mels.
+
+Main config:
+
+```json
+{
+  "hidden": 168,
+  "encoder_layers": 5,
+  "decoder_layers": 6,
+  "decoder_ff_mult": 3,
+  "kernel_size": 7,
+  "speaker_dim": 64,
+  "dropout": 0.08,
+  "n_mels": 80,
+  "sample_rate": 24000,
+  "max_frames": 1400,
+  "postnet_scale": 0.1,
+  "use_frame_pitch": true,
+  "abs_frame_bins": 512
+}
+```
+
+Acoustic parameter split:
+
+```text
+total acoustic: 3.465M
+encoder:        1.292M
+decoder:        1.211M
+postnet:        0.276M
+local context:  0.226M
+frame GRU:      0.128M
+heads/embeds/projections: remainder
+```
+
+### Vocoder
+
+The vocoder is a custom Snake-activation HiFi-GAN-style generator.
+
+Main config:
+
+```json
+{
+  "variant": "snake_v2mid",
+  "sample_rate": 24000,
+  "n_fft": 1024,
+  "hop_size": 256,
+  "win_size": 1024,
+  "num_mels": 80,
+  "fmax": 12000.0,
+  "upsample_rates": [8, 8, 2, 2],
+  "upsample_kernel_sizes": [16, 16, 4, 4],
+  "upsample_initial_channel": 144,
+  "resblock_kernel_sizes": [3, 7, 11],
+  "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+  "activation": "snake"
+}
+```
+
+The vocoder was trained with HiFi-GAN-style adversarial losses and multi-resolution spectral pressure. Discriminators are training-only and are not included in inference.
+
+## Training Data
+
+The final acoustic model was trained primarily on synthetic Mark-style teacher speech.
+
+Known final broad training mix:
+
+```text
+mixed_80k_v2_40_v1_20_old_20
+rows: 80,000
+estimated duration: ~125.2 hours
+composition:
+  40k newer generalization v2 examples
+  20k generalization v1 examples
+  20k older Mark anchor examples
+```
+
+The practical release voice is one speaker:
+
+```text
+qwen3_mark
+```
+
+## Training Recipe
+
+The final F checkpoint is a staged recovery candidate, not a single monolithic run.
+
+High-level acoustic lineage:
+
+```text
+1. Mark-focused acoustic base
+2. mixed 80k generalization training
+3. predictor-exposure heads training
+4. robust prosody bridge
+5. recovery phases A/B/C
+6. short predictor-tail cleanup
+```
+
+Known acoustic continuation after the earlier Mark base:
+
+```text
+mixed80k broad training:      ~32k selected checkpoint
+predictor exposure:            9k steps
+robust prosody bridge:       1.2k steps
+recovery phase A:             3k steps
+recovery phase B:             3k steps
+recovery phase C:           2.5k steps
+predictor tail:             0.8k steps
+known continuation total:   ~51.5k steps
+```
+
+Acoustic losses included mel reconstruction, MSE, delta/acceleration losses, duration loss, energy loss, brightness loss, pitch loss, predicted-prosody exposure, and robust-prosody exposure.
+
+Final selected files in this repo:
+
+```text
+weights/inflect_nano_v1_acoustic.pt
+weights/inflect_nano_v1_vocoder.pt
+```
+
+## Limitations
+
+This model is intentionally tiny and has clear quality limits:
+
+- Unseen text can stumble or sound unstable.
+- The voice can sound robotic, buzzy, or artifacted.
+- Long or unusual prompts are less reliable.
+- It inherits habits from synthetic Qwen3-TTS teacher data.
+- It is not a voice cloning model.
+- It is not multilingual.
+- It is not suitable for production accessibility, safety, or high-quality narration use.
+
+## Recommended Framing
+
+Use this as:
+
+> An experimental 4.63M-parameter English TTS model exploring the quality/size tradeoff for ultra-small local speech synthesis.
+
+Do not present it as SOTA or production-quality.
+
+## License
+
+Apache-2.0. The repo includes TinyTTS text frontend code; its license is included as `TINY_TTS_LICENSE`.

TINY_TTS_LICENSE

@@ -0,0 +1,152 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship made available under
+      the License, as indicated by a copyright notice that is included in
+      or attached to the work (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean, as submitted to the Licensor for inclusion
+      in the Work by the copyright owner or by an individual or Legal Entity
+      authorized to submit on behalf of the copyright owner.
+
+      "Contributor" shall mean Licensor and any Legal Entity on behalf of
+      whom a Contribution has been received by the Licensor and included
+      within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by the combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a cross-claim
+      or counterclaim in a lawsuit) alleging that the Work or any
+      Contribution embodied within the Work constitutes direct or contributory
+      patent infringement, then any patent licenses granted to You under
+      this License for that Work shall terminate as of the date such
+      litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or Derivative
+          Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work; and
+
+      (d) If the Work includes a "NOTICE" text file, You must include a
+          readable copy of the attribution notices contained within such
+          NOTICE file.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or agreed
+      to in writing, Licensor provides the Work (and each Contributor
+      provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES
+      OR CONDITIONS OF ANY KIND, either express or implied, including,
+      without limitation, any warranties or conditions of TITLE,
+      NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR
+      PURPOSE. You are solely responsible for determining the
+      appropriateness of using or reproducing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or exemplary damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or all other
+      commercial damages or losses), even if such Contributor has been
+      advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2025 tronghieuit
+
+   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.

app.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+import soundfile as sf
+import torch
+
+from inference import DEFAULT_ACOUSTIC, DEFAULT_VOCODER, load_acoustic, load_vocoder, synthesize
+
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ACOUSTIC, SPEAKERS, ACOUSTIC_PARAMS = load_acoustic(DEFAULT_ACOUSTIC, DEVICE)
+VOCODER, VOCODER_PARAMS = load_vocoder(DEFAULT_VOCODER, DEVICE)
+
+
+def generate(text: str, length_scale: float, pitch_scale: float, energy_scale: float) -> str:
+    text = (text or "").strip()
+    if not text:
+        raise gr.Error("Enter text first.")
+    if len(text) > 350:
+        raise gr.Error("Keep text under 350 characters for this tiny demo model.")
+    audio = synthesize(
+        text,
+        ACOUSTIC,
+        VOCODER,
+        SPEAKERS,
+        DEVICE,
+        length_scale=length_scale,
+        pitch_scale=pitch_scale,
+        energy_scale=energy_scale,
+    )
+    path = Path(tempfile.mkdtemp()) / "inflect_nano_v1.wav"
+    sf.write(str(path), audio, 24000, subtype="PCM_16")
+    return str(path)
+
+
+DESCRIPTION = f"""
+Experimental ultra-small English TTS stack.
+
+Inference params: {(ACOUSTIC_PARAMS + VOCODER_PARAMS) / 1_000_000:.3f}M total
+({ACOUSTIC_PARAMS / 1_000_000:.3f}M acoustic + {VOCODER_PARAMS / 1_000_000:.3f}M vocoder).
+
+This is a research/demo model, not a polished production-quality TTS system.
+"""
+
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Textbox(
+            label="Text",
+            value="Wait, are you actually being for real now? I can't believe it!",
+            lines=3,
+        ),
+        gr.Slider(0.85, 1.20, value=1.00, step=0.01, label="Length scale"),
+        gr.Slider(0.85, 1.15, value=1.00, step=0.01, label="Pitch scale"),
+        gr.Slider(0.85, 1.15, value=1.00, step=0.01, label="Energy scale"),
+    ],
+    outputs=gr.Audio(label="Generated audio", type="filepath"),
+    title="Inflect-Nano-v1",
+    description=DESCRIPTION,
+    examples=[
+        ["Please say chrysanthemum, thoroughly, proprietary, and rural without rushing through the middle syllables.", 1.0, 1.0, 1.0],
+        ["No, seriously, did Jordan leave the receipt in Albuquerque, or did Priya move it to Worcester?", 1.0, 1.0, 1.0],
+        ["The Wi-Fi password is Q7-Delta-9921, but please do not say the dash like a minus sign.", 1.0, 1.0, 1.0],
+    ],
+)
+
+
+if __name__ == "__main__":
+    demo.launch()

examples/example_01.wav

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

examples/example_02.wav

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

examples/example_03.wav

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

examples/example_04.wav

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

examples/example_05.wav

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

examples/example_06.wav

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

examples/example_07.wav

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

examples/example_08.wav

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

examples/examples.json

@@ -0,0 +1,58 @@
+[
+  {
+    "id": "example_01",
+    "text": "Wait, are you actually being for real now? I can't believe it!",
+    "file": "examples/example_01.wav",
+    "seconds": 3.296,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_02",
+    "text": "Sophia sent me 43 pictures of her doing stuff... interesting.",
+    "file": "examples/example_02.wav",
+    "seconds": 3.104,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_03",
+    "text": "Please say chrysanthemum, thoroughly, proprietary, and rural without rushing through the middle syllables.",
+    "file": "examples/example_03.wav",
+    "seconds": 5.045333333333334,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_04",
+    "text": "No, seriously, did Jordan leave the receipt in Albuquerque, or did Priya move it to Worcester?",
+    "file": "examples/example_04.wav",
+    "seconds": 4.949333333333334,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_05",
+    "text": "The Wi-Fi password is Q7-Delta-9921, but please do not say the dash like a minus sign.",
+    "file": "examples/example_05.wav",
+    "seconds": 5.941333333333334,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_06",
+    "text": "I appreciate the honesty, but that explanation sounded weirdly dramatic for a Tuesday morning.",
+    "file": "examples/example_06.wav",
+    "seconds": 4.693333333333333,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_07",
+    "text": "Could you whisper the first part, then brighten up when you say, 'we finally solved it'?",
+    "file": "examples/example_07.wav",
+    "seconds": 4.373333333333333,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  },
+  {
+    "id": "example_08",
+    "text": "The dermatologist, the anesthesiologist, and the statistician all disagreed about February.",
+    "file": "examples/example_08.wav",
+    "seconds": 4.8853333333333335,
+    "source_gallery": "INFLECT_MICRO_5M_OOD_STRESS_CANDIDATES_V1/F_RECOVERY_FINAL"
+  }
+]

examples/smoke_test.wav

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

inference.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parent
+sys.path.insert(0, str(REPO_ROOT))
+
+from tiny_tts.nn import commons
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.text.english import grapheme_to_phoneme, normalize_text
+from tiny_tts.utils import ADD_BLANK
+
+from tinytts_text_cleaning import clean_tinytts_text
+from train_hifigan_oracle_v1 import HifiGanGenerator, make_config
+from train_inflect_micro_fastspeech_v3_pitch import MicroFastSpeech, MicroFastSpeechConfig
+
+
+DEFAULT_ACOUSTIC = REPO_ROOT / "weights" / "inflect_nano_v1_acoustic.pt"
+DEFAULT_VOCODER = REPO_ROOT / "weights" / "inflect_nano_v1_vocoder.pt"
+
+
+def text_to_tokens(text: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    cleaned = clean_tinytts_text(text)
+    normalized = normalize_text(cleaned)
+    phones, tones, _ = grapheme_to_phoneme(normalized)
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+    return torch.LongTensor(phone_ids), torch.LongTensor(tone_ids), torch.LongTensor(lang_ids)
+
+
+def load_acoustic(path: Path, device: torch.device) -> tuple[MicroFastSpeech, dict[str, int], int]:
+    ckpt = torch.load(path, map_location=device, weights_only=False)
+    cfg = MicroFastSpeechConfig(**ckpt["config"])
+    model = MicroFastSpeech(cfg).to(device)
+    model.load_state_dict(ckpt["model"])
+    model.eval()
+    params = int(ckpt.get("params") or sum(p.numel() for p in model.parameters()))
+    return model, ckpt.get("speakers") or {"qwen3_mark": 0}, params
+
+
+def load_vocoder(path: Path, device: torch.device) -> tuple[HifiGanGenerator, int]:
+    ckpt = torch.load(path, map_location=device, weights_only=False)
+    cfg = make_config((ckpt.get("config") or {}).get("variant", "snake_v2mid"))
+    model = HifiGanGenerator(cfg).to(device)
+    model.load_state_dict(ckpt["generator"])
+    model.remove_weight_norm()
+    model.eval()
+    params = int(ckpt.get("generator_params") or sum(p.numel() for p in model.parameters()))
+    return model, params
+
+
+def rms_db(audio: np.ndarray) -> float:
+    return 20.0 * math.log10(float(np.sqrt(np.mean(audio**2, dtype=np.float64))) + 1e-9)
+
+
+def normalize_audio(audio: np.ndarray, target_rms_db: float = -20.0, peak_db: float = -1.0) -> np.ndarray:
+    audio = np.asarray(audio, dtype=np.float32).reshape(-1)
+    if audio.size == 0:
+        audio = np.zeros(1, dtype=np.float32)
+    audio = audio - float(audio.mean())
+    audio *= 10 ** ((target_rms_db - rms_db(audio)) / 20.0)
+    peak = float(np.max(np.abs(audio)) + 1e-9)
+    peak_limit = 10 ** (peak_db / 20.0)
+    if peak > peak_limit:
+        audio *= peak_limit / peak
+    return np.clip(audio, -1.0, 1.0)
+
+
+@torch.inference_mode()
+def synthesize(
+    text: str,
+    acoustic: MicroFastSpeech,
+    vocoder: HifiGanGenerator,
+    speakers: dict[str, int],
+    device: torch.device,
+    length_scale: float = 1.0,
+    pitch_scale: float = 1.0,
+    energy_scale: float = 1.0,
+) -> np.ndarray:
+    phone, tone, lang = text_to_tokens(text)
+    phone = phone.unsqueeze(0).to(device)
+    tone = tone.unsqueeze(0).to(device)
+    lang = lang.unsqueeze(0).to(device)
+    speaker = torch.LongTensor([int(speakers.get("qwen3_mark", 0))]).to(device)
+    mel = acoustic.infer(
+        phone,
+        tone,
+        lang,
+        speaker,
+        length_scale=float(length_scale),
+        pitch_scale=float(pitch_scale),
+        energy_scale=float(energy_scale),
+    )
+    wav = vocoder(mel).squeeze().detach().cpu().numpy()
+    return normalize_audio(wav)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="Run Inflect-Nano-v1 text-to-speech.")
+    ap.add_argument("--text", required=True)
+    ap.add_argument("--out", type=Path, default=Path("inflect_nano_v1_output.wav"))
+    ap.add_argument("--acoustic", type=Path, default=DEFAULT_ACOUSTIC)
+    ap.add_argument("--vocoder", type=Path, default=DEFAULT_VOCODER)
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--length-scale", type=float, default=1.0)
+    ap.add_argument("--pitch-scale", type=float, default=1.0)
+    ap.add_argument("--energy-scale", type=float, default=1.0)
+    args = ap.parse_args()
+
+    device = torch.device(args.device)
+    acoustic, speakers, acoustic_params = load_acoustic(args.acoustic, device)
+    vocoder, vocoder_params = load_vocoder(args.vocoder, device)
+    audio = synthesize(
+        args.text,
+        acoustic,
+        vocoder,
+        speakers,
+        device,
+        length_scale=args.length_scale,
+        pitch_scale=args.pitch_scale,
+        energy_scale=args.energy_scale,
+    )
+    args.out.parent.mkdir(parents=True, exist_ok=True)
+    sf.write(str(args.out), audio, 24000, subtype="PCM_16")
+    print(f"Wrote {args.out}")
+    print(f"Params: acoustic={acoustic_params:,} vocoder={vocoder_params:,} total={acoustic_params + vocoder_params:,}")
+
+
+if __name__ == "__main__":
+    main()

model_summary.json

@@ -0,0 +1,78 @@
+{
+  "model_name": "Inflect-Nano-v1",
+  "language": "English",
+  "speaker": "qwen3_mark",
+  "sample_rate": 24000,
+  "n_mels": 80,
+  "acoustic_params": 3465125,
+  "vocoder_generator_params": 1167077,
+  "total_inference_params": 4632202,
+  "acoustic_config": {
+    "vocab_size": 256,
+    "tone_size": 16,
+    "lang_size": 4,
+    "n_mels": 80,
+    "hidden": 168,
+    "encoder_layers": 5,
+    "decoder_layers": 6,
+    "decoder_ff_mult": 3,
+    "kernel_size": 7,
+    "speaker_count": 2,
+    "speaker_dim": 64,
+    "dropout": 0.08,
+    "sample_rate": 24000,
+    "max_frames": 1400,
+    "postnet_scale": 0.1,
+    "use_frame_pitch": true,
+    "abs_frame_bins": 512
+  },
+  "vocoder_config": {
+    "variant": "snake_v2mid",
+    "sample_rate": 24000,
+    "n_fft": 1024,
+    "hop_size": 256,
+    "win_size": 1024,
+    "num_mels": 80,
+    "fmin": 0.0,
+    "fmax": 12000.0,
+    "resblock": "1",
+    "upsample_rates": [
+      8,
+      8,
+      2,
+      2
+    ],
+    "upsample_kernel_sizes": [
+      16,
+      16,
+      4,
+      4
+    ],
+    "upsample_initial_channel": 144,
+    "resblock_kernel_sizes": [
+      3,
+      7,
+      11
+    ],
+    "resblock_dilation_sizes": [
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ],
+      [
+        1,
+        3,
+        5
+      ]
+    ],
+    "activation": "snake"
+  },
+  "acoustic_checkpoint": "D:\\Inflect-Storage\\Inflect-New-offload\\outputs\\inflect_micro_fastspeech_v10_curriculum\\mark_generalization_v2\\overnight_recovery_v1\\phaseD_predictor_tail_800\\inflect-micro-fastspeech-800.pt",
+  "vocoder_checkpoint": "D:\\Inflect-Storage\\Inflect-New-offload\\outputs\\hifigan_oracle_v1\\snake_v2mid_final_polish_v1\\hifigan-snake_v2mid-165000.pt"
+}

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchaudio
+soundfile
+numpy
+g2p_en
+transformers
+gradio

tiny_tts/__init__.py

@@ -0,0 +1,90 @@
+import os
+import torch
+import soundfile as sf
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models.synthesizer import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+from tiny_tts.infer import load_engine
+
+class TinyTTS:
+    def __init__(self, checkpoint_path=None, device=None):
+        if device is None:
+            self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            self.device = device
+            
+        if checkpoint_path is None:
+            # Look for default checkpoint in pacakage
+            pkg_dir = os.path.dirname(os.path.abspath(__file__))
+            default_ckpt = os.path.join(os.path.dirname(pkg_dir), "checkpoints", "G.pth")
+            # 2. Check HuggingFace Cache / Download
+            if not os.path.exists(default_ckpt):
+                try:
+                    from huggingface_hub import hf_hub_download
+                    print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+                    default_ckpt = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+                except ImportError:
+                    raise ImportError("huggingface_hub is required to auto-download the model. Run: pip install huggingface_hub")
+                except Exception as e:
+                    raise ValueError(f"Failed to download checkpoint from Hugging Face: {e}")
+
+            checkpoint_path = default_ckpt
+                
+        self.model = load_engine(checkpoint_path, self.device)
+
+    def speak(self, text, output_path="output.wav", speaker="MALE", speed=1.0):
+        """Synthesize text to speech and save to output_path."""
+        print(f"Synthesizing: {text}")
+
+        # Normalize text
+        normalized = normalize_text(text)
+
+        # Phonemize
+        phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+        # Convert to sequence
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        # Add blanks
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids = commons.insert_blanks(tone_ids, 0)
+            lang_ids = commons.insert_blanks(lang_ids, 0)
+
+        x = torch.LongTensor(phone_ids).unsqueeze(0).to(self.device)
+        x_lengths = torch.LongTensor([len(phone_ids)]).to(self.device)
+        tone = torch.LongTensor(tone_ids).unsqueeze(0).to(self.device)
+        language = torch.LongTensor(lang_ids).unsqueeze(0).to(self.device)
+
+        # Speaker ID
+        if speaker not in SPK2ID:
+            print(f"Warning: Speaker '{speaker}' not found, using ID 0. Available: {list(SPK2ID.keys())}")
+            sid = torch.LongTensor([0]).to(self.device)
+        else:
+            sid = torch.LongTensor([SPK2ID[speaker]]).to(self.device)
+
+        # BERT features (disabled - using zero tensors)
+        bert = torch.zeros(1024, len(phone_ids)).to(self.device).unsqueeze(0)
+        ja_bert = torch.zeros(768, len(phone_ids)).to(self.device).unsqueeze(0)
+
+        # speed > 1.0 = faster speech, < 1.0 = slower speech
+        length_scale = 1.0 / speed
+
+        with torch.no_grad():
+            audio, *_ = self.model.infer(
+                x, x_lengths, sid, tone, language, bert, ja_bert,
+                noise_scale=0.667,
+                noise_scale_w=0.8,
+                length_scale=length_scale
+            )
+
+        audio_np = audio[0, 0].cpu().numpy()
+        sf.write(output_path, audio_np, SAMPLING_RATE)
+        print(f"Saved audio to {output_path}")
+        return audio_np

tiny_tts/alignment/__init__.py

@@ -0,0 +1,16 @@
+from numpy import zeros, int32, float32
+from torch import from_numpy
+
+from .core import viterbi_decode_kernel
+
+
+def viterbi_decode(neg_cent, mask):
+    device = neg_cent.device
+    dtype = neg_cent.dtype
+    neg_cent = neg_cent.data.cpu().numpy().astype(float32)
+    path = zeros(neg_cent.shape, dtype=int32)
+
+    t_t_max = mask.sum(1)[:, 0].data.cpu().numpy().astype(int32)
+    t_s_max = mask.sum(2)[:, 0].data.cpu().numpy().astype(int32)
+    viterbi_decode_kernel(path, neg_cent, t_t_max, t_s_max)
+    return from_numpy(path).to(device=device, dtype=dtype)

tiny_tts/alignment/core.py

@@ -0,0 +1,46 @@
+import numba
+
+
+@numba.jit(
+    numba.void(
+        numba.int32[:, :, ::1],
+        numba.float32[:, :, ::1],
+        numba.int32[::1],
+        numba.int32[::1],
+    ),
+    nopython=True,
+    nogil=True,
+)
+def viterbi_decode_kernel(paths, values, t_ys, t_xs):
+    b = paths.shape[0]
+    max_neg_val = -1e9
+    for i in range(int(b)):
+        path = paths[i]
+        value = values[i]
+        t_y = t_ys[i]
+        t_x = t_xs[i]
+
+        v_prev = v_cur = 0.0
+        index = t_x - 1
+
+        for y in range(t_y):
+            for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+                if x == y:
+                    v_cur = max_neg_val
+                else:
+                    v_cur = value[y - 1, x]
+                if x == 0:
+                    if y == 0:
+                        v_prev = 0.0
+                    else:
+                        v_prev = max_neg_val
+                else:
+                    v_prev = value[y - 1, x - 1]
+                value[y, x] += max(v_prev, v_cur)
+
+        for y in range(t_y - 1, -1, -1):
+            path[y, index] = 1
+            if index != 0 and (
+                index == y or value[y - 1, index] < value[y - 1, index - 1]
+            ):
+                index = index - 1

tiny_tts/infer.py

@@ -0,0 +1,191 @@
+import os
+import sys
+import re
+import torch
+import soundfile as sf
+import argparse
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.models import VoiceSynthesizer
+from tiny_tts.text.symbols import symbols
+from tiny_tts.utils import (
+    SAMPLING_RATE, SEGMENT_FRAMES, ADD_BLANK, SPEC_CHANNELS,
+    N_SPEAKERS, SPK2ID, MODEL_PARAMS,
+)
+
+
+def load_engine(checkpoint_path, device='cuda'):
+    print(f"Loading model from {checkpoint_path}")
+    net_g = VoiceSynthesizer(
+        len(symbols),
+        SPEC_CHANNELS,
+        SEGMENT_FRAMES,
+        n_speakers=N_SPEAKERS,
+        **MODEL_PARAMS
+    ).to(device)
+
+    # Count model parameters
+    total_params = sum(p.numel() for p in net_g.parameters())
+    trainable_params = sum(p.numel() for p in net_g.parameters() if p.requires_grad)
+    print(f"Model parameters: {total_params/1e6:.2f}M total, {trainable_params/1e6:.2f}M trainable")
+
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    state_dict = checkpoint['model']
+
+    # Remove module. prefix and filter shape mismatches
+    model_state = net_g.state_dict()
+    new_state_dict = {}
+    skipped = []
+    for k, v in state_dict.items():
+        key = k[7:] if k.startswith('module.') else k
+        if key in model_state:
+            if v.shape == model_state[key].shape:
+                new_state_dict[key] = v
+            else:
+                skipped.append(f"{key}: ckpt{v.shape} vs model{model_state[key].shape}")
+        else:
+            new_state_dict[key] = v
+
+    if skipped:
+        print(f"Skipped {len(skipped)} mismatched keys:")
+        for s in skipped[:5]:
+            print(f"  {s}")
+        if len(skipped) > 5:
+            print(f"  ... and {len(skipped)-5} more")
+
+    net_g.load_state_dict(new_state_dict, strict=False)
+    net_g.eval()
+
+    # Fold weight_norm into weight tensors for faster inference (~18% speedup)
+    net_g.dec.remove_weight_norm()
+
+    return net_g
+
+
+def synthesize(text, output_path, model, speaker="MALE", device='cuda', speed=1.0):
+    print(f"Synthesizing: {text}")
+
+    # Normalize text
+    normalized = normalize_text(text)
+
+    # Phonemize
+    phones, tones, word2ph = grapheme_to_phoneme(normalized)
+
+    # Convert to sequence
+    phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+    # Add blanks
+    if ADD_BLANK:
+        phone_ids = commons.insert_blanks(phone_ids, 0)
+        tone_ids = commons.insert_blanks(tone_ids, 0)
+        lang_ids = commons.insert_blanks(lang_ids, 0)
+
+    x = torch.LongTensor(phone_ids).unsqueeze(0).to(device)
+    x_lengths = torch.LongTensor([len(phone_ids)]).to(device)
+    tone = torch.LongTensor(tone_ids).unsqueeze(0).to(device)
+    language = torch.LongTensor(lang_ids).unsqueeze(0).to(device)
+
+    # Speaker ID
+    if speaker not in SPK2ID:
+        print(f"Warning: Speaker {speaker} not found, using ID 0")
+        sid = torch.LongTensor([0]).to(device)
+    else:
+        sid = torch.LongTensor([SPK2ID[speaker]]).to(device)
+
+    # BERT features (disabled - using zero tensors)
+    bert = torch.zeros(1024, len(phone_ids)).to(device).unsqueeze(0)
+    ja_bert = torch.zeros(768, len(phone_ids)).to(device).unsqueeze(0)
+
+    # speed > 1.0 = faster speech, < 1.0 = slower speech
+    length_scale = 1.0 / speed
+
+    with torch.no_grad():
+        audio, *_ = model.infer(
+            x, x_lengths, sid, tone, language, bert, ja_bert,
+            noise_scale=0.667,
+            noise_scale_w=0.8,
+            length_scale=length_scale
+        )
+
+    audio = audio[0, 0].cpu().numpy()
+    sf.write(output_path, audio, SAMPLING_RATE)
+    print(f"Saved audio to {output_path}")
+
+
+def get_latest_checkpoint(checkpoint_dir):
+    """Finds the latest G_*.pth checkpoint in the given directory."""
+    checkpoints = [f for f in os.listdir(checkpoint_dir) if f.startswith('G_') and f.endswith('.pth')]
+    if not checkpoints:
+        return None
+
+    def get_step(filename):
+        match = re.search(r'_(\d+)\.pth', filename)
+        return int(match.group(1)) if match else -1
+
+    latest_ckpt = max(checkpoints, key=get_step)
+    return os.path.join(checkpoint_dir, latest_ckpt)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TinyTTS — English Text-to-Speech Inference")
+    parser.add_argument("--text", "-t", type=str, default="The weather is nice today, and I feel very relaxed.", help="Text to synthesize")
+    parser.add_argument("--checkpoint", "-c", type=str, default=None, help="Path to checkpoint. Auto-downloads if not provided.")
+    parser.add_argument("--output", "-o", type=str, default="output.wav", help="Output audio file path")
+    parser.add_argument("--speaker", "-s", type=str, default="MALE", help="Speaker ID")
+    parser.add_argument("--speed", type=float, default=1.0, help="Speech speed (1.0=normal, 1.5=faster, 0.7=slower)")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to use (cuda or cpu)")
+
+    args = parser.parse_args()
+
+    if args.checkpoint is None:
+        try:
+            from huggingface_hub import hf_hub_download
+            print("Downloading/Loading checkpoint from Hugging Face Hub (backtracking/tiny-tts)...")
+            args.checkpoint = hf_hub_download(repo_id="backtracking/tiny-tts", filename="G.pth")
+        except ImportError:
+            print("Error: huggingface_hub is required for auto-download. Run: pip install huggingface_hub")
+            sys.exit(1)
+        except Exception as e:
+            print(f"Error downloading checkpoint: {e}")
+            sys.exit(1)
+
+    if not os.path.exists(args.checkpoint):
+        print(f"Error: Checkpoint or directory not found at {args.checkpoint}")
+        sys.exit(1)
+
+    if os.path.isdir(args.checkpoint):
+        latest_ckpt = get_latest_checkpoint(args.checkpoint)
+        if not latest_ckpt:
+            print(f"Error: No G_*.pth checkpoints found in directory {args.checkpoint}")
+            sys.exit(1)
+        args.checkpoint = latest_ckpt
+        print(f"Auto-detected latest checkpoint: {args.checkpoint}")
+
+    # Extract step from checkpoint filename
+    ckpt_basename = os.path.basename(args.checkpoint)
+    match = re.search(r'_(\d+)\.pth', ckpt_basename)
+    step_str = match.group(1) if match else "unknown"
+
+    # Save to output folder
+    out_dir = "infer_outputs"
+    os.makedirs(out_dir, exist_ok=True)
+
+    out_name = os.path.basename(args.output)
+    name, ext = os.path.splitext(out_name)
+    model = load_engine(args.checkpoint, args.device)
+
+    if args.speaker.lower() == "all":
+        if not SPK2ID:
+            print("Error: No speakers found")
+            sys.exit(1)
+        print(f"Synthesizing for all {len(SPK2ID)} speakers...")
+        for spk in SPK2ID.keys():
+            final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{spk}{ext}")
+            synthesize(args.text, final_output, model, speaker=spk, device=args.device, speed=args.speed)
+    else:
+        final_output = os.path.join(out_dir, f"{name}_step{step_str}_spk{args.speaker}{ext}")
+        synthesize(args.text, final_output, model, speaker=args.speaker, device=args.device, speed=args.speed)
+
+if __name__ == "__main__":
+    main()

tiny_tts/infer_onnx.py

@@ -0,0 +1,199 @@
+"""
+ONNX Runtime inference engine for TinyTTS.
+
+Replaces the PyTorch VoiceSynthesizer.infer() with equivalent
+ONNX Runtime sessions + NumPy ops for the non-exported parts
+(alignment path computation).
+"""
+import os
+import numpy as np
+import soundfile as sf
+
+from tiny_tts.text.english import normalize_text, grapheme_to_phoneme
+from tiny_tts.text import phonemes_to_ids
+from tiny_tts.nn import commons
+from tiny_tts.utils.config import (
+    SAMPLING_RATE, ADD_BLANK, SPK2ID,
+)
+
+try:
+    import onnxruntime as ort
+except ImportError:
+    raise ImportError("onnxruntime is required. Run: pip install onnxruntime")
+
+
+def _build_session(path: str, use_gpu: bool = False):
+    """Create an ORT InferenceSession with optional GPU support."""
+    providers = (
+        ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        if use_gpu else
+        ["CPUExecutionProvider"]
+    )
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    opts.intra_op_num_threads = os.cpu_count() or 4
+    return ort.InferenceSession(path, sess_options=opts, providers=providers)
+
+
+def _create_length_mask_np(lengths, max_len=None):
+    """NumPy equivalent of commons.create_length_mask."""
+    if max_len is None:
+        max_len = int(lengths.max())
+    ids = np.arange(max_len, dtype=np.float32)     # [T]
+    mask = (ids[None, :] < lengths[:, None]).astype(np.float32)  # [B, T]
+    return mask
+
+
+def _compute_alignment_path_np(w_ceil, attn_mask):
+    """
+    Monotonic alignment path - vectorized via cumsum (much faster than Python loops).
+    w_ceil:    [B, 1, T_x] — integer duration per phone
+    attn_mask: [B, 1, T_y, T_x] — joint mask
+    Returns attn: [B, 1, T_y, T_x]
+    """
+    B, _, T_x = w_ceil.shape
+    T_y = attn_mask.shape[2]
+
+    # Build duration matrix: for each phone column expand the duration
+    # cumulative sum of durations gives us the end frame index for each phone
+    dur = w_ceil[:, 0, :]                            # [B, T_x]
+    cum_dur = np.cumsum(dur, axis=1)                 # [B, T_x]  — end frame (1-indexed)
+    cum_dur_prev = np.pad(cum_dur[:, :-1], ((0,0),(1,0)))  # [B, T_x] — start frame
+
+    # Frame indices: [1, T_y, 1]
+    frame_idx = np.arange(T_y, dtype=np.float32)[None, :, None]   # [1, T_y, 1]
+    # For each phone, mark frames [start, end)
+    # cum_dur_prev: [B,1,T_x], cum_dur: [B,1,T_x]
+    start = cum_dur_prev[:, None, :]                 # [B, 1, T_x]
+    end   = cum_dur[:, None, :]                      # [B, 1, T_x]
+    attn  = ((frame_idx >= start) & (frame_idx < end)).astype(np.float32)  # [B, T_y, T_x]
+    attn  = attn[:, None, :, :]                      # [B, 1, T_y, T_x]
+    return attn * attn_mask
+
+
+class OnnxTinyTTS:
+    """
+    Inference using ONNX Runtime.
+
+    Args:
+        onnx_dir: directory containing the 4 .onnx files
+        use_gpu:  if True, try CUDAExecutionProvider
+    """
+
+    def __init__(self, onnx_dir: str = "onnx", use_gpu: bool = False):
+        onnx_dir = os.path.abspath(onnx_dir)
+        print(f"Loading ONNX sessions from: {onnx_dir}")
+
+        self._enc  = _build_session(os.path.join(onnx_dir, "text_encoder.onnx"), use_gpu)
+        self._dp   = _build_session(os.path.join(onnx_dir, "duration_predictor.onnx"), use_gpu)
+        self._flow = _build_session(os.path.join(onnx_dir, "flow.onnx"), use_gpu)
+        self._dec  = _build_session(os.path.join(onnx_dir, "decoder.onnx"), use_gpu)
+
+        print("ONNX sessions ready ✅")
+
+    def _text_to_ids(self, text: str):
+        normalized = normalize_text(text)
+        phones, tones, _ = grapheme_to_phoneme(normalized)
+        phone_ids, tone_ids, lang_ids = phonemes_to_ids(phones, tones, "EN")
+
+        if ADD_BLANK:
+            phone_ids = commons.insert_blanks(phone_ids, 0)
+            tone_ids  = commons.insert_blanks(tone_ids, 0)
+            lang_ids  = commons.insert_blanks(lang_ids, 0)
+
+        return phone_ids, tone_ids, lang_ids
+
+    def speak(
+        self,
+        text: str,
+        output_path: str = "onnx_output.wav",
+        speaker: str = "female",
+        noise_scale: float = 0.667,
+        noise_scale_w: float = 0.8,
+        length_scale: float = 1.0,
+        output_sr: int = None,
+    ) -> np.ndarray:
+        """Synthesize speech and save to output_path.
+        
+        Args:
+            output_sr: If set (e.g. 22050), resample the output from 44100 Hz.
+                       Useful to reduce file size while keeping quality.
+        """
+        print(f"[ONNX] Synthesizing: {text}")
+
+        phone_ids, tone_ids, lang_ids = self._text_to_ids(text)
+        T = len(phone_ids)
+
+        # Prepare inputs as float32 / int64 arrays
+        x       = np.array(phone_ids, dtype=np.int64)[None, :]        # [1, T]
+        x_len   = np.array([T], dtype=np.int64)                        # [1]
+        tone    = np.array(tone_ids, dtype=np.int64)[None, :]          # [1, T]
+        lang    = np.array(lang_ids, dtype=np.int64)[None, :]          # [1, T]
+        bert    = np.zeros((1, 1024, T), dtype=np.float32)
+        ja_bert = np.zeros((1, 768,  T), dtype=np.float32)
+        sid_val = SPK2ID.get(speaker, 0)
+        sid     = np.array([sid_val], dtype=np.int64)                  # [1]
+
+        # ── 1. Text Encoder ──────────────────────────────────────────────
+        x_enc, m_p, logs_p, x_mask, g = self._enc.run(
+            None,
+            {
+                "phone_ids":    x,
+                "phone_lengths":x_len,
+                "tone_ids":     tone,
+                "language_ids": lang,
+                "bert":         bert,
+                "ja_bert":      ja_bert,
+                "speaker_id":   sid,
+            },
+        )
+
+        # ── 2. Duration Predictor ─────────────────────────────────────────
+        logw = self._dp.run(None, {"x": x_enc, "x_mask": x_mask, "g": g})[0]
+
+        # ── 3. Alignment Path (NumPy) ─────────────────────────────────────
+        w        = np.exp(logw) * x_mask * length_scale          # [1, 1, T]
+        w_ceil   = np.ceil(w)                                    # [1, 1, T]
+        y_len    = max(1, int(w_ceil.sum()))
+        y_lens   = np.array([y_len], dtype=np.int64)
+
+        y_mask   = _create_length_mask_np(y_lens, y_len)         # [1, T_y]
+        y_mask   = y_mask[:, None, :]                            # [1, 1, T_y]
+        # attn_mask: [1, 1, T_y, T_x]  (outer product of frame mask and phone mask)
+        attn_mask = y_mask[:, :, :, None] * x_mask[:, :, None, :]   # [1,1,T_y,T_x]
+        attn     = _compute_alignment_path_np(w_ceil, attn_mask) # [1, 1, T_y, T_x]
+
+        # Expand prior stats via alignment
+        m_p_exp    = np.matmul(attn[:, 0], m_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+        logs_p_exp = np.matmul(attn[:, 0], logs_p.transpose(0, 2, 1)).transpose(0, 2, 1)
+
+        # ── 4. Sample z_p ─────────────────────────────────────────────────
+        z_p = m_p_exp + np.random.randn(*m_p_exp.shape).astype(np.float32) * \
+              np.exp(logs_p_exp) * noise_scale
+
+        # ── 5. Flow (reverse) ─────────────────────────────────────────────
+        z = self._flow.run(
+            None,
+            {"z_p": z_p, "y_mask": y_mask.astype(np.float32), "g": g},
+        )[0]
+
+        # ── 6. Decoder ────────────────────────────────────────────────────
+        z_masked = (z * y_mask).astype(np.float32)
+        audio = self._dec.run(None, {"z": z_masked, "g": g})[0]  # [1, 1, samples]
+
+        audio_np = audio[0, 0]
+        save_sr = SAMPLING_RATE
+        if output_sr is not None and output_sr != SAMPLING_RATE:
+            try:
+                import torchaudio
+                import torch
+                wav_t = torch.from_numpy(audio_np).unsqueeze(0)
+                resampler = torchaudio.transforms.Resample(SAMPLING_RATE, output_sr)
+                audio_np = resampler(wav_t).squeeze(0).numpy()
+                save_sr = output_sr
+            except Exception as e:
+                print(f"[ONNX] Resampling failed ({e}), saving at {SAMPLING_RATE}Hz")
+
+        sf.write(output_path, audio_np, save_sr)
+        print(f"[ONNX] Saved: {output_path} ({save_sr}Hz)")
+        return audio_np

tiny_tts/models/__init__.py

@@ -0,0 +1 @@
+from .synthesizer import VoiceSynthesizer

tiny_tts/models/synthesizer.py

@@ -0,0 +1,718 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from tiny_tts.nn import commons
+from tiny_tts.nn import modules
+from tiny_tts.nn import attentions
+
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from tiny_tts.nn.commons import initialize_weights, compute_padding
+import tiny_tts.alignment as alignment
+
+
+class AttentionFlowBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        n_flows=4,
+        gin_channels=0,
+        share_parameter=False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+
+        self.wn = (
+            attentions.FeedForward(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=self.gin_channels,
+            )
+            if share_parameter
+            else None
+        )
+
+        for i in range(n_flows):
+            self.flows.append(
+                modules.TransformerCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    n_layers,
+                    n_heads,
+                    p_dropout,
+                    filter_channels,
+                    mean_only=True,
+                    wn_sharing_parameter=self.wn,
+                    gin_channels=self.gin_channels,
+                )
+            )
+            self.flows.append(modules.FlipTransform())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class VariationalDurationModel(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout,
+        n_flows=4,
+        gin_channels=0,
+    ):
+        super().__init__()
+        filter_channels = in_channels
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.log_flow = modules.LogTransform()
+        self.flows = nn.ModuleList()
+        self.flows.append(modules.AffineCoupling(2))
+        for i in range(n_flows):
+            self.flows.append(
+                modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.flows.append(modules.FlipTransform())
+
+        self.post_pre = nn.Conv1d(1, filter_channels, 1)
+        self.post_proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.post_convs = modules.DepthwiseSepConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        self.post_flows = nn.ModuleList()
+        self.post_flows.append(modules.AffineCoupling(2))
+        for i in range(4):
+            self.post_flows.append(
+                modules.ConvolutionalFlow(2, filter_channels, kernel_size, n_layers=3)
+            )
+            self.post_flows.append(modules.FlipTransform())
+
+        self.pre = nn.Conv1d(in_channels, filter_channels, 1)
+        self.proj = nn.Conv1d(filter_channels, filter_channels, 1)
+        self.convs = modules.DepthwiseSepConv(
+            filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout
+        )
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
+
+    def forward(self, x, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
+        x = torch.detach(x)
+        x = self.pre(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.convs(x, x_mask)
+        x = self.proj(x) * x_mask
+
+        if not reverse:
+            flows = self.flows
+            assert w is not None
+
+            logdet_tot_q = 0
+            h_w = self.post_pre(w)
+            h_w = self.post_convs(h_w, x_mask)
+            h_w = self.post_proj(h_w) * x_mask
+            e_q = (
+                torch.randn(w.size(0), 2, w.size(2)).to(device=x.device, dtype=x.dtype)
+                * x_mask
+            )
+            z_q = e_q
+            for flow in self.post_flows:
+                z_q, logdet_q = flow(z_q, x_mask, g=(x + h_w))
+                logdet_tot_q += logdet_q
+            z_u, z1 = torch.split(z_q, [1, 1], 1)
+            u = torch.sigmoid(z_u) * x_mask
+            z0 = (w - u) * x_mask
+            logdet_tot_q += torch.sum(
+                (F.logsigmoid(z_u) + F.logsigmoid(-z_u)) * x_mask, [1, 2]
+            )
+            logq = (
+                torch.sum(-0.5 * (math.log(2 * math.pi) + (e_q**2)) * x_mask, [1, 2])
+                - logdet_tot_q
+            )
+
+            logdet_tot = 0
+            z0, logdet = self.log_flow(z0, x_mask)
+            logdet_tot += logdet
+            z = torch.cat([z0, z1], 1)
+            for flow in flows:
+                z, logdet = flow(z, x_mask, g=x, reverse=reverse)
+                logdet_tot = logdet_tot + logdet
+            nll = (
+                torch.sum(0.5 * (math.log(2 * math.pi) + (z**2)) * x_mask, [1, 2])
+                - logdet_tot
+            )
+            return nll + logq
+        else:
+            flows = list(reversed(self.flows))
+            flows = flows[:-2] + [flows[-1]]
+            z = (
+                torch.randn(x.size(0), 2, x.size(2)).to(device=x.device, dtype=x.dtype)
+                * noise_scale
+            )
+            for flow in flows:
+                z = flow(z, x_mask, g=x, reverse=reverse)
+            z0, z1 = torch.split(z, [1, 1], 1)
+            logw = z0
+            return logw
+
+
+class DurationEstimator(nn.Module):
+    def __init__(
+        self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+
+        self.drop = nn.Dropout(p_dropout)
+        self.conv_1 = nn.Conv1d(
+            in_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_1 = modules.ChannelNorm(filter_channels)
+        self.conv_2 = nn.Conv1d(
+            filter_channels, filter_channels, kernel_size, padding=kernel_size // 2
+        )
+        self.norm_2 = modules.ChannelNorm(filter_channels)
+        self.proj = nn.Conv1d(filter_channels, 1, 1)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, in_channels, 1)
+
+    def forward(self, x, x_mask, g=None):
+        x = torch.detach(x)
+        if g is not None:
+            g = torch.detach(g)
+            x = x + self.cond(g)
+        x = self.conv_1(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_1(x)
+        x = self.drop(x)
+        x = self.conv_2(x * x_mask)
+        x = torch.relu(x)
+        x = self.norm_2(x)
+        x = self.drop(x)
+        x = self.proj(x * x_mask)
+        return x * x_mask
+
+
+class PhonemeEncoder(nn.Module):
+    def __init__(
+        self,
+        n_vocab,
+        out_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        gin_channels=0,
+        num_languages=None,
+        num_tones=None,
+    ):
+        super().__init__()
+        if num_languages is None:
+            from tiny_tts.text import num_languages
+        if num_tones is None:
+            from tiny_tts.text import num_tones
+        self.n_vocab = n_vocab
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.gin_channels = gin_channels
+        self.emb = nn.Embedding(n_vocab, hidden_channels)
+        nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.tone_emb = nn.Embedding(num_tones, hidden_channels)
+        nn.init.normal_(self.tone_emb.weight, 0.0, hidden_channels**-0.5)
+        self.language_emb = nn.Embedding(num_languages, hidden_channels)
+        nn.init.normal_(self.language_emb.weight, 0.0, hidden_channels**-0.5)
+        self.bert_proj = nn.Conv1d(1024, hidden_channels, 1)
+        self.ja_bert_proj = nn.Conv1d(768, hidden_channels, 1)
+
+        self.encoder = attentions.TransformerBlock(
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, tone, language, bert, ja_bert, g=None):
+        bert_emb = self.bert_proj(bert).transpose(1, 2)
+        ja_bert_emb = self.ja_bert_proj(ja_bert).transpose(1, 2)
+        x = (
+            self.emb(x)
+            + self.tone_emb(tone)
+            + self.language_emb(language)
+            + bert_emb
+            + ja_bert_emb
+        ) * math.sqrt(
+            self.hidden_channels
+        )
+        x = torch.transpose(x, 1, -1)
+        x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+
+        x = self.encoder(x * x_mask, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        return x, m, logs, x_mask
+
+
+class FlowBlock(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        n_flows=4,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.n_flows = n_flows
+        self.gin_channels = gin_channels
+
+        self.flows = nn.ModuleList()
+        for i in range(n_flows):
+            self.flows.append(
+                modules.FlowCouplingLayer(
+                    channels,
+                    hidden_channels,
+                    kernel_size,
+                    dilation_rate,
+                    n_layers,
+                    gin_channels=gin_channels,
+                    mean_only=True,
+                )
+            )
+            self.flows.append(modules.FlipTransform())
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        if not reverse:
+            for flow in self.flows:
+                x, _ = flow(x, x_mask, g=g, reverse=reverse)
+        else:
+            for flow in reversed(self.flows):
+                x = flow(x, x_mask, g=g, reverse=reverse)
+        return x
+
+
+class LatentEncoder(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+
+        self.pre = nn.Conv1d(in_channels, hidden_channels, 1)
+        self.enc = modules.WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            gin_channels=gin_channels,
+        )
+        self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
+
+    def forward(self, x, x_lengths, g=None, tau=1.0):
+        x_mask = torch.unsqueeze(commons.create_length_mask(x_lengths, x.size(2)), 1).to(
+            x.dtype
+        )
+        x = self.pre(x) * x_mask
+        x = self.enc(x, x_mask, g=g)
+        stats = self.proj(x) * x_mask
+        m, logs = torch.split(stats, self.out_channels, dim=1)
+        z = (m + torch.randn_like(m) * tau * torch.exp(logs)) * x_mask
+        return z, m, logs, x_mask
+
+
+class WaveformDecoder(torch.nn.Module):
+    def __init__(
+        self,
+        initial_channel,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels=0,
+    ):
+        super(WaveformDecoder, self).__init__()
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+        self.conv_pre = Conv1d(
+            initial_channel, upsample_initial_channel, 7, 1, padding=3
+        )
+        resblock = modules.ConvResBlock if resblock == "1" else modules.ConvResBlockLight
+
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                weight_norm(
+                    ConvTranspose1d(
+                        upsample_initial_channel // (2**i),
+                        upsample_initial_channel // (2 ** (i + 1)),
+                        k,
+                        u,
+                        padding=(k - u) // 2,
+                    )
+                )
+            )
+
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(resblock(ch, k, d))
+
+        self.conv_post = Conv1d(ch, 1, 7, 1, padding=3, bias=False)
+        self.ups.apply(initialize_weights)
+
+        if gin_channels != 0:
+            self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
+
+    def forward(self, x, g=None):
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.cond(g)
+
+        for i in range(self.num_upsamples):
+            x = F.leaky_relu(x, modules.LRELU_SLOPE)
+            x = self.ups[i](x)
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = F.leaky_relu(x)
+        x = self.conv_post(x)
+        x = torch.tanh(x)
+
+        return x
+
+    def remove_weight_norm(self):
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
+
+
+class StyleEncoder(nn.Module):
+    def __init__(self, spec_channels, gin_channels=0, layernorm=False):
+        super().__init__()
+        self.spec_channels = spec_channels
+        ref_enc_filters = [32, 32, 64, 64, 128, 128]
+        K = len(ref_enc_filters)
+        filters = [1] + ref_enc_filters
+        convs = [
+            weight_norm(
+                nn.Conv2d(
+                    in_channels=filters[i],
+                    out_channels=filters[i + 1],
+                    kernel_size=(3, 3),
+                    stride=(2, 2),
+                    padding=(1, 1),
+                )
+            )
+            for i in range(K)
+        ]
+        self.convs = nn.ModuleList(convs)
+
+        out_channels = self.calculate_channels(spec_channels, 3, 2, 1, K)
+        self.gru = nn.GRU(
+            input_size=ref_enc_filters[-1] * out_channels,
+            hidden_size=256 // 2,
+            batch_first=True,
+        )
+        self.proj = nn.Linear(128, gin_channels)
+        if layernorm:
+            self.layernorm = nn.LayerNorm(self.spec_channels)
+        else:
+            self.layernorm = None
+
+    def forward(self, inputs, mask=None):
+        N = inputs.size(0)
+
+        out = inputs.view(N, 1, -1, self.spec_channels)
+        if self.layernorm is not None:
+            out = self.layernorm(out)
+
+        for conv in self.convs:
+            out = conv(out)
+            out = F.relu(out)
+
+        out = out.transpose(1, 2)
+        T = out.size(1)
+        N = out.size(0)
+        out = out.contiguous().view(N, T, -1)
+
+        self.gru.flatten_parameters()
+        memory, out = self.gru(out)
+
+        return self.proj(out.squeeze(0))
+
+    def calculate_channels(self, L, kernel_size, stride, pad, n_convs):
+        for i in range(n_convs):
+            L = (L - kernel_size + 2 * pad) // stride + 1
+        return L
+
+
+class VoiceSynthesizer(nn.Module):
+    """Voice synthesis model for inference."""
+
+    def __init__(
+        self,
+        n_vocab,
+        spec_channels,
+        segment_size,
+        inter_channels,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size,
+        p_dropout,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        n_speakers=256,
+        gin_channels=256,
+        use_sdp=True,
+        n_flow_layer=4,
+        n_layers_trans_flow=6,
+        flow_share_parameter=False,
+        use_transformer_flow=True,
+        use_vc=False,
+        num_languages=None,
+        num_tones=None,
+        norm_refenc=False,
+        **kwargs
+    ):
+        super().__init__()
+        self.n_vocab = n_vocab
+        self.spec_channels = spec_channels
+        self.inter_channels = inter_channels
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.resblock = resblock
+        self.resblock_kernel_sizes = resblock_kernel_sizes
+        self.resblock_dilation_sizes = resblock_dilation_sizes
+        self.upsample_rates = upsample_rates
+        self.upsample_initial_channel = upsample_initial_channel
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.segment_size = segment_size
+        self.n_speakers = n_speakers
+        self.gin_channels = gin_channels
+        self.n_layers_trans_flow = n_layers_trans_flow
+        self.use_spk_conditioned_encoder = kwargs.get(
+            "use_spk_conditioned_encoder", True
+        )
+        self.use_sdp = use_sdp
+        self.use_noise_scaled_mas = kwargs.get("use_noise_scaled_mas", False)
+        self.mas_noise_scale_initial = kwargs.get("mas_noise_scale_initial", 0.01)
+        self.noise_scale_delta = kwargs.get("noise_scale_delta", 2e-6)
+        self.current_mas_noise_scale = self.mas_noise_scale_initial
+        if self.use_spk_conditioned_encoder and gin_channels > 0:
+            self.enc_gin_channels = gin_channels
+        else:
+            self.enc_gin_channels = 0
+        self.enc_p = PhonemeEncoder(
+            n_vocab,
+            inter_channels,
+            hidden_channels,
+            filter_channels,
+            n_heads,
+            n_layers,
+            kernel_size,
+            p_dropout,
+            gin_channels=self.enc_gin_channels,
+            num_languages=num_languages,
+            num_tones=num_tones,
+        )
+        self.dec = WaveformDecoder(
+            inter_channels,
+            resblock,
+            resblock_kernel_sizes,
+            resblock_dilation_sizes,
+            upsample_rates,
+            upsample_initial_channel,
+            upsample_kernel_sizes,
+            gin_channels=gin_channels,
+        )
+        self.enc_q = LatentEncoder(
+            spec_channels,
+            inter_channels,
+            hidden_channels,
+            5,
+            1,
+            16,
+            gin_channels=gin_channels,
+        )
+        if use_transformer_flow:
+            self.flow = AttentionFlowBlock(
+                inter_channels,
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers_trans_flow,
+                5,
+                p_dropout,
+                n_flow_layer,
+                gin_channels=gin_channels,
+                share_parameter=flow_share_parameter,
+            )
+        else:
+            self.flow = FlowBlock(
+                inter_channels,
+                hidden_channels,
+                5,
+                1,
+                n_flow_layer,
+                gin_channels=gin_channels,
+            )
+        self.sdp = VariationalDurationModel(
+            hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels
+        )
+        self.dp = DurationEstimator(
+            hidden_channels, 256, 3, 0.5, gin_channels=gin_channels
+        )
+
+        if n_speakers > 0:
+            self.emb_g = nn.Embedding(n_speakers, gin_channels)
+        else:
+            self.ref_enc = StyleEncoder(spec_channels, gin_channels, layernorm=norm_refenc)
+        self.use_vc = use_vc
+
+    def infer(
+        self,
+        x,
+        x_lengths,
+        sid,
+        tone,
+        language,
+        bert,
+        ja_bert,
+        noise_scale=0.667,
+        length_scale=1,
+        noise_scale_w=0.8,
+        max_len=None,
+        sdp_ratio=0,
+        y=None,
+        g=None,
+    ):
+        if g is None:
+            if self.n_speakers > 0:
+                g = self.emb_g(sid).unsqueeze(-1)
+            else:
+                g = self.ref_enc(y.transpose(1, 2)).unsqueeze(-1)
+        if self.use_vc:
+            g_p = None
+        else:
+            g_p = g
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, tone, language, bert, ja_bert, g=g_p
+        )
+        logw = self.sdp(x, x_mask, g=g, reverse=True, noise_scale=noise_scale_w) * (
+            sdp_ratio
+        ) + self.dp(x, x_mask, g=g) * (1 - sdp_ratio)
+        w = torch.exp(logw) * x_mask * length_scale
+
+        w_ceil = torch.ceil(w)
+        y_lengths = torch.clamp_min(torch.sum(w_ceil, [1, 2]), 1).long()
+        y_mask = torch.unsqueeze(commons.create_length_mask(y_lengths, None), 1).to(
+            x_mask.dtype
+        )
+        attn_mask = torch.unsqueeze(x_mask, 2) * torch.unsqueeze(y_mask, -1)
+        attn = commons.compute_alignment_path(w_ceil, attn_mask)
+
+        m_p = torch.matmul(attn.squeeze(1), m_p.transpose(1, 2)).transpose(
+            1, 2
+        )
+        logs_p = torch.matmul(attn.squeeze(1), logs_p.transpose(1, 2)).transpose(
+            1, 2
+        )
+
+        z_p = m_p + torch.randn_like(m_p) * torch.exp(logs_p) * noise_scale
+        z = self.flow(z_p, y_mask, g=g, reverse=True)
+        o = self.dec((z * y_mask)[:, :, :max_len], g=g)
+        return o, attn, y_mask, (z, z_p, m_p, logs_p)

tiny_tts/nn/__init__.py

@@ -0,0 +1 @@
+# Neural network building blocks

tiny_tts/nn/attentions.py

@@ -0,0 +1,424 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from . import commons
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ChannelLayerNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+@torch.jit.script
+def gated_activation(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        window_size=4,
+        isflow=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+
+        self.cond_layer_idx = self.n_layers
+        if "gin_channels" in kwargs:
+            self.gin_channels = kwargs["gin_channels"]
+            if self.gin_channels != 0:
+                self.spk_emb_linear = nn.Linear(self.gin_channels, self.hidden_channels)
+                self.cond_layer_idx = (
+                    kwargs["cond_layer_idx"] if "cond_layer_idx" in kwargs else 2
+                )
+                assert (
+                    self.cond_layer_idx < self.n_layers
+                ), "cond_layer_idx should be less than n_layers"
+        self.drop = nn.Dropout(p_dropout)
+        self.attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+
+        for i in range(self.n_layers):
+            self.attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    window_size=window_size,
+                )
+            )
+            self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FeedForward(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                )
+            )
+            self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, g=None):
+        attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            if i == self.cond_layer_idx and g is not None:
+                g = self.spk_emb_linear(g.transpose(1, 2))
+                g = g.transpose(1, 2)
+                x = x + g
+                x = x * x_mask
+            y = self.attn_layers[i](x, x, attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        filter_channels,
+        n_heads,
+        n_layers,
+        kernel_size=1,
+        p_dropout=0.0,
+        proximal_bias=False,
+        proximal_init=True,
+        **kwargs
+    ):
+        super().__init__()
+        self.hidden_channels = hidden_channels
+        self.filter_channels = filter_channels
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+
+        self.drop = nn.Dropout(p_dropout)
+        self.self_attn_layers = nn.ModuleList()
+        self.norm_layers_0 = nn.ModuleList()
+        self.encdec_attn_layers = nn.ModuleList()
+        self.norm_layers_1 = nn.ModuleList()
+        self.ffn_layers = nn.ModuleList()
+        self.norm_layers_2 = nn.ModuleList()
+        for i in range(self.n_layers):
+            self.self_attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels,
+                    hidden_channels,
+                    n_heads,
+                    p_dropout=p_dropout,
+                    proximal_bias=proximal_bias,
+                    proximal_init=proximal_init,
+                )
+            )
+            self.norm_layers_0.append(ChannelLayerNorm(hidden_channels))
+            self.encdec_attn_layers.append(
+                MultiHeadSelfAttention(
+                    hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout
+                )
+            )
+            self.norm_layers_1.append(ChannelLayerNorm(hidden_channels))
+            self.ffn_layers.append(
+                FeedForward(
+                    hidden_channels,
+                    hidden_channels,
+                    filter_channels,
+                    kernel_size,
+                    p_dropout=p_dropout,
+                    causal=True,
+                )
+            )
+            self.norm_layers_2.append(ChannelLayerNorm(hidden_channels))
+
+    def forward(self, x, x_mask, h, h_mask):
+        self_attn_mask = commons.subsequent_mask(x_mask.size(2)).to(
+            device=x.device, dtype=x.dtype
+        )
+        encdec_attn_mask = h_mask.unsqueeze(2) * x_mask.unsqueeze(-1)
+        x = x * x_mask
+        for i in range(self.n_layers):
+            y = self.self_attn_layers[i](x, x, self_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_0[i](x + y)
+
+            y = self.encdec_attn_layers[i](x, h, encdec_attn_mask)
+            y = self.drop(y)
+            x = self.norm_layers_1[i](x + y)
+
+            y = self.ffn_layers[i](x, x_mask)
+            y = self.drop(y)
+            x = self.norm_layers_2[i](x + y)
+        x = x * x_mask
+        return x
+
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(
+        self,
+        channels,
+        out_channels,
+        n_heads,
+        p_dropout=0.0,
+        window_size=None,
+        heads_share=True,
+        block_length=None,
+        proximal_bias=False,
+        proximal_init=False,
+    ):
+        super().__init__()
+        assert channels % n_heads == 0
+
+        self.channels = channels
+        self.out_channels = out_channels
+        self.n_heads = n_heads
+        self.p_dropout = p_dropout
+        self.window_size = window_size
+        self.heads_share = heads_share
+        self.block_length = block_length
+        self.proximal_bias = proximal_bias
+        self.proximal_init = proximal_init
+        self.attn = None
+
+        self.k_channels = channels // n_heads
+        self.conv_q = nn.Conv1d(channels, channels, 1)
+        self.conv_k = nn.Conv1d(channels, channels, 1)
+        self.conv_v = nn.Conv1d(channels, channels, 1)
+        self.conv_o = nn.Conv1d(channels, out_channels, 1)
+        self.drop = nn.Dropout(p_dropout)
+
+        if window_size is not None:
+            n_heads_rel = 1 if heads_share else n_heads
+            rel_stddev = self.k_channels**-0.5
+            self.emb_rel_k = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+            self.emb_rel_v = nn.Parameter(
+                torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels)
+                * rel_stddev
+            )
+
+        nn.init.xavier_uniform_(self.conv_q.weight)
+        nn.init.xavier_uniform_(self.conv_k.weight)
+        nn.init.xavier_uniform_(self.conv_v.weight)
+        if proximal_init:
+            with torch.no_grad():
+                self.conv_k.weight.copy_(self.conv_q.weight)
+                self.conv_k.bias.copy_(self.conv_q.bias)
+
+    def forward(self, x, c, attn_mask=None):
+        q = self.conv_q(x)
+        k = self.conv_k(c)
+        v = self.conv_v(c)
+
+        x, self.attn = self.attention(q, k, v, mask=attn_mask)
+
+        x = self.conv_o(x)
+        return x
+
+    def attention(self, query, key, value, mask=None):
+        b, d, t_s, t_t = (*key.size(), query.size(2))
+        query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
+        key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+        value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
+
+        scores = torch.matmul(query / math.sqrt(self.k_channels), key.transpose(-2, -1))
+        if self.window_size is not None:
+            assert (
+                t_s == t_t
+            ), "Relative attention is only available for self-attention."
+            key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
+            rel_logits = self._matmul_with_relative_keys(
+                query / math.sqrt(self.k_channels), key_relative_embeddings
+            )
+            scores_local = self._relative_position_to_absolute_position(rel_logits)
+            scores = scores + scores_local
+        if self.proximal_bias:
+            assert t_s == t_t, "Proximal bias is only available for self-attention."
+            scores = scores + self._attention_bias_proximal(t_s).to(
+                device=scores.device, dtype=scores.dtype
+            )
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e4)
+            if self.block_length is not None:
+                assert (
+                    t_s == t_t
+                ), "Local attention is only available for self-attention."
+                block_mask = (
+                    torch.ones_like(scores)
+                    .triu(-self.block_length)
+                    .tril(self.block_length)
+                )
+                scores = scores.masked_fill(block_mask == 0, -1e4)
+        p_attn = F.softmax(scores, dim=-1)
+        p_attn = self.drop(p_attn)
+        output = torch.matmul(p_attn, value)
+        if self.window_size is not None:
+            relative_weights = self._absolute_position_to_relative_position(p_attn)
+            value_relative_embeddings = self._get_relative_embeddings(
+                self.emb_rel_v, t_s
+            )
+            output = output + self._matmul_with_relative_values(
+                relative_weights, value_relative_embeddings
+            )
+        output = (
+            output.transpose(2, 3).contiguous().view(b, d, t_t)
+        )
+        return output, p_attn
+
+    def _matmul_with_relative_values(self, x, y):
+        ret = torch.matmul(x, y.unsqueeze(0))
+        return ret
+
+    def _matmul_with_relative_keys(self, x, y):
+        ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))
+        return ret
+
+    def _get_relative_embeddings(self, relative_embeddings, length):
+        2 * self.window_size + 1
+        pad_length = max(length - (self.window_size + 1), 0)
+        slice_start_position = max((self.window_size + 1) - length, 0)
+        slice_end_position = slice_start_position + 2 * length - 1
+        if pad_length > 0:
+            padded_relative_embeddings = F.pad(
+                relative_embeddings,
+                commons.flatten_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]),
+            )
+        else:
+            padded_relative_embeddings = relative_embeddings
+        used_relative_embeddings = padded_relative_embeddings[
+            :, slice_start_position:slice_end_position
+        ]
+        return used_relative_embeddings
+
+    def _relative_position_to_absolute_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = F.pad(x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))
+
+        x_flat = x.view([batch, heads, length * 2 * length])
+        x_flat = F.pad(
+            x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [0, length - 1]])
+        )
+
+        x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[
+            :, :, :length, length - 1 :
+        ]
+        return x_final
+
+    def _absolute_position_to_relative_position(self, x):
+        batch, heads, length, _ = x.size()
+        x = F.pad(
+            x, commons.flatten_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]])
+        )
+        x_flat = x.view([batch, heads, length**2 + length * (length - 1)])
+        x_flat = F.pad(x_flat, commons.flatten_pad_shape([[0, 0], [0, 0], [length, 0]]))
+        x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]
+        return x_final
+
+    def _attention_bias_proximal(self, length):
+        r = torch.arange(length, dtype=torch.float32)
+        diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)
+        return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        filter_channels,
+        kernel_size,
+        p_dropout=0.0,
+        activation=None,
+        causal=False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.p_dropout = p_dropout
+        self.activation = activation
+        self.causal = causal
+
+        if causal:
+            self.padding = self._causal_padding
+        else:
+            self.padding = self._same_padding
+
+        self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size)
+        self.conv_2 = nn.Conv1d(filter_channels, out_channels, kernel_size)
+        self.drop = nn.Dropout(p_dropout)
+
+    def forward(self, x, x_mask):
+        x = self.conv_1(self.padding(x * x_mask))
+        if self.activation == "gelu":
+            x = x * torch.sigmoid(1.702 * x)
+        else:
+            x = torch.relu(x)
+        x = self.drop(x)
+        x = self.conv_2(self.padding(x * x_mask))
+        return x * x_mask
+
+    def _causal_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = self.kernel_size - 1
+        pad_r = 0
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.flatten_pad_shape(padding))
+        return x
+
+    def _same_padding(self, x):
+        if self.kernel_size == 1:
+            return x
+        pad_l = (self.kernel_size - 1) // 2
+        pad_r = self.kernel_size // 2
+        padding = [[0, 0], [0, 0], [pad_l, pad_r]]
+        x = F.pad(x, commons.flatten_pad_shape(padding))
+        return x

tiny_tts/nn/commons.py

@@ -0,0 +1,151 @@
+import math
+import torch
+from torch.nn import functional as F
+
+
+def initialize_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def compute_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+def flatten_pad_shape(pad_shape):
+    layer = pad_shape[::-1]
+    pad_shape = [item for sublist in layer for item in sublist]
+    return pad_shape
+
+
+def insert_blanks(lst, item):
+    result = [item] * (len(lst) * 2 + 1)
+    result[1::2] = lst
+    return result
+
+
+def kl_divergence(m_p, logs_p, m_q, logs_q):
+    kl = (logs_q - logs_p) - 0.5
+    kl += (
+        0.5 * (torch.exp(2.0 * logs_p) + ((m_p - m_q) ** 2)) * torch.exp(-2.0 * logs_q)
+    )
+    return kl
+
+
+def rand_gumbel(shape):
+    uniform_samples = torch.rand(shape) * 0.99998 + 0.00001
+    return -torch.log(-torch.log(uniform_samples))
+
+
+def rand_gumbel_like(x):
+    g = rand_gumbel(x.size()).to(dtype=x.dtype, device=x.device)
+    return g
+
+
+def extract_segments(x, ids_str, segment_size=4):
+    ret = torch.zeros_like(x[:, :, :segment_size])
+    for i in range(x.size(0)):
+        idx_str = max(0, ids_str[i].item())
+        idx_end = idx_str + segment_size
+        available = x.size(2) - idx_str
+        if available >= segment_size:
+            ret[i] = x[i, :, idx_str:idx_end]
+        elif available > 0:
+            ret[i, :, :available] = x[i, :, idx_str:idx_str + available]
+    return ret
+
+
+def random_segments(x, x_lengths=None, segment_size=4):
+    b, d, t = x.size()
+    if x_lengths is None:
+        x_lengths = t
+    ids_str_max = torch.clamp(x_lengths - segment_size + 1, min=0)
+    ids_str = (torch.rand([b]).to(device=x.device) * ids_str_max).to(dtype=torch.long)
+    ret = extract_segments(x, ids_str, segment_size)
+    return ret, ids_str
+
+
+def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4):
+    position = torch.arange(length, dtype=torch.float)
+    num_timescales = channels // 2
+    log_timescale_increment = math.log(float(max_timescale) / float(min_timescale)) / (
+        num_timescales - 1
+    )
+    inv_timescales = min_timescale * torch.exp(
+        torch.arange(num_timescales, dtype=torch.float) * -log_timescale_increment
+    )
+    scaled_time = position.unsqueeze(0) * inv_timescales.unsqueeze(1)
+    signal = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], 0)
+    signal = F.pad(signal, [0, 0, 0, channels % 2])
+    signal = signal.view(1, channels, length)
+    return signal
+
+
+def add_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return x + signal.to(dtype=x.dtype, device=x.device)
+
+
+def cat_timing_signal_1d(x, min_timescale=1.0, max_timescale=1.0e4, axis=1):
+    b, channels, length = x.size()
+    signal = get_timing_signal_1d(length, channels, min_timescale, max_timescale)
+    return torch.cat([x, signal.to(dtype=x.dtype, device=x.device)], axis)
+
+
+def subsequent_mask(length):
+    mask = torch.tril(torch.ones(length, length)).unsqueeze(0).unsqueeze(0)
+    return mask
+
+
+@torch.jit.script
+def gated_activation(input_a, input_b, n_channels):
+    n_channels_int = n_channels[0]
+    in_act = input_a + input_b
+    t_act = torch.tanh(in_act[:, :n_channels_int, :])
+    s_act = torch.sigmoid(in_act[:, n_channels_int:, :])
+    acts = t_act * s_act
+    return acts
+
+
+def shift_1d(x):
+    x = F.pad(x, flatten_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]
+    return x
+
+
+def create_length_mask(length, max_length=None):
+    if max_length is None:
+        max_length = length.max()
+    x = torch.arange(max_length, dtype=length.dtype, device=length.device)
+    return x.unsqueeze(0) < length.unsqueeze(1)
+
+
+def compute_alignment_path(duration, mask):
+    b, _, t_y, t_x = mask.shape
+    cum_duration = torch.cumsum(duration, -1)
+
+    cum_duration_flat = cum_duration.view(b * t_x)
+    path = create_length_mask(cum_duration_flat, t_y).to(mask.dtype)
+    path = path.view(b, t_x, t_y)
+    path = path - F.pad(path, flatten_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
+    path = path.unsqueeze(1).transpose(2, 3) * mask
+    return path
+
+
+def clip_grad_value_(parameters, clip_value, norm_type=2):
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(filter(lambda p: p.grad is not None, parameters))
+    norm_type = float(norm_type)
+    if clip_value is not None:
+        clip_value = float(clip_value)
+
+    total_norm = 0
+    for p in parameters:
+        param_norm = p.grad.data.norm(norm_type)
+        total_norm += param_norm.item() ** norm_type
+        if clip_value is not None:
+            p.grad.data.clamp_(min=-clip_value, max=clip_value)
+    total_norm = total_norm ** (1.0 / norm_type)
+    return total_norm

tiny_tts/nn/modules.py

@@ -0,0 +1,578 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from torch.nn import Conv1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from . import commons
+from .commons import initialize_weights, compute_padding
+from .transforms import spline_transform
+from .attentions import TransformerBlock
+
+LRELU_SLOPE = 0.1
+
+
+class ChannelNorm(nn.Module):
+    def __init__(self, channels, eps=1e-5):
+        super().__init__()
+        self.channels = channels
+        self.eps = eps
+
+        self.gamma = nn.Parameter(torch.ones(channels))
+        self.beta = nn.Parameter(torch.zeros(channels))
+
+    def forward(self, x):
+        x = x.transpose(1, -1)
+        x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+        return x.transpose(1, -1)
+
+
+class ConvReluNorm(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        kernel_size,
+        n_layers,
+        p_dropout,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+        assert n_layers > 1, "Number of layers should be larger than 0."
+
+        self.conv_layers = nn.ModuleList()
+        self.norm_layers = nn.ModuleList()
+        self.conv_layers.append(
+            nn.Conv1d(
+                in_channels, hidden_channels, kernel_size, padding=kernel_size // 2
+            )
+        )
+        self.norm_layers.append(ChannelNorm(hidden_channels))
+        self.relu_drop = nn.Sequential(nn.ReLU(), nn.Dropout(p_dropout))
+        for _ in range(n_layers - 1):
+            self.conv_layers.append(
+                nn.Conv1d(
+                    hidden_channels,
+                    hidden_channels,
+                    kernel_size,
+                    padding=kernel_size // 2,
+                )
+            )
+            self.norm_layers.append(ChannelNorm(hidden_channels))
+        self.proj = nn.Conv1d(hidden_channels, out_channels, 1)
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask):
+        x_org = x
+        for i in range(self.n_layers):
+            x = self.conv_layers[i](x * x_mask)
+            x = self.norm_layers[i](x)
+            x = self.relu_drop(x)
+        x = x_org + self.proj(x)
+        return x * x_mask
+
+
+class DepthwiseSepConv(nn.Module):
+    """Dilated and Depth-Separable Convolution"""
+
+    def __init__(self, channels, kernel_size, n_layers, p_dropout=0.0):
+        super().__init__()
+        self.channels = channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.p_dropout = p_dropout
+
+        self.drop = nn.Dropout(p_dropout)
+        self.convs_sep = nn.ModuleList()
+        self.convs_1x1 = nn.ModuleList()
+        self.norms_1 = nn.ModuleList()
+        self.norms_2 = nn.ModuleList()
+        for i in range(n_layers):
+            dilation = kernel_size**i
+            padding = (kernel_size * dilation - dilation) // 2
+            self.convs_sep.append(
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    groups=channels,
+                    dilation=dilation,
+                    padding=padding,
+                )
+            )
+            self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
+            self.norms_1.append(ChannelNorm(channels))
+            self.norms_2.append(ChannelNorm(channels))
+
+    def forward(self, x, x_mask, g=None):
+        if g is not None:
+            x = x + g
+        for i in range(self.n_layers):
+            y = self.convs_sep[i](x * x_mask)
+            y = self.norms_1[i](y)
+            y = F.gelu(y)
+            y = self.convs_1x1[i](y)
+            y = self.norms_2[i](y)
+            y = F.gelu(y)
+            y = self.drop(y)
+            x = x + y
+        return x * x_mask
+
+
+class WaveNet(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        gin_channels=0,
+        p_dropout=0,
+    ):
+        super(WaveNet, self).__init__()
+        assert kernel_size % 2 == 1
+        self.hidden_channels = hidden_channels
+        self.kernel_size = (kernel_size,)
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.gin_channels = gin_channels
+        self.p_dropout = p_dropout
+
+        self.in_layers = torch.nn.ModuleList()
+        self.res_skip_layers = torch.nn.ModuleList()
+        self.drop = nn.Dropout(p_dropout)
+
+        if gin_channels != 0:
+            cond_layer = torch.nn.Conv1d(
+                gin_channels, 2 * hidden_channels * n_layers, 1
+            )
+            self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name="weight")
+
+        for i in range(n_layers):
+            dilation = dilation_rate**i
+            padding = int((kernel_size * dilation - dilation) / 2)
+            in_layer = torch.nn.Conv1d(
+                hidden_channels,
+                2 * hidden_channels,
+                kernel_size,
+                dilation=dilation,
+                padding=padding,
+            )
+            in_layer = torch.nn.utils.weight_norm(in_layer, name="weight")
+            self.in_layers.append(in_layer)
+
+            if i < n_layers - 1:
+                res_skip_channels = 2 * hidden_channels
+            else:
+                res_skip_channels = hidden_channels
+
+            res_skip_layer = torch.nn.Conv1d(hidden_channels, res_skip_channels, 1)
+            res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name="weight")
+            self.res_skip_layers.append(res_skip_layer)
+
+    def forward(self, x, x_mask, g=None, **kwargs):
+        output = torch.zeros_like(x)
+        n_channels_tensor = torch.IntTensor([self.hidden_channels])
+
+        if g is not None:
+            g = self.cond_layer(g)
+
+        for i in range(self.n_layers):
+            x_in = self.in_layers[i](x)
+            if g is not None:
+                cond_offset = i * 2 * self.hidden_channels
+                g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :]
+            else:
+                g_l = torch.zeros_like(x_in)
+
+            acts = commons.gated_activation(x_in, g_l, n_channels_tensor)
+            acts = self.drop(acts)
+
+            res_skip_acts = self.res_skip_layers[i](acts)
+            if i < self.n_layers - 1:
+                res_acts = res_skip_acts[:, : self.hidden_channels, :]
+                x = (x + res_acts) * x_mask
+                output = output + res_skip_acts[:, self.hidden_channels :, :]
+            else:
+                output = output + res_skip_acts
+        return output * x_mask
+
+    def remove_weight_norm(self):
+        if self.gin_channels != 0:
+            torch.nn.utils.remove_weight_norm(self.cond_layer)
+        for l in self.in_layers:
+            torch.nn.utils.remove_weight_norm(l)
+        for l in self.res_skip_layers:
+            torch.nn.utils.remove_weight_norm(l)
+
+
+class ConvResBlock(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ConvResBlock, self).__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=compute_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=compute_padding(kernel_size, dilation[1]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[2],
+                        padding=compute_padding(kernel_size, dilation[2]),
+                    )
+                ),
+            ]
+        )
+        self.convs1.apply(initialize_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=1,
+                        padding=compute_padding(kernel_size, 1),
+                    )
+                ),
+            ]
+        )
+        self.convs2.apply(initialize_weights)
+
+    def forward(self, x, x_mask=None):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c2(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ConvResBlockLight(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
+        super(ConvResBlockLight, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[0],
+                        padding=compute_padding(kernel_size, dilation[0]),
+                    )
+                ),
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=dilation[1],
+                        padding=compute_padding(kernel_size, dilation[1]),
+                    )
+                ),
+            ]
+        )
+        self.convs.apply(initialize_weights)
+
+    def forward(self, x, x_mask=None):
+        for c in self.convs:
+            xt = F.leaky_relu(x, LRELU_SLOPE)
+            if x_mask is not None:
+                xt = xt * x_mask
+            xt = c(xt)
+            x = xt + x
+        if x_mask is not None:
+            x = x * x_mask
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs:
+            remove_weight_norm(l)
+
+
+class LogTransform(nn.Module):
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask
+            logdet = torch.sum(-y, [1, 2])
+            return y, logdet
+        else:
+            x = torch.exp(x) * x_mask
+            return x
+
+
+class FlipTransform(nn.Module):
+    def forward(self, x, *args, reverse=False, **kwargs):
+        x = torch.flip(x, [1])
+        if not reverse:
+            logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
+            return x, logdet
+        else:
+            return x
+
+
+class AffineCoupling(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.channels = channels
+        self.m = nn.Parameter(torch.zeros(channels, 1))
+        self.logs = nn.Parameter(torch.zeros(channels, 1))
+
+    def forward(self, x, x_mask, reverse=False, **kwargs):
+        if not reverse:
+            y = self.m + torch.exp(self.logs) * x
+            y = y * x_mask
+            logdet = torch.sum(self.logs * x_mask, [1, 2])
+            return y, logdet
+        else:
+            x = (x - self.m) * torch.exp(-self.logs) * x_mask
+            return x
+
+
+class FlowCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        dilation_rate,
+        n_layers,
+        p_dropout=0,
+        gin_channels=0,
+        mean_only=False,
+    ):
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.dilation_rate = dilation_rate
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = WaveNet(
+            hidden_channels,
+            kernel_size,
+            dilation_rate,
+            n_layers,
+            p_dropout=p_dropout,
+            gin_channels=gin_channels,
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x
+
+
+class ConvolutionalFlow(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        filter_channels,
+        kernel_size,
+        n_layers,
+        num_bins=10,
+        tail_bound=5.0,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.filter_channels = filter_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.num_bins = num_bins
+        self.tail_bound = tail_bound
+        self.half_channels = in_channels // 2
+
+        self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
+        self.convs = DepthwiseSepConv(filter_channels, kernel_size, n_layers, p_dropout=0.0)
+        self.proj = nn.Conv1d(
+            filter_channels, self.half_channels * (num_bins * 3 - 1), 1
+        )
+        self.proj.weight.data.zero_()
+        self.proj.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0)
+        h = self.convs(h, x_mask, g=g)
+        h = self.proj(h) * x_mask
+
+        b, c, t = x0.shape
+        h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2)
+
+        unnormalized_widths = h[..., : self.num_bins] / math.sqrt(self.filter_channels)
+        unnormalized_heights = h[..., self.num_bins : 2 * self.num_bins] / math.sqrt(
+            self.filter_channels
+        )
+        unnormalized_derivatives = h[..., 2 * self.num_bins :]
+
+        x1, logabsdet = spline_transform(
+            x1,
+            unnormalized_widths,
+            unnormalized_heights,
+            unnormalized_derivatives,
+            inverse=reverse,
+            tails="linear",
+            tail_bound=self.tail_bound,
+        )
+
+        x = torch.cat([x0, x1], 1) * x_mask
+        logdet = torch.sum(logabsdet * x_mask, [1, 2])
+        if not reverse:
+            return x, logdet
+        else:
+            return x
+
+
+class TransformerCouplingLayer(nn.Module):
+    def __init__(
+        self,
+        channels,
+        hidden_channels,
+        kernel_size,
+        n_layers,
+        n_heads,
+        p_dropout=0,
+        filter_channels=0,
+        mean_only=False,
+        wn_sharing_parameter=None,
+        gin_channels=0,
+    ):
+        assert n_layers == 3, n_layers
+        assert channels % 2 == 0, "channels should be divisible by 2"
+        super().__init__()
+        self.channels = channels
+        self.hidden_channels = hidden_channels
+        self.kernel_size = kernel_size
+        self.n_layers = n_layers
+        self.half_channels = channels // 2
+        self.mean_only = mean_only
+
+        self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
+        self.enc = (
+            TransformerBlock(
+                hidden_channels,
+                filter_channels,
+                n_heads,
+                n_layers,
+                kernel_size,
+                p_dropout,
+                isflow=True,
+                gin_channels=gin_channels,
+            )
+            if wn_sharing_parameter is None
+            else wn_sharing_parameter
+        )
+        self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
+        self.post.weight.data.zero_()
+        self.post.bias.data.zero_()
+
+    def forward(self, x, x_mask, g=None, reverse=False):
+        x0, x1 = torch.split(x, [self.half_channels] * 2, 1)
+        h = self.pre(x0) * x_mask
+        h = self.enc(h, x_mask, g=g)
+        stats = self.post(h) * x_mask
+        if not self.mean_only:
+            m, logs = torch.split(stats, [self.half_channels] * 2, 1)
+        else:
+            m = stats
+            logs = torch.zeros_like(m)
+
+        if not reverse:
+            x1 = m + x1 * torch.exp(logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            logdet = torch.sum(logs, [1, 2])
+            return x, logdet
+        else:
+            x1 = (x1 - m) * torch.exp(-logs) * x_mask
+            x = torch.cat([x0, x1], 1)
+            return x

tiny_tts/nn/transforms.py

@@ -0,0 +1,209 @@
+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def spline_transform(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails=None,
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if tails is None:
+        spline_fn = quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unbounded_spline
+        spline_kwargs = {"tails": tails, "tail_bound": tail_bound}
+
+    outputs, logabsdet = spline_fn(
+        inputs=inputs,
+        unnormalized_widths=unnormalized_widths,
+        unnormalized_heights=unnormalized_heights,
+        unnormalized_derivatives=unnormalized_derivatives,
+        inverse=inverse,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+        **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(inputs[..., None] >= bin_locations, dim=-1) - 1
+
+
+def unbounded_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    tails="linear",
+    tail_bound=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == "linear":
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError("{} tails are not implemented.".format(tails))
+
+    (
+        outputs[inside_interval_mask],
+        logabsdet[inside_interval_mask],
+    ) = quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound,
+        right=tail_bound,
+        bottom=-tail_bound,
+        top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative,
+    )
+
+    return outputs, logabsdet
+
+
+def quadratic_spline(
+    inputs,
+    unnormalized_widths,
+    unnormalized_heights,
+    unnormalized_derivatives,
+    inverse=False,
+    left=0.0,
+    right=1.0,
+    bottom=0.0,
+    top=1.0,
+    min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+    min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+    min_derivative=DEFAULT_MIN_DERIVATIVE,
+):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError("Input to a transform is not within its domain")
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError("Minimal bin width too large for the number of bins")
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError("Minimal bin height too large for the number of bins")
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode="constant", value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode="constant", value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        ) + input_heights * (input_delta - input_derivatives)
+        b = input_heights * input_derivatives - (inputs - input_cumheights) * (
+            input_derivatives + input_derivatives_plus_one - 2 * input_delta
+        )
+        c = -input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * root.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - root).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (
+            input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta
+        )
+        denominator = input_delta + (
+            (input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+            * theta_one_minus_theta
+        )
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (
+            input_derivatives_plus_one * theta.pow(2)
+            + 2 * input_delta * theta_one_minus_theta
+            + input_derivatives * (1 - theta).pow(2)
+        )
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet

tiny_tts/text/__init__.py

@@ -0,0 +1,19 @@
+from .symbols import *
+
+
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+
+
+def phonemes_to_ids(cleaned_text, tones, language, symbol_to_id=None):
+    """Converts a list of phoneme symbols to a sequence of integer IDs."""
+    symbol_to_id_map = symbol_to_id if symbol_to_id else _symbol_to_id
+    unk_id = symbol_to_id_map.get("UNK")
+    if unk_id is None:
+        phones = [symbol_to_id_map[symbol] for symbol in cleaned_text]
+    else:
+        phones = [symbol_to_id_map.get(symbol, unk_id) for symbol in cleaned_text]
+    tone_start = language_tone_start_map[language]
+    tones = [i + tone_start for i in tones]
+    lang_id = language_id_map[language]
+    lang_ids = [lang_id for _ in phones]
+    return phones, tones, lang_ids

tiny_tts/text/english.py

@@ -0,0 +1,173 @@
+import pickle
+import os
+import re
+from g2p_en import G2p
+
+from . import symbols
+
+from .english_utils.abbreviations import expand_abbreviations
+from .english_utils.time_norm import expand_time_english
+from .english_utils.number_norm import normalize_numbers
+
+
+def distribute_phone(n_phone, n_word):
+    phones_per_word = [0] * n_word
+    for task in range(n_phone):
+        min_tasks = min(phones_per_word)
+        min_indices = [
+            i for i, x in enumerate(phones_per_word) if x == min_tasks
+        ]
+        chosen_index = min_indices[len(min_indices) // 2]
+        phones_per_word[chosen_index] += 1
+    return phones_per_word
+
+
+from transformers import AutoTokenizer
+
+current_file_path = os.path.dirname(__file__)
+CMU_DICT_PATH = os.path.join(current_file_path, "cmudict.rep")
+CACHE_PATH = os.path.join(current_file_path, "cmudict_cache.pickle")
+_g2p = G2p()
+
+arpa = {
+    "AH0", "S", "AH1", "EY2", "AE2", "EH0", "OW2", "UH0", "NG", "B",
+    "G", "AY0", "M", "AA0", "F", "AO0", "ER2", "UH1", "IY1", "AH2",
+    "DH", "IY0", "EY1", "IH0", "K", "N", "W", "IY2", "T", "AA1",
+    "ER1", "EH2", "OY0", "UH2", "UW1", "Z", "AW2", "AW1", "V", "UW2",
+    "AA2", "ER", "AW0", "UW0", "R", "OW1", "EH1", "ZH", "AE0", "IH2",
+    "IH", "Y", "JH", "P", "AY1", "EY0", "OY2", "TH", "HH", "D",
+    "ER0", "CH", "AO1", "AE1", "AO2", "OY1", "AY2", "IH1", "OW0", "L", "SH",
+}
+
+
+def map_phoneme(ph):
+    rep_map = {
+        "：": ",", "；": ",", "，": ",", "。": ".", "！": "!",
+        "？": "?", "\n": ".", "·": ",", "、": ",", "...": "…", "v": "V",
+    }
+    if ph in rep_map.keys():
+        ph = rep_map[ph]
+    if ph in symbols:
+        return ph
+    if ph not in symbols:
+        ph = "UNK"
+    return ph
+
+
+def read_dict():
+    g2p_dict = {}
+    start_line = 49
+    with open(CMU_DICT_PATH) as f:
+        line = f.readline()
+        line_index = 1
+        while line:
+            if line_index >= start_line:
+                line = line.strip()
+                word_split = line.split("  ")
+                word = word_split[0]
+
+                syllable_split = word_split[1].split(" - ")
+                g2p_dict[word] = []
+                for syllable in syllable_split:
+                    phone_split = syllable.split(" ")
+                    g2p_dict[word].append(phone_split)
+
+            line_index = line_index + 1
+            line = f.readline()
+
+    return g2p_dict
+
+
+def cache_dict(g2p_dict, file_path):
+    with open(file_path, "wb") as pickle_file:
+        pickle.dump(g2p_dict, pickle_file)
+
+
+def get_dict():
+    if os.path.exists(CACHE_PATH):
+        with open(CACHE_PATH, "rb") as pickle_file:
+            g2p_dict = pickle.load(pickle_file)
+    else:
+        g2p_dict = read_dict()
+        cache_dict(g2p_dict, CACHE_PATH)
+
+    return g2p_dict
+
+
+eng_dict = get_dict()
+
+
+def parse_phoneme(phn):
+    tone = 0
+    if re.search(r"\d$", phn):
+        tone = int(phn[-1]) + 1
+        phn = phn[:-1]
+    return phn.lower(), tone
+
+
+def parse_syllables(syllables):
+    tones = []
+    phonemes = []
+    for phn_list in syllables:
+        for i in range(len(phn_list)):
+            phn = phn_list[i]
+            phn, tone = parse_phoneme(phn)
+            phonemes.append(phn)
+            tones.append(tone)
+    return phonemes, tones
+
+
+def normalize_text(text):
+    text = text.lower()
+    text = expand_time_english(text)
+    text = normalize_numbers(text)
+    text = expand_abbreviations(text)
+    return text
+
+
+model_id = 'bert-base-uncased'
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+
+def grapheme_to_phoneme(text, pad_start_end=True, tokenized=None):
+    if tokenized is None:
+        tokenized = tokenizer.tokenize(text)
+    ph_groups = []
+    for t in tokenized:
+        if not t.startswith("#"):
+            ph_groups.append([t])
+        else:
+            ph_groups[-1].append(t.replace("#", ""))
+
+    phones = []
+    tones = []
+    word2ph = []
+    for group in ph_groups:
+        w = "".join(group)
+        phone_len = 0
+        word_len = len(group)
+        if w.upper() in eng_dict:
+            phns, tns = parse_syllables(eng_dict[w.upper()])
+            phones += phns
+            tones += tns
+            phone_len += len(phns)
+        else:
+            phone_list = list(filter(lambda p: p != " ", _g2p(w)))
+            for ph in phone_list:
+                if ph in arpa:
+                    ph, tn = parse_phoneme(ph)
+                    phones.append(ph)
+                    tones.append(tn)
+                else:
+                    phones.append(ph)
+                    tones.append(0)
+                phone_len += 1
+        aaa = distribute_phone(phone_len, word_len)
+        word2ph += aaa
+    phones = [map_phoneme(i) for i in phones]
+
+    if pad_start_end:
+        phones = ["_"] + phones + ["_"]
+        tones = [0] + tones + [0]
+        word2ph = [1] + word2ph + [1]
+    return phones, tones, word2ph

tiny_tts/text/english_utils/abbreviations.py

@@ -0,0 +1,35 @@
+import re
+
+# List of (regular expression, replacement) pairs for abbreviations in english:
+abbreviations_en = [
+    (re.compile("\\b%s\\." % x[0], re.IGNORECASE), x[1])
+    for x in [
+        ("mrs", "misess"),
+        ("mr", "mister"),
+        ("dr", "doctor"),
+        ("st", "saint"),
+        ("co", "company"),
+        ("jr", "junior"),
+        ("maj", "major"),
+        ("gen", "general"),
+        ("drs", "doctors"),
+        ("rev", "reverend"),
+        ("lt", "lieutenant"),
+        ("hon", "honorable"),
+        ("sgt", "sergeant"),
+        ("capt", "captain"),
+        ("esq", "esquire"),
+        ("ltd", "limited"),
+        ("col", "colonel"),
+        ("ft", "fort"),
+    ]
+]
+
+def expand_abbreviations(text, lang="en"):
+    if lang == "en":
+        _abbreviations = abbreviations_en
+    else:
+        raise NotImplementedError()
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text

tiny_tts/text/english_utils/number_norm.py

@@ -0,0 +1,97 @@
+""" from https://github.com/keithito/tacotron """
+
+import re
+from typing import Dict
+
+import inflect
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])")
+_decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)")
+_currency_re = re.compile(r"(£|\$|¥)([0-9\,\.]*[0-9]+)")
+_ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)")
+_number_re = re.compile(r"-?[0-9]+")
+
+
+def _remove_commas(m):
+    return m.group(1).replace(",", "")
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace(".", " point ")
+
+
+def __expand_currency(value: str, inflection: Dict[float, str]) -> str:
+    parts = value.replace(",", "").split(".")
+    if len(parts) > 2:
+        return f"{value} {inflection[2]}"  # Unexpected format
+    text = []
+    integer = int(parts[0]) if parts[0] else 0
+    if integer > 0:
+        integer_unit = inflection.get(integer, inflection[2])
+        text.append(f"{integer} {integer_unit}")
+    fraction = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if fraction > 0:
+        fraction_unit = inflection.get(fraction / 100, inflection[0.02])
+        text.append(f"{fraction} {fraction_unit}")
+    if len(text) == 0:
+        return f"zero {inflection[2]}"
+    return " ".join(text)
+
+
+def _expand_currency(m: "re.Match") -> str:
+    currencies = {
+        "$": {
+            0.01: "cent",
+            0.02: "cents",
+            1: "dollar",
+            2: "dollars",
+        },
+        "€": {
+            0.01: "cent",
+            0.02: "cents",
+            1: "euro",
+            2: "euros",
+        },
+        "£": {
+            0.01: "penny",
+            0.02: "pence",
+            1: "pound sterling",
+            2: "pounds sterling",
+        },
+        "¥": {
+            # TODO rin
+            0.02: "sen",
+            2: "yen",
+        },
+    }
+    unit = m.group(1)
+    currency = currencies[unit]
+    value = m.group(2)
+    return __expand_currency(value, currency)
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    num = int(m.group(0))
+    if 1000 < num < 3000:
+        if num == 2000:
+            return "two thousand"
+        if 2000 < num < 2010:
+            return "two thousand " + _inflect.number_to_words(num % 100)
+        if num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + " hundred"
+        return _inflect.number_to_words(num, andword="", zero="oh", group=2).replace(", ", " ")
+    return _inflect.number_to_words(num, andword="")
+
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_currency_re, _expand_currency, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    return text

tiny_tts/text/english_utils/time_norm.py

@@ -0,0 +1,47 @@
+import re
+
+import inflect
+
+_inflect = inflect.engine()
+
+_time_re = re.compile(
+    r"""\b
+                          ((0?[0-9])|(1[0-1])|(1[2-9])|(2[0-3]))  # hours
+                          :
+                          ([0-5][0-9])                            # minutes
+                          \s*(a\\.m\\.|am|pm|p\\.m\\.|a\\.m|p\\.m)? # am/pm
+                          \b""",
+    re.IGNORECASE | re.X,
+)
+
+
+def _expand_num(n: int) -> str:
+    return _inflect.number_to_words(n)
+
+
+def _expand_time_english(match: "re.Match") -> str:
+    hour = int(match.group(1))
+    past_noon = hour >= 12
+    time = []
+    if hour > 12:
+        hour -= 12
+    elif hour == 0:
+        hour = 12
+        past_noon = True
+    time.append(_expand_num(hour))
+
+    minute = int(match.group(6))
+    if minute > 0:
+        if minute < 10:
+            time.append("oh")
+        time.append(_expand_num(minute))
+    am_pm = match.group(7)
+    if am_pm is None:
+        time.append("p m" if past_noon else "a m")
+    else:
+        time.extend(list(am_pm.replace(".", "")))
+    return " ".join(time)
+
+
+def expand_time_english(text: str) -> str:
+    return re.sub(_time_re, _expand_time_english, text)

tiny_tts/text/symbols.py

@@ -0,0 +1,293 @@
+# punctuation = ["!", "?", "…", ",", ".", "'", "-"]
+punctuation = ["!", "?", "…", ",", ".", "'", "-", "¿", "¡"]
+pu_symbols = punctuation + ["SP", "UNK"]
+pad = "_"
+
+# chinese
+zh_symbols = [
+    "E",
+    "En",
+    "a",
+    "ai",
+    "an",
+    "ang",
+    "ao",
+    "b",
+    "c",
+    "ch",
+    "d",
+    "e",
+    "ei",
+    "en",
+    "eng",
+    "er",
+    "f",
+    "g",
+    "h",
+    "i",
+    "i0",
+    "ia",
+    "ian",
+    "iang",
+    "iao",
+    "ie",
+    "in",
+    "ing",
+    "iong",
+    "ir",
+    "iu",
+    "j",
+    "k",
+    "l",
+    "m",
+    "n",
+    "o",
+    "ong",
+    "ou",
+    "p",
+    "q",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "u",
+    "ua",
+    "uai",
+    "uan",
+    "uang",
+    "ui",
+    "un",
+    "uo",
+    "v",
+    "van",
+    "ve",
+    "vn",
+    "w",
+    "x",
+    "y",
+    "z",
+    "zh",
+    "AA",
+    "EE",
+    "OO",
+]
+num_zh_tones = 6
+
+# japanese
+ja_symbols = [
+    "N",
+    "a",
+    "a:",
+    "b",
+    "by",
+    "ch",
+    "d",
+    "dy",
+    "e",
+    "e:",
+    "f",
+    "g",
+    "gy",
+    "h",
+    "hy",
+    "i",
+    "i:",
+    "j",
+    "k",
+    "ky",
+    "m",
+    "my",
+    "n",
+    "ny",
+    "o",
+    "o:",
+    "p",
+    "py",
+    "q",
+    "r",
+    "ry",
+    "s",
+    "sh",
+    "t",
+    "ts",
+    "ty",
+    "u",
+    "u:",
+    "w",
+    "y",
+    "z",
+    "zy",
+]
+num_ja_tones = 1
+
+# English
+en_symbols = [
+    "aa",
+    "ae",
+    "ah",
+    "ao",
+    "aw",
+    "ay",
+    "b",
+    "ch",
+    "d",
+    "dh",
+    "eh",
+    "er",
+    "ey",
+    "f",
+    "g",
+    "hh",
+    "ih",
+    "iy",
+    "jh",
+    "k",
+    "l",
+    "m",
+    "n",
+    "ng",
+    "ow",
+    "oy",
+    "p",
+    "r",
+    "s",
+    "sh",
+    "t",
+    "th",
+    "uh",
+    "uw",
+    "V",
+    "w",
+    "y",
+    "z",
+    "zh",
+]
+num_en_tones = 4
+
+# Korean
+kr_symbols = ['ᄌ', 'ᅥ', 'ᆫ', 'ᅦ', 'ᄋ', 'ᅵ', 'ᄅ', 'ᅴ', 'ᄀ', 'ᅡ', 'ᄎ', 'ᅪ', 'ᄑ', 'ᅩ', 'ᄐ', 'ᄃ', 'ᅢ', 'ᅮ', 'ᆼ', 'ᅳ', 'ᄒ', 'ᄆ', 'ᆯ', 'ᆷ', 'ᄂ', 'ᄇ', 'ᄉ', 'ᆮ', 'ᄁ', 'ᅬ', 'ᅣ', 'ᄄ', 'ᆨ', 'ᄍ', 'ᅧ', 'ᄏ', 'ᆸ', 'ᅭ', '(', 'ᄊ', ')', 'ᅲ', 'ᅨ', 'ᄈ', 'ᅱ', 'ᅯ', 'ᅫ', 'ᅰ', 'ᅤ', '~', '\\', '[', ']', '/', '^', ':', 'ㄸ', '*']
+num_kr_tones = 1
+
+# Spanish
+es_symbols = [
+        "N",
+        "Q",
+        "a",
+        "b",
+        "d",
+        "e",
+        "f",
+        "g",
+        "h",
+        "i",
+        "j",
+        "k",
+        "l",
+        "m",
+        "n",
+        "o",
+        "p",
+        "s",
+        "t",
+        "u",
+        "v",
+        "w",
+        "x",
+        "y",
+        "z",
+        "ɑ",
+        "æ",
+        "ʃ",
+        "ʑ",
+        "ç",
+        "ɯ",
+        "ɪ",
+        "ɔ",
+        "ɛ",
+        "ɹ",
+        "ð",
+        "ə",
+        "ɫ",
+        "ɥ",
+        "ɸ",
+        "ʊ",
+        "ɾ",
+        "ʒ",
+        "θ",
+        "β",
+        "ŋ",
+        "ɦ",
+        "ɡ",
+        "r",
+        "ɲ",
+        "ʝ",
+        "ɣ",
+        "ʎ",
+        "ˈ",
+        "ˌ",
+        "ː"
+    ]
+num_es_tones = 1
+
+# French
+fr_symbols = [
+    "\u0303",
+    "œ",
+    "ø",
+    "ʁ",
+    "ɒ",
+    "ʌ",
+    "ɜ",
+    "ɐ"
+]
+num_fr_tones = 1
+
+# German
+de_symbols = [
+    "ʏ",
+    "̩"
+]
+num_de_tones = 1
+
+# Russian
+ru_symbols = [
+    "ɭ",
+    "ʲ",
+    "ɕ",
+    "\"",
+    "ɵ",
+    "^",
+    "ɬ"
+]
+num_ru_tones = 1
+
+# combine all symbols
+normal_symbols = sorted(set(zh_symbols + ja_symbols + en_symbols + kr_symbols + es_symbols + fr_symbols + de_symbols + ru_symbols))
+symbols = [pad] + normal_symbols + pu_symbols
+sil_phonemes_ids = [symbols.index(i) for i in pu_symbols]
+
+# combine all tones
+num_tones = num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones + num_ru_tones
+
+# language maps
+language_id_map = {"ZH": 0, "JP": 1, "EN": 2, "ZH_MIX_EN": 3, 'KR': 4, 'ES': 5, 'SP': 5, 'FR': 6, 'DE': 7, 'RU': 8, 'VI': 9}
+num_languages = 10
+
+language_tone_start_map = {
+    "ZH": 0,
+    "ZH_MIX_EN": 0,
+    "JP": num_zh_tones,
+    "EN": num_zh_tones + num_ja_tones,
+    'KR': num_zh_tones + num_ja_tones + num_en_tones,
+    "ES": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones,
+    "SP": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones,
+    "FR": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones,
+    "DE": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones,
+    "RU": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones,
+    "VI": num_zh_tones + num_ja_tones + num_en_tones + num_kr_tones + num_es_tones + num_fr_tones + num_de_tones + num_ru_tones,
+}
+
+if __name__ == "__main__":
+    a = set(zh_symbols)
+    b = set(en_symbols)
+    print(sorted(a & b))

tiny_tts/utils/__init__.py

@@ -0,0 +1,5 @@
+from .config import (
+    SAMPLING_RATE, FILTER_LENGTH, HOP_LENGTH, SEGMENT_FRAMES,
+    ADD_BLANK, SPEC_CHANNELS, N_SPEAKERS, SPK2ID,
+    MODEL_PARAMS, NUM_LANGUAGES, NUM_TONES,
+)

tiny_tts/utils/config.py

@@ -0,0 +1,42 @@
+# Audio
+SAMPLING_RATE = 44100
+FILTER_LENGTH = 2048
+HOP_LENGTH = 512
+SEGMENT_FRAMES = 32
+ADD_BLANK = True
+SPEC_CHANNELS = FILTER_LENGTH // 2 + 1  # 1025
+N_MEL_CHANNELS = 128                      # updated in new checkpoint
+
+# Speakers
+N_SPEAKERS = 1
+SPK2ID = {"MALE": 0}
+
+# Model — matches config.json for G_150000.pth (lighter version)
+MODEL_PARAMS = dict(
+    use_spk_conditioned_encoder=True,
+    use_noise_scaled_mas=True,
+    inter_channels=32,
+    hidden_channels=32,
+    filter_channels=128,
+    n_heads=2,
+    n_layers=3,
+    n_layers_trans_flow=3,
+    kernel_size=3,
+    p_dropout=0.1,
+    resblock="1",
+    resblock_kernel_sizes=[3, 7, 11],
+    resblock_dilation_sizes=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+    upsample_rates=[8, 8, 2, 2, 2],
+    upsample_initial_channel=64,
+    upsample_kernel_sizes=[16, 16, 8, 2, 2],
+    n_layers_q=3,
+    use_spectral_norm=False,
+    gin_channels=128,
+    use_sdp=True,
+    mas_noise_scale_initial=0.01,
+    noise_scale_delta=2e-06,
+)
+
+# Language / Tone
+NUM_LANGUAGES = 1
+NUM_TONES = 6

tinytts_text_cleaning.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import re
+
+
+_QUOTE_TRANSLATION = str.maketrans(
+    {
+        "\u2018": "'",
+        "\u2019": "'",
+        "\u201c": "",
+        "\u201d": "",
+        "\u2014": ",",
+        "\u2013": ",",
+        ";": ",",
+        ":": ",",
+        "\n": ".",
+    }
+)
+
+
+def clean_tinytts_text(text: str) -> str:
+    """Normalize text into punctuation TinyTTS actually has symbols for."""
+    text = str(text).translate(_QUOTE_TRANSLATION)
+    text = text.replace("...", "…")
+    text = re.sub(r"\s+", " ", text).strip()
+    text = re.sub(r"\s+([,.!?…])", r"\1", text)
+    text = re.sub(r"([,.!?…]){2,}", r"\1", text)
+    return text

train_hifigan_oracle_v1.py

@@ -0,0 +1,830 @@
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from torch.nn.utils import remove_weight_norm, spectral_norm, weight_norm
+from torch.utils.data import DataLoader, Dataset
+
+
+@dataclass(frozen=True)
+class HifiGanConfig:
+    variant: str
+    sample_rate: int = 24000
+    n_fft: int = 1024
+    hop_size: int = 256
+    win_size: int = 1024
+    num_mels: int = 80
+    fmin: float = 0.0
+    fmax: float = 12000.0
+    resblock: str = "1"
+    upsample_rates: tuple[int, ...] = (8, 8, 2, 2)
+    upsample_kernel_sizes: tuple[int, ...] = (16, 16, 4, 4)
+    upsample_initial_channel: int = 128
+    resblock_kernel_sizes: tuple[int, ...] = (3, 7, 11)
+    resblock_dilation_sizes: tuple[tuple[int, ...], ...] = ((1, 3, 5), (1, 3, 5), (1, 3, 5))
+    activation: str = "lrelu"
+    conditioning_channels: int = 0
+
+
+def make_config(variant: str) -> HifiGanConfig:
+    if variant == "v2":
+        return HifiGanConfig(variant="v2")
+    if variant == "v2plus":
+        return HifiGanConfig(variant="v2plus", upsample_initial_channel=160)
+    if variant == "v2wide":
+        return HifiGanConfig(variant="v2wide", upsample_initial_channel=176)
+    if variant == "snake_v2mid":
+        return HifiGanConfig(variant="snake_v2mid", upsample_initial_channel=144, activation="snake")
+    if variant == "snake_v2balanced":
+        return HifiGanConfig(variant="snake_v2balanced", upsample_initial_channel=160, activation="snake")
+    if variant == "source_snake_v2balanced":
+        return HifiGanConfig(
+            variant="source_snake_v2balanced",
+            upsample_initial_channel=160,
+            activation="snake",
+            conditioning_channels=5,
+        )
+    if variant == "v3":
+        return HifiGanConfig(
+            variant="v3",
+            resblock="2",
+            upsample_rates=(8, 8, 4),
+            upsample_kernel_sizes=(16, 16, 8),
+            upsample_initial_channel=256,
+            resblock_kernel_sizes=(3, 5, 7),
+            resblock_dilation_sizes=((1, 2), (2, 6), (3, 12)),
+        )
+    raise ValueError(f"Unknown variant: {variant}")
+
+
+def get_padding(kernel_size: int, dilation: int = 1) -> int:
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class SnakeActivation(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.log_alpha = nn.Parameter(torch.zeros(1, channels, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        alpha = self.log_alpha.exp().clamp(1e-4, 100.0)
+        return x + torch.sin(alpha * x).pow(2) / alpha
+
+
+def make_activation(channels: int, activation: str) -> nn.Module:
+    if activation == "snake":
+        return SnakeActivation(channels)
+    return nn.LeakyReLU(0.1)
+
+
+class ResBlock1(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
+        super().__init__()
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilations
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1))
+                )
+                for _ in dilations
+            ]
+        )
+        self.acts1 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+        self.acts2 = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, self.acts1, self.acts2):
+            y = a1(x)
+            y = c1(y)
+            y = a2(y)
+            y = c2(y)
+            x = x + y
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in list(self.convs1) + list(self.convs2):
+            remove_weight_norm(layer)
+
+
+class ResBlock2(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dilations: tuple[int, ...], activation: str = "lrelu"):
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilations
+            ]
+        )
+        self.acts = nn.ModuleList([make_activation(channels, activation) for _ in dilations])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv, act in zip(self.convs, self.acts):
+            y = act(x)
+            y = conv(y)
+            x = x + y
+        return x
+
+    def remove_weight_norm(self) -> None:
+        for layer in self.convs:
+            remove_weight_norm(layer)
+
+
+class HifiGanGenerator(nn.Module):
+    def __init__(self, cfg: HifiGanConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.num_kernels = len(cfg.resblock_kernel_sizes)
+        self.num_upsamples = len(cfg.upsample_rates)
+        self.conv_pre = weight_norm(
+            nn.Conv1d(cfg.num_mels + cfg.conditioning_channels, cfg.upsample_initial_channel, 7, 1, padding=3)
+        )
+        self.ups = nn.ModuleList()
+        self.up_acts = nn.ModuleList()
+        self.resblocks = nn.ModuleList()
+        resblock_cls = ResBlock1 if cfg.resblock == "1" else ResBlock2
+        for i, (rate, kernel) in enumerate(zip(cfg.upsample_rates, cfg.upsample_kernel_sizes)):
+            in_ch = cfg.upsample_initial_channel // (2**i)
+            out_ch = cfg.upsample_initial_channel // (2 ** (i + 1))
+            self.up_acts.append(make_activation(in_ch, cfg.activation))
+            self.ups.append(
+                weight_norm(
+                    nn.ConvTranspose1d(
+                        in_ch,
+                        out_ch,
+                        kernel,
+                        rate,
+                        padding=(kernel - rate) // 2,
+                    )
+                )
+            )
+            for k, d in zip(cfg.resblock_kernel_sizes, cfg.resblock_dilation_sizes):
+                self.resblocks.append(resblock_cls(out_ch, k, d, cfg.activation))
+        final_ch = cfg.upsample_initial_channel // (2 ** len(cfg.upsample_rates))
+        self.post_act = make_activation(final_ch, cfg.activation)
+        self.conv_post = weight_norm(nn.Conv1d(final_ch, 1, 7, 1, padding=3))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_pre(x)
+        for i, up in enumerate(self.ups):
+            x = self.up_acts[i](x)
+            x = up(x)
+            xs = 0.0
+            for j in range(self.num_kernels):
+                xs = xs + self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+        x = self.post_act(x)
+        x = self.conv_post(x)
+        return torch.tanh(x)
+
+    def remove_weight_norm(self) -> None:
+        remove_weight_norm(self.conv_pre)
+        for up in self.ups:
+            remove_weight_norm(up)
+        for block in self.resblocks:
+            block.remove_weight_norm()
+        remove_weight_norm(self.conv_post)
+
+
+def extract_source_features(
+    wav: torch.Tensor,
+    cfg: HifiGanConfig,
+    frames: int,
+    dropout: float = 0.0,
+    noise: float = 0.0,
+) -> torch.Tensor:
+    """Return low-rate F0/voicing features for source-conditioned generators."""
+    pitch = torchaudio.functional.detect_pitch_frequency(
+        wav.detach().cpu(),
+        sample_rate=cfg.sample_rate,
+        frame_time=cfg.hop_size / cfg.sample_rate,
+        win_length=30,
+    ).to(wav.device)
+    if pitch.ndim == 1:
+        pitch = pitch.unsqueeze(0)
+    if pitch.shape[-1] < frames:
+        pitch = F.pad(pitch, (0, frames - pitch.shape[-1]), value=0.0)
+    pitch = pitch[..., :frames]
+    voiced = ((pitch >= 55.0) & (pitch <= 420.0)).float()
+    pitch = pitch.clamp(55.0, 420.0)
+    log_f0 = ((torch.log(pitch) - math.log(140.0)) / 0.45).clamp(-3.0, 3.0) * voiced
+    if noise > 0.0:
+        log_f0 = (log_f0 + torch.randn_like(log_f0) * noise * voiced).clamp(-3.0, 3.0)
+    jump = F.pad((log_f0[..., 1:] - log_f0[..., :-1]).abs(), (1, 0))
+    confidence = torch.exp(-1.5 * jump) * voiced
+    reconstructed_f0 = torch.exp(log_f0 * 0.45 + math.log(140.0))
+    phase = torch.cumsum(2.0 * math.pi * reconstructed_f0 * (cfg.hop_size / cfg.sample_rate), dim=-1)
+    source = torch.stack(
+        [log_f0, voiced, confidence, torch.sin(phase) * confidence, torch.cos(phase) * confidence],
+        dim=1,
+    )
+    if dropout > 0.0:
+        # Drop the complete source sketch for some examples so inference remains
+        # stable when predicted F0 confidence is poor.
+        keep = (torch.rand(source.shape[0], 1, 1, device=source.device) >= dropout).to(source.dtype)
+        source = source * keep
+    return source
+
+
+class DiscriminatorP(nn.Module):
+    def __init__(self, period: int):
+        super().__init__()
+        self.period = period
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(nn.Conv2d(1, 32, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(32, 128, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(128, 512, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0))),
+                weight_norm(nn.Conv2d(1024, 1024, (5, 1), 1, padding=(2, 0))),
+            ]
+        )
+        self.conv_post = weight_norm(nn.Conv2d(1024, 1, (3, 1), 1, padding=(1, 0)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        b, c, t = x.shape
+        if t % self.period != 0:
+            x = F.pad(x, (0, self.period - (t % self.period)), mode="reflect")
+            t = x.shape[-1]
+        x = x.view(b, c, t // self.period, self.period)
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiPeriodDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList([DiscriminatorP(p) for p in (2, 3, 5, 7, 11)])
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for d in self.discriminators:
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class DiscriminatorS(nn.Module):
+    def __init__(self, use_spectral_norm: bool = False):
+        super().__init__()
+        norm = spectral_norm if use_spectral_norm else weight_norm
+        self.convs = nn.ModuleList(
+            [
+                norm(nn.Conv1d(1, 128, 15, 1, padding=7)),
+                norm(nn.Conv1d(128, 128, 41, 2, groups=4, padding=20)),
+                norm(nn.Conv1d(128, 256, 41, 2, groups=16, padding=20)),
+                norm(nn.Conv1d(256, 512, 41, 4, groups=16, padding=20)),
+                norm(nn.Conv1d(512, 1024, 41, 4, groups=16, padding=20)),
+                norm(nn.Conv1d(1024, 1024, 41, 1, groups=16, padding=20)),
+                norm(nn.Conv1d(1024, 1024, 5, 1, padding=2)),
+            ]
+        )
+        self.conv_post = norm(nn.Conv1d(1024, 1, 3, 1, padding=1))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiScaleDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.discriminators = nn.ModuleList([DiscriminatorS(True), DiscriminatorS(), DiscriminatorS()])
+        self.meanpools = nn.ModuleList([nn.AvgPool1d(4, 2, padding=2), nn.AvgPool1d(4, 2, padding=2)])
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for i, d in enumerate(self.discriminators):
+            if i:
+                y = self.meanpools[i - 1](y)
+                y_hat = self.meanpools[i - 1](y_hat)
+            y_d_r, fmap_r = d(y)
+            y_d_g, fmap_g = d(y_hat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class SpectrogramDiscriminator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        channels = (32, 64, 128, 128)
+        layers: list[nn.Module] = []
+        in_ch = 1
+        for out_ch, stride in zip(channels, ((1, 2), (2, 2), (2, 2), (2, 1))):
+            layers.append(weight_norm(nn.Conv2d(in_ch, out_ch, (5, 5), stride=stride, padding=(2, 2))))
+            in_ch = out_ch
+        self.convs = nn.ModuleList(layers)
+        self.conv_post = weight_norm(nn.Conv2d(in_ch, 1, (3, 3), padding=(1, 1)))
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        fmap = []
+        for conv in self.convs:
+            x = F.leaky_relu(conv(x), 0.1)
+            fmap.append(x)
+        x = self.conv_post(x)
+        fmap.append(x)
+        return torch.flatten(x, 1, -1), fmap
+
+
+class MultiResolutionSpectrogramDiscriminator(nn.Module):
+    def __init__(self, fft_sizes: tuple[int, ...] = (256, 512, 1024), hop_sizes: tuple[int, ...] = (64, 128, 256), win_lengths: tuple[int, ...] = (256, 512, 1024)):
+        super().__init__()
+        self.fft_sizes = fft_sizes
+        self.hop_sizes = hop_sizes
+        self.win_lengths = win_lengths
+        self.discriminators = nn.ModuleList([SpectrogramDiscriminator() for _ in fft_sizes])
+
+    def _features(self, wav: torch.Tensor, fft: int, hop: int, win_len: int) -> torch.Tensor:
+        wav = wav.squeeze(1)
+        window = torch.hann_window(win_len, device=wav.device)
+        spec = torch.stft(wav, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        mag = torch.log(spec.abs().clamp_min(1e-5))
+        mean = mag.mean(dim=(1, 2), keepdim=True)
+        std = mag.std(dim=(1, 2), keepdim=True).clamp_min(1e-4)
+        return ((mag - mean) / std).unsqueeze(1)
+
+    def forward(self, y: torch.Tensor, y_hat: torch.Tensor):
+        y_d_rs, y_d_gs, fmap_rs, fmap_gs = [], [], [], []
+        for disc, fft, hop, win_len in zip(self.discriminators, self.fft_sizes, self.hop_sizes, self.win_lengths):
+            y_feat = self._features(y, fft, hop, win_len)
+            y_hat_feat = self._features(y_hat, fft, hop, win_len)
+            y_d_r, fmap_r = disc(y_feat)
+            y_d_g, fmap_g = disc(y_hat_feat)
+            y_d_rs.append(y_d_r)
+            y_d_gs.append(y_d_g)
+            fmap_rs.append(fmap_r)
+            fmap_gs.append(fmap_g)
+        return y_d_rs, y_d_gs, fmap_rs, fmap_gs
+
+
+class MelFrontend(nn.Module):
+    def __init__(self, cfg: HifiGanConfig):
+        super().__init__()
+        self.mel = torchaudio.transforms.MelSpectrogram(
+            sample_rate=cfg.sample_rate,
+            n_fft=cfg.n_fft,
+            win_length=cfg.win_size,
+            hop_length=cfg.hop_size,
+            f_min=cfg.fmin,
+            f_max=cfg.fmax,
+            n_mels=cfg.num_mels,
+            power=1.0,
+            center=True,
+            norm="slaney",
+            mel_scale="slaney",
+        )
+
+    def forward(self, wav: torch.Tensor) -> torch.Tensor:
+        return torch.log(torch.clamp(self.mel(wav), min=1e-5))
+
+
+def load_rows(path: Path, max_rows: int, min_seconds: float, max_seconds: float) -> list[dict]:
+    rows = []
+    with path.open("r", encoding="utf-8-sig") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            row = json.loads(line)
+            audio = Path(str(row.get("target_audio") or ""))
+            text = str(row.get("target_text") or "").strip()
+            dur = float(row.get("target_duration_s") or 0.0)
+            if audio.is_file() and text and min_seconds <= (dur or 4.0) <= max_seconds:
+                rows.append({"audio": str(audio), "text": text, "duration": dur})
+                if max_rows > 0 and len(rows) >= max_rows:
+                    break
+    if not rows:
+        raise RuntimeError(f"No rows loaded from {path}")
+    return rows
+
+
+def load_audio(path: str, sample_rate: int) -> torch.Tensor:
+    wav, sr = torchaudio.load(path)
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+    if sr != sample_rate:
+        wav = torchaudio.functional.resample(wav, sr, sample_rate)
+    wav = wav.squeeze(0)
+    return wav.clamp(-1, 1)
+
+
+class AudioDataset(Dataset):
+    def __init__(self, rows: list[dict], cfg: HifiGanConfig, segment_size: int, seed: int):
+        self.rows = rows
+        self.cfg = cfg
+        self.segment_size = segment_size
+        self.rng = random.Random(seed)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, idx: int) -> torch.Tensor:
+        wav = load_audio(self.rows[idx]["audio"], self.cfg.sample_rate)
+        if wav.numel() >= self.segment_size:
+            start = self.rng.randint(0, wav.numel() - self.segment_size)
+            return wav[start : start + self.segment_size]
+        return F.pad(wav, (0, self.segment_size - wav.numel()))
+
+
+def feature_loss(fmap_r, fmap_g) -> torch.Tensor:
+    loss = 0.0
+    for dr, dg in zip(fmap_r, fmap_g):
+        for rl, gl in zip(dr, dg):
+            loss = loss + F.l1_loss(rl.detach(), gl)
+    return loss * 2
+
+
+def discriminator_loss(disc_real_outputs, disc_generated_outputs) -> torch.Tensor:
+    loss = 0.0
+    for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
+        loss = loss + torch.mean((1 - dr) ** 2) + torch.mean(dg**2)
+    return loss
+
+
+def generator_loss(disc_outputs) -> torch.Tensor:
+    loss = 0.0
+    for dg in disc_outputs:
+        loss = loss + torch.mean((1 - dg) ** 2)
+    return loss
+
+
+def stft_mag_loss(y_hat: torch.Tensor, y: torch.Tensor, fft_sizes: tuple[int, ...], hop_sizes: tuple[int, ...], win_lengths: tuple[int, ...]) -> torch.Tensor:
+    # Multi-resolution spectral loss catches buzz/shimmer that can hide behind
+    # mel loss, especially for a small generator near convergence.
+    y_hat = y_hat.squeeze(1)
+    y = y.squeeze(1)
+    total = torch.zeros((), device=y.device)
+    for fft, hop, win_len in zip(fft_sizes, hop_sizes, win_lengths):
+        window = torch.hann_window(win_len, device=y.device)
+        pred = torch.stft(y_hat, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        target = torch.stft(y, n_fft=fft, hop_length=hop, win_length=win_len, window=window, return_complex=True)
+        pred_mag = pred.abs().clamp_min(1e-7)
+        target_mag = target.abs().clamp_min(1e-7)
+        sc = torch.linalg.vector_norm(target_mag - pred_mag) / torch.linalg.vector_norm(target_mag).clamp_min(1e-7)
+        log_mag = F.l1_loss(torch.log(pred_mag), torch.log(target_mag))
+        total = total + sc + log_mag
+    return total / max(1, len(fft_sizes))
+
+
+def count_parameters(module: nn.Module) -> int:
+    return sum(p.numel() for p in module.parameters())
+
+
+def jsonable_args(args: argparse.Namespace) -> dict:
+    return {k: str(v) if isinstance(v, Path) else v for k, v in vars(args).items()}
+
+
+def save_checkpoint(
+    path: Path,
+    generator: nn.Module,
+    mpd: nn.Module,
+    msd: nn.Module,
+    optim_g,
+    optim_d,
+    cfg: HifiGanConfig,
+    step: int,
+    args,
+    mrsd: nn.Module | None = None,
+) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    payload = {
+            "generator": generator.state_dict(),
+            "mpd": mpd.state_dict(),
+            "msd": msd.state_dict(),
+            "optim_g": optim_g.state_dict(),
+            "optim_d": optim_d.state_dict(),
+            "config": asdict(cfg),
+            "step": step,
+            "args": jsonable_args(args),
+            "generator_params": count_parameters(generator),
+        }
+    if mrsd is not None:
+        payload["mrsd"] = mrsd.state_dict()
+    torch.save(payload, tmp)
+    tmp.replace(path)
+
+
+def checkpoint_step(path: Path) -> int:
+    stem = path.stem
+    tail = stem.rsplit("-", 1)[-1]
+    return int(tail) if tail.isdigit() else -1
+
+
+def prune_checkpoints(out_dir: Path, variant: str, keep: int) -> None:
+    if keep <= 0:
+        return
+    numbered = [p for p in out_dir.glob(f"hifigan-{variant}-*.pt") if checkpoint_step(p) >= 0]
+    numbered.sort(key=checkpoint_step, reverse=True)
+    for old in numbered[keep:]:
+        old.unlink(missing_ok=True)
+
+
+def latest_checkpoint(out_dir: Path) -> Path | None:
+    numbered = [p for p in out_dir.glob("hifigan-*-*.pt") if checkpoint_step(p) >= 0]
+    if numbered:
+        return max(numbered, key=checkpoint_step)
+    ckpts = sorted(out_dir.glob("hifigan-*-latest.pt"), key=lambda p: p.stat().st_mtime, reverse=True)
+    return ckpts[0] if ckpts else None
+
+
+def partial_load_state(module: nn.Module, state: dict[str, torch.Tensor]) -> tuple[int, int]:
+    current = module.state_dict()
+    patched: dict[str, torch.Tensor] = {}
+    copied = 0
+    skipped = 0
+    for name, target in current.items():
+        source = state.get(name)
+        if source is None:
+            skipped += 1
+            continue
+        if source.shape == target.shape:
+            patched[name] = source
+            copied += 1
+            continue
+        if source.ndim != target.ndim:
+            skipped += 1
+            continue
+        value = target.clone()
+        slices = tuple(slice(0, min(a, b)) for a, b in zip(target.shape, source.shape))
+        value[slices] = source[slices].to(value.device, value.dtype)
+        patched[name] = value
+        copied += 1
+    module.load_state_dict(patched, strict=False)
+    return copied, skipped
+
+
+def train(args: argparse.Namespace) -> None:
+    torch.backends.cudnn.benchmark = True
+    cfg = make_config(args.variant)
+    device = torch.device(args.device)
+    rows = load_rows(args.train_jsonl, args.max_rows, args.min_seconds, args.max_seconds)
+    rng = random.Random(args.seed)
+    rng.shuffle(rows)
+    dataset = AudioDataset(rows, cfg, args.segment_size, args.seed)
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.num_workers)
+    mel_frontend = MelFrontend(cfg).to(device)
+    generator = HifiGanGenerator(cfg).to(device)
+    mpd = MultiPeriodDiscriminator().to(device)
+    msd = MultiScaleDiscriminator().to(device)
+    mrsd = MultiResolutionSpectrogramDiscriminator().to(device) if args.spec_disc_weight > 0.0 else None
+    optim_g = torch.optim.AdamW(generator.parameters(), lr=args.lr, betas=(0.8, 0.99))
+    disc_params = list(mpd.parameters()) + list(msd.parameters())
+    if mrsd is not None:
+        disc_params += list(mrsd.parameters())
+    optim_d = torch.optim.AdamW(disc_params, lr=args.lr, betas=(0.8, 0.99))
+    start_step = 0
+    if args.init_checkpoint and not args.resume:
+        ckpt = torch.load(args.init_checkpoint, map_location=device, weights_only=False)
+        if args.partial_init:
+            copied, skipped = partial_load_state(generator, ckpt["generator"])
+            print(f"Partially initialized generator from {args.init_checkpoint}: copied={copied} skipped={skipped}")
+        else:
+            generator.load_state_dict(ckpt["generator"])
+        if "mpd" in ckpt and "msd" in ckpt:
+            mpd.load_state_dict(ckpt["mpd"])
+            msd.load_state_dict(ckpt["msd"])
+        if mrsd is not None and "mrsd" in ckpt:
+            mrsd.load_state_dict(ckpt["mrsd"])
+        can_load_disc_optim = mrsd is None or "mrsd" in ckpt
+        if not args.partial_init and not args.reset_optim and "optim_g" in ckpt:
+            optim_g.load_state_dict(ckpt["optim_g"])
+        if not args.partial_init and not args.reset_optim and can_load_disc_optim and "optim_d" in ckpt:
+            optim_d.load_state_dict(ckpt["optim_d"])
+        for group in optim_g.param_groups:
+            group["lr"] = args.lr
+        for group in optim_d.param_groups:
+            group["lr"] = args.lr
+        start_step = int(ckpt.get("step") or 0)
+        print(f"Initialized {args.init_checkpoint} at step {start_step}; lr={args.lr:g}")
+    if args.resume:
+        ckpt_path = latest_checkpoint(args.out_dir)
+        if ckpt_path:
+            ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+            generator.load_state_dict(ckpt["generator"])
+            mpd.load_state_dict(ckpt["mpd"])
+            msd.load_state_dict(ckpt["msd"])
+            if mrsd is not None and "mrsd" in ckpt:
+                mrsd.load_state_dict(ckpt["mrsd"])
+            optim_g.load_state_dict(ckpt["optim_g"])
+            optim_d.load_state_dict(ckpt["optim_d"])
+            for group in optim_g.param_groups:
+                group["lr"] = args.lr
+            for group in optim_d.param_groups:
+                group["lr"] = args.lr
+            start_step = int(ckpt.get("step") or 0)
+            print(f"Resumed {ckpt_path} at step {start_step}; lr={args.lr:g}")
+
+    args.out_dir.mkdir(parents=True, exist_ok=True)
+    prune_checkpoints(args.out_dir, args.variant, args.keep_checkpoints)
+    (args.out_dir / "config.json").write_text(
+        json.dumps(
+            {
+                "config": asdict(cfg),
+                "args": jsonable_args(args),
+                "rows": len(rows),
+                "generator_params": count_parameters(generator),
+                "mpd_params": count_parameters(mpd),
+                "msd_params": count_parameters(msd),
+                "mrsd_params": count_parameters(mrsd) if mrsd is not None else 0,
+            },
+            indent=2,
+        ),
+        encoding="utf-8",
+    )
+    print(f"Variant: {args.variant}")
+    print(f"Rows: {len(rows)}")
+    print(f"Generator params: {count_parameters(generator):,} ({count_parameters(generator)/1_000_000:.3f}M)")
+    print(f"MPD params: {count_parameters(mpd):,} MSD params: {count_parameters(msd):,} (training only)")
+    if mrsd is not None:
+        print(f"MRSD params: {count_parameters(mrsd):,} (training only)")
+    if args.steps == 0:
+        return
+
+    step = start_step
+    started = time.time()
+    try:
+        while step < args.steps:
+            for wav in loader:
+                step += 1
+                y = wav.unsqueeze(1).to(device)
+                with torch.no_grad():
+                    mel = mel_frontend(wav.to(device))
+                    if cfg.conditioning_channels:
+                        source = extract_source_features(
+                            wav.to(device),
+                            cfg,
+                            mel.shape[-1],
+                            dropout=args.source_dropout,
+                            noise=args.source_noise,
+                        )
+                        generator_input = torch.cat([mel, source], dim=1)
+                    else:
+                        generator_input = mel
+                y_hat = generator(generator_input)
+                common = min(y.shape[-1], y_hat.shape[-1])
+                y = y[..., :common]
+                y_hat = y_hat[..., :common]
+                y_mel = mel_frontend(y.squeeze(1))
+                y_hat_mel = mel_frontend(y_hat.squeeze(1))
+
+                optim_d.zero_grad(set_to_none=True)
+                y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_hat.detach())
+                y_ds_hat_r, y_ds_hat_g, _, _ = msd(y, y_hat.detach())
+                loss_disc = discriminator_loss(y_df_hat_r, y_df_hat_g) + discriminator_loss(y_ds_hat_r, y_ds_hat_g)
+                loss_spec_disc = torch.zeros((), device=device)
+                if mrsd is not None:
+                    y_dm_hat_r, y_dm_hat_g, _, _ = mrsd(y, y_hat.detach())
+                    loss_spec_disc = discriminator_loss(y_dm_hat_r, y_dm_hat_g)
+                    loss_disc = loss_disc + args.spec_disc_weight * loss_spec_disc
+                loss_disc.backward()
+                torch.nn.utils.clip_grad_norm_(disc_params, args.grad_clip)
+                optim_d.step()
+
+                optim_g.zero_grad(set_to_none=True)
+                mel_loss = F.l1_loss(y_mel, y_hat_mel) * args.mel_weight
+                y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_hat)
+                y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = msd(y, y_hat)
+                loss_fm = feature_loss(fmap_f_r, fmap_f_g) + feature_loss(fmap_s_r, fmap_s_g)
+                loss_gen = generator_loss(y_df_hat_g) + generator_loss(y_ds_hat_g)
+                loss_spec_gen = torch.zeros((), device=device)
+                loss_spec_fm = torch.zeros((), device=device)
+                if mrsd is not None:
+                    y_dm_hat_r, y_dm_hat_g, fmap_m_r, fmap_m_g = mrsd(y, y_hat)
+                    loss_spec_gen = generator_loss(y_dm_hat_g)
+                    loss_spec_fm = feature_loss(fmap_m_r, fmap_m_g)
+                wav_l1 = F.l1_loss(y_hat, y) * args.wav_weight
+                stft_loss = torch.zeros((), device=device)
+                if args.stft_weight > 0.0:
+                    stft_loss = stft_mag_loss(y_hat, y, (512, 1024, 2048), (128, 256, 512), (512, 1024, 2048)) * args.stft_weight
+                loss_g = (
+                    mel_loss
+                    + args.fm_weight * loss_fm
+                    + args.adv_weight * loss_gen
+                    + wav_l1
+                    + stft_loss
+                    + args.spec_disc_weight * loss_spec_gen
+                    + args.spec_fm_weight * loss_spec_fm
+                )
+                loss_g.backward()
+                grad_g = torch.nn.utils.clip_grad_norm_(generator.parameters(), args.grad_clip)
+                optim_g.step()
+
+                if step == 1 or step % args.log_interval == 0:
+                    elapsed = max(time.time() - started, 1e-6)
+                    speed = (step - start_step) / elapsed
+                    eta = (args.steps - step) / max(speed, 1e-6)
+                    print(
+                        f"step={step}/{args.steps} g={loss_g.item():.4f} d={loss_disc.item():.4f} "
+                        f"mel={mel_loss.item():.4f} fm={loss_fm.item():.4f} adv={loss_gen.item():.4f} "
+                        f"wav={wav_l1.item():.4f} stft={stft_loss.item():.4f} "
+                        f"sd={loss_spec_disc.item():.4f} sfm={loss_spec_fm.item():.4f} sadv={loss_spec_gen.item():.4f} "
+                        f"grad={float(grad_g):.3f} speed={speed:.3f} step/s eta={eta/60:.1f}m",
+                        flush=True,
+                    )
+                if step % args.save_interval == 0 or step >= args.steps:
+                    prune_checkpoints(args.out_dir, args.variant, max(args.keep_checkpoints - 1, 0))
+                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+                    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+                if step >= args.steps:
+                    break
+    except KeyboardInterrupt:
+        if step > start_step:
+            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-interrupt-{step}.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+            save_checkpoint(args.out_dir / f"hifigan-{args.variant}-latest.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+            print(f"Interrupted. Saved checkpoint at step {step}.", flush=True)
+        raise
+    save_checkpoint(args.out_dir / f"hifigan-{args.variant}-final.pt", generator, mpd, msd, optim_g, optim_d, cfg, step, args, mrsd)
+    print(f"Done. {args.out_dir}")
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="Train exact-ish HiFi-GAN V2/V3 oracle vocoders on corrected Mark audio.")
+    ap.add_argument("--train-jsonl", type=Path, required=True)
+    ap.add_argument("--out-dir", type=Path, required=True)
+    ap.add_argument(
+        "--variant",
+        choices=["v2", "v2plus", "v2wide", "snake_v2mid", "snake_v2balanced", "source_snake_v2balanced", "v3"],
+        required=True,
+    )
+    ap.add_argument("--steps", type=int, default=5000)
+    ap.add_argument("--max-rows", type=int, default=0)
+    ap.add_argument("--min-seconds", type=float, default=1.0)
+    ap.add_argument("--max-seconds", type=float, default=12.0)
+    ap.add_argument("--segment-size", type=int, default=8192)
+    ap.add_argument("--batch-size", type=int, default=8)
+    ap.add_argument("--num-workers", type=int, default=0)
+    ap.add_argument("--lr", type=float, default=2.0e-4)
+    ap.add_argument("--mel-weight", type=float, default=45.0)
+    ap.add_argument("--wav-weight", type=float, default=1.0)
+    ap.add_argument("--fm-weight", type=float, default=1.0)
+    ap.add_argument("--adv-weight", type=float, default=1.0)
+    ap.add_argument("--stft-weight", type=float, default=0.0)
+    ap.add_argument("--spec-disc-weight", type=float, default=0.0, help="Training-only multi-resolution spectrogram adversarial weight.")
+    ap.add_argument("--spec-fm-weight", type=float, default=0.0, help="Training-only spectrogram discriminator feature-matching weight.")
+    ap.add_argument("--source-dropout", type=float, default=0.0, help="Probability of dropping source conditioning per training example.")
+    ap.add_argument("--source-noise", type=float, default=0.0, help="Stddev of normalized log-F0 corruption for source conditioning.")
+    ap.add_argument("--grad-clip", type=float, default=1000.0)
+    ap.add_argument("--log-interval", type=int, default=50)
+    ap.add_argument("--save-interval", type=int, default=1000)
+    ap.add_argument("--keep-checkpoints", type=int, default=12)
+    ap.add_argument("--seed", type=int, default=1234)
+    ap.add_argument("--device", default="cuda")
+    ap.add_argument("--resume", action="store_true")
+    ap.add_argument("--init-checkpoint", type=Path)
+    ap.add_argument("--partial-init", action="store_true", help="Slice-copy compatible generator weights from init-checkpoint into a resized generator.")
+    ap.add_argument("--reset-optim", action="store_true", help="When initializing from a checkpoint, load model/discriminators but start fresh optimizers.")
+    args = ap.parse_args()
+    train(args)
+
+
+if __name__ == "__main__":
+    main()

weights/inflect_nano_v1_acoustic.pt

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