Duplicate from bharatgenai/sooktam2

Co-authored-by: Vansh Pundir <vanshp123@users.noreply.huggingface.co>

.gitattributes

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

.gitignore

@@ -0,0 +1,217 @@
+# Customed
+.vscode/
+tests/
+runs/
+data
+data/
+ckpts/
+wandb/
+results/
+packaged_results/
+packaged_results_backup/
+outputs/
+data/
+data
+ray_results/
+hindi_models.zip
+hindi_models
+pretrained_model.safetensors
+awscliv2.zip
+sukumo28.wav-preview-2.6.0 (1).vsix
+
+hf_inference_models/
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+ckpts/
+data/
+
+
+# C extensions
+*.so
+.cursor
+# 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/
+results/
+# 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/
+5066549580888614_chunk_5_0000076.wav
+
+hindi_models.zip
+hindi_models/whisper-large-hi-noldcil/pytorch_model.bin
+hindi_models/whisper-medium-hi_alldata_multigpu/pytorch_model.bin
+pretrained_model.safetensors
+awscliv2.zip
+src/third_party/BigVGAN/filelists/LibriTTS/train-full.txt
+src/third_party/BigVGAN/demo/examples/megalovania_24k.wav
+hindi_models/whisper-large-hi-noldcil/vocab.json
+src/third_party/BigVGAN/demo/examples/musdbhq_44k.wav
+src/third_party/BigVGAN/demo/examples/musiccaps1_44k.wav
+src/third_party/BigVGAN/demo/examples/musiccaps2_44k.wav
+(1).vsix
+src/third_party/BigVGAN/demo/examples/hifitts_44k.wav
+hindi_models/whisper-large-hi-noldcil/merges.txt
+src/third_party/BigVGAN/demo/examples/jensen_24k.wav
+src/third_party/BigVGAN/demo/examples/dance_24k.wav
+src/third_party/BigVGAN/demo/examples/queen_24k.wav
+hindi_models/whisper-large-hi-noldcil/trainer_state.json
+src/f5_tts/infer/examples/basic/basic_ref_zh.wav
+hindi_models/whisper-medium-hi_alldata_multigpu/trainer_state.json
+src/third_party/BigVGAN/demo/examples/libritts_24k.wav
+src/f5_tts/infer/examples/multi/main.flac
+src/f5_tts/infer/examples/basic/basic_ref_en.wav
+src/f5_tts/infer/examples/multi/town.flac
+hindi_models/whisper-large-hi-noldcil/preprocessor_config.json
+src/f5_tts/infer/examples/multi/country.flac

LICENSE

@@ -0,0 +1,50 @@
+BharatGen Research License (BRL)
+Last Updated: August 7, 2025
+Note: This license is subject to change. Please refer to the latest version before using any BharatGen resource.
+
+1. Scope and Acceptance
+- This license applies to any use, modification, or distribution of BharatGen Models, Derivatives, or Outputs.
+- By accessing, using, modifying, or distributing a BharatGen Model or its Derivatives, or by creating Outputs from them, you agree to be bound by this License.
+- If you accept on behalf of an organization, you represent that you have authority to do so.
+
+2. Permitted Use
+- You are granted a non-exclusive, non-transferable, royalty-free license to use, modify, and distribute BharatGen Models and Derivatives strictly for non-commercial research and academic purposes.
+- You must retain the following notice in all copies and distributions:
+
+"Licensed by BharatGen under the BharatGen Research License."
+- If you create Derivatives, you must clearly indicate that modifications were made by you.
+- You must not use BharatGen Models, Derivatives, or Outputs for any commercial purposes without prior written permission from the BharatGen legal team.
+
+3. Commercial Use & Modifications
+- Commercial use is strictly prohibited unless you obtain explicit written permission from the BharatGen legal team.
+- If you wish to release, share, publish, fine-tune, or deploy any BharatGen Models or Derivatives (including Outputs) publicly, you must request and receive written approval from the BharatGen team.
+Please reach out via our contact form or designated email.
+
+4. Limitations
+- You must not claim that your modifications, Derivatives, or Outputs are official BharatGen products.
+- You may not use BharatGen resources for any unlawful or unethical purposes.
+- Redistribution of BharatGen resources through SaaS, APIs, or any hosted service—even for free—requires permission.
+
+5. Outputs
+- You own the content (Outputs) you generate using BharatGen models, subject to the non-commercial restriction.
+- You are fully responsible for the Outputs you generate and their usage.
+
+6. Disclaimer & Liability
+- BharatGen resources are provided "AS IS", without warranties of any kind.
+- The BharatGen team are not liable for any damages arising from the use of BharatGen Models, Derivatives, or Outputs.
+
+7. Termination
+- This license is valid until terminated.
+- It will terminate immediately if you breach any of its terms. Upon termination, you must stop using and delete all copies of the Models and Derivatives.
+
+8. Jurisdiction
+This License is governed by the laws of India. Any disputes shall be subject to the courts of India.
+
+Definitions
+BharatGen Model: Any model, code, data, or related resource released under this license.
+Derivative: Any modification, fine-tuning, or adaptation of a BharatGen Model.
+Output: Any result generated by running a BharatGen Model.
+You: Any individual or organization using BharatGen resources.
+Non-commercial research: Use in academic, personal, or scientific contexts not intended to generate revenue.
+
+If you have questions or require a commercial license, please contact the BharatGen team.

README.md

