initial commit

.gitignore

@@ -0,0 +1,173 @@
+# Customed
+.vscode/
+tests/
+runs/
+data/
+ckpts/
+wandb/
+results/
+
+
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Yushen CHEN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,14 +1,186 @@
----
-title: Gradio Text To Speech App
-emoji: ⚡
-colorFrom: gray
-colorTo: purple
-sdk: gradio
-sdk_version: 5.1.0
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: This TTS app uses E2-TTS model for precise voice cloning
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching
+
+[![arXiv](https://img.shields.io/badge/arXiv-2410.06885-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2410.06885)
+[![demo](https://img.shields.io/badge/GitHub-Demo%20page-blue.svg)](https://swivid.github.io/F5-TTS/)
+[![space](https://img.shields.io/badge/🤗-Space%20demo-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+
+**F5-TTS**: Diffusion Transformer with ConvNeXt V2, faster trained and inference.
+
+**E2 TTS**: Flat-UNet Transformer, closest reproduction.
+
+**Sway Sampling**: Inference-time flow step sampling strategy, greatly improves performance
+
+## Installation
+
+Clone the repository:
+
+```bash
+git clone https://github.com/SWivid/F5-TTS.git
+cd F5-TTS
+```
+
+Install packages:
+
+```bash
+pip install -r requirements.txt
+```
+
+Install torch with your CUDA version, e.g. :
+
+```bash
+pip install torch==2.3.0+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+pip install torchaudio==2.3.0+cu118 --extra-index-url https://download.pytorch.org/whl/cu118
+```
+
+## Prepare Dataset
+
+Example data processing scripts for Emilia and Wenetspeech4TTS, and you may tailor your own one along with a Dataset class in `model/dataset.py`.
+
+```bash
+# prepare custom dataset up to your need
+# download corresponding dataset first, and fill in the path in scripts
+
+# Prepare the Emilia dataset
+python scripts/prepare_emilia.py
+
+# Prepare the Wenetspeech4TTS dataset
+python scripts/prepare_wenetspeech4tts.py
+```
+
+## Training
+
+Once your datasets are prepared, you can start the training process.
+
+```bash
+# setup accelerate config, e.g. use multi-gpu ddp, fp16
+# will be to: ~/.cache/huggingface/accelerate/default_config.yaml     
+accelerate config
+accelerate launch test_train.py
+```
+
+## Inference
+
+To run inference with pretrained models, download the checkpoints from [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS).
+
+### Single Inference
+
+You can test single inference using the following command. Before running the command, modify the config up to your need.
+
+```bash
+# modify the config up to your need,
+# e.g. fix_duration (the total length of prompt + to_generate, currently support up to 30s)
+#      nfe_step     (larger takes more time to do more precise inference ode)
+#      ode_method   (switch to 'midpoint' for better compatibility with small nfe_step, )
+#                   ( though 'midpoint' is 2nd-order ode solver, slower compared to 1st-order 'Euler')
+python test_infer_single.py
+```
+### Speech Editing
+
+To test speech editing capabilities, use the following command.
+
+```bash
+python test_infer_single_edit.py
+```
+
+### Gradio App
+
+You can launch a Gradio app (web interface) to launch a GUI for inference.
+
+First, make sure you have the dependencies installed (`pip install -r requirements.txt`). Then, install the Gradio app dependencies:
+
+```bash
+pip install -r requirements_gradio.txt
+```
+
+After installing the dependencies, launch the app (will load ckpt from Huggingface, you may set `ckpt_path` to local file in `gradio_app.py`):
+
+```bash
+python gradio_app.py
+```
+
+You can specify the port/host:
+
+```bash
+python gradio_app.py --port 7860 --host 0.0.0.0
+```
+
+Or launch a share link:
+
+```bash
+python gradio_app.py --share
+```
+
+## Evaluation
+
+### Prepare Test Datasets
+
+1. Seed-TTS test set: Download from [seed-tts-eval](https://github.com/BytedanceSpeech/seed-tts-eval).
+2. LibriSpeech test-clean: Download from [OpenSLR](http://www.openslr.org/12/).
+3. Unzip the downloaded datasets and place them in the data/ directory.
+4. Update the path for the test-clean data in `test_infer_batch.py`
+5. Our filtered LibriSpeech-PC 4-10s subset is already under data/ in this repo
+
+### Batch Inference for Test Set
+
+To run batch inference for evaluations, execute the following commands:
+
+```bash
+# batch inference for evaluations
+accelerate config  # if not set before
+bash test_infer_batch.sh
+```
+
+### Download Evaluation Model Checkpoints
+
+1. Chinese ASR Model: [Paraformer-zh](https://huggingface.co/funasr/paraformer-zh)
+2. English ASR Model: [Faster-Whisper](https://huggingface.co/Systran/faster-whisper-large-v3)
+3. WavLM Model: Download from [Google Drive](https://drive.google.com/file/d/1-aE1NfzpRCLxA4GUxX9ITI3F9LlbtEGP/view).
+
+### Objective Evaluation
+
+**Some Notes**
+
+For faster-whisper with CUDA 11:
+
+```bash
+pip install --force-reinstall ctranslate2==3.24.0
+```
+
+(Recommended) To avoid possible ASR failures, such as abnormal repetitions in output:
+
+```bash
+pip install faster-whisper==0.10.1
+```
+
+Update the path with your batch-inferenced results, and carry out WER / SIM evaluations:
+```bash
+# Evaluation for Seed-TTS test set
+python scripts/eval_seedtts_testset.py
+
+# Evaluation for LibriSpeech-PC test-clean (cross-sentence)
+python scripts/eval_librispeech_test_clean.py
+```
+
+## Acknowledgements
+
+- [E2-TTS](https://arxiv.org/abs/2406.18009) brilliant work, simple and effective
+- [Emilia](https://arxiv.org/abs/2407.05361), [WenetSpeech4TTS](https://arxiv.org/abs/2406.05763) valuable datasets
+- [lucidrains](https://github.com/lucidrains) initial CFM structure with also [bfs18](https://github.com/bfs18) for discussion
+- [SD3](https://arxiv.org/abs/2403.03206) & [Hugging Face diffusers](https://github.com/huggingface/diffusers) DiT and MMDiT code structure
+- [torchdiffeq](https://github.com/rtqichen/torchdiffeq) as ODE solver, [Vocos](https://huggingface.co/charactr/vocos-mel-24khz) as vocoder
+- [mrfakename](https://x.com/realmrfakename) huggingface space demo ~
+- [FunASR](https://github.com/modelscope/FunASR), [faster-whisper](https://github.com/SYSTRAN/faster-whisper), [UniSpeech](https://github.com/microsoft/UniSpeech) for evaluation tools
+- [ctc-forced-aligner](https://github.com/MahmoudAshraf97/ctc-forced-aligner) for speech edit test
+
+## Citation
+```
+@article{chen-etal-2024-f5tts,
+      title={F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching}, 
+      author={Yushen Chen and Zhikang Niu and Ziyang Ma and Keqi Deng and Chunhui Wang and Jian Zhao and Kai Yu and Xie Chen},
+      journal={arXiv preprint arXiv:2410.06885},
+      year={2024},
+}
+```
+## License
+
+Our code is released under MIT License.

gradio_app.py

@@ -0,0 +1,1049 @@
+# import os
+# import re
+# import torch
+# import torchaudio
+# import gradio as gr
+# import numpy as np
+# import tempfile
+# from einops import rearrange
+# from vocos import Vocos
+# from pydub import AudioSegment
+# from model import CFM, UNetT, DiT, MMDiT
+# from cached_path import cached_path
+# from model.utils import (
+#     load_checkpoint,
+#     get_tokenizer,
+#     convert_char_to_pinyin,
+#     save_spectrogram,
+# )
+# from transformers import pipeline
+# import librosa
+# import click
+
+# device = (
+#     "cuda"
+#     if torch.cuda.is_available()
+#     else "mps" if torch.backends.mps.is_available() else "cpu"
+# )
+
+# print(f"Using {device} device")
+
+# pipe = pipeline(
+#     "automatic-speech-recognition",
+#     model="openai/whisper-large-v3-turbo",
+#     torch_dtype=torch.float16,
+#     device=device,
+# )
+
+# # --------------------- Settings -------------------- #
+
+# target_sample_rate = 24000
+# n_mel_channels = 100
+# hop_length = 256
+# target_rms = 0.1
+# nfe_step = 32  # 16, 32
+# cfg_strength = 2.0
+# ode_method = "euler"
+# sway_sampling_coef = -1.0
+# speed = 1.0
+# # fix_duration = 27  # None or float (duration in seconds)
+# fix_duration = None
+
+
+# def load_model(exp_name, model_cls, model_cfg, ckpt_step):
+#     ckpt_path = str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.safetensors"))
+#     # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors
+#     vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
+#     model = CFM(
+#         transformer=model_cls(
+#             **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
+#         ),
+#         mel_spec_kwargs=dict(
+#             target_sample_rate=target_sample_rate,
+#             n_mel_channels=n_mel_channels,
+#             hop_length=hop_length,
+#         ),
+#         odeint_kwargs=dict(
+#             method=ode_method,
+#         ),
+#         vocab_char_map=vocab_char_map,
+#     ).to(device)
+
+#     model = load_checkpoint(model, ckpt_path, device, use_ema = True)
+
+#     return model
+
+
+# # load models
+# F5TTS_model_cfg = dict(
+#     dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+# )
+# E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+# F5TTS_ema_model = load_model(
+#     "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000
+# )
+# E2TTS_ema_model = load_model(
+#     "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
+# )
+
+
+# def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence):
+#     print(gen_text)
+#     if len(gen_text) > 200:
+#         raise gr.Error("Please keep your text under 200 chars.")
+#     gr.Info("Converting audio...")
+#     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+#         aseg = AudioSegment.from_file(ref_audio_orig)
+#         audio_duration = len(aseg)
+#         if audio_duration > 15000:
+#             gr.Warning("Audio is over 15s, clipping to only first 15s.")
+#             aseg = aseg[:15000]
+#         aseg.export(f.name, format="wav")
+#         ref_audio = f.name
+#     if exp_name == "F5-TTS":
+#         ema_model = F5TTS_ema_model
+#     elif exp_name == "E2-TTS":
+#         ema_model = E2TTS_ema_model
+
+#     if not ref_text.strip():
+#         gr.Info("No reference text provided, transcribing reference audio...")
+#         ref_text = outputs = pipe(
+#             ref_audio,
+#             chunk_length_s=30,
+#             batch_size=128,
+#             generate_kwargs={"task": "transcribe"},
+#             return_timestamps=False,
+#         )["text"].strip()
+#         gr.Info("Finished transcription")
+#     else:
+#         gr.Info("Using custom reference text...")
+#     audio, sr = torchaudio.load(ref_audio)
+#     if audio.shape[0] > 1:
+#         audio = torch.mean(audio, dim=0, keepdim=True)
+
+#     rms = torch.sqrt(torch.mean(torch.square(audio)))
+#     if rms < target_rms:
+#         audio = audio * target_rms / rms
+#     if sr != target_sample_rate:
+#         resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+#         audio = resampler(audio)
+#     audio = audio.to(device)
+
+#     # Prepare the text
+#     text_list = [ref_text + gen_text]
+#     final_text_list = convert_char_to_pinyin(text_list)
+
+#     # Calculate duration
+#     ref_audio_len = audio.shape[-1] // hop_length
+#     # if fix_duration is not None:
+#     #     duration = int(fix_duration * target_sample_rate / hop_length)
+#     # else:
+#     zh_pause_punc = r"。，、；：？！"
+#     ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))
+#     gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))
+#     duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+#     # inference
+#     gr.Info(f"Generating audio using {exp_name}")
+#     with torch.inference_mode():
+#         generated, _ = ema_model.sample(
+#             cond=audio,
+#             text=final_text_list,
+#             duration=duration,
+#             steps=nfe_step,
+#             cfg_strength=cfg_strength,
+#             sway_sampling_coef=sway_sampling_coef,
+#         )
+
+#     generated = generated[:, ref_audio_len:, :]
+#     generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+#     gr.Info("Running vocoder")
+#     vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+#     generated_wave = vocos.decode(generated_mel_spec.cpu())
+#     if rms < target_rms:
+#         generated_wave = generated_wave * rms / target_rms
+
+#     # wav -> numpy
+#     generated_wave = generated_wave.squeeze().cpu().numpy()
+
+#     if remove_silence:
+#         gr.Info("Removing audio silences... This may take a moment")
+#         non_silent_intervals = librosa.effects.split(generated_wave, top_db=30)
+#         non_silent_wave = np.array([])
+#         for interval in non_silent_intervals:
+#             start, end = interval
+#             non_silent_wave = np.concatenate(
+#                 [non_silent_wave, generated_wave[start:end]]
+#             )
+#         generated_wave = non_silent_wave
+
+#     # spectogram
+#     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+#         spectrogram_path = tmp_spectrogram.name
+#         save_spectrogram(generated_mel_spec[0].cpu().numpy(), spectrogram_path)
+
+#     return (target_sample_rate, generated_wave), spectrogram_path
+
+
+# with gr.Blocks() as app:
+#     gr.Markdown(
+#         """
+# # Antriksh AI
+
+# """
+#     )
+    
+#     # Image
+#     gr.Image(value="C:\\Users\\USER\\OneDrive\\Documents\\logo.jpg", width=300, height= 150 )
+
+#     ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+#     gen_text_input = gr.Textbox(label="Text to Generate (max 200 chars.)", lines=4)
+#     model_choice = gr.Radio(
+#         choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+#     )
+#     generate_btn = gr.Button("Synthesize", variant="primary")
+#     with gr.Accordion("Advanced Settings", open=False):
+#         ref_text_input = gr.Textbox(
+#             label="Reference Text",
+#             info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.",
+#             lines=2,
+#         )
+#         remove_silence = gr.Checkbox(
+#             label="Remove Silences",
+#             info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.",
+#             value=True,
+#         )
+
+#     audio_output = gr.Audio(label="Synthesized Audio")
+#     spectrogram_output = gr.Image(label="Spectrogram")
+
+#     generate_btn.click(
+#         infer,
+#         inputs=[
+#             ref_audio_input,
+#             ref_text_input,
+#             gen_text_input,
+#             model_choice,
+#             remove_silence,
+#         ],
+#         outputs=[audio_output, spectrogram_output],
+#     )
+
+
+# @click.command()
+# @click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+# @click.option("--host", "-H", default=None, help="Host to run the app on")
+# @click.option(
+#     "--share",
+#     "-s",
+#     default=True,
+#     is_flag=True,
+#     help="Share the app via Gradio share link",
+# )
+# @click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+# def main(port, host, share, api):
+#     global app
+#     print(f"Starting app...")
+#     app.queue(api_open=api).launch(
+#         server_name=host, server_port=port, share=True, show_api=api
+#     )
+
+
+# if __name__ == "__main__":
+#     main()
+
+import re
+import torch
+import torchaudio
+import gradio as gr
+import numpy as np
+import tempfile
+from einops import rearrange
+from vocos import Vocos
+from pydub import AudioSegment, silence
+from model import CFM, UNetT, DiT, MMDiT
+from cached_path import cached_path
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer,
+    convert_char_to_pinyin,
+    save_spectrogram,
+)
+from transformers import pipeline
+import click
+import soundfile as sf
+
+try:
+    import spaces
+    USING_SPACES = True
+except ImportError:
+    USING_SPACES = False
+
+def gpu_decorator(func):
+    if USING_SPACES:
+        return spaces.GPU(func)
+    else:
+        return func
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps" if torch.backends.mps.is_available() else "cpu"
+)
+
+print(f"Using {device} device")
+
+pipe = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo",
+    torch_dtype=torch.float16,
+    device=device,
+)
+vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# --------------------- Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+
+def load_model(repo_name, exp_name, model_cls, model_cfg, ckpt_step):
+    ckpt_path = str(cached_path(f"hf://SWivid/{repo_name}/{exp_name}/model_{ckpt_step}.safetensors"))
+    # ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"  # .pt | .safetensors
+    vocab_char_map, vocab_size = get_tokenizer("Emilia_ZH_EN", "pinyin")
+    model = CFM(
+        transformer=model_cls(
+            **model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels
+        ),
+        mel_spec_kwargs=dict(
+            target_sample_rate=target_sample_rate,
+            n_mel_channels=n_mel_channels,
+            hop_length=hop_length,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    model = load_checkpoint(model, ckpt_path, device, use_ema = True)
+
+    return model
+
+
+# load models
+F5TTS_model_cfg = dict(
+    dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4
+)
+E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4)
+
+F5TTS_ema_model = load_model(
+    "F5-TTS", "F5TTS_Base", DiT, F5TTS_model_cfg, 1200000
+)
+E2TTS_ema_model = load_model(
+    "E2-TTS", "E2TTS_Base", UNetT, E2TTS_model_cfg, 1200000
+)
+
+def chunk_text(text, max_chars=135):
+    """
+    Splits the input text into chunks, each with a maximum number of characters.
+
+    Args:
+        text (str): The text to be split.
+        max_chars (int): The maximum number of characters per chunk.
+
+    Returns:
+        List[str]: A list of text chunks.
+    """
+    chunks = []
+    current_chunk = ""
+    # Split the text into sentences based on punctuation followed by whitespace
+    sentences = re.split(r'(?<=[;:,.!?])\s+|(?<=[；：，。！？])', text)
+
+    for sentence in sentences:
+        if len(current_chunk.encode('utf-8')) + len(sentence.encode('utf-8')) <= max_chars:
+            current_chunk += sentence + " " if sentence and len(sentence[-1].encode('utf-8')) == 1 else sentence
+        else:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+            current_chunk = sentence + " " if sentence and len(sentence[-1].encode('utf-8')) == 1 else sentence
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
+    return chunks
+
+@gpu_decorator
+def infer_batch(ref_audio, ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration=0.15, progress=gr.Progress()):
+    if exp_name == "F5-TTS":
+        ema_model = F5TTS_ema_model
+    elif exp_name == "E2-TTS":
+        ema_model = E2TTS_ema_model
+
+    audio, sr = ref_audio
+    if audio.shape[0] > 1:
+        audio = torch.mean(audio, dim=0, keepdim=True)
+
+    rms = torch.sqrt(torch.mean(torch.square(audio)))
+    if rms < target_rms:
+        audio = audio * target_rms / rms
+    if sr != target_sample_rate:
+        resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    generated_waves = []
+    spectrograms = []
+
+    for i, gen_text in enumerate(progress.tqdm(gen_text_batches)):
+        # Prepare the text
+        if len(ref_text[-1].encode('utf-8')) == 1:
+            ref_text = ref_text + " "
+        text_list = [ref_text + gen_text]
+        final_text_list = convert_char_to_pinyin(text_list)
+
+        # Calculate duration
+        ref_audio_len = audio.shape[-1] // hop_length
+        zh_pause_punc = r"。，、；：？！"
+        ref_text_len = len(ref_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, ref_text))
+        gen_text_len = len(gen_text.encode('utf-8')) + 3 * len(re.findall(zh_pause_punc, gen_text))
+        duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = ema_model.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+
+        generated = generated[:, ref_audio_len:, :]
+        generated_mel_spec = rearrange(generated, "1 n d -> 1 d n")
+        generated_wave = vocos.decode(generated_mel_spec.cpu())
+        if rms < target_rms:
+            generated_wave = generated_wave * rms / target_rms
+
+        # wav -> numpy
+        generated_wave = generated_wave.squeeze().cpu().numpy()
+        
+        generated_waves.append(generated_wave)
+        spectrograms.append(generated_mel_spec[0].cpu().numpy())
+
+    # Combine all generated waves with cross-fading
+    if cross_fade_duration <= 0:
+        # Simply concatenate
+        final_wave = np.concatenate(generated_waves)
+    else:
+        final_wave = generated_waves[0]
+        for i in range(1, len(generated_waves)):
+            prev_wave = final_wave
+            next_wave = generated_waves[i]
+
+            # Calculate cross-fade samples, ensuring it does not exceed wave lengths
+            cross_fade_samples = int(cross_fade_duration * target_sample_rate)
+            cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))
+
+            if cross_fade_samples <= 0:
+                # No overlap possible, concatenate
+                final_wave = np.concatenate([prev_wave, next_wave])
+                continue
+
+            # Overlapping parts
+            prev_overlap = prev_wave[-cross_fade_samples:]
+            next_overlap = next_wave[:cross_fade_samples]
+
+            # Fade out and fade in
+            fade_out = np.linspace(1, 0, cross_fade_samples)
+            fade_in = np.linspace(0, 1, cross_fade_samples)
+
+            # Cross-faded overlap
+            cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in
+
+            # Combine
+            new_wave = np.concatenate([
+                prev_wave[:-cross_fade_samples],
+                cross_faded_overlap,
+                next_wave[cross_fade_samples:]
+            ])
+
+            final_wave = new_wave
+
+    # Remove silence
+    if remove_silence:
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+            sf.write(f.name, final_wave, target_sample_rate)
+            aseg = AudioSegment.from_file(f.name)
+            non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500)
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                non_silent_wave += non_silent_seg
+            aseg = non_silent_wave
+            aseg.export(f.name, format="wav")
+            final_wave, _ = torchaudio.load(f.name)
+        final_wave = final_wave.squeeze().cpu().numpy()
+
+    # Create a combined spectrogram
+    combined_spectrogram = np.concatenate(spectrograms, axis=1)
+    
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+        save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (target_sample_rate, final_wave), spectrogram_path
+
+@gpu_decorator
+def infer(ref_audio_orig, ref_text, gen_text, exp_name, remove_silence, cross_fade_duration=0.15):
+
+    print(gen_text)
+
+    gr.Info("Converting audio...")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        non_silent_segs = silence.split_on_silence(
+            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000
+        )
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            non_silent_wave += non_silent_seg
+        aseg = non_silent_wave
+
+        audio_duration = len(aseg)
+        if audio_duration > 15000:
+            gr.Warning("Audio is over 15s, clipping to only first 15s.")
+            aseg = aseg[:15000]
+        aseg.export(f.name, format="wav")
+        ref_audio = f.name
+
+    if not ref_text.strip():
+        gr.Info("No reference text provided, transcribing reference audio...")
+        ref_text = pipe(
+            ref_audio,
+            chunk_length_s=30,
+            batch_size=128,
+            generate_kwargs={"task": "transcribe"},
+            return_timestamps=False,
+        )["text"].strip()
+        gr.Info("Finished transcription")
+    else:
+        gr.Info("Using custom reference text...")
+
+    # Add the functionality to ensure it ends with ". "
+    if not ref_text.endswith(". "):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    audio, sr = torchaudio.load(ref_audio)
+
+    # Use the new chunk_text function to split gen_text
+    max_chars = int(len(ref_text.encode('utf-8')) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr))
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    print('ref_text', ref_text)
+    for i, batch_text in enumerate(gen_text_batches):
+        print(f'gen_text {i}', batch_text)
+    
+    gr.Info(f"Generating audio using {exp_name} in {len(gen_text_batches)} batches")
+    return infer_batch((audio, sr), ref_text, gen_text_batches, exp_name, remove_silence, cross_fade_duration)
+
+
+@gpu_decorator
+def generate_podcast(script, speaker1_name, ref_audio1, ref_text1, speaker2_name, ref_audio2, ref_text2, exp_name, remove_silence):
+    # Split the script into speaker blocks
+    speaker_pattern = re.compile(f"^({re.escape(speaker1_name)}|{re.escape(speaker2_name)}):", re.MULTILINE)
+    speaker_blocks = speaker_pattern.split(script)[1:]  # Skip the first empty element
+    
+    generated_audio_segments = []
+    
+    for i in range(0, len(speaker_blocks), 2):
+        speaker = speaker_blocks[i]
+        text = speaker_blocks[i+1].strip()
+        
+        # Determine which speaker is talking
+        if speaker == speaker1_name:
+            ref_audio = ref_audio1
+            ref_text = ref_text1
+        elif speaker == speaker2_name:
+            ref_audio = ref_audio2
+            ref_text = ref_text2
+        else:
+            continue  # Skip if the speaker is neither speaker1 nor speaker2
+        
+        # Generate audio for this block
+        audio, _ = infer(ref_audio, ref_text, text, exp_name, remove_silence)
+        
+        # Convert the generated audio to a numpy array
+        sr, audio_data = audio
+        
+        # Save the audio data as a WAV file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            sf.write(temp_file.name, audio_data, sr)
+            audio_segment = AudioSegment.from_wav(temp_file.name)
+        
+        generated_audio_segments.append(audio_segment)
+        
+        # Add a short pause between speakers
+        pause = AudioSegment.silent(duration=500)  # 500ms pause
+        generated_audio_segments.append(pause)
+    
+    # Concatenate all audio segments
+    final_podcast = sum(generated_audio_segments)
+    
+    # Export the final podcast
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+        podcast_path = temp_file.name
+        final_podcast.export(podcast_path, format="wav")
+    
+    return podcast_path
+
+def parse_speechtypes_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+def update_speed(new_speed):
+    global speed
+    speed = new_speed
+    return f"Speed set to: {speed}"
+
+with gr.Blocks() as app_credits:
+    gr.Markdown("""
+# Credits
+
+* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)
+* [RootingInLoad](https://github.com/RootingInLoad) for the podcast generation
+* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation
+""")
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    gen_text_input = gr.Textbox(label="Text to Generate", lines=10)
+    model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        ref_text_input = gr.Textbox(
+            label="Reference Text",
+            info="Leave blank to automatically transcribe the reference audio. If you enter text it will override automatic transcription.",
+            lines=2,
+        )
+        remove_silence = gr.Checkbox(
+            label="Remove Silences",
+            info="The model tends to produce silences, especially on longer audio. We can manually remove silences if needed. Note that this is an experimental feature and may produce strange results. This will also increase generation time.",
+            value=False,
+        )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=speed,
+            step=0.1,
+            info="Adjust the speed of the audio.",
+        )
+        cross_fade_duration_slider = gr.Slider(
+            label="Cross-Fade Duration (s)",
+            minimum=0.0,
+            maximum=1.0,
+            value=0.15,
+            step=0.01,
+            info="Set the duration of the cross-fade between audio clips.",
+        )
+    speed_slider.change(update_speed, inputs=speed_slider)
+
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    generate_btn.click(
+        infer,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            model_choice,
+            remove_silence,
+            cross_fade_duration_slider,
+        ],
+        outputs=[audio_output, spectrogram_output],
+    )
+    
+with gr.Blocks() as app_podcast:
+    gr.Markdown("# Podcast Generation")
+    speaker1_name = gr.Textbox(label="Speaker 1 Name")
+    ref_audio_input1 = gr.Audio(label="Reference Audio (Speaker 1)", type="filepath")
+    ref_text_input1 = gr.Textbox(label="Reference Text (Speaker 1)", lines=2)
+    
+    speaker2_name = gr.Textbox(label="Speaker 2 Name")
+    ref_audio_input2 = gr.Audio(label="Reference Audio (Speaker 2)", type="filepath")
+    ref_text_input2 = gr.Textbox(label="Reference Text (Speaker 2)", lines=2)
+    
+    script_input = gr.Textbox(label="Podcast Script", lines=10, 
+                                placeholder="Enter the script with speaker names at the start of each block, e.g.:\nSean: How did you start studying...\n\nMeghan: I came to my interest in technology...\nIt was a long journey...\n\nSean: That's fascinating. Can you elaborate...")
+    
+    podcast_model_choice = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+    podcast_remove_silence = gr.Checkbox(
+        label="Remove Silences",
+        value=True,
+    )
+    generate_podcast_btn = gr.Button("Generate Podcast", variant="primary")
+    podcast_output = gr.Audio(label="Generated Podcast")
+
+    def podcast_generation(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence):
+        return generate_podcast(script, speaker1, ref_audio1, ref_text1, speaker2, ref_audio2, ref_text2, model, remove_silence)
+
+    generate_podcast_btn.click(
+        podcast_generation,
+        inputs=[
+            script_input,
+            speaker1_name,
+            ref_audio_input1,
+            ref_text_input1,
+            speaker2_name,
+            ref_audio_input2,
+            ref_text_input2,
+            podcast_model_choice,
+            podcast_remove_silence,
+        ],
+        outputs=podcast_output,
+    )
+
+def parse_emotional_text(gen_text):
+    # Pattern to find (Emotion)
+    pattern = r'\((.*?)\)'
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_emotion = 'Regular'
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                segments.append({'emotion': current_emotion, 'text': text})
+        else:
+            # This is emotion
+            emotion = tokens[i].strip()
+            current_emotion = emotion
+
+    return segments
+
+with gr.Blocks() as app_emotional:
+    # New section for emotional generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the "Add Speech Type" button. Enter your text in the format shown below, and the system will generate speech using the appropriate emotions. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+
+    **Example Input:**
+
+    (Regular) Hello, I'd like to order a sandwich please. (Surprised) What do you mean you're out of bread? (Sad) I really wanted a sandwich though... (Angry) You know what, darn you and your little shop, you suck! (Whisper) I'll just go back home and cry now. (Shouting) Why me?!
+    """
+    )
+
+    gr.Markdown("Upload different audio clips for each speech type. 'Regular' emotion is mandatory. You can add additional speech types by clicking the 'Add Speech Type' button.")
+
+    # Regular speech type (mandatory)
+    with gr.Row():
+        regular_name = gr.Textbox(value='Regular', label='Speech Type Name', interactive=False)
+        regular_audio = gr.Audio(label='Regular Reference Audio', type='filepath')
+        regular_ref_text = gr.Textbox(label='Reference Text (Regular)', lines=2)
+
+    # Additional speech types (up to 99 more)
+    max_speech_types = 100
+    speech_type_names = []
+    speech_type_audios = []
+    speech_type_ref_texts = []
+    speech_type_delete_btns = []
+
+    for i in range(max_speech_types - 1):
+        with gr.Row():
+            name_input = gr.Textbox(label='Speech Type Name', visible=False)
+            audio_input = gr.Audio(label='Reference Audio', type='filepath', visible=False)
+            ref_text_input = gr.Textbox(label='Reference Text', lines=2, visible=False)
+            delete_btn = gr.Button("Delete", variant="secondary", visible=False)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_delete_btns.append(delete_btn)
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of current number of speech types
+    speech_type_count = gr.State(value=0)
+
+    # Function to add a speech type
+    def add_speech_type_fn(speech_type_count):
+        if speech_type_count < max_speech_types - 1:
+            speech_type_count += 1
+            # Prepare updates for the components
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+            for i in range(max_speech_types - 1):
+                if i < speech_type_count:
+                    name_updates.append(gr.update(visible=True))
+                    audio_updates.append(gr.update(visible=True))
+                    ref_text_updates.append(gr.update(visible=True))
+                    delete_btn_updates.append(gr.update(visible=True))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+        else:
+            # Optionally, show a warning
+            # gr.Warning("Maximum number of speech types reached.")
+            name_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            audio_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            ref_text_updates = [gr.update() for _ in range(max_speech_types - 1)]
+            delete_btn_updates = [gr.update() for _ in range(max_speech_types - 1)]
+        return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+    add_speech_type_btn.click(
+        add_speech_type_fn,
+        inputs=speech_type_count,
+        outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+    )
+
+    # Function to delete a speech type
+    def make_delete_speech_type_fn(index):
+        def delete_speech_type_fn(speech_type_count):
+            # Prepare updates
+            name_updates = []
+            audio_updates = []
+            ref_text_updates = []
+            delete_btn_updates = []
+
+            for i in range(max_speech_types - 1):
+                if i == index:
+                    name_updates.append(gr.update(visible=False, value=''))
+                    audio_updates.append(gr.update(visible=False, value=None))
+                    ref_text_updates.append(gr.update(visible=False, value=''))
+                    delete_btn_updates.append(gr.update(visible=False))
+                else:
+                    name_updates.append(gr.update())
+                    audio_updates.append(gr.update())
+                    ref_text_updates.append(gr.update())
+                    delete_btn_updates.append(gr.update())
+
+            speech_type_count = max(0, speech_type_count - 1)
+
+            return [speech_type_count] + name_updates + audio_updates + ref_text_updates + delete_btn_updates
+
+        return delete_speech_type_fn
+
+    for i, delete_btn in enumerate(speech_type_delete_btns):
+        delete_fn = make_delete_speech_type_fn(i)
+        delete_btn.click(
+            delete_fn,
+            inputs=speech_type_count,
+            outputs=[speech_type_count] + speech_type_names + speech_type_audios + speech_type_ref_texts + speech_type_delete_btns
+        )
+
+    # Text input for the prompt
+    gen_text_input_emotional = gr.Textbox(label="Text to Generate", lines=10)
+
+    # Model choice
+    model_choice_emotional = gr.Radio(
+        choices=["F5-TTS", "E2-TTS"], label="Choose TTS Model", value="F5-TTS"
+    )
+
+    with gr.Accordion("Advanced Settings", open=False):
+        remove_silence_emotional = gr.Checkbox(
+            label="Remove Silences",
+            value=True,
+        )
+
+    # Generate button
+    generate_emotional_btn = gr.Button("Generate Emotional Speech", variant="primary")
+
+    # Output audio
+    audio_output_emotional = gr.Audio(label="Synthesized Audio")
+    @gpu_decorator
+    def generate_emotional_speech(
+        regular_audio,
+        regular_ref_text,
+        gen_text,
+        *args,
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+        speech_type_audios_list = args[num_additional_speech_types:2 * num_additional_speech_types]
+        speech_type_ref_texts_list = args[2 * num_additional_speech_types:3 * num_additional_speech_types]
+        model_choice = args[3 * num_additional_speech_types]
+        remove_silence = args[3 * num_additional_speech_types + 1]
+
+        # Collect the speech types and their audios into a dict
+        speech_types = {'Regular': {'audio': regular_audio, 'ref_text': regular_ref_text}}
+
+        for name_input, audio_input, ref_text_input in zip(speech_type_names_list, speech_type_audios_list, speech_type_ref_texts_list):
+            if name_input and audio_input:
+                speech_types[name_input] = {'audio': audio_input, 'ref_text': ref_text_input}
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_emotion = 'Regular'
+
+        for segment in segments:
+            emotion = segment['emotion']
+            text = segment['text']
+
+            if emotion in speech_types:
+                current_emotion = emotion
+            else:
+                # If emotion not available, default to Regular
+                current_emotion = 'Regular'
+
+            ref_audio = speech_types[current_emotion]['audio']
+            ref_text = speech_types[current_emotion].get('ref_text', '')
+
+            # Generate speech for this segment
+            audio, _ = infer(ref_audio, ref_text, text, model_choice, remove_silence, 0)
+            sr, audio_data = audio
+
+            generated_audio_segments.append(audio_data)
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (sr, final_audio_data)
+        else:
+            gr.Warning("No audio generated.")
+            return None
+
+    generate_emotional_btn.click(
+        generate_emotional_speech,
+        inputs=[
+            regular_audio,
+            regular_ref_text,
+            gen_text_input_emotional,
+        ] + speech_type_names + speech_type_audios + speech_type_ref_texts + [
+            model_choice_emotional,
+            remove_silence_emotional,
+        ],
+        outputs=audio_output_emotional,
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(
+        gen_text,
+        regular_name,
+        *args
+    ):
+        num_additional_speech_types = max_speech_types - 1
+        speech_type_names_list = args[:num_additional_speech_types]
+
+        # Collect the speech types names
+        speech_types_available = set()
+        if regular_name:
+            speech_types_available.add(regular_name)
+        for name_input in speech_type_names_list:
+            if name_input:
+                speech_types_available.add(name_input)
+
+        # Parse the gen_text to get the speech types used
+        segments = parse_emotional_text(gen_text)
+        speech_types_in_text = set(segment['emotion'] for segment in segments)
+
+        # Check if all speech types in text are available
+        missing_speech_types = speech_types_in_text - speech_types_available
+
+        if missing_speech_types:
+            # Disable the generate button
+            return gr.update(interactive=False)
+        else:
+            # Enable the generate button
+            return gr.update(interactive=True)
+
+    gen_text_input_emotional.change(
+        validate_speech_types,
+        inputs=[gen_text_input_emotional, regular_name] + speech_type_names,
+        outputs=generate_emotional_btn
+    )
+with gr.Blocks() as app:
+    gr.Markdown(
+        """
+# Antriksh AI
+"""
+    )
+    
+    # Add the image here
+    gr.Image(
+        value="C:\\Users\\USER\\Downloads\\logo-removebg-preview.png",
+        label="AI System Logo",
+        show_label=False,
+        width=300,
+        height=150
+    )
+    
+    gr.TabbedInterface([app_tts, app_podcast, app_emotional, app_credits], ["TTS", "Podcast", "Multi-Style", "Credits"])
+    
+    
+@click.command()
+@click.option("--port", "-p", default=None, type=int, help="Port to run the app on")
+@click.option("--host", "-H", default=None, help="Host to run the app on")
+@click.option(
+    "--share",
+    "-s",
+    default=False,
+    is_flag=True,
+    help="Share the app via Gradio share link",
+)
+@click.option("--api", "-a", default=True, is_flag=True, help="Allow API access")
+def main(port, host, share, api):
+    global app
+    print(f"Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host, server_port=port, share=share, show_api=api
+    )
+
+
+if __name__ == "__main__":
+    if not USING_SPACES:
+        main()
+    else:
+        app.queue().launch()

model/__init__.py

@@ -0,0 +1,7 @@
+from model.cfm import CFM
+
+from model.backbones.unett import UNetT
+from model.backbones.dit import DiT
+from model.backbones.mmdit import MMDiT
+
+from model.trainer import Trainer

model/backbones/README.md

@@ -0,0 +1,20 @@
+## Backbones quick introduction
+
+
+### unett.py
+- flat unet transformer
+- structure same as in e2-tts & voicebox paper except using rotary pos emb
+- update: allow possible abs pos emb & convnextv2 blocks for embedded text before concat
+
+### dit.py
+- adaln-zero dit
+- embedded timestep as condition
+- concatted noised_input + masked_cond + embedded_text, linear proj in
+- possible abs pos emb & convnextv2 blocks for embedded text before concat
+- possible long skip connection (first layer to last layer)
+
+### mmdit.py
+- sd3 structure
+- timestep as condition
+- left stream: text embedded and applied a abs pos emb
+- right stream: masked_cond & noised_input concatted and with same conv pos emb as unett

model/backbones/dit.py

@@ -0,0 +1,158 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import repeat
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from model.modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    DiTBlock,
+    AdaLayerNormZero_Final,
+    precompute_freqs_cis, get_pos_embed_indices,
+)
+
+
+# Text embedding
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, conv_layers = 0, conv_mult = 2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)])
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int['b nt'], seq_len, drop_text = False):
+        batch, text_len = text.shape[0], text.shape[1]
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
+        text = F.pad(text, (0, seq_len - text_len), value = 0)
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+
+        text = self.text_embed(text) # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim = out_dim)
+
+    def forward(self, x: float['b n d'], cond: float['b n d'], text_embed: float['b n d'], drop_audio_cond = False):
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+
+        x = self.proj(torch.cat((x, cond, text_embed), dim = -1))
+        x = self.conv_pos_embed(x) + x
+        return x
+    
+
+# Transformer backbone using DiT blocks
+
+class DiT(nn.Module):
+    def __init__(self, *, 
+                 dim, depth = 8, heads = 8, dim_head = 64, dropout = 0.1, ff_mult = 4,
+                 mel_dim = 100, text_num_embeds = 256, text_dim = None, conv_layers = 0,
+                 long_skip_connection = False,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        if text_dim is None:
+            text_dim = mel_dim
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers = conv_layers)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+        
+        self.transformer_blocks = nn.ModuleList(
+            [
+                DiTBlock(
+                    dim = dim,
+                    heads = heads,
+                    dim_head = dim_head,
+                    ff_mult = ff_mult,
+                    dropout = dropout
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias = False) if long_skip_connection else None
+        
+        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward(
+        self,
+        x: float['b n d'],     # nosied input audio
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = batch)
+        
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        text_embed = self.text_embed(text, seq_len, drop_text = drop_text)
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond = drop_audio_cond)
+        
+        rope = self.rotary_embed.forward_from_seq_len(seq_len)
+
+        if self.long_skip_connection is not None:
+            residual = x
+
+        for block in self.transformer_blocks:
+            x = block(x, t, mask = mask, rope = rope)
+
+        if self.long_skip_connection is not None:
+            x = self.long_skip_connection(torch.cat((x, residual), dim = -1))
+
+        x = self.norm_out(x, t)
+        output = self.proj_out(x)
+
+        return output

model/backbones/mmdit.py

@@ -0,0 +1,136 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from einops import repeat
+
+from x_transformers.x_transformers import RotaryEmbedding
+
+from model.modules import (
+    TimestepEmbedding,
+    ConvPositionEmbedding,
+    MMDiTBlock,
+    AdaLayerNormZero_Final,
+    precompute_freqs_cis, get_pos_embed_indices,
+)
+
+
+# text embedding
+
+class TextEmbedding(nn.Module):
+    def __init__(self, out_dim, text_num_embeds):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim)  # will use 0 as filler token
+
+        self.precompute_max_pos = 1024
+        self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)
+
+    def forward(self, text: int['b nt'], drop_text = False) -> int['b nt d']:
+        text = text + 1
+        if drop_text:
+            text = torch.zeros_like(text)
+        text = self.text_embed(text)
+
+        # sinus pos emb
+        batch_start = torch.zeros((text.shape[0],), dtype=torch.long)
+        batch_text_len = text.shape[1]
+        pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)
+        text_pos_embed = self.freqs_cis[pos_idx]
+
+        text = text + text_pos_embed
+
+        return text
+
+
+# noised input & masked cond audio embedding
+
+class AudioEmbedding(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.linear = nn.Linear(2 * in_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(out_dim)
+
+    def forward(self, x: float['b n d'], cond: float['b n d'], drop_audio_cond = False):
+        if drop_audio_cond:
+            cond = torch.zeros_like(cond)
+        x = torch.cat((x, cond), dim = -1)
+        x = self.linear(x)
+        x = self.conv_pos_embed(x) + x
+        return x
+    
+
+# Transformer backbone using MM-DiT blocks
+
+class MMDiT(nn.Module):
+    def __init__(self, *, 
+                 dim, depth = 8, heads = 8, dim_head = 64, dropout = 0.1, ff_mult = 4,
+                 text_num_embeds = 256, mel_dim = 100,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        self.text_embed = TextEmbedding(dim, text_num_embeds)
+        self.audio_embed = AudioEmbedding(mel_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        self.dim = dim
+        self.depth = depth
+        
+        self.transformer_blocks = nn.ModuleList(
+            [
+                MMDiTBlock(
+                    dim = dim,
+                    heads = heads,
+                    dim_head = dim_head,
+                    dropout = dropout,
+                    ff_mult = ff_mult,
+                    context_pre_only = i == depth - 1,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm_out = AdaLayerNormZero_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward(
+        self,
+        x: float['b n d'],     # nosied input audio
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch = x.shape[0]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = batch)
+
+        # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        c = self.text_embed(text, drop_text = drop_text)
+        x = self.audio_embed(x, cond, drop_audio_cond = drop_audio_cond)
+
+        seq_len = x.shape[1]
+        text_len = text.shape[1]
+        rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)
+        rope_text = self.rotary_embed.forward_from_seq_len(text_len)
+        
+        for block in self.transformer_blocks:
+            c, x = block(x, c, t, mask = mask, rope = rope_audio, c_rope = rope_text)
+
+        x = self.norm_out(x, t)
+        output = self.proj_out(x)
+
+        return output

model/backbones/unett.py

@@ -0,0 +1,201 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Literal
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from einops import repeat, pack, unpack
+
+from x_transformers import RMSNorm
+from x_transformers.x_transformers import RotaryEmbedding
+
+from model.modules import (
+    TimestepEmbedding,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    Attention,
+    AttnProcessor,
+    FeedForward,
+    precompute_freqs_cis, get_pos_embed_indices,
+)
+
+
+# Text embedding
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, conv_layers = 0, conv_mult = 2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        if conv_layers > 0:
+            self.extra_modeling = True
+            self.precompute_max_pos = 4096  # ~44s of 24khz audio
+            self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
+            self.text_blocks = nn.Sequential(*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)])
+        else:
+            self.extra_modeling = False
+
+    def forward(self, text: int['b nt'], seq_len, drop_text = False):
+        batch, text_len = text.shape[0], text.shape[1]
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        text = text[:, :seq_len]  # curtail if character tokens are more than the mel spec tokens
+        text = F.pad(text, (0, seq_len - text_len), value = 0)
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+
+        text = self.text_embed(text) # b n -> b n d
+
+        # possible extra modeling
+        if self.extra_modeling:
+            # sinus pos emb
+            batch_start = torch.zeros((batch,), dtype=torch.long)
+            pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
+            text_pos_embed = self.freqs_cis[pos_idx]
+            text = text + text_pos_embed
+
+            # convnextv2 blocks
+            text = self.text_blocks(text)
+
+        return text
+
+
+# noised input audio and context mixing embedding
+
+class InputEmbedding(nn.Module):
+    def __init__(self, mel_dim, text_dim, out_dim):
+        super().__init__()
+        self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
+        self.conv_pos_embed = ConvPositionEmbedding(dim = out_dim)
+
+    def forward(self, x: float['b n d'], cond: float['b n d'], text_embed: float['b n d'], drop_audio_cond = False):
+        if drop_audio_cond:  # cfg for cond audio
+            cond = torch.zeros_like(cond)
+
+        x = self.proj(torch.cat((x, cond, text_embed), dim = -1))
+        x = self.conv_pos_embed(x) + x
+        return x
+
+
+# Flat UNet Transformer backbone
+
+class UNetT(nn.Module):
+    def __init__(self, *,
+                 dim, depth = 8, heads = 8, dim_head = 64, dropout = 0.1, ff_mult = 4,
+                 mel_dim = 100, text_num_embeds = 256, text_dim = None, conv_layers = 0,
+                 skip_connect_type: Literal['add', 'concat', 'none'] = 'concat',
+    ):
+        super().__init__()
+        assert depth % 2 == 0, "UNet-Transformer's depth should be even."
+
+        self.time_embed = TimestepEmbedding(dim)
+        if text_dim is None:
+            text_dim = mel_dim
+        self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers = conv_layers)
+        self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
+
+        self.rotary_embed = RotaryEmbedding(dim_head)
+
+        # transformer layers & skip connections
+
+        self.dim = dim
+        self.skip_connect_type = skip_connect_type
+        needs_skip_proj = skip_connect_type == 'concat'
+
+        self.depth = depth
+        self.layers = nn.ModuleList([])
+
+        for idx in range(depth):
+            is_later_half = idx >= (depth // 2)
+
+            attn_norm = RMSNorm(dim)
+            attn = Attention(
+                processor = AttnProcessor(),
+                dim = dim,
+                heads = heads,
+                dim_head = dim_head,
+                dropout = dropout,
+                )
+
+            ff_norm = RMSNorm(dim)
+            ff = FeedForward(dim = dim, mult = ff_mult, dropout = dropout, approximate = "tanh")
+
+            skip_proj = nn.Linear(dim * 2, dim, bias = False) if needs_skip_proj and is_later_half else None
+
+            self.layers.append(nn.ModuleList([
+                skip_proj,
+                attn_norm,
+                attn,
+                ff_norm,
+                ff,
+            ]))
+
+        self.norm_out = RMSNorm(dim)
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+    def forward(
+        self,
+        x: float['b n d'],     # nosied input audio
+        cond: float['b n d'],  # masked cond audio
+        text: int['b nt'],     # text
+        time: float['b'] | float[''],  # time step
+        drop_audio_cond,  # cfg for cond audio
+        drop_text,        # cfg for text
+        mask: bool['b n'] | None = None,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = repeat(time, ' -> b', b = batch)
+        
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        text_embed = self.text_embed(text, seq_len, drop_text = drop_text)
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond = drop_audio_cond)
+
+        # postfix time t to input x, [b n d] -> [b n+1 d]
+        x, ps = pack((t, x), 'b * d')
+        if mask is not None:
+            mask = F.pad(mask, (1, 0), value=1)
+        
+        rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)
+
+        # flat unet transformer
+        skip_connect_type = self.skip_connect_type
+        skips = []
+        for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):
+            layer = idx + 1
+
+            # skip connection logic
+            is_first_half = layer <= (self.depth // 2)
+            is_later_half = not is_first_half
+
+            if is_first_half:
+                skips.append(x)
+
+            if is_later_half:
+                skip = skips.pop()
+                if skip_connect_type == 'concat':
+                    x = torch.cat((x, skip), dim = -1)
+                    x = maybe_skip_proj(x)
+                elif skip_connect_type == 'add':
+                    x = x + skip
+
+            # attention and feedforward blocks
+            x = attn(attn_norm(x), rope = rope, mask = mask) + x
+            x = ff(ff_norm(x)) + x
+
+        assert len(skips) == 0
+
+        _, x = unpack(self.norm_out(x), ps, 'b * d')
+
+        return self.proj_out(x)

model/cfm.py

@@ -0,0 +1,279 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Callable
+from random import random
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from torchdiffeq import odeint
+
+from einops import rearrange
+
+from model.modules import MelSpec
+
+from model.utils import (
+    default, exists, 
+    list_str_to_idx, list_str_to_tensor, 
+    lens_to_mask, mask_from_frac_lengths,
+) 
+
+
+class CFM(nn.Module):
+    def __init__(
+        self,
+        transformer: nn.Module,
+        sigma = 0.,
+        odeint_kwargs: dict = dict(
+            # atol = 1e-5,
+            # rtol = 1e-5,
+            method = 'euler'  # 'midpoint'
+        ),
+        audio_drop_prob = 0.3,
+        cond_drop_prob = 0.2,
+        num_channels = None,
+        mel_spec_module: nn.Module | None = None,
+        mel_spec_kwargs: dict = dict(),
+        frac_lengths_mask: tuple[float, float] = (0.7, 1.),
+        vocab_char_map: dict[str: int] | None = None
+    ):
+        super().__init__()
+
+        self.frac_lengths_mask = frac_lengths_mask
+
+        # mel spec
+        self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs))
+        num_channels = default(num_channels, self.mel_spec.n_mel_channels)
+        self.num_channels = num_channels
+
+        # classifier-free guidance
+        self.audio_drop_prob = audio_drop_prob
+        self.cond_drop_prob = cond_drop_prob
+
+        # transformer
+        self.transformer = transformer
+        dim = transformer.dim
+        self.dim = dim
+
+        # conditional flow related
+        self.sigma = sigma
+
+        # sampling related
+        self.odeint_kwargs = odeint_kwargs
+
+        # vocab map for tokenization
+        self.vocab_char_map = vocab_char_map
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond: float['b n d'] | float['b nw'],
+        text: int['b nt'] | list[str],
+        duration: int | int['b'],
+        *,
+        lens: int['b'] | None = None,
+        steps = 32,
+        cfg_strength = 1., 
+        sway_sampling_coef = None,
+        seed: int | None = None,
+        max_duration = 4096, 
+        vocoder: Callable[[float['b d n']], float['b nw']] | None = None,
+        no_ref_audio = False,
+        duplicate_test = False,
+        t_inter = 0.1,
+        edit_mask = None,
+    ):
+        self.eval()
+
+        # raw wave
+
+        if cond.ndim == 2:
+            cond = self.mel_spec(cond)
+            cond = rearrange(cond, 'b d n -> b n d')
+            assert cond.shape[-1] == self.num_channels
+
+        batch, cond_seq_len, device = *cond.shape[:2], cond.device
+        if not exists(lens):
+            lens = torch.full((batch,), cond_seq_len, device = device, dtype = torch.long)
+
+        # text
+
+        if isinstance(text, list):
+            if exists(self.vocab_char_map):
+                text = list_str_to_idx(text, self.vocab_char_map).to(device)
+            else:
+                text = list_str_to_tensor(text).to(device)
+            assert text.shape[0] == batch
+
+        if exists(text):
+            text_lens = (text != -1).sum(dim = -1)
+            lens = torch.maximum(text_lens, lens) # make sure lengths are at least those of the text characters
+
+        # duration
+
+        cond_mask = lens_to_mask(lens)
+        if edit_mask is not None:
+            cond_mask = cond_mask & edit_mask
+
+        if isinstance(duration, int):
+            duration = torch.full((batch,), duration, device = device, dtype = torch.long)
+
+        duration = torch.maximum(lens + 1, duration) # just add one token so something is generated
+        duration = duration.clamp(max = max_duration)
+        max_duration = duration.amax()
+        
+        # duplicate test corner for inner time step oberservation
+        if duplicate_test:
+            test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2*cond_seq_len), value = 0.)
+            
+        cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value = 0.)
+        cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value = False)
+        cond_mask = rearrange(cond_mask, '... -> ... 1')
+        step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))  # allow direct control (cut cond audio) with lens passed in
+
+        if batch > 1:
+            mask = lens_to_mask(duration)
+        else:  # save memory and speed up, as single inference need no mask currently
+            mask = None
+
+        # test for no ref audio
+        if no_ref_audio:
+            cond = torch.zeros_like(cond)
+
+        # neural ode
+
+        def fn(t, x):
+            # at each step, conditioning is fixed
+            # step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))
+
+            # predict flow
+            pred = self.transformer(x = x, cond = step_cond, text = text, time = t, mask = mask, drop_audio_cond = False, drop_text = False)
+            if cfg_strength < 1e-5:
+                return pred
+            
+            null_pred = self.transformer(x = x, cond = step_cond, text = text, time = t, mask = mask, drop_audio_cond = True, drop_text = True)
+            return pred + (pred - null_pred) * cfg_strength
+
+        # noise input
+        # to make sure batch inference result is same with different batch size, and for sure single inference
+        # still some difference maybe due to convolutional layers
+        y0 = []
+        for dur in duration:
+            if exists(seed):
+                torch.manual_seed(seed)
+            y0.append(torch.randn(dur, self.num_channels, device = self.device))
+        y0 = pad_sequence(y0, padding_value = 0, batch_first = True)
+
+        t_start = 0
+
+        # duplicate test corner for inner time step oberservation
+        if duplicate_test:
+            t_start = t_inter
+            y0 = (1 - t_start) * y0 + t_start * test_cond
+            steps = int(steps * (1 - t_start))
+
+        t = torch.linspace(t_start, 1, steps, device = self.device)
+        if sway_sampling_coef is not None:
+            t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)
+
+        trajectory = odeint(fn, y0, t, **self.odeint_kwargs)
+        
+        sampled = trajectory[-1]
+        out = sampled
+        out = torch.where(cond_mask, cond, out)
+
+        if exists(vocoder):
+            out = rearrange(out, 'b n d -> b d n')
+            out = vocoder(out)
+
+        return out, trajectory
+
+    def forward(
+        self,
+        inp: float['b n d'] | float['b nw'], # mel or raw wave
+        text: int['b nt'] | list[str],
+        *,
+        lens: int['b'] | None = None,
+        noise_scheduler: str | None = None,
+    ):
+        # handle raw wave
+        if inp.ndim == 2:
+            inp = self.mel_spec(inp)
+            inp = rearrange(inp, 'b d n -> b n d')
+            assert inp.shape[-1] == self.num_channels
+
+        batch, seq_len, dtype, device, σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma
+
+        # handle text as string
+        if isinstance(text, list):
+            if exists(self.vocab_char_map):
+                text = list_str_to_idx(text, self.vocab_char_map).to(device)
+            else:
+                text = list_str_to_tensor(text).to(device)
+            assert text.shape[0] == batch
+
+        # lens and mask
+        if not exists(lens):
+            lens = torch.full((batch,), seq_len, device = device)
+        
+        mask = lens_to_mask(lens, length = seq_len)  # useless here, as collate_fn will pad to max length in batch
+
+        # get a random span to mask out for training conditionally
+        frac_lengths = torch.zeros((batch,), device = self.device).float().uniform_(*self.frac_lengths_mask)
+        rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)
+
+        if exists(mask):
+            rand_span_mask &= mask
+
+        # mel is x1
+        x1 = inp
+
+        # x0 is gaussian noise
+        x0 = torch.randn_like(x1)
+
+        # time step
+        time = torch.rand((batch,), dtype = dtype, device = self.device)
+        # TODO. noise_scheduler
+
+        # sample xt (φ_t(x) in the paper)
+        t = rearrange(time, 'b -> b 1 1')
+        φ = (1 - t) * x0 + t * x1
+        flow = x1 - x0
+
+        # only predict what is within the random mask span for infilling
+        cond = torch.where(
+            rand_span_mask[..., None],
+            torch.zeros_like(x1), x1
+        )
+
+        # transformer and cfg training with a drop rate
+        drop_audio_cond = random() < self.audio_drop_prob  # p_drop in voicebox paper
+        if random() < self.cond_drop_prob:  # p_uncond in voicebox paper
+            drop_audio_cond = True
+            drop_text = True
+        else:
+            drop_text = False
+            
+        # if want rigourously mask out padding, record in collate_fn in dataset.py, and pass in here
+        # adding mask will use more memory, thus also need to adjust batchsampler with scaled down threshold for long sequences
+        pred = self.transformer(x = φ, cond = cond, text = text, time = time, drop_audio_cond = drop_audio_cond, drop_text = drop_text)
+
+        # flow matching loss
+        loss = F.mse_loss(pred, flow, reduction = 'none')
+        loss = loss[rand_span_mask]
+
+        return loss.mean(), cond, pred

model/dataset.py

@@ -0,0 +1,241 @@
+import json
+import random
+from tqdm import tqdm
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, Sampler
+import torchaudio
+from datasets import load_dataset, load_from_disk
+from datasets import Dataset as Dataset_
+from typing import Union
+
+from einops import rearrange
+
+from model.modules import MelSpec
+
+
+class HFDataset(Dataset):
+    def __init__(
+        self,
+        hf_dataset: Dataset,
+        target_sample_rate = 24_000,
+        n_mel_channels = 100,
+        hop_length = 256,
+    ):
+        self.data = hf_dataset
+        self.target_sample_rate = target_sample_rate
+        self.hop_length = hop_length
+        self.mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length)
+        
+    def get_frame_len(self, index):
+        row = self.data[index]
+        audio = row['audio']['array']
+        sample_rate = row['audio']['sampling_rate']
+        return audio.shape[-1] / sample_rate * self.target_sample_rate / self.hop_length
+
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        row = self.data[index]
+        audio = row['audio']['array']
+
+        # logger.info(f"Audio shape: {audio.shape}")
+
+        sample_rate = row['audio']['sampling_rate']
+        duration = audio.shape[-1] / sample_rate
+
+        if duration > 30 or duration < 0.3:
+            return self.__getitem__((index + 1) % len(self.data))
+        
+        audio_tensor = torch.from_numpy(audio).float()
+        
+        if sample_rate != self.target_sample_rate:
+            resampler = torchaudio.transforms.Resample(sample_rate, self.target_sample_rate)
+            audio_tensor = resampler(audio_tensor)
+        
+        audio_tensor = rearrange(audio_tensor, 't -> 1 t')
+        
+        mel_spec = self.mel_spectrogram(audio_tensor)
+        
+        mel_spec = rearrange(mel_spec, '1 d t -> d t')
+        
+        text = row['text']
+        
+        return dict(
+            mel_spec = mel_spec,
+            text = text,
+        )
+
+class CustomDataset(Dataset):
+    def __init__(
+        self,
+        custom_dataset: Dataset,
+        durations = None,
+        target_sample_rate = 24_000,
+        hop_length = 256,
+        n_mel_channels = 100,
+        preprocessed_mel = False,
+    ):
+        self.data = custom_dataset
+        self.durations = durations
+        self.target_sample_rate = target_sample_rate
+        self.hop_length = hop_length
+        self.preprocessed_mel = preprocessed_mel
+        if not preprocessed_mel:
+            self.mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, hop_length=hop_length, n_mel_channels=n_mel_channels)
+
+    def get_frame_len(self, index):
+        if self.durations is not None:  # Please make sure the separately provided durations are correct, otherwise 99.99% OOM
+            return self.durations[index] * self.target_sample_rate / self.hop_length
+        return self.data[index]["duration"] * self.target_sample_rate / self.hop_length
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, index):
+        row = self.data[index]
+        audio_path = row["audio_path"]
+        text = row["text"]
+        duration = row["duration"]
+
+        if self.preprocessed_mel:
+            mel_spec = torch.tensor(row["mel_spec"])
+
+        else:
+            audio, source_sample_rate = torchaudio.load(audio_path)
+
+            if duration > 30 or duration < 0.3:
+                return self.__getitem__((index + 1) % len(self.data))
+            
+            if source_sample_rate != self.target_sample_rate:
+                resampler = torchaudio.transforms.Resample(source_sample_rate, self.target_sample_rate)
+                audio = resampler(audio)
+            
+            mel_spec = self.mel_spectrogram(audio)
+            mel_spec = rearrange(mel_spec, '1 d t -> d t')
+        
+        return dict(
+            mel_spec = mel_spec,
+            text = text,
+        )
+    
+
+# Dynamic Batch Sampler
+
+class DynamicBatchSampler(Sampler[list[int]]):
+    """ Extension of Sampler that will do the following:
+        1.  Change the batch size (essentially number of sequences)
+            in a batch to ensure that the total number of frames are less
+            than a certain threshold.
+        2.  Make sure the padding efficiency in the batch is high.
+    """
+
+    def __init__(self, sampler: Sampler[int], frames_threshold: int, max_samples=0, random_seed=None, drop_last: bool = False):
+        self.sampler = sampler
+        self.frames_threshold = frames_threshold
+        self.max_samples = max_samples
+
+        indices, batches = [], []
+        data_source = self.sampler.data_source
+        
+        for idx in tqdm(self.sampler, desc=f"Sorting with sampler... if slow, check whether dataset is provided with duration"):
+            indices.append((idx, data_source.get_frame_len(idx)))
+        indices.sort(key=lambda elem : elem[1])
+
+        batch = []
+        batch_frames = 0
+        for idx, frame_len in tqdm(indices, desc=f"Creating dynamic batches with {frames_threshold} audio frames per gpu"):
+            if batch_frames + frame_len <= self.frames_threshold and (max_samples == 0 or len(batch) < max_samples):
+                batch.append(idx)
+                batch_frames += frame_len
+            else:
+                if len(batch) > 0:
+                    batches.append(batch)
+                if frame_len <= self.frames_threshold:
+                    batch = [idx]
+                    batch_frames = frame_len
+                else:
+                    batch = []
+                    batch_frames = 0
+
+        if not drop_last and len(batch) > 0:
+            batches.append(batch)
+
+        del indices
+
+        # if want to have different batches between epochs, may just set a seed and log it in ckpt
+        # cuz during multi-gpu training, although the batch on per gpu not change between epochs, the formed general minibatch is different
+        # e.g. for epoch n, use (random_seed + n)
+        random.seed(random_seed)
+        random.shuffle(batches)
+
+        self.batches = batches
+
+    def __iter__(self):
+        return iter(self.batches)
+
+    def __len__(self):
+        return len(self.batches)
+
+
+# Load dataset
+def load_dataset(
+        dataset_name: str,
+        tokenizer: str,
+        dataset_type: str = "CustomDataset", 
+        audio_type: str = "raw", 
+        mel_spec_kwargs: dict = dict()
+        ) -> Union[CustomDataset, HFDataset]:
+    
+    print("Loading dataset ...")
+
+    if dataset_type == "CustomDataset":
+        if audio_type == "raw":
+            try:
+                train_dataset = load_from_disk(f"data/{dataset_name}_{tokenizer}/raw")
+            except:
+                train_dataset = Dataset_.from_file(f"data/{dataset_name}_{tokenizer}/raw.arrow")
+            preprocessed_mel = False
+        elif audio_type == "mel":
+            train_dataset = Dataset_.from_file(f"data/{dataset_name}_{tokenizer}/mel.arrow")
+            preprocessed_mel = True
+        with open(f"data/{dataset_name}_{tokenizer}/duration.json", 'r', encoding='utf-8') as f:
+            data_dict = json.load(f)
+        durations = data_dict["duration"]
+        train_dataset = CustomDataset(train_dataset, durations=durations, preprocessed_mel=preprocessed_mel, **mel_spec_kwargs)
+
+    elif dataset_type == "HFDataset":
+        print("Should manually modify the path of huggingface dataset to your need.\n" +
+              "May also the corresponding script cuz different dataset may have different format.")
+        pre, post = dataset_name.split("_")
+        train_dataset = HFDataset(load_dataset(f"{pre}/{pre}", split=f"train.{post}", cache_dir="./data"),)
+
+    return train_dataset
+
+
+# collation
+
+def collate_fn(batch):
+    mel_specs = [item['mel_spec'].squeeze(0) for item in batch]
+    mel_lengths = torch.LongTensor([spec.shape[-1] for spec in mel_specs])
+    max_mel_length = mel_lengths.amax()
+
+    padded_mel_specs = []
+    for spec in mel_specs:  # TODO. maybe records mask for attention here
+        padding = (0, max_mel_length - spec.size(-1))
+        padded_spec = F.pad(spec, padding, value = 0)
+        padded_mel_specs.append(padded_spec)
+    
+    mel_specs = torch.stack(padded_mel_specs)
+
+    text = [item['text'] for item in batch]
+    text_lengths = torch.LongTensor([len(item) for item in text])
+
+    return dict(
+        mel = mel_specs,
+        mel_lengths = mel_lengths,
+        text = text,
+        text_lengths = text_lengths,
+    )

model/ecapa_tdnn.py

@@ -0,0 +1,268 @@
+# just for speaker similarity evaluation, third-party code
+
+# From https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/
+# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+''' Res2Conv1d + BatchNorm1d + ReLU
+'''
+
+class Res2Conv1dReluBn(nn.Module):
+    '''
+    in_channels == out_channels == channels
+    '''
+
+    def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):
+        super().__init__()
+        assert channels % scale == 0, "{} % {} != 0".format(channels, scale)
+        self.scale = scale
+        self.width = channels // scale
+        self.nums = scale if scale == 1 else scale - 1
+
+        self.convs = []
+        self.bns = []
+        for i in range(self.nums):
+            self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))
+            self.bns.append(nn.BatchNorm1d(self.width))
+        self.convs = nn.ModuleList(self.convs)
+        self.bns = nn.ModuleList(self.bns)
+
+    def forward(self, x):
+        out = []
+        spx = torch.split(x, self.width, 1)
+        for i in range(self.nums):
+            if i == 0:
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            # Order: conv -> relu -> bn
+            sp = self.convs[i](sp)
+            sp = self.bns[i](F.relu(sp))
+            out.append(sp)
+        if self.scale != 1:
+            out.append(spx[self.nums])
+        out = torch.cat(out, dim=1)
+
+        return out
+
+
+''' Conv1d + BatchNorm1d + ReLU
+'''
+
+class Conv1dReluBn(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return self.bn(F.relu(self.conv(x)))
+
+
+''' The SE connection of 1D case.
+'''
+
+class SE_Connect(nn.Module):
+    def __init__(self, channels, se_bottleneck_dim=128):
+        super().__init__()
+        self.linear1 = nn.Linear(channels, se_bottleneck_dim)
+        self.linear2 = nn.Linear(se_bottleneck_dim, channels)
+
+    def forward(self, x):
+        out = x.mean(dim=2)
+        out = F.relu(self.linear1(out))
+        out = torch.sigmoid(self.linear2(out))
+        out = x * out.unsqueeze(2)
+
+        return out
+
+
+''' SE-Res2Block of the ECAPA-TDNN architecture.
+'''
+
+# def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):
+#     return nn.Sequential(
+#         Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0),
+#         Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale),
+#         Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0),
+#         SE_Connect(channels)
+#     )
+
+class SE_Res2Block(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):
+        super().__init__()
+        self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)
+        self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)
+
+        self.shortcut = None
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv1d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=1,
+            )
+
+    def forward(self, x):
+        residual = x
+        if self.shortcut:
+            residual = self.shortcut(x)
+
+        x = self.Conv1dReluBn1(x)
+        x = self.Res2Conv1dReluBn(x)
+        x = self.Conv1dReluBn2(x)
+        x = self.SE_Connect(x)
+
+        return x + residual
+
+
+''' Attentive weighted mean and standard deviation pooling.
+'''
+
+class AttentiveStatsPool(nn.Module):
+    def __init__(self, in_dim, attention_channels=128, global_context_att=False):
+        super().__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1)  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1)  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1)  # equals V and k in the paper
+
+    def forward(self, x):
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! In experiments, I find ReLU hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))
+        # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        residuals = torch.sum(alpha * (x ** 2), dim=2) - mean ** 2
+        std = torch.sqrt(residuals.clamp(min=1e-9))
+        return torch.cat([mean, std], dim=1)
+
+
+class ECAPA_TDNN(nn.Module):
+    def __init__(self, feat_dim=80, channels=512, emb_dim=192, global_context_att=False,
+                 feat_type='wavlm_large', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None):
+        super().__init__()
+
+        self.feat_type = feat_type
+        self.feature_selection = feature_selection
+        self.update_extract = update_extract
+        self.sr = sr
+        
+        torch.hub._validate_not_a_forked_repo=lambda a,b,c: True
+        try:
+            local_s3prl_path = os.path.expanduser("~/.cache/torch/hub/s3prl_s3prl_main")
+            self.feature_extract = torch.hub.load(local_s3prl_path, feat_type, source='local', config_path=config_path)
+        except:
+            self.feature_extract = torch.hub.load('s3prl/s3prl', feat_type)
+
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[23].self_attn, "fp32_attention"):
+            self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False
+        if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(self.feature_extract.model.encoder.layers[11].self_attn, "fp32_attention"):
+            self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False
+
+        self.feat_num = self.get_feat_num()
+        self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))
+
+        if feat_type != 'fbank' and feat_type != 'mfcc':
+            freeze_list = ['final_proj', 'label_embs_concat', 'mask_emb', 'project_q', 'quantizer']
+            for name, param in self.feature_extract.named_parameters():
+                for freeze_val in freeze_list:
+                    if freeze_val in name:
+                        param.requires_grad = False
+                        break
+
+        if not self.update_extract:
+            for param in self.feature_extract.parameters():
+                param.requires_grad = False
+
+        self.instance_norm = nn.InstanceNorm1d(feat_dim)
+        # self.channels = [channels] * 4 + [channels * 3]
+        self.channels = [channels] * 4 + [1536]
+
+        self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)
+        self.layer2 = SE_Res2Block(self.channels[0], self.channels[1], kernel_size=3, stride=1, padding=2, dilation=2, scale=8, se_bottleneck_dim=128)
+        self.layer3 = SE_Res2Block(self.channels[1], self.channels[2], kernel_size=3, stride=1, padding=3, dilation=3, scale=8, se_bottleneck_dim=128)
+        self.layer4 = SE_Res2Block(self.channels[2], self.channels[3], kernel_size=3, stride=1, padding=4, dilation=4, scale=8, se_bottleneck_dim=128)
+
+        # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)
+        cat_channels = channels * 3
+        self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)
+        self.pooling = AttentiveStatsPool(self.channels[-1], attention_channels=128, global_context_att=global_context_att)
+        self.bn = nn.BatchNorm1d(self.channels[-1] * 2)
+        self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)
+
+
+    def get_feat_num(self):
+        self.feature_extract.eval()
+        wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)]
+        with torch.no_grad():
+            features = self.feature_extract(wav)
+        select_feature = features[self.feature_selection]
+        if isinstance(select_feature, (list, tuple)):
+            return len(select_feature)
+        else:
+            return 1
+
+    def get_feat(self, x):
+        if self.update_extract:
+            x = self.feature_extract([sample for sample in x])
+        else:
+            with torch.no_grad():
+                if self.feat_type == 'fbank' or self.feat_type == 'mfcc':
+                    x = self.feature_extract(x) + 1e-6  # B x feat_dim x time_len
+                else:
+                    x = self.feature_extract([sample for sample in x])
+
+        if self.feat_type == 'fbank':
+            x = x.log()
+
+        if self.feat_type != "fbank" and self.feat_type != "mfcc":
+            x = x[self.feature_selection]
+            if isinstance(x, (list, tuple)):
+                x = torch.stack(x, dim=0)
+            else:
+                x = x.unsqueeze(0)
+            norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            x = (norm_weights * x).sum(dim=0)
+            x = torch.transpose(x, 1, 2) + 1e-6
+
+        x = self.instance_norm(x)
+        return x
+
+    def forward(self, x):
+        x = self.get_feat(x)
+
+        out1 = self.layer1(x)
+        out2 = self.layer2(out1)
+        out3 = self.layer3(out2)
+        out4 = self.layer4(out3)
+
+        out = torch.cat([out2, out3, out4], dim=1)
+        out = F.relu(self.conv(out))
+        out = self.bn(self.pooling(out))
+        out = self.linear(out)
+
+        return out
+
+
+def ECAPA_TDNN_SMALL(feat_dim, emb_dim=256, feat_type='wavlm_large', sr=16000, feature_selection="hidden_states", update_extract=False, config_path=None):
+    return ECAPA_TDNN(feat_dim=feat_dim, channels=512, emb_dim=emb_dim,
+                      feat_type=feat_type, sr=sr, feature_selection=feature_selection, update_extract=update_extract, config_path=config_path)

model/modules.py

@@ -0,0 +1,575 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+from typing import Optional
+import math
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torchaudio
+
+from einops import rearrange
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+
+# raw wav to mel spec
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        filter_length = 1024,
+        hop_length = 256,
+        win_length = 1024,
+        n_mel_channels = 100,
+        target_sample_rate = 24_000,
+        normalize = False,
+        power = 1,
+        norm = None,
+        center = True,
+    ):
+        super().__init__()
+        self.n_mel_channels = n_mel_channels
+
+        self.mel_stft = torchaudio.transforms.MelSpectrogram(
+            sample_rate = target_sample_rate,
+            n_fft = filter_length,
+            win_length = win_length,
+            hop_length = hop_length,
+            n_mels = n_mel_channels,
+            power = power,
+            center = center,
+            normalized = normalize,
+            norm = norm,
+        )
+
+        self.register_buffer('dummy', torch.tensor(0), persistent = False)
+
+    def forward(self, inp):
+        if len(inp.shape) == 3:
+            inp = rearrange(inp, 'b 1 nw -> b nw')
+
+        assert len(inp.shape) == 2
+
+        if self.dummy.device != inp.device:
+            self.to(inp.device)
+
+        mel = self.mel_stft(inp)
+        mel = mel.clamp(min = 1e-5).log()
+        return mel
+    
+
+# sinusoidal position embedding
+
+class SinusPositionEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x, scale=1000):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
+        emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+# convolutional position embedding
+
+class ConvPositionEmbedding(nn.Module):
+    def __init__(self, dim, kernel_size = 31, groups = 16):
+        super().__init__()
+        assert kernel_size % 2 != 0
+        self.conv1d = nn.Sequential(
+            nn.Conv1d(dim, dim, kernel_size, groups = groups, padding = kernel_size // 2),
+            nn.Mish(),
+            nn.Conv1d(dim, dim, kernel_size, groups = groups, padding = kernel_size // 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: float['b n d'], mask: bool['b n'] | None  = None):
+        if mask is not None:
+            mask = mask[..., None]
+            x = x.masked_fill(~mask, 0.)
+
+        x = rearrange(x, 'b n d -> b d n')
+        x = self.conv1d(x)
+        out = rearrange(x, 'b d n -> b n d')
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.)
+
+        return out
+
+
+# rotary positional embedding related
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.):
+    # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
+    # has some connection to NTK literature
+    # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
+    # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py
+    theta *= theta_rescale_factor ** (dim / (dim - 2))
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cos = torch.cos(freqs)  # real part
+    freqs_sin = torch.sin(freqs)  # imaginary part
+    return torch.cat([freqs_cos, freqs_sin], dim=-1)
+
+def get_pos_embed_indices(start, length, max_pos, scale=1.):
+    # length = length if isinstance(length, int) else length.max()
+    scale = scale * torch.ones_like(start, dtype=torch.float32)  # in case scale is a scalar
+    pos = start.unsqueeze(1) + (
+            torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) *
+            scale.unsqueeze(1)).long()
+    # avoid extra long error.
+    pos = torch.where(pos < max_pos, pos, max_pos - 1)
+    return pos
+
+
+# Global Response Normalization layer (Instance Normalization ?)
+
+class GRN(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=1, keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py
+# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108
+
+class ConvNeXtV2Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        intermediate_dim: int,
+        dilation: int = 1,
+    ):
+        super().__init__()
+        padding = (dilation * (7 - 1)) // 2
+        self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation)  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(intermediate_dim)
+        self.pwconv2 = nn.Linear(intermediate_dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = x.transpose(1, 2)  # b n d -> b d n
+        x = self.dwconv(x)
+        x = x.transpose(1, 2)  # b d n -> b n d
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        return residual + x
+
+
+# AdaLayerNormZero
+# return with modulated x for attn input, and params for later mlp modulation
+
+class AdaLayerNormZero(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb = None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+# AdaLayerNormZero for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+class AdaLayerNormZero_Final(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2)
+
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+# FeedForward
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, dim_out = None, mult = 4, dropout = 0., approximate: str = 'none'):
+        super().__init__()
+        inner_dim = int(dim * mult)
+        dim_out = dim_out if dim_out is not None else dim
+
+        activation = nn.GELU(approximate=approximate)
+        project_in = nn.Sequential(
+            nn.Linear(dim, inner_dim),
+            activation
+        )
+        self.ff = nn.Sequential(
+            project_in,
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim_out)
+        )
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+# Attention with possible joint part
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        processor: JointAttnProcessor | AttnProcessor,
+        dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        context_dim: Optional[int] = None, # if not None -> joint attention
+        context_pre_only = None,
+    ):
+        super().__init__()
+
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+        self.processor = processor
+
+        self.dim = dim
+        self.heads = heads
+        self.inner_dim = dim_head * heads
+        self.dropout = dropout
+
+        self.context_dim = context_dim
+        self.context_pre_only = context_pre_only
+
+        self.to_q = nn.Linear(dim, self.inner_dim)
+        self.to_k = nn.Linear(dim, self.inner_dim)
+        self.to_v = nn.Linear(dim, self.inner_dim)
+
+        if self.context_dim is not None:
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if self.context_pre_only is not None:
+                self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+
+        self.to_out = nn.ModuleList([])
+        self.to_out.append(nn.Linear(self.inner_dim, dim))
+        self.to_out.append(nn.Dropout(dropout))
+
+        if self.context_pre_only is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, dim)
+
+    def forward(
+        self,
+        x: float['b n d'], # noised input x
+        c: float['b n d'] = None,  # context c
+        mask: bool['b n'] | None = None,
+        rope = None,  # rotary position embedding for x
+        c_rope = None,  # rotary position embedding for c
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c = c, mask = mask, rope = rope, c_rope = c_rope)
+        else:
+            return self.processor(self, x, mask = mask, rope = rope)
+
+
+# Attention processor
+
+class AttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float['b n d'], # noised input x
+        mask: bool['b n'] | None = None,
+        rope = None,  # rotary position embedding
+    ) -> torch.FloatTensor:
+
+        batch_size = x.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if xpos_scale is not None else (1., 1.)
+
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        # attention
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = mask
+            attn_mask = rearrange(attn_mask, 'b n -> b 1 1 n')
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+
+        if mask is not None:
+            mask = rearrange(mask, 'b n -> b n 1')
+            x = x.masked_fill(~mask, 0.)
+
+        return x
+    
+
+# Joint Attention processor for MM-DiT
+# modified from diffusers/src/diffusers/models/attention_processor.py
+
+class JointAttnProcessor:
+    def __init__(self):
+        pass
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float['b n d'], # noised input x
+        c: float['b nt d'] = None,  # context c, here text
+        mask: bool['b n'] | None = None,
+        rope = None,  # rotary position embedding for x
+        c_rope = None,  # rotary position embedding for c
+    ) -> torch.FloatTensor:
+        residual = x
+
+        batch_size = c.shape[0]
+
+        # `sample` projections.
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # `context` projections.
+        c_query = attn.to_q_c(c)
+        c_key = attn.to_k_c(c)
+        c_value = attn.to_v_c(c)
+
+        # apply rope for context and noised input independently
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if xpos_scale is not None else (1., 1.)
+            query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+            key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+        if c_rope is not None:
+            freqs, xpos_scale = c_rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale ** -1.) if xpos_scale is not None else (1., 1.)
+            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
+            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
+
+        # attention
+        query = torch.cat([query, c_query], dim=1)
+        key = torch.cat([key, c_key], dim=1)
+        value = torch.cat([value, c_value], dim=1)
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # mask. e.g. inference got a batch with different target durations, mask out the padding
+        if mask is not None:
+            attn_mask = F.pad(mask, (0, c.shape[1]), value = True)  # no mask for c (text)
+            attn_mask = rearrange(attn_mask, 'b n -> b 1 1 n')
+            attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
+        else:
+            attn_mask = None
+
+        x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
+        x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        x = x.to(query.dtype)
+
+        # Split the attention outputs.
+        x, c = (
+            x[:, :residual.shape[1]],
+            x[:, residual.shape[1]:],
+        )
+
+        # linear proj
+        x = attn.to_out[0](x)
+        # dropout
+        x = attn.to_out[1](x)
+        if not attn.context_pre_only:
+            c = attn.to_out_c(c)
+
+        if mask is not None:
+            mask = rearrange(mask, 'b n -> b n 1')
+            x = x.masked_fill(~mask, 0.)
+            # c = c.masked_fill(~mask, 0.)  # no mask for c (text)
+
+        return x, c
+
+
+# DiT Block
+
+class DiTBlock(nn.Module):
+
+    def __init__(self, dim, heads, dim_head, ff_mult = 4, dropout = 0.1):
+        super().__init__()
+        
+        self.attn_norm = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            processor = AttnProcessor(),
+            dim = dim,
+            heads = heads,
+            dim_head = dim_head,
+            dropout = dropout,
+            )
+        
+        self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff = FeedForward(dim = dim, mult = ff_mult, dropout = dropout, approximate = "tanh")
+
+    def forward(self, x, t, mask = None, rope = None): # x: noised input, t: time embedding
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attn(x=norm, mask=mask, rope=rope)
+
+        # process attention output for input x
+        x = x + gate_msa.unsqueeze(1) * attn_output
+        
+        norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.ff(norm)
+        x = x + gate_mlp.unsqueeze(1) * ff_output
+
+        return x
+
+
+# MMDiT Block https://arxiv.org/abs/2403.03206
+
+class MMDiTBlock(nn.Module):
+    r""" 
+    modified from diffusers/src/diffusers/models/attention.py
+
+    notes.
+    _c: context related. text, cond, etc. (left part in sd3 fig2.b)
+    _x: noised input related. (right part)
+    context_pre_only: last layer only do prenorm + modulation cuz no more ffn
+    """
+
+    def __init__(self, dim, heads, dim_head, ff_mult = 4, dropout = 0.1, context_pre_only = False):
+        super().__init__()
+
+        self.context_pre_only = context_pre_only
+        
+        self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)
+        self.attn_norm_x = AdaLayerNormZero(dim)
+        self.attn = Attention(
+            processor = JointAttnProcessor(),
+            dim = dim,
+            heads = heads,
+            dim_head = dim_head,
+            dropout = dropout,
+            context_dim = dim,
+            context_pre_only = context_pre_only,
+            )
+
+        if not context_pre_only:
+            self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+            self.ff_c = FeedForward(dim = dim, mult = ff_mult, dropout = dropout, approximate = "tanh")
+        else:
+            self.ff_norm_c = None
+            self.ff_c = None
+        self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.ff_x = FeedForward(dim = dim, mult = ff_mult, dropout = dropout, approximate = "tanh")
+
+    def forward(self, x, c, t, mask = None, rope = None, c_rope = None): # x: noised input, c: context, t: time embedding
+        # pre-norm & modulation for attention input
+        if self.context_pre_only:
+            norm_c = self.attn_norm_c(c, t)
+        else:
+            norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
+
+        # attention
+        x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
+
+        # process attention output for context c
+        if self.context_pre_only:
+            c = None
+        else: # if not last layer
+            c = c + c_gate_msa.unsqueeze(1) * c_attn_output
+
+            norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
+            c_ff_output = self.ff_c(norm_c)
+            c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
+
+        # process attention output for input x
+        x = x + x_gate_msa.unsqueeze(1) * x_attn_output
+        
+        norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
+        x_ff_output = self.ff_x(norm_x)
+        x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
+
+        return c, x
+
+
+# time step conditioning embedding
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, dim, freq_embed_dim=256):
+        super().__init__()
+        self.time_embed = SinusPositionEmbedding(freq_embed_dim)
+        self.time_mlp = nn.Sequential(
+            nn.Linear(freq_embed_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+
+    def forward(self, timestep: float['b']):
+        time_hidden = self.time_embed(timestep)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

model/trainer.py

@@ -0,0 +1,245 @@
+from __future__ import annotations
+
+import os
+import gc
+from tqdm import tqdm
+import wandb
+
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader, Dataset, SequentialSampler
+from torch.optim.lr_scheduler import LinearLR, SequentialLR
+
+from einops import rearrange
+
+from accelerate import Accelerator
+from accelerate.utils import DistributedDataParallelKwargs
+
+from ema_pytorch import EMA
+
+from model import CFM
+from model.utils import exists, default
+from model.dataset import DynamicBatchSampler, collate_fn
+
+
+# trainer
+
+class Trainer:
+    def __init__(
+        self,
+        model: CFM,
+        epochs,
+        learning_rate,
+        num_warmup_updates = 20000,
+        save_per_updates = 1000, 
+        checkpoint_path = None,
+        batch_size = 32, 
+        batch_size_type: str = "sample",
+        max_samples = 32,
+        grad_accumulation_steps = 1,
+        max_grad_norm = 1.0,
+        noise_scheduler: str | None = None,
+        duration_predictor: torch.nn.Module | None = None,
+        wandb_project = "test_e2-tts",
+        wandb_run_name = "test_run",
+        wandb_resume_id: str = None,
+        last_per_steps = None,
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict()
+    ):
+        
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters = True)
+
+        self.accelerator = Accelerator(
+            log_with = "wandb",
+            kwargs_handlers = [ddp_kwargs],
+            gradient_accumulation_steps = grad_accumulation_steps,
+            **accelerate_kwargs
+        )
+        
+        if exists(wandb_resume_id):
+            init_kwargs={"wandb": {"resume": "allow", "name": wandb_run_name, 'id': wandb_resume_id}}
+        else:
+            init_kwargs={"wandb": {"resume": "allow", "name": wandb_run_name}}
+        self.accelerator.init_trackers(
+            project_name = wandb_project, 
+            init_kwargs=init_kwargs,
+            config={"epochs": epochs,
+                    "learning_rate": learning_rate,
+                    "num_warmup_updates": num_warmup_updates, 
+                    "batch_size": batch_size,
+                    "batch_size_type": batch_size_type,
+                    "max_samples": max_samples,
+                    "grad_accumulation_steps": grad_accumulation_steps,
+                    "max_grad_norm": max_grad_norm,
+                    "gpus": self.accelerator.num_processes,
+                    "noise_scheduler": noise_scheduler}
+            )
+
+        self.model = model
+
+        if self.is_main:
+            self.ema_model = EMA(
+                model,
+                include_online_model = False,
+                **ema_kwargs
+            )
+
+            self.ema_model.to(self.accelerator.device)
+
+        self.epochs = epochs
+        self.num_warmup_updates = num_warmup_updates
+        self.save_per_updates = save_per_updates
+        self.last_per_steps = default(last_per_steps, save_per_updates * grad_accumulation_steps)
+        self.checkpoint_path = default(checkpoint_path, 'ckpts/test_e2-tts')
+
+        self.batch_size = batch_size
+        self.batch_size_type = batch_size_type
+        self.max_samples = max_samples
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+
+        self.noise_scheduler = noise_scheduler
+
+        self.duration_predictor = duration_predictor
+
+        self.optimizer = AdamW(model.parameters(), lr=learning_rate)
+        self.model, self.optimizer = self.accelerator.prepare(
+            self.model, self.optimizer
+        )
+
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save_checkpoint(self, step, last=False):
+        self.accelerator.wait_for_everyone()
+        if self.is_main:
+            checkpoint = dict(
+                model_state_dict = self.accelerator.unwrap_model(self.model).state_dict(),
+                optimizer_state_dict = self.accelerator.unwrap_model(self.optimizer).state_dict(),
+                ema_model_state_dict = self.ema_model.state_dict(),
+                scheduler_state_dict = self.scheduler.state_dict(),
+                step = step
+            )
+            if not os.path.exists(self.checkpoint_path):
+                os.makedirs(self.checkpoint_path)
+            if last == True:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_last.pt")
+                print(f"Saved last checkpoint at step {step}")
+            else:
+                self.accelerator.save(checkpoint, f"{self.checkpoint_path}/model_{step}.pt")
+
+    def load_checkpoint(self):
+        if not exists(self.checkpoint_path) or not os.path.exists(self.checkpoint_path) or not os.listdir(self.checkpoint_path):
+            return 0
+        
+        self.accelerator.wait_for_everyone()
+        if "model_last.pt" in os.listdir(self.checkpoint_path):
+            latest_checkpoint = "model_last.pt"
+        else:
+            latest_checkpoint = sorted(os.listdir(self.checkpoint_path), key=lambda x: int(''.join(filter(str.isdigit, x))))[-1]
+        # checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location=self.accelerator.device)  # rather use accelerator.load_state ಥ_ಥ
+        checkpoint = torch.load(f"{self.checkpoint_path}/{latest_checkpoint}", map_location="cpu")
+        self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint['model_state_dict'])
+        self.accelerator.unwrap_model(self.optimizer).load_state_dict(checkpoint['optimizer_state_dict'])
+
+        if self.is_main:
+            self.ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+
+        if self.scheduler:
+            self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+        
+        step = checkpoint['step']
+        del checkpoint; gc.collect()
+        return step
+
+    def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):
+        
+        if exists(resumable_with_seed):
+            generator = torch.Generator()
+            generator.manual_seed(resumable_with_seed)
+        else: 
+            generator = None
+
+        if self.batch_size_type == "sample":
+            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True, 
+                                          batch_size=self.batch_size, shuffle=True, generator=generator)
+        elif self.batch_size_type == "frame":
+            self.accelerator.even_batches = False
+            sampler = SequentialSampler(train_dataset)
+            batch_sampler = DynamicBatchSampler(sampler, self.batch_size, max_samples=self.max_samples, random_seed=resumable_with_seed, drop_last=False)
+            train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, num_workers=num_workers, pin_memory=True,
+                                          batch_sampler=batch_sampler)
+        else:
+            raise ValueError(f"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}")
+        
+        #  accelerator.prepare() dispatches batches to devices;
+        #  which means the length of dataloader calculated before, should consider the number of devices
+        warmup_steps = self.num_warmup_updates * self.accelerator.num_processes  # consider a fixed warmup steps while using accelerate multi-gpu ddp
+                                                                                 # otherwise by default with split_batches=False, warmup steps change with num_processes
+        total_steps = len(train_dataloader) * self.epochs / self.grad_accumulation_steps
+        decay_steps = total_steps - warmup_steps
+        warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_steps)
+        decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_steps)
+        self.scheduler = SequentialLR(self.optimizer, 
+                                      schedulers=[warmup_scheduler, decay_scheduler],
+                                      milestones=[warmup_steps])
+        train_dataloader, self.scheduler = self.accelerator.prepare(train_dataloader, self.scheduler)  # actual steps = 1 gpu steps / gpus
+        start_step = self.load_checkpoint()
+        global_step = start_step
+
+        if exists(resumable_with_seed):
+            orig_epoch_step = len(train_dataloader)
+            skipped_epoch = int(start_step // orig_epoch_step)
+            skipped_batch = start_step % orig_epoch_step
+            skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)
+        else:
+            skipped_epoch = 0
+
+        for epoch in range(skipped_epoch, self.epochs):
+            self.model.train()
+            if exists(resumable_with_seed) and epoch == skipped_epoch:
+                progress_bar = tqdm(skipped_dataloader, desc=f"Epoch {epoch+1}/{self.epochs}", unit="step", disable=not self.accelerator.is_local_main_process, 
+                                    initial=skipped_batch, total=orig_epoch_step)
+            else:
+                progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{self.epochs}", unit="step", disable=not self.accelerator.is_local_main_process)
+
+            for batch in progress_bar:
+                with self.accelerator.accumulate(self.model):
+                    text_inputs = batch['text']
+                    mel_spec = rearrange(batch['mel'], 'b d n -> b n d')
+                    mel_lengths = batch["mel_lengths"]
+
+                    # TODO. add duration predictor training
+                    if self.duration_predictor is not None and self.accelerator.is_local_main_process:
+                        dur_loss = self.duration_predictor(mel_spec, lens=batch.get('durations'))
+                        self.accelerator.log({"duration loss": dur_loss.item()}, step=global_step)
+
+                    loss, cond, pred = self.model(mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler)
+                    self.accelerator.backward(loss)
+
+                    if self.max_grad_norm > 0 and self.accelerator.sync_gradients:
+                        self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+
+                    self.optimizer.step()
+                    self.scheduler.step()
+                    self.optimizer.zero_grad()
+
+                if self.is_main:
+                    self.ema_model.update()
+
+                global_step += 1
+
+                if self.accelerator.is_local_main_process:
+                    self.accelerator.log({"loss": loss.item(), "lr": self.scheduler.get_last_lr()[0]}, step=global_step)
+                
+                progress_bar.set_postfix(step=str(global_step), loss=loss.item())
+                
+                if global_step % (self.save_per_updates * self.grad_accumulation_steps) == 0:
+                    self.save_checkpoint(global_step)
+                
+                if global_step % self.last_per_steps == 0:
+                    self.save_checkpoint(global_step, last=True)
+        
+        self.accelerator.end_training()

model/utils.py

@@ -0,0 +1,572 @@
+from __future__ import annotations
+
+import os
+import re
+import math
+import random
+import string
+from tqdm import tqdm
+from collections import defaultdict
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pylab as plt
+
+import torch
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+import torchaudio
+
+import einx
+from einops import rearrange, reduce
+
+import jieba
+from pypinyin import lazy_pinyin, Style
+import zhconv
+from zhon.hanzi import punctuation
+from jiwer import compute_measures
+
+from funasr import AutoModel
+from faster_whisper import WhisperModel
+
+from model.ecapa_tdnn import ECAPA_TDNN_SMALL
+from model.modules import MelSpec
+
+
+# seed everything
+
+def seed_everything(seed = 0):
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# tensor helpers
+
+def lens_to_mask(
+    t: int['b'],
+    length: int | None = None
+) -> bool['b n']:
+
+    if not exists(length):
+        length = t.amax()
+
+    seq = torch.arange(length, device = t.device)
+    return einx.less('n, b -> b n', seq, t)
+
+def mask_from_start_end_indices(
+    seq_len: int['b'],
+    start: int['b'],
+    end: int['b']
+):
+    max_seq_len = seq_len.max().item()  
+    seq = torch.arange(max_seq_len, device = start.device).long()
+    return einx.greater_equal('n, b -> b n', seq, start) & einx.less('n, b -> b n', seq, end)
+
+def mask_from_frac_lengths(
+    seq_len: int['b'],
+    frac_lengths: float['b']
+):
+    lengths = (frac_lengths * seq_len).long()
+    max_start = seq_len - lengths
+
+    rand = torch.rand_like(frac_lengths)
+    start = (max_start * rand).long().clamp(min = 0)
+    end = start + lengths
+
+    return mask_from_start_end_indices(seq_len, start, end)
+
+def maybe_masked_mean(
+    t: float['b n d'],
+    mask: bool['b n'] = None
+) -> float['b d']:
+
+    if not exists(mask):
+        return t.mean(dim = 1)
+
+    t = einx.where('b n, b n d, -> b n d', mask, t, 0.)
+    num = reduce(t, 'b n d -> b d', 'sum')
+    den = reduce(mask.float(), 'b n -> b', 'sum')
+
+    return einx.divide('b d, b -> b d', num, den.clamp(min = 1.))
+
+
+# simple utf-8 tokenizer, since paper went character based
+def list_str_to_tensor(
+    text: list[str],
+    padding_value = -1
+) -> int['b nt']:
+    list_tensors = [torch.tensor([*bytes(t, 'UTF-8')]) for t in text]  # ByT5 style
+    text = pad_sequence(list_tensors, padding_value = padding_value, batch_first = True)
+    return text
+
+# char tokenizer, based on custom dataset's extracted .txt file
+def list_str_to_idx(
+    text: list[str] | list[list[str]],
+    vocab_char_map: dict[str, int],  # {char: idx}
+    padding_value = -1
+) -> int['b nt']:
+    list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text]  # pinyin or char style
+    text = pad_sequence(list_idx_tensors, padding_value = padding_value, batch_first = True)
+    return text
+
+
+# Get tokenizer
+
+def get_tokenizer(dataset_name, tokenizer: str = "pinyin"):
+    ''' 
+    tokenizer   - "pinyin" do g2p for only chinese characters, need .txt vocab_file
+                - "char" for char-wise tokenizer, need .txt vocab_file
+                - "byte" for utf-8 tokenizer
+    vocab_size  - if use "pinyin", all available pinyin types, common alphabets (also those with accent) and symbols
+                - if use "char", derived from unfiltered character & symbol counts of custom dataset
+                - if use "byte", set to 256 (unicode byte range) 
+    ''' 
+    if tokenizer in ["pinyin", "char"]:
+        with open (f"data/{dataset_name}_{tokenizer}/vocab.txt", "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                vocab_char_map[char[:-1]] = i
+        vocab_size = len(vocab_char_map)
+        assert vocab_char_map[" "] == 0, "make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char"
+
+    elif tokenizer == "byte":
+        vocab_char_map = None
+        vocab_size = 256
+
+    return vocab_char_map, vocab_size
+
+
+# convert char to pinyin
+
+def convert_char_to_pinyin(text_list, polyphone = True):
+    final_text_list = []
+    god_knows_why_en_testset_contains_zh_quote = str.maketrans({'“': '"', '”': '"', '‘': "'", '’': "'"})  # in case librispeech (orig no-pc) test-clean
+    custom_trans = str.maketrans({';': ','})  # add custom trans here, to address oov
+    for text in text_list:
+        char_list = []
+        text = text.translate(god_knows_why_en_testset_contains_zh_quote)
+        text = text.translate(custom_trans)
+        for seg in jieba.cut(text):
+            seg_byte_len = len(bytes(seg, 'UTF-8'))
+            if seg_byte_len == len(seg):  # if pure alphabets and symbols
+                if char_list and seg_byte_len > 1 and char_list[-1] not in " :'\"":
+                    char_list.append(" ")
+                char_list.extend(seg)
+            elif polyphone and seg_byte_len == 3 * len(seg):  # if pure chinese characters
+                seg = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
+                for c in seg:
+                    if c not in "。，、；：？！《》【】—…":
+                        char_list.append(" ")
+                    char_list.append(c)
+            else:  # if mixed chinese characters, alphabets and symbols
+                for c in seg:
+                    if ord(c) < 256:
+                        char_list.extend(c)
+                    else:
+                        if c not in "。，、；：？！《》【】—…":
+                            char_list.append(" ")
+                            char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
+                        else:  # if is zh punc
+                            char_list.append(c)
+        final_text_list.append(char_list)
+
+    return final_text_list
+
+
+# save spectrogram
+def save_spectrogram(spectrogram, path):
+    plt.figure(figsize=(12, 4))
+    plt.imshow(spectrogram, origin='lower', aspect='auto')
+    plt.colorbar()
+    plt.savefig(path)
+    plt.close()
+
+
+# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav
+def get_seedtts_testset_metainfo(metalst):
+    f = open(metalst); lines = f.readlines(); f.close()
+    metainfo = []
+    for line in lines:
+        if len(line.strip().split('|')) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split('|')
+        elif len(line.strip().split('|')) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split('|')
+            gt_wav = os.path.join(os.path.dirname(metalst), "wavs", utt + ".wav")
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+        metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav))
+    return metainfo
+
+
+# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav
+def get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path):
+    f = open(metalst); lines = f.readlines(); f.close()
+    metainfo = []
+    for line in lines:
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split('\t')
+
+        # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        ref_spk_id, ref_chaptr_id, _ =  ref_utt.split('-')
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + '.flac')
+
+        # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.'  # if use librispeech test-clean (no-pc)
+        gen_spk_id, gen_chaptr_id, _ =  gen_utt.split('-')
+        gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + '.flac')
+
+        metainfo.append((gen_utt, ref_txt, ref_wav, " " + gen_txt, gen_wav))
+
+    return metainfo
+
+
+# padded to max length mel batch
+def padded_mel_batch(ref_mels):
+    max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax()
+    padded_ref_mels = []
+    for mel in ref_mels:
+        padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value = 0)
+        padded_ref_mels.append(padded_ref_mel)
+    padded_ref_mels = torch.stack(padded_ref_mels)
+    padded_ref_mels = rearrange(padded_ref_mels, 'b d n -> b n d')
+    return padded_ref_mels
+
+
+# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav
+
+def get_inference_prompt(
+    metainfo, 
+    speed = 1., tokenizer = "pinyin", polyphone = True, 
+    target_sample_rate = 24000, n_mel_channels = 100, hop_length = 256, target_rms = 0.1,
+    use_truth_duration = False,
+    infer_batch_size = 1, num_buckets = 200, min_secs = 3, max_secs = 40,
+):
+    prompts_all = []
+
+    min_tokens = min_secs * target_sample_rate // hop_length
+    max_tokens = max_secs * target_sample_rate // hop_length
+
+    batch_accum = [0] * num_buckets
+    utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = \
+        ([[] for _ in range(num_buckets)] for _ in range(6))
+
+    mel_spectrogram = MelSpec(target_sample_rate=target_sample_rate, n_mel_channels=n_mel_channels, hop_length=hop_length)
+
+    for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc="Processing prompts..."):
+
+        # Audio
+        ref_audio, ref_sr = torchaudio.load(prompt_wav)
+        ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio)))
+        if ref_rms < target_rms:
+            ref_audio = ref_audio * target_rms / ref_rms
+        assert ref_audio.shape[-1] > 5000, f"Empty prompt wav: {prompt_wav}, or torchaudio backend issue."
+        if ref_sr != target_sample_rate:
+            resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)
+            ref_audio = resampler(ref_audio)
+
+        # Text
+        text = [prompt_text + gt_text]
+        if tokenizer == "pinyin":
+            text_list = convert_char_to_pinyin(text, polyphone = polyphone)
+        else:
+            text_list = text
+
+        # Duration, mel frame length
+        ref_mel_len = ref_audio.shape[-1] // hop_length
+        if use_truth_duration:
+            gt_audio, gt_sr = torchaudio.load(gt_wav)
+            if gt_sr != target_sample_rate:
+                resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate)
+                gt_audio = resampler(gt_audio)
+            total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed)
+
+            # # test vocoder resynthesis
+            # ref_audio = gt_audio
+        else:
+            zh_pause_punc = r"。，、；：？！"
+            ref_text_len = len(prompt_text) + len(re.findall(zh_pause_punc, prompt_text))
+            gen_text_len = len(gt_text) + len(re.findall(zh_pause_punc, gt_text))
+            total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)
+
+        # to mel spectrogram
+        ref_mel = mel_spectrogram(ref_audio)
+        ref_mel = rearrange(ref_mel, '1 d n -> d n')
+
+        # deal with batch
+        assert infer_batch_size > 0, "infer_batch_size should be greater than 0."
+        assert min_tokens <= total_mel_len <= max_tokens, \
+            f"Audio {utt} has duration {total_mel_len*hop_length//target_sample_rate}s out of range [{min_secs}, {max_secs}]."
+        bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets)
+
+        utts[bucket_i].append(utt)
+        ref_rms_list[bucket_i].append(ref_rms)
+        ref_mels[bucket_i].append(ref_mel)
+        ref_mel_lens[bucket_i].append(ref_mel_len)
+        total_mel_lens[bucket_i].append(total_mel_len)
+        final_text_list[bucket_i].extend(text_list)
+
+        batch_accum[bucket_i] += total_mel_len
+
+        if batch_accum[bucket_i] >= infer_batch_size:
+            # print(f"\n{len(ref_mels[bucket_i][0][0])}\n{ref_mel_lens[bucket_i]}\n{total_mel_lens[bucket_i]}")
+            prompts_all.append((
+                utts[bucket_i], 
+                ref_rms_list[bucket_i], 
+                padded_mel_batch(ref_mels[bucket_i]), 
+                ref_mel_lens[bucket_i], 
+                total_mel_lens[bucket_i], 
+                final_text_list[bucket_i]
+            ))
+            batch_accum[bucket_i] = 0
+            utts[bucket_i], ref_rms_list[bucket_i], ref_mels[bucket_i], ref_mel_lens[bucket_i], total_mel_lens[bucket_i], final_text_list[bucket_i] = [], [], [], [], [], []
+
+    # add residual
+    for bucket_i, bucket_frames in enumerate(batch_accum):
+        if bucket_frames > 0:
+            prompts_all.append((
+                utts[bucket_i], 
+                ref_rms_list[bucket_i], 
+                padded_mel_batch(ref_mels[bucket_i]), 
+                ref_mel_lens[bucket_i], 
+                total_mel_lens[bucket_i], 
+                final_text_list[bucket_i]
+            ))
+    # not only leave easy work for last workers
+    random.seed(666)
+    random.shuffle(prompts_all)
+
+    return prompts_all
+
+
+# get wav_res_ref_text of seed-tts test metalst
+# https://github.com/BytedanceSpeech/seed-tts-eval
+
+def get_seed_tts_test(metalst, gen_wav_dir, gpus):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        if len(line.strip().split('|')) == 5:
+            utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split('|')
+        elif len(line.strip().split('|')) == 4:
+            utt, prompt_text, prompt_wav, gt_text = line.strip().split('|')
+
+        if not os.path.exists(os.path.join(gen_wav_dir, utt + '.wav')):
+            continue
+        gen_wav = os.path.join(gen_wav_dir, utt + '.wav')
+        if not os.path.isabs(prompt_wav):
+            prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)
+
+        test_set_.append((gen_wav, prompt_wav, gt_text))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+    
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i*wav_per_job:(i+1)*wav_per_job]))
+
+    return test_set
+
+
+# get librispeech test-clean cross sentence test
+
+def get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = False):
+    f = open(metalst)
+    lines = f.readlines()
+    f.close()
+
+    test_set_ = []
+    for line in tqdm(lines):
+        ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split('\t')
+
+        if eval_ground_truth:
+            gen_spk_id, gen_chaptr_id, _ =  gen_utt.split('-')
+            gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + '.flac')
+        else:
+            if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + '.wav')):
+                raise FileNotFoundError(f"Generated wav not found: {gen_utt}")
+            gen_wav = os.path.join(gen_wav_dir, gen_utt + '.wav')
+
+        ref_spk_id, ref_chaptr_id, _ =  ref_utt.split('-')
+        ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + '.flac')
+
+        test_set_.append((gen_wav, ref_wav, gen_txt))
+
+    num_jobs = len(gpus)
+    if num_jobs == 1:
+        return [(gpus[0], test_set_)]
+    
+    wav_per_job = len(test_set_) // num_jobs + 1
+    test_set = []
+    for i in range(num_jobs):
+        test_set.append((gpus[i], test_set_[i*wav_per_job:(i+1)*wav_per_job]))
+
+    return test_set
+
+
+# load asr model
+
+def load_asr_model(lang, ckpt_dir = ""):
+    if lang == "zh":
+        model = AutoModel(
+            model = os.path.join(ckpt_dir, "paraformer-zh"), 
+            # vad_model = os.path.join(ckpt_dir, "fsmn-vad"), 
+            # punc_model = os.path.join(ckpt_dir, "ct-punc"),
+            # spk_model = os.path.join(ckpt_dir, "cam++"), 
+            disable_update=True,
+            )  # following seed-tts setting
+    elif lang == "en":
+        model_size = "large-v3" if ckpt_dir == "" else ckpt_dir
+        model = WhisperModel(model_size, device="cuda", compute_type="float16")
+    return model
+
+
+# WER Evaluation, the way Seed-TTS does
+
+def run_asr_wer(args):
+    rank, lang, test_set, ckpt_dir = args
+
+    if lang == "zh":
+        torch.cuda.set_device(rank)
+    elif lang == "en":
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(rank)
+    else:
+        raise NotImplementedError("lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now.")
+
+    asr_model = load_asr_model(lang, ckpt_dir = ckpt_dir)
+
+    punctuation_all = punctuation + string.punctuation
+    wers = []
+
+    for gen_wav, prompt_wav, truth in tqdm(test_set):
+        if lang == "zh":
+            res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True)
+            hypo = res[0]["text"]
+            hypo = zhconv.convert(hypo, 'zh-cn')
+        elif lang == "en":
+            segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language="en")
+            hypo = ''
+            for segment in segments:
+                hypo = hypo + ' ' + segment.text
+
+        # raw_truth = truth
+        # raw_hypo = hypo
+
+        for x in punctuation_all:
+            truth = truth.replace(x, '')
+            hypo = hypo.replace(x, '')
+
+        truth = truth.replace('  ', ' ')
+        hypo = hypo.replace('  ', ' ')
+
+        if lang == "zh":
+            truth = " ".join([x for x in truth])
+            hypo = " ".join([x for x in hypo])
+        elif lang == "en":
+            truth = truth.lower()
+            hypo = hypo.lower()
+
+        measures = compute_measures(truth, hypo)
+        wer = measures["wer"]
+
+        # ref_list = truth.split(" ")
+        # subs = measures["substitutions"] / len(ref_list)
+        # dele = measures["deletions"] / len(ref_list)
+        # inse = measures["insertions"] / len(ref_list)
+
+        wers.append(wer)
+
+    return wers
+
+
+# SIM Evaluation
+
+def run_sim(args):
+    rank, test_set, ckpt_dir = args
+    device = f"cuda:{rank}"
+
+    model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type='wavlm_large', config_path=None)
+    state_dict = torch.load(ckpt_dir, map_location=lambda storage, loc: storage)
+    model.load_state_dict(state_dict['model'], strict=False)
+
+    use_gpu=True if torch.cuda.is_available() else False
+    if use_gpu:
+        model = model.cuda(device)
+    model.eval()
+
+    sim_list = []
+    for wav1, wav2, truth in tqdm(test_set):
+
+        wav1, sr1 = torchaudio.load(wav1)
+        wav2, sr2 = torchaudio.load(wav2)
+
+        resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000)
+        resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000)
+        wav1 = resample1(wav1)
+        wav2 = resample2(wav2)
+
+        if use_gpu:
+            wav1 = wav1.cuda(device)
+            wav2 = wav2.cuda(device)
+        with torch.no_grad():
+            emb1 = model(wav1)
+            emb2 = model(wav2)
+        
+        sim = F.cosine_similarity(emb1, emb2)[0].item()
+        # print(f"VSim score between two audios: {sim:.4f} (-1.0, 1.0).")
+        sim_list.append(sim)
+    
+    return sim_list
+
+
+# filter func for dirty data with many repetitions
+
+def repetition_found(text, length = 2, tolerance = 10):
+    pattern_count = defaultdict(int)
+    for i in range(len(text) - length + 1):
+        pattern = text[i:i + length]
+        pattern_count[pattern] += 1
+    for pattern, count in pattern_count.items():
+        if count > tolerance:
+            return True
+    return False
+
+
+# load model checkpoint for inference
+
+def load_checkpoint(model, ckpt_path, device, use_ema = True):
+    from ema_pytorch import EMA
+
+    ckpt_type = ckpt_path.split(".")[-1]
+    if ckpt_type == "safetensors":
+        from safetensors.torch import load_file
+        checkpoint = load_file(ckpt_path, device=device)
+    else:
+        checkpoint = torch.load(ckpt_path, map_location=device)
+
+    if use_ema == True:
+        ema_model = EMA(model, include_online_model = False).to(device)
+        if ckpt_type == "safetensors":
+            ema_model.load_state_dict(checkpoint)
+        else:
+            ema_model.load_state_dict(checkpoint['ema_model_state_dict'])
+        ema_model.copy_params_from_ema_to_model()
+    else:
+        model.load_state_dict(checkpoint['model_state_dict'])
+        
+    return model

requirements.txt

@@ -0,0 +1,25 @@
+accelerate>=0.33.0
+datasets
+einops>=0.8.0
+einx>=0.3.0
+ema_pytorch>=0.5.2
+faster_whisper
+funasr
+jieba
+jiwer
+librosa
+matplotlib
+pypinyin
+safetensors
+# torch>=2.0
+# torchaudio>=2.3.0
+torchdiffeq
+tqdm>=4.65.0
+transformers
+vocos
+wandb
+x_transformers>=1.31.14
+zhconv
+zhon
+gradio
+cached-path

requirements_gradio.txt

@@ -0,0 +1,4 @@
+cached_path
+click
+gradio
+pydub

scripts/count_max_epoch.py

@@ -0,0 +1,32 @@
+'''ADAPTIVE BATCH SIZE'''
+print('Adaptive batch size: using grouping batch sampler, frames_per_gpu fixed fed in')
+print('  -> least padding, gather wavs with accumulated frames in a batch\n')
+
+# data
+total_hours = 95282
+mel_hop_length = 256
+mel_sampling_rate = 24000
+
+# target
+wanted_max_updates = 1000000
+
+# train params
+gpus = 8
+frames_per_gpu = 38400  # 8 * 38400 = 307200
+grad_accum = 1
+
+# intermediate
+mini_batch_frames = frames_per_gpu * grad_accum * gpus
+mini_batch_hours = mini_batch_frames * mel_hop_length / mel_sampling_rate / 3600
+updates_per_epoch = total_hours / mini_batch_hours
+steps_per_epoch = updates_per_epoch * grad_accum
+
+# result
+epochs = wanted_max_updates / updates_per_epoch
+print(f"epochs should be set to: {epochs:.0f} ({epochs/grad_accum:.1f} x gd_acum {grad_accum})")
+print(f"progress_bar should show approx. 0/{updates_per_epoch:.0f} updates")
+print(f"                      or approx. 0/{steps_per_epoch:.0f} steps")
+
+# others
+print(f"total {total_hours:.0f} hours")
+print(f"mini-batch of {mini_batch_frames:.0f} frames, {mini_batch_hours:.2f} hours per mini-batch")

scripts/count_params_gflops.py

@@ -0,0 +1,35 @@
+import sys, os
+sys.path.append(os.getcwd())
+
+from model import M2_TTS, UNetT, DiT, MMDiT
+
+import torch
+import thop
+
+
+''' ~155M '''
+# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4)
+# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4, text_dim = 512, conv_layers = 4)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4)
+# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4, long_skip_connection = True)
+# transformer =     MMDiT(dim = 512, depth = 16, heads = 16, ff_mult = 2)
+
+''' ~335M '''
+# FLOPs: 622.1 G, Params: 333.2 M
+# transformer =     UNetT(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+# FLOPs: 363.4 G, Params: 335.8 M
+transformer =       DiT(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+
+model = M2_TTS(transformer=transformer)
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+duration = 20
+frame_length = int(duration * target_sample_rate / hop_length)
+text_length = 150
+
+flops, params = thop.profile(model, inputs=(torch.randn(1, frame_length, n_mel_channels), torch.zeros(1, text_length, dtype=torch.long)))
+print(f"FLOPs: {flops / 1e9} G")
+print(f"Params: {params / 1e6} M")

scripts/eval_librispeech_test_clean.py

@@ -0,0 +1,67 @@
+# Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation)
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import multiprocessing as mp
+import numpy as np
+
+from model.utils import (
+    get_librispeech_test,
+    run_asr_wer,
+    run_sim,
+)
+
+
+eval_task = "wer"  # sim | wer
+lang = "en"
+metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
+librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
+gen_wav_dir = "PATH_TO_GENERATED"  # generated wavs
+
+gpus = [0,1,2,3,4,5,6,7]
+test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path)
+
+## In LibriSpeech, some speakers utilized varying voice characteristics for different characters in the book,
+## leading to a low similarity for the ground truth in some cases.
+# test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = True)  # eval ground truth
+
+local = False
+if local:  # use local custom checkpoint dir
+    asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+else:
+    asr_ckpt_dir = ""  # auto download to cache dir
+
+wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+
+# --------------------------- WER ---------------------------
+
+if eval_task == "wer":
+    wers = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_asr_wer, args)
+        for wers_ in results:
+            wers.extend(wers_)
+
+    wer = round(np.mean(wers)*100, 3)
+    print(f"\nTotal {len(wers)} samples")
+    print(f"WER      : {wer}%")
+
+
+# --------------------------- SIM ---------------------------
+
+if eval_task == "sim":
+    sim_list = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_sim, args)
+        for sim_ in results:
+            sim_list.extend(sim_)
+
+    sim = round(sum(sim_list)/len(sim_list), 3)
+    print(f"\nTotal {len(sim_list)} samples")
+    print(f"SIM      : {sim}")

scripts/eval_seedtts_testset.py

@@ -0,0 +1,69 @@
+# Evaluate with Seed-TTS testset
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import multiprocessing as mp
+import numpy as np
+
+from model.utils import (
+    get_seed_tts_test,
+    run_asr_wer,
+    run_sim,
+)
+
+
+eval_task = "wer"  # sim | wer
+lang = "zh"        # zh | en
+metalst = f"data/seedtts_testset/{lang}/meta.lst"  # seed-tts testset
+# gen_wav_dir = f"data/seedtts_testset/{lang}/wavs"  # ground truth wavs
+gen_wav_dir = f"PATH_TO_GENERATED"  # generated wavs
+
+
+# NOTE. paraformer-zh result will be slightly different according to the number of gpus, cuz batchsize is different
+#       zh 1.254 seems a result of 4 workers wer_seed_tts 
+gpus = [0,1,2,3,4,5,6,7]
+test_set = get_seed_tts_test(metalst, gen_wav_dir, gpus)
+
+local = False
+if local:  # use local custom checkpoint dir
+    if lang == "zh":
+        asr_ckpt_dir = "../checkpoints/funasr" # paraformer-zh dir under funasr
+    elif lang == "en":
+        asr_ckpt_dir = "../checkpoints/Systran/faster-whisper-large-v3"
+else:
+    asr_ckpt_dir = ""  # auto download to cache dir
+
+wavlm_ckpt_dir = "../checkpoints/UniSpeech/wavlm_large_finetune.pth"
+
+
+# --------------------------- WER ---------------------------
+
+if eval_task == "wer":
+    wers = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_asr_wer, args)
+        for wers_ in results:
+            wers.extend(wers_)
+
+    wer = round(np.mean(wers)*100, 3)
+    print(f"\nTotal {len(wers)} samples")
+    print(f"WER      : {wer}%")
+
+
+# --------------------------- SIM ---------------------------
+
+if eval_task == "sim":
+    sim_list = []
+
+    with mp.Pool(processes=len(gpus)) as pool:
+        args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]
+        results = pool.map(run_sim, args)
+        for sim_ in results:
+            sim_list.extend(sim_)
+
+    sim = round(sum(sim_list)/len(sim_list), 3)
+    print(f"\nTotal {len(sim_list)} samples")
+    print(f"SIM      : {sim}")

scripts/prepare_emilia.py

@@ -0,0 +1,143 @@
+# Emilia Dataset: https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07
+# if use updated new version, i.e. WebDataset, feel free to modify / draft your own script
+
+# generate audio text map for Emilia ZH & EN
+# evaluate for vocab size
+
+import sys, os
+sys.path.append(os.getcwd())
+
+from pathlib import Path
+import json
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+
+from datasets import Dataset
+from datasets.arrow_writer import ArrowWriter
+
+from model.utils import (
+    repetition_found,
+    convert_char_to_pinyin,
+)
+
+
+out_zh = {"ZH_B00041_S06226", "ZH_B00042_S09204", "ZH_B00065_S09430", "ZH_B00065_S09431", "ZH_B00066_S09327", "ZH_B00066_S09328"}
+zh_filters = ["い", "て"]
+# seems synthesized audios, or heavily code-switched
+out_en = {
+    "EN_B00013_S00913", "EN_B00042_S00120", "EN_B00055_S04111", "EN_B00061_S00693", "EN_B00061_S01494", "EN_B00061_S03375",
+    
+    "EN_B00059_S00092", "EN_B00111_S04300", "EN_B00100_S03759", "EN_B00087_S03811", "EN_B00059_S00950", "EN_B00089_S00946", "EN_B00078_S05127", "EN_B00070_S04089", "EN_B00074_S09659", "EN_B00061_S06983", "EN_B00061_S07060", "EN_B00059_S08397", "EN_B00082_S06192", "EN_B00091_S01238", "EN_B00089_S07349", "EN_B00070_S04343", "EN_B00061_S02400", "EN_B00076_S01262", "EN_B00068_S06467", "EN_B00076_S02943", "EN_B00064_S05954", "EN_B00061_S05386", "EN_B00066_S06544", "EN_B00076_S06944", "EN_B00072_S08620", "EN_B00076_S07135", "EN_B00076_S09127", "EN_B00065_S00497", "EN_B00059_S06227", "EN_B00063_S02859", "EN_B00075_S01547", "EN_B00061_S08286", "EN_B00079_S02901", "EN_B00092_S03643", "EN_B00096_S08653", "EN_B00063_S04297", "EN_B00063_S04614", "EN_B00079_S04698", "EN_B00104_S01666", "EN_B00061_S09504", "EN_B00061_S09694", "EN_B00065_S05444", "EN_B00063_S06860", "EN_B00065_S05725", "EN_B00069_S07628", "EN_B00083_S03875", "EN_B00071_S07665", "EN_B00071_S07665", "EN_B00062_S04187", "EN_B00065_S09873", "EN_B00065_S09922", "EN_B00084_S02463", "EN_B00067_S05066", "EN_B00106_S08060", "EN_B00073_S06399", "EN_B00073_S09236", "EN_B00087_S00432", "EN_B00085_S05618", "EN_B00064_S01262", "EN_B00072_S01739", "EN_B00059_S03913", "EN_B00069_S04036", "EN_B00067_S05623", "EN_B00060_S05389", "EN_B00060_S07290", "EN_B00062_S08995",
+}
+en_filters = ["ا", "い", "て"]
+
+
+def deal_with_audio_dir(audio_dir):
+    audio_jsonl = audio_dir.with_suffix(".jsonl")
+    sub_result, durations = [], []
+    vocab_set = set()
+    bad_case_zh = 0
+    bad_case_en = 0
+    with open(audio_jsonl, "r") as f:
+        lines = f.readlines()
+        for line in tqdm(lines, desc=f"{audio_jsonl.stem}"):
+            obj = json.loads(line)
+            text = obj["text"]
+            if obj['language'] == "zh":
+                if obj["wav"].split("/")[1] in out_zh or any(f in text for f in zh_filters) or repetition_found(text):
+                    bad_case_zh += 1
+                    continue
+                else:
+                    text = text.translate(str.maketrans({',': '，', '!': '！', '?': '？'}))  # not "。" cuz much code-switched
+            if obj['language'] == "en":
+                if obj["wav"].split("/")[1] in out_en or any(f in text for f in en_filters) or repetition_found(text, length=4):
+                    bad_case_en += 1
+                    continue
+            if tokenizer == "pinyin":
+                text = convert_char_to_pinyin([text], polyphone = polyphone)[0]
+            duration = obj["duration"]
+            sub_result.append({"audio_path": str(audio_dir.parent / obj["wav"]), "text": text, "duration": duration})
+            durations.append(duration)
+            vocab_set.update(list(text))
+    return sub_result, durations, vocab_set, bad_case_zh, bad_case_en
+
+
+def main():
+    assert tokenizer in ["pinyin", "char"]
+    result = []
+    duration_list = []
+    text_vocab_set = set()
+    total_bad_case_zh = 0
+    total_bad_case_en = 0
+
+    # process raw data
+    executor = ProcessPoolExecutor(max_workers=max_workers)
+    futures = []
+    for lang in langs:
+        dataset_path = Path(os.path.join(dataset_dir, lang))
+        [
+            futures.append(executor.submit(deal_with_audio_dir, audio_dir))
+            for audio_dir in dataset_path.iterdir()
+            if audio_dir.is_dir()
+        ]
+    for futures in tqdm(futures, total=len(futures)):
+        sub_result, durations, vocab_set, bad_case_zh, bad_case_en = futures.result()
+        result.extend(sub_result)
+        duration_list.extend(durations)
+        text_vocab_set.update(vocab_set)
+        total_bad_case_zh += bad_case_zh
+        total_bad_case_en += bad_case_en
+    executor.shutdown()
+
+    # save preprocessed dataset to disk
+    if not os.path.exists(f"data/{dataset_name}"):
+        os.makedirs(f"data/{dataset_name}")
+    print(f"\nSaving to data/{dataset_name} ...")
+    # dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list})  # oom
+    # dataset.save_to_disk(f"data/{dataset_name}/raw", max_shard_size="2GB")
+    with ArrowWriter(path=f"data/{dataset_name}/raw.arrow") as writer:
+        for line in tqdm(result, desc=f"Writing to raw.arrow ..."):
+            writer.write(line)
+
+    # dup a json separately saving duration in case for DynamicBatchSampler ease
+    with open(f"data/{dataset_name}/duration.json", 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)
+
+    # vocab map, i.e. tokenizer
+    # add alphabets and symbols (optional, if plan to ft on de/fr etc.)
+    # if tokenizer == "pinyin":
+    #     text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])
+    with open(f"data/{dataset_name}/vocab.txt", "w") as f:
+        for vocab in sorted(text_vocab_set):
+            f.write(vocab + "\n")
+
+    print(f"\nFor {dataset_name}, sample count: {len(result)}")
+    print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}")
+    print(f"For {dataset_name}, total {sum(duration_list)/3600:.2f} hours")
+    if "ZH" in langs: print(f"Bad zh transcription case: {total_bad_case_zh}")
+    if "EN" in langs: print(f"Bad en transcription case: {total_bad_case_en}\n")
+
+
+if __name__ == "__main__":
+
+    max_workers = 32
+
+    tokenizer = "pinyin"  # "pinyin" | "char"
+    polyphone = True
+
+    langs = ["ZH", "EN"]
+    dataset_dir = "<SOME_PATH>/Emilia_Dataset/raw"
+    dataset_name = f"Emilia_{'_'.join(langs)}_{tokenizer}"
+    print(f"\nPrepare for {dataset_name}\n")
+
+    main()
+
+    # Emilia               ZH & EN
+    # samples count       37837916   (after removal)
+    # pinyin vocab size       2543   (polyphone)
+    # total duration      95281.87   (hours)
+    # bad zh asr cnt        230435   (samples)
+    # bad eh asr cnt         37217   (samples)
+
+    # vocab size may be slightly different due to jieba tokenizer and pypinyin (e.g. way of polyphoneme)
+    # please be careful if using pretrained model, make sure the vocab.txt is same

scripts/prepare_wenetspeech4tts.py

@@ -0,0 +1,116 @@
+# generate audio text map for WenetSpeech4TTS
+# evaluate for vocab size
+
+import sys, os
+sys.path.append(os.getcwd())
+
+import json
+from tqdm import tqdm
+from concurrent.futures import ProcessPoolExecutor
+
+import torchaudio
+from datasets import Dataset
+
+from model.utils import convert_char_to_pinyin
+
+
+def deal_with_sub_path_files(dataset_path, sub_path):
+    print(f"Dealing with: {sub_path}")
+
+    text_dir = os.path.join(dataset_path, sub_path, "txts")
+    audio_dir = os.path.join(dataset_path, sub_path, "wavs")
+    text_files = os.listdir(text_dir)
+
+    audio_paths, texts, durations = [], [], []
+    for text_file in tqdm(text_files):
+        with open(os.path.join(text_dir, text_file), 'r', encoding='utf-8') as file:
+            first_line = file.readline().split("\t")
+        audio_nm = first_line[0]
+        audio_path = os.path.join(audio_dir, audio_nm + ".wav")
+        text = first_line[1].strip()
+
+        audio_paths.append(audio_path)
+
+        if tokenizer == "pinyin":
+            texts.extend(convert_char_to_pinyin([text], polyphone = polyphone))
+        elif tokenizer == "char":
+            texts.append(text)
+
+        audio, sample_rate = torchaudio.load(audio_path)
+        durations.append(audio.shape[-1] / sample_rate)
+
+    return audio_paths, texts, durations
+
+
+def main():
+    assert tokenizer in ["pinyin", "char"]
+
+    audio_path_list, text_list, duration_list = [], [], []
+    
+    executor = ProcessPoolExecutor(max_workers=max_workers)
+    futures = []
+    for dataset_path in dataset_paths:
+        sub_items = os.listdir(dataset_path)
+        sub_paths = [item for item in sub_items if os.path.isdir(os.path.join(dataset_path, item))]
+        for sub_path in sub_paths:
+            futures.append(executor.submit(deal_with_sub_path_files, dataset_path, sub_path))
+    for future in tqdm(futures, total=len(futures)):
+        audio_paths, texts, durations = future.result()
+        audio_path_list.extend(audio_paths)
+        text_list.extend(texts)
+        duration_list.extend(durations)
+    executor.shutdown()
+
+    if not os.path.exists("data"):
+        os.makedirs("data")
+
+    print(f"\nSaving to data/{dataset_name}_{tokenizer} ...")
+    dataset = Dataset.from_dict({"audio_path": audio_path_list, "text": text_list, "duration": duration_list})
+    dataset.save_to_disk(f"data/{dataset_name}_{tokenizer}/raw", max_shard_size="2GB")  # arrow format
+
+    with open(f"data/{dataset_name}_{tokenizer}/duration.json", 'w', encoding='utf-8') as f:
+        json.dump({"duration": duration_list}, f, ensure_ascii=False)  # dup a json separately saving duration in case for DynamicBatchSampler ease
+
+    print("\nEvaluating vocab size (all characters and symbols / all phonemes) ...")
+    text_vocab_set = set()
+    for text in tqdm(text_list):
+        text_vocab_set.update(list(text))
+
+    # add alphabets and symbols (optional, if plan to ft on de/fr etc.)
+    if tokenizer == "pinyin":
+        text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])
+
+    with open(f"data/{dataset_name}_{tokenizer}/vocab.txt", "w") as f:
+        for vocab in sorted(text_vocab_set):
+            f.write(vocab + "\n")
+    print(f"\nFor {dataset_name}, sample count: {len(text_list)}")
+    print(f"For {dataset_name}, vocab size is: {len(text_vocab_set)}\n")
+    
+
+if __name__ == "__main__":
+
+    max_workers = 32
+
+    tokenizer = "pinyin"  # "pinyin" | "char"
+    polyphone = True
+    dataset_choice = 1  # 1: Premium, 2: Standard, 3: Basic
+
+    dataset_name = ["WenetSpeech4TTS_Premium", "WenetSpeech4TTS_Standard", "WenetSpeech4TTS_Basic"][dataset_choice-1]
+    dataset_paths = [
+        "<SOME_PATH>/WenetSpeech4TTS/Basic",
+        "<SOME_PATH>/WenetSpeech4TTS/Standard",
+        "<SOME_PATH>/WenetSpeech4TTS/Premium",
+        ][-dataset_choice:]
+    print(f"\nChoose Dataset: {dataset_name}\n")
+
+    main()
+
+    # Results (if adding alphabets with accents and symbols):
+    # WenetSpeech4TTS       Basic     Standard     Premium
+    # samples count       3932473      1941220      407494
+    # pinyin vocab size      1349         1348        1344   (no polyphone)
+    #                           -            -        1459   (polyphone) 
+    # char   vocab size      5264         5219        5042
+    
+    # vocab size may be slightly different due to jieba tokenizer and pypinyin (e.g. way of polyphoneme)
+    # please be careful if using pretrained model, make sure the vocab.txt is same

test_infer_batch.py

@@ -0,0 +1,197 @@
+import os
+import time
+import random
+from tqdm import tqdm
+import argparse
+
+import torch
+import torchaudio
+from accelerate import Accelerator
+from einops import rearrange
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer, 
+    get_seedtts_testset_metainfo, 
+    get_librispeech_test_clean_metainfo, 
+    get_inference_prompt,
+)
+
+accelerator = Accelerator()
+device = f"cuda:{accelerator.process_index}"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+
+
+# ---------------------- infer setting ---------------------- #
+
+parser = argparse.ArgumentParser(description="batch inference")
+
+parser.add_argument('-s', '--seed', default=None, type=int)
+parser.add_argument('-d', '--dataset', default="Emilia_ZH_EN")
+parser.add_argument('-n', '--expname', required=True)
+parser.add_argument('-c', '--ckptstep', default=1200000, type=int)
+
+parser.add_argument('-nfe', '--nfestep', default=32, type=int)
+parser.add_argument('-o', '--odemethod', default="euler")
+parser.add_argument('-ss', '--swaysampling', default=-1, type=float)
+
+parser.add_argument('-t', '--testset', required=True)
+
+args = parser.parse_args()
+
+
+seed = args.seed
+dataset_name = args.dataset
+exp_name = args.expname
+ckpt_step = args.ckptstep
+ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"
+
+nfe_step = args.nfestep
+ode_method = args.odemethod
+sway_sampling_coef = args.swaysampling
+
+testset = args.testset
+
+
+infer_batch_size = 1  # max frames. 1 for ddp single inference (recommended)
+cfg_strength = 2.
+speed = 1.
+use_truth_duration = False
+no_ref_audio = False
+
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+if testset == "ls_pc_test_clean":
+    metalst = "data/librispeech_pc_test_clean_cross_sentence.lst"
+    librispeech_test_clean_path = "<SOME_PATH>/LibriSpeech/test-clean"  # test-clean path
+    metainfo = get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path)
+    
+elif testset == "seedtts_test_zh":
+    metalst = "data/seedtts_testset/zh/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+elif testset == "seedtts_test_en":
+    metalst = "data/seedtts_testset/en/meta.lst"
+    metainfo = get_seedtts_testset_metainfo(metalst)
+
+
+# path to save genereted wavs
+if seed is None: seed = random.randint(-10000, 10000)
+output_dir = f"results/{exp_name}_{ckpt_step}/{testset}/" \
+    f"seed{seed}_{ode_method}_nfe{nfe_step}" \
+    f"{f'_ss{sway_sampling_coef}' if sway_sampling_coef else ''}" \
+    f"_cfg{cfg_strength}_speed{speed}" \
+    f"{'_gt-dur' if use_truth_duration else ''}" \
+    f"{'_no-ref-audio' if no_ref_audio else ''}"
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+prompts_all = get_inference_prompt(
+    metainfo, 
+    speed = speed, 
+    tokenizer = tokenizer, 
+    target_sample_rate = target_sample_rate, 
+    n_mel_channels = n_mel_channels,
+    hop_length = hop_length,
+    target_rms = target_rms,
+    use_truth_duration = use_truth_duration,
+    infer_batch_size = infer_batch_size,
+)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer = model_cls(
+        **model_cfg,
+        text_num_embeds = vocab_size, 
+        mel_dim = n_mel_channels
+    ),
+    mel_spec_kwargs = dict(
+        target_sample_rate = target_sample_rate, 
+        n_mel_channels = n_mel_channels,
+        hop_length = hop_length,
+    ),
+    odeint_kwargs = dict(
+        method = ode_method,
+    ),
+    vocab_char_map = vocab_char_map,
+).to(device)
+
+model = load_checkpoint(model, ckpt_path, device, use_ema = use_ema)
+
+if not os.path.exists(output_dir) and accelerator.is_main_process:
+    os.makedirs(output_dir)
+
+# start batch inference
+accelerator.wait_for_everyone()
+start = time.time()
+
+with accelerator.split_between_processes(prompts_all) as prompts:
+
+    for prompt in tqdm(prompts, disable=not accelerator.is_local_main_process):
+        utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = prompt
+        ref_mels = ref_mels.to(device)
+        ref_mel_lens = torch.tensor(ref_mel_lens, dtype = torch.long).to(device)
+        total_mel_lens = torch.tensor(total_mel_lens, dtype = torch.long).to(device)
+        
+        # Inference
+        with torch.inference_mode():
+            generated, _ = model.sample(
+                cond = ref_mels,
+                text = final_text_list,
+                duration = total_mel_lens,
+                lens = ref_mel_lens,
+                steps = nfe_step,
+                cfg_strength = cfg_strength,
+                sway_sampling_coef = sway_sampling_coef,
+                no_ref_audio = no_ref_audio,
+                seed = seed,
+            )
+        # Final result
+        for i, gen in enumerate(generated):
+            gen = gen[ref_mel_lens[i]:total_mel_lens[i], :].unsqueeze(0)
+            gen_mel_spec = rearrange(gen, '1 n d -> 1 d n')
+            generated_wave = vocos.decode(gen_mel_spec.cpu())
+            if ref_rms_list[i] < target_rms:
+                generated_wave = generated_wave * ref_rms_list[i] / target_rms
+            torchaudio.save(f"{output_dir}/{utts[i]}.wav", generated_wave, target_sample_rate)
+
+accelerator.wait_for_everyone()
+if accelerator.is_main_process:
+    timediff = time.time() - start
+    print(f"Done batch inference in {timediff / 60 :.2f} minutes.")

test_infer_batch.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# e.g. F5-TTS, 16 NFE
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_zh" -nfe 16
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "seedtts_test_en" -nfe 16
+accelerate launch test_infer_batch.py -n "F5TTS_Base" -t "ls_pc_test_clean" -nfe 16
+
+# e.g. Vanilla E2 TTS, 32 NFE
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_zh" -o "midpoint" -ss 0
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "seedtts_test_en" -o "midpoint" -ss 0
+accelerate launch test_infer_batch.py -n "E2TTS_Base" -t "ls_pc_test_clean" -o "midpoint" -ss 0
+
+# etc.

test_infer_single.py

@@ -0,0 +1,163 @@
+import os
+import re
+
+import torch
+import torchaudio
+from einops import rearrange
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT, MMDiT
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer, 
+    convert_char_to_pinyin, 
+    save_spectrogram,
+)
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+ckpt_step = 1200000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.
+ode_method = 'euler'  # euler | midpoint
+sway_sampling_coef = -1.
+speed = 0.8
+fix_duration = 27  # None (will linear estimate. if code-switched, consider fix) | float (total in seconds, include ref audio) 
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"
+output_dir = "tests"
+
+ref_audio = "tests/ref_audio/rashmika_input.wav"
+ref_text = ""
+
+# gen_text = "Hi everyone! This is Rashmika Mandanna, sending you my heartfelt Diwali wishes on behalf of Antriksh AI. May this festival of lights fill your life with love, joy, and togetherness. Let’s celebrate the spirit of Diwali by spreading happiness and peace wherever we go. At Antriksh AI, we’re excited to blend innovation with creativity, and this Diwali, we wish you endless light and positivity. Enjoy the festivities and cherish every moment! With love, Rashmika Mandanna & Antriksh AI"
+
+gen_text_ = "Happy Birthday, Dhillip Kumar. Virat Kohli this side, all the best for your future endeavours! "
+
+
+# ref_audio = "tests/ref_audio/test_zh_1_ref_short.wav"
+# ref_text = "对，这就是我，万人敬仰的太乙真人。"
+# gen_text = "突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道：\"我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？\""
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer = model_cls(
+        **model_cfg,
+        text_num_embeds = vocab_size, 
+        mel_dim = n_mel_channels
+    ),
+    mel_spec_kwargs = dict(
+        target_sample_rate = target_sample_rate, 
+        n_mel_channels = n_mel_channels,
+        hop_length = hop_length,
+    ),
+    odeint_kwargs = dict(
+        method = ode_method,
+    ),
+    vocab_char_map = vocab_char_map,
+).to(device)
+
+model = load_checkpoint(model, ckpt_path, device, use_ema = use_ema)
+
+# Audio
+audio, sr = torchaudio.load(ref_audio)
+if audio.shape[0] > 1:
+    audio = torch.mean(audio, dim=0, keepdim=True)
+rms = torch.sqrt(torch.mean(torch.square(audio)))
+if rms < target_rms:
+    audio = audio * target_rms / rms
+if sr != target_sample_rate:
+    resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+    audio = resampler(audio)
+audio = audio.to(device)
+
+# Text
+text_list = [ref_text + gen_text]
+if tokenizer == "pinyin":
+    final_text_list = convert_char_to_pinyin(text_list)
+else:
+    final_text_list = [text_list]
+print(f"text  : {text_list}")
+print(f"pinyin: {final_text_list}")
+
+# Duration
+ref_audio_len = audio.shape[-1] // hop_length
+if fix_duration is not None:
+    duration = int(fix_duration * target_sample_rate / hop_length)
+else:  # simple linear scale calcul
+    zh_pause_punc = r"。，、；：？！"
+    ref_text_len = len(ref_text) + len(re.findall(zh_pause_punc, ref_text))
+    gen_text_len = len(gen_text) + len(re.findall(zh_pause_punc, gen_text))
+    duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed)
+
+# Inference
+with torch.inference_mode():
+    generated, trajectory = model.sample(
+        cond = audio,
+        text = final_text_list,
+        duration = duration,
+        steps = nfe_step,
+        cfg_strength = cfg_strength,
+        sway_sampling_coef = sway_sampling_coef,
+        seed = seed,
+    )
+print(f"Generated mel: {generated.shape}")
+
+# Final result
+generated = generated[:, ref_audio_len:, :]
+generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
+generated_wave = vocos.decode(generated_mel_spec.cpu())
+if rms < target_rms:
+    generated_wave = generated_wave * rms / target_rms
+
+save_spectrogram(generated_mel_spec[0].cpu().numpy(), f"{output_dir}/test_single_dbday.png")
+torchaudio.save(f"{output_dir}/test_single_dbday.wav", generated_wave, target_sample_rate)
+print(f"Generated wav: {generated_wave.shape}")

test_infer_single_edit.py

@@ -0,0 +1,182 @@
+import os
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from einops import rearrange
+from vocos import Vocos
+
+from model import CFM, UNetT, DiT, MMDiT
+from model.utils import (
+    load_checkpoint,
+    get_tokenizer, 
+    convert_char_to_pinyin, 
+    save_spectrogram,
+)
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+# --------------------- Dataset Settings -------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+target_rms = 0.1
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+ckpt_step = 1200000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.
+ode_method = 'euler'  # euler | midpoint
+sway_sampling_coef = -1.
+speed = 1.
+
+if exp_name == "F5TTS_Base":
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+
+elif exp_name == "E2TTS_Base":
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+ckpt_path = f"ckpts/{exp_name}/model_{ckpt_step}.pt"
+output_dir = "tests"
+
+# [leverage https://github.com/MahmoudAshraf97/ctc-forced-aligner to get char level alignment]
+# pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git
+# [write the origin_text into a file, e.g. tests/test_edit.txt]
+# ctc-forced-aligner --audio_path "tests/ref_audio/test_en_1_ref_short.wav" --text_path "tests/test_edit.txt" --language "zho" --romanize --split_size "char"
+# [result will be saved at same path of audio file]
+# [--language "zho" for Chinese, "eng" for English]
+# [if local ckpt, set --alignment_model "../checkpoints/mms-300m-1130-forced-aligner"]
+
+audio_to_edit = "tests/ref_audio/test_en_1_ref_short.wav"
+origin_text = "Some call me nature, others call me mother nature."
+target_text = "Some call me optimist, others call me realist."
+parts_to_edit = [[1.42, 2.44], [4.04, 4.9], ]  # stard_ends of "nature" & "mother nature", in seconds
+fix_duration = [1.2, 1, ]  # fix duration for "optimist" & "realist", in seconds
+
+# audio_to_edit = "tests/ref_audio/test_zh_1_ref_short.wav"
+# origin_text = "对，这就是我，万人敬仰的太乙真人。"
+# target_text = "对，那就是你，万人敬仰的太白金星。"
+# parts_to_edit = [[0.84, 1.4], [1.92, 2.4], [4.26, 6.26], ]
+# fix_duration = None  # use origin text duration
+
+
+# -------------------------------------------------#
+
+use_ema = True
+
+if not os.path.exists(output_dir):
+    os.makedirs(output_dir)
+
+# Vocoder model
+local = False
+if local:
+    vocos_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+    vocos = Vocos.from_hparams(f"{vocos_local_path}/config.yaml")
+    state_dict = torch.load(f"{vocos_local_path}/pytorch_model.bin", map_location=device)
+    vocos.load_state_dict(state_dict)
+    vocos.eval()
+else:
+    vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz")
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer = model_cls(
+        **model_cfg,
+        text_num_embeds = vocab_size, 
+        mel_dim = n_mel_channels
+    ),
+    mel_spec_kwargs = dict(
+        target_sample_rate = target_sample_rate, 
+        n_mel_channels = n_mel_channels,
+        hop_length = hop_length,
+    ),
+    odeint_kwargs = dict(
+        method = ode_method,
+    ),
+    vocab_char_map = vocab_char_map,
+).to(device)
+
+model = load_checkpoint(model, ckpt_path, device, use_ema = use_ema)
+
+# Audio
+audio, sr = torchaudio.load(audio_to_edit)
+if audio.shape[0] > 1:
+    audio = torch.mean(audio, dim=0, keepdim=True)
+rms = torch.sqrt(torch.mean(torch.square(audio)))
+if rms < target_rms:
+    audio = audio * target_rms / rms
+if sr != target_sample_rate:
+    resampler = torchaudio.transforms.Resample(sr, target_sample_rate)
+    audio = resampler(audio)
+offset = 0
+audio_ = torch.zeros(1, 0)
+edit_mask = torch.zeros(1, 0, dtype=torch.bool)
+for part in parts_to_edit:
+    start, end = part
+    part_dur = end - start if fix_duration is None else fix_duration.pop(0)
+    part_dur = part_dur * target_sample_rate
+    start = start * target_sample_rate
+    audio_ = torch.cat((audio_, audio[:, round(offset):round(start)], torch.zeros(1, round(part_dur))), dim = -1)
+    edit_mask = torch.cat((edit_mask, 
+                           torch.ones(1, round((start - offset) / hop_length), dtype = torch.bool), 
+                           torch.zeros(1, round(part_dur / hop_length), dtype = torch.bool)
+                           ), dim = -1)
+    offset = end * target_sample_rate
+# audio = torch.cat((audio_, audio[:, round(offset):]), dim = -1)
+edit_mask = F.pad(edit_mask, (0, audio.shape[-1] // hop_length - edit_mask.shape[-1] + 1), value = True)
+audio = audio.to(device)
+edit_mask = edit_mask.to(device)
+
+# Text
+text_list = [target_text]
+if tokenizer == "pinyin":
+    final_text_list = convert_char_to_pinyin(text_list)
+else:
+    final_text_list = [text_list]
+print(f"text  : {text_list}")
+print(f"pinyin: {final_text_list}")
+
+# Duration
+ref_audio_len = 0
+duration = audio.shape[-1] // hop_length
+
+# Inference
+with torch.inference_mode():
+    generated, trajectory = model.sample(
+        cond = audio,
+        text = final_text_list,
+        duration = duration,
+        steps = nfe_step,
+        cfg_strength = cfg_strength,
+        sway_sampling_coef = sway_sampling_coef,
+        seed = seed,
+        edit_mask = edit_mask,
+    )
+print(f"Generated mel: {generated.shape}")
+
+# Final result
+generated = generated[:, ref_audio_len:, :]
+generated_mel_spec = rearrange(generated, '1 n d -> 1 d n')
+generated_wave = vocos.decode(generated_mel_spec.cpu())
+if rms < target_rms:
+    generated_wave = generated_wave * rms / target_rms
+
+save_spectrogram(generated_mel_spec[0].cpu().numpy(), f"{output_dir}/test_single_edit.png")
+torchaudio.save(f"{output_dir}/test_single_edit.wav", generated_wave, target_sample_rate)
+print(f"Generated wav: {generated_wave.shape}")

test_train.py

@@ -0,0 +1,91 @@
+from model import CFM, UNetT, DiT, MMDiT, Trainer
+from model.utils import get_tokenizer
+from model.dataset import load_dataset
+
+
+# -------------------------- Dataset Settings --------------------------- #
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+
+tokenizer = "pinyin"
+dataset_name = "Emilia_ZH_EN"
+
+
+# -------------------------- Training Settings -------------------------- #
+
+exp_name = "F5TTS_Base"  # F5TTS_Base | E2TTS_Base
+
+learning_rate = 7.5e-5
+
+batch_size_per_gpu = 38400  # 8 GPUs, 8 * 38400 = 307200
+batch_size_type = "frame"  # "frame" or "sample"
+max_samples = 64  # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models
+grad_accumulation_steps = 1  # note: updates = steps / grad_accumulation_steps
+max_grad_norm = 1.
+
+epochs = 11  # use linear decay, thus epochs control the slope
+num_warmup_updates = 20000  # warmup steps
+save_per_updates = 50000  # save checkpoint per steps
+last_per_steps = 5000  # save last checkpoint per steps
+
+# model params
+if exp_name == "F5TTS_Base":
+    wandb_resume_id = None
+    model_cls = DiT
+    model_cfg = dict(dim = 1024, depth = 22, heads = 16, ff_mult = 2, text_dim = 512, conv_layers = 4)
+elif exp_name == "E2TTS_Base":
+    wandb_resume_id = None
+    model_cls = UNetT
+    model_cfg = dict(dim = 1024, depth = 24, heads = 16, ff_mult = 4)
+
+
+# ----------------------------------------------------------------------- #
+
+def main():
+
+    vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+    mel_spec_kwargs = dict(
+            target_sample_rate = target_sample_rate, 
+            n_mel_channels = n_mel_channels,
+            hop_length = hop_length,
+        )
+    
+    e2tts = CFM(
+        transformer = model_cls(
+            **model_cfg,
+            text_num_embeds = vocab_size, 
+            mel_dim = n_mel_channels
+        ),
+        mel_spec_kwargs = mel_spec_kwargs,
+        vocab_char_map = vocab_char_map,
+    )
+
+    trainer = Trainer(
+        e2tts,
+        epochs, 
+        learning_rate,
+        num_warmup_updates = num_warmup_updates,
+        save_per_updates = save_per_updates, 
+        checkpoint_path = f'ckpts/{exp_name}',
+        batch_size = batch_size_per_gpu, 
+        batch_size_type = batch_size_type,
+        max_samples = max_samples,
+        grad_accumulation_steps = grad_accumulation_steps,
+        max_grad_norm = max_grad_norm,
+        wandb_project = "CFM-TTS",
+        wandb_run_name = exp_name,
+        wandb_resume_id = wandb_resume_id,
+        last_per_steps = last_per_steps,
+    )
+
+    train_dataset = load_dataset(dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)
+    trainer.train(train_dataset, 
+                  resumable_with_seed = 666 # seed for shuffling dataset
+                  )
+
+
+if __name__ == '__main__':
+    main()