Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+logo.png filter=lfs diff=lfs merge=lfs -text

FP32/vocoder.safetensors

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

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 2026 Banaxi-Tech
+
+   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

@@ -1,3 +1,138 @@
 ---
 license: apache-2.0
+language:
+  - en
+tags:
+  - text-to-speech
+  - tts
+  - speech
+  - tacotron
+  - hifigan
+  - ljspeech
+library_name: pytorch
+pipeline_tag: text-to-speech
+datasets:
+  - lj_speech
 ---
+
+# BananaMind TTS V2
+
+![BananaMind TTS V2 logo](logo.png)
+
+BananaMind TTS V2 is a small from-scratch English single-speaker text-to-speech model trained on LJSpeech. It combines the Tacotron-lite acoustic model with a self-trained HiFi-GAN vocoder, so it no longer depends on the Griffin-Lim fallback used in V1.
+
+This release is named V2 for packaging and release purposes. In the local training code, the HiFi-GAN work was sometimes called V3.
+
+Training code: https://github.com/Banaxi-Tech/bananamind-tts-v1-training-code
+
+## What This Model Is
+
+- English-only TTS
+- Single speaker
+- Character-input Tacotron-lite acoustic model
+- Self-trained HiFi-GAN neural vocoder
+- Trained from scratch on LJSpeech
+- 22.05 kHz audio output
+- Weights provided as `safetensors` for inference
+
+## What This Model Is Not
+
+- Not voice cloning
+- No speaker embeddings
+- No reference audio conditioning
+- No multi-speaker support
+- No pretrained TTS checkpoint was used
+
+## Files
+
+- `model.safetensors`: Tacotron-lite acoustic model weights
+- `vocoder.safetensors`: default BF16 HiFi-GAN generator-only vocoder
+- `FP32/vocoder.safetensors`: FP32 HiFi-GAN generator-only vocoder
+- `full_vocoder/vocoder.pt`: full epoch twenty eight vocoder training checkpoint with generator, discriminators, optimizers, config, epoch, and step
+- `config.json`: Hugging Face custom model config for `AutoModel`
+- `configuration_bananamind_tts.py`: custom `AutoConfig` implementation
+- `modeling_bananamind_tts.py`: custom `AutoModel` implementation with HiFi-GAN loading
+- `model_config.json`: sidecar metadata with acoustic config, tokenizer, vocoder metadata, epoch, and step
+- `generate.py`: local generation example
+
+## Usage
+
+Install runtime dependencies:
+
+```bash
+pip install torch numpy safetensors transformers huggingface_hub
+```
+
+Use with Transformers remote code:
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Banaxi-Tech/BananaMind-TTS-V2",
+    trust_remote_code=True,
+)
+model.eval()
+
+with torch.inference_mode():
+    out = model.tts(
+        "Hello from BananaMind TTS version two.",
+        normalize_wav=True,
+    )
+
+model.save_wav("sample.wav", out.waveform, out.sample_rate)
+```
+
+From this folder, you can also run:
+
+```bash
+python generate.py
+```
+
+## Vocoder Options
+
+The default vocoder is `vocoder.safetensors`, a BF16 generator-only HiFi-GAN export. It is small and intended for normal inference.
+
+To switch to the FP32 generator-only vocoder:
+
+```python
+model.reload_vocoder("FP32/vocoder.safetensors", dtype="float32")
+```
+
+The full training vocoder checkpoint is included at `full_vocoder/vocoder.pt`. It is much larger because it includes the discriminators and optimizer states needed for training resume. It is not needed for normal text-to-speech generation.
+
+## Text Normalization Note
+
+The tokenizer is character based and strips unsupported characters, including digits. Write numbers as words:
+
+- Use `one plus one is two`
+- Not `1 plus 1 is 2`
+
+## Training Summary
+
+- Acoustic model type: Tacotron-lite
+- Vocoder type: HiFi-GAN generator
+- Dataset: LJSpeech
+- Sample rate: 22050 Hz
+- Acoustic epoch: 80
+- Acoustic step: 64529
+- Vocoder epoch: 28
+- Vocoder step: 42000
+- Approximate acoustic parameters: 9.49M
+- Approximate vocoder generator parameters: 3.56M
+
+## Current Limitations
+
+- Character input is less robust than phoneme input for English pronunciation.
+- Numbers and unusual symbols must be normalized before synthesis.
+- Quality is limited by LJSpeech-only training and a small single-speaker architecture.
+- This is a fixed-voice research/demo model, not a production voice system.
+
+## Intended Use
+
+This model is intended for English single-speaker text-to-speech generation with the included custom Transformers model code.
+
+## Out-of-Scope Use
+
+Do not present this model as a voice cloning model or use it to impersonate any person. It has no voice-cloning capability.

config.json

@@ -0,0 +1,191 @@
+{
+  "model_type": "bananamind_tts",
+  "architectures": [
+    "BananaMindTTSModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_bananamind_tts.BananaMindTTSConfig",
+    "AutoModel": "modeling_bananamind_tts.BananaMindTTSModel"
+  },
+  "banana_config": {
+    "project": {
+      "name": "BananaTTS-20M-Tacotron",
+      "seed": 1337
+    },
+    "dataset": {
+      "name": "MikhailT/lj-speech",
+      "fallback_name": "keithito/lj_speech",
+      "split": "train",
+      "cache_dir": "data/cache/ljspeech_22050",
+      "validation_ratio": 0.02,
+      "limit": null,
+      "percent": null,
+      "min_audio_seconds": 0.25,
+      "max_audio_seconds": 20.0
+    },
+    "text": {
+      "use_phonemes": false,
+      "keep_punctuation": true,
+      "max_tokens": 256
+    },
+    "audio": {
+      "sample_rate": 22050,
+      "n_fft": 1024,
+      "hop_length": 256,
+      "win_length": 1024,
+      "n_mels": 80,
+      "f_min": 0.0,
+      "f_max": 8000.0,
+      "power": 1.0,
+      "griffin_lim_iters": 32,
+      "mel_stats": {
+        "mean": -1.2497830139289026,
+        "std": 1.9540364596741207,
+        "min": -11.512925148010254,
+        "max": 6.246073246002197,
+        "count": 593805760,
+        "normalized_training": true
+      }
+    },
+    "model": {
+      "type": "tacotron_lite",
+      "hidden_size": 256,
+      "reduction_factor": 4,
+      "encoder_conv_layers": 3,
+      "encoder_dropout": 0.15,
+      "prenet_sizes": [
+        256,
+        128
+      ],
+      "attention_dim": 128,
+      "decoder_dim": 512,
+      "location_channels": 32,
+      "location_kernel_size": 31,
+      "postnet_channels": 512,
+      "postnet_layers": 5,
+      "dropout": 0.5,
+      "postnet_dropout": 0.5,
+      "max_decoder_steps": 1200,
+      "stop_threshold": 0.55,
+      "attention_window": 12
+    },
+    "training": {
+      "batch_size": 16,
+      "epochs": 80,
+      "learning_rate": 0.001,
+      "weight_decay": 1e-06,
+      "grad_clip": 1.0,
+      "log_interval": 20,
+      "val_interval": 500,
+      "save_interval": 1000,
+      "num_workers": 4,
+      "checkpoints_dir": "checkpoints_tacotron",
+      "runs_dir": "runs",
+      "resume": null,
+      "mel_loss_weight": 1.0,
+      "stop_loss_weight": 0.5,
+      "guided_attention_weight": 1.0,
+      "guided_attention_sigma": 0.4
+    }
+  },
+  "tokenizer": {
+    "symbols": [
+      "<pad>",
+      "<unk>",
+      "<bos>",
+      "<eos>",
+      "a",
+      "b",
+      "c",
+      "d",
+      "e",
+      "f",
+      "g",
+      "h",
+      "i",
+      "j",
+      "k",
+      "l",
+      "m",
+      "n",
+      "o",
+      "p",
+      "q",
+      "r",
+      "s",
+      "t",
+      "u",
+      "v",
+      "w",
+      "x",
+      "y",
+      "z",
+      " ",
+      "'",
+      ".",
+      ",",
+      "!",
+      "?",
+      ";",
+      ":",
+      "-"
+    ],
+    "use_phonemes": false,
+    "keep_punctuation": true
+  },
+  "vocoder": {
+    "enabled": true,
+    "type": "hifigan",
+    "default_weights": "vocoder.safetensors",
+    "default_dtype": "bfloat16",
+    "fp32_weights": "FP32/vocoder.safetensors",
+    "full_checkpoint": "full_vocoder/vocoder.pt",
+    "epoch": 28,
+    "step": 42000,
+    "generator_parameters": 3555649,
+    "config": {
+      "type": "hifigan",
+      "checkpoint": "checkpoints_vocoder/vocoder_latest.pt",
+      "initial_channels": 256,
+      "upsample_rates": [
+        8,
+        8,
+        2,
+        2
+      ],
+      "upsample_kernel_sizes": [
+        16,
+        16,
+        4,
+        4
+      ],
+      "resblock_kernel_sizes": [
+        3,
+        7,
+        11
+      ],
+      "resblock_dilation_sizes": [
+        [
+          1,
+          3,
+          5
+        ],
+        [
+          1,
+          3,
+          5
+        ],
+        [
+          1,
+          3,
+          5
+        ]
+      ],
+      "leaky_relu_slope": 0.1
+    }
+  },
+  "epoch": 80,
+  "step": 64529,
+  "torch_dtype": "float32",
+  "library_name": "transformers"
+}

configuration_bananamind_tts.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class BananaMindTTSConfig(PretrainedConfig):
+    model_type = "bananamind_tts"
+
+    def __init__(
+        self,
+        banana_config: dict | None = None,
+        tokenizer: dict | None = None,
+        vocoder: dict | None = None,
+        epoch: int | None = None,
+        step: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.banana_config = banana_config or {}
+        self.tokenizer = tokenizer or {}
+        self.vocoder = vocoder or {}
+        self.epoch = epoch
+        self.step = step

full_vocoder/vocoder.pt

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

generate.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from transformers import AutoModel
+
+
+MODEL_DIR = Path(__file__).resolve().parent
+
+
+def main() -> None:
+    model = AutoModel.from_pretrained(MODEL_DIR, trust_remote_code=True)
+    model.eval()
+    with torch.inference_mode():
+        out = model.tts(
+            "This is BananaMind TTS version two with a HiFi GAN vocoder.",
+            normalize_wav=True,
+        )
+    model.save_wav(MODEL_DIR / "sample.wav", out.waveform, out.sample_rate)
+    print(f"Saved {MODEL_DIR / 'sample.wav'}")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

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

model_config.json

@@ -0,0 +1,183 @@
+{
+  "model_type": "tacotron_lite",
+  "step": 64529,
+  "epoch": 80,
+  "config": {
+    "project": {
+      "name": "BananaTTS-20M-Tacotron",
+      "seed": 1337
+    },
+    "dataset": {
+      "name": "MikhailT/lj-speech",
+      "fallback_name": "keithito/lj_speech",
+      "split": "train",
+      "cache_dir": "data/cache/ljspeech_22050",
+      "validation_ratio": 0.02,
+      "limit": null,
+      "percent": null,
+      "min_audio_seconds": 0.25,
+      "max_audio_seconds": 20.0
+    },
+    "text": {
+      "use_phonemes": false,
+      "keep_punctuation": true,
+      "max_tokens": 256
+    },
+    "audio": {
+      "sample_rate": 22050,
+      "n_fft": 1024,
+      "hop_length": 256,
+      "win_length": 1024,
+      "n_mels": 80,
+      "f_min": 0.0,
+      "f_max": 8000.0,
+      "power": 1.0,
+      "griffin_lim_iters": 32,
+      "mel_stats": {
+        "mean": -1.2497830139289026,
+        "std": 1.9540364596741207,
+        "min": -11.512925148010254,
+        "max": 6.246073246002197,
+        "count": 593805760,
+        "normalized_training": true
+      }
+    },
+    "model": {
+      "type": "tacotron_lite",
+      "hidden_size": 256,
+      "reduction_factor": 4,
+      "encoder_conv_layers": 3,
+      "encoder_dropout": 0.15,
+      "prenet_sizes": [
+        256,
+        128
+      ],
+      "attention_dim": 128,
+      "decoder_dim": 512,
+      "location_channels": 32,
+      "location_kernel_size": 31,
+      "postnet_channels": 512,
+      "postnet_layers": 5,
+      "dropout": 0.5,
+      "postnet_dropout": 0.5,
+      "max_decoder_steps": 1200,
+      "stop_threshold": 0.55,
+      "attention_window": 12
+    },
+    "training": {
+      "batch_size": 16,
+      "epochs": 80,
+      "learning_rate": 0.001,
+      "weight_decay": 1e-06,
+      "grad_clip": 1.0,
+      "log_interval": 20,
+      "val_interval": 500,
+      "save_interval": 1000,
+      "num_workers": 4,
+      "checkpoints_dir": "checkpoints_tacotron",
+      "runs_dir": "runs",
+      "resume": null,
+      "mel_loss_weight": 1.0,
+      "stop_loss_weight": 0.5,
+      "guided_attention_weight": 1.0,
+      "guided_attention_sigma": 0.4
+    }
+  },
+  "tokenizer": {
+    "symbols": [
+      "<pad>",
+      "<unk>",
+      "<bos>",
+      "<eos>",
+      "a",
+      "b",
+      "c",
+      "d",
+      "e",
+      "f",
+      "g",
+      "h",
+      "i",
+      "j",
+      "k",
+      "l",
+      "m",
+      "n",
+      "o",
+      "p",
+      "q",
+      "r",
+      "s",
+      "t",
+      "u",
+      "v",
+      "w",
+      "x",
+      "y",
+      "z",
+      " ",
+      "'",
+      ".",
+      ",",
+      "!",
+      "?",
+      ";",
+      ":",
+      "-"
+    ],
+    "use_phonemes": false,
+    "keep_punctuation": true
+  },
+  "weights": "model.safetensors",
+  "vocoder": {
+    "enabled": true,
+    "type": "hifigan",
+    "default_weights": "vocoder.safetensors",
+    "default_dtype": "bfloat16",
+    "fp32_weights": "FP32/vocoder.safetensors",
+    "full_checkpoint": "full_vocoder/vocoder.pt",
+    "epoch": 28,
+    "step": 42000,
+    "generator_parameters": 3555649,
+    "config": {
+      "type": "hifigan",
+      "checkpoint": "checkpoints_vocoder/vocoder_latest.pt",
+      "initial_channels": 256,
+      "upsample_rates": [
+        8,
+        8,
+        2,
+        2
+      ],
+      "upsample_kernel_sizes": [
+        16,
+        16,
+        4,
+        4
+      ],
+      "resblock_kernel_sizes": [
+        3,
+        7,
+        11
+      ],
+      "resblock_dilation_sizes": [
+        [
+          1,
+          3,
+          5
+        ],
+        [
+          1,
+          3,
+          5
+        ],
+        [
+          1,
+          3,
+          5
+        ]
+      ],
+      "leaky_relu_slope": 0.1
+    }
+  }
+}

modeling_bananamind_tts.py

@@ -0,0 +1,551 @@
+from __future__ import annotations
+
+import math
+import re
+import string
+import wave
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+
+from .configuration_bananamind_tts import BananaMindTTSConfig
+
+
+PAD = "<pad>"
+UNK = "<unk>"
+BOS = "<bos>"
+EOS = "<eos>"
+ALLOWED = set(string.ascii_lowercase + " '.,!?;:-")
+
+
+def normalize_text(text: str, keep_punctuation: bool = True) -> str:
+    text = text.lower()
+    text = text.replace("\u2019", "'").replace("\u2018", "'")
+    text = text.replace("\u201c", '"').replace("\u201d", '"')
+    text = text.replace("&", " and ")
+    allowed = ALLOWED if keep_punctuation else set(string.ascii_lowercase + " '")
+    text = "".join(ch if ch in allowed else " " for ch in text)
+    return re.sub(r"\s+", " ", text).strip()
+
+
+@dataclass(frozen=True)
+class TextTokenizer:
+    symbols: list[str]
+    keep_punctuation: bool = True
+
+    @property
+    def symbol_to_id(self) -> dict[str, int]:
+        return {s: i for i, s in enumerate(self.symbols)}
+
+    @property
+    def pad_id(self) -> int:
+        return self.symbol_to_id[PAD]
+
+    @property
+    def bos_id(self) -> int:
+        return self.symbol_to_id[BOS]
+
+    @property
+    def eos_id(self) -> int:
+        return self.symbol_to_id[EOS]
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.symbols)
+
+    def encode(self, text: str) -> list[int]:
+        normalized = normalize_text(text, keep_punctuation=self.keep_punctuation)
+        stoi = self.symbol_to_id
+        ids = [stoi.get(ch, stoi[UNK]) for ch in normalized]
+        return [self.bos_id] + ids + [self.eos_id]
+
+
+def tokenizer_from_config(config: BananaMindTTSConfig) -> TextTokenizer:
+    data = config.tokenizer
+    return TextTokenizer(symbols=list(data["symbols"]), keep_punctuation=bool(data.get("keep_punctuation", True)))
+
+
+def make_padding_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor:
+    max_len = int(max_len or lengths.max().item())
+    return torch.arange(max_len, device=lengths.device).unsqueeze(0) >= lengths.unsqueeze(1)
+
+
+class Prenet(nn.Module):
+    def __init__(self, in_dim: int, sizes: list[int], dropout: float):
+        super().__init__()
+        layers: list[nn.Module] = []
+        last = in_dim
+        for size in sizes:
+            layers.extend([nn.Linear(last, size), nn.ReLU(), nn.Dropout(dropout)])
+            last = size
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dropout: float):
+        super().__init__()
+        self.conv = nn.Conv1d(channels, channels, kernel_size, padding=kernel_size // 2)
+        self.bn = nn.BatchNorm1d(channels)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dropout(F.relu(self.bn(self.conv(x))))
+
+
+class Encoder(nn.Module):
+    def __init__(self, vocab_size: int, hidden_size: int, pad_id: int, dropout: float, conv_layers: int):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, hidden_size, padding_idx=pad_id)
+        self.convs = nn.ModuleList([ConvBlock(hidden_size, 5, dropout) for _ in range(conv_layers)])
+        self.lstm = nn.LSTM(hidden_size, hidden_size // 2, num_layers=1, batch_first=True, bidirectional=True)
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        x = self.embedding(tokens).transpose(1, 2)
+        for conv in self.convs:
+            x = conv(x)
+        self.lstm.flatten_parameters()
+        out, _ = self.lstm(x.transpose(1, 2))
+        return out
+
+
+class LocationSensitiveAttention(nn.Module):
+    def __init__(self, query_dim: int, memory_dim: int, attention_dim: int, location_channels: int, location_kernel: int):
+        super().__init__()
+        self.query_layer = nn.Linear(query_dim, attention_dim, bias=False)
+        self.memory_layer = nn.Linear(memory_dim, attention_dim, bias=False)
+        self.location_conv = nn.Conv1d(2, location_channels, location_kernel, padding=location_kernel // 2, bias=False)
+        self.location_layer = nn.Linear(location_channels, attention_dim, bias=False)
+        self.v = nn.Linear(attention_dim, 1, bias=True)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        memory: torch.Tensor,
+        processed_memory: torch.Tensor,
+        attention_weights: torch.Tensor,
+        attention_cum: torch.Tensor,
+        mask: torch.Tensor,
+        window_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        location = torch.stack([attention_weights, attention_cum], dim=1)
+        location = self.location_conv(location).transpose(1, 2)
+        energies = self.v(
+            torch.tanh(self.query_layer(query).unsqueeze(1) + processed_memory + self.location_layer(location))
+        ).squeeze(-1)
+        energies = energies.masked_fill(mask, -1e4)
+        if window_mask is not None:
+            energies = energies.masked_fill(window_mask, -1e4)
+        weights = F.softmax(energies, dim=-1)
+        context = torch.bmm(weights.unsqueeze(1), memory).squeeze(1)
+        return context, weights
+
+
+class Postnet(nn.Module):
+    def __init__(self, n_mels: int, channels: int, layers: int, dropout: float):
+        super().__init__()
+        modules: list[nn.Module] = []
+        in_channels = n_mels
+        for idx in range(layers):
+            out_channels = n_mels if idx == layers - 1 else channels
+            modules.extend([nn.Conv1d(in_channels, out_channels, 5, padding=2), nn.BatchNorm1d(out_channels)])
+            if idx < layers - 1:
+                modules.append(nn.Tanh())
+            modules.append(nn.Dropout(dropout))
+            in_channels = out_channels
+        self.net = nn.Sequential(*modules)
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        return self.net(mel.transpose(1, 2)).transpose(1, 2)
+
+
+class HiFiGANResBlock(nn.Module):
+    def __init__(self, channels: int, kernel_size: int, dilations: list[int], slope: float = 0.1):
+        super().__init__()
+        self.slope = slope
+        self.convs1 = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    channels,
+                    channels,
+                    kernel_size,
+                    padding=(kernel_size * dilation - dilation) // 2,
+                    dilation=dilation,
+                )
+                for dilation in dilations
+            ]
+        )
+        self.convs2 = nn.ModuleList(
+            [nn.Conv1d(channels, channels, kernel_size, padding=(kernel_size - 1) // 2) for _ in dilations]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for conv1, conv2 in zip(self.convs1, self.convs2, strict=True):
+            residual = x
+            y = F.leaky_relu(x, self.slope)
+            y = conv1(y)
+            y = F.leaky_relu(y, self.slope)
+            y = conv2(y)
+            x = y + residual
+        return x
+
+
+class HiFiGANGenerator(nn.Module):
+    def __init__(self, n_mels: int, config: dict[str, Any]):
+        super().__init__()
+        channels = int(config.get("initial_channels", 256))
+        upsample_rates = [int(x) for x in config.get("upsample_rates", [8, 8, 2, 2])]
+        upsample_kernels = [int(x) for x in config.get("upsample_kernel_sizes", [16, 16, 4, 4])]
+        resblock_kernels = [int(x) for x in config.get("resblock_kernel_sizes", [3, 7, 11])]
+        resblock_dilations = config.get("resblock_dilation_sizes", [[1, 3, 5], [1, 3, 5], [1, 3, 5]])
+        self.slope = float(config.get("leaky_relu_slope", 0.1))
+        self.conv_pre = nn.Conv1d(n_mels, channels, kernel_size=7, padding=3)
+        self.ups = nn.ModuleList()
+        self.resblocks = nn.ModuleList()
+        current_channels = channels
+        for rate, kernel in zip(upsample_rates, upsample_kernels, strict=True):
+            next_channels = current_channels // 2
+            self.ups.append(
+                nn.ConvTranspose1d(
+                    current_channels,
+                    next_channels,
+                    kernel_size=kernel,
+                    stride=rate,
+                    padding=(kernel - rate) // 2,
+                )
+            )
+            for kernel_size, dilations in zip(resblock_kernels, resblock_dilations, strict=True):
+                self.resblocks.append(HiFiGANResBlock(next_channels, int(kernel_size), [int(d) for d in dilations], self.slope))
+            current_channels = next_channels
+        self.num_resblocks = len(resblock_kernels)
+        self.conv_post = nn.Conv1d(current_channels, 1, kernel_size=7, padding=3)
+
+    def forward(self, mel: torch.Tensor) -> torch.Tensor:
+        x = self.conv_pre(mel)
+        for idx, upsample in enumerate(self.ups):
+            x = F.leaky_relu(x, self.slope)
+            x = upsample(x)
+            block_offset = idx * self.num_resblocks
+            xs = [self.resblocks[block_offset + block_idx](x) for block_idx in range(self.num_resblocks)]
+            x = torch.stack(xs, dim=0).mean(dim=0)
+        x = F.leaky_relu(x, self.slope)
+        return torch.tanh(self.conv_post(x)).squeeze(1)
+
+
+@dataclass
+class TTSOutput:
+    waveform: torch.Tensor
+    sample_rate: int
+    mel: torch.Tensor
+    alignments: torch.Tensor
+    stop_logits: torch.Tensor
+
+
+def _resolve_asset(config: BananaMindTTSConfig, filename: str) -> Path:
+    model_id = getattr(config, "_name_or_path", "") or ""
+    if model_id and Path(model_id).exists():
+        return Path(model_id) / filename
+    local = Path(filename)
+    if local.exists():
+        return local
+    try:
+        from huggingface_hub import hf_hub_download
+
+        return Path(hf_hub_download(repo_id=model_id, filename=filename))
+    except Exception as exc:
+        raise FileNotFoundError(f"Could not resolve model asset {filename!r} for {model_id!r}") from exc
+
+
+class BananaMindTTSModel(PreTrainedModel):
+    config_class = BananaMindTTSConfig
+    base_model_prefix = ""
+    main_input_name = "tokens"
+    _tied_weights_keys: list[str] = []
+    all_tied_weights_keys: dict[str, str] = {}
+    _keys_to_ignore_on_load_missing = [r"vocoder_generator\..*"]
+
+    def __init__(self, config: BananaMindTTSConfig):
+        super().__init__(config)
+        self.tokenizer = tokenizer_from_config(config)
+        banana_config = config.banana_config
+        model_config = banana_config["model"]
+        audio_config = banana_config["audio"]
+        hidden = int(model_config.get("hidden_size", 256))
+        dropout = float(model_config.get("dropout", 0.5))
+        prenet_sizes = list(model_config.get("prenet_sizes", [256, 128]))
+        attention_dim = int(model_config.get("attention_dim", 128))
+        decoder_dim = int(model_config.get("decoder_dim", 512))
+        n_mels = int(audio_config["n_mels"])
+
+        self.n_mels = n_mels
+        self.reduction_factor = int(model_config.get("reduction_factor", 4))
+        self.encoder = Encoder(
+            self.tokenizer.vocab_size,
+            hidden,
+            self.tokenizer.pad_id,
+            float(model_config.get("encoder_dropout", 0.15)),
+            int(model_config.get("encoder_conv_layers", 3)),
+        )
+        self.prenet = Prenet(n_mels, prenet_sizes, dropout)
+        self.attention_rnn = nn.GRUCell(prenet_sizes[-1] + hidden, decoder_dim)
+        self.attention = LocationSensitiveAttention(
+            decoder_dim,
+            hidden,
+            attention_dim,
+            int(model_config.get("location_channels", 32)),
+            int(model_config.get("location_kernel_size", 31)),
+        )
+        self.decoder_rnn = nn.GRUCell(decoder_dim + hidden, decoder_dim)
+        proj_dim = decoder_dim + hidden
+        self.mel_proj = nn.Linear(proj_dim, n_mels * self.reduction_factor)
+        self.stop_proj = nn.Linear(proj_dim, self.reduction_factor)
+        self.postnet = Postnet(
+            n_mels,
+            int(model_config.get("postnet_channels", 512)),
+            int(model_config.get("postnet_layers", 5)),
+            float(model_config.get("postnet_dropout", 0.5)),
+        )
+        self.vocoder_generator: HiFiGANGenerator | None = None
+        self.vocoder_dtype = torch.float32
+        self._vocoder_loaded = False
+
+    def ensure_default_vocoder(self) -> None:
+        if self._vocoder_loaded:
+            return
+        self._vocoder_loaded = True
+        vocoder_info = getattr(self.config, "vocoder", {}) or {}
+        if not vocoder_info.get("enabled", True):
+            return
+        filename = str(vocoder_info.get("default_weights", "vocoder.safetensors"))
+        dtype_name = str(vocoder_info.get("default_dtype", "bfloat16")).lower()
+        self.vocoder_dtype = torch.bfloat16 if dtype_name in {"bf16", "bfloat16"} else torch.float16 if dtype_name in {"fp16", "float16"} else torch.float32
+        banana_config = self.config.banana_config
+        audio_config = banana_config["audio"]
+        vocoder_config = vocoder_info.get("config") or banana_config.get("vocoder") or {}
+        model_device = next(self.parameters()).device
+        self.vocoder_generator = HiFiGANGenerator(int(audio_config["n_mels"]), vocoder_config).to(
+            device=model_device,
+            dtype=self.vocoder_dtype,
+        )
+        try:
+            from safetensors.torch import load_file
+
+            state = load_file(str(_resolve_asset(self.config, filename)), device="cpu")
+            self.vocoder_generator.load_state_dict(state)
+            self.vocoder_generator.eval()
+        except Exception:
+            self.vocoder_generator = None
+            raise
+
+    def reload_vocoder(self, weights: str = "vocoder.safetensors", dtype: str = "bfloat16"):
+        self.config.vocoder["default_weights"] = weights
+        self.config.vocoder["default_dtype"] = dtype
+        self.vocoder_generator = None
+        self._vocoder_loaded = False
+        self.ensure_default_vocoder()
+        return self
+
+    def initialize_decoder_states(self, memory: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        batch, text_len, hidden = memory.shape
+        attention_hidden = memory.new_zeros(batch, self.attention_rnn.hidden_size)
+        decoder_hidden = memory.new_zeros(batch, self.decoder_rnn.hidden_size)
+        context = memory.new_zeros(batch, hidden)
+        attention_weights = memory.new_zeros(batch, text_len)
+        attention_weights[:, 0] = 1.0
+        attention_cum = memory.new_zeros(batch, text_len)
+        attention_cum[:, 0] = 1.0
+        return attention_hidden, decoder_hidden, context, attention_weights, attention_cum
+
+    def decode_step(
+        self,
+        decoder_input: torch.Tensor,
+        memory: torch.Tensor,
+        processed_memory: torch.Tensor,
+        mask: torch.Tensor,
+        attention_hidden: torch.Tensor,
+        decoder_hidden: torch.Tensor,
+        context: torch.Tensor,
+        attention_weights: torch.Tensor,
+        attention_cum: torch.Tensor,
+        window_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        prenet_out = self.prenet(decoder_input)
+        attention_hidden = self.attention_rnn(torch.cat([prenet_out, context], dim=-1), attention_hidden)
+        context, attention_weights = self.attention(attention_hidden, memory, processed_memory, attention_weights, attention_cum, mask, window_mask)
+        attention_cum = attention_cum + attention_weights
+        decoder_hidden = self.decoder_rnn(torch.cat([attention_hidden, context], dim=-1), decoder_hidden)
+        proj_input = torch.cat([decoder_hidden, context], dim=-1)
+        mel = self.mel_proj(proj_input).view(proj_input.shape[0], self.reduction_factor, self.n_mels)
+        stop = self.stop_proj(proj_input)
+        return mel, stop, attention_hidden, decoder_hidden, context, attention_weights, attention_cum
+
+    @torch.no_grad()
+    def infer(
+        self,
+        tokens: torch.Tensor,
+        token_lens: torch.Tensor,
+        max_steps: int = 1200,
+        stop_threshold: float = 0.55,
+        min_steps: int = 30,
+        attention_window: int = 12,
+    ):
+        memory = self.encoder(tokens)
+        mask = make_padding_mask(token_lens, tokens.shape[1])
+        processed_memory = self.attention.memory_layer(memory)
+        states = self.initialize_decoder_states(memory)
+        decoder_input = memory.new_zeros(tokens.shape[0], self.n_mels)
+        mel_outputs: list[torch.Tensor] = []
+        stop_outputs: list[torch.Tensor] = []
+        alignments: list[torch.Tensor] = []
+        max_decoder_steps = max(1, max_steps // self.reduction_factor)
+        min_decoder_steps = max(1, min_steps // self.reduction_factor)
+        prev_attention_index = torch.zeros(tokens.shape[0], dtype=torch.long, device=tokens.device)
+        for step in range(max_decoder_steps):
+            window_mask = None
+            if attention_window > 0:
+                positions = torch.arange(tokens.shape[1], device=tokens.device).unsqueeze(0)
+                left = (prev_attention_index - 1).clamp_min(0).unsqueeze(1)
+                right = (prev_attention_index + attention_window).clamp_max(tokens.shape[1] - 1).unsqueeze(1)
+                window_mask = (positions < left) | (positions > right) | mask
+            mel, stop, *states = self.decode_step(decoder_input, memory, processed_memory, mask, *states, window_mask=window_mask)
+            mel_outputs.append(mel)
+            stop_outputs.append(stop)
+            alignments.append(states[-2])
+            prev_attention_index = states[-2].argmax(dim=-1).clamp_min(prev_attention_index)
+            decoder_input = mel[:, -1]
+            if step >= min_decoder_steps and torch.sigmoid(stop[:, -1]).min().item() > stop_threshold:
+                break
+        mel = torch.stack(mel_outputs, dim=1).reshape(tokens.shape[0], -1, self.n_mels)
+        stop_logits = torch.stack(stop_outputs, dim=1).reshape(tokens.shape[0], -1)
+        align = torch.stack(alignments, dim=1)
+        mel_postnet = mel + self.postnet(mel)
+        return mel_postnet, align, stop_logits
+
+    def denormalize_mel(self, mel: torch.Tensor) -> torch.Tensor:
+        stats = self.config.banana_config.get("audio", {}).get("mel_stats")
+        if not stats or not stats.get("normalized_training", False):
+            return mel
+        mel = mel * max(float(stats["std"]), 1e-5) + float(stats["mean"])
+        return mel.clamp(float(stats.get("min", -12.0)), float(stats.get("max", 8.0)))
+
+    @torch.no_grad()
+    def vocode(self, mel: torch.Tensor) -> torch.Tensor:
+        audio_config = self.config.banana_config["audio"]
+        self.ensure_default_vocoder()
+        if self.vocoder_generator is None:
+            return griffin_lim(mel.detach().cpu(), audio_config)
+        vocoder_device = next(self.vocoder_generator.parameters()).device
+        mel_batch = mel.to(device=vocoder_device, dtype=self.vocoder_dtype).transpose(0, 1).unsqueeze(0)
+        wav = self.vocoder_generator(mel_batch).float()
+        target_samples = max(1, (mel.shape[0] - 1) * int(audio_config["hop_length"]))
+        if wav.shape[-1] < target_samples:
+            wav = F.pad(wav, (0, target_samples - wav.shape[-1]))
+        return wav[..., :target_samples].squeeze(0).detach().cpu()
+
+    @torch.no_grad()
+    def tts(
+        self,
+        text: str,
+        max_steps: int | None = None,
+        stop_threshold: float | None = None,
+        attention_window: int | None = None,
+        normalize_wav: bool = True,
+    ) -> TTSOutput:
+        device = next(self.parameters()).device
+        model_config = self.config.banana_config["model"]
+        audio_config = self.config.banana_config["audio"]
+        token_ids = self.tokenizer.encode(text)
+        tokens = torch.tensor(token_ids, dtype=torch.long, device=device).unsqueeze(0)
+        token_lens = torch.tensor([len(token_ids)], dtype=torch.long, device=device)
+        mel, alignments, stop_logits = self.infer(
+            tokens,
+            token_lens,
+            max_steps=int(max_steps or model_config.get("max_decoder_steps", 1200)),
+            stop_threshold=float(stop_threshold or model_config.get("stop_threshold", 0.55)),
+            min_steps=max(20, len(token_ids) * 3),
+            attention_window=int(attention_window if attention_window is not None else model_config.get("attention_window", 12)),
+        )
+        mel = self.denormalize_mel(mel[0].detach())
+        wav = self.vocode(mel)
+        if normalize_wav:
+            peak = wav.abs().max()
+            if peak > 0:
+                wav = wav / peak * 0.95
+        return TTSOutput(
+            waveform=wav,
+            sample_rate=int(audio_config["sample_rate"]),
+            mel=mel.detach().cpu(),
+            alignments=alignments.detach().cpu(),
+            stop_logits=stop_logits.detach().cpu(),
+        )
+
+    def forward(self, tokens: torch.Tensor, token_lens: torch.Tensor, **kwargs):
+        return self.infer(tokens, token_lens, **kwargs)
+
+    @staticmethod
+    def save_wav(path: str | Path, waveform: torch.Tensor, sample_rate: int) -> None:
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        wav_np = waveform.detach().cpu().clamp(-1.0, 1.0).numpy()
+        wav_i16 = np.clip(wav_np * 32767.0, -32768, 32767).astype(np.int16)
+        with wave.open(str(path), "wb") as f:
+            f.setnchannels(1)
+            f.setsampwidth(2)
+            f.setframerate(sample_rate)
+            f.writeframes(wav_i16.tobytes())
+
+
+def hz_to_mel(freq: torch.Tensor) -> torch.Tensor:
+    return 2595.0 * torch.log10(1.0 + freq / 700.0)
+
+
+def mel_to_hz(mel: torch.Tensor) -> torch.Tensor:
+    return 700.0 * (10.0 ** (mel / 2595.0) - 1.0)
+
+
+def mel_filterbank(sample_rate: int, n_fft: int, n_mels: int, f_min: float, f_max: float) -> torch.Tensor:
+    f_max = min(float(f_max), sample_rate / 2)
+    m_min = hz_to_mel(torch.tensor(float(f_min)))
+    m_max = hz_to_mel(torch.tensor(float(f_max)))
+    m_points = torch.linspace(m_min, m_max, n_mels + 2)
+    f_points = mel_to_hz(m_points)
+    bins = torch.floor((n_fft + 1) * f_points / sample_rate).long()
+    fb = torch.zeros(n_mels, n_fft // 2 + 1)
+    for i in range(n_mels):
+        left, center, right = bins[i].item(), bins[i + 1].item(), bins[i + 2].item()
+        center = max(center, left + 1)
+        right = max(right, center + 1)
+        for j in range(left, min(center, fb.shape[1])):
+            fb[i, j] = (j - left) / max(1, center - left)
+        for j in range(center, min(right, fb.shape[1])):
+            fb[i, j] = (right - j) / max(1, right - center)
+    return fb
+
+
+def griffin_lim(mel: torch.Tensor, config: dict[str, Any]) -> torch.Tensor:
+    n_fft = int(config["n_fft"])
+    hop_length = int(config["hop_length"])
+    win_length = int(config["win_length"])
+    n_iters = int(config.get("griffin_lim_iters", 32))
+    window = torch.hann_window(win_length)
+    fb = mel_filterbank(int(config["sample_rate"]), n_fft, int(config["n_mels"]), float(config.get("f_min", 0.0)), float(config.get("f_max", 8000.0)))
+    mag = torch.linalg.pinv(fb) @ torch.exp(mel.float()).transpose(0, 1)
+    mag = mag.clamp_min(1e-6)
+    spec = torch.polar(mag, 2 * math.pi * torch.rand_like(mag))
+    wav_len = max(1, (mel.shape[0] - 1) * hop_length)
+    wav = torch.istft(spec, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=wav_len)
+    for _ in range(n_iters):
+        rebuilt = torch.stft(wav, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, return_complex=True)
+        spec = mag * torch.exp(1j * rebuilt.angle())
+        wav = torch.istft(spec, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=wav_len)
+    wav = torch.nan_to_num(wav.detach().cpu())
+    peak = wav.abs().max()
+    return wav / peak if peak > 1.0 else wav.clamp(-1.0, 1.0)

vocoder.safetensors

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