@@ -0,0 +1,174 @@
+---
+language:
+- en
+- hi
+- mr
+- gu
+- ta
+- te
+- kn
+- bn
+- ml
+- or
+- ur
+- pa
+pipeline_tag: text-to-speech
+library_name: transformers
+tags:
+- text-to-speech
+- tts
+- multilingual
+- indic
+- f5-tts
+- sooktam2
+---
+
+<p align="center">
+  <img src="https://cdn-avatars.huggingface.co/v1/production/uploads/67b462a1f4f414c2b3e2bc2f/EnVeNWEIeZ6yF6ueZ7E3Y.jpeg" width="140" alt="BharatGen Logo"/>
+</p>
+
+<h1 align="center">Sooktam-2 🇮🇳</h1>
+<p align="center"><em>"विविधता में ही भारत की शक्ति है, और हर भाषा उस शक्ति की आवाज़ है।"</em></p>
+<p align="center"><b>Sovereign AI · Built in Bharat · For Bharat</b></p>
+
+---
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1YvgkOL7mM7vcOE8IHOhHD9PprYUh5bvb)
+
+## The Story
+
+India is not one voice - it is a symphony. Tamil, Bengali, Urdu, Hindi, Kannada - each a living civilisation, spoken daily by hundreds of millions. Yet for too long, AI treated them as afterthoughts. Models built elsewhere, for someone else, leaving Bharat to make do with approximations of its own languages.
+
+**BharatGen was built to end that.** We are India's sovereign AI initiative - weaving the country's languages, cultures, and voices into technology that is truly Indian. Not adapted. Not translated. *Built from the ground up, for Bharat.*
+
+**Sooktam-2** is our answer to India's need for a voice. A Text-to-Speech model that speaks 12 languages (11 indian languages + 1 indian english) with the phonetic precision, prosody, and cultural soul they deserve - so that every Indian, in every state, can hear AI speak *their* language, in *their* accent, and feel at home.
+
+This is **GenAI for Bharat, by Bharat.**
+
+---
+
+## What is Sooktam-2?
+
+Sooktam-2 is a sovereign multilingual Text-to-Speech model built by BharatGen. It synthesises natural, expressive speech across India's major languages using reference-guided voice conditioning - preserving the speaker's voice, accent, and cultural cadence.
+
+**Represented Languages - 12**
+
+`Hindi` · `Marathi` · `Gujarati` · `Tamil` · `Telugu` · `Kannada` · `Bengali` · `Malayalam` · `Odia` · `Urdu` · `Punjabi` · `Indian English`
+
+**Key Capabilities**
+- Reference-guided voice cloning
+- Multilingual Indic speech synthesis
+- Natural prosody and expressive delivery
+- Language-aware CLS tokenization for accurate Indic phonetics
+- Production-quality audio output, deployment-ready at scale
+
+---
+
+## Quickstart
+
+- Python version = 3.10
+
+```bash
+git clone https://huggingface.co/bharatgenai/sooktam2
+cd sooktam2
+sh setup-cls.sh
+```
+
+---
+
+## Python Inference
+
+```python
+import os
+from transformers import AutoModel
+
+# --- Model ID ---
+MODEL_ID = "bharatgenai/sooktam2"
+
+# --- Your reference audio and target text ---
+REF_AUDIO = "reference.wav"   # A short, clean voice clip (3–10 sec)
+REF_TEXT  = "सर, मैं तब से यह कह रहा हूँ कि मैंने अपना टिकट कैंसल कर दिया है, लेकिन अब तक मेरे पैसे वापस नहीं आए हैं। आप इस मामले को देखेंगे भी या नहीं?"
+GEN_TEXT  = "यह एक टेस्ट वाक्य है जिसे आवाज़ में बदलना है।"
+
+# --- Output ---
+OUT_DIR = "outputs"
+OUT_WAV = os.path.join(OUT_DIR, "sooktam_cls.wav")
+
+# --- Load model (auto-downloads checkpoint + vocab from HuggingFace) ---
+model = AutoModel.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+# CLS tokenization is handled inside utils_infer via cls_tokenizer_v2
+wav, sr, _ = model.infer(
+    ref_file=REF_AUDIO,
+    ref_text=REF_TEXT,
+    gen_text=GEN_TEXT,
+    tokenizer="cls",
+    cls_language="hindi",
+    file_wave=OUT_WAV,
+)
+
+print("Saved:", OUT_WAV, "sample_rate:", sr, "samples:", len(wav))
+```
+
+> The model and vocab download automatically from HuggingFace on first run. No manual checkpoint hunting required.
+
+---
+
+## Hugging Face AutoModel
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "bharatgenai/sooktam2",
+    trust_remote_code=True,
+)
+
+wav, sr, _ = model.infer(
+    ref_file="ref.wav",
+    ref_text="Your reference transcript.",
+    gen_text="Text you want to synthesise.",
+    tokenizer="cls",
+    cls_language="hindi",
+)
+```
+
+---
+
+## License
+This post-trained checkpoint is released under the BharatGen non-commercial license.
+
+Please refer to the [LICENSE](./LICENSE) file for detailed terms and conditions.
+
+---
+
+## Contributors
+- Yash
+- Supreet
+- Isha
+- Vansh
+- Pranav
+
+For any questions or feedback, please contact: contact@bharatgen.com
+
+---
+
+## BharatGen - Sovereign AI for a Sovereign Nation
+
+BharatGen is India's initiative to build AI that is Indian in its roots, inclusive in its reach, and sovereign in its design. We believe that a nation of India's civilisational depth - of Sanskrit and Tamil, of Tagore and Kabir, of a billion daily conversations - should not have to borrow its voice from elsewhere.
+
+India's languages are not a niche. They are the world's richest linguistic heritage. And now, they have a model built for them.
+
+We are just getting started.
+
+---
+
+<p align="center">
+  <a href="https://bharatgen.com">bharatgen.com</a> · <a href="https://huggingface.co/bharatgenai/sooktam2">HuggingFace ↗</a>
+  <br/><br/>
+  <b>जय हिन्द · जय भारत 🇮🇳</b>
+</p>

config.json

@@ -0,0 +1,13 @@
+{
+  "model_type": "f5_tts",
+  "model_name": "F5TTS_v1_Base",
+  "ckpt_file": "model_1250000.pt",
+  "vocab_file": "vocab.txt",
+  "ode_method": "euler",
+  "use_ema": true,
+  "auto_map": {
+    "AutoConfig": "hf_auto.F5TTSConfig",
+    "AutoModel": "hf_auto.F5TTSAutoModel",
+    "AutoTokenizer": "hf_auto.F5TTSTokenizer"
+  }
+}

hf_auto.py

@@ -0,0 +1,18 @@
+"""HF Auto classes entrypoint for remote loading (used by AutoModel/AutoTokenizer)."""
+
+import os
+import sys
+
+_ROOT = os.path.dirname(os.path.abspath(__file__))
+_SRC = os.path.join(_ROOT, "src")
+if _SRC not in sys.path:
+    sys.path.insert(0, _SRC)
+
+from f5_tts.hf_auto import (  # noqa: E402
+    F5TTSConfig,
+    F5TTSAutoModel,
+    F5TTSTokenizer,
+    register_f5tts_auto,
+)
+
+__all__ = ["F5TTSConfig", "F5TTSAutoModel", "F5TTSTokenizer", "register_f5tts_auto"]

model_1250000.pt

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

pyproject.toml

@@ -0,0 +1,74 @@
+[build-system]
+requires = ["setuptools >= 61.0", "setuptools-scm>=8.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "f5-tts"
+version = "1.1.9"
+description = "F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching"
+readme = "README.md"
+license = {text = "MIT License"}
+classifiers = [
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+
+]
+requires-python = ">=3.9,<3.12"
+dependencies = [
+    "accelerate>=0.33.0",
+    "bitsandbytes>0.37.0; platform_machine!='arm64' and platform_system!='Darwin'",
+    "cached_path",
+    "click",
+    "datasets",
+    "ema-pytorch==0.7.9",
+    "gradio>=5.0.0",
+    "hydra-core>=1.3.0",
+    "jiwer",
+    "jieba",
+    "zhon",
+    "pytorch-wpe==0.0.1",
+    "torch-complex==0.4.4",
+    "torchcodec==0.9.1",
+    "librosa",
+    "matplotlib",
+    "numpy<=1.26.4; python_version<='3.10'",
+    "pydantic<=2.10.6",
+    "pydub",
+    "pypinyin",
+    "safetensors",
+    "soundfile",
+    "tomli",
+    "torch==2.6.0",
+    "torchaudio==2.6.0",
+    "torchdiffeq==0.2.5",
+    "torchelastic==0.2.2",
+    "torchvision==0.21.0",
+    "tqdm>=4.65.0",
+    "transformers",
+    "transformers_stream_generator",
+    "unidecode",
+    "vocos",
+    "wandb",
+    "x_transformers>=1.31.14",
+]
+
+[project.optional-dependencies]
+eval = [
+    "faster_whisper==1.2.1",
+    "funasr",
+    "jiwer",
+    "modelscope",
+    "s3prl",
+    "zhconv",
+    "zhon",
+]
+
+[project.urls]
+Homepage = "https://github.com/SWivid/F5-TTS"
+
+[project.scripts]
+"f5-tts_infer-cli" = "f5_tts.infer.infer_cli:main"
+"f5-tts_infer-gradio" = "f5_tts.infer.infer_gradio:main"
+"f5-tts_finetune-cli" = "f5_tts.train.finetune_cli:main"
+"f5-tts_finetune-gradio" = "f5_tts.train.finetune_gradio:main"

ref.wav

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

setup-cls.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+set -e
+
+
+# apt update
+# apt-get install libavformat-dev libavdevice-dev -y
+# if gpu is available
+if [ -n "$(lspci | grep NVIDIA)" ]; then
+    apt-get install libcudnn9-cuda-12 -y
+fi
+pip install torch==2.6.0 torchaudio==2.6.0 torchvision==0.21.0 transformers==4.41.2 huggingface-hub==0.24.7 indic-nlp-library==0.92
+
+
+echo "Installing dependencies for prepare_cls_version.py..."
+pip install \
+    datasets==4.5.0 \
+    tqdm \
+    indo-arabic-transliteration==0.1.5 \
+    indic-unified-parser==1.0.6 \
+    indic-numtowords==1.1.0 \
+    indic-nlp-library==0.92 \
+    git+https://github.com/libindic/indic-trans.git@0287fa62289968f0ce06cbe2df61cfadf4088c75 \
+    urduhack==1.1.1 \
+    keras==2.15.0 \
+    tensorflow==2.15.0 \
+    tensorflow-addons==0.23.0 \
+    fastapi==0.128.7 \
+    uvicorn==0.40.0
+
+pip install indic_unified_parser==1.0.6 indo-arabic-transliteration==0.1.5 indic-numtowords==1.1.0
+pip install click==8.0.1
+pip install -e . --no-cache-dir

sooktam.safetensors

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

src/f5_tts/api.py

@@ -0,0 +1,174 @@
+import random
+import sys
+from importlib.resources import files
+
+import soundfile as sf
+import tqdm
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    save_spectrogram,
+    transcribe,
+)
+from f5_tts.model.utils import seed_everything
+
+
+class F5TTS:
+    def __init__(
+        self,
+        model="F5TTS_v1_Base",
+        ckpt_file="",
+        vocab_file="",
+        ode_method="euler",
+        use_ema=True,
+        vocoder_local_path=None,
+        device=None,
+        hf_cache_dir=None,
+    ):
+        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+        model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+        model_arc = model_cfg.model.arch
+
+        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+        self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate
+
+        self.ode_method = ode_method
+        self.use_ema = use_ema
+
+        if device is not None:
+            self.device = device
+        else:
+            import torch
+
+            self.device = (
+                "cuda"
+                if torch.cuda.is_available()
+                else "xpu"
+                if torch.xpu.is_available()
+                else "mps"
+                if torch.backends.mps.is_available()
+                else "cpu"
+            )
+
+        # Load models
+        self.vocoder = load_vocoder(
+            self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir
+        )
+
+        repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
+
+        # override for previous models
+        if model == "F5TTS_Base":
+            if self.mel_spec_type == "vocos":
+                ckpt_step = 1200000
+            elif self.mel_spec_type == "bigvgan":
+                model = "F5TTS_Base_bigvgan"
+                ckpt_type = "pt"
+        elif model == "E2TTS_Base":
+            repo_name = "E2-TTS"
+            ckpt_step = 1200000
+
+        if not ckpt_file:
+            ckpt_file = str(
+                cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}", cache_dir=hf_cache_dir)
+            )
+        self.ema_model = load_model(
+            model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device
+        )
+
+    def transcribe(self, ref_audio, language=None):
+        return transcribe(ref_audio, language)
+
+    def export_wav(self, wav, file_wave, remove_silence=False):
+        sf.write(file_wave, wav, self.target_sample_rate)
+
+        if remove_silence:
+            remove_silence_for_generated_wav(file_wave)
+
+    def export_spectrogram(self, spec, file_spec):
+        save_spectrogram(spec, file_spec)
+
+    def infer(
+        self,
+        ref_file,
+        ref_text,
+        gen_text,
+        show_info=print,
+        progress=tqdm,
+        target_rms=0.1,
+        cross_fade_duration=0.15,
+        sway_sampling_coef=-1,
+        cfg_strength=2,
+        nfe_step=32,
+        speed=1.0,
+        fix_duration=None,
+        remove_silence=False,
+        file_wave=None,
+        file_spec=None,
+        seed=None,
+        tokenizer="pinyin",
+        cls_language=None,
+        cls_server_url=None,
+        cls_timeout=5.0,
+        cls_tokenizer_fn=None,
+    ):
+        if seed is None:
+            seed = random.randint(0, sys.maxsize)
+        seed_everything(seed)
+        self.seed = seed
+
+        ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text)
+
+        wav, sr, spec = infer_process(
+            ref_file,
+            ref_text,
+            gen_text,
+            self.ema_model,
+            self.vocoder,
+            self.mel_spec_type,
+            show_info=show_info,
+            progress=progress,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=speed,
+            fix_duration=fix_duration,
+            device=self.device,
+            tokenizer=tokenizer,
+            cls_language=cls_language,
+            cls_server_url=cls_server_url,
+            cls_timeout=cls_timeout,
+            cls_tokenizer_fn=cls_tokenizer_fn,
+        )
+
+        if file_wave is not None:
+            self.export_wav(wav, file_wave, remove_silence)
+
+        if file_spec is not None:
+            self.export_spectrogram(spec, file_spec)
+
+        return wav, sr, spec
+
+
+if __name__ == "__main__":
+    f5tts = F5TTS()
+
+    wav, sr, spec = f5tts.infer(
+        ref_file=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")),
+        ref_text="some call me nature, others call me mother nature.",
+        gen_text="""I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring. Respect me and I'll nurture you; ignore me and you shall face the consequences.""",
+        file_wave=str(files("f5_tts").joinpath("../../tests/api_out.wav")),
+        file_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
+        seed=None,
+    )
+
+    print("seed :", f5tts.seed)

src/f5_tts/configs/E2TTS_Base.yaml

@@ -0,0 +1,22 @@
+model:
+  name: E2TTS_Base
+  tokenizer: pinyin
+  tokenizer_path: null
+  backbone: UNetT
+  arch:
+    dim: 1024
+    depth: 24
+    heads: 16
+    ff_mult: 4
+    text_mask_padding: False
+    pe_attn_head: 1
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/configs/E2TTS_Small.yaml

@@ -0,0 +1,22 @@
+model:
+  name: E2TTS_Small
+  tokenizer: pinyin
+  tokenizer_path: null
+  backbone: UNetT
+  arch:
+    dim: 768
+    depth: 18
+    heads: 12
+    ff_mult: 4
+    text_mask_padding: False
+    pe_attn_head: 1
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/configs/F5TTS_Base.yaml

@@ -0,0 +1,27 @@
+model:
+  name: F5TTS_Base
+  tokenizer: pinyin
+  tokenizer_path: null
+  backbone: DiT
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: False
+    conv_layers: 4
+    pe_attn_head: 1
+    attn_backend: torch
+    attn_mask_enabled: False
+    checkpoint_activations: False
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/configs/F5TTS_Small.yaml

@@ -0,0 +1,27 @@
+model:
+  name: F5TTS_Small
+  tokenizer: pinyin
+  tokenizer_path: null
+  backbone: DiT
+  arch:
+    dim: 768
+    depth: 18
+    heads: 12
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: False
+    conv_layers: 4
+    pe_attn_head: 1
+    attn_backend: torch
+    attn_mask_enabled: False
+    checkpoint_activations: False
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/configs/F5TTS_v1_Base.yaml

@@ -0,0 +1,28 @@
+model:
+  name: F5TTS_v1_Base
+  tokenizer: pinyin
+  tokenizer_path: null
+  backbone: DiT
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: True
+    qk_norm: null
+    conv_layers: 4
+    pe_attn_head: null
+    attn_backend: torch
+    attn_mask_enabled: False
+    checkpoint_activations: False
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/configs/F5TTS_v1_Base_frame.yaml

@@ -0,0 +1,28 @@
+model:
+  name: F5TTS_v1_Base
+  tokenizer: char
+  tokenizer_path: null
+  backbone: DiT
+  arch:
+    dim: 1024
+    depth: 22
+    heads: 16
+    ff_mult: 2
+    text_dim: 512
+    text_mask_padding: True
+    qk_norm: null
+    conv_layers: 4
+    pe_attn_head: null
+    attn_backend: torch
+    attn_mask_enabled: False
+    checkpoint_activations: False
+  mel_spec:
+    target_sample_rate: 24000
+    n_mel_channels: 100
+    hop_length: 256
+    win_length: 1024
+    n_fft: 1024
+    mel_spec_type: vocos
+  vocoder:
+    is_local: False
+    local_path: null

src/f5_tts/hf_auto.py

@@ -0,0 +1,228 @@
+"""Hugging Face AutoModel integration for F5-TTS (inference-only)."""
+
+from __future__ import annotations
+
+import os
+from typing import Any, List, Optional
+
+import torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoConfig, AutoModel, AutoTokenizer
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.utils import logging
+
+from f5_tts.api import F5TTS
+
+logger = logging.get_logger(__name__)
+
+
+class F5TTSConfig(PretrainedConfig):
+    model_type = "f5_tts"
+
+    def __init__(
+        self,
+        model_name: str = "F5TTS_v1_Base",
+        ckpt_file: str = "",
+        vocab_file: str = "",
+        ode_method: str = "euler",
+        use_ema: bool = True,
+        vocoder_local_path: Optional[str] = None,
+        device: Optional[str] = None,
+        hf_cache_dir: Optional[str] = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.ckpt_file = ckpt_file
+        self.vocab_file = vocab_file
+        self.ode_method = ode_method
+        self.use_ema = use_ema
+        self.vocoder_local_path = vocoder_local_path
+        self.device = device
+        self.hf_cache_dir = hf_cache_dir
+
+        if "auto_map" not in kwargs:
+            # Keep AutoTokenizer as a string to satisfy Hub config validators.
+            self.auto_map = {
+                "AutoConfig": "hf_auto.F5TTSConfig",
+                "AutoModel": "hf_auto.F5TTSAutoModel",
+                "AutoTokenizer": "hf_auto.F5TTSTokenizer",
+            }
+
+
+class F5TTSTokenizer(PreTrainedTokenizer):
+    """Minimal character-level tokenizer backed by vocab.txt (inference helper)."""
+
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(self, vocab_file: str, **kwargs) -> None:
+        self.vocab_file = vocab_file
+        tokens = self._load_vocab_tokens(vocab_file)
+        self.vocab = {tok: idx for idx, tok in enumerate(tokens)}
+        self.ids_to_tokens = {idx: tok for tok, idx in self.vocab.items()}
+
+        if kwargs.get("unk_token") is None:
+            kwargs["unk_token"] = "<unk>"
+        super().__init__(**kwargs)
+
+        if self.unk_token not in self.vocab:
+            unk_id = len(self.vocab)
+            self.vocab[self.unk_token] = unk_id
+            self.ids_to_tokens[unk_id] = self.unk_token
+
+    @staticmethod
+    def _load_vocab_tokens(path: str) -> List[str]:
+        with open(path, "r", encoding="utf-8") as handle:
+            return [line.rstrip("\n") for line in handle]
+
+    def get_vocab(self) -> dict:
+        return dict(self.vocab)
+
+    def _tokenize(self, text: str) -> List[str]:
+        return list(text)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self.ids_to_tokens.get(index, self.unk_token)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
+        os.makedirs(save_directory, exist_ok=True)
+        filename = (filename_prefix + "-" if filename_prefix else "") + "vocab.txt"
+        path = os.path.join(save_directory, filename)
+        with open(path, "w", encoding="utf-8") as handle:
+            for idx in range(len(self.ids_to_tokens)):
+                handle.write(f"{self.ids_to_tokens[idx]}\n")
+        return (path,)
+
+
+def load_tokenizer(
+    repo_or_path: str = "bharatgenai/sooktam2",
+    vocab_file: str = "vocab.txt",
+    cache_dir: Optional[str] = None,
+    revision: Optional[str] = None,
+    token: Optional[str] = None,
+    local_files_only: bool = False,
+) -> F5TTSTokenizer:
+    """Load the character-level tokenizer from a local folder or Hugging Face."""
+    resolved = F5TTSAutoModel._resolve_file(
+        vocab_file,
+        repo_or_path,
+        cache_dir,
+        revision,
+        token,
+        local_files_only,
+    )
+    return F5TTSTokenizer(resolved)
+
+
+class F5TTSAutoModel(PreTrainedModel):
+    config_class = F5TTSConfig
+
+    def __init__(self, config: F5TTSConfig, ckpt_file: str = "", vocab_file: str = "", **kwargs) -> None:
+        super().__init__(config)
+        self._dummy = torch.nn.Parameter(torch.zeros(1), requires_grad=False)
+
+        self.tts = F5TTS(
+            model=config.model_name,
+            ckpt_file=ckpt_file or config.ckpt_file,
+            vocab_file=vocab_file or config.vocab_file,
+            ode_method=config.ode_method,
+            use_ema=config.use_ema,
+            vocoder_local_path=config.vocoder_local_path,
+            device=config.device,
+            hf_cache_dir=config.hf_cache_dir,
+        )
+
+    @staticmethod
+    def _resolve_file(
+        filename: str,
+        repo_or_path: Optional[str],
+        cache_dir: Optional[str],
+        revision: Optional[str],
+        token: Optional[str],
+        local_files_only: bool,
+    ) -> str:
+        if not filename:
+            return ""
+        if os.path.isfile(filename):
+            return filename
+        if repo_or_path and os.path.isdir(repo_or_path):
+            candidate = os.path.join(repo_or_path, filename)
+            if os.path.isfile(candidate):
+                return candidate
+        if not repo_or_path:
+            return filename
+        return hf_hub_download(
+            repo_id=repo_or_path,
+            filename=filename,
+            cache_dir=cache_dir,
+            revision=revision,
+            token=token,
+            local_files_only=local_files_only,
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Optional[str], *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        if config is None:
+            config_kwargs = {
+                "cache_dir": kwargs.get("cache_dir"),
+                "revision": kwargs.get("revision"),
+                "token": kwargs.get("token"),
+                "local_files_only": kwargs.get("local_files_only", False),
+                "trust_remote_code": kwargs.get("trust_remote_code"),
+            }
+            try:
+                config = F5TTSConfig.from_pretrained(pretrained_model_name_or_path, **config_kwargs)
+            except Exception:  # noqa: BLE001
+                logger.warning("F5TTSConfig not found, using defaults.")
+                config = F5TTSConfig()
+
+        ckpt_file = kwargs.pop("ckpt_file", None) or config.ckpt_file
+        vocab_file = kwargs.pop("vocab_file", None) or config.vocab_file
+
+        cache_dir = kwargs.get("cache_dir") or config.hf_cache_dir
+        revision = kwargs.get("revision")
+        token = kwargs.get("token")
+        local_files_only = kwargs.get("local_files_only", False)
+
+        ckpt_file = cls._resolve_file(
+            ckpt_file,
+            pretrained_model_name_or_path,
+            cache_dir,
+            revision,
+            token,
+            local_files_only,
+        )
+        vocab_file = cls._resolve_file(
+            vocab_file,
+            pretrained_model_name_or_path,
+            cache_dir,
+            revision,
+            token,
+            local_files_only,
+        )
+
+        return cls(config, ckpt_file=ckpt_file, vocab_file=vocab_file)
+
+    def forward(self, *args, **kwargs):  # noqa: D401
+        raise NotImplementedError("Use .infer(...) or .tts.infer(...) for generation.")
+
+    def infer(self, *args, **kwargs):
+        return self.tts.infer(*args, **kwargs)
+
+    def save_pretrained(self, save_directory: str, **kwargs):
+        os.makedirs(save_directory, exist_ok=True)
+        self.config.save_pretrained(save_directory)
+
+
+def register_f5tts_auto() -> None:
+    """Register F5-TTS with Hugging Face AutoConfig/AutoModel/AutoTokenizer (local usage)."""
+    AutoConfig.register(F5TTSConfig.model_type, F5TTSConfig)
+    AutoModel.register(F5TTSConfig, F5TTSAutoModel)
+    AutoTokenizer.register(F5TTSConfig, F5TTSTokenizer)

src/f5_tts/infer/README.md

@@ -0,0 +1,177 @@
+# Inference
+
+The pretrained model checkpoints can be reached at [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS) and [🤖 Model Scope](https://www.modelscope.cn/models/SWivid/F5-TTS_Emilia-ZH-EN), or will be automatically downloaded when running inference scripts.
+
+**More checkpoints with whole community efforts can be found in [SHARED.md](SHARED.md), supporting more languages.**
+
+Currently support **30s for a single** generation, which is the **total length** (same logic if `fix_duration`) including both prompt and output audio. However, `infer_cli` and `infer_gradio` will automatically do chunk generation for longer text. Long reference audio will be **clip short to ~12s**.
+
+To avoid possible inference failures, make sure you have seen through the following instructions.
+
+- Use reference audio <12s and leave proper silence space (e.g. 1s) at the end. Otherwise there is a risk of truncating in the middle of word, leading to suboptimal generation.
+- <ins>Uppercased letters</ins> (best with form like K.F.C.) will be uttered letter by letter, and lowercased letters used for common words. 
+- Add some spaces (blank: " ") or punctuations (e.g. "," ".") <ins>to explicitly introduce some pauses</ins>.
+- If English punctuation marks the end of a sentence, make sure there is a space " " after it. Otherwise not regarded as when chunk.
+- <ins>Preprocess numbers</ins> to Chinese letters if you want to have them read in Chinese, otherwise in English.
+- If the generation output is blank (pure silence), <ins>check for FFmpeg installation</ins>.
+- Try <ins>turn off `use_ema` if using an early-stage</ins> finetuned checkpoint (which goes just few updates).
+
+
+## Gradio App
+
+Currently supported features:
+
+- Basic TTS with Chunk Inference
+- Multi-Style / Multi-Speaker Generation
+- Voice Chat powered by Qwen2.5-3B-Instruct
+- [Custom inference with more language support](SHARED.md)
+
+The cli command `f5-tts_infer-gradio` equals to `python src/f5_tts/infer/infer_gradio.py`, which launches a Gradio APP (web interface) for inference.
+
+The script will load model checkpoints from Huggingface. You can also manually download files and update the path to `load_model()` in `infer_gradio.py`. Currently only load TTS models first, will load ASR model to do transcription if `ref_text` not provided, will load LLM model if use Voice Chat.
+
+More flags options:
+
+```bash
+# Automatically launch the interface in the default web browser
+f5-tts_infer-gradio --inbrowser
+
+# Set the root path of the application, if it's not served from the root ("/") of the domain
+# For example, if the application is served at "https://example.com/myapp"
+f5-tts_infer-gradio --root_path "/myapp"
+```
+
+Could also be used as a component for larger application:
+```python
+import gradio as gr
+from f5_tts.infer.infer_gradio import app
+
+with gr.Blocks() as main_app:
+    gr.Markdown("# This is an example of using F5-TTS within a bigger Gradio app")
+
+    # ... other Gradio components
+
+    app.render()
+
+main_app.launch()
+```
+
+
+## CLI Inference
+
+The cli command `f5-tts_infer-cli` equals to `python src/f5_tts/infer/infer_cli.py`, which is a command line tool for inference.
+
+The script will load model checkpoints from Huggingface. You can also manually download files and use `--ckpt_file` to specify the model you want to load, or directly update in `infer_cli.py`.
+
+For change vocab.txt use `--vocab_file` to provide your `vocab.txt` file.
+
+Basically you can inference with flags:
+```bash
+# Leave --ref_text "" will have ASR model transcribe (extra GPU memory usage)
+f5-tts_infer-cli \
+--model F5TTS_v1_Base \
+--ref_audio "ref_audio.wav" \
+--ref_text "The content, subtitle or transcription of reference audio." \
+--gen_text "Some text you want TTS model generate for you."
+
+# Use BigVGAN as vocoder. Currently only support F5TTS_Base. 
+f5-tts_infer-cli --model F5TTS_Base --vocoder_name bigvgan --load_vocoder_from_local
+
+# Use custom path checkpoint, e.g.
+f5-tts_infer-cli --ckpt_file ckpts/F5TTS_v1_Base/model_1250000.safetensors
+
+# More instructions
+f5-tts_infer-cli --help
+```
+
+And a `.toml` file would help with more flexible usage.
+
+```bash
+f5-tts_infer-cli -c custom.toml
+```
+
+For example, you can use `.toml` to pass in variables, refer to `src/f5_tts/infer/examples/basic/basic.toml`:
+
+```toml
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/basic/basic_ref_en.wav"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = "Some call me nature, others call me mother nature."
+gen_text = "I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring."
+# File with text to generate. Ignores the text above.
+gen_file = ""
+remove_silence = false
+output_dir = "tests"
+```
+
+You can also leverage `.toml` file to do multi-style generation, refer to `src/f5_tts/infer/examples/multi/story.toml`.
+
+```toml
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/multi/main.flac"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = ""
+gen_text = ""
+# File with text to generate. Ignores the text above.
+gen_file = "infer/examples/multi/story.txt"
+remove_silence = true
+output_dir = "tests"
+
+[voices.town]
+ref_audio = "infer/examples/multi/town.flac"
+ref_text = ""
+
+[voices.country]
+ref_audio = "infer/examples/multi/country.flac"
+ref_text = ""
+```
+You should mark the voice with `[main]` `[town]` `[country]` whenever you want to change voice, refer to `src/f5_tts/infer/examples/multi/story.txt`.
+
+## API Usage
+
+```python
+from importlib.resources import files
+from f5_tts.api import F5TTS
+
+f5tts = F5TTS()
+wav, sr, spec = f5tts.infer(
+    ref_file=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")),
+    ref_text="some call me nature, others call me mother nature.",
+    gen_text="""I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring. Respect me and I'll nurture you; ignore me and you shall face the consequences.""",
+    file_wave=str(files("f5_tts").joinpath("../../tests/api_out.wav")),
+    file_spec=str(files("f5_tts").joinpath("../../tests/api_out.png")),
+    seed=None,
+)
+```
+Check [api.py](../api.py) for more details.
+
+## TensorRT-LLM Deployment
+
+See [detailed instructions](../runtime/triton_trtllm/README.md) for more information.
+
+## Socket Real-time Service
+
+Real-time voice output with chunk stream:
+
+```bash
+# Start socket server
+python src/f5_tts/socket_server.py
+
+# If PyAudio not installed
+sudo apt-get install portaudio19-dev
+pip install pyaudio
+
+# Communicate with socket client
+python src/f5_tts/socket_client.py
+```
+
+## Speech Editing
+
+To test speech editing capabilities, use the following command:
+
+```bash
+python src/f5_tts/infer/speech_edit.py
+```
+

src/f5_tts/infer/SHARED.md

@@ -0,0 +1,193 @@
+<!-- omit in toc -->
+# Shared Model Cards
+
+<!-- omit in toc -->
+### **Prerequisites of using**
+- This document is serving as a quick lookup table for the community training/finetuning result, with various language support.
+- The models in this repository are open source and are based on voluntary contributions from contributors.
+- The use of models must be conditioned on respect for the respective creators. The convenience brought comes from their efforts.
+
+<!-- omit in toc -->
+### **Welcome to share here**
+- Have a pretrained/finetuned result: model checkpoint (pruned best to facilitate inference, i.e. leave only `ema_model_state_dict`) and corresponding vocab file (for tokenization).
+- Host a public [huggingface model repository](https://huggingface.co/new) and upload the model related files.
+- Make a pull request adding a model card to the current page, i.e. `src\f5_tts\infer\SHARED.md`.
+
+<!-- omit in toc -->
+### Supported Languages
+- [Multilingual](#multilingual)
+    - [F5-TTS v1 v0 Base @ zh \& en @ F5-TTS](#f5-tts-v1-v0-base--zh--en--f5-tts)
+- [English](#english)
+- [Finnish](#finnish)
+    - [F5-TTS Base @ fi @ AsmoKoskinen](#f5-tts-base--fi--asmokoskinen)
+- [French](#french)
+    - [F5-TTS Base @ fr @ RASPIAUDIO](#f5-tts-base--fr--raspiaudio)
+- [German](#german)
+    - [F5-TTS Base @ de @ hvoss-techfak](#f5-tts-base--de--hvoss-techfak)
+- [Hindi](#hindi)
+    - [F5-TTS Small @ hi @ SPRINGLab](#f5-tts-small--hi--springlab)
+- [Italian](#italian)
+    - [F5-TTS Base @ it @ alien79](#f5-tts-base--it--alien79)
+- [Japanese](#japanese)
+    - [F5-TTS Base @ ja @ Jmica](#f5-tts-base--ja--jmica)
+- [Mandarin](#mandarin)
+- [Russian](#russian)
+    - [F5-TTS Base @ ru @ HotDro4illa](#f5-tts-base--ru--hotdro4illa)
+- [Spanish](#spanish)
+    - [F5-TTS Base @ es @ jpgallegoar](#f5-tts-base--es--jpgallegoar)
+
+
+## Multilingual
+
+#### F5-TTS v1 v0 Base @ zh & en @ F5-TTS
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS v1 Base|[ckpt & vocab](https://huggingface.co/SWivid/F5-TTS/tree/main/F5TTS_v1_Base)|[Emilia 95K zh&en](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07)|cc-by-nc-4.0|
+
+```bash
+Model: hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors
+# A Variant Model: hf://SWivid/F5-TTS/F5TTS_v1_Base_no_zero_init/model_1250000.safetensors
+Vocab: hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "conv_layers": 4}
+```
+
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/SWivid/F5-TTS/tree/main/F5TTS_Base)|[Emilia 95K zh&en](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07)|cc-by-nc-4.0|
+
+```bash
+Model: hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors
+Vocab: hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+*Other infos, e.g. Author info, Github repo, Link to some sampled results, Usage instruction, Tutorial (Blog, Video, etc.) ...*
+
+
+## English
+
+
+## Finnish
+
+#### F5-TTS Base @ fi @ AsmoKoskinen
+|Model|🤗Hugging Face|Data|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/AsmoKoskinen/F5-TTS_Finnish_Model)|[Common Voice](https://huggingface.co/datasets/mozilla-foundation/common_voice_17_0), [Vox Populi](https://huggingface.co/datasets/facebook/voxpopuli)|cc-by-nc-4.0|
+
+```bash
+Model: hf://AsmoKoskinen/F5-TTS_Finnish_Model/model_common_voice_fi_vox_populi_fi_20241206.safetensors
+Vocab: hf://AsmoKoskinen/F5-TTS_Finnish_Model/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+
+## French
+
+#### F5-TTS Base @ fr @ RASPIAUDIO
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/RASPIAUDIO/F5-French-MixedSpeakers-reduced)|[LibriVox](https://librivox.org/)|cc-by-nc-4.0|
+
+```bash
+Model: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/model_last_reduced.pt
+Vocab: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- [Online Inference with Hugging Face Space](https://huggingface.co/spaces/RASPIAUDIO/f5-tts_french).
+- [Tutorial video to train a new language model](https://www.youtube.com/watch?v=UO4usaOojys).
+- [Discussion about this training can be found here](https://github.com/SWivid/F5-TTS/issues/434).
+
+
+## German
+
+#### F5-TTS Base @ de @ hvoss-techfak
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/hvoss-techfak/F5-TTS-German)|[Mozilla Common Voice 19.0](https://commonvoice.mozilla.org/en/datasets) & 800 hours Crowdsourced |cc-by-nc-4.0|
+
+```bash
+Model: hf://hvoss-techfak/F5-TTS-German/model_f5tts_german.pt
+Vocab: hf://hvoss-techfak/F5-TTS-German/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- Finetuned by [@hvoss-techfak](https://github.com/hvoss-techfak)
+
+
+## Hindi
+
+#### F5-TTS Small @ hi @ SPRINGLab
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Small|[ckpt & vocab](https://huggingface.co/SPRINGLab/F5-Hindi-24KHz)|[IndicTTS Hi](https://huggingface.co/datasets/SPRINGLab/IndicTTS-Hindi) & [IndicVoices-R Hi](https://huggingface.co/datasets/SPRINGLab/IndicVoices-R_Hindi) |cc-by-4.0|
+
+```bash
+Model: hf://SPRINGLab/F5-Hindi-24KHz/model_2500000.safetensors
+Vocab: hf://SPRINGLab/F5-Hindi-24KHz/vocab.txt
+Config: {"dim": 768, "depth": 18, "heads": 12, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- Authors: SPRING Lab, Indian Institute of Technology, Madras
+- Website: https://asr.iitm.ac.in/
+
+
+## Italian
+
+#### F5-TTS Base @ it @ alien79
+|Model|🤗Hugging Face|Data|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/alien79/F5-TTS-italian)|[ylacombe/cml-tts](https://huggingface.co/datasets/ylacombe/cml-tts) |cc-by-nc-4.0|
+
+```bash
+Model: hf://alien79/F5-TTS-italian/model_159600.safetensors
+Vocab: hf://alien79/F5-TTS-italian/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+- Trained by [Mithril Man](https://github.com/MithrilMan)
+- Model details on [hf project home](https://huggingface.co/alien79/F5-TTS-italian)
+- Open to collaborations to further improve the model
+
+
+## Japanese
+
+#### F5-TTS Base @ ja @ Jmica
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/Jmica/F5TTS/tree/main/JA_21999120)|[Emilia 1.7k JA](https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07) & [Galgame Dataset 5.4k](https://huggingface.co/datasets/OOPPEENN/Galgame_Dataset)|cc-by-nc-4.0|
+
+```bash
+Model: hf://Jmica/F5TTS/JA_21999120/model_21999120.pt
+Vocab: hf://Jmica/F5TTS/JA_21999120/vocab_japanese.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+
+
+## Mandarin
+
+
+## Russian
+
+#### F5-TTS Base @ ru @ HotDro4illa
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/hotstone228/F5-TTS-Russian)|[Common voice](https://huggingface.co/datasets/mozilla-foundation/common_voice_17_0)|cc-by-nc-4.0|
+
+```bash
+Model: hf://hotstone228/F5-TTS-Russian/model_last.safetensors
+Vocab: hf://hotstone228/F5-TTS-Russian/vocab.txt
+Config: {"dim": 1024, "depth": 22, "heads": 16, "ff_mult": 2, "text_dim": 512, "text_mask_padding": False, "conv_layers": 4, "pe_attn_head": 1}
+```
+- Finetuned by [HotDro4illa](https://github.com/HotDro4illa)
+- Any improvements are welcome
+
+
+## Spanish
+
+#### F5-TTS Base @ es @ jpgallegoar
+|Model|🤗Hugging Face|Data (Hours)|Model License|
+|:---:|:------------:|:-----------:|:-------------:|
+|F5-TTS Base|[ckpt & vocab](https://huggingface.co/jpgallegoar/F5-Spanish)|[Voxpopuli](https://huggingface.co/datasets/facebook/voxpopuli) & Crowdsourced & TEDx, 218 hours|cc0-1.0|
+
+- @jpgallegoar [GitHub repo](https://github.com/jpgallegoar/Spanish-F5), Jupyter Notebook and Gradio usage for Spanish model.

src/f5_tts/infer/cls_tokenizer_v2.py

@@ -0,0 +1,181 @@
+import re
+import string
+
+try:
+    from indic_unified_parser.uparser import wordparse
+except Exception:  # noqa: BLE001
+    wordparse = None
+
+script_ranges = {
+    # Indo-Aryan
+    "devanagari": [("\u0900", "\u097F")],   # Hindi, Marathi
+    "arabic": [("\u0600", "\u06FF")],       # Urdu
+    "gurmukhi": [("\u0A00", "\u0A7F")],     # Punjabi
+    "gujarati": [("\u0A80", "\u0AFF")],     # Gujarati
+    "bengali": [("\u0980", "\u09FF")],      # Bengali
+    "odia": [("\u0B00", "\u0B7F")],         # Odia
+
+    # Dravidian
+    "tamil": [("\u0B80", "\u0BFF")],        # Tamil
+    "telugu": [("\u0C00", "\u0C7F")],       # Telugu
+    "kannada": [("\u0C80", "\u0CFF")],      # Kannada
+    "malayalam": [("\u0D00", "\u0D7F")],    # Malayalam
+
+    # English + digits + common punctuation
+    "latin_basic": [("\u0020", "\u007E")]   # English letters, digits, ASCII symbols
+}
+
+def has_non_indic_script(text):
+    """
+    Check if text contains any other characters from the specified Indic/Urdu scripts.
+    
+    Returns True if any character falls in the Unicode ranges outside of Hindi/Marathi (Devanagari),
+    Gujarati, Punjabi (Gurmukhi), Urdu (Arabic), English (Latin) False otherwise.
+    """
+    for char in text:
+        for lang, ranges in script_ranges.items():
+            for start, end in ranges:
+                if start <= char <= end:
+                    return False
+    return True
+
+non_problematic_chars = set()
+
+for lang, ranges in script_ranges.items():
+    for start, end in ranges:
+        for char in range(ord(start), ord(end)+1):
+            parsed = None
+            try:
+                parsed = wordparse(chr(char), 0, 0, 1)
+            except Exception as e:
+                pass
+            if parsed is not None and isinstance(parsed, str) and parsed.strip() != "":
+                non_problematic_chars.add(chr(char))
+
+def get_transliteration(text, language):
+    if language.lower() == "urdu":
+        # We will not add bias with transliteration as of now
+        # text = ml_transliterate(text, from_script="ur-PK", to_script="hi-IN")
+        pass
+    elif language.lower() == "punjabi":
+        # GURUMUKHI TO DEVNAGRI IS NOT THERE - SKIPPING FOR NOW, WILL REVIST IF CLS SHOWS EVIDENCE OF IMPROVEMENT
+        # text = script_convert(text, from_script="pa-IN", to_script="hi-IN")
+        pass
+    return text
+
+def normalize_indic_nasals(text):
+    # Combined pattern and replacement using capturing groups for script-specific anusvara and consonant groups
+    pattern = (
+        r'(ं|ং|ં|ਂ|ಂ|ം|ଂ|ం|ஂ)'  # anusvara chars for Devanagari, Bengali, Gujarati, Punjabi, Kannada, Malayalam, Odia, Telugu, Tamil
+        r'([कखगघङचछजझञटठडढणतथदधनपफबभम'
+        r'কখগঘঙচছজঝঞটঠডঢণতথদধনপফবভম'
+        r'કખગઘઙચછજઝઞટઠડઢણતથદધનપફબભમ'
+        r'ਕਖਗਘਙਚਛਜਝਞਟਠਡਢਣਤਥਦਧਨਪਫਬਭਮ'
+        r'ಕಖಗಘಙಚಛಜಝಞಟಠಡಢಣತಥದಧನಪಫಬಭಮ'
+        r'കഖഗഘങചഛജഝഞടഠഡഢണതഥദധനപഫബഭമ'
+        r'କଖଗଘଙଚଛଜଝଞଟଠଡଢଣତଥଦଧନପଫବଭମ'
+        r'కఖగఘఙచఛజఝఞటఠడఢణతథదధనపఫబభమ'
+        r'கஙசஜஞடணதநபம])'
+    )
+    
+    replacement = lambda m: {
+        # Mapping anusvara to conjunct nasal for each script block
+        'ं': {'क': 'ङ्', 'ख': 'ङ्', 'ग': 'ङ्', 'घ': 'ङ्', 'ङ': 'ङ्',
+              'च': 'ञ्', 'छ': 'ञ्', 'ज': 'ञ्', 'झ': 'ञ्', 'ञ': 'ञ्',
+              'ट': 'ण्', 'ठ': 'ण्', 'ड': 'ण्', 'ढ': 'ण्', 'ण': 'ण्',
+              'त': 'न्', 'थ': 'न्', 'द': 'न्', 'ध': 'न्', 'न': 'न्',
+              'प': 'म्', 'फ': 'म्', 'ब': 'म्', 'भ': 'म्', 'म': 'म्'},
+        'ং': {'ক': 'ঙ্', 'খ': 'ঙ্', 'গ': 'ঙ্', 'ঘ': 'ঙ্', 'ঙ': 'ঙ্',
+              'চ': 'ঞ্', 'ছ': 'ঞ্', 'জ': 'ঞ্', 'ঝ': 'ঞ্', 'ঞ': 'ঞ্',
+              'ট': 'ণ্', 'ঠ': 'ণ্', 'ড': 'ণ্', 'ঢ': 'ণ্', 'ণ': 'ণ্',
+              'ত': 'ন্', 'থ': 'ন্', 'দ': 'ন্', 'ধ': 'ন্', 'ন': 'ন্',
+              'প': 'ম্', 'ফ': 'ম্', 'ব': 'ম্', 'ভ': 'ম্', 'ম': 'ম্'},
+        'ં': {'ક': 'ઙ્', 'ખ': 'ઙ્', 'ગ': 'ઙ્', 'ઘ': 'ઙ્', 'ઙ': 'ઙ્',
+              'ચ': 'ઞ્', 'છ': 'ઞ્', 'જ': 'ઞ્', 'ઝ': 'ઞ્', 'ઞ': 'ઞ્',
+              'ટ': 'ણ્', 'ઠ': 'ણ્', 'ડ': 'ણ્', 'ઢ': 'ણ્', 'ણ': 'ણ્',
+              'ત': 'ન્', 'થ': 'ન્', 'દ': 'ન્', 'ધ': 'ન્', 'ન': 'ન્',
+              'પ': 'મ્', 'ફ': 'મ્', 'બ': 'મ્', 'ભ': 'મ્', 'મ': 'મ્'},
+        'ਂ': {'ਕ': 'ਙ੍', 'ਖ': 'ਙ੍', 'ਗ': 'ਙ੍', 'ਘ': 'ਙ੍', 'ਙ': 'ਙ੍',
+              'ਚ': 'ਞ੍', 'ਛ': 'ਞ੍', 'ਜ': 'ਞ੍', 'ਝ': 'ਞ੍', 'ਞ': 'ਞ੍',
+              'ਟ': 'ਣ੍', 'ਠ': 'ਣ੍', 'ਡ': 'ਣ੍', 'ਢ': 'ਣ੍', 'ਣ': 'ਣ੍',
+              'ਤ': 'ਨ੍', 'ਥ': 'ਨ੍', 'ਦ': 'ਨ੍', 'ਧ': 'ਨ੍', 'ਨ': 'ਨ੍',
+              'ਪ': 'ਮ੍', 'ਫ': 'ਮ੍', 'ਬ': 'ਮ੍', 'ਭ': 'ਮ੍', 'ਮ': 'ਮ੍'},
+        'ಂ': {'ಕ': 'ಙ್', 'ಖ': 'ಙ್', 'ಗ': 'ಙ್', 'ಘ': 'ಙ್', 'ಙ': 'ಙ್',
+              'ಚ': 'ಞ್', 'ಛ': 'ಞ್', 'ಜ': 'ಞ್', 'ಝ': 'ಞ್', 'ಞ': 'ಞ್',
+              'ಟ': 'ಣ್', 'ಠ': 'ಣ್', 'ಡ': 'ಣ್', 'ಢ': 'ಣ್', 'ಣ': 'ಣ್',
+              'ತ': 'ನ್', 'ಥ': 'ನ್', 'ದ': 'ನ್', 'ಧ': 'ನ್', 'ನ': 'ನ್',
+              'ಪ': 'ಮ್', 'ಫ': 'ಮ್', 'ಬ': 'ಮ್', 'ಭ': 'ಮ್', 'ಮ': 'ಮ್'},
+        'ം': {'ക': 'ങ്', 'ഖ': 'ങ്', 'ഗ': 'ങ്', 'ഘ': 'ങ്', 'ങ': 'ങ്',
+              'ച': 'ഞ്', 'ഛ': 'ഞ്', 'ജ': 'ഞ്', 'ഝ': 'ഞ്', 'ഞ': 'ഞ്',
+              'ട': 'ണ്', 'ഠ': 'ണ്', 'ഡ': 'ണ്', 'ഢ': 'ണ്', 'ണ': 'ണ',
+              'ത': 'ന്', 'ഥ': 'ന്', 'ദ': 'ന്', 'ധ': 'ന്', 'ന': 'ന്',
+              'പ': 'മ്', 'ഫ': 'മ്', 'ബ': 'മ്', 'ഭ': 'മ്', 'മ': 'മ്'},
+        'ଂ': {'କ': 'ଙ୍', 'ଖ': 'ଙ୍', 'ଗ': 'ଙ୍', 'ଘ': 'ଙ୍', 'ଙ': 'ଙ୍',
+              'ଚ': 'ଞ୍', 'ଛ': 'ଞ୍', 'ଜ': 'ଞ୍', 'ଝ': 'ଞ୍', 'ଞ': 'ଞ୍',
+              'ଟ': 'ଣ୍', 'ଠ': 'ଣ୍', 'ଡ': 'ଣ୍', 'ଢ': 'ଣ୍', 'ଣ': 'ଣ୍',
+              'ତ': 'ନ୍', 'ଥ': 'ନ୍', 'ଦ': 'ନ୍', 'ଧ': 'ନ୍', 'ନ': 'ନ୍',
+              'ପ': 'ମ୍', 'ଫ': 'ମ୍', 'ବ': 'ମ୍', 'ଭ': 'ମ୍', 'ମ': 'ମ୍'},
+        'ం': {'క': 'ఙ్', 'ఖ': 'ఙ్', 'గ': 'ఙ్', 'ఘ': 'ఙ్', 'ఙ': 'ఙ్',
+              'చ': 'ఞ్', 'ఛ': 'ఞ్', 'జ': 'ఞ్', 'ఝ': 'ఞ్', 'ఞ': 'ఞ్',
+              'ట': 'ణ్', 'ఠ': 'ణ్', 'డ': 'ణ్', 'ఢ': 'ణ్', 'ణ': 'ణ్',
+              'త': 'న్', 'థ': 'న్', 'ద': 'న్', 'ధ': 'న్', 'న': 'న్',
+              'ప': 'మ్', 'ఫ': 'మ్', 'బ': 'మ్', 'భ': 'మ్', 'మ': 'మ్'},
+        'ஂ': {'க': 'ங்', 'ங': 'ங்',
+              'ச': 'ஞ்', 'ஜ': 'ஞ்', 'ஞ': 'ஞ்',
+              'ட': 'ண்', 'ண': 'ண்',
+              'த': 'ந்', 'ந': 'ந்',
+              'ப': 'ம்', 'ம': 'ம்'}
+    }[m.group(1)][m.group(2)] + m.group(2)
+
+    return re.sub(pattern, replacement, text)
+
+def process_segment(segment, state, language):
+    if state == "problematic":
+        return list(segment)
+    elif state == "english":
+        return [f"en_{char}" for char in segment]
+    else:
+        try:
+            return wordparse(segment, 0, 0, 1).split()
+        except Exception as e:
+            return list(segment)
+
+def get_cls_token_list(text, language):
+    cls_token_list = []
+    state = "text"
+    if has_non_indic_script(text):
+        raise Exception("Non-indic script found in text.")
+    for word in text.split():
+        segment = ""
+        for char in word:
+            if char in string.ascii_letters:
+                curr_state = "english"
+            elif char in non_problematic_chars:
+                curr_state = "text"
+            else:
+                curr_state = "problematic"
+            if state != curr_state:
+                cls_token_list.extend(process_segment(segment, state, language))
+                segment = ""
+            segment += char
+            state = curr_state
+        if segment:
+            cls_token_list.extend(process_segment(segment, state, language))
+        cls_token_list.append(" ")
+    return cls_token_list[:-1]
+
+def get_cls_for_out_of_mapping(text):
+    cls_token_list = []
+    for word in text.split():
+        for char in word:
+            processed_char = char
+            if char in string.ascii_letters:
+                processed_char = f"en_{char}"
+            cls_token_list.append(processed_char)
+        cls_token_list.append(" ")
+    return cls_token_list[:-1]
+
+def cls_tokenize_text(text: str, language: str):
+    if wordparse is None:
+        raise RuntimeError("indic_unified_parser is required for CLS tokenization but is not installed.")
+    return get_cls_token_list(normalize_indic_nasals(get_transliteration(text.lower(), language)), language)

src/f5_tts/infer/examples/basic/basic.toml

@@ -0,0 +1,11 @@
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/basic/basic_ref_en.wav"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = "Some call me nature, others call me mother nature."
+gen_text = "I don't really care what you call me. I've been a silent spectator, watching species evolve, empires rise and fall. But always remember, I am mighty and enduring."
+# File with text to generate. Ignores the text above.
+gen_file = ""
+remove_silence = false
+output_dir = "tests"
+output_file = "infer_cli_basic.wav"

src/f5_tts/infer/examples/multi/story.toml

@@ -0,0 +1,20 @@
+# F5TTS_v1_Base | E2TTS_Base
+model = "F5TTS_v1_Base"
+ref_audio = "infer/examples/multi/main.flac"
+# If an empty "", transcribes the reference audio automatically.
+ref_text = ""
+gen_text = ""
+# File with text to generate. Ignores the text above.
+gen_file = "infer/examples/multi/story.txt"
+remove_silence = true
+output_dir = "tests"
+output_file = "infer_cli_story.wav"
+
+[voices.town]
+ref_audio = "infer/examples/multi/town.flac"
+ref_text = ""
+speed = 0.8  # will ignore global speed
+
+[voices.country]
+ref_audio = "infer/examples/multi/country.flac"
+ref_text = ""

src/f5_tts/infer/examples/multi/story.txt

@@ -0,0 +1 @@
+A Town Mouse and a Country Mouse were acquaintances, and the Country Mouse one day invited his friend to come and see him at his home in the fields. The Town Mouse came, and they sat down to a dinner of barleycorns and roots, the latter of which had a distinctly earthy flavour. The fare was not much to the taste of the guest, and presently he broke out with [town] "My poor dear friend, you live here no better than the ants! Now, you should just see how I fare! My larder is a regular horn of plenty. You must come and stay with me, and I promise you you shall live on the fat of the land." [main] So when he returned to town he took the Country Mouse with him, and showed him into a larder containing flour and oatmeal and figs and honey and dates. The Country Mouse had never seen anything like it, and sat down to enjoy the luxuries his friend provided: but before they had well begun, the door of the larder opened and someone came in. The two Mice scampered off and hid themselves in a narrow and exceedingly uncomfortable hole. Presently, when all was quiet, they ventured out again; but someone else came in, and off they scuttled again. This was too much for the visitor. [country] "Goodbye," [main] said he, [country] "I'm off. You live in the lap of luxury, I can see, but you are surrounded by dangers; whereas at home I can enjoy my simple dinner of roots and corn in peace."

src/f5_tts/infer/examples/vocab.txt

@@ -0,0 +1,2545 @@
+ 
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+_
+a
+a1
+ai1
+ai2
+ai3
+ai4
+an1
+an3
+an4
+ang1
+ang2
+ang4
+ao1
+ao2
+ao3
+ao4
+b
+ba
+ba1
+ba2
+ba3
+ba4
+bai1
+bai2
+bai3
+bai4
+ban1
+ban2
+ban3
+ban4
+bang1
+bang2
+bang3
+bang4
+bao1
+bao2
+bao3
+bao4
+bei
+bei1
+bei2
+bei3
+bei4
+ben1
+ben2
+ben3
+ben4
+beng
+beng1
+beng2
+beng3
+beng4
+bi1
+bi2
+bi3
+bi4
+bian1
+bian2
+bian3
+bian4
+biao1
+biao2
+biao3
+bie1
+bie2
+bie3
+bie4
+bin1
+bin4
+bing1
+bing2
+bing3
+bing4
+bo
+bo1
+bo2
+bo3
+bo4
+bu2
+bu3
+bu4
+c
+ca1
+cai1
+cai2
+cai3
+cai4
+can1
+can2
+can3
+can4
+cang1
+cang2
+cao1
+cao2
+cao3
+ce4
+cen1
+cen2
+ceng1
+ceng2
+ceng4
+cha1
+cha2
+cha3
+cha4
+chai1
+chai2
+chan1
+chan2
+chan3
+chan4
+chang1
+chang2
+chang3
+chang4
+chao1
+chao2
+chao3
+che1
+che2
+che3
+che4
+chen1
+chen2
+chen3
+chen4
+cheng1
+cheng2
+cheng3
+cheng4
+chi1
+chi2
+chi3
+chi4
+chong1
+chong2
+chong3
+chong4
+chou1
+chou2
+chou3
+chou4
+chu1
+chu2
+chu3
+chu4
+chua1
+chuai1
+chuai2
+chuai3
+chuai4
+chuan1
+chuan2
+chuan3
+chuan4
+chuang1
+chuang2
+chuang3
+chuang4
+chui1
+chui2
+chun1
+chun2
+chun3
+chuo1
+chuo4
+ci1
+ci2
+ci3
+ci4
+cong1
+cong2
+cou4
+cu1
+cu4
+cuan1
+cuan2
+cuan4
+cui1
+cui3
+cui4
+cun1
+cun2
+cun4
+cuo1
+cuo2
+cuo4
+d
+da
+da1
+da2
+da3
+da4
+dai1
+dai2
+dai3
+dai4
+dan1
+dan2
+dan3
+dan4
+dang1
+dang2
+dang3
+dang4
+dao1
+dao2
+dao3
+dao4
+de
+de1
+de2
+dei3
+den4
+deng1
+deng2
+deng3
+deng4
+di1
+di2
+di3
+di4
+dia3
+dian1
+dian2
+dian3
+dian4
+diao1
+diao3
+diao4
+die1
+die2
+die4
+ding1
+ding2
+ding3
+ding4
+diu1
+dong1
+dong3
+dong4
+dou1
+dou2
+dou3
+dou4
+du1
+du2
+du3
+du4
+duan1
+duan2
+duan3
+duan4
+dui1
+dui4
+dun1
+dun3
+dun4
+duo1
+duo2
+duo3
+duo4
+e
+e1
+e2
+e3
+e4
+ei2
+en1
+en4
+er
+er2
+er3
+er4
+f
+fa1
+fa2
+fa3
+fa4
+fan1
+fan2
+fan3
+fan4
+fang1
+fang2
+fang3
+fang4
+fei1
+fei2
+fei3
+fei4
+fen1
+fen2
+fen3
+fen4
+feng1
+feng2
+feng3
+feng4
+fo2
+fou2
+fou3
+fu1
+fu2
+fu3
+fu4
+g
+ga1
+ga2
+ga3
+ga4
+gai1
+gai2
+gai3
+gai4
+gan1
+gan2
+gan3
+gan4
+gang1
+gang2
+gang3
+gang4
+gao1
+gao2
+gao3
+gao4
+ge1
+ge2
+ge3
+ge4
+gei2
+gei3
+gen1
+gen2
+gen3
+gen4
+geng1
+geng3
+geng4
+gong1
+gong3
+gong4
+gou1
+gou2
+gou3
+gou4
+gu
+gu1
+gu2
+gu3
+gu4
+gua1
+gua2
+gua3
+gua4
+guai1
+guai2
+guai3
+guai4
+guan1
+guan2
+guan3
+guan4
+guang1
+guang2
+guang3
+guang4
+gui1
+gui2
+gui3
+gui4
+gun3
+gun4
+guo1
+guo2
+guo3
+guo4
+h
+ha1
+ha2
+ha3
+hai1
+hai2
+hai3
+hai4
+han1
+han2
+han3
+han4
+hang1
+hang2
+hang4
+hao1
+hao2
+hao3
+hao4
+he1
+he2
+he4
+hei1
+hen2
+hen3
+hen4
+heng1
+heng2
+heng4
+hong1
+hong2
+hong3
+hong4
+hou1
+hou2
+hou3
+hou4
+hu1
+hu2
+hu3
+hu4
+hua1
+hua2
+hua4
+huai2
+huai4
+huan1
+huan2
+huan3
+huan4
+huang1
+huang2
+huang3
+huang4
+hui1
+hui2
+hui3
+hui4
+hun1
+hun2
+hun4
+huo
+huo1
+huo2
+huo3
+huo4
+i
+j
+ji1
+ji2
+ji3
+ji4
+jia
+jia1
+jia2
+jia3
+jia4
+jian1
+jian2
+jian3
+jian4
+jiang1
+jiang2
+jiang3
+jiang4
+jiao1
+jiao2
+jiao3
+jiao4
+jie1
+jie2
+jie3
+jie4
+jin1
+jin2
+jin3
+jin4
+jing1
+jing2
+jing3
+jing4
+jiong3
+jiu1
+jiu2
+jiu3
+jiu4
+ju1
+ju2
+ju3
+ju4
+juan1
+juan2
+juan3
+juan4
+jue1
+jue2
+jue4
+jun1
+jun4
+k
+ka1
+ka2
+ka3
+kai1
+kai2
+kai3
+kai4
+kan1
+kan2
+kan3
+kan4
+kang1
+kang2
+kang4
+kao1
+kao2
+kao3
+kao4
+ke1
+ke2
+ke3
+ke4
+ken3
+keng1
+kong1
+kong3
+kong4
+kou1
+kou2
+kou3
+kou4
+ku1
+ku2
+ku3
+ku4
+kua1
+kua3
+kua4
+kuai3
+kuai4
+kuan1
+kuan2
+kuan3
+kuang1
+kuang2
+kuang4
+kui1
+kui2
+kui3
+kui4
+kun1
+kun3
+kun4
+kuo4
+l
+la
+la1
+la2
+la3
+la4
+lai2
+lai4
+lan2
+lan3
+lan4
+lang1
+lang2
+lang3
+lang4
+lao1
+lao2
+lao3
+lao4
+le
+le1
+le4
+lei
+lei1
+lei2
+lei3
+lei4
+leng1
+leng2
+leng3
+leng4
+li
+li1
+li2
+li3
+li4
+lia3
+lian2
+lian3
+lian4
+liang2
+liang3
+liang4
+liao1
+liao2
+liao3
+liao4
+lie1
+lie2
+lie3
+lie4
+lin1
+lin2
+lin3
+lin4
+ling2
+ling3
+ling4
+liu1
+liu2
+liu3
+liu4
+long1
+long2
+long3
+long4
+lou1
+lou2
+lou3
+lou4
+lu1
+lu2
+lu3
+lu4
+luan2
+luan3
+luan4
+lun1
+lun2
+lun4
+luo1
+luo2
+luo3
+luo4
+lv2
+lv3
+lv4
+lve3
+lve4
+m
+ma
+ma1
+ma2
+ma3
+ma4
+mai2
+mai3
+mai4
+man1
+man2
+man3
+man4
+mang2
+mang3
+mao1
+mao2
+mao3
+mao4
+me
+mei2
+mei3
+mei4
+men
+men1
+men2
+men4
+meng
+meng1
+meng2
+meng3
+meng4
+mi1
+mi2
+mi3
+mi4
+mian2
+mian3
+mian4
+miao1
+miao2
+miao3
+miao4
+mie1
+mie4
+min2
+min3
+ming2
+ming3
+ming4
+miu4
+mo1
+mo2
+mo3
+mo4
+mou1
+mou2
+mou3
+mu2
+mu3
+mu4
+n
+n2
+na1
+na2
+na3
+na4
+nai2
+nai3
+nai4
+nan1
+nan2
+nan3
+nan4
+nang1
+nang2
+nang3
+nao1
+nao2
+nao3
+nao4
+ne
+ne2
+ne4
+nei3
+nei4
+nen4
+neng2
+ni1
+ni2
+ni3
+ni4
+nian1
+nian2
+nian3
+nian4
+niang2
+niang4
+niao2
+niao3
+niao4
+nie1
+nie4
+nin2
+ning2
+ning3
+ning4
+niu1
+niu2
+niu3
+niu4
+nong2
+nong4
+nou4
+nu2
+nu3
+nu4
+nuan3
+nuo2
+nuo4
+nv2
+nv3
+nve4
+o
+o1
+o2
+ou1
+ou2
+ou3
+ou4
+p
+pa1
+pa2
+pa4
+pai1
+pai2
+pai3
+pai4
+pan1
+pan2
+pan4
+pang1
+pang2
+pang4
+pao1
+pao2
+pao3
+pao4
+pei1
+pei2
+pei4
+pen1
+pen2
+pen4
+peng1
+peng2
+peng3
+peng4
+pi1
+pi2
+pi3
+pi4
+pian1
+pian2
+pian4
+piao1
+piao2
+piao3
+piao4
+pie1
+pie2
+pie3
+pin1
+pin2
+pin3
+pin4
+ping1
+ping2
+po1
+po2
+po3
+po4
+pou1
+pu1
+pu2
+pu3
+pu4
+q
+qi1
+qi2
+qi3
+qi4
+qia1
+qia3
+qia4
+qian1
+qian2
+qian3
+qian4
+qiang1
+qiang2
+qiang3
+qiang4
+qiao1
+qiao2
+qiao3
+qiao4
+qie1
+qie2
+qie3
+qie4
+qin1
+qin2
+qin3
+qin4
+qing1
+qing2
+qing3
+qing4
+qiong1
+qiong2
+qiu1
+qiu2
+qiu3
+qu1
+qu2
+qu3
+qu4
+quan1
+quan2
+quan3
+quan4
+que1
+que2
+que4
+qun2
+r
+ran2
+ran3
+rang1
+rang2
+rang3
+rang4
+rao2
+rao3
+rao4
+re2
+re3
+re4
+ren2
+ren3
+ren4
+reng1
+reng2
+ri4
+rong1
+rong2
+rong3
+rou2
+rou4
+ru2
+ru3
+ru4
+ruan2
+ruan3
+rui3
+rui4
+run4
+ruo4
+s
+sa1
+sa2
+sa3
+sa4
+sai1
+sai4
+san1
+san2
+san3
+san4
+sang1
+sang3
+sang4
+sao1
+sao2
+sao3
+sao4
+se4
+sen1
+seng1
+sha1
+sha2
+sha3
+sha4
+shai1
+shai2
+shai3
+shai4
+shan1
+shan3
+shan4
+shang
+shang1
+shang3
+shang4
+shao1
+shao2
+shao3
+shao4
+she1
+she2
+she3
+she4
+shei2
+shen1
+shen2
+shen3
+shen4
+sheng1
+sheng2
+sheng3
+sheng4
+shi
+shi1
+shi2
+shi3
+shi4
+shou1
+shou2
+shou3
+shou4
+shu1
+shu2
+shu3
+shu4
+shua1
+shua2
+shua3
+shua4
+shuai1
+shuai3
+shuai4
+shuan1
+shuan4
+shuang1
+shuang3
+shui2
+shui3
+shui4
+shun3
+shun4
+shuo1
+shuo4
+si1
+si2
+si3
+si4
+song1
+song3
+song4
+sou1
+sou3
+sou4
+su1
+su2
+su4
+suan1
+suan4
+sui1
+sui2
+sui3
+sui4
+sun1
+sun3
+suo
+suo1
+suo2
+suo3
+t
+ta1
+ta2
+ta3
+ta4
+tai1
+tai2
+tai4
+tan1
+tan2
+tan3
+tan4
+tang1
+tang2
+tang3
+tang4
+tao1
+tao2
+tao3
+tao4
+te4
+teng2
+ti1
+ti2
+ti3
+ti4
+tian1
+tian2
+tian3
+tiao1
+tiao2
+tiao3
+tiao4
+tie1
+tie2
+tie3
+tie4
+ting1
+ting2
+ting3
+tong1
+tong2
+tong3
+tong4
+tou
+tou1
+tou2
+tou4
+tu1
+tu2
+tu3
+tu4
+tuan1
+tuan2
+tui1
+tui2
+tui3
+tui4
+tun1
+tun2
+tun4
+tuo1
+tuo2
+tuo3
+tuo4
+u
+v
+w
+wa
+wa1
+wa2
+wa3
+wa4
+wai1
+wai3
+wai4
+wan1
+wan2
+wan3
+wan4
+wang1
+wang2
+wang3
+wang4
+wei1
+wei2
+wei3
+wei4
+wen1
+wen2
+wen3
+wen4
+weng1
+weng4
+wo1
+wo2
+wo3
+wo4
+wu1
+wu2
+wu3
+wu4
+x
+xi1
+xi2
+xi3
+xi4
+xia1
+xia2
+xia4
+xian1
+xian2
+xian3
+xian4
+xiang1
+xiang2
+xiang3
+xiang4
+xiao1
+xiao2
+xiao3
+xiao4
+xie1
+xie2
+xie3
+xie4
+xin1
+xin2
+xin4
+xing1
+xing2
+xing3
+xing4
+xiong1
+xiong2
+xiu1
+xiu3
+xiu4
+xu
+xu1
+xu2
+xu3
+xu4
+xuan1
+xuan2
+xuan3
+xuan4
+xue1
+xue2
+xue3
+xue4
+xun1
+xun2
+xun4
+y
+ya
+ya1
+ya2
+ya3
+ya4
+yan1
+yan2
+yan3
+yan4
+yang1
+yang2
+yang3
+yang4
+yao1
+yao2
+yao3
+yao4
+ye1
+ye2
+ye3
+ye4
+yi
+yi1
+yi2
+yi3
+yi4
+yin1
+yin2
+yin3
+yin4
+ying1
+ying2
+ying3
+ying4
+yo1
+yong1
+yong2
+yong3
+yong4
+you1
+you2
+you3
+you4
+yu1
+yu2
+yu3
+yu4
+yuan1
+yuan2
+yuan3
+yuan4
+yue1
+yue4
+yun1
+yun2
+yun3
+yun4
+z
+za1
+za2
+za3
+zai1
+zai3
+zai4
+zan1
+zan2
+zan3
+zan4
+zang1
+zang4
+zao1
+zao2
+zao3
+zao4
+ze2
+ze4
+zei2
+zen3
+zeng1
+zeng4
+zha1
+zha2
+zha3
+zha4
+zhai1
+zhai2
+zhai3
+zhai4
+zhan1
+zhan2
+zhan3
+zhan4
+zhang1
+zhang2
+zhang3
+zhang4
+zhao1
+zhao2
+zhao3
+zhao4
+zhe
+zhe1
+zhe2
+zhe3
+zhe4
+zhen1
+zhen2
+zhen3
+zhen4
+zheng1
+zheng2
+zheng3
+zheng4
+zhi1
+zhi2
+zhi3
+zhi4
+zhong1
+zhong2
+zhong3
+zhong4
+zhou1
+zhou2
+zhou3
+zhou4
+zhu1
+zhu2
+zhu3
+zhu4
+zhua1
+zhua2
+zhua3
+zhuai1
+zhuai3
+zhuai4
+zhuan1
+zhuan2
+zhuan3
+zhuan4
+zhuang1
+zhuang4
+zhui1
+zhui4
+zhun1
+zhun2
+zhun3
+zhuo1
+zhuo2
+zi
+zi1
+zi2
+zi3
+zi4
+zong1
+zong2
+zong3
+zong4
+zou1
+zou2
+zou3
+zou4
+zu1
+zu2
+zu3
+zuan1
+zuan3
+zuan4
+zui2
+zui3
+zui4
+zun1
+zuo
+zuo1
+zuo2
+zuo3
+zuo4
+{
+~
+¡
+¢
+£
+¥
+§
+¨
+©
+«
+®
+¯
+°
+±
+²
+³
+´
+µ
+·
+¹
+º
+»
+¼
+½
+¾
+¿
+À
+Á
+Â
+Ã
+Ä
+Å
+Æ
+Ç
+È
+É
+Ê
+Í
+Î
+Ñ
+Ó
+Ö
+×
+Ø
+Ú
+Ü
+Ý
+Þ
+ß
+à
+á
+â
+ã
+ä
+å
+æ
+ç
+è
+é
+ê
+ë
+ì
+í
+î
+ï
+ð
+ñ
+ò
+ó
+ô
+õ
+ö
+ø
+ù
+ú
+û
+ü
+ý
+Ā
+ā
+ă
+ą
+ć
+Č
+č
+Đ
+đ
+ē
+ė
+ę
+ě
+ĝ
+ğ
+ħ
+ī
+į
+İ
+ı
+Ł
+ł
+ń
+ņ
+ň
+ŋ
+Ō
+ō
+ő
+œ
+ř
+Ś
+ś
+Ş
+ş
+Š
+š
+Ť
+ť
+ũ
+ū
+ź
+Ż
+ż
+Ž
+ž
+ơ
+ư
+ǎ
+ǐ
+ǒ
+ǔ
+ǚ
+ș
+ț
+ɑ
+ɔ
+ɕ
+ə
+ɛ
+ɜ
+ɡ
+ɣ
+ɪ
+ɫ
+ɴ
+ɹ
+ɾ
+ʃ
+ʊ
+ʌ
+ʒ
+ʔ
+ʰ
+ʷ
+ʻ
+ʾ
+ʿ
+ˈ
+ː
+˙
+˜
+ˢ
+́
+̅
+Α
+Β
+Δ
+Ε
+Θ
+Κ
+Λ
+Μ
+Ξ
+Π
+Σ
+Τ
+Φ
+Χ
+Ψ
+Ω
+ά
+έ
+ή
+ί
+α
+β
+γ
+δ
+ε
+ζ
+η
+θ
+ι
+κ
+λ
+μ
+ν
+ξ
+ο
+π
+ρ
+ς
+σ
+τ
+υ
+φ
+χ
+ψ
+ω
+ϊ
+ό
+ύ
+ώ
+ϕ
+ϵ
+Ё
+А
+Б
+В
+Г
+Д
+Е
+Ж
+З
+И
+Й
+К
+Л
+М
+Н
+О
+П
+Р
+С
+Т
+У
+Ф
+Х
+Ц
+Ч
+Ш
+Щ
+Ы
+Ь
+Э
+Ю
+Я
+а
+б
+в
+г
+д
+е
+ж
+з
+и
+й
+к
+л
+м
+н
+о
+п
+р
+с
+т
+у
+ф
+х
+ц
+ч
+ш
+щ
+ъ
+ы
+ь
+э
+ю
+я
+ё
+і
+ְ
+ִ
+ֵ
+ֶ
+ַ
+ָ
+ֹ
+ּ
+־
+ׁ
+א
+ב
+ג
+ד
+ה
+ו
+ז
+ח
+ט
+י
+כ
+ל
+ם
+מ
+ן
+נ
+ס
+ע
+פ
+ק
+ר
+ש
+ת
+أ
+ب
+ة
+ت
+ج
+ح
+د
+ر
+ز
+س
+ص
+ط
+ع
+ق
+ك
+ل
+م
+ن
+ه
+و
+ي
+َ
+ُ
+ِ
+ْ
+ก
+ข
+ง
+จ
+ต
+ท
+น
+ป
+ย
+ร
+ว
+ส
+ห
+อ
+ฮ
+ั
+า
+ี
+ึ
+โ
+ใ
+ไ
+่
+้
+์
+ḍ
+Ḥ
+ḥ
+ṁ
+ṃ
+ṅ
+ṇ
+Ṛ
+ṛ
+Ṣ
+ṣ
+Ṭ
+ṭ
+ạ
+ả
+Ấ
+ấ
+ầ
+ậ
+ắ
+ằ
+ẻ
+ẽ
+ế
+ề
+ể
+ễ
+ệ
+ị
+ọ
+ỏ
+ố
+ồ
+ộ
+ớ
+ờ
+ở
+ụ
+ủ
+ứ
+ữ
+ἀ
+ἁ
+Ἀ
+ἐ
+ἔ
+ἰ
+ἱ
+ὀ
+ὁ
+ὐ
+ὲ
+ὸ
+���
+᾽
+ῆ
+ῇ
+ῶ
+‎
+‑
+‒
+–
+—
+―
+‖
+†
+‡
+•
+…
+‧
+‬
+′
+″
+⁄
+⁡
+⁰
+⁴
+⁵
+⁶
+⁷
+⁸
+⁹
+₁
+₂
+₃
+€
+₱
+₹
+₽
+℃
+ℏ
+ℓ
+№
+ℝ
+™
+⅓
+⅔
+⅛
+→
+∂
+∈
+∑
+−
+∗
+√
+∞
+∫
+≈
+≠
+≡
+≤
+≥
+⋅
+⋯
+█
+♪
+⟨
+⟩
+、
+。
+《
+》
+「
+」
+【
+】
+あ
+う
+え
+お
+か
+が
+き
+ぎ
+く
+ぐ
+け
+げ
+こ
+ご
+さ
+し
+じ
+す
+ず
+せ
+ぜ
+そ
+ぞ
+た
+だ
+ち
+っ
+つ
+で
+と
+ど
+な
+に
+ね
+の
+は
+ば
+ひ
+ぶ
+へ
+べ
+ま
+み
+む
+め
+も
+ゃ
+や
+ゆ
+ょ
+よ
+ら
+り
+る
+れ
+ろ
+わ
+を
+ん
+ァ
+ア
+ィ
+イ
+ウ
+ェ
+エ
+オ
+カ
+ガ
+キ
+ク
+ケ
+ゲ
+コ
+ゴ
+サ
+ザ
+シ
+ジ
+ス
+ズ
+セ
+ゾ
+タ
+ダ
+チ
+ッ
+ツ
+テ
+デ
+ト
+ド
+ナ
+ニ
+ネ
+ノ
+バ
+パ
+ビ
+ピ
+フ
+プ
+ヘ
+ベ
+ペ
+ホ
+ボ
+ポ
+マ
+ミ
+ム
+メ
+モ
+ャ
+ヤ
+ュ
+ユ
+ョ
+ヨ
+ラ
+リ
+ル
+レ
+ロ
+ワ
+ン
+・
+ー
+ㄋ
+ㄍ
+ㄎ
+ㄏ
+ㄓ
+ㄕ
+ㄚ
+ㄜ
+ㄟ
+ㄤ
+ㄥ
+ㄧ
+ㄱ
+ㄴ
+ㄷ
+ㄹ
+ㅁ
+ㅂ
+ㅅ
+ㅈ
+ㅍ
+ㅎ
+ㅏ
+ㅓ
+ㅗ
+ㅜ
+ㅡ
+ㅣ
+㗎
+가
+각
+간
+갈
+감
+갑
+갓
+갔
+강
+같
+개
+거
+건
+걸
+겁
+것
+겉
+게
+겠
+겨
+결
+겼
+경
+계
+고
+곤
+골
+곱
+공
+과
+관
+광
+교
+구
+국
+굴
+귀
+귄
+그
+근
+글
+금
+기
+긴
+길
+까
+깍
+깔
+깜
+깨
+께
+꼬
+꼭
+꽃
+꾸
+꿔
+끔
+끗
+끝
+끼
+나
+난
+날
+남
+납
+내
+냐
+냥
+너
+넘
+넣
+네
+녁
+년
+녕
+노
+녹
+놀
+누
+눈
+느
+는
+늘
+니
+님
+닙
+다
+닥
+단
+달
+닭
+당
+대
+더
+덕
+던
+덥
+데
+도
+독
+동
+돼
+됐
+되
+된
+될
+두
+둑
+둥
+드
+들
+등
+디
+따
+딱
+딸
+땅
+때
+떤
+떨
+떻
+또
+똑
+뚱
+뛰
+뜻
+띠
+라
+락
+란
+람
+랍
+랑
+래
+랜
+러
+런
+럼
+렇
+레
+려
+력
+렵
+렸
+로
+록
+롬
+루
+르
+른
+를
+름
+릉
+리
+릴
+림
+마
+막
+만
+많
+말
+맑
+맙
+맛
+매
+머
+먹
+멍
+메
+면
+명
+몇
+모
+목
+몸
+못
+무
+문
+물
+뭐
+뭘
+미
+민
+밌
+밑
+바
+박
+밖
+반
+받
+발
+밤
+밥
+방
+배
+백
+밸
+뱀
+버
+번
+벌
+벚
+베
+벼
+벽
+별
+병
+보
+복
+본
+볼
+봐
+봤
+부
+분
+불
+비
+빔
+빛
+빠
+빨
+뼈
+뽀
+뿅
+쁘
+사
+산
+살
+삼
+샀
+상
+새
+색
+생
+서
+선
+설
+섭
+섰
+성
+세
+셔
+션
+셨
+소
+속
+손
+송
+수
+숙
+순
+술
+숫
+숭
+숲
+쉬
+쉽
+스
+슨
+습
+슷
+시
+식
+신
+실
+싫
+심
+십
+싶
+싸
+써
+쓰
+쓴
+씌
+씨
+씩
+씬
+아
+악
+안
+않
+알
+야
+약
+얀
+양
+얘
+어
+언
+얼
+엄
+업
+없
+었
+엉
+에
+여
+역
+연
+염
+엽
+영
+옆
+예
+옛
+오
+온
+올
+옷
+옹
+와
+왔
+왜
+요
+욕
+용
+우
+운
+울
+웃
+워
+원
+월
+웠
+위
+윙
+유
+육
+윤
+으
+은
+을
+음
+응
+의
+이
+익
+인
+일
+읽
+임
+입
+있
+자
+작
+잔
+잖
+잘
+잡
+잤
+장
+재
+저
+전
+점
+정
+제
+져
+졌
+조
+족
+좀
+종
+좋
+죠
+주
+준
+줄
+중
+줘
+즈
+즐
+즘
+지
+진
+집
+짜
+짝
+쩌
+쪼
+쪽
+쫌
+쭈
+쯔
+찌
+찍
+차
+착
+찾
+책
+처
+천
+철
+체
+쳐
+쳤
+초
+촌
+추
+출
+춤
+춥
+춰
+치
+친
+칠
+침
+칩
+칼
+커
+켓
+코
+콩
+쿠
+퀴
+크
+큰
+큽
+키
+킨
+타
+태
+터
+턴
+털
+테
+토
+통
+투
+트
+특
+튼
+틀
+티
+팀
+파
+팔
+패
+페
+펜
+펭
+평
+포
+폭
+표
+품
+풍
+프
+플
+피
+필
+하
+학
+한
+할
+함
+합
+항
+해
+햇
+했
+행
+허
+험
+형
+혜
+호
+혼
+홀
+화
+회
+획
+후
+휴
+흐
+흔
+희
+히
+힘
+ﷺ
+ﷻ
+！
+，
+？
+�
+𠮶

src/f5_tts/infer/hf_infer.py

@@ -0,0 +1,98 @@
+"""Inference helpers intended for Hugging Face usage (no HTTP server required)."""
+
+from __future__ import annotations
+
+import os
+from functools import lru_cache
+from typing import Any, Callable, Optional, Tuple
+
+from f5_tts.api import F5TTS
+
+ENV_DEFAULTS = {
+    "model": os.environ.get("F5TTS_MODEL", "F5TTS_v1_Base"),
+    "ckpt_file": os.environ.get("F5TTS_CKPT", ""),
+    "vocab_file": os.environ.get("F5TTS_VOCAB", ""),
+    "ode_method": os.environ.get("F5TTS_ODE_METHOD", "euler"),
+    "use_ema": os.environ.get("F5TTS_USE_EMA", "true").lower() != "false",
+    "vocoder_local_path": os.environ.get("F5TTS_VOCODER_PATH"),
+    "device": os.environ.get("F5TTS_DEVICE"),
+    "hf_cache_dir": os.environ.get("F5TTS_HF_CACHE_DIR"),
+}
+
+
+@lru_cache(maxsize=2)
+def load_tts(
+    model: str = ENV_DEFAULTS["model"],
+    ckpt_file: str = ENV_DEFAULTS["ckpt_file"],
+    vocab_file: str = ENV_DEFAULTS["vocab_file"],
+    ode_method: str = ENV_DEFAULTS["ode_method"],
+    use_ema: bool = ENV_DEFAULTS["use_ema"],
+    vocoder_local_path: Optional[str] = ENV_DEFAULTS["vocoder_local_path"],
+    device: Optional[str] = ENV_DEFAULTS["device"],
+    hf_cache_dir: Optional[str] = ENV_DEFAULTS["hf_cache_dir"],
+) -> F5TTS:
+    """Load and cache an F5TTS model for inference."""
+    return F5TTS(
+        model=model,
+        ckpt_file=ckpt_file,
+        vocab_file=vocab_file,
+        ode_method=ode_method,
+        use_ema=use_ema,
+        vocoder_local_path=vocoder_local_path,
+        device=device,
+        hf_cache_dir=hf_cache_dir,
+    )
+
+
+def synthesize(
+    tts: F5TTS,
+    ref_audio_path: str,
+    ref_text: str,
+    gen_text: str,
+    *,
+    target_rms: float = 0.1,
+    cross_fade_duration: float = 0.15,
+    sway_sampling_coef: float = -1.0,
+    cfg_strength: float = 2.0,
+    nfe_step: int = 32,
+    speed: float = 1.0,
+    fix_duration: Optional[float] = None,
+    remove_silence: bool = False,
+    seed: Optional[int] = None,
+    tokenizer: str = "pinyin",
+    cls_language: Optional[str] = None,
+    cls_tokenizer_fn: Optional[Callable[[str, str], list]] = None,
+    cls_server_url: Optional[str] = None,
+    cls_timeout: float = 5.0,
+    file_wave: Optional[str] = None,
+    file_spec: Optional[str] = None,
+    show_info=None,
+    progress=None,
+) -> Tuple[Any, int, Optional[Any]]:
+    """Run inference and return (wav, sample_rate, spectrogram)."""
+    if show_info is None:
+        show_info = lambda *args, **kwargs: None
+
+    return tts.infer(
+        ref_file=ref_audio_path,
+        ref_text=ref_text,
+        gen_text=gen_text,
+        show_info=show_info,
+        progress=progress,
+        target_rms=target_rms,
+        cross_fade_duration=cross_fade_duration,
+        sway_sampling_coef=sway_sampling_coef,
+        cfg_strength=cfg_strength,
+        nfe_step=nfe_step,
+        speed=speed,
+        fix_duration=fix_duration,
+        remove_silence=remove_silence,
+        file_wave=file_wave,
+        file_spec=file_spec,
+        seed=seed,
+        tokenizer=tokenizer,
+        cls_language=cls_language,
+        cls_tokenizer_fn=cls_tokenizer_fn,
+        cls_server_url=cls_server_url,
+        cls_timeout=cls_timeout,
+    )

src/f5_tts/infer/infer_api.py

@@ -0,0 +1,367 @@
+"""FastAPI server for F5-TTS inference.
+
+Launch with a custom checkpoint:
+    python src/f5_tts/infer/infer_api.py --ckpt-file ckpts/my_model.safetensors --vocab-file ckpts/vocab.txt
+
+The API exposes:
+- GET  /health         -> basic readiness info
+- POST /v1/tts         -> synthesize speech (JSON body)
+"""
+
+import base64
+import io
+import os
+import tempfile
+import threading
+from functools import lru_cache
+from typing import Optional
+
+import click
+import soundfile as sf
+import uvicorn
+from fastapi import FastAPI, HTTPException, Query
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field, model_validator
+
+from f5_tts.api import F5TTS
+from f5_tts.infer.utils_infer import save_spectrogram
+
+# Allow configuration through environment variables for quick overrides
+ENV_DEFAULTS = {
+    "model": os.environ.get("F5TTS_API_MODEL", "F5TTS_v1_Base"),
+    "ckpt_file": os.environ.get(
+        "F5TTS_API_CKPT",
+        "/workspace/personal/team_folders/F5-TTS-common/ckpts/F5TTS_v1_Base_vocos_cls_speech_db_wer_filtered_12_langs_train_finetune_cls/"
+        "model_1250000.pt",
+    ),
+    "vocab_file": os.environ.get(
+        "F5TTS_API_VOCAB",
+        "/workspace/personal/team_folders/F5-TTS-common/ckpts/F5TTS_v1_Base_vocos_cls_speech_db_wer_filtered_12_langs_train_finetune_cls/"
+        "vocab.txt",
+    ),
+    "ode_method": os.environ.get("F5TTS_API_ODE_METHOD", "euler"),
+    "use_ema": os.environ.get("F5TTS_API_USE_EMA", "true").lower() != "false",
+    "vocoder_local_path": os.environ.get("F5TTS_API_VOCODER_PATH"),
+    "device": os.environ.get("F5TTS_API_DEVICE"),
+    "hf_cache_dir": os.environ.get("F5TTS_API_HF_CACHE_DIR"),
+    "en_model": os.environ.get("F5TTS_API_EN_MODEL", os.environ.get("F5TTS_API_MODEL", "F5TTS_v1_Base")),
+    "en_ckpt_file": os.environ.get(
+        "F5TTS_API_EN_CKPT",
+        "/workspace/personal/team_folders/vansh.pundir/F5-TTS/ckpts/"
+        "F5TTS_v1_Base_12_lang_vocos_char_speech_db_only_TTS_12_langs_eval_v3_char_dedup_validation/"
+        "model_550000.pt",
+    ),
+    "en_vocab_file": os.environ.get(
+        "F5TTS_API_EN_VOCAB",
+        "/workspace/personal/team_folders/vansh.pundir/F5-TTS/ckpts/"
+        "F5TTS_v1_Base_12_lang_vocos_char_speech_db_only_TTS_12_langs_eval_v3_char_dedup_validation/"
+        "vocab.txt",
+    ),
+    "cls_url": os.environ.get("F5TTS_CLS_URL", "http://localhost:8061/process"),
+    "cls_timeout": float(os.environ.get("F5TTS_CLS_TIMEOUT", "5.0")),
+}
+
+
+class InferenceRequest(BaseModel):
+    ref_audio_path: Optional[str] = Field(
+        default=None, description="Path to reference audio reachable by the server."
+    )
+    ref_audio_base64: Optional[str] = Field(
+        default=None, description="Base64-encoded reference audio (recommended: WAV/FLAC)."
+    )
+    ref_text: str = Field(
+        default="",
+        description="Transcript of the reference audio. Leave blank to auto-transcribe (requires ASR).",
+    )
+    gen_text: str = Field(..., description="Text to synthesize.")
+    target_rms: float = Field(default=0.1, description="Minimum RMS applied to reference audio.")
+    cross_fade_duration: float = Field(default=0.15, description="Seconds to overlap between chunks.")
+    sway_sampling_coef: float = Field(default=-1.0, description="Sway sampling coefficient.")
+    cfg_strength: float = Field(default=2.0, description="Classifier-free guidance strength.")
+    nfe_step: int = Field(default=32, description="Number of function evaluations.")
+    speed: float = Field(default=1.0, description="Generation speed multiplier.")
+    fix_duration: Optional[float] = Field(
+        default=None, description="Force output duration (seconds). Leave None for automatic."
+    )
+    remove_silence: bool = Field(default=False, description="Remove leading/trailing silence from output.")
+    seed: Optional[int] = Field(default=None, description="Set for deterministic output.")
+    return_spectrogram: bool = Field(default=False, description="Also return spectrogram as base64 PNG.")
+    tokenizer: Optional[str] = Field(
+        default=None,
+        description="Optional tokenizer override: char | cls | pinyin. If omitted, uses legacy pinyin behavior.",
+    )
+    cls_language: Optional[str] = Field(
+        default=None,
+        description="CLS language name (e.g., hindi, english). Used only when tokenizer=cls.",
+    )
+
+    @model_validator(mode="after")
+    def ensure_audio_source(self):
+        if not self.ref_audio_path and not self.ref_audio_base64:
+            raise ValueError("Provide either ref_audio_path or ref_audio_base64.")
+        if not self.gen_text or not self.gen_text.strip():
+            raise ValueError("gen_text cannot be empty.")
+        return self
+
+
+def _encode_wav_base64(wav, sample_rate: int) -> str:
+    """Encode waveform to a base64 WAV string."""
+    with io.BytesIO() as buffer:
+        sf.write(buffer, wav, sample_rate, format="WAV")
+        return base64.b64encode(buffer.getvalue()).decode("ascii")
+
+
+def _encode_spec_base64(spec) -> str:
+    """Save spectrogram to a temp file and encode it as base64 PNG."""
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
+        tmp_path = tmp.name
+    try:
+        save_spectrogram(spec, tmp_path)
+        with open(tmp_path, "rb") as img:
+            return base64.b64encode(img.read()).decode("ascii")
+    finally:
+        os.remove(tmp_path)
+
+
+def _write_temp_audio(data: bytes) -> str:
+    """Persist uploaded audio bytes to a temp file for downstream processing."""
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        tmp.write(data)
+        return tmp.name
+
+
+@lru_cache(maxsize=4)
+def _load_model(
+    model: str = ENV_DEFAULTS["model"],
+    ckpt_file: str = ENV_DEFAULTS["ckpt_file"],
+    vocab_file: str = ENV_DEFAULTS["vocab_file"],
+    ode_method: str = ENV_DEFAULTS["ode_method"],
+    use_ema: bool = ENV_DEFAULTS["use_ema"],
+    vocoder_local_path: Optional[str] = ENV_DEFAULTS["vocoder_local_path"],
+    device: Optional[str] = ENV_DEFAULTS["device"],
+    hf_cache_dir: Optional[str] = ENV_DEFAULTS["hf_cache_dir"],
+):
+    """Cache TTS models by configuration to avoid reloading across requests."""
+    return F5TTS(
+        model=model,
+        ckpt_file=ckpt_file,
+        vocab_file=vocab_file,
+        ode_method=ode_method,
+        use_ema=use_ema,
+        vocoder_local_path=vocoder_local_path,
+        device=device,
+        hf_cache_dir=hf_cache_dir,
+    )
+
+
+def create_app(
+    model: str = ENV_DEFAULTS["model"],
+    ckpt_file: str = ENV_DEFAULTS["ckpt_file"],
+    vocab_file: str = ENV_DEFAULTS["vocab_file"],
+    en_model: str = ENV_DEFAULTS["en_model"],
+    en_ckpt_file: str = ENV_DEFAULTS["en_ckpt_file"],
+    en_vocab_file: str = ENV_DEFAULTS["en_vocab_file"],
+    ode_method: str = ENV_DEFAULTS["ode_method"],
+    use_ema: bool = ENV_DEFAULTS["use_ema"],
+    vocoder_local_path: Optional[str] = ENV_DEFAULTS["vocoder_local_path"],
+    device: Optional[str] = ENV_DEFAULTS["device"],
+    hf_cache_dir: Optional[str] = ENV_DEFAULTS["hf_cache_dir"],
+):
+    """Build a FastAPI app wired to a single F5TTS instance."""
+    tts_hi = _load_model(
+        model=model,
+        ckpt_file=ckpt_file,
+        vocab_file=vocab_file,
+        ode_method=ode_method,
+        use_ema=use_ema,
+        vocoder_local_path=vocoder_local_path,
+        device=device,
+        hf_cache_dir=hf_cache_dir,
+    )
+    infer_lock_hi = threading.Lock()
+    infer_lock_en = threading.Lock()
+
+    app = FastAPI(title="F5-TTS API", version="1.0")
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    @app.get("/health")
+    def health():
+        return {
+            "status": "ok",
+            "device": tts_hi.device,
+            "mel_spec_type": tts_hi.mel_spec_type,
+            "use_ema": tts_hi.use_ema,
+            "supported_langs": ["hi", "en"],
+        }
+
+    @app.post("/v1/tts")
+    def infer(payload: InferenceRequest, lang: str = Query("hi", description="Language code: hi|en")):
+        lang_key = (lang or "hi").strip().lower()
+        if lang_key == "hi":
+            tts = tts_hi
+            infer_lock = infer_lock_hi
+        elif lang_key == "en":
+            tts = _load_model(
+                model=en_model,
+                ckpt_file=en_ckpt_file,
+                vocab_file=en_vocab_file,
+                ode_method=ode_method,
+                use_ema=use_ema,
+                vocoder_local_path=vocoder_local_path,
+                device=device,
+                hf_cache_dir=hf_cache_dir,
+            )
+            infer_lock = infer_lock_en
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Unsupported lang '{lang}'. Use 'hi' for Hindi or 'en' for English.",
+            )
+
+        if lang_key == "hi":
+            tokenizer_used = "cls"
+        elif lang_key == "en":
+            tokenizer_used = "char"
+        else:
+            raise HTTPException(
+                status_code=400,
+                detail="Unsupported lang for hard-coded tokenizer. Use 'hi' or 'en'.",
+            )
+
+        cls_language = None
+        if tokenizer_used == "cls":
+            if payload.cls_language and payload.cls_language.strip():
+                cls_language = payload.cls_language.strip().lower()
+            else:
+                cls_language = "hindi" if lang_key == "hi" else "english" if lang_key == "en" else None
+            if not cls_language:
+                raise HTTPException(
+                    status_code=400,
+                    detail="cls_language is required when tokenizer=cls and lang is not hi/en.",
+                )
+
+        cleanup_path = None
+        if payload.ref_audio_path:
+            ref_audio = payload.ref_audio_path
+            if not os.path.exists(ref_audio):
+                raise HTTPException(status_code=400, detail=f"ref_audio_path not found: {ref_audio}")
+        else:
+            try:
+                audio_bytes = base64.b64decode(payload.ref_audio_base64)
+            except Exception as exc:  # noqa: BLE001
+                raise HTTPException(status_code=400, detail=f"Invalid ref_audio_base64: {exc}") from exc
+            ref_audio = _write_temp_audio(audio_bytes)
+            cleanup_path = ref_audio
+
+        try:
+            with infer_lock:
+                try:
+                    wav, sr, spec = tts.infer(
+                        ref_file=ref_audio,
+                        ref_text=payload.ref_text,
+                        gen_text=payload.gen_text,
+                        show_info=lambda *args, **kwargs: None,
+                        progress=None,
+                        target_rms=payload.target_rms,
+                        cross_fade_duration=payload.cross_fade_duration,
+                        sway_sampling_coef=payload.sway_sampling_coef,
+                        cfg_strength=payload.cfg_strength,
+                        nfe_step=payload.nfe_step,
+                        speed=payload.speed,
+                        fix_duration=payload.fix_duration,
+                        remove_silence=payload.remove_silence,
+                        seed=payload.seed,
+                        tokenizer=tokenizer_used,
+                        cls_language=cls_language,
+                        cls_server_url=ENV_DEFAULTS["cls_url"],
+                        cls_timeout=ENV_DEFAULTS["cls_timeout"],
+                    )
+                except Exception as exc:  # noqa: BLE001
+                    if tokenizer_used == "cls":
+                        raise HTTPException(
+                            status_code=502,
+                            detail=f"CLS tokenization failed: {exc}",
+                        ) from exc
+                    raise
+        finally:
+            if cleanup_path and os.path.exists(cleanup_path):
+                os.remove(cleanup_path)
+
+        response = {
+            "audio_base64": _encode_wav_base64(wav, sr),
+            "sample_rate": sr,
+            "seed": getattr(tts, "seed", payload.seed),
+        }
+        if payload.return_spectrogram and spec is not None:
+            response["spectrogram_base64"] = _encode_spec_base64(spec)
+
+        return response
+
+    return app
+
+
+app = create_app()
+
+
+@click.command()
+@click.option("--model", default=ENV_DEFAULTS["model"], show_default=True, help="Model config name to load.")
+@click.option("--ckpt-file", default=ENV_DEFAULTS["ckpt_file"], show_default=True, help="Checkpoint file path.")
+@click.option("--vocab-file", default=ENV_DEFAULTS["vocab_file"], show_default=True, help="Custom vocab file path.")
+@click.option("--ode-method", default=ENV_DEFAULTS["ode_method"], show_default=True, help="ODE method for sampler.")
+@click.option(
+    "--use-ema/--no-use-ema",
+    default=ENV_DEFAULTS["use_ema"],
+    show_default=True,
+    help="Load EMA weights from checkpoint.",
+)
+@click.option(
+    "--vocoder-local-path",
+    default=ENV_DEFAULTS["vocoder_local_path"],
+    show_default=True,
+    help="Local vocoder directory (skips HF download).",
+)
+@click.option("--device", default=ENV_DEFAULTS["device"], show_default=True, help="Force device: cpu|cuda|mps|xpu.")
+@click.option(
+    "--hf-cache-dir",
+    default=ENV_DEFAULTS["hf_cache_dir"],
+    show_default=True,
+    help="HuggingFace cache directory override.",
+)
+@click.option("--host", default="0.0.0.0", show_default=True, help="API host.")
+@click.option("--port", default=8060, show_default=True, help="API port.", type=int)
+@click.option("--root-path", default="", show_default=True, help="Set FastAPI root_path when behind a proxy.")
+def main(
+    model,
+    ckpt_file,
+    vocab_file,
+    ode_method,
+    use_ema,
+    vocoder_local_path,
+    device,
+    hf_cache_dir,
+    host,
+    port,
+    root_path,
+):
+    """Run the FastAPI server for HTTP inference."""
+    api_app = create_app(
+        model=model,
+        ckpt_file=ckpt_file,
+        vocab_file=vocab_file,
+        ode_method=ode_method,
+        use_ema=use_ema,
+        vocoder_local_path=vocoder_local_path,
+        device=device,
+        hf_cache_dir=hf_cache_dir,
+    )
+    uvicorn.run(api_app, host=host, port=port, root_path=root_path)
+
+
+if __name__ == "__main__":
+    main()

src/f5_tts/infer/infer_cli.py

@@ -0,0 +1,383 @@
+import argparse
+import codecs
+import os
+import re
+from datetime import datetime
+from importlib.resources import files
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import tomli
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+from unidecode import unidecode
+
+from f5_tts.infer.utils_infer import (
+    cfg_strength,
+    cross_fade_duration,
+    device,
+    fix_duration,
+    infer_process,
+    load_model,
+    load_vocoder,
+    mel_spec_type,
+    nfe_step,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    speed,
+    sway_sampling_coef,
+    target_rms,
+)
+
+
+parser = argparse.ArgumentParser(
+    prog="python3 infer-cli.py",
+    description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
+    epilog="Specify options above to override one or more settings from config.",
+)
+parser.add_argument(
+    "-c",
+    "--config",
+    type=str,
+    default=os.path.join(files("f5_tts").joinpath("infer/examples/basic"), "basic.toml"),
+    help="The configuration file, default see infer/examples/basic/basic.toml",
+)
+
+
+# Note. Not to provide default value here in order to read default from config file
+
+parser.add_argument(
+    "-m",
+    "--model",
+    type=str,
+    help="The model name: F5TTS_v1_Base | F5TTS_Base | E2TTS_Base | etc.",
+)
+parser.add_argument(
+    "-mc",
+    "--model_cfg",
+    type=str,
+    help="The path to F5-TTS model config file .yaml",
+)
+parser.add_argument(
+    "-p",
+    "--ckpt_file",
+    type=str,
+    help="The path to model checkpoint .pt, leave blank to use default",
+)
+parser.add_argument(
+    "-v",
+    "--vocab_file",
+    type=str,
+    help="The path to vocab file .txt, leave blank to use default",
+)
+parser.add_argument(
+    "-r",
+    "--ref_audio",
+    type=str,
+    help="The reference audio file.",
+)
+parser.add_argument(
+    "-s",
+    "--ref_text",
+    type=str,
+    help="The transcript/subtitle for the reference audio",
+)
+parser.add_argument(
+    "-t",
+    "--gen_text",
+    type=str,
+    help="The text to make model synthesize a speech",
+)
+parser.add_argument(
+    "-f",
+    "--gen_file",
+    type=str,
+    help="The file with text to generate, will ignore --gen_text",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    help="The path to output folder",
+)
+parser.add_argument(
+    "-w",
+    "--output_file",
+    type=str,
+    help="The name of output file",
+)
+parser.add_argument(
+    "--save_chunk",
+    action="store_true",
+    help="To save each audio chunks during inference",
+)
+parser.add_argument(
+    "--no_legacy_text",
+    action="store_false",
+    help="Not to use lossy ASCII transliterations of unicode text in saved file names.",
+)
+parser.add_argument(
+    "--remove_silence",
+    action="store_true",
+    help="To remove long silence found in ouput",
+)
+parser.add_argument(
+    "--load_vocoder_from_local",
+    action="store_true",
+    help="To load vocoder from local dir, default to ../checkpoints/vocos-mel-24khz",
+)
+parser.add_argument(
+    "--vocoder_name",
+    type=str,
+    choices=["vocos", "bigvgan"],
+    help=f"Used vocoder name: vocos | bigvgan, default {mel_spec_type}",
+)
+parser.add_argument(
+    "--target_rms",
+    type=float,
+    help=f"Target output speech loudness normalization value, default {target_rms}",
+)
+parser.add_argument(
+    "--cross_fade_duration",
+    type=float,
+    help=f"Duration of cross-fade between audio segments in seconds, default {cross_fade_duration}",
+)
+parser.add_argument(
+    "--nfe_step",
+    type=int,
+    help=f"The number of function evaluation (denoising steps), default {nfe_step}",
+)
+parser.add_argument(
+    "--cfg_strength",
+    type=float,
+    help=f"Classifier-free guidance strength, default {cfg_strength}",
+)
+parser.add_argument(
+    "--sway_sampling_coef",
+    type=float,
+    help=f"Sway Sampling coefficient, default {sway_sampling_coef}",
+)
+parser.add_argument(
+    "--speed",
+    type=float,
+    help=f"The speed of the generated audio, default {speed}",
+)
+parser.add_argument(
+    "--fix_duration",
+    type=float,
+    help=f"Fix the total duration (ref and gen audios) in seconds, default {fix_duration}",
+)
+parser.add_argument(
+    "--device",
+    type=str,
+    help="Specify the device to run on",
+)
+args = parser.parse_args()
+
+
+# config file
+
+config = tomli.load(open(args.config, "rb"))
+
+
+# command-line interface parameters
+
+model = args.model or config.get("model", "F5TTS_v1_Base")
+ckpt_file = args.ckpt_file or config.get("ckpt_file", "")
+vocab_file = args.vocab_file or config.get("vocab_file", "")
+
+ref_audio = args.ref_audio or config.get("ref_audio", "infer/examples/basic/basic_ref_en.wav")
+ref_text = (
+    args.ref_text
+    if args.ref_text is not None
+    else config.get("ref_text", "Some call me nature, others call me mother nature.")
+)
+gen_text = args.gen_text or config.get("gen_text", "Here we generate something just for test.")
+gen_file = args.gen_file or config.get("gen_file", "")
+
+output_dir = args.output_dir or config.get("output_dir", "tests")
+output_file = args.output_file or config.get(
+    "output_file", f"infer_cli_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}.wav"
+)
+
+save_chunk = args.save_chunk or config.get("save_chunk", False)
+use_legacy_text = args.no_legacy_text or config.get("no_legacy_text", False)  # no_legacy_text is a store_false arg
+if save_chunk and use_legacy_text:
+    print(
+        "\nWarning to --save_chunk: lossy ASCII transliterations of unicode text for legacy (.wav) file names, --no_legacy_text to disable.\n"
+    )
+
+remove_silence = args.remove_silence or config.get("remove_silence", False)
+load_vocoder_from_local = args.load_vocoder_from_local or config.get("load_vocoder_from_local", False)
+
+vocoder_name = args.vocoder_name or config.get("vocoder_name", mel_spec_type)
+target_rms = args.target_rms or config.get("target_rms", target_rms)
+cross_fade_duration = args.cross_fade_duration or config.get("cross_fade_duration", cross_fade_duration)
+nfe_step = args.nfe_step or config.get("nfe_step", nfe_step)
+cfg_strength = args.cfg_strength or config.get("cfg_strength", cfg_strength)
+sway_sampling_coef = args.sway_sampling_coef or config.get("sway_sampling_coef", sway_sampling_coef)
+speed = args.speed or config.get("speed", speed)
+fix_duration = args.fix_duration or config.get("fix_duration", fix_duration)
+device = args.device or config.get("device", device)
+
+
+# patches for pip pkg user
+if "infer/examples/" in ref_audio:
+    ref_audio = str(files("f5_tts").joinpath(f"{ref_audio}"))
+if "infer/examples/" in gen_file:
+    gen_file = str(files("f5_tts").joinpath(f"{gen_file}"))
+if "voices" in config:
+    for voice in config["voices"]:
+        voice_ref_audio = config["voices"][voice]["ref_audio"]
+        if "infer/examples/" in voice_ref_audio:
+            config["voices"][voice]["ref_audio"] = str(files("f5_tts").joinpath(f"{voice_ref_audio}"))
+
+
+# ignore gen_text if gen_file provided
+
+if gen_file:
+    gen_text = codecs.open(gen_file, "r", "utf-8").read()
+
+
+# output path
+
+wave_path = Path(output_dir) / output_file
+# spectrogram_path = Path(output_dir) / "infer_cli_out.png"
+if save_chunk:
+    output_chunk_dir = os.path.join(output_dir, f"{Path(output_file).stem}_chunks")
+    if not os.path.exists(output_chunk_dir):
+        os.makedirs(output_chunk_dir)
+
+
+# load vocoder
+
+if vocoder_name == "vocos":
+    vocoder_local_path = "../checkpoints/vocos-mel-24khz"
+elif vocoder_name == "bigvgan":
+    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+
+vocoder = load_vocoder(
+    vocoder_name=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path, device=device
+)
+
+
+# load TTS model
+
+model_cfg = OmegaConf.load(
+    args.model_cfg or config.get("model_cfg", str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+)
+model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+model_arc = model_cfg.model.arch
+
+repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
+
+if model != "F5TTS_Base":
+    assert vocoder_name == model_cfg.model.mel_spec.mel_spec_type
+
+# override for previous models
+if model == "F5TTS_Base":
+    if vocoder_name == "vocos":
+        ckpt_step = 1200000
+    elif vocoder_name == "bigvgan":
+        model = "F5TTS_Base_bigvgan"
+        ckpt_type = "pt"
+elif model == "E2TTS_Base":
+    repo_name = "E2-TTS"
+    ckpt_step = 1200000
+
+if not ckpt_file:
+    ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}"))
+
+print(f"Using {model}...")
+ema_model = load_model(
+    model_cls, model_arc, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file, device=device
+)
+
+
+# inference process
+
+
+def main():
+    main_voice = {"ref_audio": ref_audio, "ref_text": ref_text}
+    if "voices" not in config:
+        voices = {"main": main_voice}
+    else:
+        voices = config["voices"]
+        voices["main"] = main_voice
+    for voice in voices:
+        print("Voice:", voice)
+        print("ref_audio ", voices[voice]["ref_audio"])
+        voices[voice]["ref_audio"], voices[voice]["ref_text"] = preprocess_ref_audio_text(
+            voices[voice]["ref_audio"], voices[voice]["ref_text"]
+        )
+        print("ref_audio_", voices[voice]["ref_audio"], "\n\n")
+
+    generated_audio_segments = []
+    reg1 = r"(?=\[\w+\])"
+    chunks = re.split(reg1, gen_text)
+    reg2 = r"\[(\w+)\]"
+    for text in chunks:
+        if not text.strip():
+            continue
+        match = re.match(reg2, text)
+        if match:
+            voice = match[1]
+        else:
+            print("No voice tag found, using main.")
+            voice = "main"
+        if voice not in voices:
+            print(f"Voice {voice} not found, using main.")
+            voice = "main"
+        text = re.sub(reg2, "", text)
+        ref_audio_ = voices[voice]["ref_audio"]
+        ref_text_ = voices[voice]["ref_text"]
+        local_speed = voices[voice].get("speed", speed)
+        gen_text_ = text.strip()
+        print(f"Voice: {voice}")
+        audio_segment, final_sample_rate, spectrogram = infer_process(
+            ref_audio_,
+            ref_text_,
+            gen_text_,
+            ema_model,
+            vocoder,
+            mel_spec_type=vocoder_name,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=local_speed,
+            fix_duration=fix_duration,
+            device=device,
+        )
+        generated_audio_segments.append(audio_segment)
+
+        if save_chunk:
+            if len(gen_text_) > 200:
+                gen_text_ = gen_text_[:200] + " ... "
+            if use_legacy_text:
+                gen_text_ = unidecode(gen_text_)
+            sf.write(
+                os.path.join(output_chunk_dir, f"{len(generated_audio_segments) - 1}_{gen_text_}.wav"),
+                audio_segment,
+                final_sample_rate,
+            )
+
+    if generated_audio_segments:
+        final_wave = np.concatenate(generated_audio_segments)
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(wave_path, "wb") as f:
+            sf.write(f.name, final_wave, final_sample_rate)
+            # Remove silence
+            if remove_silence:
+                remove_silence_for_generated_wav(f.name)
+            print(f.name)
+
+
+if __name__ == "__main__":
+    main()

src/f5_tts/infer/infer_cli_resized_vocab.py

@@ -0,0 +1,521 @@
+import argparse
+import codecs
+import os
+import re
+import torch
+from datetime import datetime
+from importlib.resources import files
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import tomli
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+from unidecode import unidecode
+
+from f5_tts.infer.utils_infer import (
+    cfg_strength,
+    cross_fade_duration,
+    device,
+    fix_duration,
+    infer_process,
+    load_vocoder,
+    mel_spec_type,
+    nfe_step,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    speed,
+    sway_sampling_coef,
+    target_rms,
+)
+from f5_tts.model import CFM
+
+
+# -------------------------- Patch load_model with vocab resize -------------------------- #
+
+def load_model_with_vocab_resize(
+    model_cls,
+    model_cfg,
+    ckpt_path,
+    mel_spec_type="vocos",
+    vocab_file="",
+    use_ema=True,
+    device=None,
+):
+    """Load model with vocab size mismatch handling"""
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    print(f"\nLoading F5-TTS model:")
+    print(f"  Checkpoint: {ckpt_path}")
+    print(f"  Vocab file: {vocab_file}")
+    print(f"  Device: {device}")
+    
+    # Load vocab
+    vocab_char_map = None
+    vocab_size = None
+    
+    if vocab_file:
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            vocab_char_map = {}
+            for i, char in enumerate(f):
+                char = char.rstrip('\n')
+                vocab_char_map[char] = i
+        vocab_size = len(vocab_char_map)
+        print(f"  Vocab size: {vocab_size}")
+    
+    # Setup mel spec kwargs
+    if mel_spec_type == "vocos":
+        mel_spec_kwargs = dict(
+            n_fft=1024,
+            hop_length=256,
+            win_length=1024,
+            n_mel_channels=100,
+            target_sample_rate=24000,
+            mel_spec_type="vocos",
+        )
+    elif mel_spec_type == "bigvgan":
+        mel_spec_kwargs = dict(
+            n_fft=1024,
+            hop_length=256,
+            win_length=1024,
+            n_mel_channels=100,
+            target_sample_rate=24000,
+            mel_spec_type="bigvgan",
+        )
+    else:
+        raise ValueError(f"Unknown mel_spec_type: {mel_spec_type}")
+    
+    # Update model config with vocab size
+    if vocab_size is not None:
+        model_cfg = dict(model_cfg)  # Make a copy
+        model_cfg["text_num_embeds"] = vocab_size
+        model_cfg["mel_dim"] = mel_spec_kwargs["n_mel_channels"]
+    
+    # Initialize model
+    model = CFM(
+        transformer=model_cls(**model_cfg),
+        mel_spec_kwargs=mel_spec_kwargs,
+        vocab_char_map=vocab_char_map,
+    )
+    
+    model = model.to(device)
+    
+    # Load checkpoint with vocab resize support
+    print(f"\nLoading checkpoint: {ckpt_path}")
+    
+    # Handle different checkpoint formats
+    if ckpt_path.endswith('.safetensors'):
+        from safetensors.torch import load_file
+        state_dict = load_file(ckpt_path)
+    else:
+        checkpoint = torch.load(ckpt_path, map_location=device)
+        
+        # Get state dict
+        if use_ema and "ema_model_state_dict" in checkpoint:
+            state_dict = checkpoint["ema_model_state_dict"]
+        elif "model_state_dict" in checkpoint:
+            state_dict = checkpoint["model_state_dict"]
+        else:
+            state_dict = checkpoint
+    
+    # Remove ema_model prefix if exists
+    if any(k.startswith("ema_model.") for k in state_dict.keys()):
+        state_dict = {
+            k.replace("ema_model.", ""): v 
+            for k, v in state_dict.items() 
+            if k not in ["initted", "step"]
+        }
+    
+    # Handle vocab size mismatch
+    embed_key = 'transformer.text_embed.text_embed.weight'
+    if embed_key in state_dict:
+        current_embed_weight = model.transformer.text_embed.text_embed.weight
+        new_vocab_size, embed_dim = current_embed_weight.shape
+        old_weight = state_dict[embed_key]
+        old_vocab_size = old_weight.shape[0]
+        
+        if old_vocab_size != new_vocab_size:
+            print(f"\n⚠️  Vocab size mismatch detected!")
+            print(f"   Checkpoint vocab: {old_vocab_size}")
+            print(f"   Current vocab: {new_vocab_size}")
+            print(f"   Resizing embedding layer...")
+            
+            # Create new embedding weight
+            new_weight = torch.zeros(
+                new_vocab_size, embed_dim, 
+                device=old_weight.device, 
+                dtype=old_weight.dtype
+            )
+            
+            # Copy overlapping weights
+            min_vocab_size = min(old_vocab_size, new_vocab_size)
+            new_weight[:min_vocab_size, :] = old_weight[:min_vocab_size, :]
+            
+            # Handle different cases
+            if new_vocab_size > old_vocab_size:
+                # More tokens - initialize new ones
+                num_new_tokens = new_vocab_size - old_vocab_size
+                sample_indices = torch.randint(1, old_vocab_size, (num_new_tokens,))
+                new_weight[old_vocab_size:, :] = old_weight[sample_indices, :]
+                print(f"   ✓ Initialized {num_new_tokens} new embeddings")
+            else:
+                # Fewer tokens - just use subset
+                num_missing = old_vocab_size - new_vocab_size
+                print(f"   ✓ Using subset of embeddings ({new_vocab_size}/{old_vocab_size}, {num_missing} tokens dropped)")
+            
+            state_dict[embed_key] = new_weight
+    
+    # Load with strict=False
+    missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+    
+    if missing_keys:
+        print(f"⚠️  Missing keys: {missing_keys}")
+    if unexpected_keys:
+        print(f"⚠️  Unexpected keys: {unexpected_keys}")
+    
+    print("✓ Model loaded successfully\n")
+    
+    model.eval()
+    return model
+
+
+
+# -------------------------- Argument Parsing -------------------------- #
+
+parser = argparse.ArgumentParser(
+    prog="python3 infer-cli.py",
+    description="Commandline interface for E2/F5 TTS with Advanced Batch Processing.",
+    epilog="Specify options above to override one or more settings from config.",
+)
+parser.add_argument(
+    "-c",
+    "--config",
+    type=str,
+    default=os.path.join(files("f5_tts").joinpath("infer/examples/basic"), "basic.toml"),
+    help="The configuration file, default see infer/examples/basic/basic.toml",
+)
+
+parser.add_argument(
+    "-m",
+    "--model",
+    type=str,
+    help="The model name: F5TTS_v1_Base | F5TTS_Base | E2TTS_Base | etc.",
+)
+parser.add_argument(
+    "-mc",
+    "--model_cfg",
+    type=str,
+    help="The path to F5-TTS model config file .yaml",
+)
+parser.add_argument(
+    "-p",
+    "--ckpt_file",
+    type=str,
+    help="The path to model checkpoint .pt, leave blank to use default",
+)
+parser.add_argument(
+    "-v",
+    "--vocab_file",
+    type=str,
+    help="The path to vocab file .txt, leave blank to use default",
+)
+parser.add_argument(
+    "-r",
+    "--ref_audio",
+    type=str,
+    help="The reference audio file.",
+)
+parser.add_argument(
+    "-s",
+    "--ref_text",
+    type=str,
+    help="The transcript/subtitle for the reference audio",
+)
+parser.add_argument(
+    "-t",
+    "--gen_text",
+    type=str,
+    help="The text to make model synthesize a speech",
+)
+parser.add_argument(
+    "-f",
+    "--gen_file",
+    type=str,
+    help="The file with text to generate, will ignore --gen_text",
+)
+parser.add_argument(
+    "-o",
+    "--output_dir",
+    type=str,
+    help="The path to output folder",
+)
+parser.add_argument(
+    "-w",
+    "--output_file",
+    type=str,
+    help="The name of output file",
+)
+parser.add_argument(
+    "--save_chunk",
+    action="store_true",
+    help="To save each audio chunks during inference",
+)
+parser.add_argument(
+    "--no_legacy_text",
+    action="store_false",
+    help="Not to use lossy ASCII transliterations of unicode text in saved file names.",
+)
+parser.add_argument(
+    "--remove_silence",
+    action="store_true",
+    help="To remove long silence found in ouput",
+)
+parser.add_argument(
+    "--load_vocoder_from_local",
+    action="store_true",
+    help="To load vocoder from local dir, default to ../checkpoints/vocos-mel-24khz",
+)
+parser.add_argument(
+    "--vocoder_name",
+    type=str,
+    choices=["vocos", "bigvgan"],
+    help=f"Used vocoder name: vocos | bigvgan, default {mel_spec_type}",
+)
+parser.add_argument(
+    "--target_rms",
+    type=float,
+    help=f"Target output speech loudness normalization value, default {target_rms}",
+)
+parser.add_argument(
+    "--cross_fade_duration",
+    type=float,
+    help=f"Duration of cross-fade between audio segments in seconds, default {cross_fade_duration}",
+)
+parser.add_argument(
+    "--nfe_step",
+    type=int,
+    help=f"The number of function evaluation (denoising steps), default {nfe_step}",
+)
+parser.add_argument(
+    "--cfg_strength",
+    type=float,
+    help=f"Classifier-free guidance strength, default {cfg_strength}",
+)
+parser.add_argument(
+    "--sway_sampling_coef",
+    type=float,
+    help=f"Sway Sampling coefficient, default {sway_sampling_coef}",
+)
+parser.add_argument(
+    "--speed",
+    type=float,
+    help=f"The speed of the generated audio, default {speed}",
+)
+parser.add_argument(
+    "--fix_duration",
+    type=float,
+    help=f"Fix the total duration (ref and gen audios) in seconds, default {fix_duration}",
+)
+parser.add_argument(
+    "--device",
+    type=str,
+    help="Specify the device to run on",
+)
+args = parser.parse_args()
+
+# config file
+config = tomli.load(open(args.config, "rb"))
+
+# command-line interface parameters
+model = args.model or config.get("model", "F5TTS_v1_Base")
+ckpt_file = args.ckpt_file or config.get("ckpt_file", "")
+vocab_file = args.vocab_file or config.get("vocab_file", "")
+
+ref_audio = args.ref_audio or config.get("ref_audio", "infer/examples/basic/basic_ref_en.wav")
+ref_text = (
+    args.ref_text
+    if args.ref_text is not None
+    else config.get("ref_text", "Some call me nature, others call me mother nature.")
+)
+gen_text = args.gen_text or config.get("gen_text", "Here we generate something just for test.")
+gen_file = args.gen_file or config.get("gen_file", "")
+
+output_dir = args.output_dir or config.get("output_dir", "tests")
+output_file = args.output_file or config.get(
+    "output_file", f"infer_cli_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}.wav"
+)
+
+save_chunk = args.save_chunk or config.get("save_chunk", False)
+use_legacy_text = args.no_legacy_text or config.get("no_legacy_text", False)
+if save_chunk and use_legacy_text:
+    print(
+        "\nWarning to --save_chunk: lossy ASCII transliterations of unicode text for legacy (.wav) file names, --no_legacy_text to disable.\n"
+    )
+
+remove_silence = args.remove_silence or config.get("remove_silence", False)
+load_vocoder_from_local = args.load_vocoder_from_local or config.get("load_vocoder_from_local", False)
+
+vocoder_name = args.vocoder_name or config.get("vocoder_name", mel_spec_type)
+target_rms = args.target_rms or config.get("target_rms", target_rms)
+cross_fade_duration = args.cross_fade_duration or config.get("cross_fade_duration", cross_fade_duration)
+nfe_step = args.nfe_step or config.get("nfe_step", nfe_step)
+cfg_strength = args.cfg_strength or config.get("cfg_strength", cfg_strength)
+sway_sampling_coef = args.sway_sampling_coef or config.get("sway_sampling_coef", sway_sampling_coef)
+speed = args.speed or config.get("speed", speed)
+fix_duration = args.fix_duration or config.get("fix_duration", fix_duration)
+device = args.device or config.get("device", device)
+
+# patches for pip pkg user
+if "infer/examples/" in ref_audio:
+    ref_audio = str(files("f5_tts").joinpath(f"{ref_audio}"))
+if "infer/examples/" in gen_file:
+    gen_file = str(files("f5_tts").joinpath(f"{gen_file}"))
+if "voices" in config:
+    for voice in config["voices"]:
+        voice_ref_audio = config["voices"][voice]["ref_audio"]
+        if "infer/examples/" in voice_ref_audio:
+            config["voices"][voice]["ref_audio"] = str(files("f5_tts").joinpath(f"{voice_ref_audio}"))
+
+# ignore gen_text if gen_file provided
+if gen_file:
+    gen_text = codecs.open(gen_file, "r", "utf-8").read()
+
+# output path
+wave_path = Path(output_dir) / output_file
+if save_chunk:
+    output_chunk_dir = os.path.join(output_dir, f"{Path(output_file).stem}_chunks")
+    if not os.path.exists(output_chunk_dir):
+        os.makedirs(output_chunk_dir)
+
+# load vocoder
+if vocoder_name == "vocos":
+    vocoder_local_path = "../checkpoints/vocos-mel-24khz"
+elif vocoder_name == "bigvgan":
+    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+
+vocoder = load_vocoder(
+    vocoder_name=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path, device=device
+)
+
+# load TTS model
+model_cfg = OmegaConf.load(
+    args.model_cfg or config.get("model_cfg", str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+)
+model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+model_arc = model_cfg.model.arch
+
+repo_name, ckpt_step, ckpt_type = "F5-TTS", 1250000, "safetensors"
+
+if model != "F5TTS_Base":
+    assert vocoder_name == model_cfg.model.mel_spec.mel_spec_type
+
+# override for previous models
+if model == "F5TTS_Base":
+    if vocoder_name == "vocos":
+        ckpt_step = 1200000
+    elif vocoder_name == "bigvgan":
+        model = "F5TTS_Base_bigvgan"
+        ckpt_type = "pt"
+elif model == "E2TTS_Base":
+    repo_name = "E2-TTS"
+    ckpt_step = 1200000
+
+if not ckpt_file:
+    ckpt_file = str(cached_path(f"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}"))
+
+print(f"Using {model}...")
+
+# Use patched load function with vocab resize support
+ema_model = load_model_with_vocab_resize(
+    model_cls, model_arc, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file, device=device
+)
+
+
+# inference process
+def main():
+    main_voice = {"ref_audio": ref_audio, "ref_text": ref_text}
+    if "voices" not in config:
+        voices = {"main": main_voice}
+    else:
+        voices = config["voices"]
+        voices["main"] = main_voice
+    
+    for voice in voices:
+        print("Voice:", voice)
+        print("ref_audio ", voices[voice]["ref_audio"])
+        voices[voice]["ref_audio"], voices[voice]["ref_text"] = preprocess_ref_audio_text(
+            voices[voice]["ref_audio"], voices[voice]["ref_text"]
+        )
+        print("ref_audio_", voices[voice]["ref_audio"], "\n\n")
+
+    generated_audio_segments = []
+    reg1 = r"(?=\[\w+\])"
+    chunks = re.split(reg1, gen_text)
+    reg2 = r"\[(\w+)\]"
+    
+    for text in chunks:
+        if not text.strip():
+            continue
+        match = re.match(reg2, text)
+        if match:
+            voice = match[1]
+        else:
+            print("No voice tag found, using main.")
+            voice = "main"
+        if voice not in voices:
+            print(f"Voice {voice} not found, using main.")
+            voice = "main"
+        text = re.sub(reg2, "", text)
+        ref_audio_ = voices[voice]["ref_audio"]
+        ref_text_ = voices[voice]["ref_text"]
+        local_speed = voices[voice].get("speed", speed)
+        gen_text_ = text.strip()
+        print(f"Voice: {voice}")
+        
+        audio_segment, final_sample_rate, spectrogram = infer_process(
+            ref_audio_,
+            ref_text_,
+            gen_text_,
+            ema_model,
+            vocoder,
+            mel_spec_type=vocoder_name,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=local_speed,
+            fix_duration=fix_duration,
+            device=device,
+        )
+        generated_audio_segments.append(audio_segment)
+
+        if save_chunk:
+            if len(gen_text_) > 200:
+                gen_text_ = gen_text_[:200] + " ... "
+            if use_legacy_text:
+                gen_text_ = unidecode(gen_text_)
+            sf.write(
+                os.path.join(output_chunk_dir, f"{len(generated_audio_segments) - 1}_{gen_text_}.wav"),
+                audio_segment,
+                final_sample_rate,
+            )
+
+    if generated_audio_segments:
+        final_wave = np.concatenate(generated_audio_segments)
+
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+
+        with open(wave_path, "wb") as f:
+            sf.write(f.name, final_wave, final_sample_rate)
+            # Remove silence
+            if remove_silence:
+                remove_silence_for_generated_wav(f.name)
+            print(f.name)
+
+if __name__ == "__main__":
+    main()

src/f5_tts/infer/infer_gradio.py

@@ -0,0 +1,1121 @@
+# ruff: noqa: E402
+# Above allows ruff to ignore E402: module level import not at top of file
+
+import gc
+import json
+import os
+import re
+import tempfile
+from collections import OrderedDict
+from functools import lru_cache
+from importlib.resources import files
+
+import click
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+from cached_path import cached_path
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+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
+
+
+from f5_tts.infer.utils_infer import (
+    infer_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+    remove_silence_for_generated_wav,
+    save_spectrogram,
+    tempfile_kwargs,
+)
+from f5_tts.model import DiT, UNetT
+
+
+DEFAULT_TTS_MODEL = "F5-TTS_v1"
+tts_model_choice = DEFAULT_TTS_MODEL
+
+DEFAULT_TTS_MODEL_CFG = [
+    "hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors",
+    "hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt",
+    json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),
+]
+
+
+# load models
+
+vocoder = load_vocoder()
+
+
+def load_f5tts():
+    ckpt_path = str(cached_path(DEFAULT_TTS_MODEL_CFG[0]))
+    F5TTS_model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
+    return load_model(DiT, F5TTS_model_cfg, ckpt_path)
+
+
+def load_e2tts():
+    ckpt_path = str(cached_path("hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors"))
+    E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4, text_mask_padding=False, pe_attn_head=1)
+    return load_model(UNetT, E2TTS_model_cfg, ckpt_path)
+
+
+def load_custom(ckpt_path: str, vocab_path="", model_cfg=None):
+    ckpt_path, vocab_path = ckpt_path.strip(), vocab_path.strip()
+    if ckpt_path.startswith("hf://"):
+        ckpt_path = str(cached_path(ckpt_path))
+    if vocab_path.startswith("hf://"):
+        vocab_path = str(cached_path(vocab_path))
+    if model_cfg is None:
+        model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])
+    elif isinstance(model_cfg, str):
+        model_cfg = json.loads(model_cfg)
+    return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path)
+
+
+F5TTS_ema_model = load_f5tts()
+E2TTS_ema_model = load_e2tts() if USING_SPACES else None
+custom_ema_model, pre_custom_path = None, ""
+
+chat_model_state = None
+chat_tokenizer_state = None
+
+
+@gpu_decorator
+def chat_model_inference(messages, model, tokenizer):
+    """Generate response using Qwen"""
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.95,
+    )
+
+    generated_ids = [
+        output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+    return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+
+@gpu_decorator
+def load_text_from_file(file):
+    if file:
+        with open(file, "r", encoding="utf-8") as f:
+            text = f.read().strip()
+    else:
+        text = ""
+    return gr.update(value=text)
+
+
+@lru_cache(maxsize=1000)  # NOTE. need to ensure params of infer() hashable
+@gpu_decorator
+def infer(
+    ref_audio_orig,
+    ref_text,
+    gen_text,
+    model,
+    remove_silence,
+    seed,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    speed=1,
+    show_info=gr.Info,
+):
+    if not ref_audio_orig:
+        gr.Warning("Please provide reference audio.")
+        return gr.update(), gr.update(), ref_text
+
+    # Set inference seed
+    if seed < 0 or seed > 2**31 - 1:
+        gr.Warning("Seed must in range 0 ~ 2147483647. Using random seed instead.")
+        seed = np.random.randint(0, 2**31 - 1)
+    torch.manual_seed(seed)
+    used_seed = seed
+
+    if not gen_text.strip():
+        gr.Warning("Please enter text to generate or upload a text file.")
+        return gr.update(), gr.update(), ref_text
+
+    ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)
+
+    if model == DEFAULT_TTS_MODEL:
+        ema_model = F5TTS_ema_model
+    elif model == "E2-TTS":
+        global E2TTS_ema_model
+        if E2TTS_ema_model is None:
+            show_info("Loading E2-TTS model...")
+            E2TTS_ema_model = load_e2tts()
+        ema_model = E2TTS_ema_model
+    elif isinstance(model, tuple) and model[0] == "Custom":
+        assert not USING_SPACES, "Only official checkpoints allowed in Spaces."
+        global custom_ema_model, pre_custom_path
+        if pre_custom_path != model[1]:
+            show_info("Loading Custom TTS model...")
+            custom_ema_model = load_custom(model[1], vocab_path=model[2], model_cfg=model[3])
+            pre_custom_path = model[1]
+        ema_model = custom_ema_model
+
+    final_wave, final_sample_rate, combined_spectrogram = infer_process(
+        ref_audio,
+        ref_text,
+        gen_text,
+        ema_model,
+        vocoder,
+        cross_fade_duration=cross_fade_duration,
+        nfe_step=nfe_step,
+        speed=speed,
+        show_info=show_info,
+        progress=gr.Progress(),
+    )
+
+    # Remove silence
+    if remove_silence:
+        with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
+            temp_path = f.name
+        try:
+            sf.write(temp_path, final_wave, final_sample_rate)
+            remove_silence_for_generated_wav(f.name)
+            final_wave, _ = torchaudio.load(f.name)
+        finally:
+            os.unlink(temp_path)
+        final_wave = final_wave.squeeze().cpu().numpy()
+
+    # Save the spectrogram
+    with tempfile.NamedTemporaryFile(suffix=".png", **tempfile_kwargs) as tmp_spectrogram:
+        spectrogram_path = tmp_spectrogram.name
+    save_spectrogram(combined_spectrogram, spectrogram_path)
+
+    return (final_sample_rate, final_wave), spectrogram_path, ref_text, used_seed
+
+
+with gr.Blocks() as app_tts:
+    gr.Markdown("# Batched TTS")
+    ref_audio_input = gr.Audio(label="Reference Audio", type="filepath")
+    with gr.Row():
+        gen_text_input = gr.Textbox(
+            label="Text to Generate",
+            lines=10,
+            max_lines=40,
+            scale=4,
+        )
+        gen_text_file = gr.File(label="Load Text to Generate from File (.txt)", file_types=[".txt"], scale=1)
+    generate_btn = gr.Button("Synthesize", variant="primary")
+    with gr.Accordion("Advanced Settings", open=False):
+        with gr.Row():
+            ref_text_input = gr.Textbox(
+                label="Reference Text",
+                info="Leave blank to automatically transcribe the reference audio. If you enter text or upload a file, it will override automatic transcription.",
+                lines=2,
+                scale=4,
+            )
+            ref_text_file = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"], scale=1)
+        with gr.Row():
+            randomize_seed = gr.Checkbox(
+                label="Randomize Seed",
+                info="Check to use a random seed for each generation. Uncheck to use the seed specified.",
+                value=True,
+                scale=3,
+            )
+            seed_input = gr.Number(show_label=False, value=0, precision=0, scale=1)
+            with gr.Column(scale=4):
+                remove_silence = gr.Checkbox(
+                    label="Remove Silences",
+                    info="If undesired long silence(s) produced, turn on to automatically detect and crop.",
+                    value=False,
+                )
+        speed_slider = gr.Slider(
+            label="Speed",
+            minimum=0.3,
+            maximum=2.0,
+            value=1.0,
+            step=0.1,
+            info="Adjust the speed of the audio.",
+        )
+        nfe_slider = gr.Slider(
+            label="NFE Steps",
+            minimum=4,
+            maximum=64,
+            value=32,
+            step=2,
+            info="Set the number of denoising steps.",
+        )
+        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.",
+        )
+
+    audio_output = gr.Audio(label="Synthesized Audio")
+    spectrogram_output = gr.Image(label="Spectrogram")
+
+    @gpu_decorator
+    def basic_tts(
+        ref_audio_input,
+        ref_text_input,
+        gen_text_input,
+        remove_silence,
+        randomize_seed,
+        seed_input,
+        cross_fade_duration_slider,
+        nfe_slider,
+        speed_slider,
+    ):
+        if randomize_seed:
+            seed_input = np.random.randint(0, 2**31 - 1)
+
+        audio_out, spectrogram_path, ref_text_out, used_seed = infer(
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            tts_model_choice,
+            remove_silence,
+            seed=seed_input,
+            cross_fade_duration=cross_fade_duration_slider,
+            nfe_step=nfe_slider,
+            speed=speed_slider,
+        )
+        return audio_out, spectrogram_path, ref_text_out, used_seed
+
+    gen_text_file.upload(
+        load_text_from_file,
+        inputs=[gen_text_file],
+        outputs=[gen_text_input],
+    )
+
+    ref_text_file.upload(
+        load_text_from_file,
+        inputs=[ref_text_file],
+        outputs=[ref_text_input],
+    )
+
+    ref_audio_input.clear(
+        lambda: [None, None],
+        None,
+        [ref_text_input, ref_text_file],
+    )
+
+    generate_btn.click(
+        basic_tts,
+        inputs=[
+            ref_audio_input,
+            ref_text_input,
+            gen_text_input,
+            remove_silence,
+            randomize_seed,
+            seed_input,
+            cross_fade_duration_slider,
+            nfe_slider,
+            speed_slider,
+        ],
+        outputs=[audio_output, spectrogram_output, ref_text_input, seed_input],
+    )
+
+
+def parse_speechtypes_text(gen_text):
+    # Pattern to find {str} or {"name": str, "seed": int, "speed": float}
+    pattern = r"(\{.*?\})"
+
+    # Split the text by the pattern
+    tokens = re.split(pattern, gen_text)
+
+    segments = []
+
+    current_type_dict = {
+        "name": "Regular",
+        "seed": -1,
+        "speed": 1.0,
+    }
+
+    for i in range(len(tokens)):
+        if i % 2 == 0:
+            # This is text
+            text = tokens[i].strip()
+            if text:
+                current_type_dict["text"] = text
+                segments.append(current_type_dict)
+        else:
+            # This is type
+            type_str = tokens[i].strip()
+            try:  # if type dict
+                current_type_dict = json.loads(type_str)
+            except json.decoder.JSONDecodeError:
+                type_str = type_str[1:-1]  # remove brace {}
+                current_type_dict = {"name": type_str, "seed": -1, "speed": 1.0}
+
+    return segments
+
+
+with gr.Blocks() as app_multistyle:
+    # New section for multistyle generation
+    gr.Markdown(
+        """
+    # Multiple Speech-Type Generation
+
+    This section allows you to generate multiple speech types or multiple people's voices. Enter your text in the format shown below, or upload a .txt file with the same format. The system will generate speech using the appropriate type. If unspecified, the model will use the regular speech type. The current speech type will be used until the next speech type is specified.
+    """
+    )
+
+    with gr.Row():
+        gr.Markdown(
+            """
+            **Example Input:** <br>
+            {Regular} Hello, I'd like to order a sandwich please. <br>
+            {Surprised} What do you mean you're out of bread? <br>
+            {Sad} I really wanted a sandwich though... <br>
+            {Angry} You know what, darn you and your little shop! <br>
+            {Whisper} I'll just go back home and cry now. <br>
+            {Shouting} Why me?!
+            """
+        )
+
+        gr.Markdown(
+            """
+            **Example Input 2:** <br>
+            {"name": "Speaker1_Happy", "seed": -1, "speed": 1} Hello, I'd like to order a sandwich please. <br>
+            {"name": "Speaker2_Regular", "seed": -1, "speed": 1} Sorry, we're out of bread. <br>
+            {"name": "Speaker1_Sad", "seed": -1, "speed": 1} I really wanted a sandwich though... <br>
+            {"name": "Speaker2_Whisper", "seed": -1, "speed": 1} I'll give you the last one I was hiding.
+            """
+        )
+
+    gr.Markdown(
+        'Upload different audio clips for each speech type. The first speech type is mandatory. You can add additional speech types by clicking the "Add Speech Type" button.'
+    )
+
+    # Regular speech type (mandatory)
+    with gr.Row(variant="compact") as regular_row:
+        with gr.Column(scale=1, min_width=160):
+            regular_name = gr.Textbox(value="Regular", label="Speech Type Name")
+            regular_insert = gr.Button("Insert Label", variant="secondary")
+        with gr.Column(scale=3):
+            regular_audio = gr.Audio(label="Regular Reference Audio", type="filepath")
+        with gr.Column(scale=3):
+            regular_ref_text = gr.Textbox(label="Reference Text (Regular)", lines=4)
+            with gr.Row():
+                regular_seed_slider = gr.Slider(
+                    show_label=False, minimum=-1, maximum=999, value=-1, step=1, info="Seed, -1 for random"
+                )
+                regular_speed_slider = gr.Slider(
+                    show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info="Adjust the speed"
+                )
+        with gr.Column(scale=1, min_width=160):
+            regular_ref_text_file = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"])
+
+    # Regular speech type (max 100)
+    max_speech_types = 100
+    speech_type_rows = [regular_row]
+    speech_type_names = [regular_name]
+    speech_type_audios = [regular_audio]
+    speech_type_ref_texts = [regular_ref_text]
+    speech_type_ref_text_files = [regular_ref_text_file]
+    speech_type_seeds = [regular_seed_slider]
+    speech_type_speeds = [regular_speed_slider]
+    speech_type_delete_btns = [None]
+    speech_type_insert_btns = [regular_insert]
+
+    # Additional speech types (99 more)
+    for i in range(max_speech_types - 1):
+        with gr.Row(variant="compact", visible=False) as row:
+            with gr.Column(scale=1, min_width=160):
+                name_input = gr.Textbox(label="Speech Type Name")
+                insert_btn = gr.Button("Insert Label", variant="secondary")
+                delete_btn = gr.Button("Delete Type", variant="stop")
+            with gr.Column(scale=3):
+                audio_input = gr.Audio(label="Reference Audio", type="filepath")
+            with gr.Column(scale=3):
+                ref_text_input = gr.Textbox(label="Reference Text", lines=4)
+                with gr.Row():
+                    seed_input = gr.Slider(
+                        show_label=False, minimum=-1, maximum=999, value=-1, step=1, info="Seed. -1 for random"
+                    )
+                    speed_input = gr.Slider(
+                        show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info="Adjust the speed"
+                    )
+            with gr.Column(scale=1, min_width=160):
+                ref_text_file_input = gr.File(label="Load Reference Text from File (.txt)", file_types=[".txt"])
+        speech_type_rows.append(row)
+        speech_type_names.append(name_input)
+        speech_type_audios.append(audio_input)
+        speech_type_ref_texts.append(ref_text_input)
+        speech_type_ref_text_files.append(ref_text_file_input)
+        speech_type_seeds.append(seed_input)
+        speech_type_speeds.append(speed_input)
+        speech_type_delete_btns.append(delete_btn)
+        speech_type_insert_btns.append(insert_btn)
+
+    # Global logic for all speech types
+    for i in range(max_speech_types):
+        speech_type_audios[i].clear(
+            lambda: [None, None],
+            None,
+            [speech_type_ref_texts[i], speech_type_ref_text_files[i]],
+        )
+        speech_type_ref_text_files[i].upload(
+            load_text_from_file,
+            inputs=[speech_type_ref_text_files[i]],
+            outputs=[speech_type_ref_texts[i]],
+        )
+
+    # Button to add speech type
+    add_speech_type_btn = gr.Button("Add Speech Type")
+
+    # Keep track of autoincrement of speech types, no roll back
+    speech_type_count = 1
+
+    # Function to add a speech type
+    def add_speech_type_fn():
+        row_updates = [gr.update() for _ in range(max_speech_types)]
+        global speech_type_count
+        if speech_type_count < max_speech_types:
+            row_updates[speech_type_count] = gr.update(visible=True)
+            speech_type_count += 1
+        else:
+            gr.Warning("Exhausted maximum number of speech types. Consider restart the app.")
+        return row_updates
+
+    add_speech_type_btn.click(add_speech_type_fn, outputs=speech_type_rows)
+
+    # Function to delete a speech type
+    def delete_speech_type_fn():
+        return gr.update(visible=False), None, None, None, None
+
+    # Update delete button clicks and ref text file changes
+    for i in range(1, len(speech_type_delete_btns)):
+        speech_type_delete_btns[i].click(
+            delete_speech_type_fn,
+            outputs=[
+                speech_type_rows[i],
+                speech_type_names[i],
+                speech_type_audios[i],
+                speech_type_ref_texts[i],
+                speech_type_ref_text_files[i],
+            ],
+        )
+
+    # Text input for the prompt
+    with gr.Row():
+        gen_text_input_multistyle = gr.Textbox(
+            label="Text to Generate",
+            lines=10,
+            max_lines=40,
+            scale=4,
+            placeholder="Enter the script with speaker names (or emotion types) at the start of each block, e.g.:\n\n{Regular} Hello, I'd like to order a sandwich please.\n{Surprised} What do you mean you're out of bread?\n{Sad} I really wanted a sandwich though...\n{Angry} You know what, darn you and your little shop!\n{Whisper} I'll just go back home and cry now.\n{Shouting} Why me?!",
+        )
+        gen_text_file_multistyle = gr.File(label="Load Text to Generate from File (.txt)", file_types=[".txt"], scale=1)
+
+    def make_insert_speech_type_fn(index):
+        def insert_speech_type_fn(current_text, speech_type_name, speech_type_seed, speech_type_speed):
+            current_text = current_text or ""
+            if not speech_type_name:
+                gr.Warning("Please enter speech type name before insert.")
+                return current_text
+            speech_type_dict = {
+                "name": speech_type_name,
+                "seed": speech_type_seed,
+                "speed": speech_type_speed,
+            }
+            updated_text = current_text + json.dumps(speech_type_dict) + " "
+            return updated_text
+
+        return insert_speech_type_fn
+
+    for i, insert_btn in enumerate(speech_type_insert_btns):
+        insert_fn = make_insert_speech_type_fn(i)
+        insert_btn.click(
+            insert_fn,
+            inputs=[gen_text_input_multistyle, speech_type_names[i], speech_type_seeds[i], speech_type_speeds[i]],
+            outputs=gen_text_input_multistyle,
+        )
+
+    with gr.Accordion("Advanced Settings", open=True):
+        with gr.Row():
+            with gr.Column():
+                show_cherrypick_multistyle = gr.Checkbox(
+                    label="Show Cherry-pick Interface",
+                    info="Turn on to show interface, picking seeds from previous generations.",
+                    value=False,
+                )
+            with gr.Column():
+                remove_silence_multistyle = gr.Checkbox(
+                    label="Remove Silences",
+                    info="Turn on to automatically detect and crop long silences.",
+                    value=True,
+                )
+
+    # Generate button
+    generate_multistyle_btn = gr.Button("Generate Multi-Style Speech", variant="primary")
+
+    # Output audio
+    audio_output_multistyle = gr.Audio(label="Synthesized Audio")
+
+    # Used seed gallery
+    cherrypick_interface_multistyle = gr.Textbox(
+        label="Cherry-pick Interface",
+        lines=10,
+        max_lines=40,
+        show_copy_button=True,
+        interactive=False,
+        visible=False,
+    )
+
+    # Logic control to show/hide the cherrypick interface
+    show_cherrypick_multistyle.change(
+        lambda is_visible: gr.update(visible=is_visible),
+        show_cherrypick_multistyle,
+        cherrypick_interface_multistyle,
+    )
+
+    # Function to load text to generate from file
+    gen_text_file_multistyle.upload(
+        load_text_from_file,
+        inputs=[gen_text_file_multistyle],
+        outputs=[gen_text_input_multistyle],
+    )
+
+    @gpu_decorator
+    def generate_multistyle_speech(
+        gen_text,
+        *args,
+    ):
+        speech_type_names_list = args[:max_speech_types]
+        speech_type_audios_list = args[max_speech_types : 2 * max_speech_types]
+        speech_type_ref_texts_list = args[2 * max_speech_types : 3 * max_speech_types]
+        remove_silence = args[3 * max_speech_types]
+        # Collect the speech types and their audios into a dict
+        speech_types = OrderedDict()
+
+        ref_text_idx = 0
+        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}
+            else:
+                speech_types[f"@{ref_text_idx}@"] = {"audio": "", "ref_text": ""}
+            ref_text_idx += 1
+
+        # Parse the gen_text into segments
+        segments = parse_speechtypes_text(gen_text)
+
+        # For each segment, generate speech
+        generated_audio_segments = []
+        current_type_name = "Regular"
+        inference_meta_data = ""
+
+        for segment in segments:
+            name = segment["name"]
+            seed_input = segment["seed"]
+            speed = segment["speed"]
+            text = segment["text"]
+
+            if name in speech_types:
+                current_type_name = name
+            else:
+                gr.Warning(f"Type {name} is not available, will use Regular as default.")
+                current_type_name = "Regular"
+
+            try:
+                ref_audio = speech_types[current_type_name]["audio"]
+            except KeyError:
+                gr.Warning(f"Please provide reference audio for type {current_type_name}.")
+                return [None] + [speech_types[name]["ref_text"] for name in speech_types] + [None]
+            ref_text = speech_types[current_type_name].get("ref_text", "")
+
+            if seed_input == -1:
+                seed_input = np.random.randint(0, 2**31 - 1)
+
+            # Generate or retrieve speech for this segment
+            audio_out, _, ref_text_out, used_seed = infer(
+                ref_audio,
+                ref_text,
+                text,
+                tts_model_choice,
+                remove_silence,
+                seed=seed_input,
+                cross_fade_duration=0,
+                speed=speed,
+                show_info=print,  # no pull to top when generating
+            )
+            sr, audio_data = audio_out
+
+            generated_audio_segments.append(audio_data)
+            speech_types[current_type_name]["ref_text"] = ref_text_out
+            inference_meta_data += json.dumps(dict(name=name, seed=used_seed, speed=speed)) + f" {text}\n"
+
+        # Concatenate all audio segments
+        if generated_audio_segments:
+            final_audio_data = np.concatenate(generated_audio_segments)
+            return (
+                [(sr, final_audio_data)]
+                + [speech_types[name]["ref_text"] for name in speech_types]
+                + [inference_meta_data]
+            )
+        else:
+            gr.Warning("No audio generated.")
+            return [None] + [speech_types[name]["ref_text"] for name in speech_types] + [None]
+
+    generate_multistyle_btn.click(
+        generate_multistyle_speech,
+        inputs=[
+            gen_text_input_multistyle,
+        ]
+        + speech_type_names
+        + speech_type_audios
+        + speech_type_ref_texts
+        + [
+            remove_silence_multistyle,
+        ],
+        outputs=[audio_output_multistyle] + speech_type_ref_texts + [cherrypick_interface_multistyle],
+    )
+
+    # Validation function to disable Generate button if speech types are missing
+    def validate_speech_types(gen_text, regular_name, *args):
+        speech_type_names_list = args
+
+        # 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_speechtypes_text(gen_text)
+        speech_types_in_text = set(segment["name"] 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_multistyle.change(
+        validate_speech_types,
+        inputs=[gen_text_input_multistyle, regular_name] + speech_type_names,
+        outputs=generate_multistyle_btn,
+    )
+
+
+with gr.Blocks() as app_chat:
+    gr.Markdown(
+        """
+# Voice Chat
+Have a conversation with an AI using your reference voice!
+1. Upload a reference audio clip and optionally its transcript (via text or .txt file).
+2. Load the chat model.
+3. Record your message through your microphone or type it.
+4. The AI will respond using the reference voice.
+"""
+    )
+
+    chat_model_name_list = [
+        "Qwen/Qwen2.5-3B-Instruct",
+        "microsoft/Phi-4-mini-instruct",
+    ]
+
+    @gpu_decorator
+    def load_chat_model(chat_model_name):
+        show_info = gr.Info
+        global chat_model_state, chat_tokenizer_state
+        if chat_model_state is not None:
+            chat_model_state = None
+            chat_tokenizer_state = None
+            gc.collect()
+            torch.cuda.empty_cache()
+
+        show_info(f"Loading chat model: {chat_model_name}")
+        chat_model_state = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype="auto", device_map="auto")
+        chat_tokenizer_state = AutoTokenizer.from_pretrained(chat_model_name)
+        show_info(f"Chat model {chat_model_name} loaded successfully!")
+
+        return gr.update(visible=False), gr.update(visible=True)
+
+    if USING_SPACES:
+        load_chat_model(chat_model_name_list[0])
+
+    chat_model_name_input = gr.Dropdown(
+        choices=chat_model_name_list,
+        value=chat_model_name_list[0],
+        label="Chat Model Name",
+        info="Enter the name of a HuggingFace chat model",
+        allow_custom_value=not USING_SPACES,
+    )
+    load_chat_model_btn = gr.Button("Load Chat Model", variant="primary", visible=not USING_SPACES)
+    chat_interface_container = gr.Column(visible=USING_SPACES)
+
+    chat_model_name_input.change(
+        lambda: gr.update(visible=True),
+        None,
+        load_chat_model_btn,
+        show_progress="hidden",
+    )
+    load_chat_model_btn.click(
+        load_chat_model, inputs=[chat_model_name_input], outputs=[load_chat_model_btn, chat_interface_container]
+    )
+
+    with chat_interface_container:
+        with gr.Row():
+            with gr.Column():
+                ref_audio_chat = gr.Audio(label="Reference Audio", type="filepath")
+            with gr.Column():
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Row():
+                        ref_text_chat = gr.Textbox(
+                            label="Reference Text",
+                            info="Optional: Leave blank to auto-transcribe",
+                            lines=2,
+                            scale=3,
+                        )
+                        ref_text_file_chat = gr.File(
+                            label="Load Reference Text from File (.txt)", file_types=[".txt"], scale=1
+                        )
+                    with gr.Row():
+                        randomize_seed_chat = gr.Checkbox(
+                            label="Randomize Seed",
+                            value=True,
+                            info="Uncheck to use the seed specified.",
+                            scale=3,
+                        )
+                        seed_input_chat = gr.Number(show_label=False, value=0, precision=0, scale=1)
+                    remove_silence_chat = gr.Checkbox(
+                        label="Remove Silences",
+                        value=True,
+                    )
+                    system_prompt_chat = gr.Textbox(
+                        label="System Prompt",
+                        value="You are not an AI assistant, you are whoever the user says you are. You must stay in character. Keep your responses concise since they will be spoken out loud.",
+                        lines=2,
+                    )
+
+        chatbot_interface = gr.Chatbot(label="Conversation", type="messages")
+
+        with gr.Row():
+            with gr.Column():
+                audio_input_chat = gr.Microphone(
+                    label="Speak your message",
+                    type="filepath",
+                )
+                audio_output_chat = gr.Audio(autoplay=True)
+            with gr.Column():
+                text_input_chat = gr.Textbox(
+                    label="Type your message",
+                    lines=1,
+                )
+                send_btn_chat = gr.Button("Send Message")
+                clear_btn_chat = gr.Button("Clear Conversation")
+
+        # Modify process_audio_input to generate user input
+        @gpu_decorator
+        def process_audio_input(conv_state, audio_path, text):
+            """Handle audio or text input from user"""
+
+            if not audio_path and not text.strip():
+                return conv_state
+
+            if audio_path:
+                text = preprocess_ref_audio_text(audio_path, text)[1]
+            if not text.strip():
+                return conv_state
+
+            conv_state.append({"role": "user", "content": text})
+            return conv_state
+
+        # Use model and tokenizer from state to get text response
+        @gpu_decorator
+        def generate_text_response(conv_state, system_prompt):
+            """Generate text response from AI"""
+
+            system_prompt_state = [{"role": "system", "content": system_prompt}]
+            response = chat_model_inference(system_prompt_state + conv_state, chat_model_state, chat_tokenizer_state)
+
+            conv_state.append({"role": "assistant", "content": response})
+            return conv_state
+
+        @gpu_decorator
+        def generate_audio_response(conv_state, ref_audio, ref_text, remove_silence, randomize_seed, seed_input):
+            """Generate TTS audio for AI response"""
+            if not conv_state or not ref_audio:
+                return None, ref_text, seed_input
+
+            last_ai_response = conv_state[-1]["content"]
+            if not last_ai_response or conv_state[-1]["role"] != "assistant":
+                return None, ref_text, seed_input
+
+            if randomize_seed:
+                seed_input = np.random.randint(0, 2**31 - 1)
+
+            audio_result, _, ref_text_out, used_seed = infer(
+                ref_audio,
+                ref_text,
+                last_ai_response,
+                tts_model_choice,
+                remove_silence,
+                seed=seed_input,
+                cross_fade_duration=0.15,
+                speed=1.0,
+                show_info=print,  # show_info=print no pull to top when generating
+            )
+            return audio_result, ref_text_out, used_seed
+
+        def clear_conversation():
+            """Reset the conversation"""
+            return [], None
+
+        ref_text_file_chat.upload(
+            load_text_from_file,
+            inputs=[ref_text_file_chat],
+            outputs=[ref_text_chat],
+        )
+
+        for user_operation in [audio_input_chat.stop_recording, text_input_chat.submit, send_btn_chat.click]:
+            user_operation(
+                process_audio_input,
+                inputs=[chatbot_interface, audio_input_chat, text_input_chat],
+                outputs=[chatbot_interface],
+            ).then(
+                generate_text_response,
+                inputs=[chatbot_interface, system_prompt_chat],
+                outputs=[chatbot_interface],
+            ).then(
+                generate_audio_response,
+                inputs=[
+                    chatbot_interface,
+                    ref_audio_chat,
+                    ref_text_chat,
+                    remove_silence_chat,
+                    randomize_seed_chat,
+                    seed_input_chat,
+                ],
+                outputs=[audio_output_chat, ref_text_chat, seed_input_chat],
+            ).then(
+                lambda: [None, None],
+                None,
+                [audio_input_chat, text_input_chat],
+            )
+
+        # Handle clear button or system prompt change and reset conversation
+        for user_operation in [clear_btn_chat.click, system_prompt_chat.change, chatbot_interface.clear]:
+            user_operation(
+                clear_conversation,
+                outputs=[chatbot_interface, audio_output_chat],
+            )
+
+
+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 initial chunk generation and podcast app exploration
+* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation & voice chat
+""")
+
+
+with gr.Blocks() as app:
+    gr.Markdown(
+        f"""
+# F5-TTS Demo Space
+
+This is {"a local web UI for [F5-TTS](https://github.com/SWivid/F5-TTS)" if not USING_SPACES else "an online demo for [F5-TTS](https://github.com/SWivid/F5-TTS)"} with advanced batch processing support. This app supports the following TTS models:
+
+* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)
+* [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)
+
+The checkpoints currently support English and Chinese.
+
+If you're having issues, try converting your reference audio to WAV or MP3, clipping it to 12s with  ✂  in the bottom right corner (otherwise might have non-optimal auto-trimmed result).
+
+**NOTE: Reference text will be automatically transcribed with Whisper if not provided. For best results, keep your reference clips short (<12s). Ensure the audio is fully uploaded before generating.**
+"""
+    )
+
+    last_used_custom = files("f5_tts").joinpath("infer/.cache/last_used_custom_model_info_v1.txt")
+
+    def load_last_used_custom():
+        try:
+            custom = []
+            with open(last_used_custom, "r", encoding="utf-8") as f:
+                for line in f:
+                    custom.append(line.strip())
+            return custom
+        except FileNotFoundError:
+            last_used_custom.parent.mkdir(parents=True, exist_ok=True)
+            return DEFAULT_TTS_MODEL_CFG
+
+    def switch_tts_model(new_choice):
+        global tts_model_choice
+        if new_choice == "Custom":  # override in case webpage is refreshed
+            custom_ckpt_path, custom_vocab_path, custom_model_cfg = load_last_used_custom()
+            tts_model_choice = ("Custom", custom_ckpt_path, custom_vocab_path, custom_model_cfg)
+            return (
+                gr.update(visible=True, value=custom_ckpt_path),
+                gr.update(visible=True, value=custom_vocab_path),
+                gr.update(visible=True, value=custom_model_cfg),
+            )
+        else:
+            tts_model_choice = new_choice
+            return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+    def set_custom_model(custom_ckpt_path, custom_vocab_path, custom_model_cfg):
+        global tts_model_choice
+        tts_model_choice = ("Custom", custom_ckpt_path, custom_vocab_path, custom_model_cfg)
+        with open(last_used_custom, "w", encoding="utf-8") as f:
+            f.write(custom_ckpt_path + "\n" + custom_vocab_path + "\n" + custom_model_cfg + "\n")
+
+    with gr.Row():
+        if not USING_SPACES:
+            choose_tts_model = gr.Radio(
+                choices=[DEFAULT_TTS_MODEL, "E2-TTS", "Custom"], label="Choose TTS Model", value=DEFAULT_TTS_MODEL
+            )
+        else:
+            choose_tts_model = gr.Radio(
+                choices=[DEFAULT_TTS_MODEL, "E2-TTS"], label="Choose TTS Model", value=DEFAULT_TTS_MODEL
+            )
+        custom_ckpt_path = gr.Dropdown(
+            choices=[DEFAULT_TTS_MODEL_CFG[0]],
+            value=load_last_used_custom()[0],
+            allow_custom_value=True,
+            label="Model: local_path | hf://user_id/repo_id/model_ckpt",
+            visible=False,
+        )
+        custom_vocab_path = gr.Dropdown(
+            choices=[DEFAULT_TTS_MODEL_CFG[1]],
+            value=load_last_used_custom()[1],
+            allow_custom_value=True,
+            label="Vocab: local_path | hf://user_id/repo_id/vocab_file",
+            visible=False,
+        )
+        custom_model_cfg = gr.Dropdown(
+            choices=[
+                DEFAULT_TTS_MODEL_CFG[2],
+                json.dumps(
+                    dict(
+                        dim=1024,
+                        depth=22,
+                        heads=16,
+                        ff_mult=2,
+                        text_dim=512,
+                        text_mask_padding=False,
+                        conv_layers=4,
+                        pe_attn_head=1,
+                    )
+                ),
+                json.dumps(
+                    dict(
+                        dim=768,
+                        depth=18,
+                        heads=12,
+                        ff_mult=2,
+                        text_dim=512,
+                        text_mask_padding=False,
+                        conv_layers=4,
+                        pe_attn_head=1,
+                    )
+                ),
+            ],
+            value=load_last_used_custom()[2],
+            allow_custom_value=True,
+            label="Config: in a dictionary form",
+            visible=False,
+        )
+
+    choose_tts_model.change(
+        switch_tts_model,
+        inputs=[choose_tts_model],
+        outputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
+        show_progress="hidden",
+    )
+    custom_ckpt_path.change(
+        set_custom_model,
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
+        show_progress="hidden",
+    )
+    custom_vocab_path.change(
+        set_custom_model,
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
+        show_progress="hidden",
+    )
+    custom_model_cfg.change(
+        set_custom_model,
+        inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],
+        show_progress="hidden",
+    )
+
+    gr.TabbedInterface(
+        [app_tts, app_multistyle, app_chat, app_credits],
+        ["Basic-TTS", "Multi-Speech", "Voice-Chat", "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")
+@click.option(
+    "--root_path",
+    "-r",
+    default=None,
+    type=str,
+    help='The root path (or "mount point") of the application, if it\'s not served from the root ("/") of the domain. Often used when the application is behind a reverse proxy that forwards requests to the application, e.g. set "/myapp" or full URL for application served at "https://example.com/myapp".',
+)
+@click.option(
+    "--inbrowser",
+    "-i",
+    is_flag=True,
+    default=False,
+    help="Automatically launch the interface in the default web browser",
+)
+def main(port, host, share, api, root_path, inbrowser):
+    global app
+    print("Starting app...")
+    app.queue(api_open=api).launch(
+        server_name=host,
+        server_port=port,
+        share=share,
+        show_api=api,
+        root_path=root_path,
+        inbrowser=inbrowser,
+    )
+
+
+if __name__ == "__main__":
+    if not USING_SPACES:
+        main()
+    else:
+        app.queue().launch()

src/f5_tts/infer/speech_edit.py

@@ -0,0 +1,205 @@
+import os
+
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # for MPS device compatibility
+
+from importlib.resources import files
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from cached_path import cached_path
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+
+from f5_tts.infer.utils_infer import load_checkpoint, load_vocoder, save_spectrogram
+from f5_tts.model import CFM
+from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer
+
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "xpu"
+    if torch.xpu.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+
+# ---------------------- infer setting ---------------------- #
+
+seed = None  # int | None
+
+exp_name = "F5TTS_v1_Base"  # F5TTS_v1_Base | E2TTS_Base
+ckpt_step = 1250000
+
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+ode_method = "euler"  # euler | midpoint
+sway_sampling_coef = -1.0
+speed = 1.0
+target_rms = 0.1
+
+
+model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{exp_name}.yaml")))
+model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+model_arc = model_cfg.model.arch
+
+dataset_name = model_cfg.datasets.name
+tokenizer = model_cfg.model.tokenizer
+
+mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+target_sample_rate = model_cfg.model.mel_spec.target_sample_rate
+n_mel_channels = model_cfg.model.mel_spec.n_mel_channels
+hop_length = model_cfg.model.mel_spec.hop_length
+win_length = model_cfg.model.mel_spec.win_length
+n_fft = model_cfg.model.mel_spec.n_fft
+
+
+# ckpt_path = str(files("f5_tts").joinpath("../../")) + f"/ckpts/{exp_name}/model_{ckpt_step}.safetensors"
+ckpt_path = str(cached_path(f"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.safetensors"))
+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 "src/f5_tts/infer/examples/basic/basic_ref_en.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 = str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.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 = "src/f5_tts/infer/examples/basic/basic_ref_zh.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 mel_spec_type == "vocos":
+    vocoder_local_path = "../checkpoints/charactr/vocos-mel-24khz"
+elif mel_spec_type == "bigvgan":
+    vocoder_local_path = "../checkpoints/bigvgan_v2_24khz_100band_256x"
+vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)
+
+# Tokenizer
+vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)
+
+# Model
+model = CFM(
+    transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+    mel_spec_kwargs=dict(
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        n_mel_channels=n_mel_channels,
+        target_sample_rate=target_sample_rate,
+        mel_spec_type=mel_spec_type,
+    ),
+    odeint_kwargs=dict(
+        method=ode_method,
+    ),
+    vocab_char_map=vocab_char_map,
+).to(device)
+
+dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+model = load_checkpoint(model, ckpt_path, device, dtype=dtype, 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.to(torch.float32)
+    generated = generated[:, ref_audio_len:, :]
+    gen_mel_spec = generated.permute(0, 2, 1)
+    if mel_spec_type == "vocos":
+        generated_wave = vocoder.decode(gen_mel_spec).cpu()
+    elif mel_spec_type == "bigvgan":
+        generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()
+
+    if rms < target_rms:
+        generated_wave = generated_wave * rms / target_rms
+
+    save_spectrogram(gen_mel_spec[0].cpu().numpy(), f"{output_dir}/speech_edit_out.png")
+    torchaudio.save(f"{output_dir}/speech_edit_out.wav", generated_wave, target_sample_rate)
+    print(f"Generated wav: {generated_wave.shape}")

src/f5_tts/infer/utils_infer.py

@@ -0,0 +1,765 @@
+# A unified script for inference process
+# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format
+import os
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+
+os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"  # for MPS device compatibility
+sys.path.append(f"{os.path.dirname(os.path.abspath(__file__))}/../../third_party/BigVGAN/")
+
+import hashlib
+import re
+import tempfile
+from importlib.resources import files
+
+import matplotlib
+
+
+matplotlib.use("Agg")
+
+import matplotlib.pylab as plt
+import numpy as np
+import torch
+import torchaudio
+import tqdm
+import requests
+from huggingface_hub import hf_hub_download
+from pydub import AudioSegment, silence
+from transformers import pipeline
+from vocos import Vocos
+
+from f5_tts.model import CFM
+from f5_tts.model.modules import MelSpec
+from f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer
+from f5_tts.infer.cls_tokenizer_v2 import cls_tokenize_text
+
+
+_ref_audio_cache = {}
+_ref_text_cache = {}
+
+
+def _load_state_dict_resilient(model: torch.nn.Module, state_dict: dict):
+    """
+    Load a state_dict while tolerating shape mismatches (e.g. when switching vocabularies).
+    Any parameters whose shapes do not match the current model will be skipped so that
+    those layers keep their freshly initialized weights instead of raising.
+    """
+    model_state = model.state_dict()
+    filtered_state = {}
+    mismatched = {}
+
+    for key, weight in state_dict.items():
+        target = model_state.get(key)
+        if target is None:
+            continue
+        if hasattr(target, "shape") and hasattr(weight, "shape") and target.shape != weight.shape:
+            # If only the vocab size dimension differs, align by slicing/padding so embeddings still load.
+            if (
+                len(target.shape) == len(weight.shape)
+                and target.shape[1:] == weight.shape[1:]
+                and key.endswith("text_embed.weight")
+            ):
+                new_weight = target.clone()
+                rows = min(target.shape[0], weight.shape[0])
+                new_weight[:rows] = weight[:rows]
+                filtered_state[key] = new_weight
+                if target.shape[0] != weight.shape[0]:
+                    print(
+                        f"Info: resized {key} from {tuple(weight.shape)} to {tuple(target.shape)} "
+                        f"(copied {rows} rows)."
+                    )
+                continue
+            mismatched[key] = (tuple(weight.shape), tuple(target.shape))
+            continue
+        filtered_state[key] = weight
+
+    missing, unexpected = model.load_state_dict(filtered_state, strict=False)
+    if mismatched:
+        mismatch_info = ", ".join(f"{k}: {src} -> {dst}" for k, (src, dst) in mismatched.items())
+        print(f"Warning: skipped loading parameters with shape mismatch ({mismatch_info}).")
+    if missing:
+        print(f"Warning: missing parameters not loaded from checkpoint: {missing}")
+    if unexpected:
+        print(f"Warning: unexpected parameters in checkpoint were ignored: {unexpected}")
+
+device = (
+    "cuda"
+    if torch.cuda.is_available()
+    else "xpu"
+    if torch.xpu.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
+
+tempfile_kwargs = {"delete_on_close": False} if sys.version_info >= (3, 12) else {"delete": False}
+
+# -----------------------------------------
+
+target_sample_rate = 24000
+n_mel_channels = 100
+hop_length = 256
+win_length = 1024
+n_fft = 1024
+mel_spec_type = "vocos"
+target_rms = 0.1
+cross_fade_duration = 0.15
+ode_method = "euler"
+nfe_step = 32  # 16, 32
+cfg_strength = 2.0
+sway_sampling_coef = -1.0
+speed = 1.0
+fix_duration = None
+
+# -----------------------------------------
+
+
+# chunk text into smaller pieces
+
+
+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
+
+
+def tokenize_texts(
+    text_list,
+    tokenizer="pinyin",
+    cls_language=None,
+    cls_server_url=None,
+    cls_timeout=5.0,
+    cls_tokenizer_fn=None,
+):
+    tokenizer = (tokenizer or "pinyin").strip().lower()
+    if tokenizer == "pinyin":
+        return convert_char_to_pinyin(text_list)
+    if tokenizer == "char":
+        return [list(t) for t in text_list]
+    if tokenizer == "cls":
+        if not cls_language:
+            raise ValueError("cls_language must be set when tokenizer='cls'.")
+        if cls_tokenizer_fn is not None:
+            results = []
+            for text in text_list:
+                cls_tokens = cls_tokenizer_fn(text, cls_language)
+                if not isinstance(cls_tokens, list) or len(cls_tokens) == 0:
+                    raise RuntimeError("CLS tokenizer function returned empty tokens.")
+                results.append(cls_tokens)
+            return results
+        if cls_server_url:
+            results = []
+            for text in text_list:
+                try:
+                    resp = requests.post(
+                        cls_server_url,
+                        json={"text": text, "language": cls_language},
+                        timeout=cls_timeout,
+                    )
+                except Exception as exc:  # noqa: BLE001
+                    raise RuntimeError(f"CLS server request failed: {exc}") from exc
+                if resp.status_code != 200:
+                    raise RuntimeError(f"CLS server error {resp.status_code}: {resp.text}")
+                try:
+                    data = resp.json()
+                except Exception as exc:  # noqa: BLE001
+                    raise RuntimeError(f"CLS server returned non-JSON response: {exc}") from exc
+                cls_tokens = data.get("cls")
+                if not isinstance(cls_tokens, list) or len(cls_tokens) == 0:
+                    raise RuntimeError("CLS server returned empty tokens.")
+                results.append(cls_tokens)
+            return results
+        results = []
+        for text in text_list:
+            cls_tokens = cls_tokenize_text(text, cls_language)
+            if not isinstance(cls_tokens, list) or len(cls_tokens) == 0:
+                raise RuntimeError("CLS tokenizer returned empty tokens.")
+            results.append(cls_tokens)
+        return results
+    raise ValueError(f"Unsupported tokenizer: {tokenizer}")
+
+
+# load vocoder
+def load_vocoder(vocoder_name="vocos", is_local=False, local_path="", device=device, hf_cache_dir=None):
+    if vocoder_name == "vocos":
+        # vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz").to(device)
+        if is_local:
+            print(f"Load vocos from local path {local_path}")
+            config_path = f"{local_path}/config.yaml"
+            model_path = f"{local_path}/pytorch_model.bin"
+        else:
+            print("Download Vocos from huggingface charactr/vocos-mel-24khz")
+            repo_id = "charactr/vocos-mel-24khz"
+            config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="config.yaml")
+            model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename="pytorch_model.bin")
+        vocoder = Vocos.from_hparams(config_path)
+        state_dict = torch.load(model_path, map_location="cpu", weights_only=True)
+        from vocos.feature_extractors import EncodecFeatures
+
+        if isinstance(vocoder.feature_extractor, EncodecFeatures):
+            encodec_parameters = {
+                "feature_extractor.encodec." + key: value
+                for key, value in vocoder.feature_extractor.encodec.state_dict().items()
+            }
+            state_dict.update(encodec_parameters)
+        vocoder.load_state_dict(state_dict)
+        vocoder = vocoder.eval().to(device)
+    elif vocoder_name == "bigvgan":
+        try:
+            from third_party.BigVGAN import bigvgan
+        except ImportError:
+            print("You need to follow the README to init submodule and change the BigVGAN source code.")
+        if is_local:
+            # download generator from https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x/tree/main
+            vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)
+        else:
+            vocoder = bigvgan.BigVGAN.from_pretrained(
+                "nvidia/bigvgan_v2_24khz_100band_256x", use_cuda_kernel=False, cache_dir=hf_cache_dir
+            )
+
+        vocoder.remove_weight_norm()
+        vocoder = vocoder.eval().to(device)
+    return vocoder
+
+
+# load asr pipeline
+
+asr_pipe = None
+
+
+def initialize_asr_pipeline(device: str = device, dtype=None):
+    if dtype is None:
+        dtype = (
+            torch.float16
+            if "cuda" in device
+            and torch.cuda.get_device_properties(device).major >= 7
+            and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+            else torch.float32
+        )
+    global asr_pipe
+    asr_pipe = pipeline( 
+        "automatic-speech-recognition",
+        model="openai/whisper-large-v3-turbo",
+        torch_dtype=dtype,
+        device=device,
+    )
+
+
+# transcribe
+
+
+def transcribe(ref_audio, language=None):
+    global asr_pipe
+    if asr_pipe is None:
+        initialize_asr_pipeline(device=device)
+    return asr_pipe(
+        ref_audio,
+        chunk_length_s=30,
+        batch_size=128,
+        generate_kwargs={"task": "transcribe", "language": language} if language else {"task": "transcribe"},
+        return_timestamps=False,
+    )["text"].strip()
+
+
+# load model checkpoint for inference
+
+
+def load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):
+    if dtype is None:
+        dtype = torch.float32
+        # dtype = (
+        #     torch.float16
+        #     if "cuda" in device
+        #     and torch.cuda.get_device_properties(device).major >= 6
+        #     and not torch.cuda.get_device_name().endswith("[ZLUDA]")
+        #     else torch.float32
+        # )
+    model = model.to(dtype)
+
+    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, weights_only=True)
+
+    if use_ema:
+        if ckpt_type == "safetensors":
+            checkpoint = {"ema_model_state_dict": checkpoint}
+        checkpoint["model_state_dict"] = {
+            k.replace("ema_model.", ""): v
+            for k, v in checkpoint["ema_model_state_dict"].items()
+            if k not in ["initted", "step"]
+        }
+
+        # patch for backward compatibility, 305e3ea
+        for key in ["mel_spec.mel_stft.mel_scale.fb", "mel_spec.mel_stft.spectrogram.window"]:
+            if key in checkpoint["model_state_dict"]:
+                del checkpoint["model_state_dict"][key]
+
+        _load_state_dict_resilient(model, checkpoint["model_state_dict"])
+    else:
+        if ckpt_type == "safetensors":
+            checkpoint = {"model_state_dict": checkpoint}
+        _load_state_dict_resilient(model, checkpoint["model_state_dict"])
+
+    del checkpoint
+    torch.cuda.empty_cache()
+
+    return model.to(device)
+
+
+# load model for inference
+
+
+def load_model(
+    model_cls,
+    model_cfg,
+    ckpt_path,
+    mel_spec_type=mel_spec_type,
+    vocab_file="",
+    ode_method=ode_method,
+    use_ema=True,
+    device=device,
+):
+    if vocab_file == "":
+        vocab_file = str(files("f5_tts").joinpath("infer/examples/vocab.txt"))
+    tokenizer = "custom"
+
+    print("\nvocab : ", vocab_file)
+    print("token : ", tokenizer)
+    print("model : ", ckpt_path, "\n")
+
+    vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer)
+    model = CFM(
+        transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),
+        mel_spec_kwargs=dict(
+            n_fft=n_fft,
+            hop_length=hop_length,
+            win_length=win_length,
+            n_mel_channels=n_mel_channels,
+            target_sample_rate=target_sample_rate,
+            mel_spec_type=mel_spec_type,
+        ),
+        odeint_kwargs=dict(
+            method=ode_method,
+        ),
+        vocab_char_map=vocab_char_map,
+    ).to(device)
+
+    dtype = torch.float32 if mel_spec_type == "bigvgan" else None
+    model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)
+
+    return model
+
+
+def remove_silence_edges(audio, silence_threshold=-42):
+    # Remove silence from the start
+    non_silent_start_idx = silence.detect_leading_silence(audio, silence_threshold=silence_threshold)
+    audio = audio[non_silent_start_idx:]
+
+    # Remove silence from the end
+    non_silent_end_duration = audio.duration_seconds
+    for ms in reversed(audio):
+        if ms.dBFS > silence_threshold:
+            break
+        non_silent_end_duration -= 0.001
+    trimmed_audio = audio[: int(non_silent_end_duration * 1000)]
+
+    return trimmed_audio
+
+
+# preprocess reference audio and text
+
+
+def preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print):
+    show_info("Converting audio...")
+
+    # Compute a hash of the reference audio file
+    with open(ref_audio_orig, "rb") as audio_file:
+        audio_data = audio_file.read()
+        audio_hash = hashlib.md5(audio_data).hexdigest()
+
+    global _ref_audio_cache
+
+    if audio_hash in _ref_audio_cache:
+        show_info("Using cached preprocessed reference audio...")
+        ref_audio = _ref_audio_cache[audio_hash]
+
+    else:  # first pass, do preprocess
+        with tempfile.NamedTemporaryFile(suffix=".wav", **tempfile_kwargs) as f:
+            temp_path = f.name
+
+        aseg = AudioSegment.from_file(ref_audio_orig)
+
+        # 1. try to find long silence for clipping
+        non_silent_segs = silence.split_on_silence(
+            aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000, seek_step=10
+        )
+        non_silent_wave = AudioSegment.silent(duration=0)
+        for non_silent_seg in non_silent_segs:
+            if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                show_info("Audio is over 12s, clipping short. (1)")
+                break
+            non_silent_wave += non_silent_seg
+
+        # 2. try to find short silence for clipping if 1. failed
+        if len(non_silent_wave) > 12000:
+            non_silent_segs = silence.split_on_silence(
+                aseg, min_silence_len=100, silence_thresh=-40, keep_silence=1000, seek_step=10
+            )
+            non_silent_wave = AudioSegment.silent(duration=0)
+            for non_silent_seg in non_silent_segs:
+                if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:
+                    show_info("Audio is over 12s, clipping short. (2)")
+                    break
+                non_silent_wave += non_silent_seg
+
+        aseg = non_silent_wave
+
+        # 3. if no proper silence found for clipping
+        if len(aseg) > 12000:
+            aseg = aseg[:12000]
+            show_info("Audio is over 12s, clipping short. (3)")
+
+        aseg = remove_silence_edges(aseg) + AudioSegment.silent(duration=50)
+        aseg.export(temp_path, format="wav")
+        ref_audio = temp_path
+
+        # Cache the processed reference audio
+        _ref_audio_cache[audio_hash] = ref_audio
+
+    if not ref_text.strip():
+        global _ref_text_cache
+        if audio_hash in _ref_text_cache:
+            # Use cached asr transcription
+            show_info("Using cached reference text...")
+            ref_text = _ref_text_cache[audio_hash]
+        else:
+            show_info("No reference text provided, transcribing reference audio...")
+            ref_text = transcribe(ref_audio)
+            # Cache the transcribed text (not caching custom ref_text, enabling users to do manual tweak)
+            _ref_text_cache[audio_hash] = ref_text
+    else:
+        show_info("Using custom reference text...")
+
+    # Ensure ref_text ends with a proper sentence-ending punctuation
+    if not ref_text.endswith(". ") and not ref_text.endswith("。"):
+        if ref_text.endswith("."):
+            ref_text += " "
+        else:
+            ref_text += ". "
+
+    print("\nref_text  ", ref_text)
+
+    return ref_audio, ref_text
+
+
+# infer process: chunk text -> infer batches [i.e. infer_batch_process()]
+
+
+def infer_process(
+    ref_audio,
+    ref_text,
+    gen_text,
+    model_obj,
+    vocoder,
+    mel_spec_type=mel_spec_type,
+    show_info=print,
+    progress=tqdm,
+    target_rms=target_rms,
+    cross_fade_duration=cross_fade_duration,
+    nfe_step=nfe_step,
+    cfg_strength=cfg_strength,
+    sway_sampling_coef=sway_sampling_coef,
+    speed=speed,
+    fix_duration=fix_duration,
+    device=device,
+    tokenizer="pinyin",
+    cls_language=None,
+    cls_server_url=None,
+    cls_timeout=5.0,
+    cls_tokenizer_fn=None,
+):
+    # Split the input text into batches
+    audio, sr = torchaudio.load(ref_audio)
+    max_chars = int(len(ref_text.encode("utf-8")) / (audio.shape[-1] / sr) * (22 - audio.shape[-1] / sr) * speed)
+    gen_text_batches = chunk_text(gen_text, max_chars=max_chars)
+    for i, gen_text in enumerate(gen_text_batches):
+        print(f"gen_text {i}", gen_text)
+    print("\n")
+
+    show_info(f"Generating audio in {len(gen_text_batches)} batches...")
+    return next(
+        infer_batch_process(
+            (audio, sr),
+            ref_text,
+            gen_text_batches,
+            model_obj,
+            vocoder,
+            mel_spec_type=mel_spec_type,
+            progress=progress,
+            target_rms=target_rms,
+            cross_fade_duration=cross_fade_duration,
+            nfe_step=nfe_step,
+            cfg_strength=cfg_strength,
+            sway_sampling_coef=sway_sampling_coef,
+            speed=speed,
+            fix_duration=fix_duration,
+            device=device,
+            tokenizer=tokenizer,
+            cls_language=cls_language,
+            cls_server_url=cls_server_url,
+            cls_timeout=cls_timeout,
+            cls_tokenizer_fn=cls_tokenizer_fn,
+        )
+    )
+
+
+# infer batches
+
+
+def infer_batch_process(
+    ref_audio,
+    ref_text,
+    gen_text_batches,
+    model_obj,
+    vocoder,
+    mel_spec_type="vocos",
+    progress=tqdm,
+    target_rms=0.1,
+    cross_fade_duration=0.15,
+    nfe_step=32,
+    cfg_strength=2.0,
+    sway_sampling_coef=-1,
+    speed=1,
+    fix_duration=None,
+    device=None,
+    streaming=False,
+    chunk_size=2048,
+    tokenizer="pinyin",
+    cls_language=None,
+    cls_server_url=None,
+    cls_timeout=5.0,
+    cls_tokenizer_fn=None,
+):
+    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(orig_freq=sr, new_freq=target_sample_rate)
+        audio = resampler(audio)
+    audio = audio.to(device)
+
+    mel_spectrogram = MelSpec(
+        n_fft=n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        n_mel_channels=n_mel_channels,
+        target_sample_rate=target_sample_rate,
+        mel_spec_type=mel_spec_type,
+    )
+    ref_mel = mel_spectrogram(audio)
+    if ref_mel.dim() == 3:
+        ref_mel = ref_mel[0]
+    ref_mel_len = ref_mel.shape[-1]
+
+    generated_waves = []
+    spectrograms = []
+
+    if len(ref_text[-1].encode("utf-8")) == 1:
+        ref_text = ref_text + " "
+
+    def process_batch(gen_text):
+        local_speed = speed
+        if len(gen_text.encode("utf-8")) < 10:
+            local_speed = 0.3
+
+        # Prepare the text
+        text_list = [ref_text + gen_text]
+        final_text_list = tokenize_texts(
+            text_list,
+            tokenizer=tokenizer,
+            cls_language=cls_language,
+            cls_server_url=cls_server_url,
+            cls_timeout=cls_timeout,
+            cls_tokenizer_fn=cls_tokenizer_fn,
+        )
+
+        ref_audio_len = audio.shape[-1] // hop_length
+        if fix_duration is not None:
+            duration = int(fix_duration * target_sample_rate / hop_length)
+        else:
+            # Calculate duration
+            ref_text_len = len(ref_text.encode("utf-8"))
+            gen_text_len = len(gen_text.encode("utf-8"))
+            duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / local_speed)
+
+        # inference
+        with torch.inference_mode():
+            generated, _ = model_obj.sample(
+                cond=audio,
+                text=final_text_list,
+                duration=duration,
+                steps=nfe_step,
+                cfg_strength=cfg_strength,
+                sway_sampling_coef=sway_sampling_coef,
+            )
+            del _
+
+            generated = generated.to(torch.float32)  # generated mel spectrogram [B, T, n_mel]
+            gen = generated[0]
+
+            # Align ref mel inside generated mel to find the boundary, then cut the robotic head.
+            ref_mel_local = ref_mel[:, :ref_mel_len]
+            gen_mel = gen.transpose(0, 1)  # [n_mel, T]
+            if ref_mel_local.shape[0] != gen_mel.shape[0] and ref_mel_local.shape[1] == gen_mel.shape[0]:
+                ref_mel_local = ref_mel_local.transpose(0, 1)
+            cut_idx = 0
+            if gen_mel.shape[1] > ref_mel_local.shape[1]:
+                ref_len = ref_mel_local.shape[1]
+                gen_unfold = gen_mel.unfold(1, ref_len, 1)  # [n_mel, T-ref_len+1, ref_len]
+                diff = gen_unfold - ref_mel_local.unsqueeze(1)
+                mse = torch.mean(diff * diff, dim=(0, 2))
+                best = int(torch.argmin(mse).item())
+                cut_idx = max(0, int((best + ref_len) * hop_length))
+
+            gen = gen.unsqueeze(0)
+            gen_mel_spec = gen.permute(0, 2, 1).to(torch.float32)
+            if mel_spec_type == "vocos":
+                generated_wave = vocoder.decode(gen_mel_spec).cpu()
+            elif mel_spec_type == "bigvgan":
+                generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()
+            if rms < target_rms:
+                generated_wave = generated_wave * rms / target_rms
+
+            # wav -> numpy
+            wave = generated_wave.squeeze(0).to(torch.float32)
+            if cut_idx > 0 and cut_idx < wave.numel():
+                wave = wave[cut_idx:]
+            generated_wave = wave.unsqueeze(0)
+            generated_wave = generated_wave.squeeze().cpu().numpy()
+
+            if streaming:
+                for j in range(0, len(generated_wave), chunk_size):
+                    yield generated_wave[j : j + chunk_size], target_sample_rate
+            else:
+                generated_cpu = gen_mel_spec[0].cpu().numpy()
+                del generated
+                yield generated_wave, generated_cpu
+
+    if streaming:
+        for gen_text in progress.tqdm(gen_text_batches) if progress is not None else gen_text_batches:
+            for chunk in process_batch(gen_text):
+                yield chunk
+    else:
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(process_batch, gen_text) for gen_text in gen_text_batches]
+            for future in progress.tqdm(futures) if progress is not None else futures:
+                result = future.result()
+                if result:
+                    generated_wave, generated_mel_spec = next(result)
+                    generated_waves.append(generated_wave)
+                    spectrograms.append(generated_mel_spec)
+
+        if generated_waves:
+            if cross_fade_duration <= 0:
+                # Simply concatenate
+                final_wave = np.concatenate(generated_waves)
+            else:
+                # Combine all generated waves with cross-fading
+                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
+
+            # Create a combined spectrogram
+            combined_spectrogram = np.concatenate(spectrograms, axis=1)
+
+            yield final_wave, target_sample_rate, combined_spectrogram
+
+        else:
+            yield None, target_sample_rate, None
+
+
+# remove silence from generated wav
+
+
+def remove_silence_for_generated_wav(filename):
+    aseg = AudioSegment.from_file(filename)
+    non_silent_segs = silence.split_on_silence(
+        aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500, seek_step=10
+    )
+    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(filename, format="wav")
+
+
+# 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()

src/f5_tts/model/__init__.py

@@ -0,0 +1,5 @@
+from f5_tts.model.backbones.dit import DiT
+from f5_tts.model.backbones.mmdit import MMDiT
+from f5_tts.model.backbones.unett import UNetT
+from f5_tts.model.cfm import CFM
+__all__ = ["CFM", "UNetT", "DiT", "MMDiT"]

src/f5_tts/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
+- 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
+- stable diffusion 3 block 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

src/f5_tts/model/backbones/dit.py

@@ -0,0 +1,318 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    AdaLayerNorm_Final,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    DiTBlock,
+    TimestepEmbedding,
+    get_pos_embed_indices,
+    precompute_freqs_cis,
+)
+
+
+# Text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(
+        self, text_num_embeds, text_dim, mask_padding=True, average_upsampling=False, conv_layers=0, conv_mult=2
+    ):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        self.mask_padding = mask_padding  # mask filler and batch padding tokens or not
+        self.average_upsampling = average_upsampling  # zipvoice-style text late average upsampling (after text encoder)
+        if average_upsampling:
+            assert mask_padding, "text_embedding_average_upsampling requires text_mask_padding to be True"
+
+        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 average_upsample_text_by_mask(self, text, text_mask, audio_mask):
+        batch, text_len, text_dim = text.shape
+
+        if audio_mask is None:
+            audio_mask = torch.ones_like(text_mask, dtype=torch.bool)
+        valid_mask = audio_mask & text_mask
+        audio_lens = audio_mask.sum(dim=1)  # [batch]
+        valid_lens = valid_mask.sum(dim=1)  # [batch]
+
+        upsampled_text = torch.zeros_like(text)
+
+        for i in range(batch):
+            audio_len = audio_lens[i].item()
+            valid_len = valid_lens[i].item()
+
+            if valid_len == 0:
+                continue
+
+            valid_ind = torch.where(valid_mask[i])[0]
+            valid_data = text[i, valid_ind, :]  # [valid_len, text_dim]
+
+            base_repeat = audio_len // valid_len
+            remainder = audio_len % valid_len
+
+            indices = []
+            for j in range(valid_len):
+                repeat_count = base_repeat + (1 if j >= valid_len - remainder else 0)
+                indices.extend([j] * repeat_count)
+
+            indices = torch.tensor(indices[:audio_len], device=text.device, dtype=torch.long)
+            upsampled = valid_data[indices]  # [audio_len, text_dim]
+
+            upsampled_text[i, :audio_len, :] = upsampled
+
+        return upsampled_text
+
+    def forward(self, text: int["b nt"], seq_len, drop_text=False, audio_mask: bool["b n"] | None = None):  # noqa: F722
+        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
+        batch, text_len = text.shape[0], text.shape[1]
+        text = F.pad(text, (0, seq_len - text_len), value=0)  # (opt.) if not self.average_upsampling:
+        if self.mask_padding:
+            text_mask = text == 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,), device=text.device, 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
+            if self.mask_padding:
+                text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)
+                for block in self.text_blocks:
+                    text = block(text)
+                    text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)
+            else:
+                text = self.text_blocks(text)
+
+        if self.average_upsampling:
+            text = self.average_upsample_text_by_mask(text, ~text_mask, audio_mask)
+
+        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):  # noqa: F722
+        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,
+        text_mask_padding=True,
+        text_embedding_average_upsampling=False,
+        qk_norm=None,
+        conv_layers=0,
+        pe_attn_head=None,
+        attn_backend="torch",  # "torch" | "flash_attn"
+        attn_mask_enabled=False,
+        long_skip_connection=False,
+        checkpoint_activations=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,
+            mask_padding=text_mask_padding,
+            average_upsampling=text_embedding_average_upsampling,
+            conv_layers=conv_layers,
+        )
+        self.text_cond, self.text_uncond = None, None  # text cache
+        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,
+                    qk_norm=qk_norm,
+                    pe_attn_head=pe_attn_head,
+                    attn_backend=attn_backend,
+                    attn_mask_enabled=attn_mask_enabled,
+                )
+                for _ in range(depth)
+            ]
+        )
+        self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None
+
+        self.norm_out = AdaLayerNorm_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+        self.checkpoint_activations = checkpoint_activations
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Zero-out AdaLN layers in DiT blocks:
+        for block in self.transformer_blocks:
+            nn.init.constant_(block.attn_norm.linear.weight, 0)
+            nn.init.constant_(block.attn_norm.linear.bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.norm_out.linear.weight, 0)
+        nn.init.constant_(self.norm_out.linear.bias, 0)
+        nn.init.constant_(self.proj_out.weight, 0)
+        nn.init.constant_(self.proj_out.bias, 0)
+
+    def ckpt_wrapper(self, module):
+        # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py
+        def ckpt_forward(*inputs):
+            outputs = module(*inputs)
+            return outputs
+
+        return ckpt_forward
+
+    def get_input_embed(
+        self,
+        x,  # b n d
+        cond,  # b n d
+        text,  # b nt
+        drop_audio_cond: bool = False,
+        drop_text: bool = False,
+        cache: bool = True,
+        audio_mask: bool["b n"] | None = None,  # noqa: F722
+    ):
+        seq_len = x.shape[1]
+        if cache:
+            if drop_text:
+                if self.text_uncond is None:
+                    self.text_uncond = self.text_embed(text, seq_len, drop_text=True, audio_mask=audio_mask)
+                text_embed = self.text_uncond
+            else:
+                if self.text_cond is None:
+                    self.text_cond = self.text_embed(text, seq_len, drop_text=False, audio_mask=audio_mask)
+                text_embed = self.text_cond
+        else:
+            text_embed = self.text_embed(text, seq_len, drop_text=drop_text, audio_mask=audio_mask)
+
+        # TODO YASH: bfloat conv which leads to flash attention does not support fp32 error happens
+        x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
+
+        return x
+
+    def clear_cache(self):
+        self.text_cond, self.text_uncond = None, None
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        drop_audio_cond: bool = False,  # cfg for cond audio
+        drop_text: bool = False,  # cfg for text
+        cfg_infer: bool = False,  # cfg inference, pack cond & uncond forward
+        cache: bool = False,
+        skip_flash_attn: bool = False,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning time, text: text, x: noised audio + cond audio + text
+        t = self.time_embed(time)
+
+        if cfg_infer:  # pack cond & uncond forward: b n d -> 2b n d
+            x_cond = self.get_input_embed(
+                x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache, audio_mask=mask
+            )
+            x_uncond = self.get_input_embed(
+                x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache, audio_mask=mask
+            )
+            x = torch.cat((x_cond, x_uncond), dim=0)
+            t = torch.cat((t, t), dim=0)
+            mask = torch.cat((mask, mask), dim=0) if mask is not None else None
+        else:
+            x = self.get_input_embed(
+                x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache, audio_mask=mask
+            )
+
+        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:
+            if self.checkpoint_activations:
+                # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint
+                x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False, skip_flash_attn=skip_flash_attn)
+            else:
+                x = block(x, t, mask=mask, rope=rope, skip_flash_attn=skip_flash_attn)
+
+        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

src/f5_tts/model/backbones/mmdit.py

@@ -0,0 +1,212 @@
+"""
+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 x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    AdaLayerNorm_Final,
+    ConvPositionEmbedding,
+    MMDiTBlock,
+    TimestepEmbedding,
+    get_pos_embed_indices,
+    precompute_freqs_cis,
+)
+
+
+# text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, out_dim, text_num_embeds, mask_padding=True):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim)  # will use 0 as filler token
+
+        self.mask_padding = mask_padding  # mask filler and batch padding tokens or not
+
+        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"]:  # noqa: F722
+        text = text + 1  # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
+        if self.mask_padding:
+            text_mask = text == 0
+
+        if drop_text:  # cfg for text
+            text = torch.zeros_like(text)
+
+        text = self.text_embed(text)  # b nt -> b nt d
+
+        # 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
+
+        if self.mask_padding:
+            text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)
+
+        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):  # noqa: F722
+        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,
+        mel_dim=100,
+        text_num_embeds=256,
+        text_mask_padding=True,
+        qk_norm=None,
+    ):
+        super().__init__()
+
+        self.time_embed = TimestepEmbedding(dim)
+        self.text_embed = TextEmbedding(dim, text_num_embeds, mask_padding=text_mask_padding)
+        self.text_cond, self.text_uncond = None, None  # text cache
+        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,
+                    qk_norm=qk_norm,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm_out = AdaLayerNorm_Final(dim)  # final modulation
+        self.proj_out = nn.Linear(dim, mel_dim)
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # Zero-out AdaLN layers in MMDiT blocks:
+        for block in self.transformer_blocks:
+            nn.init.constant_(block.attn_norm_x.linear.weight, 0)
+            nn.init.constant_(block.attn_norm_x.linear.bias, 0)
+            nn.init.constant_(block.attn_norm_c.linear.weight, 0)
+            nn.init.constant_(block.attn_norm_c.linear.bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.norm_out.linear.weight, 0)
+        nn.init.constant_(self.norm_out.linear.bias, 0)
+        nn.init.constant_(self.proj_out.weight, 0)
+        nn.init.constant_(self.proj_out.bias, 0)
+
+    def get_input_embed(
+        self,
+        x,  # b n d
+        cond,  # b n d
+        text,  # b nt
+        drop_audio_cond: bool = False,
+        drop_text: bool = False,
+        cache: bool = True,
+    ):
+        if cache:
+            if drop_text:
+                if self.text_uncond is None:
+                    self.text_uncond = self.text_embed(text, drop_text=True)
+                c = self.text_uncond
+            else:
+                if self.text_cond is None:
+                    self.text_cond = self.text_embed(text, drop_text=False)
+                c = self.text_cond
+        else:
+            c = self.text_embed(text, drop_text=drop_text)
+        x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)
+
+        return x, c
+
+    def clear_cache(self):
+        self.text_cond, self.text_uncond = None, None
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        drop_audio_cond: bool = False,  # cfg for cond audio
+        drop_text: bool = False,  # cfg for text
+        cfg_infer: bool = False,  # cfg inference, pack cond & uncond forward
+        cache: bool = False,
+    ):
+        batch = x.shape[0]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        if cfg_infer:  # pack cond & uncond forward: b n d -> 2b n d
+            x_cond, c_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache)
+            x_uncond, c_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache)
+            x = torch.cat((x_cond, x_uncond), dim=0)
+            c = torch.cat((c_cond, c_uncond), dim=0)
+            t = torch.cat((t, t), dim=0)
+            mask = torch.cat((mask, mask), dim=0) if mask is not None else None
+        else:
+            x, c = self.get_input_embed(
+                x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache
+            )
+
+        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

src/f5_tts/model/backbones/unett.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 Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from x_transformers import RMSNorm
+from x_transformers.x_transformers import RotaryEmbedding
+
+from f5_tts.model.modules import (
+    Attention,
+    AttnProcessor,
+    ConvNeXtV2Block,
+    ConvPositionEmbedding,
+    FeedForward,
+    TimestepEmbedding,
+    get_pos_embed_indices,
+    precompute_freqs_cis,
+)
+
+
+# Text embedding
+
+
+class TextEmbedding(nn.Module):
+    def __init__(self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2):
+        super().__init__()
+        self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim)  # use 0 as filler token
+
+        self.mask_padding = mask_padding  # mask filler and batch padding tokens or not
+
+        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):  # noqa: F722
+        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
+        batch, text_len = text.shape[0], text.shape[1]
+        text = F.pad(text, (0, seq_len - text_len), value=0)
+        if self.mask_padding:
+            text_mask = text == 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
+            if self.mask_padding:
+                text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)
+                for block in self.text_blocks:
+                    text = block(text)
+                    text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)
+            else:
+                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):  # noqa: F722
+        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,
+        text_mask_padding=True,
+        qk_norm=None,
+        conv_layers=0,
+        pe_attn_head=None,
+        attn_backend="torch",  # "torch" | "flash_attn"
+        attn_mask_enabled=False,
+        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, mask_padding=text_mask_padding, conv_layers=conv_layers
+        )
+        self.text_cond, self.text_uncond = None, None  # text cache
+        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(
+                    pe_attn_head=pe_attn_head,
+                    attn_backend=attn_backend,
+                    attn_mask_enabled=attn_mask_enabled,
+                ),
+                dim=dim,
+                heads=heads,
+                dim_head=dim_head,
+                dropout=dropout,
+                qk_norm=qk_norm,
+            )
+
+            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 get_input_embed(
+        self,
+        x,  # b n d
+        cond,  # b n d
+        text,  # b nt
+        drop_audio_cond: bool = False,
+        drop_text: bool = False,
+        cache: bool = True,
+    ):
+        seq_len = x.shape[1]
+        if cache:
+            if drop_text:
+                if self.text_uncond is None:
+                    self.text_uncond = self.text_embed(text, seq_len, drop_text=True)
+                text_embed = self.text_uncond
+            else:
+                if self.text_cond is None:
+                    self.text_cond = self.text_embed(text, seq_len, drop_text=False)
+                text_embed = self.text_cond
+        else:
+            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)
+
+        return x
+
+    def clear_cache(self):
+        self.text_cond, self.text_uncond = None, None
+
+    def forward(
+        self,
+        x: float["b n d"],  # nosied input audio  # noqa: F722
+        cond: float["b n d"],  # masked cond audio  # noqa: F722
+        text: int["b nt"],  # text  # noqa: F722
+        time: float["b"] | float[""],  # time step  # noqa: F821 F722
+        mask: bool["b n"] | None = None,  # noqa: F722
+        drop_audio_cond: bool = False,  # cfg for cond audio
+        drop_text: bool = False,  # cfg for text
+        cfg_infer: bool = False,  # cfg inference, pack cond & uncond forward
+        cache: bool = False,
+    ):
+        batch, seq_len = x.shape[0], x.shape[1]
+        if time.ndim == 0:
+            time = time.repeat(batch)
+
+        # t: conditioning time, c: context (text + masked cond audio), x: noised input audio
+        t = self.time_embed(time)
+        if cfg_infer:  # pack cond & uncond forward: b n d -> 2b n d
+            x_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache)
+            x_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache)
+            x = torch.cat((x_cond, x_uncond), dim=0)
+            t = torch.cat((t, t), dim=0)
+            mask = torch.cat((mask, mask), dim=0) if mask is not None else None
+        else:
+            x = self.get_input_embed(x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache)
+
+        # postfix time t to input x, [b n d] -> [b n+1 d]
+        x = torch.cat([t.unsqueeze(1), x], dim=1)  # pack t to x
+        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 = self.norm_out(x)[:, 1:, :]  # unpack t from x
+
+        return self.proj_out(x)

src/f5_tts/model/cfm.py

@@ -0,0 +1,320 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+
+from __future__ import annotations
+
+from random import random
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+from torchdiffeq import odeint
+
+from f5_tts.model.modules import MelSpec
+from f5_tts.model.utils import (
+    default,
+    exists,
+    get_epss_timesteps,
+    lens_to_mask,
+    list_str_to_idx,
+    list_str_to_tensor,
+    mask_from_frac_lengths,
+)
+
+
+class CFM(nn.Module):
+    def __init__(
+        self,
+        transformer: nn.Module,
+        sigma=0.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.0),
+        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"],  # noqa: F722
+        text: int["b nt"] | list[str],  # noqa: F722
+        duration: int | int["b"],  # noqa: F821
+        *,
+        lens: int["b"] | None = None,  # noqa: F821
+        steps=32,
+        cfg_strength=1.0,
+        sway_sampling_coef=None,
+        seed: int | None = None,
+        max_duration=4096,
+        vocoder: Callable[[float["b d n"]], float["b nw"]] | None = None,  # noqa: F722
+        use_epss=True,
+        no_ref_audio=False,
+        duplicate_test=False,
+        t_inter=0.1,
+        edit_mask=None,
+    ):
+        try:
+            self.eval()
+            # raw wave
+
+            if cond.ndim == 2:
+                cond = self.mel_spec(cond)
+                cond = cond.permute(0, 2, 1)
+                assert cond.shape[-1] == self.num_channels
+
+            cond = cond.to(next(self.parameters()).dtype)
+
+            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
+
+            # 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(
+                torch.maximum((text != -1).sum(dim=-1), lens) + 1, duration
+            )  # duration at least text/audio prompt length plus 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.0)
+
+            cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)
+            if no_ref_audio:
+                cond = torch.zeros_like(cond)
+
+            cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)
+            cond_mask = cond_mask.unsqueeze(-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
+
+            # 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 (cond)
+                if cfg_strength < 1e-5:
+                    pred = self.transformer(
+                        x=x,
+                        cond=step_cond,
+                        text=text,
+                        time=t,
+                        mask=mask,
+                        drop_audio_cond=False,
+                        drop_text=False,
+                        cache=True,
+                        skip_flash_attn=True
+                    )
+                    return pred
+
+                # predict flow (cond and uncond), for classifier-free guidance
+                pred_cfg = self.transformer(
+                    x=x,
+                    cond=step_cond,
+                    text=text,
+                    time=t,
+                    mask=mask,
+                    cfg_infer=True,
+                    cache=True,
+                    skip_flash_attn=True
+                )
+                pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)
+                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, dtype=step_cond.dtype))
+            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))
+
+            if t_start == 0 and use_epss:  # use Empirically Pruned Step Sampling for low NFE
+                t = get_epss_timesteps(steps, device=self.device, dtype=step_cond.dtype)
+            else:
+                t = torch.linspace(t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype)
+            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)
+            self.transformer.clear_cache()
+
+            sampled = trajectory[-1]
+            out = sampled
+            # out = torch.where(cond_mask, cond, out)
+
+            if exists(vocoder):
+                out = out.permute(0, 2, 1)
+                out = vocoder(out)
+
+            return out, trajectory
+
+        except Exception as e:
+            print(f"cond shape: {cond.shape}, text shape: {text.shape if torch.is_tensor(text) else 'N/A'}, duration: {duration.shape}, lens: {lens.shape}")
+            print(f"cond: {cond}, text: {text}, duration: {duration}, lens: {lens}")
+            raise e
+
+    def forward(
+        self,
+        inp: float["b n d"] | float["b nw"],  # mel or raw wave  # noqa: F722
+        text: int["b nt"] | list[str],  # noqa: F722
+        *,
+        lens: int["b"] | None = None,  # noqa: F821
+        noise_scheduler: str | None = None,
+    ):
+        try:
+            # handle raw wave
+            # print(f"inp: {inp}, text: {text}, lens: {lens}")
+            if inp.ndim == 2:
+                inp = self.mel_spec(inp)
+                inp = inp.permute(0, 2, 1)
+                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
+            audio_frames_in_batch = mask.sum().item()
+            max_allowed_frames = mask.shape[0] * mask.shape[1]
+            perc_frame_utilised = audio_frames_in_batch / max_allowed_frames
+
+            # 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 = time.unsqueeze(-1).unsqueeze(-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
+
+            
+            # apply mask will use more memory; might adjust batchsize or batchsampler long sequence threshold
+            pred = self.transformer(
+                x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, mask=mask
+            )
+            
+            # flow matching loss
+            loss = F.mse_loss(pred, flow, reduction="none")
+            loss = loss[rand_span_mask]
+
+            return loss.mean(), cond, pred, audio_frames_in_batch, max_allowed_frames, perc_frame_utilised, batch, seq_len
+        except Exception as e:
+            print(f"input shape: {inp.shape}, text shape: {text.shape if torch.is_tensor(text) else 'N/A'}")
+            print(f"input: {input}, text: {text}")
+            raise e

src/f5_tts/model/modules.py

@@ -0,0 +1,790 @@
+"""
+ein notation:
+b - batch
+n - sequence
+nt - text sequence
+nw - raw wave length
+d - dimension
+"""
+# flake8: noqa
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+import torchaudio
+from librosa.filters import mel as librosa_mel_fn
+from torch import nn
+from x_transformers.x_transformers import apply_rotary_pos_emb
+
+from f5_tts.model.utils import is_package_available
+
+
+# raw wav to mel spec
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
+
+
+def get_bigvgan_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+    fmin=0,
+    fmax=None,
+    center=False,
+):  # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
+    device = waveform.device
+    key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
+
+    if key not in mel_basis_cache:
+        mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
+        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)  # TODO: why they need .float()?
+        hann_window_cache[key] = torch.hann_window(win_length).to(device)
+
+    mel_basis = mel_basis_cache[key]
+    hann_window = hann_window_cache[key]
+
+    padding = (n_fft - hop_length) // 2
+    waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
+
+    spec = torch.stft(
+        waveform,
+        n_fft,
+        hop_length=hop_length,
+        win_length=win_length,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+    mel_spec = torch.matmul(mel_basis, spec)
+    mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
+
+    return mel_spec
+
+
+def get_vocos_mel_spectrogram(
+    waveform,
+    n_fft=1024,
+    n_mel_channels=100,
+    target_sample_rate=24000,
+    hop_length=256,
+    win_length=1024,
+):
+    mel_stft = torchaudio.transforms.MelSpectrogram(
+        sample_rate=target_sample_rate,
+        n_fft=n_fft,
+        win_length=win_length,
+        hop_length=hop_length,
+        n_mels=n_mel_channels,
+        power=1,
+        center=True,
+        normalized=False,
+        norm=None,
+    ).to(waveform.device)
+    if len(waveform.shape) == 3:
+        waveform = waveform.squeeze(1)  # 'b 1 nw -> b nw'
+
+    assert len(waveform.shape) == 2
+
+    mel = mel_stft(waveform)
+    mel = mel.clamp(min=1e-5).log()
+    return mel
+
+
+class MelSpec(nn.Module):
+    def __init__(
+        self,
+        n_fft=1024,
+        hop_length=256,
+        win_length=1024,
+        n_mel_channels=100,
+        target_sample_rate=24_000,
+        mel_spec_type="vocos",
+    ):
+        super().__init__()
+        assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
+
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+        self.n_mel_channels = n_mel_channels
+        self.target_sample_rate = target_sample_rate
+
+        if mel_spec_type == "vocos":
+            self.extractor = get_vocos_mel_spectrogram
+        elif mel_spec_type == "bigvgan":
+            self.extractor = get_bigvgan_mel_spectrogram
+
+        self.register_buffer("dummy", torch.tensor(0), persistent=False)
+
+    def forward(self, wav):
+        if self.dummy.device != wav.device:
+            self.to(wav.device)
+
+        # Handle stereo/multi-channel audio
+        if len(wav.shape) == 3:  # (B, C, T)
+            if wav.shape[1] > 1:
+                wav = torch.mean(wav, dim=1, keepdim=True)
+        elif len(wav.shape) == 2:  # (C, T) or (B, T)
+            # We assume if it's 2D and shape[0] > 1 and not usually a batch size, it might be stereo
+            # But better to be explicit in what extractors expect.
+            # Most extractors here expect (B, T) after potential squeeze.
+            pass
+
+        mel = self.extractor(
+            waveform=wav,
+            n_fft=self.n_fft,
+            n_mel_channels=self.n_mel_channels,
+            target_sample_rate=self.target_sample_rate,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+        )
+
+        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.0)
+
+        x = x.permute(0, 2, 1)
+        x = self.conv1d(x)
+        out = x.permute(0, 2, 1)
+
+        if mask is not None:
+            out = out.masked_fill(~mask, 0.0)
+
+        return out
+
+
+# rotary positional embedding related
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):
+    # 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.0):
+    # 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
+
+# RMSNorm
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.native_rms_norm = float(torch.__version__[:3]) >= 2.4
+
+    def forward(self, x):
+        if self.native_rms_norm:
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)
+        else:
+            variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            if self.weight.dtype in [torch.float16, torch.bfloat16]:
+                x = x.to(self.weight.dtype)
+            x = x * self.weight
+
+        return x
+
+
+# AdaLayerNorm
+# return with modulated x for attn input, and params for later mlp modulation
+
+
+class AdaLayerNorm(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
+
+
+# AdaLayerNorm for final layer
+# return only with modulated x for attn input, cuz no more mlp modulation
+
+
+class AdaLayerNorm_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.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: bool = False,
+        qk_norm: Optional[str] = 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 qk_norm is None:
+            self.q_norm = None
+            self.k_norm = None
+        elif qk_norm == "rms_norm":
+            self.q_norm = RMSNorm(dim_head, eps=1e-6)
+            self.k_norm = RMSNorm(dim_head, eps=1e-6)
+        else:
+            raise ValueError(f"Unimplemented qk_norm: {qk_norm}")
+
+        if self.context_dim is not None:
+            self.to_q_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_k_c = nn.Linear(context_dim, self.inner_dim)
+            self.to_v_c = nn.Linear(context_dim, self.inner_dim)
+            if qk_norm is None:
+                self.c_q_norm = None
+                self.c_k_norm = None
+            elif qk_norm == "rms_norm":
+                self.c_q_norm = RMSNorm(dim_head, eps=1e-6)
+                self.c_k_norm = RMSNorm(dim_head, eps=1e-6)
+
+        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_dim is not None and not self.context_pre_only:
+            self.to_out_c = nn.Linear(self.inner_dim, context_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
+        skip_flash_attn: bool = False,
+    ) -> torch.Tensor:
+        if c is not None:
+            return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, skip_flash_attn=skip_flash_attn)
+        else:
+            return self.processor(self, x, mask=mask, rope=rope, skip_flash_attn=skip_flash_attn)
+    
+
+class AttnProcessor:
+    def __init__(
+        self,
+        pe_attn_head: int | None = None,  # number of attention head to apply rope, None for all
+        attn_backend: str = "torch",  # "torch" or "flash_attn"
+        attn_mask_enabled: bool = True,
+    ):
+        if attn_backend == "flash_attn":
+            assert is_package_available("flash_attn"), "Please install flash-attn first."
+            from flash_attn.bert_padding import pad_input, unpad_input
+            from flash_attn import flash_attn_varlen_func, flash_attn_func  
+
+        self.pe_attn_head = pe_attn_head
+        self.attn_backend = attn_backend
+        self.attn_mask_enabled = attn_mask_enabled
+
+    def __call__(
+        self,
+        attn: Attention,
+        x: float["b n d"],  # noised input x
+        mask: bool["b n"] | None = None,
+        rope=None,  # rotary position embedding,
+        skip_flash_attn: bool = False,
+    ) -> torch.FloatTensor:
+        batch_size = x.shape[0]
+
+        # `sample` projections
+        query = attn.to_q(x)
+        key = attn.to_k(x)
+        value = attn.to_v(x)
+
+        # 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)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+
+        # apply rotary position embedding
+        if rope is not None:
+            freqs, xpos_scale = rope
+            q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
+
+            if self.pe_attn_head is not None:
+                pn = self.pe_attn_head
+                query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)
+                key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)
+            else:
+                query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
+                key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
+
+        if skip_flash_attn or self.attn_backend == "torch":
+            # mask. e.g. inference got a batch with different target durations, mask out the padding
+            if self.attn_mask_enabled and mask is not None:
+                attn_mask = mask
+                attn_mask = attn_mask.unsqueeze(1).unsqueeze(1)  # '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)
+
+        elif self.attn_backend == "flash_attn":
+            query = query.transpose(1, 2)  # [b, h, n, d] -> [b, n, h, d]
+            key = key.transpose(1, 2)
+            value = value.transpose(1, 2)
+            if self.attn_mask_enabled and mask is not None:
+                query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask)
+                key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask)
+                value, _, _, _, _ = unpad_input(value, mask)
+                x = flash_attn_varlen_func(
+                    query,
+                    key,
+                    value,
+                    q_cu_seqlens,
+                    k_cu_seqlens,
+                    q_max_seqlen_in_batch,
+                    k_max_seqlen_in_batch,
+                )
+                x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch)
+                x = x.reshape(batch_size, -1, attn.heads * head_dim)
+            else:
+                x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False)
+                x = x.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 = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.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)
+
+        # 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)
+        c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        # qk norm
+        if attn.q_norm is not None:
+            query = attn.q_norm(query)
+        if attn.k_norm is not None:
+            key = attn.k_norm(key)
+        if attn.c_q_norm is not None:
+            c_query = attn.c_q_norm(c_query)
+        if attn.c_k_norm is not None:
+            c_key = attn.c_k_norm(c_key)
+
+        # 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.0) if xpos_scale is not None else (1.0, 1.0)
+            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.0) if xpos_scale is not None else (1.0, 1.0)
+            c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
+            c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
+
+        # joint attention
+        query = torch.cat([query, c_query], dim=2)
+        key = torch.cat([key, c_key], dim=2)
+        value = torch.cat([value, c_value], dim=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 = attn_mask.unsqueeze(1).unsqueeze(1)  # '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 = mask.unsqueeze(-1)
+            x = x.masked_fill(~mask, 0.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,
+        qk_norm=None,
+        pe_attn_head=None,
+        attn_backend="torch",  # "torch" or "flash_attn"
+        attn_mask_enabled=True,
+    ):
+        super().__init__()
+
+        self.attn_norm = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=AttnProcessor(
+                pe_attn_head=pe_attn_head,
+                attn_backend=attn_backend,
+                attn_mask_enabled=attn_mask_enabled,
+            ),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            qk_norm=qk_norm,
+        )
+
+        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, skip_flash_attn: bool = False,):  # 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, skip_flash_attn=skip_flash_attn)
+
+        # 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_dim=None, context_pre_only=False, qk_norm=None
+    ):
+        super().__init__()
+        if context_dim is None:
+            context_dim = dim
+        self.context_pre_only = context_pre_only
+
+        self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim)
+        self.attn_norm_x = AdaLayerNorm(dim)
+        self.attn = Attention(
+            processor=JointAttnProcessor(),
+            dim=dim,
+            heads=heads,
+            dim_head=dim_head,
+            dropout=dropout,
+            context_dim=context_dim,
+            context_pre_only=context_pre_only,
+            qk_norm=qk_norm,
+        )
+
+        if not context_pre_only:
+            self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6)
+            self.ff_c = FeedForward(dim=context_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_hidden = time_hidden.to(timestep.dtype)
+        time = self.time_mlp(time_hidden)  # b d
+        return time

src/f5_tts/model/utils.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import os
+import random
+from collections import defaultdict
+from importlib.resources import files
+from typing import List, Union
+
+import jieba
+import torch
+from pypinyin import Style, lazy_pinyin
+from torch.nn.utils.rnn import pad_sequence
+import requests
+import json
+import socket
+import subprocess
+from datetime import datetime
+
+
+# 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
+
+
+def is_package_available(package_name: str) -> bool:
+    try:
+        import importlib
+
+        package_exists = importlib.util.find_spec(package_name) is not None
+        return package_exists
+    except Exception:
+        return False
+
+
+# tensor helpers
+
+
+def lens_to_mask(t: int["b"], length: int | None = None) -> bool["b n"]:  # noqa: F722 F821
+    if not exists(length):
+        length = t.amax()
+
+    seq = torch.arange(length, device=t.device)
+    return seq[None, :] < t[:, None]
+
+
+def mask_from_start_end_indices(seq_len: int["b"], start: int["b"], end: int["b"]):  # noqa: F722 F821
+    max_seq_len = seq_len.max().item()
+    seq = torch.arange(max_seq_len, device=start.device).long()
+    start_mask = seq[None, :] >= start[:, None]
+    end_mask = seq[None, :] < end[:, None]
+    return start_mask & end_mask
+
+
+def mask_from_frac_lengths(seq_len: int["b"], frac_lengths: float["b"]):  # noqa: F722 F821
+    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"]:  # noqa: F722
+    if not exists(mask):
+        return t.mean(dim=1)
+
+    t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))
+    num = t.sum(dim=1)
+    den = mask.float().sum(dim=1)
+
+    return num / den.clamp(min=1.0)
+
+
+# simple utf-8 tokenizer, since paper went character based
+def list_str_to_tensor(text: list[str], padding_value=-1) -> int["b nt"]:  # noqa: F722
+    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"]:  # noqa: F722
+    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, extra_vocab_path: str = None):
+    """
+    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
+                - "custom" if you're directly passing in a path to the vocab.txt you want to use
+    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)
+    extra_vocab_path - path to txt file containing additional characters (one per line) to expand vocabulary
+    """
+    
+    if tokenizer in ["pinyin", "char", "cls"]:
+        tokenizer_path = os.path.join(files("f5_tts").joinpath("../../data"), f"{dataset_name}_{tokenizer}/vocab.txt")
+        with open(tokenizer_path, "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
+        
+    elif tokenizer == "custom":
+        with open(dataset_name, "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)
+    
+    # Load and merge extra vocabulary from txt file
+    if extra_vocab_path is not None and os.path.exists(extra_vocab_path):
+        if vocab_char_map is not None:  # Only extend if not byte tokenizer
+            current_vocab_size = len(vocab_char_map)
+            with open(extra_vocab_path, "r", encoding="utf-8") as f:
+                for char in f:
+                    char = char.strip()
+                    if char and char not in vocab_char_map:
+                        vocab_char_map[char] = len(vocab_char_map)
+            vocab_size = len(vocab_char_map)
+            print(f"Extended vocabulary with {vocab_size - current_vocab_size} new tokens from {extra_vocab_path}")
+    
+    return vocab_char_map, vocab_size
+
+
+
+def get_tokenizer(dataset_name, tokenizer: str | List[str], extra_vocab_path: str = None):
+    """
+    Return a dictionary of tokenizers if tokenizer is a list, each key is the tokenizer name and value is a tuple of (vocab_char_map, vocab_size)
+    Otherwise, return a dictionary with single tokenizer entry
+    """
+    
+    if isinstance(tokenizer, list):
+        tokenizers_dict = {}
+        for t_name in tokenizer:
+            vocab_char_map, vocab_size = _get_tokenizer(dataset_name, t_name, extra_vocab_path)
+            tokenizers_dict[t_name] = (vocab_char_map, vocab_size)  # Fixed: was using 't' instead of 't_name'
+        return tokenizers_dict
+    else:
+        vocab_char_map, vocab_size = _get_tokenizer(dataset_name, tokenizer, extra_vocab_path)
+        return  vocab_char_map, vocab_size
+
+def save_vocab(vocab_char_map, save_path):
+    """Save vocabulary to file"""
+    if vocab_char_map is None:
+        return
+        
+    print(f"\nSaving vocabulary to: {save_path}")
+    
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    
+    # Sort by index to maintain order
+    vocab_items = sorted(vocab_char_map.items(), key=lambda x: x[1])
+    
+    with open(save_path, 'w', encoding='utf-8') as f:
+        for char, idx in vocab_items:
+            f.write(f"{char}\n")
+    
+    print(f"✓ Saved {len(vocab_char_map)} tokens to vocab.txt")
+
+def send_slack_notification(message, webhook_url=None, title="Training Notification"):
+    """Send a notification to a Slack channel via webhook."""
+    if webhook_url is None:
+        webhook_url = os.getenv("SLACK_WEBHOOK_URL")
+    
+    if not webhook_url:
+        return False
+    
+    hostname = socket.gethostname()
+    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    
+    payload = {
+        "text": f"*{title}*\n*Time:* {timestamp}\n*Host:* {hostname}\n*Message:* {message}"
+    }
+    
+    try:
+        response = requests.post(
+            webhook_url, 
+            data=json.dumps(payload),
+            headers={'Content-Type': 'application/json'}
+        )
+        return response.status_code == 200
+    except Exception as e:
+        print(f"Failed to send Slack notification: {e}")
+        return False
+
+def track_with_dvc(path):
+    """Track a file or directory with DVC."""
+    try:
+        # Check if dvc is initialized
+        if not os.path.exists(".dvc"):
+            print("DVC not initialized. Skipping tracking.")
+            return False
+            
+        print(f"Tracking with DVC: {path}")
+        subprocess.run(["dvc", "add", path], check=True, capture_output=True)
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"DVC add failed for {path}: {e.stderr.decode() if e.stderr else str(e)}")
+        return False
+    except Exception as e:
+        print(f"Failed to track with DVC: {e}")
+        return False
+
+# convert char to pinyin
+
+
+def convert_char_to_pinyin(text_list, polyphone=True):
+    if jieba.dt.initialized is False:
+        jieba.default_logger.setLevel(50)  # CRITICAL
+        jieba.initialize()
+
+    final_text_list = []
+    custom_trans = str.maketrans(
+        {";": ",", "“": '"', "”": '"', "‘": "'", "’": "'"}
+    )  # add custom trans here, to address oov
+
+    def is_chinese(c):
+        return (
+            "\u3100" <= c <= "\u9fff"  # common chinese characters
+        )
+
+    for text in text_list:
+        char_list = []
+        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 east asian characters
+                seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)
+                for i, c in enumerate(seg):
+                    if is_chinese(c):
+                        char_list.append(" ")
+                    char_list.append(seg_[i])
+            else:  # if mixed characters, alphabets and symbols
+                for c in seg:
+                    if ord(c) < 256:
+                        char_list.extend(c)
+                    elif is_chinese(c):
+                        char_list.append(" ")
+                        char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))
+                    else:
+                        char_list.append(c)
+        final_text_list.append(char_list)
+
+    return final_text_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
+
+
+# get the empirically pruned step for sampling
+
+
+def get_epss_timesteps(n, device, dtype):
+    dt = 1 / 32
+    predefined_timesteps = {
+        5: [0, 2, 4, 8, 16, 32],
+        6: [0, 2, 4, 6, 8, 16, 32],
+        7: [0, 2, 4, 6, 8, 16, 24, 32],
+        10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],
+        12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+        16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],
+    }
+    t = predefined_timesteps.get(n, [])
+    if not t:
+        return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)
+    return dt * torch.tensor(t, device=device, dtype=dtype)

src/f5_tts/socket_client.py

@@ -0,0 +1,63 @@
+import asyncio
+import logging
+import socket
+import time
+
+import numpy as np
+import pyaudio
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+async def listen_to_F5TTS(text, server_ip="localhost", server_port=9998):
+    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    await asyncio.get_event_loop().run_in_executor(None, client_socket.connect, (server_ip, int(server_port)))
+
+    start_time = time.time()
+    first_chunk_time = None
+
+    async def play_audio_stream():
+        nonlocal first_chunk_time
+        p = pyaudio.PyAudio()
+        stream = p.open(format=pyaudio.paFloat32, channels=1, rate=24000, output=True, frames_per_buffer=2048)
+
+        try:
+            while True:
+                data = await asyncio.get_event_loop().run_in_executor(None, client_socket.recv, 8192)
+                if not data:
+                    break
+                if data == b"END":
+                    logger.info("End of audio received.")
+                    break
+
+                audio_array = np.frombuffer(data, dtype=np.float32)
+                stream.write(audio_array.tobytes())
+
+                if first_chunk_time is None:
+                    first_chunk_time = time.time()
+
+        finally:
+            stream.stop_stream()
+            stream.close()
+            p.terminate()
+
+        logger.info(f"Total time taken: {time.time() - start_time:.4f} seconds")
+
+    try:
+        data_to_send = f"{text}".encode("utf-8")
+        await asyncio.get_event_loop().run_in_executor(None, client_socket.sendall, data_to_send)
+        await play_audio_stream()
+
+    except Exception as e:
+        logger.error(f"Error in listen_to_F5TTS: {e}")
+
+    finally:
+        client_socket.close()
+
+
+if __name__ == "__main__":
+    text_to_send = "As a Reader assistant, I'm familiar with new technology. which are key to its improved performance in terms of both training speed and inference efficiency. Let's break down the components"
+
+    asyncio.run(listen_to_F5TTS(text_to_send))

src/f5_tts/socket_server.py

@@ -0,0 +1,268 @@
+import argparse
+import gc
+import logging
+import queue
+import socket
+import struct
+import threading
+import traceback
+import wave
+from importlib.resources import files
+
+import numpy as np
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+from hydra.utils import get_class
+from omegaconf import OmegaConf
+
+from f5_tts.infer.utils_infer import (
+    chunk_text,
+    infer_batch_process,
+    load_model,
+    load_vocoder,
+    preprocess_ref_audio_text,
+)
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class AudioFileWriterThread(threading.Thread):
+    """Threaded file writer to avoid blocking the TTS streaming process."""
+
+    def __init__(self, output_file, sampling_rate):
+        super().__init__()
+        self.output_file = output_file
+        self.sampling_rate = sampling_rate
+        self.queue = queue.Queue()
+        self.stop_event = threading.Event()
+        self.audio_data = []
+
+    def run(self):
+        """Process queued audio data and write it to a file."""
+        logger.info("AudioFileWriterThread started.")
+        with wave.open(self.output_file, "wb") as wf:
+            wf.setnchannels(1)
+            wf.setsampwidth(2)
+            wf.setframerate(self.sampling_rate)
+
+            while not self.stop_event.is_set() or not self.queue.empty():
+                try:
+                    chunk = self.queue.get(timeout=0.1)
+                    if chunk is not None:
+                        chunk = np.int16(chunk * 32767)
+                        self.audio_data.append(chunk)
+                        wf.writeframes(chunk.tobytes())
+                except queue.Empty:
+                    continue
+
+    def add_chunk(self, chunk):
+        """Add a new chunk to the queue."""
+        self.queue.put(chunk)
+
+    def stop(self):
+        """Stop writing and ensure all queued data is written."""
+        self.stop_event.set()
+        self.join()
+        logger.info("Audio writing completed.")
+
+
+class TTSStreamingProcessor:
+    def __init__(self, model, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32):
+        self.device = device or (
+            "cuda"
+            if torch.cuda.is_available()
+            else "xpu"
+            if torch.xpu.is_available()
+            else "mps"
+            if torch.backends.mps.is_available()
+            else "cpu"
+        )
+        model_cfg = OmegaConf.load(str(files("f5_tts").joinpath(f"configs/{model}.yaml")))
+        self.model_cls = get_class(f"f5_tts.model.{model_cfg.model.backbone}")
+        self.model_arc = model_cfg.model.arch
+        self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type
+        self.sampling_rate = model_cfg.model.mel_spec.target_sample_rate
+
+        self.model = self.load_ema_model(ckpt_file, vocab_file, dtype)
+        self.vocoder = self.load_vocoder_model()
+
+        self.update_reference(ref_audio, ref_text)
+        self._warm_up()
+        self.file_writer_thread = None
+        self.first_package = True
+
+    def load_ema_model(self, ckpt_file, vocab_file, dtype):
+        return load_model(
+            self.model_cls,
+            self.model_arc,
+            ckpt_path=ckpt_file,
+            mel_spec_type=self.mel_spec_type,
+            vocab_file=vocab_file,
+            ode_method="euler",
+            use_ema=True,
+            device=self.device,
+        ).to(self.device, dtype=dtype)
+
+    def load_vocoder_model(self):
+        return load_vocoder(vocoder_name=self.mel_spec_type, is_local=False, local_path=None, device=self.device)
+
+    def update_reference(self, ref_audio, ref_text):
+        self.ref_audio, self.ref_text = preprocess_ref_audio_text(ref_audio, ref_text)
+        self.audio, self.sr = torchaudio.load(self.ref_audio)
+
+        ref_audio_duration = self.audio.shape[-1] / self.sr
+        ref_text_byte_len = len(self.ref_text.encode("utf-8"))
+        self.max_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration))
+        self.few_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 2)
+        self.min_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 4)
+
+    def _warm_up(self):
+        logger.info("Warming up the model...")
+        gen_text = "Warm-up text for the model."
+        for _ in infer_batch_process(
+            (self.audio, self.sr),
+            self.ref_text,
+            [gen_text],
+            self.model,
+            self.vocoder,
+            progress=None,
+            device=self.device,
+            streaming=True,
+        ):
+            pass
+        logger.info("Warm-up completed.")
+
+    def generate_stream(self, text, conn):
+        text_batches = chunk_text(text, max_chars=self.max_chars)
+        if self.first_package:
+            text_batches = chunk_text(text_batches[0], max_chars=self.few_chars) + text_batches[1:]
+            text_batches = chunk_text(text_batches[0], max_chars=self.min_chars) + text_batches[1:]
+            self.first_package = False
+
+        audio_stream = infer_batch_process(
+            (self.audio, self.sr),
+            self.ref_text,
+            text_batches,
+            self.model,
+            self.vocoder,
+            progress=None,
+            device=self.device,
+            streaming=True,
+            chunk_size=2048,
+        )
+
+        # Reset the file writer thread
+        if self.file_writer_thread is not None:
+            self.file_writer_thread.stop()
+        self.file_writer_thread = AudioFileWriterThread("output.wav", self.sampling_rate)
+        self.file_writer_thread.start()
+
+        for audio_chunk, _ in audio_stream:
+            if len(audio_chunk) > 0:
+                logger.info(f"Generated audio chunk of size: {len(audio_chunk)}")
+
+                # Send audio chunk via socket
+                conn.sendall(struct.pack(f"{len(audio_chunk)}f", *audio_chunk))
+
+                # Write to file asynchronously
+                self.file_writer_thread.add_chunk(audio_chunk)
+
+        logger.info("Finished sending audio stream.")
+        conn.sendall(b"END")  # Send end signal
+
+        # Ensure all audio data is written before exiting
+        self.file_writer_thread.stop()
+
+
+def handle_client(conn, processor):
+    try:
+        with conn:
+            conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+            while True:
+                data = conn.recv(1024)
+                if not data:
+                    processor.first_package = True
+                    break
+                data_str = data.decode("utf-8").strip()
+                logger.info(f"Received text: {data_str}")
+
+                try:
+                    processor.generate_stream(data_str, conn)
+                except Exception as inner_e:
+                    logger.error(f"Error during processing: {inner_e}")
+                    traceback.print_exc()
+                    break
+    except Exception as e:
+        logger.error(f"Error handling client: {e}")
+        traceback.print_exc()
+
+
+def start_server(host, port, processor):
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind((host, port))
+        s.listen()
+        logger.info(f"Server started on {host}:{port}")
+        while True:
+            conn, addr = s.accept()
+            logger.info(f"Connected by {addr}")
+            handle_client(conn, processor)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--host", default="0.0.0.0")
+    parser.add_argument("--port", default=9998)
+
+    parser.add_argument(
+        "--model",
+        default="F5TTS_v1_Base",
+        help="The model name, e.g. F5TTS_v1_Base",
+    )
+    parser.add_argument(
+        "--ckpt_file",
+        default=str(hf_hub_download(repo_id="SWivid/F5-TTS", filename="F5TTS_v1_Base/model_1250000.safetensors")),
+        help="Path to the model checkpoint file",
+    )
+    parser.add_argument(
+        "--vocab_file",
+        default="",
+        help="Path to the vocab file if customized",
+    )
+
+    parser.add_argument(
+        "--ref_audio",
+        default=str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav")),
+        help="Reference audio to provide model with speaker characteristics",
+    )
+    parser.add_argument(
+        "--ref_text",
+        default="",
+        help="Reference audio subtitle, leave empty to auto-transcribe",
+    )
+
+    parser.add_argument("--device", default=None, help="Device to run the model on")
+    parser.add_argument("--dtype", default=torch.float32, help="Data type to use for model inference")
+
+    args = parser.parse_args()
+
+    try:
+        # Initialize the processor with the model and vocoder
+        processor = TTSStreamingProcessor(
+            model=args.model,
+            ckpt_file=args.ckpt_file,
+            vocab_file=args.vocab_file,
+            ref_audio=args.ref_audio,
+            ref_text=args.ref_text,
+            device=args.device,
+            dtype=args.dtype,
+        )
+
+        # Start the server
+        start_server(args.host, args.port, processor)
+
+    except KeyboardInterrupt:
+        gc.collect()

test.py

@@ -0,0 +1,34 @@
+import os
+from transformers import AutoModel
+
+# --- Paths / model id (adjust if needed) ---
+REPO_DIR = "."
+MODEL_ID =  "bharatgenai/sooktam2"
+
+REF_AUDIO = "./ref.wav"
+REF_TEXT = "सर, मैं तब से यह कह रहा हूँ कि मैंने अपना टिकट कैंसल कर दिया है, लेकिन अब तक मेरे पैसे वापस नहीं आए हैं। आप इस मामले को देखेंगे भी या नहीं?"
+GEN_TEXT = "यह एक टेस्ट वाक्य है जिसे आवाज़ में बदलना है।"
+
+OUT_DIR = os.path.join(REPO_DIR, "outputs")
+OUT_WAV = os.path.join(OUT_DIR, "sooktam_cls.wav")
+
+# CLS tokenization is handled inside utils_infer via cls_tokenizer_v2.
+
+# --- Load TTS model via AutoModel (auto-download ckpt + vocab from HF) ---
+model = AutoModel.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+)
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+wav, sr, _ = model.infer(
+    ref_file=REF_AUDIO,
+    ref_text=REF_TEXT,
+    gen_text=GEN_TEXT,
+    tokenizer="cls",
+    cls_language="hindi",
+    file_wave=OUT_WAV,
+)
+
+print("Saved:", OUT_WAV, "sample_rate:", sr, "samples:", len(wav))

vocab.txt

@@ -0,0 +1,543 @@
+ 
+'
+-
+.
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+?
+M
+a
+aa
+ae
+ax
+b
+bh
+c
+ch
+d
+dh
+dx
+dxh
+dxhq
+dxq
+e
+ee
+ei
+f
+g
+gh
+gq
+h
+hq
+i
+ii
+j
+jh
+k
+kh
+khq
+kq
+l
+lx
+m
+mq
+n
+nd
+ng
+nj
+nk
+nx
+o
+ou
+p
+ph
+q
+r
+rq
+rqw
+rx
+s
+sh
+sx
+t
+th
+tx
+txh
+u
+uu
+v
+w
+x
+y
+z
+č
+š
+α
+ٮ
+ऌ
+।
+०
+१
+२
+३
+४
+५
+७
+८
+९
+ॲ
+ॵ
+ੳ
+ં
+ન
+ર
+ી
+ુ
+૦
+૧
+૨
+૩
+૪
+૫
+૬
+૭
+૮
+૯
+‍
+‪
+₂
+一
+下
+不
+世
+中
+之
+亮
+人
+他
+以
+們
+其
+別
+北
+去
+可
+周
+哇
+哪
+器
+因
+在
+城
+場
+士
+外
+多
+大
+天
+如
+媽
+孫
+宗
+宰
+察
+專
+尿
+市
+師
+年
+影
+很
+後
+性
+我
+方
+明
+是
+時
+普
+最
+本
+果
+業
+次
+此
+武
+比
+漂
+為
+無
+特
+現
+生
+產
+界
+當
+病
+的
+知
+者
+能
+若
+茶
+計
+設
+誮
+警
+豪
+賽
+足
+跑
+迷
+逃
+這
+通
+道
+那
+間
+阿
+電
+面
+骚
+骡
+髋
+麼
+�
+!
+"
+#
+$
+%
+&
+(
+)
+*
++
+,
+/
+:
+;
+<
+=
+>
+@
+[
+\
+]
+^
+_
+`
+ag
+ai
+au
+en_a
+en_b
+en_c
+en_d
+en_e
+en_f
+en_g
+en_h
+en_i
+en_j
+en_k
+en_l
+en_m
+en_n
+en_o
+en_p
+en_q
+en_r
+en_s
+en_t
+en_u
+en_v
+en_w
+en_x
+en_y
+en_z
+ln
+lw
+nn
+nw
+oo
+rw
+zh
+{
+|
+}
+~
+
+£
+§
+­
+°
+·
+¸
+º
+»
+½
+×
+à
+â
+ã
+ç
+è
+é
+î
+ï
+ô
+ö
+ú
+û
+ü
+ý
+þ
+ā
+ʼ
+ˉ
+˚
+̄
+̮
+а
+д
+ж
+й
+к
+о
+ҙ
+ҡ
+ү
+ҳ
+؀
+؁
+،
+؍
+؎
+ؐ
+ؑ
+ؒ
+ؓ
+ؔ
+ؕ
+؛
+؟
+ء
+آ
+أ
+ؤ
+إ
+ئ
+ا
+ب
+ة
+ت
+ث
+ج
+ح
+خ
+د
+ذ
+ر
+ز
+س
+ش
+ص
+ض
+ط
+ظ
+ع
+غ
+ـ
+ف
+ق
+ك
+ل
+م
+ن
+ه
+و
+ى
+ي
+ً
+ٌ
+ٍ
+َ
+ُ
+ِ
+ّ
+ْ
+ٓ
+ٔ
+ٕ
+ٖ
+ٗ
+٘
+ٙ
+ٚ
+ٛ
+٬
+ٰ
+ٱ
+ٴ
+ٶ
+ٸ
+ٹ
+ٻ
+ټ
+ٽ
+پ
+ڀ
+ځ
+ڃ
+ڄ
+څ
+چ
+ڇ
+ڈ
+ڊ
+ڌ
+ڑ
+ړ
+ژ
+ڙ
+ښ
+ڦ
+ک
+ڪ
+ګ
+ڭ
+گ
+ڱ
+ڳ
+ں
+ڻ
+ھ
+ہ
+ۂ
+ۃ
+ۄ
+ۆ
+ی
+ۍ
+ێ
+ې
+ے
+ۓ
+۔
+ە
+ۖ
+ۡ
+ۤ
+۪
+۰
+۱
+۲
+۳
+۴
+۵
+۶
+۷
+۸
+۹
+ऺ
+ऻ
+ॐ
+॑
+॒
+॓
+॔
+ॕ
+ॖ
+॥
+६
+ঌ
+০
+১
+২
+৩
+৪
+৫
+৬
+৭
+৮
+৯
+৲
+৷
+ੑ
+੦
+੧
+੨
+੩
+੪
+੫
+੭
+ੲ
+ૐ
+ૢ
+ୖ
+ୢ
+୦
+୧
+୨
+୩
+୪
+୫
+୬
+୭
+୮
+୯
+௕
+ౕ
+ౖ
+౦
+౩
+ೕ
+ೖ
+೦
+೧
+೨
+೩
+೪
+೫
+೬
+೭
+೮
+೯
+ഌ
+ൡ
+൪
+ൿ
+᠎
+ṃ
+​
+‌
+‏
+‐
+–
+—
+‘
+’
+“
+”
+•
+․
+…
+‮
+›
+⁰
+₹
+⇒
+−
+ﯹ
+ﷺ
+️
+