[ { "path": ".github/ISSUE_TEMPLATE/bug_report.yml", "content": "name: \"Bug Report\"\ndescription: |\n Please provide as much details to help address the issue more efficiently, including input, output, logs and screenshots.\nlabels:\n - bug\nbody:\n - type: checkboxes\n attributes:\n label: Checks\n description: \"To ensure timely help, please confirm the following:\"\n options:\n - label: This template is only for bug reports, usage problems go with 'Help Wanted'.\n required: true\n - label: I have thoroughly reviewed the project documentation but couldn't find information to solve my problem.\n required: true\n - label: I have searched for existing issues, including closed ones, and couldn't find a solution.\n required: true\n - label: I am using English to submit this issue to facilitate community communication.\n required: true\n - type: textarea\n attributes:\n label: Environment Details\n description: \"Provide details including OS, GPU info, Python version, any relevant software or dependencies, and trainer setting.\"\n placeholder: e.g., CentOS Linux 7, 4 * RTX 3090, Python 3.10, torch==2.3.0+cu118, cuda 11.8, config yaml is ...\n validations:\n required: true\n - type: textarea\n attributes:\n label: Steps to Reproduce\n description: |\n Include detailed steps, screenshots, and logs. Use the correct markdown syntax for code blocks.\n placeholder: |\n 1. Create a new conda environment.\n 2. Clone the repository, install as local editable and properly set up.\n 3. Run the command: `accelerate launch src/f5_tts/train/train.py`.\n 4. Have following error message... (attach logs).\n validations:\n required: true\n - type: textarea\n attributes:\n label: ✔️ Expected Behavior\n placeholder: Describe in detail what you expected to happen.\n validations:\n required: false\n - type: textarea\n attributes:\n label: ❌ Actual Behavior\n placeholder: Describe in detail what actually happened.\n validations:\n required: false" }, { "path": ".github/ISSUE_TEMPLATE/config.yml", "content": "blank_issues_enabled: false\n" }, { "path": ".github/ISSUE_TEMPLATE/feature_request.yml", "content": "name: \"Feature Request\"\ndescription: |\n Some constructive suggestions and new ideas regarding current repo.\nlabels:\n - enhancement\nbody:\n - type: checkboxes\n attributes:\n label: Checks\n description: \"To help us grasp quickly, please confirm the following:\"\n options:\n - label: This template is only for feature request.\n required: true\n - label: I have thoroughly reviewed the project documentation but couldn't find any relevant information that meets my needs.\n required: true\n - label: I have searched for existing issues, including closed ones, and found not discussion yet.\n required: true\n - label: I am using English to submit this issue to facilitate community communication.\n required: true\n - type: textarea\n attributes:\n label: 1. Is this request related to a challenge you're experiencing? Tell us your story.\n description: |\n Describe the specific problem or scenario you're facing in detail. For example:\n *\"I was trying to use [feature] for [specific task], but encountered [issue]. This was frustrating because....\"*\n placeholder: Please describe the situation in as much detail as possible.\n validations:\n required: true\n\n - type: textarea\n attributes:\n label: 2. What is your suggested solution?\n description: |\n Provide a clear description of the feature or enhancement you'd like to propose. \n How would this feature solve your issue or improve the project?\n placeholder: Describe your idea or proposed solution here.\n validations:\n required: true\n\n - type: textarea\n attributes:\n label: 3. Additional context or comments\n description: |\n Any other relevant information, links, documents, or screenshots that provide clarity. \n Use this section for anything not covered above.\n placeholder: Add any extra details here.\n validations:\n required: false\n\n - type: checkboxes\n attributes:\n label: 4. Can you help us with this feature?\n description: |\n Let us know if you're interested in contributing. This is not a commitment but a way to express interest in collaboration.\n options:\n - label: I am interested in contributing to this feature.\n required: false\n\n - type: markdown\n attributes:\n value: |\n **Note:** Please submit only one request per issue to keep discussions focused and manageable." }, { "path": ".github/ISSUE_TEMPLATE/help_wanted.yml", "content": "name: \"Help Wanted\"\ndescription: |\n Please provide as much details to help address the issue more efficiently, including input, output, logs and screenshots.\nlabels:\n - help wanted\nbody:\n - type: checkboxes\n attributes:\n label: Checks\n description: \"To ensure timely help, please confirm the following:\"\n options:\n - label: This template is only for usage issues encountered.\n required: true\n - label: I have thoroughly reviewed the project documentation but couldn't find information to solve my problem.\n required: true\n - label: I have searched for existing issues, including closed ones, and couldn't find a solution.\n required: true\n - label: I am using English to submit this issue to facilitate community communication.\n required: true\n - type: textarea\n attributes:\n label: Environment Details\n description: \"Provide details such as OS, Python version, and any relevant software or dependencies.\"\n placeholder: |\n e.g., macOS 13.5, Python 3.10, torch==2.3.0, Gradio 4.44.1\n If training or finetuning related, provide detailed configuration including GPU info and training setup.\n validations:\n required: true\n - type: textarea\n attributes:\n label: Steps to Reproduce\n description: |\n Include detailed steps, screenshots, and logs. Provide used prompt wav and text. Use the correct markdown syntax for code blocks.\n placeholder: |\n 1. Create a new conda environment.\n 2. Clone the repository and install as pip package.\n 3. Run the command: `f5-tts_infer-gradio` with no ref_text provided.\n 4. Stuck there with the following message... (attach logs and also error msg e.g. after ctrl-c).\n 5. Prompt & generated wavs are [change suffix to .mp4 to enable direct upload or pack all to .zip].\n 6. Reference audio's transcription or provided ref_text is `xxx`, and text to generate is `xxx`.\n validations:\n required: true\n - type: textarea\n attributes:\n label: ✔️ Expected Behavior\n placeholder: Describe what you expected to happen in detail, e.g. output a generated audio.\n validations:\n required: false\n - type: textarea\n attributes:\n label: ❌ Actual Behavior\n placeholder: Describe what actually happened in detail, failure messages, etc.\n validations:\n required: false" }, { "path": ".github/ISSUE_TEMPLATE/question.yml", "content": "name: \"Question\"\ndescription: |\n Research question or pure inquiry about the project, usage issue goes with \"help wanted\".\nlabels:\n - question\nbody:\n - type: checkboxes\n attributes:\n label: Checks\n description: \"To help us grasp quickly, please confirm the following:\"\n options:\n - label: This template is only for research question, not usage problems, feature requests or bug reports.\n required: true\n - label: I have thoroughly reviewed the project documentation and read the related paper(s).\n required: true\n - label: I have searched for existing issues, including closed ones, no similar questions.\n required: true\n - label: I am using English to submit this issue to facilitate community communication.\n required: true\n - type: textarea\n attributes:\n label: Question details\n description: |\n Question details, clearly stated using proper markdown syntax.\n validations:\n required: true\n" }, { "path": ".github/workflows/pre-commit.yaml", "content": "name: pre-commit\n\non:\n pull_request:\n push:\n branches: [main]\n\njobs:\n pre-commit:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v3\n - uses: actions/setup-python@v3\n - uses: pre-commit/action@v3.0.1\n" }, { "path": ".github/workflows/publish-docker-image.yaml", "content": "name: Create and publish a Docker image\r\n\r\n# Configures this workflow to run every time a change is pushed to the branch called `release`.\r\non:\r\n push:\r\n branches: ['main']\r\n\r\n# Defines two custom environment variables for the workflow. These are used for the Container registry domain, and a name for the Docker image that this workflow builds.\r\nenv:\r\n REGISTRY: ghcr.io\r\n IMAGE_NAME: ${{ github.repository }}\r\n\r\n# There is a single job in this workflow. It's configured to run on the latest available version of Ubuntu.\r\njobs:\r\n build-and-push-image:\r\n runs-on: ubuntu-latest\r\n # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.\r\n permissions:\r\n contents: read\r\n packages: write\r\n # \r\n steps:\r\n - name: Checkout repository\r\n uses: actions/checkout@v4\r\n - name: Free Up GitHub Actions Ubuntu Runner Disk Space 🔧\r\n uses: jlumbroso/free-disk-space@main\r\n with:\r\n # This might remove tools that are actually needed, if set to \"true\" but frees about 6 GB\r\n tool-cache: false\r\n\r\n # All of these default to true, but feel free to set to \"false\" if necessary for your workflow\r\n android: true\r\n dotnet: true\r\n haskell: true\r\n large-packages: false\r\n swap-storage: false\r\n docker-images: false\r\n # Uses the `docker/login-action` action to log in to the Container registry registry using the account and password that will publish the packages. Once published, the packages are scoped to the account defined here.\r\n - name: Log in to the Container registry\r\n uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1\r\n with:\r\n registry: ${{ env.REGISTRY }}\r\n username: ${{ github.actor }}\r\n password: ${{ secrets.GITHUB_TOKEN }}\r\n # This step uses [docker/metadata-action](https://github.com/docker/metadata-action#about) to extract tags and labels that will be applied to the specified image. The `id` \"meta\" allows the output of this step to be referenced in a subsequent step. The `images` value provides the base name for the tags and labels.\r\n - name: Extract metadata (tags, labels) for Docker\r\n id: meta\r\n uses: docker/metadata-action@9ec57ed1fcdbf14dcef7dfbe97b2010124a938b7\r\n with:\r\n images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}\r\n # This step uses the `docker/build-push-action` action to build the image, based on your repository's `Dockerfile`. If the build succeeds, it pushes the image to GitHub Packages.\r\n # It uses the `context` parameter to define the build's context as the set of files located in the specified path. For more information, see \"[Usage](https://github.com/docker/build-push-action#usage)\" in the README of the `docker/build-push-action` repository.\r\n # It uses the `tags` and `labels` parameters to tag and label the image with the output from the \"meta\" step.\r\n - name: Build and push Docker image\r\n uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4\r\n with:\r\n context: .\r\n push: true\r\n tags: ${{ steps.meta.outputs.tags }}\r\n labels: ${{ steps.meta.outputs.labels }}\r\n" }, { "path": ".github/workflows/publish-pypi.yaml", "content": "# This workflow uses actions that are not certified by GitHub.\n# They are provided by a third-party and are governed by\n# separate terms of service, privacy policy, and support\n# documentation.\n\n# GitHub recommends pinning actions to a commit SHA.\n# To get a newer version, you will need to update the SHA.\n# You can also reference a tag or branch, but the action may change without warning.\n\nname: Upload Python Package\n\non:\n release:\n types: [published]\n\npermissions:\n contents: read\n\njobs:\n release-build:\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v4\n\n - uses: actions/setup-python@v5\n with:\n python-version: \"3.x\"\n\n - name: Build release distributions\n run: |\n # NOTE: put your own distribution build steps here.\n python -m pip install build\n python -m build\n\n - name: Upload distributions\n uses: actions/upload-artifact@v4\n with:\n name: release-dists\n path: dist/\n\n pypi-publish:\n runs-on: ubuntu-latest\n\n needs:\n - release-build\n\n permissions:\n # IMPORTANT: this permission is mandatory for trusted publishing\n id-token: write\n\n # Dedicated environments with protections for publishing are strongly recommended.\n environment:\n name: pypi\n # OPTIONAL: uncomment and update to include your PyPI project URL in the deployment status:\n # url: https://pypi.org/p/YOURPROJECT\n\n steps:\n - name: Retrieve release distributions\n uses: actions/download-artifact@v4\n with:\n name: release-dists\n path: dist/\n\n - name: Publish release distributions to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n" }, { "path": ".gitignore", "content": "# Customed\n.vscode/\ntests/\nruns/\ndata/\nckpts/\nwandb/\nresults/\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/latest/usage/project/#working-with-version-control\n.pdm.toml\n.pdm-python\n.pdm-build/\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/\n" }, { "path": ".gitmodules", "content": "[submodule \"src/third_party/BigVGAN\"]\n\tpath = src/third_party/BigVGAN\n\turl = https://github.com/NVIDIA/BigVGAN.git\n" }, { "path": ".pre-commit-config.yaml", "content": "repos:\n - repo: https://github.com/astral-sh/ruff-pre-commit\n # Ruff version.\n rev: v0.11.2\n hooks:\n - id: ruff\n name: ruff linter\n args: [--fix]\n - id: ruff-format\n name: ruff formatter\n - id: ruff\n name: ruff sorter\n args: [--select, I, --fix]\n - repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v5.0.0\n hooks:\n - id: check-yaml\n" }, { "path": "Dockerfile", "content": "FROM pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel\n\nUSER root\n\nARG DEBIAN_FRONTEND=noninteractive\n\nLABEL github_repo=\"https://github.com/SWivid/F5-TTS\"\n\nRUN set -x \\\n && apt-get update \\\n && apt-get -y install wget curl man git less openssl libssl-dev unzip unar build-essential aria2 tmux vim \\\n && apt-get install -y openssh-server sox libsox-fmt-all libsox-fmt-mp3 libsndfile1-dev ffmpeg \\\n && apt-get install -y librdmacm1 libibumad3 librdmacm-dev libibverbs1 libibverbs-dev ibverbs-utils ibverbs-providers \\\n && rm -rf /var/lib/apt/lists/* \\\n && apt-get clean\n \nWORKDIR /workspace\n\nRUN git clone https://github.com/SWivid/F5-TTS.git \\\n && cd F5-TTS \\\n && git submodule update --init --recursive \\\n && pip install -e . --no-cache-dir\n\nENV SHELL=/bin/bash\n\nVOLUME /root/.cache/huggingface/hub/\n\nEXPOSE 7860\n\nWORKDIR /workspace/F5-TTS\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2024 Yushen CHEN\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching\n\n[![python](https://img.shields.io/badge/Python-3.10-brightgreen)](https://github.com/SWivid/F5-TTS)\n[![arXiv](https://img.shields.io/badge/arXiv-2410.06885-b31b1b.svg?logo=arXiv)](https://arxiv.org/abs/2410.06885)\n[![demo](https://img.shields.io/badge/GitHub-Demo-orange.svg)](https://swivid.github.io/F5-TTS/)\n[![hfspace](https://img.shields.io/badge/🤗-HF%20Space-yellow)](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)\n[![msspace](https://img.shields.io/badge/🤖-MS%20Space-blue)](https://modelscope.cn/studios/AI-ModelScope/E2-F5-TTS)\n[![lab](https://img.shields.io/badge/🏫-X--LANCE-grey?labelColor=lightgrey)](https://x-lance.sjtu.edu.cn/)\n[![lab](https://img.shields.io/badge/🏫-SII-grey?labelColor=lightgrey)](https://www.sii.edu.cn/)\n[![lab](https://img.shields.io/badge/🏫-PCL-grey?labelColor=lightgrey)](https://www.pcl.ac.cn)\n\n\n**F5-TTS**: Diffusion Transformer with ConvNeXt V2, faster trained and inference.\n\n**E2 TTS**: Flat-UNet Transformer, closest reproduction from [paper](https://arxiv.org/abs/2406.18009).\n\n**Sway Sampling**: Inference-time flow step sampling strategy, greatly improves performance\n\n### Thanks to all the contributors !\n\n## News\n- **2025/03/12**: 🔥 F5-TTS v1 base model with better training and inference performance. [Few demo](https://swivid.github.io/F5-TTS_updates).\n- **2024/10/08**: F5-TTS & E2 TTS base models on [🤗 Hugging Face](https://huggingface.co/SWivid/F5-TTS), [🤖 Model Scope](https://www.modelscope.cn/models/SWivid/F5-TTS_Emilia-ZH-EN), [🟣 Wisemodel](https://wisemodel.cn/models/SJTU_X-LANCE/F5-TTS_Emilia-ZH-EN).\n\n## Installation\n\n### Create a separate environment if needed\n\n```bash\n# Create a conda env with python_version>=3.10 (you could also use virtualenv)\nconda create -n f5-tts python=3.11\nconda activate f5-tts\n\n# Install FFmpeg if you haven't yet\nconda install ffmpeg\n```\n\n### Install PyTorch with matched device\n\n
\nNVIDIA GPU\n\n> ```bash\n> # Install pytorch with your CUDA version, e.g.\n> pip install torch==2.8.0+cu128 torchaudio==2.8.0+cu128 --extra-index-url https://download.pytorch.org/whl/cu128\n> \n> # And also possible previous versions, e.g.\n> pip install torch==2.4.0+cu124 torchaudio==2.4.0+cu124 --extra-index-url https://download.pytorch.org/whl/cu124\n> # etc.\n> ```\n\n
\n\n
\nAMD GPU\n\n> ```bash\n> # Install pytorch with your ROCm version (Linux only), e.g.\n> pip install torch==2.5.1+rocm6.2 torchaudio==2.5.1+rocm6.2 --extra-index-url https://download.pytorch.org/whl/rocm6.2\n> ```\n\n
\n\n
\nIntel GPU\n\n> ```bash\n> # Install pytorch with your XPU version, e.g.\n> # Intel® Deep Learning Essentials or Intel® oneAPI Base Toolkit must be installed\n> pip install torch torchaudio --index-url https://download.pytorch.org/whl/test/xpu\n> \n> # Intel GPU support is also available through IPEX (Intel® Extension for PyTorch)\n> # IPEX does not require the Intel® Deep Learning Essentials or Intel® oneAPI Base Toolkit\n> # See: https://pytorch-extension.intel.com/installation?request=platform\n> ```\n\n
\n\n
\nApple Silicon\n\n> ```bash\n> # Install the stable pytorch, e.g.\n> pip install torch torchaudio\n> ```\n\n
\n\n### Then you can choose one from below:\n\n> ### 1. As a pip package (if just for inference)\n> \n> ```bash\n> pip install f5-tts\n> ```\n> \n> ### 2. Local editable (if also do training, finetuning)\n> \n> ```bash\n> git clone https://github.com/SWivid/F5-TTS.git\n> cd F5-TTS\n> # git submodule update --init --recursive # (optional, if use bigvgan as vocoder)\n> pip install -e .\n> ```\n\n### Docker usage also available\n```bash\n# Build from Dockerfile\ndocker build -t f5tts:v1 .\n\n# Run from GitHub Container Registry\ndocker container run --rm -it --gpus=all --mount 'type=volume,source=f5-tts,target=/root/.cache/huggingface/hub/' -p 7860:7860 ghcr.io/swivid/f5-tts:main\n\n# Quickstart if you want to just run the web interface (not CLI)\ndocker container run --rm -it --gpus=all --mount 'type=volume,source=f5-tts,target=/root/.cache/huggingface/hub/' -p 7860:7860 ghcr.io/swivid/f5-tts:main f5-tts_infer-gradio --host 0.0.0.0\n```\n\n### Runtime\n\nDeployment solution with Triton and TensorRT-LLM.\n\n#### Benchmark Results\nDecoding on a single L20 GPU, using 26 different prompt_audio & target_text pairs, 16 NFE.\n\n| Model | Concurrency | Avg Latency | RTF | Mode |\n|---------------------|----------------|-------------|--------|-----------------|\n| F5-TTS Base (Vocos) | 2 | 253 ms | 0.0394 | Client-Server |\n| F5-TTS Base (Vocos) | 1 (Batch_size) | - | 0.0402 | Offline TRT-LLM |\n| F5-TTS Base (Vocos) | 1 (Batch_size) | - | 0.1467 | Offline Pytorch |\n\nSee [detailed instructions](src/f5_tts/runtime/triton_trtllm/README.md) for more information.\n\n\n## Inference\n\n- In order to achieve desired performance, take a moment to read [detailed guidance](src/f5_tts/infer).\n- By properly searching the keywords of problem encountered, [issues](https://github.com/SWivid/F5-TTS/issues?q=is%3Aissue) are very helpful.\n\n### 1. Gradio App\n\nCurrently supported features:\n\n- Basic TTS with Chunk Inference\n- Multi-Style / Multi-Speaker Generation\n- Voice Chat powered by Qwen2.5-3B-Instruct\n- [Custom inference with more language support](src/f5_tts/infer/SHARED.md)\n\n```bash\n# Launch a Gradio app (web interface)\nf5-tts_infer-gradio\n\n# Specify the port/host\nf5-tts_infer-gradio --port 7860 --host 0.0.0.0\n\n# Launch a share link\nf5-tts_infer-gradio --share\n```\n\n
\nNVIDIA device docker compose file example\n\n```yaml\nservices:\n f5-tts:\n image: ghcr.io/swivid/f5-tts:main\n ports:\n - \"7860:7860\"\n environment:\n GRADIO_SERVER_PORT: 7860\n entrypoint: [\"f5-tts_infer-gradio\", \"--port\", \"7860\", \"--host\", \"0.0.0.0\"]\n deploy:\n resources:\n reservations:\n devices:\n - driver: nvidia\n count: 1\n capabilities: [gpu]\n\nvolumes:\n f5-tts:\n driver: local\n```\n\n
\n\n### 2. CLI Inference\n\n```bash\n# Run with flags\n# Leave --ref_text \"\" will have ASR model transcribe (extra GPU memory usage)\nf5-tts_infer-cli --model F5TTS_v1_Base \\\n--ref_audio \"provide_prompt_wav_path_here.wav\" \\\n--ref_text \"The content, subtitle or transcription of reference audio.\" \\\n--gen_text \"Some text you want TTS model generate for you.\"\n\n# Run with default setting. src/f5_tts/infer/examples/basic/basic.toml\nf5-tts_infer-cli\n# Or with your own .toml file\nf5-tts_infer-cli -c custom.toml\n\n# Multi voice. See src/f5_tts/infer/README.md\nf5-tts_infer-cli -c src/f5_tts/infer/examples/multi/story.toml\n```\n\n\n## Training\n\n### 1. With Hugging Face Accelerate\n\nRefer to [training & finetuning guidance](src/f5_tts/train) for best practice.\n\n### 2. With Gradio App\n\n```bash\n# Quick start with Gradio web interface\nf5-tts_finetune-gradio\n```\n\nRead [training & finetuning guidance](src/f5_tts/train) for more instructions.\n\n\n## [Evaluation](src/f5_tts/eval)\n\n\n## Development\n\nUse pre-commit to ensure code quality (will run linters and formatters automatically):\n\n```bash\npip install pre-commit\npre-commit install\n```\n\nWhen making a pull request, before each commit, run: \n\n```bash\npre-commit run --all-files\n```\n\nNote: Some model components have linting exceptions for E722 to accommodate tensor notation.\n\n\n## Acknowledgements\n\n- [E2-TTS](https://arxiv.org/abs/2406.18009) brilliant work, simple and effective\n- [Emilia](https://arxiv.org/abs/2407.05361), [WenetSpeech4TTS](https://arxiv.org/abs/2406.05763), [LibriTTS](https://arxiv.org/abs/1904.02882), [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) valuable datasets\n- [lucidrains](https://github.com/lucidrains) initial CFM structure with also [bfs18](https://github.com/bfs18) for discussion\n- [SD3](https://arxiv.org/abs/2403.03206) & [Hugging Face diffusers](https://github.com/huggingface/diffusers) DiT and MMDiT code structure\n- [torchdiffeq](https://github.com/rtqichen/torchdiffeq) as ODE solver, [Vocos](https://huggingface.co/charactr/vocos-mel-24khz) and [BigVGAN](https://github.com/NVIDIA/BigVGAN) as vocoder\n- [FunASR](https://github.com/modelscope/FunASR), [faster-whisper](https://github.com/SYSTRAN/faster-whisper), [UniSpeech](https://github.com/microsoft/UniSpeech), [SpeechMOS](https://github.com/tarepan/SpeechMOS) for evaluation tools\n- [ctc-forced-aligner](https://github.com/MahmoudAshraf97/ctc-forced-aligner) for speech edit test\n- [mrfakename](https://x.com/realmrfakename) huggingface space demo ~\n- [f5-tts-mlx](https://github.com/lucasnewman/f5-tts-mlx/tree/main) Implementation with MLX framework by [Lucas Newman](https://github.com/lucasnewman)\n- [F5-TTS-ONNX](https://github.com/DakeQQ/F5-TTS-ONNX) ONNX Runtime version by [DakeQQ](https://github.com/DakeQQ)\n- [Yuekai Zhang](https://github.com/yuekaizhang) Triton and TensorRT-LLM support ~\n\n## Citation\nIf our work and codebase is useful for you, please cite as:\n```\n@article{chen-etal-2024-f5tts,\n title={F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching}, \n author={Yushen Chen and Zhikang Niu and Ziyang Ma and Keqi Deng and Chunhui Wang and Jian Zhao and Kai Yu and Xie Chen},\n journal={arXiv preprint arXiv:2410.06885},\n year={2024},\n}\n```\n## License\n\nOur code is released under MIT License. The pre-trained models are licensed under the CC-BY-NC license due to the training data Emilia, which is an in-the-wild dataset. Sorry for any inconvenience this may cause.\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"setuptools >= 61.0\", \"setuptools-scm>=8.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[project]\nname = \"f5-tts\"\nversion = \"1.1.17\"\ndescription = \"F5-TTS: A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching\"\nreadme = \"README.md\"\nlicense = {text = \"MIT License\"}\nclassifiers = [\n \"License :: OSI Approved :: MIT License\",\n \"Operating System :: OS Independent\",\n \"Programming Language :: Python :: 3\",\n]\ndependencies = [\n \"accelerate>=0.33.0\",\n \"bitsandbytes>0.37.0; platform_machine!='arm64' and platform_system!='Darwin'\",\n \"cached_path\",\n \"click\",\n \"datasets\",\n \"ema_pytorch>=0.5.2\",\n \"gradio>=6.0.0\",\n \"hydra-core>=1.3.0\",\n \"librosa\",\n \"matplotlib\",\n \"numpy<=1.26.4; python_version<='3.10'\",\n \"pydub\",\n \"pypinyin\",\n \"rjieba\",\n \"safetensors\",\n \"soundfile\",\n \"tomli\",\n \"torch>=2.0.0\",\n \"torchaudio>=2.0.0\",\n \"torchcodec\",\n \"torchdiffeq\",\n \"tqdm>=4.65.0\",\n \"transformers\",\n \"transformers_stream_generator\",\n \"unidecode\",\n \"vocos\",\n \"wandb\",\n \"x_transformers>=1.31.14\",\n]\n\n[project.optional-dependencies]\neval = [\n \"faster_whisper==0.10.1\",\n \"funasr\",\n \"jiwer\",\n \"modelscope\",\n \"zhconv\",\n \"zhon\",\n]\n\n[project.urls]\nHomepage = \"https://github.com/SWivid/F5-TTS\"\n\n[project.scripts]\n\"f5-tts_infer-cli\" = \"f5_tts.infer.infer_cli:main\"\n\"f5-tts_infer-gradio\" = \"f5_tts.infer.infer_gradio:main\"\n\"f5-tts_finetune-cli\" = \"f5_tts.train.finetune_cli:main\"\n\"f5-tts_finetune-gradio\" = \"f5_tts.train.finetune_gradio:main\"\n" }, { "path": "ruff.toml", "content": "line-length = 120\ntarget-version = \"py310\"\n\n[lint]\n# Only ignore variables with names starting with \"_\".\ndummy-variable-rgx = \"^_.*$\"\n\n[lint.isort]\nforce-single-line = false\nlines-after-imports = 2\n" }, { "path": "src/f5_tts/api.py", "content": "import random\nimport sys\nfrom importlib.resources import files\n\nimport soundfile as sf\nimport tqdm\nfrom cached_path import cached_path\nfrom hydra.utils import get_class\nfrom omegaconf import OmegaConf\n\nfrom f5_tts.infer.utils_infer import (\n infer_process,\n load_model,\n load_vocoder,\n preprocess_ref_audio_text,\n remove_silence_for_generated_wav,\n save_spectrogram,\n transcribe,\n)\nfrom f5_tts.model.utils import seed_everything\n\n\nclass F5TTS:\n def __init__(\n self,\n model=\"F5TTS_v1_Base\",\n ckpt_file=\"\",\n vocab_file=\"\",\n ode_method=\"euler\",\n use_ema=True,\n vocoder_local_path=None,\n device=None,\n hf_cache_dir=None,\n ):\n model_cfg = OmegaConf.load(str(files(\"f5_tts\").joinpath(f\"configs/{model}.yaml\")))\n model_cls = get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\n model_arc = model_cfg.model.arch\n\n self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type\n self.target_sample_rate = model_cfg.model.mel_spec.target_sample_rate\n\n self.ode_method = ode_method\n self.use_ema = use_ema\n\n if device is not None:\n self.device = device\n else:\n import torch\n\n self.device = (\n \"cuda\"\n if torch.cuda.is_available()\n else \"xpu\"\n if torch.xpu.is_available()\n else \"mps\"\n if torch.backends.mps.is_available()\n else \"cpu\"\n )\n\n # Load models\n self.vocoder = load_vocoder(\n self.mel_spec_type, vocoder_local_path is not None, vocoder_local_path, self.device, hf_cache_dir\n )\n\n repo_name, ckpt_step, ckpt_type = \"F5-TTS\", 1250000, \"safetensors\"\n\n # override for previous models\n if model == \"F5TTS_Base\":\n if self.mel_spec_type == \"vocos\":\n ckpt_step = 1200000\n elif self.mel_spec_type == \"bigvgan\":\n model = \"F5TTS_Base_bigvgan\"\n ckpt_type = \"pt\"\n elif model == \"E2TTS_Base\":\n repo_name = \"E2-TTS\"\n ckpt_step = 1200000\n\n if not ckpt_file:\n ckpt_file = str(\n cached_path(f\"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}\", cache_dir=hf_cache_dir)\n )\n self.ema_model = load_model(\n model_cls, model_arc, ckpt_file, self.mel_spec_type, vocab_file, self.ode_method, self.use_ema, self.device\n )\n\n def transcribe(self, ref_audio, language=None):\n return transcribe(ref_audio, language)\n\n def export_wav(self, wav, file_wave, remove_silence=False):\n sf.write(file_wave, wav, self.target_sample_rate)\n\n if remove_silence:\n remove_silence_for_generated_wav(file_wave)\n\n def export_spectrogram(self, spec, file_spec):\n save_spectrogram(spec, file_spec)\n\n def infer(\n self,\n ref_file,\n ref_text,\n gen_text,\n show_info=print,\n progress=tqdm,\n target_rms=0.1,\n cross_fade_duration=0.15,\n sway_sampling_coef=-1,\n cfg_strength=2,\n nfe_step=32,\n speed=1.0,\n fix_duration=None,\n remove_silence=False,\n file_wave=None,\n file_spec=None,\n seed=None,\n ):\n if seed is None:\n seed = random.randint(0, sys.maxsize)\n seed_everything(seed)\n self.seed = seed\n\n ref_file, ref_text = preprocess_ref_audio_text(ref_file, ref_text, show_info=show_info)\n\n wav, sr, spec = infer_process(\n ref_file,\n ref_text,\n gen_text,\n self.ema_model,\n self.vocoder,\n self.mel_spec_type,\n show_info=show_info,\n progress=progress,\n target_rms=target_rms,\n cross_fade_duration=cross_fade_duration,\n nfe_step=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n speed=speed,\n fix_duration=fix_duration,\n device=self.device,\n )\n\n if file_wave is not None:\n self.export_wav(wav, file_wave, remove_silence)\n\n if file_spec is not None:\n self.export_spectrogram(spec, file_spec)\n\n return wav, sr, spec\n\n\nif __name__ == \"__main__\":\n f5tts = F5TTS()\n\n wav, sr, spec = f5tts.infer(\n ref_file=str(files(\"f5_tts\").joinpath(\"infer/examples/basic/basic_ref_en.wav\")),\n ref_text=\"Some call me nature, others call me mother nature.\",\n 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.\",\n file_wave=str(files(\"f5_tts\").joinpath(\"../../tests/api_out.wav\")),\n file_spec=str(files(\"f5_tts\").joinpath(\"../../tests/api_out.png\")),\n seed=None,\n )\n\n print(\"seed :\", f5tts.seed)\n" }, { "path": "src/f5_tts/configs/E2TTS_Base.yaml", "content": "hydra:\n run:\n dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}\n\ndatasets:\n name: Emilia_ZH_EN # dataset name\n batch_size_per_gpu: 38400 # 8 GPUs, 8 * 38400 = 307200\n batch_size_type: frame # frame | sample\n max_samples: 64 # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models\n num_workers: 16\n\noptim:\n epochs: 11\n learning_rate: 7.5e-5\n num_warmup_updates: 20000 # warmup updates\n grad_accumulation_steps: 1 # note: updates = steps / grad_accumulation_steps\n max_grad_norm: 1.0 # gradient clipping\n bnb_optimizer: False # use bnb 8bit AdamW optimizer or not\n\nmodel:\n name: E2TTS_Base\n tokenizer: pinyin\n tokenizer_path: null # if 'custom' tokenizer, define the path want to use (should be vocab.txt)\n backbone: UNetT\n arch:\n dim: 1024\n depth: 24\n heads: 16\n ff_mult: 4\n text_mask_padding: False\n pe_attn_head: 1\n mel_spec:\n target_sample_rate: 24000\n n_mel_channels: 100\n hop_length: 256\n win_length: 1024\n n_fft: 1024\n mel_spec_type: vocos # vocos | bigvgan\n vocoder:\n is_local: False # use local offline ckpt or not\n local_path: null # local vocoder path\n\nckpts:\n logger: wandb # wandb | tensorboard | null\n wandb_project: CFM-TTS # wandb project name\n wandb_run_name: ${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name} # wandb run name\n wandb_resume_id: null # wandb run id for resuming, null to auto-detect from checkpoint\n log_samples: True # infer random sample per save checkpoint. wip, normal to fail with extra long samples\n save_per_updates: 50000 # save checkpoint per updates\n keep_last_n_checkpoints: -1 # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n last_per_updates: 5000 # save last checkpoint per updates\n save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}" }, { "path": "src/f5_tts/configs/E2TTS_Small.yaml", "content": "hydra:\n run:\n dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}\n\ndatasets:\n name: Emilia_ZH_EN\n batch_size_per_gpu: 38400 # 8 GPUs, 8 * 38400 = 307200\n batch_size_type: frame # frame | sample\n max_samples: 64 # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models\n num_workers: 16\n\noptim:\n epochs: 11\n learning_rate: 7.5e-5\n num_warmup_updates: 20000 # warmup updates\n grad_accumulation_steps: 1 # note: updates = steps / grad_accumulation_steps\n max_grad_norm: 1.0\n bnb_optimizer: False \n\nmodel:\n name: E2TTS_Small\n tokenizer: pinyin\n tokenizer_path: null # if 'custom' tokenizer, define the path want to use (should be vocab.txt)\n backbone: UNetT\n arch:\n dim: 768\n depth: 20\n heads: 12\n ff_mult: 4\n text_mask_padding: False\n pe_attn_head: 1\n mel_spec:\n target_sample_rate: 24000\n n_mel_channels: 100\n hop_length: 256\n win_length: 1024\n n_fft: 1024\n mel_spec_type: vocos # vocos | bigvgan\n vocoder:\n is_local: False # use local offline ckpt or not\n local_path: null # local vocoder path\n\nckpts:\n logger: wandb # wandb | tensorboard | null\n wandb_project: CFM-TTS # wandb project name\n wandb_run_name: ${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name} # wandb run name\n wandb_resume_id: null # wandb run id for resuming, null to auto-detect from checkpoint\n log_samples: True # infer random sample per save checkpoint. wip, normal to fail with extra long samples\n save_per_updates: 50000 # save checkpoint per updates\n keep_last_n_checkpoints: -1 # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n last_per_updates: 5000 # save last checkpoint per updates\n save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}" }, { "path": "src/f5_tts/configs/F5TTS_Base.yaml", "content": "hydra:\n run:\n dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}\n\ndatasets:\n name: Emilia_ZH_EN # dataset name\n batch_size_per_gpu: 38400 # 8 GPUs, 8 * 38400 = 307200\n batch_size_type: frame # frame | sample\n max_samples: 64 # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models\n num_workers: 16\n\noptim:\n epochs: 11\n learning_rate: 7.5e-5\n num_warmup_updates: 20000 # warmup updates\n grad_accumulation_steps: 1 # note: updates = steps / grad_accumulation_steps\n max_grad_norm: 1.0 # gradient clipping\n bnb_optimizer: False # use bnb 8bit AdamW optimizer or not\n\nmodel:\n name: F5TTS_Base # model name\n tokenizer: pinyin # tokenizer type\n tokenizer_path: null # if 'custom' tokenizer, define the path want to use (should be vocab.txt)\n backbone: DiT\n arch:\n dim: 1024\n depth: 22\n heads: 16\n ff_mult: 2\n text_dim: 512\n text_mask_padding: False\n conv_layers: 4\n pe_attn_head: 1\n attn_backend: torch # torch | flash_attn\n attn_mask_enabled: False\n checkpoint_activations: False # recompute activations and save memory for extra compute\n mel_spec:\n target_sample_rate: 24000\n n_mel_channels: 100\n hop_length: 256\n win_length: 1024\n n_fft: 1024\n mel_spec_type: vocos # vocos | bigvgan\n vocoder:\n is_local: False # use local offline ckpt or not\n local_path: null # local vocoder path\n\nckpts:\n logger: wandb # wandb | tensorboard | null\n wandb_project: CFM-TTS # wandb project name\n wandb_run_name: ${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name} # wandb run name\n wandb_resume_id: null # wandb run id for resuming, null to auto-detect from checkpoint\n log_samples: True # infer random sample per save checkpoint. wip, normal to fail with extra long samples\n save_per_updates: 50000 # save checkpoint per updates\n keep_last_n_checkpoints: -1 # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n last_per_updates: 5000 # save last checkpoint per updates\n save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}" }, { "path": "src/f5_tts/configs/F5TTS_Small.yaml", "content": "hydra:\n run:\n dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}\n\ndatasets:\n name: Emilia_ZH_EN\n batch_size_per_gpu: 38400 # 8 GPUs, 8 * 38400 = 307200\n batch_size_type: frame # frame | sample\n max_samples: 64 # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models\n num_workers: 16\n\noptim:\n epochs: 11 # only suitable for Emilia, if you want to train it on LibriTTS, set epoch 686\n learning_rate: 7.5e-5\n num_warmup_updates: 20000 # warmup updates\n grad_accumulation_steps: 1 # note: updates = steps / grad_accumulation_steps\n max_grad_norm: 1.0 # gradient clipping\n bnb_optimizer: False # use bnb 8bit AdamW optimizer or not\n\nmodel:\n name: F5TTS_Small\n tokenizer: pinyin\n tokenizer_path: null # if 'custom' tokenizer, define the path want to use (should be vocab.txt)\n backbone: DiT\n arch:\n dim: 768\n depth: 18\n heads: 12\n ff_mult: 2\n text_dim: 512\n text_mask_padding: False\n conv_layers: 4\n pe_attn_head: 1\n attn_backend: torch # torch | flash_attn\n attn_mask_enabled: False\n checkpoint_activations: False # recompute activations and save memory for extra compute\n mel_spec:\n target_sample_rate: 24000\n n_mel_channels: 100\n hop_length: 256\n win_length: 1024\n n_fft: 1024\n mel_spec_type: vocos # vocos | bigvgan\n vocoder:\n is_local: False # use local offline ckpt or not\n local_path: null # local vocoder path\n\nckpts:\n logger: wandb # wandb | tensorboard | null\n wandb_project: CFM-TTS # wandb project name\n wandb_run_name: ${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name} # wandb run name\n wandb_resume_id: null # wandb run id for resuming, null to auto-detect from checkpoint\n log_samples: True # infer random sample per save checkpoint. wip, normal to fail with extra long samples\n save_per_updates: 50000 # save checkpoint per updates\n keep_last_n_checkpoints: -1 # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n last_per_updates: 5000 # save last checkpoint per updates\n save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}\n" }, { "path": "src/f5_tts/configs/F5TTS_v1_Base.yaml", "content": "hydra:\n run:\n dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}\n\ndatasets:\n name: Emilia_ZH_EN # dataset name\n batch_size_per_gpu: 38400 # 8 GPUs, 8 * 38400 = 307200\n batch_size_type: frame # frame | sample\n max_samples: 64 # max sequences per batch if use frame-wise batch_size. we set 32 for small models, 64 for base models\n num_workers: 16\n\noptim:\n epochs: 11\n learning_rate: 7.5e-5\n num_warmup_updates: 20000 # warmup updates\n grad_accumulation_steps: 1 # note: updates = steps / grad_accumulation_steps\n max_grad_norm: 1.0 # gradient clipping\n bnb_optimizer: False # use bnb 8bit AdamW optimizer or not\n\nmodel:\n name: F5TTS_v1_Base # model name\n tokenizer: pinyin # tokenizer type\n tokenizer_path: null # if 'custom' tokenizer, define the path want to use (should be vocab.txt)\n backbone: DiT\n arch:\n dim: 1024\n depth: 22\n heads: 16\n ff_mult: 2\n text_dim: 512\n text_mask_padding: True\n qk_norm: null # null | rms_norm\n conv_layers: 4\n pe_attn_head: null\n attn_backend: torch # torch | flash_attn\n attn_mask_enabled: False\n checkpoint_activations: False # recompute activations and save memory for extra compute\n mel_spec:\n target_sample_rate: 24000\n n_mel_channels: 100\n hop_length: 256\n win_length: 1024\n n_fft: 1024\n mel_spec_type: vocos # vocos | bigvgan\n vocoder:\n is_local: False # use local offline ckpt or not\n local_path: null # local vocoder path\n\nckpts:\n logger: wandb # wandb | tensorboard | null\n wandb_project: CFM-TTS # wandb project name\n wandb_run_name: ${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name} # wandb run name\n wandb_resume_id: null # wandb run id for resuming, null to auto-detect from checkpoint\n log_samples: True # infer random sample per save checkpoint. wip, normal to fail with extra long samples\n save_per_updates: 50000 # save checkpoint per updates\n keep_last_n_checkpoints: -1 # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n last_per_updates: 5000 # save last checkpoint per updates\n save_dir: ckpts/${model.name}_${model.mel_spec.mel_spec_type}_${model.tokenizer}_${datasets.name}" }, { "path": "src/f5_tts/eval/README.md", "content": "\n# Evaluation\n\nInstall packages for evaluation:\n\n```bash\npip install -e .[eval]\n```\n\n> [!IMPORTANT]\n> For [faster-whisper](https://github.com/SYSTRAN/faster-whisper), for various compatibilities: \n> `pip install ctranslate2==4.5.0` if CUDA 12 and cuDNN 9; \n> `pip install ctranslate2==4.4.0` if CUDA 12 and cuDNN 8; \n> `pip install ctranslate2==3.24.0` if CUDA 11 and cuDNN 8.\n\n## Generating Samples for Evaluation\n\n### Prepare Test Datasets\n\n1. *Seed-TTS testset*: Download from [seed-tts-eval](https://github.com/BytedanceSpeech/seed-tts-eval).\n2. *LibriSpeech test-clean*: Download from [OpenSLR](http://www.openslr.org/12/).\n3. Unzip the downloaded datasets and place them in the `data/` directory.\n4. Our filtered LibriSpeech-PC 4-10s subset: `data/librispeech_pc_test_clean_cross_sentence.lst`\n\n### Batch Inference for Test Set\n\nTo run batch inference for evaluations, execute the following commands:\n\n```bash\n# if not setup accelerate config yet\naccelerate config\n\n# if only perform inference\nbash src/f5_tts/eval/eval_infer_batch.sh --infer-only\n\n# if inference and with corresponding evaluation, setup the following tools first\nbash src/f5_tts/eval/eval_infer_batch.sh\n```\n\n## Objective Evaluation on Generated Results\n\n### Download Evaluation Model Checkpoints\n\n1. Chinese ASR Model: [Paraformer-zh](https://huggingface.co/funasr/paraformer-zh)\n2. English ASR Model: [Faster-Whisper](https://huggingface.co/Systran/faster-whisper-large-v3)\n3. WavLM Model: Download from [Google Drive](https://drive.google.com/file/d/1-aE1NfzpRCLxA4GUxX9ITI3F9LlbtEGP/view).\n\n> [!NOTE] \n> ASR model will be automatically downloaded if `--local` not set for evaluation scripts. \n> Otherwise, you should update the `asr_ckpt_dir` path values in `eval_librispeech_test_clean.py` or `eval_seedtts_testset.py`.\n> \n> WavLM model must be downloaded and your `wavlm_ckpt_dir` path updated in `eval_librispeech_test_clean.py` and `eval_seedtts_testset.py`.\n\n### Objective Evaluation Examples\n\nUpdate the path with your batch-inferenced results, and carry out WER / SIM / UTMOS evaluations:\n```bash\n# Evaluation [WER] for Seed-TTS test [ZH] set\npython src/f5_tts/eval/eval_seedtts_testset.py --eval_task wer --lang zh --gen_wav_dir --gpu_nums 8\n\n# Evaluation [SIM] for LibriSpeech-PC test-clean (cross-sentence)\npython src/f5_tts/eval/eval_librispeech_test_clean.py --eval_task sim --gen_wav_dir --librispeech_test_clean_path \n\n# Evaluation [UTMOS]. --ext: Audio extension\npython src/f5_tts/eval/eval_utmos.py --audio_dir --ext wav\n```\n\n> [!NOTE] \n> Evaluation results can also be found in `_*_results.jsonl` files saved in ``/``.\n" }, { "path": "src/f5_tts/eval/ecapa_tdnn.py", "content": "# just for speaker similarity evaluation, third-party code\n\n# From https://github.com/microsoft/UniSpeech/blob/main/downstreams/speaker_verification/models/\n# part of the code is borrowed from https://github.com/lawlict/ECAPA-TDNN\n\nimport os\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\n\"\"\" Res2Conv1d + BatchNorm1d + ReLU\n\"\"\"\n\n\nclass Res2Conv1dReluBn(nn.Module):\n \"\"\"\n in_channels == out_channels == channels\n \"\"\"\n\n def __init__(self, channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True, scale=4):\n super().__init__()\n assert channels % scale == 0, \"{} % {} != 0\".format(channels, scale)\n self.scale = scale\n self.width = channels // scale\n self.nums = scale if scale == 1 else scale - 1\n\n self.convs = []\n self.bns = []\n for i in range(self.nums):\n self.convs.append(nn.Conv1d(self.width, self.width, kernel_size, stride, padding, dilation, bias=bias))\n self.bns.append(nn.BatchNorm1d(self.width))\n self.convs = nn.ModuleList(self.convs)\n self.bns = nn.ModuleList(self.bns)\n\n def forward(self, x):\n out = []\n spx = torch.split(x, self.width, 1)\n for i in range(self.nums):\n if i == 0:\n sp = spx[i]\n else:\n sp = sp + spx[i]\n # Order: conv -> relu -> bn\n sp = self.convs[i](sp)\n sp = self.bns[i](F.relu(sp))\n out.append(sp)\n if self.scale != 1:\n out.append(spx[self.nums])\n out = torch.cat(out, dim=1)\n\n return out\n\n\n\"\"\" Conv1d + BatchNorm1d + ReLU\n\"\"\"\n\n\nclass Conv1dReluBn(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, bias=True):\n super().__init__()\n self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, dilation, bias=bias)\n self.bn = nn.BatchNorm1d(out_channels)\n\n def forward(self, x):\n return self.bn(F.relu(self.conv(x)))\n\n\n\"\"\" The SE connection of 1D case.\n\"\"\"\n\n\nclass SE_Connect(nn.Module):\n def __init__(self, channels, se_bottleneck_dim=128):\n super().__init__()\n self.linear1 = nn.Linear(channels, se_bottleneck_dim)\n self.linear2 = nn.Linear(se_bottleneck_dim, channels)\n\n def forward(self, x):\n out = x.mean(dim=2)\n out = F.relu(self.linear1(out))\n out = torch.sigmoid(self.linear2(out))\n out = x * out.unsqueeze(2)\n\n return out\n\n\n\"\"\" SE-Res2Block of the ECAPA-TDNN architecture.\n\"\"\"\n\n# def SE_Res2Block(channels, kernel_size, stride, padding, dilation, scale):\n# return nn.Sequential(\n# Conv1dReluBn(channels, 512, kernel_size=1, stride=1, padding=0),\n# Res2Conv1dReluBn(512, kernel_size, stride, padding, dilation, scale=scale),\n# Conv1dReluBn(512, channels, kernel_size=1, stride=1, padding=0),\n# SE_Connect(channels)\n# )\n\n\nclass SE_Res2Block(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, scale, se_bottleneck_dim):\n super().__init__()\n self.Conv1dReluBn1 = Conv1dReluBn(in_channels, out_channels, kernel_size=1, stride=1, padding=0)\n self.Res2Conv1dReluBn = Res2Conv1dReluBn(out_channels, kernel_size, stride, padding, dilation, scale=scale)\n self.Conv1dReluBn2 = Conv1dReluBn(out_channels, out_channels, kernel_size=1, stride=1, padding=0)\n self.SE_Connect = SE_Connect(out_channels, se_bottleneck_dim)\n\n self.shortcut = None\n if in_channels != out_channels:\n self.shortcut = nn.Conv1d(\n in_channels=in_channels,\n out_channels=out_channels,\n kernel_size=1,\n )\n\n def forward(self, x):\n residual = x\n if self.shortcut:\n residual = self.shortcut(x)\n\n x = self.Conv1dReluBn1(x)\n x = self.Res2Conv1dReluBn(x)\n x = self.Conv1dReluBn2(x)\n x = self.SE_Connect(x)\n\n return x + residual\n\n\n\"\"\" Attentive weighted mean and standard deviation pooling.\n\"\"\"\n\n\nclass AttentiveStatsPool(nn.Module):\n def __init__(self, in_dim, attention_channels=128, global_context_att=False):\n super().__init__()\n self.global_context_att = global_context_att\n\n # Use Conv1d with stride == 1 rather than Linear, then we don't need to transpose inputs.\n if global_context_att:\n self.linear1 = nn.Conv1d(in_dim * 3, attention_channels, kernel_size=1) # equals W and b in the paper\n else:\n self.linear1 = nn.Conv1d(in_dim, attention_channels, kernel_size=1) # equals W and b in the paper\n self.linear2 = nn.Conv1d(attention_channels, in_dim, kernel_size=1) # equals V and k in the paper\n\n def forward(self, x):\n if self.global_context_att:\n context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)\n context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)\n x_in = torch.cat((x, context_mean, context_std), dim=1)\n else:\n x_in = x\n\n # DON'T use ReLU here! In experiments, I find ReLU hard to converge.\n alpha = torch.tanh(self.linear1(x_in))\n # alpha = F.relu(self.linear1(x_in))\n alpha = torch.softmax(self.linear2(alpha), dim=2)\n mean = torch.sum(alpha * x, dim=2)\n residuals = torch.sum(alpha * (x**2), dim=2) - mean**2\n std = torch.sqrt(residuals.clamp(min=1e-9))\n return torch.cat([mean, std], dim=1)\n\n\nclass ECAPA_TDNN(nn.Module):\n def __init__(\n self,\n feat_dim=80,\n channels=512,\n emb_dim=192,\n global_context_att=False,\n feat_type=\"wavlm_large\",\n sr=16000,\n feature_selection=\"hidden_states\",\n update_extract=False,\n config_path=None,\n ):\n super().__init__()\n\n self.feat_type = feat_type\n self.feature_selection = feature_selection\n self.update_extract = update_extract\n self.sr = sr\n\n torch.hub._validate_not_a_forked_repo = lambda a, b, c: True\n try:\n local_s3prl_path = os.path.expanduser(\"~/.cache/torch/hub/s3prl_s3prl_main\")\n self.feature_extract = torch.hub.load(local_s3prl_path, feat_type, source=\"local\", config_path=config_path)\n except: # noqa: E722\n self.feature_extract = torch.hub.load(\"s3prl/s3prl\", feat_type)\n\n if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(\n self.feature_extract.model.encoder.layers[23].self_attn, \"fp32_attention\"\n ):\n self.feature_extract.model.encoder.layers[23].self_attn.fp32_attention = False\n if len(self.feature_extract.model.encoder.layers) == 24 and hasattr(\n self.feature_extract.model.encoder.layers[11].self_attn, \"fp32_attention\"\n ):\n self.feature_extract.model.encoder.layers[11].self_attn.fp32_attention = False\n\n self.feat_num = self.get_feat_num()\n self.feature_weight = nn.Parameter(torch.zeros(self.feat_num))\n\n if feat_type != \"fbank\" and feat_type != \"mfcc\":\n freeze_list = [\"final_proj\", \"label_embs_concat\", \"mask_emb\", \"project_q\", \"quantizer\"]\n for name, param in self.feature_extract.named_parameters():\n for freeze_val in freeze_list:\n if freeze_val in name:\n param.requires_grad = False\n break\n\n if not self.update_extract:\n for param in self.feature_extract.parameters():\n param.requires_grad = False\n\n self.instance_norm = nn.InstanceNorm1d(feat_dim)\n # self.channels = [channels] * 4 + [channels * 3]\n self.channels = [channels] * 4 + [1536]\n\n self.layer1 = Conv1dReluBn(feat_dim, self.channels[0], kernel_size=5, padding=2)\n self.layer2 = SE_Res2Block(\n self.channels[0],\n self.channels[1],\n kernel_size=3,\n stride=1,\n padding=2,\n dilation=2,\n scale=8,\n se_bottleneck_dim=128,\n )\n self.layer3 = SE_Res2Block(\n self.channels[1],\n self.channels[2],\n kernel_size=3,\n stride=1,\n padding=3,\n dilation=3,\n scale=8,\n se_bottleneck_dim=128,\n )\n self.layer4 = SE_Res2Block(\n self.channels[2],\n self.channels[3],\n kernel_size=3,\n stride=1,\n padding=4,\n dilation=4,\n scale=8,\n se_bottleneck_dim=128,\n )\n\n # self.conv = nn.Conv1d(self.channels[-1], self.channels[-1], kernel_size=1)\n cat_channels = channels * 3\n self.conv = nn.Conv1d(cat_channels, self.channels[-1], kernel_size=1)\n self.pooling = AttentiveStatsPool(\n self.channels[-1], attention_channels=128, global_context_att=global_context_att\n )\n self.bn = nn.BatchNorm1d(self.channels[-1] * 2)\n self.linear = nn.Linear(self.channels[-1] * 2, emb_dim)\n\n def get_feat_num(self):\n self.feature_extract.eval()\n wav = [torch.randn(self.sr).to(next(self.feature_extract.parameters()).device)]\n with torch.no_grad():\n features = self.feature_extract(wav)\n select_feature = features[self.feature_selection]\n if isinstance(select_feature, (list, tuple)):\n return len(select_feature)\n else:\n return 1\n\n def get_feat(self, x):\n if self.update_extract:\n x = self.feature_extract([sample for sample in x])\n else:\n with torch.no_grad():\n if self.feat_type == \"fbank\" or self.feat_type == \"mfcc\":\n x = self.feature_extract(x) + 1e-6 # B x feat_dim x time_len\n else:\n x = self.feature_extract([sample for sample in x])\n\n if self.feat_type == \"fbank\":\n x = x.log()\n\n if self.feat_type != \"fbank\" and self.feat_type != \"mfcc\":\n x = x[self.feature_selection]\n if isinstance(x, (list, tuple)):\n x = torch.stack(x, dim=0)\n else:\n x = x.unsqueeze(0)\n norm_weights = F.softmax(self.feature_weight, dim=-1).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)\n x = (norm_weights * x).sum(dim=0)\n x = torch.transpose(x, 1, 2) + 1e-6\n\n x = self.instance_norm(x)\n return x\n\n def forward(self, x):\n x = self.get_feat(x)\n\n out1 = self.layer1(x)\n out2 = self.layer2(out1)\n out3 = self.layer3(out2)\n out4 = self.layer4(out3)\n\n out = torch.cat([out2, out3, out4], dim=1)\n out = F.relu(self.conv(out))\n out = self.bn(self.pooling(out))\n out = self.linear(out)\n\n return out\n\n\ndef ECAPA_TDNN_SMALL(\n feat_dim,\n emb_dim=256,\n feat_type=\"wavlm_large\",\n sr=16000,\n feature_selection=\"hidden_states\",\n update_extract=False,\n config_path=None,\n):\n return ECAPA_TDNN(\n feat_dim=feat_dim,\n channels=512,\n emb_dim=emb_dim,\n feat_type=feat_type,\n sr=sr,\n feature_selection=feature_selection,\n update_extract=update_extract,\n config_path=config_path,\n )\n" }, { "path": "src/f5_tts/eval/eval_infer_batch.py", "content": "import os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport argparse\nimport time\nfrom importlib.resources import files\n\nimport torch\nimport torchaudio\nfrom accelerate import Accelerator\nfrom hydra.utils import get_class\nfrom omegaconf import OmegaConf\nfrom tqdm import tqdm\n\nfrom f5_tts.eval.utils_eval import (\n get_inference_prompt,\n get_librispeech_test_clean_metainfo,\n get_seedtts_testset_metainfo,\n)\nfrom f5_tts.infer.utils_infer import load_checkpoint, load_vocoder\nfrom f5_tts.model import CFM\nfrom f5_tts.model.utils import get_tokenizer\n\n\naccelerator = Accelerator()\ndevice = f\"cuda:{accelerator.process_index}\"\n\n\nuse_ema = True\ntarget_rms = 0.1\n\n\nrel_path = str(files(\"f5_tts\").joinpath(\"../../\"))\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"batch inference\")\n\n parser.add_argument(\"-s\", \"--seed\", default=None, type=int)\n parser.add_argument(\"-n\", \"--expname\", required=True)\n parser.add_argument(\"-c\", \"--ckptstep\", default=1250000, type=int)\n\n parser.add_argument(\"-nfe\", \"--nfestep\", default=32, type=int)\n parser.add_argument(\"-o\", \"--odemethod\", default=\"euler\")\n parser.add_argument(\"-ss\", \"--swaysampling\", default=-1, type=float)\n\n parser.add_argument(\"-t\", \"--testset\", required=True)\n parser.add_argument(\n \"-p\", \"--librispeech_test_clean_path\", default=f\"{rel_path}/data/LibriSpeech/test-clean\", type=str\n )\n\n parser.add_argument(\"--local\", action=\"store_true\", help=\"Use local vocoder checkpoint directory\")\n\n args = parser.parse_args()\n\n seed = args.seed\n exp_name = args.expname\n ckpt_step = args.ckptstep\n\n nfe_step = args.nfestep\n ode_method = args.odemethod\n sway_sampling_coef = args.swaysampling\n\n testset = args.testset\n\n infer_batch_size = 1 # max frames. 1 for ddp single inference (recommended)\n cfg_strength = 2.0\n speed = 1.0\n use_truth_duration = False\n no_ref_audio = False\n\n model_cfg = OmegaConf.load(str(files(\"f5_tts\").joinpath(f\"configs/{exp_name}.yaml\")))\n model_cls = get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\n model_arc = model_cfg.model.arch\n\n dataset_name = model_cfg.datasets.name\n tokenizer = model_cfg.model.tokenizer\n\n mel_spec_type = model_cfg.model.mel_spec.mel_spec_type\n target_sample_rate = model_cfg.model.mel_spec.target_sample_rate\n n_mel_channels = model_cfg.model.mel_spec.n_mel_channels\n hop_length = model_cfg.model.mel_spec.hop_length\n win_length = model_cfg.model.mel_spec.win_length\n n_fft = model_cfg.model.mel_spec.n_fft\n\n if testset == \"ls_pc_test_clean\":\n metalst = rel_path + \"/data/librispeech_pc_test_clean_cross_sentence.lst\"\n librispeech_test_clean_path = args.librispeech_test_clean_path\n metainfo = get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path)\n\n elif testset == \"seedtts_test_zh\":\n metalst = rel_path + \"/data/seedtts_testset/zh/meta.lst\"\n metainfo = get_seedtts_testset_metainfo(metalst)\n\n elif testset == \"seedtts_test_en\":\n metalst = rel_path + \"/data/seedtts_testset/en/meta.lst\"\n metainfo = get_seedtts_testset_metainfo(metalst)\n\n # path to save genereted wavs\n output_dir = (\n f\"{rel_path}/\"\n f\"results/{exp_name}_{ckpt_step}/{testset}/\"\n f\"seed{seed}_{ode_method}_nfe{nfe_step}_{mel_spec_type}\"\n f\"{f'_ss{sway_sampling_coef}' if sway_sampling_coef else ''}\"\n f\"_cfg{cfg_strength}_speed{speed}\"\n f\"{'_gt-dur' if use_truth_duration else ''}\"\n f\"{'_no-ref-audio' if no_ref_audio else ''}\"\n )\n\n # -------------------------------------------------#\n\n prompts_all = get_inference_prompt(\n metainfo,\n speed=speed,\n tokenizer=tokenizer,\n target_sample_rate=target_sample_rate,\n n_mel_channels=n_mel_channels,\n hop_length=hop_length,\n mel_spec_type=mel_spec_type,\n target_rms=target_rms,\n use_truth_duration=use_truth_duration,\n infer_batch_size=infer_batch_size,\n )\n\n # Vocoder model\n local = args.local\n if mel_spec_type == \"vocos\":\n vocoder_local_path = \"../checkpoints/charactr/vocos-mel-24khz\"\n elif mel_spec_type == \"bigvgan\":\n vocoder_local_path = \"../checkpoints/bigvgan_v2_24khz_100band_256x\"\n vocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)\n\n # Tokenizer\n vocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)\n\n # Model\n model = CFM(\n transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),\n mel_spec_kwargs=dict(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n ),\n odeint_kwargs=dict(\n method=ode_method,\n ),\n vocab_char_map=vocab_char_map,\n ).to(device)\n\n ckpt_prefix = rel_path + f\"/ckpts/{exp_name}/model_{ckpt_step}\"\n if os.path.exists(ckpt_prefix + \".pt\"):\n ckpt_path = ckpt_prefix + \".pt\"\n elif os.path.exists(ckpt_prefix + \".safetensors\"):\n ckpt_path = ckpt_prefix + \".safetensors\"\n else:\n print(\"Loading from self-organized training checkpoints rather than released pretrained.\")\n ckpt_prefix = rel_path + f\"/{model_cfg.ckpts.save_dir}/model_{ckpt_step}\"\n if os.path.exists(ckpt_prefix + \".pt\"):\n ckpt_path = ckpt_prefix + \".pt\"\n elif os.path.exists(ckpt_prefix + \".safetensors\"):\n ckpt_path = ckpt_prefix + \".safetensors\"\n else:\n raise ValueError(\"The checkpoint does not exist or cannot be found in given location.\")\n\n dtype = torch.float32 if mel_spec_type == \"bigvgan\" else None\n model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)\n\n if not os.path.exists(output_dir) and accelerator.is_main_process:\n os.makedirs(output_dir)\n\n # start batch inference\n accelerator.wait_for_everyone()\n start = time.time()\n\n with accelerator.split_between_processes(prompts_all) as prompts:\n for prompt in tqdm(prompts, disable=not accelerator.is_local_main_process):\n utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = prompt\n ref_mels = ref_mels.to(device)\n ref_mel_lens = torch.tensor(ref_mel_lens, dtype=torch.long).to(device)\n total_mel_lens = torch.tensor(total_mel_lens, dtype=torch.long).to(device)\n\n # Inference\n with torch.inference_mode():\n generated, _ = model.sample(\n cond=ref_mels,\n text=final_text_list,\n duration=total_mel_lens,\n lens=ref_mel_lens,\n steps=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n no_ref_audio=no_ref_audio,\n seed=seed,\n )\n # Final result\n for i, gen in enumerate(generated):\n gen = gen[ref_mel_lens[i] : total_mel_lens[i], :].unsqueeze(0)\n gen_mel_spec = gen.permute(0, 2, 1).to(torch.float32)\n if mel_spec_type == \"vocos\":\n generated_wave = vocoder.decode(gen_mel_spec).cpu()\n elif mel_spec_type == \"bigvgan\":\n generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()\n\n if ref_rms_list[i] < target_rms:\n generated_wave = generated_wave * ref_rms_list[i] / target_rms\n torchaudio.save(f\"{output_dir}/{utts[i]}.wav\", generated_wave, target_sample_rate)\n\n accelerator.wait_for_everyone()\n if accelerator.is_main_process:\n timediff = time.time() - start\n print(f\"Done batch inference in {timediff / 60:.2f} minutes.\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/eval/eval_infer_batch.sh", "content": "#!/bin/bash\nset -e\nexport PYTHONWARNINGS=\"ignore::UserWarning,ignore::FutureWarning\"\n\n# Configuration parameters\nMODEL_NAME=\"F5TTS_v1_Base\"\nSEEDS=(0 1 2)\nCKPTSTEPS=(1250000)\nTASKS=(\"seedtts_test_zh\" \"seedtts_test_en\" \"ls_pc_test_clean\")\nLS_TEST_CLEAN_PATH=\"data/LibriSpeech/test-clean\"\nGPUS=\"[0,1,2,3,4,5,6,7]\"\nOFFLINE_MODE=false\n\n# Parse arguments\nif [ $OFFLINE_MODE = true ]; then\n LOCAL=\"--local\"\nelse\n LOCAL=\"\"\nfi\nINFER_ONLY=false\nwhile [[ $# -gt 0 ]]; do\n case $1 in\n --infer-only)\n INFER_ONLY=true\n shift\n ;;\n *)\n echo \"======== Unknown parameter: $1\"\n exit 1\n ;;\n esac\ndone\n\necho \"======== Starting F5-TTS batch evaluation task...\"\nif [ \"$INFER_ONLY\" = true ]; then\n echo \"======== Mode: Execute infer tasks only\"\nelse\n echo \"======== Mode: Execute full pipeline (infer + eval)\"\nfi\n\n# Function: Execute eval tasks\nexecute_eval_tasks() {\n local ckptstep=$1\n local seed=$2\n local task_name=$3\n \n local gen_wav_dir=\"results/${MODEL_NAME}_${ckptstep}/${task_name}/seed${seed}_euler_nfe32_vocos_ss-1_cfg2.0_speed1.0\"\n \n echo \">>>>>>>> Starting eval task: ckptstep=${ckptstep}, seed=${seed}, task=${task_name}\"\n \n case $task_name in\n \"seedtts_test_zh\")\n python src/f5_tts/eval/eval_seedtts_testset.py -e wer -l zh -g \"$gen_wav_dir\" -n \"$GPUS\" $LOCAL\n python src/f5_tts/eval/eval_seedtts_testset.py -e sim -l zh -g \"$gen_wav_dir\" -n \"$GPUS\" $LOCAL\n python src/f5_tts/eval/eval_utmos.py --audio_dir \"$gen_wav_dir\"\n ;;\n \"seedtts_test_en\")\n python src/f5_tts/eval/eval_seedtts_testset.py -e wer -l en -g \"$gen_wav_dir\" -n \"$GPUS\" $LOCAL\n python src/f5_tts/eval/eval_seedtts_testset.py -e sim -l en -g \"$gen_wav_dir\" -n \"$GPUS\" $LOCAL\n python src/f5_tts/eval/eval_utmos.py --audio_dir \"$gen_wav_dir\"\n ;;\n \"ls_pc_test_clean\")\n python src/f5_tts/eval/eval_librispeech_test_clean.py -e wer -g \"$gen_wav_dir\" -n \"$GPUS\" -p \"$LS_TEST_CLEAN_PATH\" $LOCAL\n python src/f5_tts/eval/eval_librispeech_test_clean.py -e sim -g \"$gen_wav_dir\" -n \"$GPUS\" -p \"$LS_TEST_CLEAN_PATH\" $LOCAL\n python src/f5_tts/eval/eval_utmos.py --audio_dir \"$gen_wav_dir\"\n ;;\n esac\n \n echo \">>>>>>>> Completed eval task: ckptstep=${ckptstep}, seed=${seed}, task=${task_name}\"\n}\n\n# Main execution loop\nfor ckptstep in \"${CKPTSTEPS[@]}\"; do\n echo \"======== Processing ckptstep: ${ckptstep}\"\n \n for seed in \"${SEEDS[@]}\"; do\n echo \"-------- Processing seed: ${seed}\"\n \n # Store eval task PIDs for current seed (if not infer-only mode)\n if [ \"$INFER_ONLY\" = false ]; then\n declare -a eval_pids\n fi\n \n # Execute each infer task sequentially\n for task in \"${TASKS[@]}\"; do\n echo \">>>>>>>> Executing infer task: accelerate launch src/f5_tts/eval/eval_infer_batch.py -s ${seed} -n \\\"${MODEL_NAME}\\\" -t \\\"${task}\\\" -c ${ckptstep} $LOCAL\"\n \n # Execute infer task (foreground execution, wait for completion)\n accelerate launch src/f5_tts/eval/eval_infer_batch.py -s ${seed} -n \"${MODEL_NAME}\" -t \"${task}\" -c ${ckptstep} -p \"${LS_TEST_CLEAN_PATH}\" $LOCAL\n \n # If not infer-only mode, launch corresponding eval task\n if [ \"$INFER_ONLY\" = false ]; then\n # Launch corresponding eval task (background execution, non-blocking for next infer)\n execute_eval_tasks $ckptstep $seed $task &\n eval_pids+=($!)\n fi\n done\n \n # If not infer-only mode, wait for all eval tasks of current seed to complete\n if [ \"$INFER_ONLY\" = false ]; then\n echo \">>>>>>>> All infer tasks for seed ${seed} completed, waiting for corresponding eval tasks to finish...\"\n \n for pid in \"${eval_pids[@]}\"; do\n wait $pid\n done\n \n unset eval_pids # Clean up array\n fi\n echo \"-------- All eval tasks for seed ${seed} completed\"\n done\n \n echo \"======== Completed ckptstep: ${ckptstep}\"\n echo\ndone\n\necho \"======== All tasks completed!\"" }, { "path": "src/f5_tts/eval/eval_infer_batch_example.sh", "content": "#!/bin/bash\n\n# e.g. F5-TTS, 16 NFE\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"F5TTS_v1_Base\" -t \"seedtts_test_zh\" -nfe 16\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"F5TTS_v1_Base\" -t \"seedtts_test_en\" -nfe 16\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"F5TTS_v1_Base\" -t \"ls_pc_test_clean\" -nfe 16 -p data/LibriSpeech/test-clean\n\n# e.g. Vanilla E2 TTS, 32 NFE\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"E2TTS_Base\" -c 1200000 -t \"seedtts_test_zh\" -o \"midpoint\" -ss 0\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"E2TTS_Base\" -c 1200000 -t \"seedtts_test_en\" -o \"midpoint\" -ss 0\naccelerate launch src/f5_tts/eval/eval_infer_batch.py -s 0 -n \"E2TTS_Base\" -c 1200000 -t \"ls_pc_test_clean\" -o \"midpoint\" -ss 0 -p data/LibriSpeech/test-clean\n\n# e.g. evaluate F5-TTS 16 NFE result on Seed-TTS test-zh\npython src/f5_tts/eval/eval_seedtts_testset.py -e wer -l zh --gen_wav_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe16_vocos_ss-1_cfg2.0_speed1.0 --gpu_nums 8\npython src/f5_tts/eval/eval_seedtts_testset.py -e sim -l zh --gen_wav_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe16_vocos_ss-1_cfg2.0_speed1.0 --gpu_nums 8\npython src/f5_tts/eval/eval_utmos.py --audio_dir results/F5TTS_v1_Base_1250000/seedtts_test_zh/seed0_euler_nfe16_vocos_ss-1_cfg2.0_speed1.0\n\n# etc.\n" }, { "path": "src/f5_tts/eval/eval_librispeech_test_clean.py", "content": "# Evaluate with Librispeech test-clean, ~3s prompt to generate 4-10s audio (the way of valle/voicebox evaluation)\n\nimport argparse\nimport ast\nimport json\nimport os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport multiprocessing as mp\nfrom importlib.resources import files\n\nimport numpy as np\n\nfrom f5_tts.eval.utils_eval import get_librispeech_test, run_asr_wer, run_sim\n\n\nrel_path = str(files(\"f5_tts\").joinpath(\"../../\"))\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"-e\", \"--eval_task\", type=str, default=\"wer\", choices=[\"sim\", \"wer\"])\n parser.add_argument(\"-l\", \"--lang\", type=str, default=\"en\")\n parser.add_argument(\"-g\", \"--gen_wav_dir\", type=str, required=True)\n parser.add_argument(\"-p\", \"--librispeech_test_clean_path\", type=str, required=True)\n parser.add_argument(\n \"-n\", \"--gpu_nums\", type=str, default=\"8\", help=\"Number of GPUs to use (e.g., 8) or GPU list (e.g., [0,1,2,3])\"\n )\n parser.add_argument(\"--local\", action=\"store_true\", help=\"Use local custom checkpoint directory\")\n return parser.parse_args()\n\n\ndef parse_gpu_nums(gpu_nums_str):\n try:\n if gpu_nums_str.startswith(\"[\") and gpu_nums_str.endswith(\"]\"):\n gpu_list = ast.literal_eval(gpu_nums_str)\n if isinstance(gpu_list, list):\n return gpu_list\n return list(range(int(gpu_nums_str)))\n except (ValueError, SyntaxError):\n raise argparse.ArgumentTypeError(\n f\"Invalid GPU specification: {gpu_nums_str}. Use a number (e.g., 8) or a list (e.g., [0,1,2,3])\"\n )\n\n\ndef main():\n args = get_args()\n eval_task = args.eval_task\n lang = args.lang\n librispeech_test_clean_path = args.librispeech_test_clean_path # test-clean path\n gen_wav_dir = args.gen_wav_dir\n metalst = rel_path + \"/data/librispeech_pc_test_clean_cross_sentence.lst\"\n\n gpus = parse_gpu_nums(args.gpu_nums)\n test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path)\n\n ## In LibriSpeech, some speakers utilized varying voice characteristics for different characters in the book,\n ## leading to a low similarity for the ground truth in some cases.\n # test_set = get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth = True) # eval ground truth\n\n local = args.local\n if local: # use local custom checkpoint dir\n asr_ckpt_dir = \"../checkpoints/Systran/faster-whisper-large-v3\"\n else:\n asr_ckpt_dir = \"\" # auto download to cache dir\n wavlm_ckpt_dir = \"../checkpoints/UniSpeech/wavlm_large_finetune.pth\"\n\n # --------------------------------------------------------------------------\n\n full_results = []\n metrics = []\n\n if eval_task == \"wer\":\n with mp.Pool(processes=len(gpus)) as pool:\n args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]\n results = pool.map(run_asr_wer, args)\n for r in results:\n full_results.extend(r)\n elif eval_task == \"sim\":\n with mp.Pool(processes=len(gpus)) as pool:\n args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]\n results = pool.map(run_sim, args)\n for r in results:\n full_results.extend(r)\n else:\n raise ValueError(f\"Unknown metric type: {eval_task}\")\n\n result_path = f\"{gen_wav_dir}/_{eval_task}_results.jsonl\"\n with open(result_path, \"w\") as f:\n for line in full_results:\n metrics.append(line[eval_task])\n f.write(json.dumps(line, ensure_ascii=False) + \"\\n\")\n metric = round(np.mean(metrics), 5)\n f.write(f\"\\n{eval_task.upper()}: {metric}\\n\")\n\n print(f\"\\nTotal {len(metrics)} samples\")\n print(f\"{eval_task.upper()}: {metric}\")\n print(f\"{eval_task.upper()} results saved to {result_path}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/eval/eval_seedtts_testset.py", "content": "# Evaluate with Seed-TTS testset\n\nimport argparse\nimport ast\nimport json\nimport os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport multiprocessing as mp\nfrom importlib.resources import files\n\nimport numpy as np\n\nfrom f5_tts.eval.utils_eval import get_seed_tts_test, run_asr_wer, run_sim\n\n\nrel_path = str(files(\"f5_tts\").joinpath(\"../../\"))\n\n\ndef get_args():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"-e\", \"--eval_task\", type=str, default=\"wer\", choices=[\"sim\", \"wer\"])\n parser.add_argument(\"-l\", \"--lang\", type=str, default=\"en\", choices=[\"zh\", \"en\"])\n parser.add_argument(\"-g\", \"--gen_wav_dir\", type=str, required=True)\n parser.add_argument(\n \"-n\", \"--gpu_nums\", type=str, default=\"8\", help=\"Number of GPUs to use (e.g., 8) or GPU list (e.g., [0,1,2,3])\"\n )\n parser.add_argument(\"--local\", action=\"store_true\", help=\"Use local custom checkpoint directory\")\n return parser.parse_args()\n\n\ndef parse_gpu_nums(gpu_nums_str):\n try:\n if gpu_nums_str.startswith(\"[\") and gpu_nums_str.endswith(\"]\"):\n gpu_list = ast.literal_eval(gpu_nums_str)\n if isinstance(gpu_list, list):\n return gpu_list\n return list(range(int(gpu_nums_str)))\n except (ValueError, SyntaxError):\n raise argparse.ArgumentTypeError(\n f\"Invalid GPU specification: {gpu_nums_str}. Use a number (e.g., 8) or a list (e.g., [0,1,2,3])\"\n )\n\n\ndef main():\n args = get_args()\n eval_task = args.eval_task\n lang = args.lang\n gen_wav_dir = args.gen_wav_dir\n metalst = rel_path + f\"/data/seedtts_testset/{lang}/meta.lst\" # seed-tts testset\n\n # NOTE. paraformer-zh result will be slightly different according to the number of gpus, cuz batchsize is different\n # zh 1.254 seems a result of 4 workers wer_seed_tts\n gpus = parse_gpu_nums(args.gpu_nums)\n test_set = get_seed_tts_test(metalst, gen_wav_dir, gpus)\n\n local = args.local\n if local: # use local custom checkpoint dir\n if lang == \"zh\":\n asr_ckpt_dir = \"../checkpoints/funasr\" # paraformer-zh dir under funasr\n elif lang == \"en\":\n asr_ckpt_dir = \"../checkpoints/Systran/faster-whisper-large-v3\"\n else:\n asr_ckpt_dir = \"\" # auto download to cache dir\n wavlm_ckpt_dir = \"../checkpoints/UniSpeech/wavlm_large_finetune.pth\"\n\n # --------------------------------------------------------------------------\n\n full_results = []\n metrics = []\n\n if eval_task == \"wer\":\n with mp.Pool(processes=len(gpus)) as pool:\n args = [(rank, lang, sub_test_set, asr_ckpt_dir) for (rank, sub_test_set) in test_set]\n results = pool.map(run_asr_wer, args)\n for r in results:\n full_results.extend(r)\n elif eval_task == \"sim\":\n with mp.Pool(processes=len(gpus)) as pool:\n args = [(rank, sub_test_set, wavlm_ckpt_dir) for (rank, sub_test_set) in test_set]\n results = pool.map(run_sim, args)\n for r in results:\n full_results.extend(r)\n else:\n raise ValueError(f\"Unknown metric type: {eval_task}\")\n\n result_path = f\"{gen_wav_dir}/_{eval_task}_results.jsonl\"\n with open(result_path, \"w\") as f:\n for line in full_results:\n metrics.append(line[eval_task])\n f.write(json.dumps(line, ensure_ascii=False) + \"\\n\")\n metric = round(np.mean(metrics), 5)\n f.write(f\"\\n{eval_task.upper()}: {metric}\\n\")\n\n print(f\"\\nTotal {len(metrics)} samples\")\n print(f\"{eval_task.upper()}: {metric}\")\n print(f\"{eval_task.upper()} results saved to {result_path}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/eval/eval_utmos.py", "content": "import argparse\nimport json\nfrom pathlib import Path\n\nimport librosa\nimport torch\nfrom tqdm import tqdm\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"UTMOS Evaluation\")\n parser.add_argument(\"--audio_dir\", type=str, required=True, help=\"Audio file path.\")\n parser.add_argument(\"--ext\", type=str, default=\"wav\", help=\"Audio extension.\")\n args = parser.parse_args()\n\n device = \"cuda\" if torch.cuda.is_available() else \"xpu\" if torch.xpu.is_available() else \"cpu\"\n\n predictor = torch.hub.load(\"tarepan/SpeechMOS:v1.2.0\", \"utmos22_strong\", trust_repo=True)\n predictor = predictor.to(device)\n\n audio_paths = list(Path(args.audio_dir).rglob(f\"*.{args.ext}\"))\n utmos_score = 0\n\n utmos_result_path = Path(args.audio_dir) / \"_utmos_results.jsonl\"\n with open(utmos_result_path, \"w\", encoding=\"utf-8\") as f:\n for audio_path in tqdm(audio_paths, desc=\"Processing\"):\n wav, sr = librosa.load(audio_path, sr=None, mono=True)\n wav_tensor = torch.from_numpy(wav).to(device).unsqueeze(0)\n score = predictor(wav_tensor, sr)\n line = {}\n line[\"wav\"], line[\"utmos\"] = str(audio_path.stem), score.item()\n utmos_score += score.item()\n f.write(json.dumps(line, ensure_ascii=False) + \"\\n\")\n avg_score = utmos_score / len(audio_paths) if len(audio_paths) > 0 else 0\n f.write(f\"\\nUTMOS: {avg_score:.4f}\\n\")\n\n print(f\"UTMOS: {avg_score:.4f}\")\n print(f\"UTMOS results saved to {utmos_result_path}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/eval/utils_eval.py", "content": "import math\nimport os\nimport random\nimport string\nfrom pathlib import Path\n\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nfrom tqdm import tqdm\n\nfrom f5_tts.eval.ecapa_tdnn import ECAPA_TDNN_SMALL\nfrom f5_tts.model.modules import MelSpec\nfrom f5_tts.model.utils import convert_char_to_pinyin\n\n\n# seedtts testset metainfo: utt, prompt_text, prompt_wav, gt_text, gt_wav\ndef get_seedtts_testset_metainfo(metalst):\n f = open(metalst)\n lines = f.readlines()\n f.close()\n metainfo = []\n for line in lines:\n if len(line.strip().split(\"|\")) == 5:\n utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split(\"|\")\n elif len(line.strip().split(\"|\")) == 4:\n utt, prompt_text, prompt_wav, gt_text = line.strip().split(\"|\")\n gt_wav = os.path.join(os.path.dirname(metalst), \"wavs\", utt + \".wav\")\n if not os.path.isabs(prompt_wav):\n prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)\n metainfo.append((utt, prompt_text, prompt_wav, gt_text, gt_wav))\n return metainfo\n\n\n# librispeech test-clean metainfo: gen_utt, ref_txt, ref_wav, gen_txt, gen_wav\ndef get_librispeech_test_clean_metainfo(metalst, librispeech_test_clean_path):\n f = open(metalst)\n lines = f.readlines()\n f.close()\n metainfo = []\n for line in lines:\n ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split(\"\\t\")\n\n # ref_txt = ref_txt[0] + ref_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc)\n ref_spk_id, ref_chaptr_id, _ = ref_utt.split(\"-\")\n ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + \".flac\")\n\n # gen_txt = gen_txt[0] + gen_txt[1:].lower() + '.' # if use librispeech test-clean (no-pc)\n gen_spk_id, gen_chaptr_id, _ = gen_utt.split(\"-\")\n gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + \".flac\")\n\n metainfo.append((gen_utt, ref_txt, ref_wav, \" \" + gen_txt, gen_wav))\n\n return metainfo\n\n\n# padded to max length mel batch\ndef padded_mel_batch(ref_mels):\n max_mel_length = torch.LongTensor([mel.shape[-1] for mel in ref_mels]).amax()\n padded_ref_mels = []\n for mel in ref_mels:\n padded_ref_mel = F.pad(mel, (0, max_mel_length - mel.shape[-1]), value=0)\n padded_ref_mels.append(padded_ref_mel)\n padded_ref_mels = torch.stack(padded_ref_mels)\n padded_ref_mels = padded_ref_mels.permute(0, 2, 1)\n return padded_ref_mels\n\n\n# get prompts from metainfo containing: utt, prompt_text, prompt_wav, gt_text, gt_wav\n\n\ndef get_inference_prompt(\n metainfo,\n speed=1.0,\n tokenizer=\"pinyin\",\n polyphone=True,\n target_sample_rate=24000,\n n_fft=1024,\n win_length=1024,\n n_mel_channels=100,\n hop_length=256,\n mel_spec_type=\"vocos\",\n target_rms=0.1,\n use_truth_duration=False,\n infer_batch_size=1,\n num_buckets=200,\n min_secs=3,\n max_secs=40,\n):\n prompts_all = []\n\n min_tokens = min_secs * target_sample_rate // hop_length\n max_tokens = max_secs * target_sample_rate // hop_length\n\n batch_accum = [0] * num_buckets\n utts, ref_rms_list, ref_mels, ref_mel_lens, total_mel_lens, final_text_list = (\n [[] for _ in range(num_buckets)] for _ in range(6)\n )\n\n mel_spectrogram = MelSpec(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n )\n\n for utt, prompt_text, prompt_wav, gt_text, gt_wav in tqdm(metainfo, desc=\"Processing prompts...\"):\n # Audio\n ref_audio, ref_sr = torchaudio.load(prompt_wav)\n ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio)))\n if ref_rms < target_rms:\n ref_audio = ref_audio * target_rms / ref_rms\n assert ref_audio.shape[-1] > 5000, f\"Empty prompt wav: {prompt_wav}, or torchaudio backend issue.\"\n if ref_sr != target_sample_rate:\n resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)\n ref_audio = resampler(ref_audio)\n\n # Text\n if len(prompt_text[-1].encode(\"utf-8\")) == 1:\n prompt_text = prompt_text + \" \"\n text = [prompt_text + gt_text]\n if tokenizer == \"pinyin\":\n text_list = convert_char_to_pinyin(text, polyphone=polyphone)\n else:\n text_list = text\n\n # to mel spectrogram\n ref_mel = mel_spectrogram(ref_audio)\n ref_mel = ref_mel.squeeze(0)\n\n # Duration, mel frame length\n ref_mel_len = ref_mel.shape[-1]\n\n if use_truth_duration:\n gt_audio, gt_sr = torchaudio.load(gt_wav)\n if gt_sr != target_sample_rate:\n resampler = torchaudio.transforms.Resample(gt_sr, target_sample_rate)\n gt_audio = resampler(gt_audio)\n total_mel_len = ref_mel_len + int(gt_audio.shape[-1] / hop_length / speed)\n\n # # test vocoder resynthesis\n # ref_audio = gt_audio\n else:\n ref_text_len = len(prompt_text.encode(\"utf-8\"))\n gen_text_len = len(gt_text.encode(\"utf-8\"))\n total_mel_len = ref_mel_len + int(ref_mel_len / ref_text_len * gen_text_len / speed)\n\n # deal with batch\n assert infer_batch_size > 0, \"infer_batch_size should be greater than 0.\"\n assert min_tokens <= total_mel_len <= max_tokens, (\n f\"Audio {utt} has duration {total_mel_len * hop_length // target_sample_rate}s out of range [{min_secs}, {max_secs}].\"\n )\n bucket_i = math.floor((total_mel_len - min_tokens) / (max_tokens - min_tokens + 1) * num_buckets)\n\n utts[bucket_i].append(utt)\n ref_rms_list[bucket_i].append(ref_rms)\n ref_mels[bucket_i].append(ref_mel)\n ref_mel_lens[bucket_i].append(ref_mel_len)\n total_mel_lens[bucket_i].append(total_mel_len)\n final_text_list[bucket_i].extend(text_list)\n\n batch_accum[bucket_i] += total_mel_len\n\n if batch_accum[bucket_i] >= infer_batch_size:\n # print(f\"\\n{len(ref_mels[bucket_i][0][0])}\\n{ref_mel_lens[bucket_i]}\\n{total_mel_lens[bucket_i]}\")\n prompts_all.append(\n (\n utts[bucket_i],\n ref_rms_list[bucket_i],\n padded_mel_batch(ref_mels[bucket_i]),\n ref_mel_lens[bucket_i],\n total_mel_lens[bucket_i],\n final_text_list[bucket_i],\n )\n )\n batch_accum[bucket_i] = 0\n (\n utts[bucket_i],\n ref_rms_list[bucket_i],\n ref_mels[bucket_i],\n ref_mel_lens[bucket_i],\n total_mel_lens[bucket_i],\n final_text_list[bucket_i],\n ) = [], [], [], [], [], []\n\n # add residual\n for bucket_i, bucket_frames in enumerate(batch_accum):\n if bucket_frames > 0:\n prompts_all.append(\n (\n utts[bucket_i],\n ref_rms_list[bucket_i],\n padded_mel_batch(ref_mels[bucket_i]),\n ref_mel_lens[bucket_i],\n total_mel_lens[bucket_i],\n final_text_list[bucket_i],\n )\n )\n # not only leave easy work for last workers\n random.seed(666)\n random.shuffle(prompts_all)\n\n return prompts_all\n\n\n# get wav_res_ref_text of seed-tts test metalst\n# https://github.com/BytedanceSpeech/seed-tts-eval\n\n\ndef get_seed_tts_test(metalst, gen_wav_dir, gpus):\n f = open(metalst)\n lines = f.readlines()\n f.close()\n\n test_set_ = []\n for line in tqdm(lines):\n if len(line.strip().split(\"|\")) == 5:\n utt, prompt_text, prompt_wav, gt_text, gt_wav = line.strip().split(\"|\")\n elif len(line.strip().split(\"|\")) == 4:\n utt, prompt_text, prompt_wav, gt_text = line.strip().split(\"|\")\n\n if not os.path.exists(os.path.join(gen_wav_dir, utt + \".wav\")):\n continue\n gen_wav = os.path.join(gen_wav_dir, utt + \".wav\")\n if not os.path.isabs(prompt_wav):\n prompt_wav = os.path.join(os.path.dirname(metalst), prompt_wav)\n\n test_set_.append((gen_wav, prompt_wav, gt_text))\n\n num_jobs = len(gpus)\n if num_jobs == 1:\n return [(gpus[0], test_set_)]\n\n wav_per_job = len(test_set_) // num_jobs + 1\n test_set = []\n for i in range(num_jobs):\n test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))\n\n return test_set\n\n\n# get librispeech test-clean cross sentence test\n\n\ndef get_librispeech_test(metalst, gen_wav_dir, gpus, librispeech_test_clean_path, eval_ground_truth=False):\n f = open(metalst)\n lines = f.readlines()\n f.close()\n\n test_set_ = []\n for line in tqdm(lines):\n ref_utt, ref_dur, ref_txt, gen_utt, gen_dur, gen_txt = line.strip().split(\"\\t\")\n\n if eval_ground_truth:\n gen_spk_id, gen_chaptr_id, _ = gen_utt.split(\"-\")\n gen_wav = os.path.join(librispeech_test_clean_path, gen_spk_id, gen_chaptr_id, gen_utt + \".flac\")\n else:\n if not os.path.exists(os.path.join(gen_wav_dir, gen_utt + \".wav\")):\n raise FileNotFoundError(f\"Generated wav not found: {gen_utt}\")\n gen_wav = os.path.join(gen_wav_dir, gen_utt + \".wav\")\n\n ref_spk_id, ref_chaptr_id, _ = ref_utt.split(\"-\")\n ref_wav = os.path.join(librispeech_test_clean_path, ref_spk_id, ref_chaptr_id, ref_utt + \".flac\")\n\n test_set_.append((gen_wav, ref_wav, gen_txt))\n\n num_jobs = len(gpus)\n if num_jobs == 1:\n return [(gpus[0], test_set_)]\n\n wav_per_job = len(test_set_) // num_jobs + 1\n test_set = []\n for i in range(num_jobs):\n test_set.append((gpus[i], test_set_[i * wav_per_job : (i + 1) * wav_per_job]))\n\n return test_set\n\n\n# load asr model\n\n\ndef load_asr_model(lang, ckpt_dir=\"\"):\n if lang == \"zh\":\n from funasr import AutoModel\n\n model = AutoModel(\n model=os.path.join(ckpt_dir, \"paraformer-zh\"),\n # vad_model = os.path.join(ckpt_dir, \"fsmn-vad\"),\n # punc_model = os.path.join(ckpt_dir, \"ct-punc\"),\n # spk_model = os.path.join(ckpt_dir, \"cam++\"),\n disable_update=True,\n ) # following seed-tts setting\n elif lang == \"en\":\n from faster_whisper import WhisperModel\n\n model_size = \"large-v3\" if ckpt_dir == \"\" else ckpt_dir\n model = WhisperModel(model_size, device=\"cuda\", compute_type=\"float16\")\n return model\n\n\n# WER Evaluation, the way Seed-TTS does\n\n\ndef run_asr_wer(args):\n rank, lang, test_set, ckpt_dir = args\n\n if lang == \"zh\":\n import zhconv\n\n torch.cuda.set_device(rank)\n elif lang == \"en\":\n os.environ[\"CUDA_VISIBLE_DEVICES\"] = str(rank)\n else:\n raise NotImplementedError(\n \"lang support only 'zh' (funasr paraformer-zh), 'en' (faster-whisper-large-v3), for now.\"\n )\n\n asr_model = load_asr_model(lang, ckpt_dir=ckpt_dir)\n\n from zhon.hanzi import punctuation\n\n punctuation_all = punctuation + string.punctuation\n wer_results = []\n\n from jiwer import process_words\n\n for gen_wav, prompt_wav, truth in tqdm(test_set):\n if lang == \"zh\":\n res = asr_model.generate(input=gen_wav, batch_size_s=300, disable_pbar=True)\n hypo = res[0][\"text\"]\n hypo = zhconv.convert(hypo, \"zh-cn\")\n elif lang == \"en\":\n segments, _ = asr_model.transcribe(gen_wav, beam_size=5, language=\"en\")\n hypo = \"\"\n for segment in segments:\n hypo = hypo + \" \" + segment.text\n\n raw_truth = truth\n raw_hypo = hypo\n\n for x in punctuation_all:\n truth = truth.replace(x, \"\")\n hypo = hypo.replace(x, \"\")\n\n truth = truth.replace(\" \", \" \")\n hypo = hypo.replace(\" \", \" \")\n\n if lang == \"zh\":\n truth = \" \".join([x for x in truth])\n hypo = \" \".join([x for x in hypo])\n elif lang == \"en\":\n truth = truth.lower()\n hypo = hypo.lower()\n\n measures = process_words(truth, hypo)\n wer = measures.wer\n\n # ref_list = truth.split(\" \")\n # subs = measures.substitutions / len(ref_list)\n # dele = measures.deletions / len(ref_list)\n # inse = measures.insertions / len(ref_list)\n\n wer_results.append(\n {\n \"wav\": Path(gen_wav).stem,\n \"truth\": raw_truth,\n \"hypo\": raw_hypo,\n \"wer\": wer,\n }\n )\n\n return wer_results\n\n\n# SIM Evaluation\n\n\ndef run_sim(args):\n rank, test_set, ckpt_dir = args\n device = f\"cuda:{rank}\"\n\n model = ECAPA_TDNN_SMALL(feat_dim=1024, feat_type=\"wavlm_large\", config_path=None)\n state_dict = torch.load(ckpt_dir, weights_only=True, map_location=lambda storage, loc: storage)\n model.load_state_dict(state_dict[\"model\"], strict=False)\n\n use_gpu = True if torch.cuda.is_available() else False\n if use_gpu:\n model = model.cuda(device)\n model.eval()\n\n sim_results = []\n for gen_wav, prompt_wav, truth in tqdm(test_set):\n wav1, sr1 = torchaudio.load(gen_wav)\n wav2, sr2 = torchaudio.load(prompt_wav)\n\n if use_gpu:\n wav1 = wav1.cuda(device)\n wav2 = wav2.cuda(device)\n\n if sr1 != 16000:\n resample1 = torchaudio.transforms.Resample(orig_freq=sr1, new_freq=16000)\n if use_gpu:\n resample1 = resample1.cuda(device)\n wav1 = resample1(wav1)\n if sr2 != 16000:\n resample2 = torchaudio.transforms.Resample(orig_freq=sr2, new_freq=16000)\n if use_gpu:\n resample2 = resample2.cuda(device)\n wav2 = resample2(wav2)\n\n with torch.no_grad():\n emb1 = model(wav1)\n emb2 = model(wav2)\n\n sim = F.cosine_similarity(emb1, emb2)[0].item()\n # print(f\"VSim score between two audios: {sim:.4f} (-1.0, 1.0).\")\n sim_results.append(\n {\n \"wav\": Path(gen_wav).stem,\n \"sim\": sim,\n }\n )\n\n return sim_results\n" }, { "path": "src/f5_tts/infer/README.md", "content": "# Inference\n\nThe 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.\n\n**More checkpoints with whole community efforts can be found in [SHARED.md](SHARED.md), supporting more languages.**\n\nCurrently 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**.\n\nTo avoid possible inference failures, make sure you have seen through the following instructions.\n\n- 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.\n- Uppercased letters (best with form like K.F.C.) will be uttered letter by letter, and lowercased letters used for common words. \n- Add some spaces (blank: \" \") or punctuations (e.g. \",\" \".\") to explicitly introduce some pauses.\n- If English punctuation marks the end of a sentence, make sure there is a space \" \" after it. Otherwise not regarded as when chunk.\n- Preprocess numbers to Chinese letters if you want to have them read in Chinese, otherwise in English.\n- If the generation output is blank (pure silence), check for FFmpeg installation.\n- Try turn off `use_ema` if using an early-stage finetuned checkpoint (which goes just few updates).\n\n\n## Gradio App\n\nCurrently supported features:\n\n- Basic TTS with Chunk Inference\n- Multi-Style / Multi-Speaker Generation\n- Voice Chat powered by Qwen2.5-3B-Instruct\n- [Custom inference with more language support](SHARED.md)\n\nThe 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.\n\nThe 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.\n\nMore flags options:\n\n```bash\n# Automatically launch the interface in the default web browser\nf5-tts_infer-gradio --inbrowser\n\n# Set the root path of the application, if it's not served from the root (\"/\") of the domain\n# For example, if the application is served at \"https://example.com/myapp\"\nf5-tts_infer-gradio --root_path \"/myapp\"\n```\n\nCould also be used as a component for larger application:\n```python\nimport gradio as gr\nfrom f5_tts.infer.infer_gradio import app\n\nwith gr.Blocks() as main_app:\n gr.Markdown(\"# This is an example of using F5-TTS within a bigger Gradio app\")\n\n # ... other Gradio components\n\n app.render()\n\nmain_app.launch()\n```\n\n\n## CLI Inference\n\nThe cli command `f5-tts_infer-cli` equals to `python src/f5_tts/infer/infer_cli.py`, which is a command line tool for inference.\n\nThe 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`.\n\nFor change vocab.txt use `--vocab_file` to provide your `vocab.txt` file.\n\nBasically you can inference with flags:\n```bash\n# Leave --ref_text \"\" will have ASR model transcribe (extra GPU memory usage)\nf5-tts_infer-cli \\\n--model F5TTS_v1_Base \\\n--ref_audio \"ref_audio.wav\" \\\n--ref_text \"The content, subtitle or transcription of reference audio.\" \\\n--gen_text \"Some text you want TTS model generate for you.\"\n\n# Use BigVGAN as vocoder. Currently only support F5TTS_Base. \nf5-tts_infer-cli --model F5TTS_Base --vocoder_name bigvgan --load_vocoder_from_local\n\n# Use custom path checkpoint, e.g.\nf5-tts_infer-cli --ckpt_file ckpts/F5TTS_v1_Base/model_1250000.safetensors\n\n# More instructions\nf5-tts_infer-cli --help\n```\n\nAnd a `.toml` file would help with more flexible usage.\n\n```bash\nf5-tts_infer-cli -c custom.toml\n```\n\nFor example, you can use `.toml` to pass in variables, refer to `src/f5_tts/infer/examples/basic/basic.toml`:\n\n```toml\n# F5TTS_v1_Base | E2TTS_Base\nmodel = \"F5TTS_v1_Base\"\nref_audio = \"infer/examples/basic/basic_ref_en.wav\"\n# If an empty \"\", transcribes the reference audio automatically.\nref_text = \"Some call me nature, others call me mother nature.\"\ngen_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.\"\n# File with text to generate. Ignores the text above.\ngen_file = \"\"\nremove_silence = false\noutput_dir = \"tests\"\n```\n\nYou can also leverage `.toml` file to do multi-style generation, refer to `src/f5_tts/infer/examples/multi/story.toml`.\n\n```toml\n# F5TTS_v1_Base | E2TTS_Base\nmodel = \"F5TTS_v1_Base\"\nref_audio = \"infer/examples/multi/main.flac\"\n# If an empty \"\", transcribes the reference audio automatically.\nref_text = \"\"\ngen_text = \"\"\n# File with text to generate. Ignores the text above.\ngen_file = \"infer/examples/multi/story.txt\"\nremove_silence = true\noutput_dir = \"tests\"\n\n[voices.town]\nref_audio = \"infer/examples/multi/town.flac\"\nref_text = \"\"\n\n[voices.country]\nref_audio = \"infer/examples/multi/country.flac\"\nref_text = \"\"\n```\nYou should mark the voice with `[main]` `[town]` `[country]` whenever you want to change voice, refer to `src/f5_tts/infer/examples/multi/story.txt`.\n\n## API Usage\n\n```python\nfrom importlib.resources import files\nfrom f5_tts.api import F5TTS\n\nf5tts = F5TTS()\nwav, sr, spec = f5tts.infer(\n ref_file=str(files(\"f5_tts\").joinpath(\"infer/examples/basic/basic_ref_en.wav\")),\n ref_text=\"some call me nature, others call me mother nature.\",\n 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.\"\"\",\n file_wave=str(files(\"f5_tts\").joinpath(\"../../tests/api_out.wav\")),\n file_spec=str(files(\"f5_tts\").joinpath(\"../../tests/api_out.png\")),\n seed=None,\n)\n```\nCheck [api.py](../api.py) for more details.\n\n## TensorRT-LLM Deployment\n\nSee [detailed instructions](../runtime/triton_trtllm/README.md) for more information.\n\n## Socket Real-time Service\n\nReal-time voice output with chunk stream:\n\n```bash\n# Start socket server\npython src/f5_tts/socket_server.py\n\n# If PyAudio not installed\nsudo apt-get install portaudio19-dev\npip install pyaudio\n\n# Communicate with socket client\npython src/f5_tts/socket_client.py\n```\n\n## Speech Editing\n\nTo test speech editing capabilities, use the following command:\n\n```bash\npython src/f5_tts/infer/speech_edit.py\n```\n\n" }, { "path": "src/f5_tts/infer/SHARED.md", "content": "\n# Shared Model Cards\n\n\n### **Prerequisites of using**\n- This document is serving as a quick lookup table for the community training/finetuning result, with various language support.\n- The models in this repository are open source and are based on voluntary contributions from contributors.\n- The use of models must be conditioned on respect for the respective creators. The convenience brought comes from their efforts.\n\n\n### **Welcome to share here**\n- 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).\n- Host a public [huggingface model repository](https://huggingface.co/new) and upload the model related files.\n- Make a pull request adding a model card to the current page, i.e. `src\\f5_tts\\infer\\SHARED.md`.\n\n\n### Supported Languages\n- [Multilingual](#multilingual)\n - [F5-TTS v1 v0 Base @ zh \\& en @ F5-TTS](#f5-tts-v1-v0-base--zh--en--f5-tts)\n- [Arabic](#arabic)\n - [F5-TTS Small @ ar & en @ SILMA AI](#f5-tts-small--ar--en--silma-ai)\n- [English](#english)\n- [Finnish](#finnish)\n - [F5-TTS Base @ fi @ AsmoKoskinen](#f5-tts-base--fi--asmokoskinen)\n- [French](#french)\n - [F5-TTS Base @ fr @ RASPIAUDIO](#f5-tts-base--fr--raspiaudio)\n- [German](#german)\n - [F5-TTS Base @ de @ hvoss-techfak](#f5-tts-base--de--hvoss-techfak)\n- [Hindi](#hindi)\n - [F5-TTS Small @ hi @ SPRINGLab](#f5-tts-small--hi--springlab)\n- [Italian](#italian)\n - [F5-TTS Base @ it @ alien79](#f5-tts-base--it--alien79)\n- [Japanese](#japanese)\n - [F5-TTS Base @ ja @ Jmica](#f5-tts-base--ja--jmica)\n- [Latvian](#latvian)\n - [F5-TTS Base @ lv @ RaivisDejus](#f5-tts-base--lv--raivisdejus)\n- [Mandarin](#mandarin)\n- [Russian](#russian)\n - [F5-TTS Base @ ru @ HotDro4illa](#f5-tts-base--ru--hotdro4illa)\n- [Spanish](#spanish)\n - [F5-TTS Base @ es @ jpgallegoar](#f5-tts-base--es--jpgallegoar)\n\n\n## Multilingual\n\n#### F5-TTS v1 v0 Base @ zh & en @ F5-TTS\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors\n# A Variant Model: hf://SWivid/F5-TTS/F5TTS_v1_Base_no_zero_init/model_1250000.safetensors\nVocab: hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"conv_layers\": 4}\n```\n\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.safetensors\nVocab: hf://SWivid/F5-TTS/F5TTS_Base/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n*Other infos, e.g. Author info, Github repo, Link to some sampled results, Usage instruction, Tutorial (Blog, Video, etc.) ...*\n\n\n## Arabic\n\n#### F5-TTS Small @ ar & en @ SILMA AI\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|F5-TTS Small|[ckpt & vocab](https://huggingface.co/silma-ai/silma-tts)| Tens of thousands EN/AR |Apache-2.0|\n\n- Pretrained by [SILMA.AI](https://silma.ai)\n- [GitHub repo](https://github.com/SILMA-AI/silma-tts), Inference code\n\n\n## English\n\n\n## Finnish\n\n#### F5-TTS Base @ fi @ AsmoKoskinen\n|Model|🤗Hugging Face|Data|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://AsmoKoskinen/F5-TTS_Finnish_Model/model_common_voice_fi_vox_populi_fi_20241206.safetensors\nVocab: hf://AsmoKoskinen/F5-TTS_Finnish_Model/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n\n## French\n\n#### F5-TTS Base @ fr @ RASPIAUDIO\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|F5-TTS Base|[ckpt & vocab](https://huggingface.co/RASPIAUDIO/F5-French-MixedSpeakers-reduced)|[LibriVox](https://librivox.org/)|cc-by-nc-4.0|\n\n```bash\nModel: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/model_last_reduced.pt\nVocab: hf://RASPIAUDIO/F5-French-MixedSpeakers-reduced/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n- [Online Inference with Hugging Face Space](https://huggingface.co/spaces/RASPIAUDIO/f5-tts_french).\n- [Tutorial video to train a new language model](https://www.youtube.com/watch?v=UO4usaOojys).\n- [Discussion about this training can be found here](https://github.com/SWivid/F5-TTS/issues/434).\n\n\n## German\n\n#### F5-TTS Base @ de @ hvoss-techfak\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://hvoss-techfak/F5-TTS-German/model_f5tts_german.pt\nVocab: hf://hvoss-techfak/F5-TTS-German/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n- Finetuned by [@hvoss-techfak](https://github.com/hvoss-techfak)\n\n\n## Hindi\n\n#### F5-TTS Small @ hi @ SPRINGLab\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://SPRINGLab/F5-Hindi-24KHz/model_2500000.safetensors\nVocab: hf://SPRINGLab/F5-Hindi-24KHz/vocab.txt\nConfig: {\"dim\": 768, \"depth\": 18, \"heads\": 12, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n- Authors: SPRING Lab, Indian Institute of Technology, Madras\n- Website: https://asr.iitm.ac.in/\n\n\n## Italian\n\n#### F5-TTS Base @ it @ alien79\n|Model|🤗Hugging Face|Data|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://alien79/F5-TTS-italian/model_159600.safetensors\nVocab: hf://alien79/F5-TTS-italian/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n- Trained by [Mithril Man](https://github.com/MithrilMan)\n- Model details on [hf project home](https://huggingface.co/alien79/F5-TTS-italian)\n- Open to collaborations to further improve the model\n\n\n## Japanese\n\n#### F5-TTS Base @ ja @ Jmica\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://Jmica/F5TTS/JA_21999120/model_21999120.pt\nVocab: hf://Jmica/F5TTS/JA_21999120/vocab_japanese.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n\n## Latvian\n\n#### F5-TTS Base @ lv @ RaivisDejus\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|F5-TTS Base|[ckpt & vocab](https://huggingface.co/RaivisDejus/F5-TTS-Latvian)|[Common voice](https://datacollective.mozillafoundation.org/datasets/cmj8u3pec00flnxxbntvfb4as)|cc0-1.0|\n\n```bash\nModel: hf://RaivisDejus/F5-TTS-Latvian/model.safetensors\nVocab: hf://RaivisDejus/F5-TTS-Latvian/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n\n\n## Mandarin\n\n\n## Russian\n\n#### F5-TTS Base @ ru @ HotDro4illa\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n```bash\nModel: hf://hotstone228/F5-TTS-Russian/model_last.safetensors\nVocab: hf://hotstone228/F5-TTS-Russian/vocab.txt\nConfig: {\"dim\": 1024, \"depth\": 22, \"heads\": 16, \"ff_mult\": 2, \"text_dim\": 512, \"text_mask_padding\": False, \"conv_layers\": 4, \"pe_attn_head\": 1}\n```\n- Finetuned by [HotDro4illa](https://github.com/HotDro4illa)\n- Any improvements are welcome\n\n\n## Spanish\n\n#### F5-TTS Base @ es @ jpgallegoar\n|Model|🤗Hugging Face|Data (Hours)|Model License|\n|:---:|:------------:|:-----------:|:-------------:|\n|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|\n\n- @jpgallegoar [GitHub repo](https://github.com/jpgallegoar/Spanish-F5), Jupyter Notebook and Gradio usage for Spanish model.\n" }, { "path": "src/f5_tts/infer/examples/basic/basic.toml", "content": "# F5TTS_v1_Base | E2TTS_Base\nmodel = \"F5TTS_v1_Base\"\nref_audio = \"infer/examples/basic/basic_ref_en.wav\"\n# If an empty \"\", transcribes the reference audio automatically.\nref_text = \"Some call me nature, others call me mother nature.\"\ngen_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.\"\n# File with text to generate. Ignores the text above.\ngen_file = \"\"\nremove_silence = false\noutput_dir = \"tests\"\noutput_file = \"infer_cli_basic.wav\"\n" }, { "path": "src/f5_tts/infer/examples/multi/story.toml", "content": "# F5TTS_v1_Base | E2TTS_Base\nmodel = \"F5TTS_v1_Base\"\nref_audio = \"infer/examples/multi/main.flac\"\n# If an empty \"\", transcribes the reference audio automatically.\nref_text = \"\"\ngen_text = \"\"\n# File with text to generate. Ignores the text above.\ngen_file = \"infer/examples/multi/story.txt\"\nremove_silence = true\noutput_dir = \"tests\"\noutput_file = \"infer_cli_story.wav\"\n\n[voices.town]\nref_audio = \"infer/examples/multi/town.flac\"\nref_text = \"\"\nspeed = 0.8 # will ignore global speed\n\n[voices.country]\nref_audio = \"infer/examples/multi/country.flac\"\nref_text = \"\"\n" }, { "path": "src/f5_tts/infer/examples/multi/story.txt", "content": "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.\"" }, { "path": "src/f5_tts/infer/examples/vocab.txt", "content": " \n!\n\"\n#\n$\n%\n&\n'\n(\n)\n*\n+\n,\n-\n.\n/\n0\n1\n2\n3\n4\n5\n6\n7\n8\n9\n:\n;\n=\n>\n?\n@\nA\nB\nC\nD\nE\nF\nG\nH\nI\nJ\nK\nL\nM\nN\nO\nP\nQ\nR\nS\nT\nU\nV\nW\nX\nY\nZ\n[\n\\\n]\n_\na\na1\nai1\nai2\nai3\nai4\nan1\nan3\nan4\nang1\nang2\nang4\nao1\nao2\nao3\nao4\nb\nba\nba1\nba2\nba3\nba4\nbai1\nbai2\nbai3\nbai4\nban1\nban2\nban3\nban4\nbang1\nbang2\nbang3\nbang4\nbao1\nbao2\nbao3\nbao4\nbei\nbei1\nbei2\nbei3\nbei4\nben1\nben2\nben3\nben4\nbeng\nbeng1\nbeng2\nbeng3\nbeng4\nbi1\nbi2\nbi3\nbi4\nbian1\nbian2\nbian3\nbian4\nbiao1\nbiao2\nbiao3\nbie1\nbie2\nbie3\nbie4\nbin1\nbin4\nbing1\nbing2\nbing3\nbing4\nbo\nbo1\nbo2\nbo3\nbo4\nbu2\nbu3\nbu4\nc\nca1\ncai1\ncai2\ncai3\ncai4\ncan1\ncan2\ncan3\ncan4\ncang1\ncang2\ncao1\ncao2\ncao3\nce4\ncen1\ncen2\nceng1\nceng2\nceng4\ncha1\ncha2\ncha3\ncha4\nchai1\nchai2\nchan1\nchan2\nchan3\nchan4\nchang1\nchang2\nchang3\nchang4\nchao1\nchao2\nchao3\nche1\nche2\nche3\nche4\nchen1\nchen2\nchen3\nchen4\ncheng1\ncheng2\ncheng3\ncheng4\nchi1\nchi2\nchi3\nchi4\nchong1\nchong2\nchong3\nchong4\nchou1\nchou2\nchou3\nchou4\nchu1\nchu2\nchu3\nchu4\nchua1\nchuai1\nchuai2\nchuai3\nchuai4\nchuan1\nchuan2\nchuan3\nchuan4\nchuang1\nchuang2\nchuang3\nchuang4\nchui1\nchui2\nchun1\nchun2\nchun3\nchuo1\nchuo4\nci1\nci2\nci3\nci4\ncong1\ncong2\ncou4\ncu1\ncu4\ncuan1\ncuan2\ncuan4\ncui1\ncui3\ncui4\ncun1\ncun2\ncun4\ncuo1\ncuo2\ncuo4\nd\nda\nda1\nda2\nda3\nda4\ndai1\ndai2\ndai3\ndai4\ndan1\ndan2\ndan3\ndan4\ndang1\ndang2\ndang3\ndang4\ndao1\ndao2\ndao3\ndao4\nde\nde1\nde2\ndei3\nden4\ndeng1\ndeng2\ndeng3\ndeng4\ndi1\ndi2\ndi3\ndi4\ndia3\ndian1\ndian2\ndian3\ndian4\ndiao1\ndiao3\ndiao4\ndie1\ndie2\ndie4\nding1\nding2\nding3\nding4\ndiu1\ndong1\ndong3\ndong4\ndou1\ndou2\ndou3\ndou4\ndu1\ndu2\ndu3\ndu4\nduan1\nduan2\nduan3\nduan4\ndui1\ndui4\ndun1\ndun3\ndun4\nduo1\nduo2\nduo3\nduo4\ne\ne1\ne2\ne3\ne4\nei2\nen1\nen4\ner\ner2\ner3\ner4\nf\nfa1\nfa2\nfa3\nfa4\nfan1\nfan2\nfan3\nfan4\nfang1\nfang2\nfang3\nfang4\nfei1\nfei2\nfei3\nfei4\nfen1\nfen2\nfen3\nfen4\nfeng1\nfeng2\nfeng3\nfeng4\nfo2\nfou2\nfou3\nfu1\nfu2\nfu3\nfu4\ng\nga1\nga2\nga3\nga4\ngai1\ngai2\ngai3\ngai4\ngan1\ngan2\ngan3\ngan4\ngang1\ngang2\ngang3\ngang4\ngao1\ngao2\ngao3\ngao4\nge1\nge2\nge3\nge4\ngei2\ngei3\ngen1\ngen2\ngen3\ngen4\ngeng1\ngeng3\ngeng4\ngong1\ngong3\ngong4\ngou1\ngou2\ngou3\ngou4\ngu\ngu1\ngu2\ngu3\ngu4\ngua1\ngua2\ngua3\ngua4\nguai1\nguai2\nguai3\nguai4\nguan1\nguan2\nguan3\nguan4\nguang1\nguang2\nguang3\nguang4\ngui1\ngui2\ngui3\ngui4\ngun3\ngun4\nguo1\nguo2\nguo3\nguo4\nh\nha1\nha2\nha3\nhai1\nhai2\nhai3\nhai4\nhan1\nhan2\nhan3\nhan4\nhang1\nhang2\nhang4\nhao1\nhao2\nhao3\nhao4\nhe1\nhe2\nhe4\nhei1\nhen2\nhen3\nhen4\nheng1\nheng2\nheng4\nhong1\nhong2\nhong3\nhong4\nhou1\nhou2\nhou3\nhou4\nhu1\nhu2\nhu3\nhu4\nhua1\nhua2\nhua4\nhuai2\nhuai4\nhuan1\nhuan2\nhuan3\nhuan4\nhuang1\nhuang2\nhuang3\nhuang4\nhui1\nhui2\nhui3\nhui4\nhun1\nhun2\nhun4\nhuo\nhuo1\nhuo2\nhuo3\nhuo4\ni\nj\nji1\nji2\nji3\nji4\njia\njia1\njia2\njia3\njia4\njian1\njian2\njian3\njian4\njiang1\njiang2\njiang3\njiang4\njiao1\njiao2\njiao3\njiao4\njie1\njie2\njie3\njie4\njin1\njin2\njin3\njin4\njing1\njing2\njing3\njing4\njiong3\njiu1\njiu2\njiu3\njiu4\nju1\nju2\nju3\nju4\njuan1\njuan2\njuan3\njuan4\njue1\njue2\njue4\njun1\njun4\nk\nka1\nka2\nka3\nkai1\nkai2\nkai3\nkai4\nkan1\nkan2\nkan3\nkan4\nkang1\nkang2\nkang4\nkao1\nkao2\nkao3\nkao4\nke1\nke2\nke3\nke4\nken3\nkeng1\nkong1\nkong3\nkong4\nkou1\nkou2\nkou3\nkou4\nku1\nku2\nku3\nku4\nkua1\nkua3\nkua4\nkuai3\nkuai4\nkuan1\nkuan2\nkuan3\nkuang1\nkuang2\nkuang4\nkui1\nkui2\nkui3\nkui4\nkun1\nkun3\nkun4\nkuo4\nl\nla\nla1\nla2\nla3\nla4\nlai2\nlai4\nlan2\nlan3\nlan4\nlang1\nlang2\nlang3\nlang4\nlao1\nlao2\nlao3\nlao4\nle\nle1\nle4\nlei\nlei1\nlei2\nlei3\nlei4\nleng1\nleng2\nleng3\nleng4\nli\nli1\nli2\nli3\nli4\nlia3\nlian2\nlian3\nlian4\nliang2\nliang3\nliang4\nliao1\nliao2\nliao3\nliao4\nlie1\nlie2\nlie3\nlie4\nlin1\nlin2\nlin3\nlin4\nling2\nling3\nling4\nliu1\nliu2\nliu3\nliu4\nlong1\nlong2\nlong3\nlong4\nlou1\nlou2\nlou3\nlou4\nlu1\nlu2\nlu3\nlu4\nluan2\nluan3\nluan4\nlun1\nlun2\nlun4\nluo1\nluo2\nluo3\nluo4\nlv2\nlv3\nlv4\nlve3\nlve4\nm\nma\nma1\nma2\nma3\nma4\nmai2\nmai3\nmai4\nman1\nman2\nman3\nman4\nmang2\nmang3\nmao1\nmao2\nmao3\nmao4\nme\nmei2\nmei3\nmei4\nmen\nmen1\nmen2\nmen4\nmeng\nmeng1\nmeng2\nmeng3\nmeng4\nmi1\nmi2\nmi3\nmi4\nmian2\nmian3\nmian4\nmiao1\nmiao2\nmiao3\nmiao4\nmie1\nmie4\nmin2\nmin3\nming2\nming3\nming4\nmiu4\nmo1\nmo2\nmo3\nmo4\nmou1\nmou2\nmou3\nmu2\nmu3\nmu4\nn\nn2\nna1\nna2\nna3\nna4\nnai2\nnai3\nnai4\nnan1\nnan2\nnan3\nnan4\nnang1\nnang2\nnang3\nnao1\nnao2\nnao3\nnao4\nne\nne2\nne4\nnei3\nnei4\nnen4\nneng2\nni1\nni2\nni3\nni4\nnian1\nnian2\nnian3\nnian4\nniang2\nniang4\nniao2\nniao3\nniao4\nnie1\nnie4\nnin2\nning2\nning3\nning4\nniu1\nniu2\nniu3\nniu4\nnong2\nnong4\nnou4\nnu2\nnu3\nnu4\nnuan3\nnuo2\nnuo4\nnv2\nnv3\nnve4\no\no1\no2\nou1\nou2\nou3\nou4\np\npa1\npa2\npa4\npai1\npai2\npai3\npai4\npan1\npan2\npan4\npang1\npang2\npang4\npao1\npao2\npao3\npao4\npei1\npei2\npei4\npen1\npen2\npen4\npeng1\npeng2\npeng3\npeng4\npi1\npi2\npi3\npi4\npian1\npian2\npian4\npiao1\npiao2\npiao3\npiao4\npie1\npie2\npie3\npin1\npin2\npin3\npin4\nping1\nping2\npo1\npo2\npo3\npo4\npou1\npu1\npu2\npu3\npu4\nq\nqi1\nqi2\nqi3\nqi4\nqia1\nqia3\nqia4\nqian1\nqian2\nqian3\nqian4\nqiang1\nqiang2\nqiang3\nqiang4\nqiao1\nqiao2\nqiao3\nqiao4\nqie1\nqie2\nqie3\nqie4\nqin1\nqin2\nqin3\nqin4\nqing1\nqing2\nqing3\nqing4\nqiong1\nqiong2\nqiu1\nqiu2\nqiu3\nqu1\nqu2\nqu3\nqu4\nquan1\nquan2\nquan3\nquan4\nque1\nque2\nque4\nqun2\nr\nran2\nran3\nrang1\nrang2\nrang3\nrang4\nrao2\nrao3\nrao4\nre2\nre3\nre4\nren2\nren3\nren4\nreng1\nreng2\nri4\nrong1\nrong2\nrong3\nrou2\nrou4\nru2\nru3\nru4\nruan2\nruan3\nrui3\nrui4\nrun4\nruo4\ns\nsa1\nsa2\nsa3\nsa4\nsai1\nsai4\nsan1\nsan2\nsan3\nsan4\nsang1\nsang3\nsang4\nsao1\nsao2\nsao3\nsao4\nse4\nsen1\nseng1\nsha1\nsha2\nsha3\nsha4\nshai1\nshai2\nshai3\nshai4\nshan1\nshan3\nshan4\nshang\nshang1\nshang3\nshang4\nshao1\nshao2\nshao3\nshao4\nshe1\nshe2\nshe3\nshe4\nshei2\nshen1\nshen2\nshen3\nshen4\nsheng1\nsheng2\nsheng3\nsheng4\nshi\nshi1\nshi2\nshi3\nshi4\nshou1\nshou2\nshou3\nshou4\nshu1\nshu2\nshu3\nshu4\nshua1\nshua2\nshua3\nshua4\nshuai1\nshuai3\nshuai4\nshuan1\nshuan4\nshuang1\nshuang3\nshui2\nshui3\nshui4\nshun3\nshun4\nshuo1\nshuo4\nsi1\nsi2\nsi3\nsi4\nsong1\nsong3\nsong4\nsou1\nsou3\nsou4\nsu1\nsu2\nsu4\nsuan1\nsuan4\nsui1\nsui2\nsui3\nsui4\nsun1\nsun3\nsuo\nsuo1\nsuo2\nsuo3\nt\nta1\nta2\nta3\nta4\ntai1\ntai2\ntai4\ntan1\ntan2\ntan3\ntan4\ntang1\ntang2\ntang3\ntang4\ntao1\ntao2\ntao3\ntao4\nte4\nteng2\nti1\nti2\nti3\nti4\ntian1\ntian2\ntian3\ntiao1\ntiao2\ntiao3\ntiao4\ntie1\ntie2\ntie3\ntie4\nting1\nting2\nting3\ntong1\ntong2\ntong3\ntong4\ntou\ntou1\ntou2\ntou4\ntu1\ntu2\ntu3\ntu4\ntuan1\ntuan2\ntui1\ntui2\ntui3\ntui4\ntun1\ntun2\ntun4\ntuo1\ntuo2\ntuo3\ntuo4\nu\nv\nw\nwa\nwa1\nwa2\nwa3\nwa4\nwai1\nwai3\nwai4\nwan1\nwan2\nwan3\nwan4\nwang1\nwang2\nwang3\nwang4\nwei1\nwei2\nwei3\nwei4\nwen1\nwen2\nwen3\nwen4\nweng1\nweng4\nwo1\nwo2\nwo3\nwo4\nwu1\nwu2\nwu3\nwu4\nx\nxi1\nxi2\nxi3\nxi4\nxia1\nxia2\nxia4\nxian1\nxian2\nxian3\nxian4\nxiang1\nxiang2\nxiang3\nxiang4\nxiao1\nxiao2\nxiao3\nxiao4\nxie1\nxie2\nxie3\nxie4\nxin1\nxin2\nxin4\nxing1\nxing2\nxing3\nxing4\nxiong1\nxiong2\nxiu1\nxiu3\nxiu4\nxu\nxu1\nxu2\nxu3\nxu4\nxuan1\nxuan2\nxuan3\nxuan4\nxue1\nxue2\nxue3\nxue4\nxun1\nxun2\nxun4\ny\nya\nya1\nya2\nya3\nya4\nyan1\nyan2\nyan3\nyan4\nyang1\nyang2\nyang3\nyang4\nyao1\nyao2\nyao3\nyao4\nye1\nye2\nye3\nye4\nyi\nyi1\nyi2\nyi3\nyi4\nyin1\nyin2\nyin3\nyin4\nying1\nying2\nying3\nying4\nyo1\nyong1\nyong2\nyong3\nyong4\nyou1\nyou2\nyou3\nyou4\nyu1\nyu2\nyu3\nyu4\nyuan1\nyuan2\nyuan3\nyuan4\nyue1\nyue4\nyun1\nyun2\nyun3\nyun4\nz\nza1\nza2\nza3\nzai1\nzai3\nzai4\nzan1\nzan2\nzan3\nzan4\nzang1\nzang4\nzao1\nzao2\nzao3\nzao4\nze2\nze4\nzei2\nzen3\nzeng1\nzeng4\nzha1\nzha2\nzha3\nzha4\nzhai1\nzhai2\nzhai3\nzhai4\nzhan1\nzhan2\nzhan3\nzhan4\nzhang1\nzhang2\nzhang3\nzhang4\nzhao1\nzhao2\nzhao3\nzhao4\nzhe\nzhe1\nzhe2\nzhe3\nzhe4\nzhen1\nzhen2\nzhen3\nzhen4\nzheng1\nzheng2\nzheng3\nzheng4\nzhi1\nzhi2\nzhi3\nzhi4\nzhong1\nzhong2\nzhong3\nzhong4\nzhou1\nzhou2\nzhou3\nzhou4\nzhu1\nzhu2\nzhu3\nzhu4\nzhua1\nzhua2\nzhua3\nzhuai1\nzhuai3\nzhuai4\nzhuan1\nzhuan2\nzhuan3\nzhuan4\nzhuang1\nzhuang4\nzhui1\nzhui4\nzhun1\nzhun2\nzhun3\nzhuo1\nzhuo2\nzi\nzi1\nzi2\nzi3\nzi4\nzong1\nzong2\nzong3\nzong4\nzou1\nzou2\nzou3\nzou4\nzu1\nzu2\nzu3\nzuan1\nzuan3\nzuan4\nzui2\nzui3\nzui4\nzun1\nzuo\nzuo1\nzuo2\nzuo3\nzuo4\n{\n~\n¡\n¢\n£\n¥\n§\n¨\n©\n«\n®\n¯\n°\n±\n²\n³\n´\nµ\n·\n¹\nº\n»\n¼\n½\n¾\n¿\nÀ\nÁ\nÂ\nÃ\nÄ\nÅ\nÆ\nÇ\nÈ\nÉ\nÊ\nÍ\nÎ\nÑ\nÓ\nÖ\n×\nØ\nÚ\nÜ\nÝ\nÞ\nß\nà\ná\nâ\nã\nä\nå\næ\nç\nè\né\nê\në\nì\ní\nî\nï\nð\nñ\nò\nó\nô\nõ\nö\nø\nù\nú\nû\nü\ný\nĀ\nā\nă\ną\nć\nČ\nč\nĐ\nđ\nē\nė\nę\ně\nĝ\nğ\nħ\nī\nį\nİ\nı\nŁ\nł\nń\nņ\nň\nŋ\nŌ\nō\nő\nœ\nř\nŚ\nś\nŞ\nş\nŠ\nš\nŤ\nť\nũ\nū\nź\nŻ\nż\nŽ\nž\nơ\nư\nǎ\nǐ\nǒ\nǔ\nǚ\nș\nț\nɑ\nɔ\nɕ\nə\nɛ\nɜ\nɡ\nɣ\nɪ\nɫ\nɴ\nɹ\nɾ\nʃ\nʊ\nʌ\nʒ\nʔ\nʰ\nʷ\nʻ\nʾ\nʿ\nˈ\nː\n˙\n˜\nˢ\ń\n̅\nΑ\nΒ\nΔ\nΕ\nΘ\nΚ\nΛ\nΜ\nΞ\nΠ\nΣ\nΤ\nΦ\nΧ\nΨ\nΩ\nά\nέ\nή\nί\nα\nβ\nγ\nδ\nε\nζ\nη\nθ\nι\nκ\nλ\nμ\nν\nξ\nο\nπ\nρ\nς\nσ\nτ\nυ\nφ\nχ\nψ\nω\nϊ\nό\nύ\nώ\nϕ\nϵ\nЁ\nА\nБ\nВ\nГ\nД\nЕ\nЖ\nЗ\nИ\nЙ\nК\nЛ\nМ\nН\nО\nП\nР\nС\nТ\nУ\nФ\nХ\nЦ\nЧ\nШ\nЩ\nЫ\nЬ\nЭ\nЮ\nЯ\nа\nб\nв\nг\nд\nе\nж\nз\nи\nй\nк\nл\nм\nн\nо\nп\nр\nс\nт\nу\nф\nх\nц\nч\nш\nщ\nъ\nы\nь\nэ\nю\nя\nё\nі\nְ\nִ\nֵ\nֶ\nַ\nָ\nֹ\nּ\n־\nׁ\nא\nב\nג\nד\nה\nו\nז\nח\nט\nי\nכ\nל\nם\nמ\nן\nנ\nס\nע\nפ\nק\nר\nש\nת\nأ\nب\nة\nت\nج\nح\nد\nر\nز\nس\nص\nط\nع\nق\nك\nل\nم\nن\nه\nو\nي\nَ\nُ\nِ\nْ\nก\nข\nง\nจ\nต\nท\nน\nป\nย\nร\nว\nส\nห\nอ\nฮ\nั\nา\nี\nึ\nโ\nใ\nไ\n่\n้\n์\nḍ\nḤ\nḥ\nṁ\nṃ\nṅ\nṇ\nṚ\nṛ\nṢ\nṣ\nṬ\nṭ\nạ\nả\nẤ\nấ\nầ\nậ\nắ\nằ\nẻ\nẽ\nế\nề\nể\nễ\nệ\nị\nọ\nỏ\nố\nồ\nộ\nớ\nờ\nở\nụ\nủ\nứ\nữ\nἀ\nἁ\nἈ\nἐ\nἔ\nἰ\nἱ\nὀ\nὁ\nὐ\nὲ\nὸ\nᾶ\n᾽\nῆ\nῇ\nῶ\n‎\n‑\n‒\n–\n—\n―\n‖\n†\n‡\n•\n…\n‧\n‬\n′\n″\n⁄\n⁡\n⁰\n⁴\n⁵\n⁶\n⁷\n⁸\n⁹\n₁\n₂\n₃\n€\n₱\n₹\n₽\n℃\nℏ\nℓ\n№\nℝ\n™\n⅓\n⅔\n⅛\n→\n∂\n∈\n∑\n−\n∗\n√\n∞\n∫\n≈\n≠\n≡\n≤\n≥\n⋅\n⋯\n█\n♪\n⟨\n⟩\n、\n。\n《\n》\n「\n」\n【\n】\nあ\nう\nえ\nお\nか\nが\nき\nぎ\nく\nぐ\nけ\nげ\nこ\nご\nさ\nし\nじ\nす\nず\nせ\nぜ\nそ\nぞ\nた\nだ\nち\nっ\nつ\nで\nと\nど\nな\nに\nね\nの\nは\nば\nひ\nぶ\nへ\nべ\nま\nみ\nむ\nめ\nも\nゃ\nや\nゆ\nょ\nよ\nら\nり\nる\nれ\nろ\nわ\nを\nん\nァ\nア\nィ\nイ\nウ\nェ\nエ\nオ\nカ\nガ\nキ\nク\nケ\nゲ\nコ\nゴ\nサ\nザ\nシ\nジ\nス\nズ\nセ\nゾ\nタ\nダ\nチ\nッ\nツ\nテ\nデ\nト\nド\nナ\nニ\nネ\nノ\nバ\nパ\nビ\nピ\nフ\nプ\nヘ\nベ\nペ\nホ\nボ\nポ\nマ\nミ\nム\nメ\nモ\nャ\nヤ\nュ\nユ\nョ\nヨ\nラ\nリ\nル\nレ\nロ\nワ\nン\n・\nー\nㄋ\nㄍ\nㄎ\nㄏ\nㄓ\nㄕ\nㄚ\nㄜ\nㄟ\nㄤ\nㄥ\nㄧ\nㄱ\nㄴ\nㄷ\nㄹ\nㅁ\nㅂ\nㅅ\nㅈ\nㅍ\nㅎ\nㅏ\nㅓ\nㅗ\nㅜ\nㅡ\nㅣ\n㗎\n가\n각\n간\n갈\n감\n갑\n갓\n갔\n강\n같\n개\n거\n건\n걸\n겁\n것\n겉\n게\n겠\n겨\n결\n겼\n경\n계\n고\n곤\n골\n곱\n공\n과\n관\n광\n교\n구\n국\n굴\n귀\n귄\n그\n근\n글\n금\n기\n긴\n길\n까\n깍\n깔\n깜\n깨\n께\n꼬\n꼭\n꽃\n꾸\n꿔\n끔\n끗\n끝\n끼\n나\n난\n날\n남\n납\n내\n냐\n냥\n너\n넘\n넣\n네\n녁\n년\n녕\n노\n녹\n놀\n누\n눈\n느\n는\n늘\n니\n님\n닙\n다\n닥\n단\n달\n닭\n당\n대\n더\n덕\n던\n덥\n데\n도\n독\n동\n돼\n됐\n되\n된\n될\n두\n둑\n둥\n드\n들\n등\n디\n따\n딱\n딸\n땅\n때\n떤\n떨\n떻\n또\n똑\n뚱\n뛰\n뜻\n띠\n라\n락\n란\n람\n랍\n랑\n래\n랜\n러\n런\n럼\n렇\n레\n려\n력\n렵\n렸\n로\n록\n롬\n루\n르\n른\n를\n름\n릉\n리\n릴\n림\n마\n막\n만\n많\n말\n맑\n맙\n맛\n매\n머\n먹\n멍\n메\n면\n명\n몇\n모\n목\n몸\n못\n무\n문\n물\n뭐\n뭘\n미\n민\n밌\n밑\n바\n박\n밖\n반\n받\n발\n밤\n밥\n방\n배\n백\n밸\n뱀\n버\n번\n벌\n벚\n베\n벼\n벽\n별\n병\n보\n복\n본\n볼\n봐\n봤\n부\n분\n불\n비\n빔\n빛\n빠\n빨\n뼈\n뽀\n뿅\n쁘\n사\n산\n살\n삼\n샀\n상\n새\n색\n생\n서\n선\n설\n섭\n섰\n성\n세\n셔\n션\n셨\n소\n속\n손\n송\n수\n숙\n순\n술\n숫\n숭\n숲\n쉬\n쉽\n스\n슨\n습\n슷\n시\n식\n신\n실\n싫\n심\n십\n싶\n싸\n써\n쓰\n쓴\n씌\n씨\n씩\n씬\n아\n악\n안\n않\n알\n야\n약\n얀\n양\n얘\n어\n언\n얼\n엄\n업\n없\n었\n엉\n에\n여\n역\n연\n염\n엽\n영\n옆\n예\n옛\n오\n온\n올\n옷\n옹\n와\n왔\n왜\n요\n욕\n용\n우\n운\n울\n웃\n워\n원\n월\n웠\n위\n윙\n유\n육\n윤\n으\n은\n을\n음\n응\n의\n이\n익\n인\n일\n읽\n임\n입\n있\n자\n작\n잔\n잖\n잘\n잡\n잤\n장\n재\n저\n전\n점\n정\n제\n져\n졌\n조\n족\n좀\n종\n좋\n죠\n주\n준\n줄\n중\n줘\n즈\n즐\n즘\n지\n진\n집\n짜\n짝\n쩌\n쪼\n쪽\n쫌\n쭈\n쯔\n찌\n찍\n차\n착\n찾\n책\n처\n천\n철\n체\n쳐\n쳤\n초\n촌\n추\n출\n춤\n춥\n춰\n치\n친\n칠\n침\n칩\n칼\n커\n켓\n코\n콩\n쿠\n퀴\n크\n큰\n큽\n키\n킨\n타\n태\n터\n턴\n털\n테\n토\n통\n투\n트\n특\n튼\n틀\n티\n팀\n파\n팔\n패\n페\n펜\n펭\n평\n포\n폭\n표\n품\n풍\n프\n플\n피\n필\n하\n학\n한\n할\n함\n합\n항\n해\n햇\n했\n행\n허\n험\n형\n혜\n호\n혼\n홀\n화\n회\n획\n후\n휴\n흐\n흔\n희\n히\n힘\nﷺ\nﷻ\n！\n，\n？\n�\n𠮶\n" }, { "path": "src/f5_tts/infer/infer_cli.py", "content": "import argparse\nimport codecs\nimport os\nimport re\nfrom datetime import datetime\nfrom importlib.resources import files\nfrom pathlib import Path\n\nimport numpy as np\nimport soundfile as sf\nimport tomli\nfrom cached_path import cached_path\nfrom hydra.utils import get_class\nfrom omegaconf import OmegaConf\nfrom unidecode import unidecode\n\nfrom f5_tts.infer.utils_infer import (\n cfg_strength,\n cross_fade_duration,\n device,\n fix_duration,\n infer_process,\n load_model,\n load_vocoder,\n mel_spec_type,\n nfe_step,\n preprocess_ref_audio_text,\n remove_silence_for_generated_wav,\n speed,\n sway_sampling_coef,\n target_rms,\n)\n\n\nparser = argparse.ArgumentParser(\n prog=\"python3 infer-cli.py\",\n description=\"Commandline interface for E2/F5 TTS with Advanced Batch Processing.\",\n epilog=\"Specify options above to override one or more settings from config.\",\n)\nparser.add_argument(\n \"-c\",\n \"--config\",\n type=str,\n default=os.path.join(files(\"f5_tts\").joinpath(\"infer/examples/basic\"), \"basic.toml\"),\n help=\"The configuration file, default see infer/examples/basic/basic.toml\",\n)\n\n\n# Note. Not to provide default value here in order to read default from config file\n\nparser.add_argument(\n \"-m\",\n \"--model\",\n type=str,\n help=\"The model name: F5TTS_v1_Base | F5TTS_Base | E2TTS_Base | etc.\",\n)\nparser.add_argument(\n \"-mc\",\n \"--model_cfg\",\n type=str,\n help=\"The path to F5-TTS model config file .yaml\",\n)\nparser.add_argument(\n \"-p\",\n \"--ckpt_file\",\n type=str,\n help=\"The path to model checkpoint .pt, leave blank to use default\",\n)\nparser.add_argument(\n \"-v\",\n \"--vocab_file\",\n type=str,\n help=\"The path to vocab file .txt, leave blank to use default\",\n)\nparser.add_argument(\n \"-r\",\n \"--ref_audio\",\n type=str,\n help=\"The reference audio file.\",\n)\nparser.add_argument(\n \"-s\",\n \"--ref_text\",\n type=str,\n help=\"The transcript/subtitle for the reference audio\",\n)\nparser.add_argument(\n \"-t\",\n \"--gen_text\",\n type=str,\n help=\"The text to make model synthesize a speech\",\n)\nparser.add_argument(\n \"-f\",\n \"--gen_file\",\n type=str,\n help=\"The file with text to generate, will ignore --gen_text\",\n)\nparser.add_argument(\n \"-o\",\n \"--output_dir\",\n type=str,\n help=\"The path to output folder\",\n)\nparser.add_argument(\n \"-w\",\n \"--output_file\",\n type=str,\n help=\"The name of output file\",\n)\nparser.add_argument(\n \"--save_chunk\",\n action=\"store_true\",\n help=\"To save each audio chunks during inference\",\n)\nparser.add_argument(\n \"--no_legacy_text\",\n action=\"store_false\",\n help=\"Not to use lossy ASCII transliterations of unicode text in saved file names.\",\n)\nparser.add_argument(\n \"--remove_silence\",\n action=\"store_true\",\n help=\"To remove long silence found in ouput\",\n)\nparser.add_argument(\n \"--load_vocoder_from_local\",\n action=\"store_true\",\n help=\"To load vocoder from local dir, default to ../checkpoints/vocos-mel-24khz\",\n)\nparser.add_argument(\n \"--vocoder_name\",\n type=str,\n choices=[\"vocos\", \"bigvgan\"],\n help=f\"Used vocoder name: vocos | bigvgan, default {mel_spec_type}\",\n)\nparser.add_argument(\n \"--target_rms\",\n type=float,\n help=f\"Target output speech loudness normalization value, default {target_rms}\",\n)\nparser.add_argument(\n \"--cross_fade_duration\",\n type=float,\n help=f\"Duration of cross-fade between audio segments in seconds, default {cross_fade_duration}\",\n)\nparser.add_argument(\n \"--nfe_step\",\n type=int,\n help=f\"The number of function evaluation (denoising steps), default {nfe_step}\",\n)\nparser.add_argument(\n \"--cfg_strength\",\n type=float,\n help=f\"Classifier-free guidance strength, default {cfg_strength}\",\n)\nparser.add_argument(\n \"--sway_sampling_coef\",\n type=float,\n help=f\"Sway Sampling coefficient, default {sway_sampling_coef}\",\n)\nparser.add_argument(\n \"--speed\",\n type=float,\n help=f\"The speed of the generated audio, default {speed}\",\n)\nparser.add_argument(\n \"--fix_duration\",\n type=float,\n help=f\"Fix the total duration (ref and gen audios) in seconds, default {fix_duration}\",\n)\nparser.add_argument(\n \"--device\",\n type=str,\n help=\"Specify the device to run on\",\n)\nargs = parser.parse_args()\n\n\n# config file\n\nconfig = tomli.load(open(args.config, \"rb\"))\n\n\n# command-line interface parameters\n\nmodel = args.model or config.get(\"model\", \"F5TTS_v1_Base\")\nckpt_file = args.ckpt_file or config.get(\"ckpt_file\", \"\")\nvocab_file = args.vocab_file or config.get(\"vocab_file\", \"\")\n\nref_audio = args.ref_audio or config.get(\"ref_audio\", \"infer/examples/basic/basic_ref_en.wav\")\nref_text = (\n args.ref_text\n if args.ref_text is not None\n else config.get(\"ref_text\", \"Some call me nature, others call me mother nature.\")\n)\ngen_text = args.gen_text or config.get(\"gen_text\", \"Here we generate something just for test.\")\ngen_file = args.gen_file or config.get(\"gen_file\", \"\")\n\noutput_dir = args.output_dir or config.get(\"output_dir\", \"tests\")\noutput_file = args.output_file or config.get(\n \"output_file\", f\"infer_cli_{datetime.now().strftime(r'%Y%m%d_%H%M%S')}.wav\"\n)\n\nsave_chunk = args.save_chunk or config.get(\"save_chunk\", False)\nuse_legacy_text = args.no_legacy_text or config.get(\"no_legacy_text\", False) # no_legacy_text is a store_false arg\nif save_chunk and use_legacy_text:\n print(\n \"\\nWarning to --save_chunk: lossy ASCII transliterations of unicode text for legacy (.wav) file names, --no_legacy_text to disable.\\n\"\n )\n\nremove_silence = args.remove_silence or config.get(\"remove_silence\", False)\nload_vocoder_from_local = args.load_vocoder_from_local or config.get(\"load_vocoder_from_local\", False)\n\nvocoder_name = args.vocoder_name or config.get(\"vocoder_name\", mel_spec_type)\ntarget_rms = args.target_rms or config.get(\"target_rms\", target_rms)\ncross_fade_duration = args.cross_fade_duration or config.get(\"cross_fade_duration\", cross_fade_duration)\nnfe_step = args.nfe_step or config.get(\"nfe_step\", nfe_step)\ncfg_strength = args.cfg_strength or config.get(\"cfg_strength\", cfg_strength)\nsway_sampling_coef = args.sway_sampling_coef or config.get(\"sway_sampling_coef\", sway_sampling_coef)\nspeed = args.speed or config.get(\"speed\", speed)\nfix_duration = args.fix_duration or config.get(\"fix_duration\", fix_duration)\ndevice = args.device or config.get(\"device\", device)\n\n\n# patches for pip pkg user\nif \"infer/examples/\" in ref_audio:\n ref_audio = str(files(\"f5_tts\").joinpath(f\"{ref_audio}\"))\nif \"infer/examples/\" in gen_file:\n gen_file = str(files(\"f5_tts\").joinpath(f\"{gen_file}\"))\nif \"voices\" in config:\n for voice in config[\"voices\"]:\n voice_ref_audio = config[\"voices\"][voice][\"ref_audio\"]\n if \"infer/examples/\" in voice_ref_audio:\n config[\"voices\"][voice][\"ref_audio\"] = str(files(\"f5_tts\").joinpath(f\"{voice_ref_audio}\"))\n\n\n# ignore gen_text if gen_file provided\n\nif gen_file:\n gen_text = codecs.open(gen_file, \"r\", \"utf-8\").read()\n\n\n# output path\n\nwave_path = Path(output_dir) / output_file\n# spectrogram_path = Path(output_dir) / \"infer_cli_out.png\"\nif save_chunk:\n output_chunk_dir = os.path.join(output_dir, f\"{Path(output_file).stem}_chunks\")\n if not os.path.exists(output_chunk_dir):\n os.makedirs(output_chunk_dir)\n\n\n# load vocoder\n\nif vocoder_name == \"vocos\":\n vocoder_local_path = \"../checkpoints/vocos-mel-24khz\"\nelif vocoder_name == \"bigvgan\":\n vocoder_local_path = \"../checkpoints/bigvgan_v2_24khz_100band_256x\"\n\nvocoder = load_vocoder(\n vocoder_name=vocoder_name, is_local=load_vocoder_from_local, local_path=vocoder_local_path, device=device\n)\n\n\n# load TTS model\n\nmodel_cfg = OmegaConf.load(\n args.model_cfg or config.get(\"model_cfg\", str(files(\"f5_tts\").joinpath(f\"configs/{model}.yaml\")))\n)\nmodel_cls = get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\nmodel_arc = model_cfg.model.arch\n\nrepo_name, ckpt_step, ckpt_type = \"F5-TTS\", 1250000, \"safetensors\"\n\nif model != \"F5TTS_Base\":\n assert vocoder_name == model_cfg.model.mel_spec.mel_spec_type\n\n# override for previous models\nif model == \"F5TTS_Base\":\n if vocoder_name == \"vocos\":\n ckpt_step = 1200000\n elif vocoder_name == \"bigvgan\":\n model = \"F5TTS_Base_bigvgan\"\n ckpt_type = \"pt\"\nelif model == \"E2TTS_Base\":\n repo_name = \"E2-TTS\"\n ckpt_step = 1200000\n\nif not ckpt_file:\n ckpt_file = str(cached_path(f\"hf://SWivid/{repo_name}/{model}/model_{ckpt_step}.{ckpt_type}\"))\nelif ckpt_file.startswith(\"hf://\"):\n ckpt_file = str(cached_path(ckpt_file))\n\nif vocab_file.startswith(\"hf://\"):\n vocab_file = str(cached_path(vocab_file))\n\nprint(f\"Using {model}...\")\nema_model = load_model(\n model_cls, model_arc, ckpt_file, mel_spec_type=vocoder_name, vocab_file=vocab_file, device=device\n)\n\n\n# inference process\n\n\ndef main():\n main_voice = {\"ref_audio\": ref_audio, \"ref_text\": ref_text}\n if \"voices\" not in config:\n voices = {\"main\": main_voice}\n else:\n voices = config[\"voices\"]\n voices[\"main\"] = main_voice\n for voice in voices:\n print(\"Voice:\", voice)\n print(\"ref_audio \", voices[voice][\"ref_audio\"])\n voices[voice][\"ref_audio\"], voices[voice][\"ref_text\"] = preprocess_ref_audio_text(\n voices[voice][\"ref_audio\"], voices[voice][\"ref_text\"]\n )\n print(\"ref_audio_\", voices[voice][\"ref_audio\"], \"\\n\\n\")\n\n generated_audio_segments = []\n reg1 = r\"(?=\\[\\w+\\])\"\n chunks = re.split(reg1, gen_text)\n reg2 = r\"\\[(\\w+)\\]\"\n for text in chunks:\n if not text.strip():\n continue\n match = re.match(reg2, text)\n if match:\n voice = match[1]\n else:\n print(\"No voice tag found, using main.\")\n voice = \"main\"\n if voice not in voices:\n print(f\"Voice {voice} not found, using main.\")\n voice = \"main\"\n text = re.sub(reg2, \"\", text)\n ref_audio_ = voices[voice][\"ref_audio\"]\n ref_text_ = voices[voice][\"ref_text\"]\n local_speed = voices[voice].get(\"speed\", speed)\n gen_text_ = text.strip()\n print(f\"Voice: {voice}\")\n audio_segment, final_sample_rate, spectrogram = infer_process(\n ref_audio_,\n ref_text_,\n gen_text_,\n ema_model,\n vocoder,\n mel_spec_type=vocoder_name,\n target_rms=target_rms,\n cross_fade_duration=cross_fade_duration,\n nfe_step=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n speed=local_speed,\n fix_duration=fix_duration,\n device=device,\n )\n generated_audio_segments.append(audio_segment)\n\n if save_chunk:\n if len(gen_text_) > 200:\n gen_text_ = gen_text_[:200] + \" ... \"\n if use_legacy_text:\n gen_text_ = unidecode(gen_text_)\n sf.write(\n os.path.join(output_chunk_dir, f\"{len(generated_audio_segments) - 1}_{gen_text_}.wav\"),\n audio_segment,\n final_sample_rate,\n )\n\n if generated_audio_segments:\n final_wave = np.concatenate(generated_audio_segments)\n\n if not os.path.exists(output_dir):\n os.makedirs(output_dir)\n\n with open(wave_path, \"wb\") as f:\n sf.write(f.name, final_wave, final_sample_rate)\n # Remove silence\n if remove_silence:\n remove_silence_for_generated_wav(f.name)\n print(f.name)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/infer/infer_gradio.py", "content": "# ruff: noqa: E402\n# Above allows ruff to ignore E402: module level import not at top of file\n\nimport gc\nimport json\nimport os\nimport re\nimport tempfile\nfrom collections import OrderedDict\nfrom functools import lru_cache\nfrom importlib.resources import files\n\nimport click\nimport gradio as gr\nimport numpy as np\nimport soundfile as sf\nimport torch\nimport torchaudio\nfrom cached_path import cached_path\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n\ntry:\n import spaces\n\n USING_SPACES = True\nexcept ImportError:\n USING_SPACES = False\n\n\ndef gpu_decorator(func):\n if USING_SPACES:\n return spaces.GPU(func)\n else:\n return func\n\n\nfrom f5_tts.infer.utils_infer import (\n infer_process,\n load_model,\n load_vocoder,\n preprocess_ref_audio_text,\n remove_silence_for_generated_wav,\n save_spectrogram,\n tempfile_kwargs,\n)\nfrom f5_tts.model import DiT, UNetT\n\n\nDEFAULT_TTS_MODEL = \"F5-TTS_v1\"\ntts_model_choice = DEFAULT_TTS_MODEL\n\nDEFAULT_TTS_MODEL_CFG = [\n \"hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors\",\n \"hf://SWivid/F5-TTS/F5TTS_v1_Base/vocab.txt\",\n json.dumps(dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)),\n]\n\n\n# load models\n\nvocoder = load_vocoder()\n\n\ndef load_f5tts():\n ckpt_path = str(cached_path(DEFAULT_TTS_MODEL_CFG[0]))\n F5TTS_model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])\n return load_model(DiT, F5TTS_model_cfg, ckpt_path)\n\n\ndef load_e2tts():\n ckpt_path = str(cached_path(\"hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.safetensors\"))\n E2TTS_model_cfg = dict(dim=1024, depth=24, heads=16, ff_mult=4, text_mask_padding=False, pe_attn_head=1)\n return load_model(UNetT, E2TTS_model_cfg, ckpt_path)\n\n\ndef load_custom(ckpt_path: str, vocab_path=\"\", model_cfg=None):\n ckpt_path, vocab_path = ckpt_path.strip(), vocab_path.strip()\n if ckpt_path.startswith(\"hf://\"):\n ckpt_path = str(cached_path(ckpt_path))\n if vocab_path.startswith(\"hf://\"):\n vocab_path = str(cached_path(vocab_path))\n if model_cfg is None:\n model_cfg = json.loads(DEFAULT_TTS_MODEL_CFG[2])\n elif isinstance(model_cfg, str):\n model_cfg = json.loads(model_cfg)\n return load_model(DiT, model_cfg, ckpt_path, vocab_file=vocab_path)\n\n\nF5TTS_ema_model = load_f5tts()\nE2TTS_ema_model = load_e2tts() if USING_SPACES else None\ncustom_ema_model, pre_custom_path = None, \"\"\n\nchat_model_state = None\nchat_tokenizer_state = None\n\n\n@gpu_decorator\ndef chat_model_inference(messages, model, tokenizer):\n \"\"\"Generate response using Qwen\"\"\"\n text = tokenizer.apply_chat_template(\n messages,\n tokenize=False,\n add_generation_prompt=True,\n )\n\n model_inputs = tokenizer([text], return_tensors=\"pt\").to(model.device)\n generated_ids = model.generate(\n **model_inputs,\n max_new_tokens=512,\n temperature=0.7,\n top_p=0.95,\n )\n\n generated_ids = [\n output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n ]\n return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n\n\n@gpu_decorator\ndef load_text_from_file(file):\n if file:\n with open(file, \"r\", encoding=\"utf-8\") as f:\n text = f.read().strip()\n else:\n text = \"\"\n return gr.update(value=text)\n\n\n@lru_cache(maxsize=1000) # NOTE. need to ensure params of infer() hashable\n@gpu_decorator\ndef infer(\n ref_audio_orig,\n ref_text,\n gen_text,\n model,\n remove_silence,\n seed,\n cross_fade_duration=0.15,\n nfe_step=32,\n speed=1,\n show_info=gr.Info,\n):\n if not ref_audio_orig:\n gr.Warning(\"Please provide reference audio.\")\n return gr.update(), gr.update(), ref_text\n\n # Set inference seed\n if seed < 0 or seed > 2**31 - 1:\n gr.Warning(\"Seed must in range 0 ~ 2147483647. Using random seed instead.\")\n seed = np.random.randint(0, 2**31 - 1)\n torch.manual_seed(seed)\n used_seed = seed\n\n if not gen_text.strip():\n gr.Warning(\"Please enter text to generate or upload a text file.\")\n return gr.update(), gr.update(), ref_text\n\n ref_audio, ref_text = preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=show_info)\n\n if model == DEFAULT_TTS_MODEL:\n ema_model = F5TTS_ema_model\n elif model == \"E2-TTS\":\n global E2TTS_ema_model\n if E2TTS_ema_model is None:\n show_info(\"Loading E2-TTS model...\")\n E2TTS_ema_model = load_e2tts()\n ema_model = E2TTS_ema_model\n elif isinstance(model, tuple) and model[0] == \"Custom\":\n assert not USING_SPACES, \"Only official checkpoints allowed in Spaces.\"\n global custom_ema_model, pre_custom_path\n if pre_custom_path != model[1]:\n show_info(\"Loading Custom TTS model...\")\n custom_ema_model = load_custom(model[1], vocab_path=model[2], model_cfg=model[3])\n pre_custom_path = model[1]\n ema_model = custom_ema_model\n\n final_wave, final_sample_rate, combined_spectrogram = infer_process(\n ref_audio,\n ref_text,\n gen_text,\n ema_model,\n vocoder,\n cross_fade_duration=cross_fade_duration,\n nfe_step=nfe_step,\n speed=speed,\n show_info=show_info,\n progress=gr.Progress(),\n )\n\n # Remove silence\n if remove_silence:\n with tempfile.NamedTemporaryFile(suffix=\".wav\", **tempfile_kwargs) as f:\n temp_path = f.name\n try:\n sf.write(temp_path, final_wave, final_sample_rate)\n remove_silence_for_generated_wav(f.name)\n final_wave, _ = torchaudio.load(f.name)\n finally:\n os.unlink(temp_path)\n final_wave = final_wave.squeeze().cpu().numpy()\n\n # Save the spectrogram\n with tempfile.NamedTemporaryFile(suffix=\".png\", **tempfile_kwargs) as tmp_spectrogram:\n spectrogram_path = tmp_spectrogram.name\n save_spectrogram(combined_spectrogram, spectrogram_path)\n\n return (final_sample_rate, final_wave), spectrogram_path, ref_text, used_seed\n\n\nwith gr.Blocks() as app_tts:\n gr.Markdown(\"# Batched TTS\")\n ref_audio_input = gr.Audio(label=\"Reference Audio\", type=\"filepath\")\n with gr.Row():\n gen_text_input = gr.Textbox(\n label=\"Text to Generate\",\n lines=10,\n max_lines=40,\n scale=4,\n )\n gen_text_file = gr.File(label=\"Load Text to Generate from File (.txt)\", file_types=[\".txt\"], scale=1)\n generate_btn = gr.Button(\"Synthesize\", variant=\"primary\")\n with gr.Accordion(\"Advanced Settings\", open=True) as adv_settn:\n with gr.Row():\n ref_text_input = gr.Textbox(\n label=\"Reference Text\",\n info=\"Leave blank to automatically transcribe the reference audio. If you enter text or upload a file, it will override automatic transcription.\",\n lines=2,\n scale=4,\n )\n ref_text_file = gr.File(label=\"Load Reference Text from File (.txt)\", file_types=[\".txt\"], scale=1)\n with gr.Row():\n randomize_seed = gr.Checkbox(\n label=\"Randomize Seed\",\n info=\"Check to use a random seed for each generation. Uncheck to use the seed specified.\",\n value=True,\n scale=3,\n )\n seed_input = gr.Number(show_label=False, value=0, precision=0, scale=1)\n with gr.Column(scale=4):\n remove_silence = gr.Checkbox(\n label=\"Remove Silences\",\n info=\"If undesired long silence(s) produced, turn on to automatically detect and crop.\",\n value=False,\n )\n speed_slider = gr.Slider(\n label=\"Speed\",\n minimum=0.3,\n maximum=2.0,\n value=1.0,\n step=0.1,\n info=\"Adjust the speed of the audio.\",\n )\n nfe_slider = gr.Slider(\n label=\"NFE Steps\",\n minimum=4,\n maximum=64,\n value=32,\n step=2,\n info=\"Set the number of denoising steps.\",\n )\n cross_fade_duration_slider = gr.Slider(\n label=\"Cross-Fade Duration (s)\",\n minimum=0.0,\n maximum=1.0,\n value=0.15,\n step=0.01,\n info=\"Set the duration of the cross-fade between audio clips.\",\n )\n\n def collapse_accordion():\n return gr.Accordion(open=False)\n\n # Workaround for https://github.com/SWivid/F5-TTS/issues/1239#issuecomment-3677987413\n # i.e. to set gr.Accordion(open=True) by default, then collapse manually Blocks loaded\n app_tts.load(\n fn=collapse_accordion,\n inputs=None,\n outputs=adv_settn,\n )\n\n audio_output = gr.Audio(label=\"Synthesized Audio\")\n spectrogram_output = gr.Image(label=\"Spectrogram\")\n\n @gpu_decorator\n def basic_tts(\n ref_audio_input,\n ref_text_input,\n gen_text_input,\n remove_silence,\n randomize_seed,\n seed_input,\n cross_fade_duration_slider,\n nfe_slider,\n speed_slider,\n ):\n if randomize_seed:\n seed_input = np.random.randint(0, 2**31 - 1)\n\n audio_out, spectrogram_path, ref_text_out, used_seed = infer(\n ref_audio_input,\n ref_text_input,\n gen_text_input,\n tts_model_choice,\n remove_silence,\n seed=seed_input,\n cross_fade_duration=cross_fade_duration_slider,\n nfe_step=nfe_slider,\n speed=speed_slider,\n )\n return audio_out, spectrogram_path, ref_text_out, used_seed\n\n gen_text_file.upload(\n load_text_from_file,\n inputs=[gen_text_file],\n outputs=[gen_text_input],\n )\n\n ref_text_file.upload(\n load_text_from_file,\n inputs=[ref_text_file],\n outputs=[ref_text_input],\n )\n\n ref_audio_input.clear(\n lambda: [None, None],\n None,\n [ref_text_input, ref_text_file],\n )\n\n generate_btn.click(\n basic_tts,\n inputs=[\n ref_audio_input,\n ref_text_input,\n gen_text_input,\n remove_silence,\n randomize_seed,\n seed_input,\n cross_fade_duration_slider,\n nfe_slider,\n speed_slider,\n ],\n outputs=[audio_output, spectrogram_output, ref_text_input, seed_input],\n )\n\n\ndef parse_speechtypes_text(gen_text):\n # Pattern to find {str} or {\"name\": str, \"seed\": int, \"speed\": float}\n pattern = r\"(\\{.*?\\})\"\n\n # Split the text by the pattern\n tokens = re.split(pattern, gen_text)\n\n segments = []\n\n current_type_dict = {\n \"name\": \"Regular\",\n \"seed\": -1,\n \"speed\": 1.0,\n }\n\n for i in range(len(tokens)):\n if i % 2 == 0:\n # This is text\n text = tokens[i].strip()\n if text:\n current_type_dict[\"text\"] = text\n segments.append(current_type_dict)\n else:\n # This is type\n type_str = tokens[i].strip()\n try: # if type dict\n current_type_dict = json.loads(type_str)\n except json.decoder.JSONDecodeError:\n type_str = type_str[1:-1] # remove brace {}\n current_type_dict = {\"name\": type_str, \"seed\": -1, \"speed\": 1.0}\n\n return segments\n\n\nwith gr.Blocks() as app_multistyle:\n # New section for multistyle generation\n gr.Markdown(\n \"\"\"\n # Multiple Speech-Type Generation\n\n 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.\n \"\"\"\n )\n\n with gr.Row():\n gr.Markdown(\n \"\"\"\n **Example Input:**
\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?!\n \"\"\"\n )\n\n gr.Markdown(\n \"\"\"\n **Example Input 2:**
\n {\"name\": \"Speaker1_Happy\", \"seed\": -1, \"speed\": 1} Hello, I'd like to order a sandwich please.
\n {\"name\": \"Speaker2_Regular\", \"seed\": -1, \"speed\": 1} Sorry, we're out of bread.
\n {\"name\": \"Speaker1_Sad\", \"seed\": -1, \"speed\": 1} I really wanted a sandwich though...
\n {\"name\": \"Speaker2_Whisper\", \"seed\": -1, \"speed\": 1} I'll give you the last one I was hiding.\n \"\"\"\n )\n\n gr.Markdown(\n '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.'\n )\n\n # Regular speech type (mandatory)\n with gr.Row(variant=\"compact\") as regular_row:\n with gr.Column(scale=1, min_width=160):\n regular_name = gr.Textbox(value=\"Regular\", label=\"Speech Type Name\")\n regular_insert = gr.Button(\"Insert Label\", variant=\"secondary\")\n with gr.Column(scale=3):\n regular_audio = gr.Audio(label=\"Regular Reference Audio\", type=\"filepath\")\n with gr.Column(scale=3):\n regular_ref_text = gr.Textbox(label=\"Reference Text (Regular)\", lines=4)\n with gr.Row():\n regular_seed_slider = gr.Slider(\n show_label=False, minimum=-1, maximum=999, value=-1, step=1, info=\"Seed, -1 for random\"\n )\n regular_speed_slider = gr.Slider(\n show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info=\"Adjust the speed\"\n )\n with gr.Column(scale=1, min_width=160):\n regular_ref_text_file = gr.File(label=\"Load Reference Text from File (.txt)\", file_types=[\".txt\"])\n\n # Regular speech type (max 100)\n max_speech_types = 100\n speech_type_rows = [regular_row]\n speech_type_names = [regular_name]\n speech_type_audios = [regular_audio]\n speech_type_ref_texts = [regular_ref_text]\n speech_type_ref_text_files = [regular_ref_text_file]\n speech_type_seeds = [regular_seed_slider]\n speech_type_speeds = [regular_speed_slider]\n speech_type_delete_btns = [None]\n speech_type_insert_btns = [regular_insert]\n\n # Additional speech types (99 more)\n for i in range(max_speech_types - 1):\n with gr.Row(variant=\"compact\", visible=False) as row:\n with gr.Column(scale=1, min_width=160):\n name_input = gr.Textbox(label=\"Speech Type Name\")\n insert_btn = gr.Button(\"Insert Label\", variant=\"secondary\")\n delete_btn = gr.Button(\"Delete Type\", variant=\"stop\")\n with gr.Column(scale=3):\n audio_input = gr.Audio(label=\"Reference Audio\", type=\"filepath\")\n with gr.Column(scale=3):\n ref_text_input = gr.Textbox(label=\"Reference Text\", lines=4)\n with gr.Row():\n seed_input = gr.Slider(\n show_label=False, minimum=-1, maximum=999, value=-1, step=1, info=\"Seed. -1 for random\"\n )\n speed_input = gr.Slider(\n show_label=False, minimum=0.3, maximum=2.0, value=1.0, step=0.1, info=\"Adjust the speed\"\n )\n with gr.Column(scale=1, min_width=160):\n ref_text_file_input = gr.File(label=\"Load Reference Text from File (.txt)\", file_types=[\".txt\"])\n speech_type_rows.append(row)\n speech_type_names.append(name_input)\n speech_type_audios.append(audio_input)\n speech_type_ref_texts.append(ref_text_input)\n speech_type_ref_text_files.append(ref_text_file_input)\n speech_type_seeds.append(seed_input)\n speech_type_speeds.append(speed_input)\n speech_type_delete_btns.append(delete_btn)\n speech_type_insert_btns.append(insert_btn)\n\n # Global logic for all speech types\n for i in range(max_speech_types):\n speech_type_audios[i].clear(\n lambda: [None, None],\n None,\n [speech_type_ref_texts[i], speech_type_ref_text_files[i]],\n )\n speech_type_ref_text_files[i].upload(\n load_text_from_file,\n inputs=[speech_type_ref_text_files[i]],\n outputs=[speech_type_ref_texts[i]],\n )\n\n # Button to add speech type\n add_speech_type_btn = gr.Button(\"Add Speech Type\")\n\n # Keep track of autoincrement of speech types, no roll back\n speech_type_count = 1\n\n # Function to add a speech type\n def add_speech_type_fn():\n row_updates = [gr.update() for _ in range(max_speech_types)]\n global speech_type_count\n if speech_type_count < max_speech_types:\n row_updates[speech_type_count] = gr.update(visible=True)\n speech_type_count += 1\n else:\n gr.Warning(\"Exhausted maximum number of speech types. Consider restart the app.\")\n return row_updates\n\n add_speech_type_btn.click(add_speech_type_fn, outputs=speech_type_rows)\n\n # Function to delete a speech type\n def delete_speech_type_fn():\n return gr.update(visible=False), None, None, None, None\n\n # Update delete button clicks and ref text file changes\n for i in range(1, len(speech_type_delete_btns)):\n speech_type_delete_btns[i].click(\n delete_speech_type_fn,\n outputs=[\n speech_type_rows[i],\n speech_type_names[i],\n speech_type_audios[i],\n speech_type_ref_texts[i],\n speech_type_ref_text_files[i],\n ],\n )\n\n # Text input for the prompt\n with gr.Row():\n gen_text_input_multistyle = gr.Textbox(\n label=\"Text to Generate\",\n lines=10,\n max_lines=40,\n scale=4,\n 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?!\",\n )\n gen_text_file_multistyle = gr.File(label=\"Load Text to Generate from File (.txt)\", file_types=[\".txt\"], scale=1)\n\n def make_insert_speech_type_fn(index):\n def insert_speech_type_fn(current_text, speech_type_name, speech_type_seed, speech_type_speed):\n current_text = current_text or \"\"\n if not speech_type_name:\n gr.Warning(\"Please enter speech type name before insert.\")\n return current_text\n speech_type_dict = {\n \"name\": speech_type_name,\n \"seed\": speech_type_seed,\n \"speed\": speech_type_speed,\n }\n updated_text = current_text + json.dumps(speech_type_dict) + \" \"\n return updated_text\n\n return insert_speech_type_fn\n\n for i, insert_btn in enumerate(speech_type_insert_btns):\n insert_fn = make_insert_speech_type_fn(i)\n insert_btn.click(\n insert_fn,\n inputs=[gen_text_input_multistyle, speech_type_names[i], speech_type_seeds[i], speech_type_speeds[i]],\n outputs=gen_text_input_multistyle,\n )\n\n with gr.Accordion(\"Advanced Settings\", open=True):\n with gr.Row():\n with gr.Column():\n show_cherrypick_multistyle = gr.Checkbox(\n label=\"Show Cherry-pick Interface\",\n info=\"Turn on to show interface, picking seeds from previous generations.\",\n value=False,\n )\n with gr.Column():\n remove_silence_multistyle = gr.Checkbox(\n label=\"Remove Silences\",\n info=\"Turn on to automatically detect and crop long silences.\",\n value=True,\n )\n\n # Generate button\n generate_multistyle_btn = gr.Button(\"Generate Multi-Style Speech\", variant=\"primary\")\n\n # Output audio\n audio_output_multistyle = gr.Audio(label=\"Synthesized Audio\")\n\n # Used seed gallery\n cherrypick_interface_multistyle = gr.Textbox(\n label=\"Cherry-pick Interface\",\n lines=10,\n max_lines=40,\n buttons=[\"copy\"], # show_copy_button=True if gradio<6.0\n interactive=False,\n visible=False,\n )\n\n # Logic control to show/hide the cherrypick interface\n show_cherrypick_multistyle.change(\n lambda is_visible: gr.update(visible=is_visible),\n show_cherrypick_multistyle,\n cherrypick_interface_multistyle,\n )\n\n # Function to load text to generate from file\n gen_text_file_multistyle.upload(\n load_text_from_file,\n inputs=[gen_text_file_multistyle],\n outputs=[gen_text_input_multistyle],\n )\n\n @gpu_decorator\n def generate_multistyle_speech(\n gen_text,\n *args,\n ):\n speech_type_names_list = args[:max_speech_types]\n speech_type_audios_list = args[max_speech_types : 2 * max_speech_types]\n speech_type_ref_texts_list = args[2 * max_speech_types : 3 * max_speech_types]\n remove_silence = args[3 * max_speech_types]\n # Collect the speech types and their audios into a dict\n speech_types = OrderedDict()\n\n ref_text_idx = 0\n for name_input, audio_input, ref_text_input in zip(\n speech_type_names_list, speech_type_audios_list, speech_type_ref_texts_list\n ):\n if name_input and audio_input:\n speech_types[name_input] = {\"audio\": audio_input, \"ref_text\": ref_text_input}\n else:\n speech_types[f\"@{ref_text_idx}@\"] = {\"audio\": \"\", \"ref_text\": \"\"}\n ref_text_idx += 1\n\n # Parse the gen_text into segments\n segments = parse_speechtypes_text(gen_text)\n\n # For each segment, generate speech\n generated_audio_segments = []\n current_type_name = \"Regular\"\n inference_meta_data = \"\"\n\n for segment in segments:\n name = segment[\"name\"]\n seed_input = segment[\"seed\"]\n speed = segment[\"speed\"]\n text = segment[\"text\"]\n\n if name in speech_types:\n current_type_name = name\n else:\n gr.Warning(f\"Type {name} is not available, will use Regular as default.\")\n current_type_name = \"Regular\"\n\n try:\n ref_audio = speech_types[current_type_name][\"audio\"]\n except KeyError:\n gr.Warning(f\"Please provide reference audio for type {current_type_name}.\")\n return [None] + [speech_types[name][\"ref_text\"] for name in speech_types] + [None]\n ref_text = speech_types[current_type_name].get(\"ref_text\", \"\")\n\n if seed_input == -1:\n seed_input = np.random.randint(0, 2**31 - 1)\n\n # Generate or retrieve speech for this segment\n audio_out, _, ref_text_out, used_seed = infer(\n ref_audio,\n ref_text,\n text,\n tts_model_choice,\n remove_silence,\n seed=seed_input,\n cross_fade_duration=0,\n speed=speed,\n show_info=print, # no pull to top when generating\n )\n sr, audio_data = audio_out\n\n generated_audio_segments.append(audio_data)\n speech_types[current_type_name][\"ref_text\"] = ref_text_out\n inference_meta_data += json.dumps(dict(name=name, seed=used_seed, speed=speed)) + f\" {text}\\n\"\n\n # Concatenate all audio segments\n if generated_audio_segments:\n final_audio_data = np.concatenate(generated_audio_segments)\n return (\n [(sr, final_audio_data)]\n + [speech_types[name][\"ref_text\"] for name in speech_types]\n + [inference_meta_data]\n )\n else:\n gr.Warning(\"No audio generated.\")\n return [None] + [speech_types[name][\"ref_text\"] for name in speech_types] + [None]\n\n generate_multistyle_btn.click(\n generate_multistyle_speech,\n inputs=[\n gen_text_input_multistyle,\n ]\n + speech_type_names\n + speech_type_audios\n + speech_type_ref_texts\n + [\n remove_silence_multistyle,\n ],\n outputs=[audio_output_multistyle] + speech_type_ref_texts + [cherrypick_interface_multistyle],\n )\n\n # Validation function to disable Generate button if speech types are missing\n def validate_speech_types(gen_text, regular_name, *args):\n speech_type_names_list = args\n\n # Collect the speech types names\n speech_types_available = set()\n if regular_name:\n speech_types_available.add(regular_name)\n for name_input in speech_type_names_list:\n if name_input:\n speech_types_available.add(name_input)\n\n # Parse the gen_text to get the speech types used\n segments = parse_speechtypes_text(gen_text)\n speech_types_in_text = set(segment[\"name\"] for segment in segments)\n\n # Check if all speech types in text are available\n missing_speech_types = speech_types_in_text - speech_types_available\n\n if missing_speech_types:\n # Disable the generate button\n return gr.update(interactive=False)\n else:\n # Enable the generate button\n return gr.update(interactive=True)\n\n gen_text_input_multistyle.change(\n validate_speech_types,\n inputs=[gen_text_input_multistyle, regular_name] + speech_type_names,\n outputs=generate_multistyle_btn,\n )\n\n\nwith gr.Blocks() as app_chat:\n gr.Markdown(\n \"\"\"\n# Voice Chat\nHave a conversation with an AI using your reference voice!\n1. Upload a reference audio clip and optionally its transcript (via text or .txt file).\n2. Load the chat model.\n3. Record your message through your microphone or type it.\n4. The AI will respond using the reference voice.\n\"\"\"\n )\n\n chat_model_name_list = [\n \"Qwen/Qwen2.5-3B-Instruct\",\n \"microsoft/Phi-4-mini-instruct\",\n ]\n\n @gpu_decorator\n def load_chat_model(chat_model_name):\n show_info = gr.Info\n global chat_model_state, chat_tokenizer_state\n if chat_model_state is not None:\n chat_model_state = None\n chat_tokenizer_state = None\n gc.collect()\n torch.cuda.empty_cache()\n\n show_info(f\"Loading chat model: {chat_model_name}\")\n chat_model_state = AutoModelForCausalLM.from_pretrained(chat_model_name, torch_dtype=\"auto\", device_map=\"auto\")\n chat_tokenizer_state = AutoTokenizer.from_pretrained(chat_model_name)\n show_info(f\"Chat model {chat_model_name} loaded successfully!\")\n\n return gr.update(visible=False), gr.update(visible=True)\n\n if USING_SPACES:\n load_chat_model(chat_model_name_list[0])\n\n chat_model_name_input = gr.Dropdown(\n choices=chat_model_name_list,\n value=chat_model_name_list[0],\n label=\"Chat Model Name\",\n info=\"Enter the name of a HuggingFace chat model\",\n allow_custom_value=not USING_SPACES,\n )\n load_chat_model_btn = gr.Button(\"Load Chat Model\", variant=\"primary\", visible=not USING_SPACES)\n chat_interface_container = gr.Column(visible=USING_SPACES)\n\n chat_model_name_input.change(\n lambda: gr.update(visible=True),\n None,\n load_chat_model_btn,\n show_progress=\"hidden\",\n )\n load_chat_model_btn.click(\n load_chat_model, inputs=[chat_model_name_input], outputs=[load_chat_model_btn, chat_interface_container]\n )\n\n with chat_interface_container:\n with gr.Row():\n with gr.Column():\n ref_audio_chat = gr.Audio(label=\"Reference Audio\", type=\"filepath\")\n with gr.Column():\n with gr.Accordion(\"Advanced Settings\", open=False):\n with gr.Row():\n ref_text_chat = gr.Textbox(\n label=\"Reference Text\",\n info=\"Optional: Leave blank to auto-transcribe\",\n lines=2,\n scale=3,\n )\n ref_text_file_chat = gr.File(\n label=\"Load Reference Text from File (.txt)\", file_types=[\".txt\"], scale=1\n )\n with gr.Row():\n randomize_seed_chat = gr.Checkbox(\n label=\"Randomize Seed\",\n value=True,\n info=\"Uncheck to use the seed specified.\",\n scale=3,\n )\n seed_input_chat = gr.Number(show_label=False, value=0, precision=0, scale=1)\n remove_silence_chat = gr.Checkbox(\n label=\"Remove Silences\",\n value=True,\n )\n system_prompt_chat = gr.Textbox(\n label=\"System Prompt\",\n 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.\",\n lines=2,\n )\n\n chatbot_interface = gr.Chatbot(\n label=\"Conversation\"\n ) # type=\"messages\" hard-coded and no need to pass in since gradio 6.0\n\n with gr.Row():\n with gr.Column():\n audio_input_chat = gr.Microphone(\n label=\"Speak your message\",\n type=\"filepath\",\n )\n audio_output_chat = gr.Audio(autoplay=True)\n with gr.Column():\n text_input_chat = gr.Textbox(\n label=\"Type your message\",\n lines=1,\n )\n send_btn_chat = gr.Button(\"Send Message\")\n clear_btn_chat = gr.Button(\"Clear Conversation\")\n\n # Modify process_audio_input to generate user input\n @gpu_decorator\n def process_audio_input(conv_state, audio_path, text):\n \"\"\"Handle audio or text input from user\"\"\"\n\n if not audio_path and not text.strip():\n return conv_state\n\n if audio_path:\n text = preprocess_ref_audio_text(audio_path, text)[1]\n if not text.strip():\n return conv_state\n\n conv_state.append({\"role\": \"user\", \"content\": text})\n return conv_state\n\n # Use model and tokenizer from state to get text response\n @gpu_decorator\n def generate_text_response(conv_state, system_prompt):\n \"\"\"Generate text response from AI\"\"\"\n for single_state in conv_state:\n if isinstance(single_state[\"content\"], list):\n assert len(single_state[\"content\"]) == 1 and single_state[\"content\"][0][\"type\"] == \"text\"\n single_state[\"content\"] = single_state[\"content\"][0][\"text\"]\n\n system_prompt_state = [{\"role\": \"system\", \"content\": system_prompt}]\n response = chat_model_inference(system_prompt_state + conv_state, chat_model_state, chat_tokenizer_state)\n\n conv_state.append({\"role\": \"assistant\", \"content\": response})\n return conv_state\n\n @gpu_decorator\n def generate_audio_response(conv_state, ref_audio, ref_text, remove_silence, randomize_seed, seed_input):\n \"\"\"Generate TTS audio for AI response\"\"\"\n if not conv_state or not ref_audio:\n return None, ref_text, seed_input\n\n last_ai_response = conv_state[-1][\"content\"][0][\"text\"]\n if not last_ai_response or conv_state[-1][\"role\"] != \"assistant\":\n return None, ref_text, seed_input\n\n if randomize_seed:\n seed_input = np.random.randint(0, 2**31 - 1)\n\n audio_result, _, ref_text_out, used_seed = infer(\n ref_audio,\n ref_text,\n last_ai_response,\n tts_model_choice,\n remove_silence,\n seed=seed_input,\n cross_fade_duration=0.15,\n speed=1.0,\n show_info=print, # show_info=print no pull to top when generating\n )\n return audio_result, ref_text_out, used_seed\n\n def clear_conversation():\n \"\"\"Reset the conversation\"\"\"\n return [], None\n\n ref_text_file_chat.upload(\n load_text_from_file,\n inputs=[ref_text_file_chat],\n outputs=[ref_text_chat],\n )\n\n for user_operation in [audio_input_chat.stop_recording, text_input_chat.submit, send_btn_chat.click]:\n user_operation(\n process_audio_input,\n inputs=[chatbot_interface, audio_input_chat, text_input_chat],\n outputs=[chatbot_interface],\n ).then(\n generate_text_response,\n inputs=[chatbot_interface, system_prompt_chat],\n outputs=[chatbot_interface],\n ).then(\n generate_audio_response,\n inputs=[\n chatbot_interface,\n ref_audio_chat,\n ref_text_chat,\n remove_silence_chat,\n randomize_seed_chat,\n seed_input_chat,\n ],\n outputs=[audio_output_chat, ref_text_chat, seed_input_chat],\n ).then(\n lambda: [None, None],\n None,\n [audio_input_chat, text_input_chat],\n )\n\n # Handle clear button or system prompt change and reset conversation\n for user_operation in [clear_btn_chat.click, system_prompt_chat.change, chatbot_interface.clear]:\n user_operation(\n clear_conversation,\n outputs=[chatbot_interface, audio_output_chat],\n )\n\n\nwith gr.Blocks() as app_credits:\n gr.Markdown(\"\"\"\n# Credits\n\n* [mrfakename](https://github.com/fakerybakery) for the original [online demo](https://huggingface.co/spaces/mrfakename/E2-F5-TTS)\n* [RootingInLoad](https://github.com/RootingInLoad) for initial chunk generation and podcast app exploration\n* [jpgallegoar](https://github.com/jpgallegoar) for multiple speech-type generation & voice chat\n\"\"\")\n\n\nwith gr.Blocks() as app:\n gr.Markdown(\n f\"\"\"\n# F5-TTS Demo Space\n\nThis 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:\n\n* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)\n* [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)\n\nThe checkpoints currently support English and Chinese.\n\nIf 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).\n\n**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.**\n\"\"\"\n )\n\n last_used_custom = files(\"f5_tts\").joinpath(\"infer/.cache/last_used_custom_model_info_v1.txt\")\n\n def load_last_used_custom():\n try:\n custom = []\n with open(last_used_custom, \"r\", encoding=\"utf-8\") as f:\n for line in f:\n custom.append(line.strip())\n return custom\n except FileNotFoundError:\n last_used_custom.parent.mkdir(parents=True, exist_ok=True)\n return DEFAULT_TTS_MODEL_CFG\n\n def switch_tts_model(new_choice):\n global tts_model_choice\n if new_choice == \"Custom\": # override in case webpage is refreshed\n custom_ckpt_path, custom_vocab_path, custom_model_cfg = load_last_used_custom()\n tts_model_choice = (\"Custom\", custom_ckpt_path, custom_vocab_path, custom_model_cfg)\n return (\n gr.update(visible=True, value=custom_ckpt_path),\n gr.update(visible=True, value=custom_vocab_path),\n gr.update(visible=True, value=custom_model_cfg),\n )\n else:\n tts_model_choice = new_choice\n return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)\n\n def set_custom_model(custom_ckpt_path, custom_vocab_path, custom_model_cfg):\n global tts_model_choice\n tts_model_choice = (\"Custom\", custom_ckpt_path, custom_vocab_path, custom_model_cfg)\n with open(last_used_custom, \"w\", encoding=\"utf-8\") as f:\n f.write(custom_ckpt_path + \"\\n\" + custom_vocab_path + \"\\n\" + custom_model_cfg + \"\\n\")\n\n with gr.Row():\n if not USING_SPACES:\n choose_tts_model = gr.Radio(\n choices=[DEFAULT_TTS_MODEL, \"E2-TTS\", \"Custom\"], label=\"Choose TTS Model\", value=DEFAULT_TTS_MODEL\n )\n else:\n choose_tts_model = gr.Radio(\n choices=[DEFAULT_TTS_MODEL, \"E2-TTS\"], label=\"Choose TTS Model\", value=DEFAULT_TTS_MODEL\n )\n custom_ckpt_path = gr.Dropdown(\n choices=[DEFAULT_TTS_MODEL_CFG[0]],\n value=load_last_used_custom()[0],\n allow_custom_value=True,\n label=\"Model: local_path | hf://user_id/repo_id/model_ckpt\",\n visible=False,\n )\n custom_vocab_path = gr.Dropdown(\n choices=[DEFAULT_TTS_MODEL_CFG[1]],\n value=load_last_used_custom()[1],\n allow_custom_value=True,\n label=\"Vocab: local_path | hf://user_id/repo_id/vocab_file\",\n visible=False,\n )\n custom_model_cfg = gr.Dropdown(\n choices=[\n DEFAULT_TTS_MODEL_CFG[2],\n json.dumps(\n dict(\n dim=1024,\n depth=22,\n heads=16,\n ff_mult=2,\n text_dim=512,\n text_mask_padding=False,\n conv_layers=4,\n pe_attn_head=1,\n )\n ),\n json.dumps(\n dict(\n dim=768,\n depth=18,\n heads=12,\n ff_mult=2,\n text_dim=512,\n text_mask_padding=False,\n conv_layers=4,\n pe_attn_head=1,\n )\n ),\n ],\n value=load_last_used_custom()[2],\n allow_custom_value=True,\n label=\"Config: in a dictionary form\",\n visible=False,\n )\n\n choose_tts_model.change(\n switch_tts_model,\n inputs=[choose_tts_model],\n outputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],\n show_progress=\"hidden\",\n )\n custom_ckpt_path.change(\n set_custom_model,\n inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],\n show_progress=\"hidden\",\n )\n custom_vocab_path.change(\n set_custom_model,\n inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],\n show_progress=\"hidden\",\n )\n custom_model_cfg.change(\n set_custom_model,\n inputs=[custom_ckpt_path, custom_vocab_path, custom_model_cfg],\n show_progress=\"hidden\",\n )\n\n gr.TabbedInterface(\n [app_tts, app_multistyle, app_chat, app_credits],\n [\"Basic-TTS\", \"Multi-Speech\", \"Voice-Chat\", \"Credits\"],\n )\n\n\n@click.command()\n@click.option(\"--port\", \"-p\", default=None, type=int, help=\"Port to run the app on\")\n@click.option(\"--host\", \"-H\", default=None, help=\"Host to run the app on\")\n@click.option(\n \"--share\",\n \"-s\",\n default=False,\n is_flag=True,\n help=\"Share the app via Gradio share link\",\n)\n@click.option(\"--api\", \"-a\", default=True, is_flag=True, help=\"Allow API access\")\n@click.option(\n \"--root_path\",\n \"-r\",\n default=None,\n type=str,\n 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\".',\n)\n@click.option(\n \"--inbrowser\",\n \"-i\",\n is_flag=True,\n default=False,\n help=\"Automatically launch the interface in the default web browser\",\n)\ndef main(port, host, share, api, root_path, inbrowser):\n global app\n print(\"Starting app...\")\n app.queue(api_open=api).launch(\n server_name=host,\n server_port=port,\n share=share,\n root_path=root_path,\n inbrowser=inbrowser,\n )\n\n\nif __name__ == \"__main__\":\n if not USING_SPACES:\n main()\n else:\n app.queue().launch()\n" }, { "path": "src/f5_tts/infer/speech_edit.py", "content": "import os\n\n\nos.environ[\"PYTORCH_ENABLE_MPS_FALLBACK\"] = \"1\" # for MPS device compatibility\n\nfrom importlib.resources import files\n\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nfrom cached_path import cached_path\nfrom hydra.utils import get_class\nfrom omegaconf import OmegaConf\n\nfrom f5_tts.infer.utils_infer import load_checkpoint, load_vocoder, save_spectrogram\nfrom f5_tts.model import CFM\nfrom f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer\n\n\ndevice = (\n \"cuda\"\n if torch.cuda.is_available()\n else \"xpu\"\n if torch.xpu.is_available()\n else \"mps\"\n if torch.backends.mps.is_available()\n else \"cpu\"\n)\n\n\n# ---------------------- infer setting ---------------------- #\n\nseed = None # int | None\n\nexp_name = \"F5TTS_v1_Base\" # F5TTS_v1_Base | E2TTS_Base\nckpt_step = 1250000\n\nnfe_step = 32 # 16, 32\ncfg_strength = 2.0\node_method = \"euler\" # euler | midpoint\nsway_sampling_coef = -1.0\nspeed = 1.0\ntarget_rms = 0.1\n\n\nmodel_cfg = OmegaConf.load(str(files(\"f5_tts\").joinpath(f\"configs/{exp_name}.yaml\")))\nmodel_cls = get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\nmodel_arc = model_cfg.model.arch\n\ndataset_name = model_cfg.datasets.name\ntokenizer = model_cfg.model.tokenizer\n\nmel_spec_type = model_cfg.model.mel_spec.mel_spec_type\ntarget_sample_rate = model_cfg.model.mel_spec.target_sample_rate\nn_mel_channels = model_cfg.model.mel_spec.n_mel_channels\nhop_length = model_cfg.model.mel_spec.hop_length\nwin_length = model_cfg.model.mel_spec.win_length\nn_fft = model_cfg.model.mel_spec.n_fft\n\n\n# ckpt_path = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/ckpts/{exp_name}/model_{ckpt_step}.safetensors\"\nckpt_path = str(cached_path(f\"hf://SWivid/F5-TTS/{exp_name}/model_{ckpt_step}.safetensors\"))\noutput_dir = \"tests\"\n\n\n# [leverage https://github.com/MahmoudAshraf97/ctc-forced-aligner to get char level alignment]\n# pip install git+https://github.com/MahmoudAshraf97/ctc-forced-aligner.git\n# [write the origin_text into a file, e.g. tests/test_edit.txt]\n# 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\"\n# [result will be saved at same path of audio file]\n# [--language \"zho\" for Chinese, \"eng\" for English]\n# [if local ckpt, set --alignment_model \"../checkpoints/mms-300m-1130-forced-aligner\"]\n\naudio_to_edit = str(files(\"f5_tts\").joinpath(\"infer/examples/basic/basic_ref_en.wav\"))\norigin_text = \"Some call me nature, others call me mother nature.\"\ntarget_text = \"Some call me optimist, others call me realist.\"\nparts_to_edit = [\n [1.42, 2.44],\n [4.04, 4.9],\n] # stard_ends of \"nature\" & \"mother nature\", in seconds\nfix_duration = [\n 1.2,\n 1,\n] # fix duration for \"optimist\" & \"realist\", in seconds\n\n# audio_to_edit = \"src/f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n# origin_text = \"对，这就是我，万人敬仰的太乙真人。\"\n# target_text = \"对，那就是你，万人敬仰的太白金星。\"\n# parts_to_edit = [[0.84, 1.4], [1.92, 2.4], [4.26, 6.26], ]\n# fix_duration = None # use origin text duration\n\n# audio_to_edit = \"src/f5_tts/infer/examples/basic/basic_ref_zh.wav\"\n# origin_text = \"对，这就是我，万人敬仰的太乙真人。\"\n# target_text = \"对，这就是你，万人敬仰的李白金星。\"\n# parts_to_edit = [[1.500, 2.784], [4.083, 6.760]]\n# fix_duration = [1.284, 2.677]\n\n\n# -------------------------------------------------#\n\nuse_ema = True\n\nif not os.path.exists(output_dir):\n os.makedirs(output_dir)\n\n# Vocoder model\nlocal = False\nif mel_spec_type == \"vocos\":\n vocoder_local_path = \"../checkpoints/charactr/vocos-mel-24khz\"\nelif mel_spec_type == \"bigvgan\":\n vocoder_local_path = \"../checkpoints/bigvgan_v2_24khz_100band_256x\"\nvocoder = load_vocoder(vocoder_name=mel_spec_type, is_local=local, local_path=vocoder_local_path)\n\n# Tokenizer\nvocab_char_map, vocab_size = get_tokenizer(dataset_name, tokenizer)\n\n# Model\nmodel = CFM(\n transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=n_mel_channels),\n mel_spec_kwargs=dict(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n ),\n odeint_kwargs=dict(\n method=ode_method,\n ),\n vocab_char_map=vocab_char_map,\n).to(device)\n\ndtype = torch.float32 if mel_spec_type == \"bigvgan\" else None\nmodel = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)\n\n# Audio\naudio, sr = torchaudio.load(audio_to_edit)\nif audio.shape[0] > 1:\n audio = torch.mean(audio, dim=0, keepdim=True)\nrms = torch.sqrt(torch.mean(torch.square(audio)))\nif rms < target_rms:\n audio = audio * target_rms / rms\nif sr != target_sample_rate:\n resampler = torchaudio.transforms.Resample(sr, target_sample_rate)\n audio = resampler(audio)\n\n# Convert to mel spectrogram FIRST (on clean original audio)\n# This avoids boundary artifacts from mel windows straddling zeros and real audio\naudio = audio.to(device)\nwith torch.inference_mode():\n original_mel = model.mel_spec(audio) # (batch, n_mel, n_frames)\n original_mel = original_mel.permute(0, 2, 1) # (batch, n_frames, n_mel)\n\n# Build mel_cond and edit_mask at FRAME level\n# Insert zero frames in mel domain instead of zero samples in wav domain\noffset_frame = 0\nmel_cond = torch.zeros(1, 0, n_mel_channels, device=device)\nedit_mask = torch.zeros(1, 0, dtype=torch.bool, device=device)\nfix_dur_list = fix_duration.copy() if fix_duration is not None else None\n\nfor part in parts_to_edit:\n start, end = part\n part_dur_sec = end - start if fix_dur_list is None else fix_dur_list.pop(0)\n\n # Convert to frames (this is the authoritative unit)\n start_frame = round(start * target_sample_rate / hop_length)\n end_frame = round(end * target_sample_rate / hop_length)\n part_dur_frames = round(part_dur_sec * target_sample_rate / hop_length)\n\n # Number of frames for the kept (non-edited) region\n keep_frames = start_frame - offset_frame\n\n # Build mel_cond: original mel frames + zero frames for edit region\n mel_cond = torch.cat(\n (\n mel_cond,\n original_mel[:, offset_frame:start_frame, :],\n torch.zeros(1, part_dur_frames, n_mel_channels, device=device),\n ),\n dim=1,\n )\n edit_mask = torch.cat(\n (\n edit_mask,\n torch.ones(1, keep_frames, dtype=torch.bool, device=device),\n torch.zeros(1, part_dur_frames, dtype=torch.bool, device=device),\n ),\n dim=-1,\n )\n offset_frame = end_frame\n\n# Append remaining mel frames after last edit\nmel_cond = torch.cat((mel_cond, original_mel[:, offset_frame:, :]), dim=1)\nedit_mask = F.pad(edit_mask, (0, mel_cond.shape[1] - edit_mask.shape[-1]), value=True)\n\n# Text\ntext_list = [target_text]\nif tokenizer == \"pinyin\":\n final_text_list = convert_char_to_pinyin(text_list)\nelse:\n final_text_list = [text_list]\nprint(f\"text : {text_list}\")\nprint(f\"pinyin: {final_text_list}\")\n\n# Duration - use mel_cond length (not raw audio length)\nduration = mel_cond.shape[1]\n\n# Inference - pass mel_cond directly (not wav)\nwith torch.inference_mode():\n generated, trajectory = model.sample(\n cond=mel_cond, # Now passing mel directly, not wav\n text=final_text_list,\n duration=duration,\n steps=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n seed=seed,\n edit_mask=edit_mask,\n )\n print(f\"Generated mel: {generated.shape}\")\n\n # Final result\n generated = generated.to(torch.float32)\n gen_mel_spec = generated.permute(0, 2, 1)\n if mel_spec_type == \"vocos\":\n generated_wave = vocoder.decode(gen_mel_spec).cpu()\n elif mel_spec_type == \"bigvgan\":\n generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()\n\n if rms < target_rms:\n generated_wave = generated_wave * rms / target_rms\n\n save_spectrogram(gen_mel_spec[0].cpu().numpy(), f\"{output_dir}/speech_edit_out.png\")\n torchaudio.save(f\"{output_dir}/speech_edit_out.wav\", generated_wave, target_sample_rate)\n print(f\"Generated wav: {generated_wave.shape}\")\n" }, { "path": "src/f5_tts/infer/utils_infer.py", "content": "# A unified script for inference process\n# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format\nimport os\nimport sys\nfrom concurrent.futures import ThreadPoolExecutor\n\n\nos.environ[\"PYTORCH_ENABLE_MPS_FALLBACK\"] = \"1\" # for MPS device compatibility\nsys.path.append(f\"{os.path.dirname(os.path.abspath(__file__))}/../../third_party/BigVGAN/\")\n\nimport hashlib\nimport re\nimport tempfile\nfrom importlib.resources import files\n\nimport matplotlib\n\n\nmatplotlib.use(\"Agg\")\n\nimport matplotlib.pylab as plt\nimport numpy as np\nimport torch\nimport torchaudio\nimport tqdm\nfrom huggingface_hub import hf_hub_download\nfrom pydub import AudioSegment, silence\nfrom transformers import pipeline\nfrom vocos import Vocos\n\nfrom f5_tts.model import CFM\nfrom f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer\n\n\n_ref_audio_cache = {}\n_ref_text_cache = {}\n\ndevice = (\n \"cuda\"\n if torch.cuda.is_available()\n else \"xpu\"\n if torch.xpu.is_available()\n else \"mps\"\n if torch.backends.mps.is_available()\n else \"cpu\"\n)\n\ntempfile_kwargs = {\"delete_on_close\": False} if sys.version_info >= (3, 12) else {\"delete\": False}\n\n# -----------------------------------------\n\ntarget_sample_rate = 24000\nn_mel_channels = 100\nhop_length = 256\nwin_length = 1024\nn_fft = 1024\nmel_spec_type = \"vocos\"\ntarget_rms = 0.1\ncross_fade_duration = 0.15\node_method = \"euler\"\nnfe_step = 32 # 16, 32\ncfg_strength = 2.0\nsway_sampling_coef = -1.0\nspeed = 1.0\nfix_duration = None\n\n# -----------------------------------------\n\n\n# chunk text into smaller pieces\n\n\ndef chunk_text(text, max_chars=135):\n \"\"\"\n Splits the input text into chunks, each with a maximum number of characters.\n\n Args:\n text (str): The text to be split.\n max_chars (int): The maximum number of characters per chunk.\n\n Returns:\n List[str]: A list of text chunks.\n \"\"\"\n chunks = []\n current_chunk = \"\"\n # Split the text into sentences based on punctuation followed by whitespace\n sentences = re.split(r\"(?<=[;:,.!?])\\s+|(?<=[；：，。！？])\", text)\n\n for sentence in sentences:\n if len(current_chunk.encode(\"utf-8\")) + len(sentence.encode(\"utf-8\")) <= max_chars:\n current_chunk += sentence + \" \" if sentence and len(sentence[-1].encode(\"utf-8\")) == 1 else sentence\n else:\n if current_chunk:\n chunks.append(current_chunk.strip())\n current_chunk = sentence + \" \" if sentence and len(sentence[-1].encode(\"utf-8\")) == 1 else sentence\n\n if current_chunk:\n chunks.append(current_chunk.strip())\n\n return chunks\n\n\n# load vocoder\ndef load_vocoder(vocoder_name=\"vocos\", is_local=False, local_path=\"\", device=device, hf_cache_dir=None):\n if vocoder_name == \"vocos\":\n # vocoder = Vocos.from_pretrained(\"charactr/vocos-mel-24khz\").to(device)\n if is_local:\n print(f\"Load vocos from local path {local_path}\")\n config_path = f\"{local_path}/config.yaml\"\n model_path = f\"{local_path}/pytorch_model.bin\"\n else:\n print(\"Download Vocos from huggingface charactr/vocos-mel-24khz\")\n repo_id = \"charactr/vocos-mel-24khz\"\n config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"config.yaml\")\n model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"pytorch_model.bin\")\n vocoder = Vocos.from_hparams(config_path)\n state_dict = torch.load(model_path, map_location=\"cpu\", weights_only=True)\n from vocos.feature_extractors import EncodecFeatures\n\n if isinstance(vocoder.feature_extractor, EncodecFeatures):\n encodec_parameters = {\n \"feature_extractor.encodec.\" + key: value\n for key, value in vocoder.feature_extractor.encodec.state_dict().items()\n }\n state_dict.update(encodec_parameters)\n vocoder.load_state_dict(state_dict)\n vocoder = vocoder.eval().to(device)\n elif vocoder_name == \"bigvgan\":\n try:\n from third_party.BigVGAN import bigvgan\n except ImportError:\n print(\"You need to follow the README to init submodule and change the BigVGAN source code.\")\n if is_local:\n # download generator from https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x/tree/main\n vocoder = bigvgan.BigVGAN.from_pretrained(local_path, use_cuda_kernel=False)\n else:\n vocoder = bigvgan.BigVGAN.from_pretrained(\n \"nvidia/bigvgan_v2_24khz_100band_256x\", use_cuda_kernel=False, cache_dir=hf_cache_dir\n )\n\n vocoder.remove_weight_norm()\n vocoder = vocoder.eval().to(device)\n return vocoder\n\n\n# load asr pipeline\n\nasr_pipe = None\n\n\ndef initialize_asr_pipeline(device: str = device, dtype=None):\n if dtype is None:\n dtype = (\n torch.float16\n if \"cuda\" in device\n and torch.cuda.get_device_properties(device).major >= 7\n and not torch.cuda.get_device_name().endswith(\"[ZLUDA]\")\n else torch.float32\n )\n global asr_pipe\n asr_pipe = pipeline(\n \"automatic-speech-recognition\",\n model=\"openai/whisper-large-v3-turbo\",\n torch_dtype=dtype,\n device=device,\n )\n\n\n# transcribe\n\n\ndef transcribe(ref_audio, language=None):\n global asr_pipe\n if asr_pipe is None:\n initialize_asr_pipeline(device=device)\n return asr_pipe(\n ref_audio,\n chunk_length_s=30,\n batch_size=128,\n generate_kwargs={\"task\": \"transcribe\", \"language\": language} if language else {\"task\": \"transcribe\"},\n return_timestamps=False,\n )[\"text\"].strip()\n\n\n# load model checkpoint for inference\n\n\ndef load_checkpoint(model, ckpt_path, device: str, dtype=None, use_ema=True):\n if dtype is None:\n dtype = (\n torch.float16\n if \"cuda\" in device\n and torch.cuda.get_device_properties(device).major >= 7\n and not torch.cuda.get_device_name().endswith(\"[ZLUDA]\")\n else torch.float32\n )\n model = model.to(dtype)\n\n ckpt_type = ckpt_path.split(\".\")[-1]\n if ckpt_type == \"safetensors\":\n from safetensors.torch import load_file\n\n checkpoint = load_file(ckpt_path, device=device)\n else:\n checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)\n\n if use_ema:\n if ckpt_type == \"safetensors\":\n checkpoint = {\"ema_model_state_dict\": checkpoint}\n checkpoint[\"model_state_dict\"] = {\n k.replace(\"ema_model.\", \"\"): v\n for k, v in checkpoint[\"ema_model_state_dict\"].items()\n if k not in [\"initted\", \"step\"]\n }\n\n # patch for backward compatibility, 305e3ea\n for key in [\"mel_spec.mel_stft.mel_scale.fb\", \"mel_spec.mel_stft.spectrogram.window\"]:\n if key in checkpoint[\"model_state_dict\"]:\n del checkpoint[\"model_state_dict\"][key]\n\n model.load_state_dict(checkpoint[\"model_state_dict\"])\n else:\n if ckpt_type == \"safetensors\":\n checkpoint = {\"model_state_dict\": checkpoint}\n model.load_state_dict(checkpoint[\"model_state_dict\"])\n\n del checkpoint\n torch.cuda.empty_cache()\n\n return model.to(device)\n\n\n# load model for inference\n\n\ndef load_model(\n model_cls,\n model_cfg,\n ckpt_path,\n mel_spec_type=mel_spec_type,\n vocab_file=\"\",\n ode_method=ode_method,\n use_ema=True,\n device=device,\n):\n if vocab_file == \"\":\n vocab_file = str(files(\"f5_tts\").joinpath(\"infer/examples/vocab.txt\"))\n tokenizer = \"custom\"\n\n print(\"\\nvocab : \", vocab_file)\n print(\"token : \", tokenizer)\n print(\"model : \", ckpt_path, \"\\n\")\n\n vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer)\n model = CFM(\n transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),\n mel_spec_kwargs=dict(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n ),\n odeint_kwargs=dict(\n method=ode_method,\n ),\n vocab_char_map=vocab_char_map,\n ).to(device)\n\n dtype = torch.float32 if mel_spec_type == \"bigvgan\" else None\n model = load_checkpoint(model, ckpt_path, device, dtype=dtype, use_ema=use_ema)\n\n return model\n\n\ndef remove_silence_edges(audio, silence_threshold=-42):\n # Remove silence from the start\n non_silent_start_idx = silence.detect_leading_silence(audio, silence_threshold=silence_threshold)\n audio = audio[non_silent_start_idx:]\n\n # Remove silence from the end\n non_silent_end_duration = audio.duration_seconds\n for ms in reversed(audio):\n if ms.dBFS > silence_threshold:\n break\n non_silent_end_duration -= 0.001\n trimmed_audio = audio[: int(non_silent_end_duration * 1000)]\n\n return trimmed_audio\n\n\n# preprocess reference audio and text\n\n\ndef preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print):\n show_info(\"Converting audio...\")\n\n # Compute a hash of the reference audio file\n with open(ref_audio_orig, \"rb\") as audio_file:\n audio_data = audio_file.read()\n audio_hash = hashlib.md5(audio_data).hexdigest()\n\n global _ref_audio_cache\n\n if audio_hash in _ref_audio_cache:\n show_info(\"Using cached preprocessed reference audio...\")\n ref_audio = _ref_audio_cache[audio_hash]\n\n else: # first pass, do preprocess\n with tempfile.NamedTemporaryFile(suffix=\".wav\", **tempfile_kwargs) as f:\n temp_path = f.name\n\n aseg = AudioSegment.from_file(ref_audio_orig)\n\n # 1. try to find long silence for clipping\n non_silent_segs = silence.split_on_silence(\n aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000, seek_step=10\n )\n non_silent_wave = AudioSegment.silent(duration=0)\n for non_silent_seg in non_silent_segs:\n if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:\n show_info(\"Audio is over 12s, clipping short. (1)\")\n break\n non_silent_wave += non_silent_seg\n\n # 2. try to find short silence for clipping if 1. failed\n if len(non_silent_wave) > 12000:\n non_silent_segs = silence.split_on_silence(\n aseg, min_silence_len=100, silence_thresh=-40, keep_silence=1000, seek_step=10\n )\n non_silent_wave = AudioSegment.silent(duration=0)\n for non_silent_seg in non_silent_segs:\n if len(non_silent_wave) > 6000 and len(non_silent_wave + non_silent_seg) > 12000:\n show_info(\"Audio is over 12s, clipping short. (2)\")\n break\n non_silent_wave += non_silent_seg\n\n aseg = non_silent_wave\n\n # 3. if no proper silence found for clipping\n if len(aseg) > 12000:\n aseg = aseg[:12000]\n show_info(\"Audio is over 12s, clipping short. (3)\")\n\n aseg = remove_silence_edges(aseg) + AudioSegment.silent(duration=50)\n aseg.export(temp_path, format=\"wav\")\n ref_audio = temp_path\n\n # Cache the processed reference audio\n _ref_audio_cache[audio_hash] = ref_audio\n\n if not ref_text.strip():\n global _ref_text_cache\n if audio_hash in _ref_text_cache:\n # Use cached asr transcription\n show_info(\"Using cached reference text...\")\n ref_text = _ref_text_cache[audio_hash]\n else:\n show_info(\"No reference text provided, transcribing reference audio...\")\n ref_text = transcribe(ref_audio)\n # Cache the transcribed text (not caching custom ref_text, enabling users to do manual tweak)\n _ref_text_cache[audio_hash] = ref_text\n else:\n show_info(\"Using custom reference text...\")\n\n # Ensure ref_text ends with a proper sentence-ending punctuation\n if not ref_text.endswith(\". \") and not ref_text.endswith(\"。\"):\n if ref_text.endswith(\".\"):\n ref_text += \" \"\n else:\n ref_text += \". \"\n\n print(\"\\nref_text \", ref_text)\n\n return ref_audio, ref_text\n\n\n# infer process: chunk text -> infer batches [i.e. infer_batch_process()]\n\n\ndef infer_process(\n ref_audio,\n ref_text,\n gen_text,\n model_obj,\n vocoder,\n mel_spec_type=mel_spec_type,\n show_info=print,\n progress=tqdm,\n target_rms=target_rms,\n cross_fade_duration=cross_fade_duration,\n nfe_step=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n speed=speed,\n fix_duration=fix_duration,\n device=device,\n):\n # Split the input text into batches\n audio, sr = torchaudio.load(ref_audio)\n max_chars = int(len(ref_text.encode(\"utf-8\")) / (audio.shape[-1] / sr) * (22 - audio.shape[-1] / sr) * speed)\n gen_text_batches = chunk_text(gen_text, max_chars=max_chars)\n for i, gen_text in enumerate(gen_text_batches):\n print(f\"gen_text {i}\", gen_text)\n print(\"\\n\")\n\n show_info(f\"Generating audio in {len(gen_text_batches)} batches...\")\n return next(\n infer_batch_process(\n (audio, sr),\n ref_text,\n gen_text_batches,\n model_obj,\n vocoder,\n mel_spec_type=mel_spec_type,\n progress=progress,\n target_rms=target_rms,\n cross_fade_duration=cross_fade_duration,\n nfe_step=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n speed=speed,\n fix_duration=fix_duration,\n device=device,\n )\n )\n\n\n# infer batches\n\n\ndef infer_batch_process(\n ref_audio,\n ref_text,\n gen_text_batches,\n model_obj,\n vocoder,\n mel_spec_type=\"vocos\",\n progress=tqdm,\n target_rms=0.1,\n cross_fade_duration=0.15,\n nfe_step=32,\n cfg_strength=2.0,\n sway_sampling_coef=-1,\n speed=1,\n fix_duration=None,\n device=None,\n streaming=False,\n chunk_size=2048,\n):\n audio, sr = ref_audio\n if audio.shape[0] > 1:\n audio = torch.mean(audio, dim=0, keepdim=True)\n\n rms = torch.sqrt(torch.mean(torch.square(audio)))\n if rms < target_rms:\n audio = audio * target_rms / rms\n if sr != target_sample_rate:\n resampler = torchaudio.transforms.Resample(sr, target_sample_rate)\n audio = resampler(audio)\n audio = audio.to(device)\n\n generated_waves = []\n spectrograms = []\n\n if len(ref_text[-1].encode(\"utf-8\")) == 1:\n ref_text = ref_text + \" \"\n\n def process_batch(gen_text):\n local_speed = speed\n if len(gen_text.encode(\"utf-8\")) < 10:\n local_speed = 0.3\n\n # Prepare the text\n text_list = [ref_text + gen_text]\n final_text_list = convert_char_to_pinyin(text_list)\n\n ref_audio_len = audio.shape[-1] // hop_length\n if fix_duration is not None:\n duration = int(fix_duration * target_sample_rate / hop_length)\n else:\n # Calculate duration\n ref_text_len = len(ref_text.encode(\"utf-8\"))\n gen_text_len = len(gen_text.encode(\"utf-8\"))\n duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / local_speed)\n\n # inference\n with torch.inference_mode():\n generated, _ = model_obj.sample(\n cond=audio,\n text=final_text_list,\n duration=duration,\n steps=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n )\n del _\n\n generated = generated.to(torch.float32) # generated mel spectrogram\n generated = generated[:, ref_audio_len:, :]\n generated = generated.permute(0, 2, 1)\n if mel_spec_type == \"vocos\":\n generated_wave = vocoder.decode(generated)\n elif mel_spec_type == \"bigvgan\":\n generated_wave = vocoder(generated)\n if rms < target_rms:\n generated_wave = generated_wave * rms / target_rms\n\n # wav -> numpy\n generated_wave = generated_wave.squeeze().cpu().numpy()\n\n if streaming:\n for j in range(0, len(generated_wave), chunk_size):\n yield generated_wave[j : j + chunk_size], target_sample_rate\n else:\n generated_cpu = generated[0].cpu().numpy()\n del generated\n yield generated_wave, generated_cpu\n\n if streaming:\n for gen_text in progress.tqdm(gen_text_batches) if progress is not None else gen_text_batches:\n for chunk in process_batch(gen_text):\n yield chunk\n else:\n with ThreadPoolExecutor() as executor:\n futures = [executor.submit(process_batch, gen_text) for gen_text in gen_text_batches]\n for future in progress.tqdm(futures) if progress is not None else futures:\n result = future.result()\n if result:\n generated_wave, generated_mel_spec = next(result)\n generated_waves.append(generated_wave)\n spectrograms.append(generated_mel_spec)\n\n if generated_waves:\n if cross_fade_duration <= 0:\n # Simply concatenate\n final_wave = np.concatenate(generated_waves)\n else:\n # Combine all generated waves with cross-fading\n final_wave = generated_waves[0]\n for i in range(1, len(generated_waves)):\n prev_wave = final_wave\n next_wave = generated_waves[i]\n\n # Calculate cross-fade samples, ensuring it does not exceed wave lengths\n cross_fade_samples = int(cross_fade_duration * target_sample_rate)\n cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave))\n\n if cross_fade_samples <= 0:\n # No overlap possible, concatenate\n final_wave = np.concatenate([prev_wave, next_wave])\n continue\n\n # Overlapping parts\n prev_overlap = prev_wave[-cross_fade_samples:]\n next_overlap = next_wave[:cross_fade_samples]\n\n # Fade out and fade in\n fade_out = np.linspace(1, 0, cross_fade_samples)\n fade_in = np.linspace(0, 1, cross_fade_samples)\n\n # Cross-faded overlap\n cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in\n\n # Combine\n new_wave = np.concatenate(\n [prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]]\n )\n\n final_wave = new_wave\n\n # Create a combined spectrogram\n combined_spectrogram = np.concatenate(spectrograms, axis=1)\n\n yield final_wave, target_sample_rate, combined_spectrogram\n\n else:\n yield None, target_sample_rate, None\n\n\n# remove silence from generated wav\n\n\ndef remove_silence_for_generated_wav(filename):\n aseg = AudioSegment.from_file(filename)\n non_silent_segs = silence.split_on_silence(\n aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500, seek_step=10\n )\n non_silent_wave = AudioSegment.silent(duration=0)\n for non_silent_seg in non_silent_segs:\n non_silent_wave += non_silent_seg\n aseg = non_silent_wave\n aseg.export(filename, format=\"wav\")\n\n\n# save spectrogram\n\n\ndef save_spectrogram(spectrogram, path):\n plt.figure(figsize=(12, 4))\n plt.imshow(spectrogram, origin=\"lower\", aspect=\"auto\")\n plt.colorbar()\n plt.savefig(path)\n plt.close()\n" }, { "path": "src/f5_tts/model/__init__.py", "content": "from f5_tts.model.backbones.dit import DiT\nfrom f5_tts.model.backbones.mmdit import MMDiT\nfrom f5_tts.model.backbones.unett import UNetT\nfrom f5_tts.model.cfm import CFM\nfrom f5_tts.model.trainer import Trainer\n\n\n__all__ = [\"CFM\", \"UNetT\", \"DiT\", \"MMDiT\", \"Trainer\"]\n" }, { "path": "src/f5_tts/model/backbones/README.md", "content": "## Backbones quick introduction\n\n\n### unett.py\n- flat unet transformer\n- structure same as in e2-tts & voicebox paper except using rotary pos emb\n- possible abs pos emb & convnextv2 blocks for embedded text before concat\n\n### dit.py\n- adaln-zero dit\n- embedded timestep as condition\n- concatted noised_input + masked_cond + embedded_text, linear proj in\n- possible abs pos emb & convnextv2 blocks for embedded text before concat\n- possible long skip connection (first layer to last layer)\n\n### mmdit.py\n- stable diffusion 3 block structure\n- timestep as condition\n- left stream: text embedded and applied a abs pos emb\n- right stream: masked_cond & noised_input concatted and with same conv pos emb as unett\n" }, { "path": "src/f5_tts/model/backbones/dit.py", "content": "\"\"\"\nein notation:\nb - batch\nn - sequence\nnt - text sequence\nnw - raw wave length\nd - dimension\n\"\"\"\n# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom x_transformers.x_transformers import RotaryEmbedding\n\nfrom f5_tts.model.modules import (\n AdaLayerNorm_Final,\n ConvNeXtV2Block,\n ConvPositionEmbedding,\n DiTBlock,\n TimestepEmbedding,\n precompute_freqs_cis,\n)\n\n\n# Text embedding\n\n\nclass TextEmbedding(nn.Module):\n def __init__(\n self, text_num_embeds, text_dim, mask_padding=True, average_upsampling=False, conv_layers=0, conv_mult=2\n ):\n super().__init__()\n self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token\n\n self.mask_padding = mask_padding # mask filler and batch padding tokens or not\n self.average_upsampling = average_upsampling # zipvoice-style text late average upsampling (after text encoder)\n if average_upsampling:\n assert mask_padding, \"text_embedding_average_upsampling requires text_mask_padding to be True\"\n\n if conv_layers > 0:\n self.extra_modeling = True\n self.precompute_max_pos = 8192 # 8192 is ~87.38s of 24khz audio; 4096 is ~43.69s of 24khz audio\n self.register_buffer(\"freqs_cis\", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)\n self.text_blocks = nn.Sequential(\n *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]\n )\n else:\n self.extra_modeling = False\n\n def average_upsample_text_by_mask(self, text, text_mask, target_lens):\n batch, max_seq_len, text_dim = text.shape\n text_lens = text_mask.sum(dim=1) # [batch]\n\n upsampled_text = torch.zeros_like(text)\n\n for i in range(batch):\n text_len = int(text_lens[i].item())\n audio_len = int(target_lens[i].item())\n\n if text_len == 0 or audio_len <= 0:\n continue\n\n valid_ind = torch.where(text_mask[i])[0]\n valid_data = text[i, valid_ind, :] # [text_len, text_dim]\n\n base_repeat = audio_len // text_len\n remainder = audio_len % text_len\n\n indices = []\n for j in range(text_len):\n repeat_count = base_repeat + (1 if j >= text_len - remainder else 0)\n indices.extend([j] * repeat_count)\n\n indices = torch.tensor(indices[:audio_len], device=text.device, dtype=torch.long)\n upsampled = valid_data[indices] # [audio_len, text_dim]\n\n upsampled_text[i, :audio_len, :] = upsampled\n\n return upsampled_text\n\n def forward(self, text: int[\"b nt\"], seq_len, drop_text=False):\n text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()\n valid_pos_mask = None\n if torch.is_tensor(seq_len):\n seq_len = seq_len.to(device=text.device, dtype=torch.long)\n max_seq_len = int(seq_len.max().item())\n else:\n max_seq_len = int(seq_len)\n\n text = text[:, :max_seq_len] # curtail if character tokens are more than the mel spec tokens\n text = F.pad(text, (0, max_seq_len - text.shape[1]), value=0)\n\n if torch.is_tensor(seq_len):\n seq_pos = torch.arange(max_seq_len, device=text.device).unsqueeze(0)\n valid_pos_mask = seq_pos < seq_len.unsqueeze(1)\n text = text.masked_fill(~valid_pos_mask, 0)\n\n if self.mask_padding:\n text_mask = text == 0\n\n if drop_text: # cfg for text\n text = torch.zeros_like(text)\n\n text = self.text_embed(text) # b n -> b n d\n if valid_pos_mask is not None:\n # Keep short-sample tail strictly zero (equivalent to per-sample pad_sequence(..., 0)).\n text = text.masked_fill(~valid_pos_mask.unsqueeze(-1), 0.0)\n\n # possible extra modeling\n if self.extra_modeling:\n # sinus pos emb; for variable seq lengths, only add positions within each sample's valid range.\n freqs = self.freqs_cis[:max_seq_len, :]\n if valid_pos_mask is not None:\n freqs = freqs.unsqueeze(0) * valid_pos_mask.unsqueeze(-1).to(freqs.dtype)\n text = text + freqs\n\n # convnextv2 blocks\n if self.mask_padding:\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n for block in self.text_blocks:\n text = block(text)\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n else:\n text = self.text_blocks(text)\n\n if self.average_upsampling:\n if torch.is_tensor(seq_len):\n target_lens = seq_len.to(device=text.device, dtype=torch.long)\n else:\n target_lens = torch.full((text.shape[0],), int(seq_len), device=text.device, dtype=torch.long)\n\n text = self.average_upsample_text_by_mask(text, ~text_mask, target_lens)\n\n return text\n\n\n# noised input audio and context mixing embedding\n\n\nclass InputEmbedding(nn.Module):\n def __init__(self, mel_dim, text_dim, out_dim):\n super().__init__()\n self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)\n self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)\n\n def forward(\n self,\n x: float[\"b n d\"],\n cond: float[\"b n d\"],\n text_embed: float[\"b n d\"],\n drop_audio_cond=False,\n audio_mask: bool[\"b n\"] | None = None,\n ):\n if drop_audio_cond: # cfg for cond audio\n cond = torch.zeros_like(cond)\n\n x = self.proj(torch.cat((x, cond, text_embed), dim=-1))\n x = self.conv_pos_embed(x, mask=audio_mask) + x\n return x\n\n\n# Transformer backbone using DiT blocks\n\n\nclass DiT(nn.Module):\n def __init__(\n self,\n *,\n dim,\n depth=8,\n heads=8,\n dim_head=64,\n dropout=0.1,\n ff_mult=4,\n mel_dim=100,\n text_num_embeds=256,\n text_dim=None,\n text_mask_padding=True,\n text_embedding_average_upsampling=False,\n qk_norm=None,\n conv_layers=0,\n pe_attn_head=None,\n attn_backend=\"torch\", # \"torch\" | \"flash_attn\"\n attn_mask_enabled=False,\n long_skip_connection=False,\n checkpoint_activations=False,\n ):\n super().__init__()\n\n self.time_embed = TimestepEmbedding(dim)\n if text_dim is None:\n text_dim = mel_dim\n self.text_embed = TextEmbedding(\n text_num_embeds,\n text_dim,\n mask_padding=text_mask_padding,\n average_upsampling=text_embedding_average_upsampling,\n conv_layers=conv_layers,\n )\n self.text_cond, self.text_uncond = None, None # text cache\n self.input_embed = InputEmbedding(mel_dim, text_dim, dim)\n\n self.rotary_embed = RotaryEmbedding(dim_head)\n\n self.dim = dim\n self.depth = depth\n\n self.transformer_blocks = nn.ModuleList(\n [\n DiTBlock(\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n ff_mult=ff_mult,\n dropout=dropout,\n qk_norm=qk_norm,\n pe_attn_head=pe_attn_head,\n attn_backend=attn_backend,\n attn_mask_enabled=attn_mask_enabled,\n )\n for _ in range(depth)\n ]\n )\n self.long_skip_connection = nn.Linear(dim * 2, dim, bias=False) if long_skip_connection else None\n\n self.norm_out = AdaLayerNorm_Final(dim) # final modulation\n self.proj_out = nn.Linear(dim, mel_dim)\n\n self.checkpoint_activations = checkpoint_activations\n\n self.initialize_weights()\n\n def initialize_weights(self):\n # Zero-out AdaLN layers in DiT blocks:\n for block in self.transformer_blocks:\n nn.init.constant_(block.attn_norm.linear.weight, 0)\n nn.init.constant_(block.attn_norm.linear.bias, 0)\n\n # Zero-out output layers:\n nn.init.constant_(self.norm_out.linear.weight, 0)\n nn.init.constant_(self.norm_out.linear.bias, 0)\n nn.init.constant_(self.proj_out.weight, 0)\n nn.init.constant_(self.proj_out.bias, 0)\n\n def ckpt_wrapper(self, module):\n # https://github.com/chuanyangjin/fast-DiT/blob/main/models.py\n def ckpt_forward(*inputs):\n outputs = module(*inputs)\n return outputs\n\n return ckpt_forward\n\n def get_input_embed(\n self,\n x, # b n d\n cond, # b n d\n text, # b nt\n drop_audio_cond: bool = False,\n drop_text: bool = False,\n cache: bool = True,\n audio_mask: bool[\"b n\"] | None = None,\n ):\n if self.text_uncond is None or self.text_cond is None or not cache:\n if audio_mask is None:\n seq_len = x.shape[1]\n else:\n seq_len = audio_mask.sum(dim=1) # per-sample valid speech length\n text_embed = self.text_embed(text, seq_len=seq_len, drop_text=drop_text)\n if cache:\n if drop_text:\n self.text_uncond = text_embed\n else:\n self.text_cond = text_embed\n\n if cache:\n if drop_text:\n text_embed = self.text_uncond\n else:\n text_embed = self.text_cond\n\n x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond, audio_mask=audio_mask)\n\n return x\n\n def clear_cache(self):\n self.text_cond, self.text_uncond = None, None\n\n def forward(\n self,\n x: float[\"b n d\"], # nosied input audio\n cond: float[\"b n d\"], # masked cond audio\n text: int[\"b nt\"], # text\n time: float[\"b\"] | float[\"\"], # time step\n mask: bool[\"b n\"] | None = None,\n drop_audio_cond: bool = False, # cfg for cond audio\n drop_text: bool = False, # cfg for text\n cfg_infer: bool = False, # cfg inference, pack cond & uncond forward\n cache: bool = False,\n ):\n batch, seq_len = x.shape[0], x.shape[1]\n if time.ndim == 0:\n time = time.repeat(batch)\n\n # t: conditioning time, text: text, x: noised audio + cond audio + text\n t = self.time_embed(time)\n if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d\n x_cond = self.get_input_embed(\n x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache, audio_mask=mask\n )\n x_uncond = self.get_input_embed(\n x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache, audio_mask=mask\n )\n x = torch.cat((x_cond, x_uncond), dim=0)\n t = torch.cat((t, t), dim=0)\n mask = torch.cat((mask, mask), dim=0) if mask is not None else None\n else:\n x = self.get_input_embed(\n x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache, audio_mask=mask\n )\n\n rope = self.rotary_embed.forward_from_seq_len(seq_len)\n\n if self.long_skip_connection is not None:\n residual = x\n\n for block in self.transformer_blocks:\n if self.checkpoint_activations:\n # https://pytorch.org/docs/stable/checkpoint.html#torch.utils.checkpoint.checkpoint\n x = torch.utils.checkpoint.checkpoint(self.ckpt_wrapper(block), x, t, mask, rope, use_reentrant=False)\n else:\n x = block(x, t, mask=mask, rope=rope)\n\n if self.long_skip_connection is not None:\n x = self.long_skip_connection(torch.cat((x, residual), dim=-1))\n\n x = self.norm_out(x, t)\n output = self.proj_out(x)\n\n return output\n" }, { "path": "src/f5_tts/model/backbones/mmdit.py", "content": "\"\"\"\nein notation:\nb - batch\nn - sequence\nnt - text sequence\nnw - raw wave length\nd - dimension\n\"\"\"\n# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nimport torch\nfrom torch import nn\nfrom x_transformers.x_transformers import RotaryEmbedding\n\nfrom f5_tts.model.modules import (\n AdaLayerNorm_Final,\n ConvPositionEmbedding,\n MMDiTBlock,\n TimestepEmbedding,\n get_pos_embed_indices,\n precompute_freqs_cis,\n)\n\n\n# text embedding\n\n\nclass TextEmbedding(nn.Module):\n def __init__(self, out_dim, text_num_embeds, mask_padding=True):\n super().__init__()\n self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim) # will use 0 as filler token\n\n self.mask_padding = mask_padding # mask filler and batch padding tokens or not\n\n self.precompute_max_pos = 1024\n self.register_buffer(\"freqs_cis\", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)\n\n def forward(self, text: int[\"b nt\"], drop_text=False) -> int[\"b nt d\"]:\n text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()\n if self.mask_padding:\n text_mask = text == 0\n\n if drop_text: # cfg for text\n text = torch.zeros_like(text)\n\n text = self.text_embed(text) # b nt -> b nt d\n\n # sinus pos emb\n batch_start = torch.zeros((text.shape[0],), dtype=torch.long)\n batch_text_len = text.shape[1]\n pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)\n text_pos_embed = self.freqs_cis[pos_idx]\n\n text = text + text_pos_embed\n\n if self.mask_padding:\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n\n return text\n\n\n# noised input & masked cond audio embedding\n\n\nclass AudioEmbedding(nn.Module):\n def __init__(self, in_dim, out_dim):\n super().__init__()\n self.linear = nn.Linear(2 * in_dim, out_dim)\n self.conv_pos_embed = ConvPositionEmbedding(out_dim)\n\n def forward(self, x: float[\"b n d\"], cond: float[\"b n d\"], drop_audio_cond=False):\n if drop_audio_cond:\n cond = torch.zeros_like(cond)\n x = torch.cat((x, cond), dim=-1)\n x = self.linear(x)\n x = self.conv_pos_embed(x) + x\n return x\n\n\n# Transformer backbone using MM-DiT blocks\n\n\nclass MMDiT(nn.Module):\n def __init__(\n self,\n *,\n dim,\n depth=8,\n heads=8,\n dim_head=64,\n dropout=0.1,\n ff_mult=4,\n mel_dim=100,\n text_num_embeds=256,\n text_mask_padding=True,\n qk_norm=None,\n checkpoint_activations=False,\n attn_backend=\"torch\",\n attn_mask_enabled=False,\n ):\n super().__init__()\n\n self.time_embed = TimestepEmbedding(dim)\n self.text_embed = TextEmbedding(dim, text_num_embeds, mask_padding=text_mask_padding)\n self.text_cond, self.text_uncond = None, None # text cache\n self.audio_embed = AudioEmbedding(mel_dim, dim)\n\n self.rotary_embed = RotaryEmbedding(dim_head)\n\n self.dim = dim\n self.depth = depth\n\n self.transformer_blocks = nn.ModuleList(\n [\n MMDiTBlock(\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n dropout=dropout,\n ff_mult=ff_mult,\n context_pre_only=i == depth - 1,\n qk_norm=qk_norm,\n attn_backend=attn_backend,\n attn_mask_enabled=attn_mask_enabled,\n )\n for i in range(depth)\n ]\n )\n self.norm_out = AdaLayerNorm_Final(dim) # final modulation\n self.proj_out = nn.Linear(dim, mel_dim)\n\n self.checkpoint_activations = checkpoint_activations\n\n self.initialize_weights()\n\n def initialize_weights(self):\n # Zero-out AdaLN layers in MMDiT blocks:\n for block in self.transformer_blocks:\n nn.init.constant_(block.attn_norm_x.linear.weight, 0)\n nn.init.constant_(block.attn_norm_x.linear.bias, 0)\n nn.init.constant_(block.attn_norm_c.linear.weight, 0)\n nn.init.constant_(block.attn_norm_c.linear.bias, 0)\n\n # Zero-out output layers:\n nn.init.constant_(self.norm_out.linear.weight, 0)\n nn.init.constant_(self.norm_out.linear.bias, 0)\n nn.init.constant_(self.proj_out.weight, 0)\n nn.init.constant_(self.proj_out.bias, 0)\n\n def ckpt_wrapper(self, module):\n def ckpt_forward(*inputs):\n outputs = module(*inputs)\n return outputs\n\n return ckpt_forward\n\n def get_input_embed(\n self,\n x, # b n d\n cond, # b n d\n text, # b nt\n drop_audio_cond: bool = False,\n drop_text: bool = False,\n cache: bool = True,\n ):\n if cache:\n if drop_text:\n if self.text_uncond is None:\n self.text_uncond = self.text_embed(text, drop_text=True)\n c = self.text_uncond\n else:\n if self.text_cond is None:\n self.text_cond = self.text_embed(text, drop_text=False)\n c = self.text_cond\n else:\n c = self.text_embed(text, drop_text=drop_text)\n x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)\n\n return x, c\n\n def clear_cache(self):\n self.text_cond, self.text_uncond = None, None\n\n def forward(\n self,\n x: float[\"b n d\"], # nosied input audio\n cond: float[\"b n d\"], # masked cond audio\n text: int[\"b nt\"], # text\n time: float[\"b\"] | float[\"\"], # time step\n mask: bool[\"b n\"] | None = None,\n drop_audio_cond: bool = False, # cfg for cond audio\n drop_text: bool = False, # cfg for text\n cfg_infer: bool = False, # cfg inference, pack cond & uncond forward\n cache: bool = False,\n ):\n batch = x.shape[0]\n if time.ndim == 0:\n time = time.repeat(batch)\n\n # t: conditioning (time), c: context (text + masked cond audio), x: noised input audio\n t = self.time_embed(time)\n c_mask = (text + 1) != 0 # True = valid, False = padding (-1 tokens)\n if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d\n x_cond, c_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache)\n x_uncond, c_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache)\n x = torch.cat((x_cond, x_uncond), dim=0)\n c = torch.cat((c_cond, c_uncond), dim=0)\n t = torch.cat((t, t), dim=0)\n mask = torch.cat((mask, mask), dim=0) if mask is not None else None\n c_mask = torch.cat((c_mask, c_mask), dim=0)\n else:\n x, c = self.get_input_embed(\n x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache\n )\n\n seq_len = x.shape[1]\n text_len = text.shape[1]\n rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)\n rope_text = self.rotary_embed.forward_from_seq_len(text_len)\n\n for block in self.transformer_blocks:\n if self.checkpoint_activations:\n c, x = torch.utils.checkpoint.checkpoint(\n self.ckpt_wrapper(block), x, c, t, mask, rope_audio, rope_text, c_mask, use_reentrant=False\n )\n else:\n c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text, c_mask=c_mask)\n\n x = self.norm_out(x, t)\n output = self.proj_out(x)\n\n return output\n" }, { "path": "src/f5_tts/model/backbones/unett.py", "content": "\"\"\"\nein notation:\nb - batch\nn - sequence\nnt - text sequence\nnw - raw wave length\nd - dimension\n\"\"\"\n# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nfrom typing import Literal\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom x_transformers import RMSNorm\nfrom x_transformers.x_transformers import RotaryEmbedding\n\nfrom f5_tts.model.modules import (\n Attention,\n AttnProcessor,\n ConvNeXtV2Block,\n ConvPositionEmbedding,\n FeedForward,\n TimestepEmbedding,\n get_pos_embed_indices,\n precompute_freqs_cis,\n)\n\n\n# Text embedding\n\n\nclass TextEmbedding(nn.Module):\n def __init__(self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2):\n super().__init__()\n self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token\n\n self.mask_padding = mask_padding # mask filler and batch padding tokens or not\n\n if conv_layers > 0:\n self.extra_modeling = True\n self.precompute_max_pos = 4096 # ~44s of 24khz audio\n self.register_buffer(\"freqs_cis\", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)\n self.text_blocks = nn.Sequential(\n *[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]\n )\n else:\n self.extra_modeling = False\n\n def forward(self, text: int[\"b nt\"], seq_len, drop_text=False):\n text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()\n text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens\n batch, text_len = text.shape[0], text.shape[1]\n text = F.pad(text, (0, seq_len - text_len), value=0)\n if self.mask_padding:\n text_mask = text == 0\n\n if drop_text: # cfg for text\n text = torch.zeros_like(text)\n\n text = self.text_embed(text) # b n -> b n d\n\n # possible extra modeling\n if self.extra_modeling:\n # sinus pos emb\n batch_start = torch.zeros((batch,), dtype=torch.long)\n pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)\n text_pos_embed = self.freqs_cis[pos_idx]\n text = text + text_pos_embed\n\n # convnextv2 blocks\n if self.mask_padding:\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n for block in self.text_blocks:\n text = block(text)\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n else:\n text = self.text_blocks(text)\n\n return text\n\n\n# noised input audio and context mixing embedding\n\n\nclass InputEmbedding(nn.Module):\n def __init__(self, mel_dim, text_dim, out_dim):\n super().__init__()\n self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)\n self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)\n\n def forward(self, x: float[\"b n d\"], cond: float[\"b n d\"], text_embed: float[\"b n d\"], drop_audio_cond=False):\n if drop_audio_cond: # cfg for cond audio\n cond = torch.zeros_like(cond)\n\n x = self.proj(torch.cat((x, cond, text_embed), dim=-1))\n x = self.conv_pos_embed(x) + x\n return x\n\n\n# Flat UNet Transformer backbone\n\n\nclass UNetT(nn.Module):\n def __init__(\n self,\n *,\n dim,\n depth=8,\n heads=8,\n dim_head=64,\n dropout=0.1,\n ff_mult=4,\n mel_dim=100,\n text_num_embeds=256,\n text_dim=None,\n text_mask_padding=True,\n qk_norm=None,\n conv_layers=0,\n pe_attn_head=None,\n attn_backend=\"torch\", # \"torch\" | \"flash_attn\"\n attn_mask_enabled=False,\n skip_connect_type: Literal[\"add\", \"concat\", \"none\"] = \"concat\",\n ):\n super().__init__()\n assert depth % 2 == 0, \"UNet-Transformer's depth should be even.\"\n\n self.time_embed = TimestepEmbedding(dim)\n if text_dim is None:\n text_dim = mel_dim\n self.text_embed = TextEmbedding(\n text_num_embeds, text_dim, mask_padding=text_mask_padding, conv_layers=conv_layers\n )\n self.text_cond, self.text_uncond = None, None # text cache\n self.input_embed = InputEmbedding(mel_dim, text_dim, dim)\n\n self.rotary_embed = RotaryEmbedding(dim_head)\n\n # transformer layers & skip connections\n\n self.dim = dim\n self.skip_connect_type = skip_connect_type\n needs_skip_proj = skip_connect_type == \"concat\"\n\n self.depth = depth\n self.layers = nn.ModuleList([])\n\n for idx in range(depth):\n is_later_half = idx >= (depth // 2)\n\n attn_norm = RMSNorm(dim)\n attn = Attention(\n processor=AttnProcessor(\n pe_attn_head=pe_attn_head,\n attn_backend=attn_backend,\n attn_mask_enabled=attn_mask_enabled,\n ),\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n dropout=dropout,\n qk_norm=qk_norm,\n )\n\n ff_norm = RMSNorm(dim)\n ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate=\"tanh\")\n\n skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None\n\n self.layers.append(\n nn.ModuleList(\n [\n skip_proj,\n attn_norm,\n attn,\n ff_norm,\n ff,\n ]\n )\n )\n\n self.norm_out = RMSNorm(dim)\n self.proj_out = nn.Linear(dim, mel_dim)\n\n def get_input_embed(\n self,\n x, # b n d\n cond, # b n d\n text, # b nt\n drop_audio_cond: bool = False,\n drop_text: bool = False,\n cache: bool = True,\n ):\n seq_len = x.shape[1]\n if cache:\n if drop_text:\n if self.text_uncond is None:\n self.text_uncond = self.text_embed(text, seq_len, drop_text=True)\n text_embed = self.text_uncond\n else:\n if self.text_cond is None:\n self.text_cond = self.text_embed(text, seq_len, drop_text=False)\n text_embed = self.text_cond\n else:\n text_embed = self.text_embed(text, seq_len, drop_text=drop_text)\n\n x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)\n\n return x\n\n def clear_cache(self):\n self.text_cond, self.text_uncond = None, None\n\n def forward(\n self,\n x: float[\"b n d\"], # nosied input audio\n cond: float[\"b n d\"], # masked cond audio\n text: int[\"b nt\"], # text\n time: float[\"b\"] | float[\"\"], # time step\n mask: bool[\"b n\"] | None = None,\n drop_audio_cond: bool = False, # cfg for cond audio\n drop_text: bool = False, # cfg for text\n cfg_infer: bool = False, # cfg inference, pack cond & uncond forward\n cache: bool = False,\n ):\n batch, seq_len = x.shape[0], x.shape[1]\n if time.ndim == 0:\n time = time.repeat(batch)\n\n # t: conditioning time, c: context (text + masked cond audio), x: noised input audio\n t = self.time_embed(time)\n if cfg_infer: # pack cond & uncond forward: b n d -> 2b n d\n x_cond = self.get_input_embed(x, cond, text, drop_audio_cond=False, drop_text=False, cache=cache)\n x_uncond = self.get_input_embed(x, cond, text, drop_audio_cond=True, drop_text=True, cache=cache)\n x = torch.cat((x_cond, x_uncond), dim=0)\n t = torch.cat((t, t), dim=0)\n mask = torch.cat((mask, mask), dim=0) if mask is not None else None\n else:\n x = self.get_input_embed(x, cond, text, drop_audio_cond=drop_audio_cond, drop_text=drop_text, cache=cache)\n\n # postfix time t to input x, [b n d] -> [b n+1 d]\n x = torch.cat([t.unsqueeze(1), x], dim=1) # pack t to x\n if mask is not None:\n mask = F.pad(mask, (1, 0), value=1)\n\n rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)\n\n # flat unet transformer\n skip_connect_type = self.skip_connect_type\n skips = []\n for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):\n layer = idx + 1\n\n # skip connection logic\n is_first_half = layer <= (self.depth // 2)\n is_later_half = not is_first_half\n\n if is_first_half:\n skips.append(x)\n\n if is_later_half:\n skip = skips.pop()\n if skip_connect_type == \"concat\":\n x = torch.cat((x, skip), dim=-1)\n x = maybe_skip_proj(x)\n elif skip_connect_type == \"add\":\n x = x + skip\n\n # attention and feedforward blocks\n x = attn(attn_norm(x), rope=rope, mask=mask) + x\n x = ff(ff_norm(x)) + x\n\n assert len(skips) == 0\n\n x = self.norm_out(x)[:, 1:, :] # unpack t from x\n\n return self.proj_out(x)\n" }, { "path": "src/f5_tts/model/cfm.py", "content": "\"\"\"\nein notation:\nb - batch\nn - sequence\nnt - text sequence\nnw - raw wave length\nd - dimension\n\"\"\"\n# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nfrom random import random\nfrom typing import Callable\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torchdiffeq import odeint\n\nfrom f5_tts.model.modules import MelSpec\nfrom f5_tts.model.utils import (\n default,\n exists,\n get_epss_timesteps,\n lens_to_mask,\n list_str_to_idx,\n list_str_to_tensor,\n mask_from_frac_lengths,\n)\n\n\nclass CFM(nn.Module):\n def __init__(\n self,\n transformer: nn.Module,\n sigma=0.0,\n odeint_kwargs: dict = dict(\n # atol = 1e-5,\n # rtol = 1e-5,\n method=\"euler\" # 'midpoint'\n ),\n audio_drop_prob=0.3,\n cond_drop_prob=0.2,\n num_channels=None,\n mel_spec_module: nn.Module | None = None,\n mel_spec_kwargs: dict = dict(),\n frac_lengths_mask: tuple[float, float] = (0.7, 1.0),\n vocab_char_map: dict[str:int] | None = None,\n ):\n super().__init__()\n\n self.frac_lengths_mask = frac_lengths_mask\n\n # mel spec\n self.mel_spec = default(mel_spec_module, MelSpec(**mel_spec_kwargs))\n num_channels = default(num_channels, self.mel_spec.n_mel_channels)\n self.num_channels = num_channels\n\n # classifier-free guidance\n self.audio_drop_prob = audio_drop_prob\n self.cond_drop_prob = cond_drop_prob\n\n # transformer\n self.transformer = transformer\n dim = transformer.dim\n self.dim = dim\n\n # conditional flow related\n self.sigma = sigma\n\n # sampling related\n self.odeint_kwargs = odeint_kwargs\n\n # vocab map for tokenization\n self.vocab_char_map = vocab_char_map\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n @torch.no_grad()\n def sample(\n self,\n cond: float[\"b n d\"] | float[\"b nw\"],\n text: int[\"b nt\"] | list[str],\n duration: int | int[\"b\"],\n *,\n lens: int[\"b\"] | None = None,\n steps=32,\n cfg_strength=1.0,\n sway_sampling_coef=None,\n seed: int | None = None,\n max_duration=65536,\n vocoder: Callable[[float[\"b d n\"]], float[\"b nw\"]] | None = None,\n use_epss=True,\n no_ref_audio=False,\n duplicate_test=False,\n t_inter=0.1,\n edit_mask=None,\n ):\n self.eval()\n # raw wave\n\n if cond.ndim == 2:\n cond = self.mel_spec(cond)\n cond = cond.permute(0, 2, 1)\n assert cond.shape[-1] == self.num_channels\n\n cond = cond.to(next(self.parameters()).dtype)\n\n batch, cond_seq_len, device = *cond.shape[:2], cond.device\n if not exists(lens):\n lens = torch.full((batch,), cond_seq_len, device=device, dtype=torch.long)\n\n # text\n\n if isinstance(text, list):\n if exists(self.vocab_char_map):\n text = list_str_to_idx(text, self.vocab_char_map).to(device)\n else:\n text = list_str_to_tensor(text).to(device)\n assert text.shape[0] == batch\n\n # duration\n\n cond_mask = lens_to_mask(lens)\n if edit_mask is not None:\n cond_mask = cond_mask & edit_mask\n\n if isinstance(duration, int):\n duration = torch.full((batch,), duration, device=device, dtype=torch.long)\n\n duration = torch.maximum(\n torch.maximum((text != -1).sum(dim=-1), lens) + 1, duration\n ) # duration at least text/audio prompt length plus one token, so something is generated\n duration = duration.clamp(max=max_duration)\n max_duration = duration.amax()\n\n # duplicate test corner for inner time step oberservation\n if duplicate_test:\n test_cond = F.pad(cond, (0, 0, cond_seq_len, max_duration - 2 * cond_seq_len), value=0.0)\n\n cond = F.pad(cond, (0, 0, 0, max_duration - cond_seq_len), value=0.0)\n if no_ref_audio:\n cond = torch.zeros_like(cond)\n\n cond_mask = F.pad(cond_mask, (0, max_duration - cond_mask.shape[-1]), value=False)\n cond_mask = cond_mask.unsqueeze(-1)\n step_cond = torch.where(\n cond_mask, cond, torch.zeros_like(cond)\n ) # allow direct control (cut cond audio) with lens passed in\n\n if batch > 1:\n mask = lens_to_mask(duration)\n else: # save memory and speed up, as single inference need no mask currently\n mask = None\n\n # neural ode\n\n def fn(t, x):\n # at each step, conditioning is fixed\n # step_cond = torch.where(cond_mask, cond, torch.zeros_like(cond))\n\n # predict flow (cond)\n if cfg_strength < 1e-5:\n pred = self.transformer(\n x=x,\n cond=step_cond,\n text=text,\n time=t,\n mask=mask,\n drop_audio_cond=False,\n drop_text=False,\n cache=True,\n )\n return pred\n\n # predict flow (cond and uncond), for classifier-free guidance\n pred_cfg = self.transformer(\n x=x,\n cond=step_cond,\n text=text,\n time=t,\n mask=mask,\n cfg_infer=True,\n cache=True,\n )\n pred, null_pred = torch.chunk(pred_cfg, 2, dim=0)\n return pred + (pred - null_pred) * cfg_strength\n\n # noise input\n # to make sure batch inference result is same with different batch size, and for sure single inference\n # still some difference maybe due to convolutional layers\n y0 = []\n for dur in duration:\n if exists(seed):\n torch.manual_seed(seed)\n y0.append(torch.randn(dur, self.num_channels, device=self.device, dtype=step_cond.dtype))\n y0 = pad_sequence(y0, padding_value=0, batch_first=True)\n\n t_start = 0\n\n # duplicate test corner for inner time step oberservation\n if duplicate_test:\n t_start = t_inter\n y0 = (1 - t_start) * y0 + t_start * test_cond\n steps = int(steps * (1 - t_start))\n\n if t_start == 0 and use_epss: # use Empirically Pruned Step Sampling for low NFE\n t = get_epss_timesteps(steps, device=self.device, dtype=step_cond.dtype)\n else:\n t = torch.linspace(t_start, 1, steps + 1, device=self.device, dtype=step_cond.dtype)\n if sway_sampling_coef is not None:\n t = t + sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t)\n\n trajectory = odeint(fn, y0, t, **self.odeint_kwargs)\n self.transformer.clear_cache()\n\n sampled = trajectory[-1]\n out = sampled\n out = torch.where(cond_mask, cond, out)\n\n if exists(vocoder):\n out = out.permute(0, 2, 1)\n out = vocoder(out)\n\n return out, trajectory\n\n def forward(\n self,\n inp: float[\"b n d\"] | float[\"b nw\"], # mel or raw wave\n text: int[\"b nt\"] | list[str],\n *,\n lens: int[\"b\"] | None = None,\n noise_scheduler: str | None = None,\n ):\n # handle raw wave\n if inp.ndim == 2:\n inp = self.mel_spec(inp)\n inp = inp.permute(0, 2, 1)\n assert inp.shape[-1] == self.num_channels\n\n batch, seq_len, dtype, device, _σ1 = *inp.shape[:2], inp.dtype, self.device, self.sigma\n\n # handle text as string\n if isinstance(text, list):\n if exists(self.vocab_char_map):\n text = list_str_to_idx(text, self.vocab_char_map).to(device)\n else:\n text = list_str_to_tensor(text).to(device)\n assert text.shape[0] == batch\n\n # lens and mask\n if not exists(lens): # if lens not acquired by trainer from collate_fn\n lens = torch.full((batch,), seq_len, device=device)\n mask = lens_to_mask(lens, length=seq_len)\n\n # get a random span to mask out for training conditionally\n frac_lengths = torch.zeros((batch,), device=self.device).float().uniform_(*self.frac_lengths_mask)\n rand_span_mask = mask_from_frac_lengths(lens, frac_lengths)\n\n if exists(mask):\n rand_span_mask &= mask\n\n # mel is x1\n x1 = inp\n\n # x0 is gaussian noise\n x0 = torch.randn_like(x1)\n\n # time step\n time = torch.rand((batch,), dtype=dtype, device=self.device)\n # TODO. noise_scheduler\n\n # sample xt (φ_t(x) in the paper)\n t = time.unsqueeze(-1).unsqueeze(-1)\n φ = (1 - t) * x0 + t * x1\n flow = x1 - x0\n\n # only predict what is within the random mask span for infilling\n cond = torch.where(rand_span_mask[..., None], torch.zeros_like(x1), x1)\n\n # transformer and cfg training with a drop rate\n drop_audio_cond = random() < self.audio_drop_prob # p_drop in voicebox paper\n if random() < self.cond_drop_prob: # p_uncond in voicebox paper\n drop_audio_cond = True\n drop_text = True\n else:\n drop_text = False\n\n # apply mask will use more memory; might adjust batchsize or batchsampler long sequence threshold\n pred = self.transformer(\n x=φ, cond=cond, text=text, time=time, drop_audio_cond=drop_audio_cond, drop_text=drop_text, mask=mask\n )\n\n # flow matching loss\n loss = F.mse_loss(pred, flow, reduction=\"none\")\n loss = loss[rand_span_mask]\n\n return loss.mean(), cond, pred\n" }, { "path": "src/f5_tts/model/dataset.py", "content": "import json\nfrom importlib.resources import files\n\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nfrom datasets import Dataset as Dataset_\nfrom datasets import load_from_disk\nfrom torch import nn\nfrom torch.utils.data import Dataset, Sampler\nfrom tqdm import tqdm\n\nfrom f5_tts.model.modules import MelSpec\nfrom f5_tts.model.utils import default\n\n\nclass HFDataset(Dataset):\n def __init__(\n self,\n hf_dataset: Dataset,\n target_sample_rate=24_000,\n n_mel_channels=100,\n hop_length=256,\n n_fft=1024,\n win_length=1024,\n mel_spec_type=\"vocos\",\n ):\n self.data = hf_dataset\n self.target_sample_rate = target_sample_rate\n self.hop_length = hop_length\n\n self.mel_spectrogram = MelSpec(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n )\n\n def get_frame_len(self, index):\n row = self.data[index]\n audio = row[\"audio\"][\"array\"]\n sample_rate = row[\"audio\"][\"sampling_rate\"]\n return audio.shape[-1] / sample_rate * self.target_sample_rate / self.hop_length\n\n def __len__(self):\n return len(self.data)\n\n def __getitem__(self, index):\n row = self.data[index]\n audio = row[\"audio\"][\"array\"]\n\n # logger.info(f\"Audio shape: {audio.shape}\")\n\n sample_rate = row[\"audio\"][\"sampling_rate\"]\n duration = audio.shape[-1] / sample_rate\n\n if duration > 30 or duration < 0.3:\n return self.__getitem__((index + 1) % len(self.data))\n\n audio_tensor = torch.from_numpy(audio).float()\n\n if sample_rate != self.target_sample_rate:\n resampler = torchaudio.transforms.Resample(sample_rate, self.target_sample_rate)\n audio_tensor = resampler(audio_tensor)\n\n audio_tensor = audio_tensor.unsqueeze(0) # 't -> 1 t')\n\n mel_spec = self.mel_spectrogram(audio_tensor)\n\n mel_spec = mel_spec.squeeze(0) # '1 d t -> d t'\n\n text = row[\"text\"]\n\n return dict(\n mel_spec=mel_spec,\n text=text,\n )\n\n\nclass CustomDataset(Dataset):\n def __init__(\n self,\n custom_dataset: Dataset,\n durations=None,\n target_sample_rate=24_000,\n hop_length=256,\n n_mel_channels=100,\n n_fft=1024,\n win_length=1024,\n mel_spec_type=\"vocos\",\n preprocessed_mel=False,\n mel_spec_module: nn.Module | None = None,\n ):\n self.data = custom_dataset\n self.durations = durations\n self.target_sample_rate = target_sample_rate\n self.hop_length = hop_length\n self.n_fft = n_fft\n self.win_length = win_length\n self.mel_spec_type = mel_spec_type\n self.preprocessed_mel = preprocessed_mel\n\n if not preprocessed_mel:\n self.mel_spectrogram = default(\n mel_spec_module,\n MelSpec(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n ),\n )\n\n def get_frame_len(self, index):\n if (\n self.durations is not None\n ): # Please make sure the separately provided durations are correct, otherwise 99.99% OOM\n return self.durations[index] * self.target_sample_rate / self.hop_length\n return self.data[index][\"duration\"] * self.target_sample_rate / self.hop_length\n\n def __len__(self):\n return len(self.data)\n\n def __getitem__(self, index):\n while True:\n row = self.data[index]\n audio_path = row[\"audio_path\"]\n text = row[\"text\"]\n duration = row[\"duration\"]\n\n # filter by given length\n if 0.3 <= duration <= 30:\n break # valid\n\n index = (index + 1) % len(self.data)\n\n if self.preprocessed_mel:\n mel_spec = torch.tensor(row[\"mel_spec\"])\n else:\n audio, source_sample_rate = torchaudio.load(audio_path)\n\n # make sure mono input\n if audio.shape[0] > 1:\n audio = torch.mean(audio, dim=0, keepdim=True)\n\n # resample if necessary\n if source_sample_rate != self.target_sample_rate:\n resampler = torchaudio.transforms.Resample(source_sample_rate, self.target_sample_rate)\n audio = resampler(audio)\n\n # to mel spectrogram\n mel_spec = self.mel_spectrogram(audio)\n mel_spec = mel_spec.squeeze(0) # '1 d t -> d t'\n\n return {\n \"mel_spec\": mel_spec,\n \"text\": text,\n }\n\n\n# Dynamic Batch Sampler\nclass DynamicBatchSampler(Sampler[list[int]]):\n \"\"\"Extension of Sampler that will do the following:\n 1. Change the batch size (essentially number of sequences)\n in a batch to ensure that the total number of frames are less\n than a certain threshold.\n 2. Make sure the padding efficiency in the batch is high.\n 3. Shuffle batches each epoch while maintaining reproducibility.\n \"\"\"\n\n def __init__(\n self, sampler: Sampler[int], frames_threshold: int, max_samples=0, random_seed=None, drop_residual: bool = False\n ):\n self.sampler = sampler\n self.frames_threshold = frames_threshold\n self.max_samples = max_samples\n self.random_seed = random_seed\n self.epoch = 0\n\n indices, batches = [], []\n data_source = self.sampler.data_source\n\n for idx in tqdm(\n self.sampler, desc=\"Sorting with sampler... if slow, check whether dataset is provided with duration\"\n ):\n indices.append((idx, data_source.get_frame_len(idx)))\n indices.sort(key=lambda elem: elem[1])\n\n batch = []\n batch_frames = 0\n for idx, frame_len in tqdm(\n indices, desc=f\"Creating dynamic batches with {frames_threshold} audio frames per gpu\"\n ):\n if batch_frames + frame_len <= self.frames_threshold and (max_samples == 0 or len(batch) < max_samples):\n batch.append(idx)\n batch_frames += frame_len\n else:\n if len(batch) > 0:\n batches.append(batch)\n if frame_len <= self.frames_threshold:\n batch = [idx]\n batch_frames = frame_len\n else:\n batch = []\n batch_frames = 0\n\n if not drop_residual and len(batch) > 0:\n batches.append(batch)\n\n del indices\n self.batches = batches\n\n # Ensure even batches with accelerate BatchSamplerShard cls under frame_per_batch setting\n self.drop_last = True\n\n def set_epoch(self, epoch: int) -> None:\n \"\"\"Sets the epoch for this sampler.\"\"\"\n self.epoch = epoch\n\n def __iter__(self):\n # Use both random_seed and epoch for deterministic but different shuffling per epoch\n if self.random_seed is not None:\n g = torch.Generator()\n g.manual_seed(self.random_seed + self.epoch)\n # Use PyTorch's random permutation for better reproducibility across PyTorch versions\n indices = torch.randperm(len(self.batches), generator=g).tolist()\n batches = [self.batches[i] for i in indices]\n else:\n batches = self.batches\n return iter(batches)\n\n def __len__(self):\n return len(self.batches)\n\n\n# Load dataset\n\n\ndef load_dataset(\n dataset_name: str,\n tokenizer: str = \"pinyin\",\n dataset_type: str = \"CustomDataset\",\n audio_type: str = \"raw\",\n mel_spec_module: nn.Module | None = None,\n mel_spec_kwargs: dict = dict(),\n) -> CustomDataset | HFDataset:\n \"\"\"\n dataset_type - \"CustomDataset\" if you want to use tokenizer name and default data path to load for train_dataset\n - \"CustomDatasetPath\" if you just want to pass the full path to a preprocessed dataset without relying on tokenizer\n \"\"\"\n\n print(\"Loading dataset ...\")\n\n if dataset_type == \"CustomDataset\":\n rel_data_path = str(files(\"f5_tts\").joinpath(f\"../../data/{dataset_name}_{tokenizer}\"))\n if audio_type == \"raw\":\n try:\n train_dataset = load_from_disk(f\"{rel_data_path}/raw\")\n except: # noqa: E722\n train_dataset = Dataset_.from_file(f\"{rel_data_path}/raw.arrow\")\n preprocessed_mel = False\n elif audio_type == \"mel\":\n train_dataset = Dataset_.from_file(f\"{rel_data_path}/mel.arrow\")\n preprocessed_mel = True\n with open(f\"{rel_data_path}/duration.json\", \"r\", encoding=\"utf-8\") as f:\n data_dict = json.load(f)\n durations = data_dict[\"duration\"]\n train_dataset = CustomDataset(\n train_dataset,\n durations=durations,\n preprocessed_mel=preprocessed_mel,\n mel_spec_module=mel_spec_module,\n **mel_spec_kwargs,\n )\n\n elif dataset_type == \"CustomDatasetPath\":\n try:\n train_dataset = load_from_disk(f\"{dataset_name}/raw\")\n except: # noqa: E722\n train_dataset = Dataset_.from_file(f\"{dataset_name}/raw.arrow\")\n\n with open(f\"{dataset_name}/duration.json\", \"r\", encoding=\"utf-8\") as f:\n data_dict = json.load(f)\n durations = data_dict[\"duration\"]\n train_dataset = CustomDataset(\n train_dataset, durations=durations, preprocessed_mel=preprocessed_mel, **mel_spec_kwargs\n )\n\n elif dataset_type == \"HFDataset\":\n print(\n \"Should manually modify the path of huggingface dataset to your need.\\n\"\n + \"May also the corresponding script cuz different dataset may have different format.\"\n )\n pre, post = dataset_name.split(\"_\")\n train_dataset = HFDataset(\n load_dataset(f\"{pre}/{pre}\", split=f\"train.{post}\", cache_dir=str(files(\"f5_tts\").joinpath(\"../../data\"))),\n )\n\n return train_dataset\n\n\n# collation\n\n\ndef collate_fn(batch):\n mel_specs = [item[\"mel_spec\"].squeeze(0) for item in batch]\n mel_lengths = torch.LongTensor([spec.shape[-1] for spec in mel_specs])\n max_mel_length = mel_lengths.amax()\n\n padded_mel_specs = []\n for spec in mel_specs:\n padding = (0, max_mel_length - spec.size(-1))\n padded_spec = F.pad(spec, padding, value=0)\n padded_mel_specs.append(padded_spec)\n\n mel_specs = torch.stack(padded_mel_specs)\n\n text = [item[\"text\"] for item in batch]\n text_lengths = torch.LongTensor([len(item) for item in text])\n\n return dict(\n mel=mel_specs,\n mel_lengths=mel_lengths, # records for padding mask\n text=text,\n text_lengths=text_lengths,\n )\n" }, { "path": "src/f5_tts/model/modules.py", "content": "\"\"\"\nein notation:\nb - batch\nn - sequence\nnt - text sequence\nnw - raw wave length\nd - dimension\n\"\"\"\n# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nimport math\nimport warnings\nfrom typing import Optional\n\nimport torch\nimport torch.nn.functional as F\nimport torchaudio\nfrom librosa.filters import mel as librosa_mel_fn\nfrom torch import nn\nfrom x_transformers.x_transformers import apply_rotary_pos_emb\n\nfrom f5_tts.model.utils import is_package_available\n\n\n# raw wav to mel spec\n\n\nmel_basis_cache = {}\nhann_window_cache = {}\n\n\ndef get_bigvgan_mel_spectrogram(\n waveform,\n n_fft=1024,\n n_mel_channels=100,\n target_sample_rate=24000,\n hop_length=256,\n win_length=1024,\n fmin=0,\n fmax=None,\n center=False,\n): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main\n device = waveform.device\n key = f\"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}\"\n\n if key not in mel_basis_cache:\n mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)\n mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()?\n hann_window_cache[key] = torch.hann_window(win_length).to(device)\n\n mel_basis = mel_basis_cache[key]\n hann_window = hann_window_cache[key]\n\n padding = (n_fft - hop_length) // 2\n waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode=\"reflect\").squeeze(1)\n\n spec = torch.stft(\n waveform,\n n_fft,\n hop_length=hop_length,\n win_length=win_length,\n window=hann_window,\n center=center,\n pad_mode=\"reflect\",\n normalized=False,\n onesided=True,\n return_complex=True,\n )\n spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)\n\n mel_spec = torch.matmul(mel_basis, spec)\n mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))\n\n return mel_spec\n\n\ndef get_vocos_mel_spectrogram(\n waveform,\n n_fft=1024,\n n_mel_channels=100,\n target_sample_rate=24000,\n hop_length=256,\n win_length=1024,\n):\n mel_stft = torchaudio.transforms.MelSpectrogram(\n sample_rate=target_sample_rate,\n n_fft=n_fft,\n win_length=win_length,\n hop_length=hop_length,\n n_mels=n_mel_channels,\n power=1,\n center=True,\n normalized=False,\n norm=None,\n ).to(waveform.device)\n if len(waveform.shape) == 3:\n waveform = waveform.squeeze(1) # 'b 1 nw -> b nw'\n\n assert len(waveform.shape) == 2\n\n mel = mel_stft(waveform)\n mel = mel.clamp(min=1e-5).log()\n return mel\n\n\nclass MelSpec(nn.Module):\n def __init__(\n self,\n n_fft=1024,\n hop_length=256,\n win_length=1024,\n n_mel_channels=100,\n target_sample_rate=24_000,\n mel_spec_type=\"vocos\",\n ):\n super().__init__()\n assert mel_spec_type in [\"vocos\", \"bigvgan\"], print(\"We only support two extract mel backend: vocos or bigvgan\")\n\n self.n_fft = n_fft\n self.hop_length = hop_length\n self.win_length = win_length\n self.n_mel_channels = n_mel_channels\n self.target_sample_rate = target_sample_rate\n\n if mel_spec_type == \"vocos\":\n self.extractor = get_vocos_mel_spectrogram\n elif mel_spec_type == \"bigvgan\":\n self.extractor = get_bigvgan_mel_spectrogram\n\n self.register_buffer(\"dummy\", torch.tensor(0), persistent=False)\n\n def forward(self, wav):\n if self.dummy.device != wav.device:\n self.to(wav.device)\n\n mel = self.extractor(\n waveform=wav,\n n_fft=self.n_fft,\n n_mel_channels=self.n_mel_channels,\n target_sample_rate=self.target_sample_rate,\n hop_length=self.hop_length,\n win_length=self.win_length,\n )\n\n return mel\n\n\n# sinusoidal position embedding\n\n\nclass SinusPositionEmbedding(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.dim = dim\n\n def forward(self, x, scale=1000):\n device = x.device\n half_dim = self.dim // 2\n emb = math.log(10000) / (half_dim - 1)\n emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)\n emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)\n emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n return emb\n\n\n# convolutional position embedding\n\n\nclass ConvPositionEmbedding(nn.Module):\n def __init__(self, dim, kernel_size=31, groups=16):\n super().__init__()\n assert kernel_size % 2 != 0\n self.conv1d = nn.Sequential(\n nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),\n nn.Mish(),\n nn.Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2),\n nn.Mish(),\n )\n self.layer_need_mask_idx = [i for i, layer in enumerate(self.conv1d) if isinstance(layer, nn.Conv1d)]\n\n def forward(self, x: float[\"b n d\"], mask: bool[\"b n\"] | None = None):\n if mask is not None:\n mask = mask.unsqueeze(1) # [B 1 N]\n x = x.permute(0, 2, 1) # [B D N]\n\n if mask is not None:\n x = x.masked_fill(~mask, 0.0)\n for i, block in enumerate(self.conv1d):\n x = block(x)\n if mask is not None and i in self.layer_need_mask_idx:\n x = x.masked_fill(~mask, 0.0)\n\n x = x.permute(0, 2, 1) # [B N D]\n\n return x\n\n\n# rotary positional embedding related\n\n\ndef precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):\n # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning\n # has some connection to NTK literature\n # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/\n # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py\n theta *= theta_rescale_factor ** (dim / (dim - 2))\n freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n t = torch.arange(end, device=freqs.device) # type: ignore\n freqs = torch.outer(t, freqs).float() # type: ignore\n freqs_cos = torch.cos(freqs) # real part\n freqs_sin = torch.sin(freqs) # imaginary part\n return torch.cat([freqs_cos, freqs_sin], dim=-1)\n\n\ndef get_pos_embed_indices(start, length, max_pos, scale=1.0):\n # length = length if isinstance(length, int) else length.max()\n scale = scale * torch.ones_like(start, dtype=torch.float32) # in case scale is a scalar\n pos = (\n start.unsqueeze(1)\n + (torch.arange(length, device=start.device, dtype=torch.float32).unsqueeze(0) * scale.unsqueeze(1)).long()\n )\n # avoid extra long error.\n pos = torch.where(pos < max_pos, pos, max_pos - 1)\n return pos\n\n\n# Global Response Normalization layer (Instance Normalization ?)\n\n\nclass GRN(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.gamma = nn.Parameter(torch.zeros(1, 1, dim))\n self.beta = nn.Parameter(torch.zeros(1, 1, dim))\n\n def forward(self, x):\n Gx = torch.norm(x, p=2, dim=1, keepdim=True)\n Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)\n return self.gamma * (x * Nx) + self.beta + x\n\n\n# ConvNeXt-V2 Block https://github.com/facebookresearch/ConvNeXt-V2/blob/main/models/convnextv2.py\n# ref: https://github.com/bfs18/e2_tts/blob/main/rfwave/modules.py#L108\n\n\nclass ConvNeXtV2Block(nn.Module):\n def __init__(\n self,\n dim: int,\n intermediate_dim: int,\n dilation: int = 1,\n ):\n super().__init__()\n padding = (dilation * (7 - 1)) // 2\n self.dwconv = nn.Conv1d(\n dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation\n ) # depthwise conv\n self.norm = nn.LayerNorm(dim, eps=1e-6)\n self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers\n self.act = nn.GELU()\n self.grn = GRN(intermediate_dim)\n self.pwconv2 = nn.Linear(intermediate_dim, dim)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n residual = x\n x = x.transpose(1, 2) # b n d -> b d n\n x = self.dwconv(x)\n x = x.transpose(1, 2) # b d n -> b n d\n x = self.norm(x)\n x = self.pwconv1(x)\n x = self.act(x)\n x = self.grn(x)\n x = self.pwconv2(x)\n return residual + x\n\n\n# RMSNorm\n\n\nclass RMSNorm(nn.Module):\n def __init__(self, dim: int, eps: float):\n super().__init__()\n self.eps = eps\n self.weight = nn.Parameter(torch.ones(dim))\n self.native_rms_norm = float(torch.__version__[:3]) >= 2.4\n\n def forward(self, x):\n if self.native_rms_norm:\n if self.weight.dtype in [torch.float16, torch.bfloat16]:\n x = x.to(self.weight.dtype)\n x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps)\n else:\n variance = x.to(torch.float32).pow(2).mean(-1, keepdim=True)\n x = x * torch.rsqrt(variance + self.eps)\n if self.weight.dtype in [torch.float16, torch.bfloat16]:\n x = x.to(self.weight.dtype)\n x = x * self.weight\n\n return x\n\n\n# AdaLayerNorm\n# return with modulated x for attn input, and params for later mlp modulation\n\n\nclass AdaLayerNorm(nn.Module):\n def __init__(self, dim):\n super().__init__()\n\n self.silu = nn.SiLU()\n self.linear = nn.Linear(dim, dim * 6)\n\n self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n\n def forward(self, x, emb=None):\n emb = self.linear(self.silu(emb))\n shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)\n\n x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]\n return x, gate_msa, shift_mlp, scale_mlp, gate_mlp\n\n\n# AdaLayerNorm for final layer\n# return only with modulated x for attn input, cuz no more mlp modulation\n\n\nclass AdaLayerNorm_Final(nn.Module):\n def __init__(self, dim):\n super().__init__()\n\n self.silu = nn.SiLU()\n self.linear = nn.Linear(dim, dim * 2)\n\n self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n\n def forward(self, x, emb):\n emb = self.linear(self.silu(emb))\n scale, shift = torch.chunk(emb, 2, dim=1)\n\n x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]\n return x\n\n\n# FeedForward\n\n\nclass FeedForward(nn.Module):\n def __init__(self, dim, dim_out=None, mult=4, dropout=0.0, approximate: str = \"none\"):\n super().__init__()\n inner_dim = int(dim * mult)\n dim_out = dim_out if dim_out is not None else dim\n\n activation = nn.GELU(approximate=approximate)\n project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation)\n self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out))\n\n def forward(self, x):\n return self.ff(x)\n\n\n# Attention with possible joint part\n# modified from diffusers/src/diffusers/models/attention_processor.py\n\n\nclass Attention(nn.Module):\n def __init__(\n self,\n processor: JointAttnProcessor | AttnProcessor,\n dim: int,\n heads: int = 8,\n dim_head: int = 64,\n dropout: float = 0.0,\n context_dim: Optional[int] = None, # if not None -> joint attention\n context_pre_only: bool = False,\n qk_norm: Optional[str] = None,\n ):\n super().__init__()\n\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n self.processor = processor\n\n self.dim = dim\n self.heads = heads\n self.inner_dim = dim_head * heads\n self.dropout = dropout\n\n self.context_dim = context_dim\n self.context_pre_only = context_pre_only\n\n self.to_q = nn.Linear(dim, self.inner_dim)\n self.to_k = nn.Linear(dim, self.inner_dim)\n self.to_v = nn.Linear(dim, self.inner_dim)\n\n if qk_norm is None:\n self.q_norm = None\n self.k_norm = None\n elif qk_norm == \"rms_norm\":\n self.q_norm = RMSNorm(dim_head, eps=1e-6)\n self.k_norm = RMSNorm(dim_head, eps=1e-6)\n else:\n raise ValueError(f\"Unimplemented qk_norm: {qk_norm}\")\n\n if self.context_dim is not None:\n self.to_q_c = nn.Linear(context_dim, self.inner_dim)\n self.to_k_c = nn.Linear(context_dim, self.inner_dim)\n self.to_v_c = nn.Linear(context_dim, self.inner_dim)\n if qk_norm is None:\n self.c_q_norm = None\n self.c_k_norm = None\n elif qk_norm == \"rms_norm\":\n self.c_q_norm = RMSNorm(dim_head, eps=1e-6)\n self.c_k_norm = RMSNorm(dim_head, eps=1e-6)\n\n self.to_out = nn.ModuleList([])\n self.to_out.append(nn.Linear(self.inner_dim, dim))\n self.to_out.append(nn.Dropout(dropout))\n\n if self.context_dim is not None and not self.context_pre_only:\n self.to_out_c = nn.Linear(self.inner_dim, context_dim)\n\n def forward(\n self,\n x: float[\"b n d\"], # noised input x\n c: float[\"b n d\"] = None, # context c\n mask: bool[\"b n\"] | None = None,\n rope=None, # rotary position embedding for x\n c_rope=None, # rotary position embedding for c\n c_mask: bool[\"b nt\"] | None = None, # text mask\n ) -> torch.Tensor:\n if c is not None:\n return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask)\n else:\n return self.processor(self, x, mask=mask, rope=rope)\n\n\n# Attention processor\n\nif is_package_available(\"flash_attn\"):\n from flash_attn import flash_attn_func, flash_attn_varlen_func\n from flash_attn.bert_padding import pad_input, unpad_input\n\n\nclass AttnProcessor:\n def __init__(\n self,\n pe_attn_head: int | None = None, # number of attention head to apply rope, None for all\n attn_backend: str = \"torch\", # \"torch\" or \"flash_attn\"\n attn_mask_enabled: bool = True,\n ):\n if attn_backend == \"flash_attn\":\n assert is_package_available(\"flash_attn\"), \"Please install flash-attn first.\"\n if attn_backend == \"torch\" and attn_mask_enabled:\n warnings.warn(\n \"attn_mask_enabled=True with attn_backend='torch' can consume large GPU memory. \"\n \"Please switch attn_backend to 'flash_attn'.\",\n UserWarning,\n )\n\n self.pe_attn_head = pe_attn_head\n self.attn_backend = attn_backend\n self.attn_mask_enabled = attn_mask_enabled\n\n def __call__(\n self,\n attn: Attention,\n x: float[\"b n d\"], # noised input x\n mask: bool[\"b n\"] | None = None,\n rope=None, # rotary position embedding\n ) -> torch.FloatTensor:\n batch_size = x.shape[0]\n\n # `sample` projections\n query = attn.to_q(x)\n key = attn.to_k(x)\n value = attn.to_v(x)\n\n # attention\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # qk norm\n if attn.q_norm is not None:\n query = attn.q_norm(query)\n if attn.k_norm is not None:\n key = attn.k_norm(key)\n\n # apply rotary position embedding\n if rope is not None:\n freqs, xpos_scale = rope\n q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)\n\n if self.pe_attn_head is not None:\n pn = self.pe_attn_head\n query[:, :pn, :, :] = apply_rotary_pos_emb(query[:, :pn, :, :], freqs, q_xpos_scale)\n key[:, :pn, :, :] = apply_rotary_pos_emb(key[:, :pn, :, :], freqs, k_xpos_scale)\n else:\n query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)\n key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)\n\n if self.attn_backend == \"torch\":\n # mask. e.g. inference got a batch with different target durations, mask out the padding\n if self.attn_mask_enabled and mask is not None:\n attn_mask = mask\n attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'\n attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])\n else:\n attn_mask = None\n x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)\n x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n\n elif self.attn_backend == \"flash_attn\":\n query = query.transpose(1, 2) # [b, h, n, d] -> [b, n, h, d]\n key = key.transpose(1, 2)\n value = value.transpose(1, 2)\n if self.attn_mask_enabled and mask is not None:\n query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask)\n key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask)\n value, _, _, _, _ = unpad_input(value, mask)\n x = flash_attn_varlen_func(\n query,\n key,\n value,\n q_cu_seqlens,\n k_cu_seqlens,\n q_max_seqlen_in_batch,\n k_max_seqlen_in_batch,\n )\n x = pad_input(x, indices, batch_size, q_max_seqlen_in_batch)\n x = x.reshape(batch_size, -1, attn.heads * head_dim)\n else:\n x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False)\n x = x.reshape(batch_size, -1, attn.heads * head_dim)\n\n x = x.to(query.dtype)\n\n # linear proj\n x = attn.to_out[0](x)\n # dropout\n x = attn.to_out[1](x)\n\n if mask is not None:\n mask = mask.unsqueeze(-1)\n x = x.masked_fill(~mask, 0.0)\n\n return x\n\n\n# Joint Attention processor for MM-DiT\n# modified from diffusers/src/diffusers/models/attention_processor.py\n\n\nclass JointAttnProcessor:\n def __init__(\n self,\n attn_backend: str = \"torch\", # \"torch\" or \"flash_attn\"\n attn_mask_enabled: bool = True,\n ):\n if attn_backend == \"flash_attn\":\n assert is_package_available(\"flash_attn\"), \"Please install flash-attn first.\"\n if attn_backend == \"torch\" and attn_mask_enabled:\n warnings.warn(\n \"attn_mask_enabled=True with attn_backend='torch' can consume large GPU memory. \"\n \"Please switch attn_backend to 'flash_attn'.\",\n UserWarning,\n )\n\n self.attn_backend = attn_backend\n self.attn_mask_enabled = attn_mask_enabled\n\n def __call__(\n self,\n attn: Attention,\n x: float[\"b n d\"], # noised input x\n c: float[\"b nt d\"] = None, # context c, here text\n mask: bool[\"b n\"] | None = None,\n rope=None, # rotary position embedding for x\n c_rope=None, # rotary position embedding for c\n c_mask: bool[\"b nt\"] | None = None, # text mask\n ) -> torch.FloatTensor:\n residual = x\n audio_mask = mask\n\n batch_size = c.shape[0]\n\n # `sample` projections\n query = attn.to_q(x)\n key = attn.to_k(x)\n value = attn.to_v(x)\n\n # `context` projections\n c_query = attn.to_q_c(c)\n c_key = attn.to_k_c(c)\n c_value = attn.to_v_c(c)\n\n # attention\n inner_dim = key.shape[-1]\n head_dim = inner_dim // attn.heads\n query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n c_query = c_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n c_key = c_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n c_value = c_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)\n\n # qk norm\n if attn.q_norm is not None:\n query = attn.q_norm(query)\n if attn.k_norm is not None:\n key = attn.k_norm(key)\n if attn.c_q_norm is not None:\n c_query = attn.c_q_norm(c_query)\n if attn.c_k_norm is not None:\n c_key = attn.c_k_norm(c_key)\n\n # apply rope for context and noised input independently\n if rope is not None:\n freqs, xpos_scale = rope\n q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)\n query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)\n key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)\n if c_rope is not None:\n freqs, xpos_scale = c_rope\n q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)\n c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)\n c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)\n\n # joint attention\n query = torch.cat([query, c_query], dim=2)\n key = torch.cat([key, c_key], dim=2)\n value = torch.cat([value, c_value], dim=2)\n\n # build combined mask for joint attention: audio mask + text mask\n if self.attn_mask_enabled and mask is not None:\n if c_mask is not None:\n mask = torch.cat([mask, c_mask], dim=1)\n else:\n mask = F.pad(mask, (0, c.shape[1]), value=True)\n\n if self.attn_backend == \"torch\":\n # mask. e.g. inference got a batch with different target durations, mask out the padding\n if self.attn_mask_enabled and mask is not None:\n attn_mask = mask\n attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'\n attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])\n else:\n attn_mask = None\n x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)\n x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)\n\n elif self.attn_backend == \"flash_attn\":\n query = query.transpose(1, 2) # [b, h, n, d] -> [b, n, h, d]\n key = key.transpose(1, 2)\n value = value.transpose(1, 2)\n if self.attn_mask_enabled and mask is not None:\n total_seq_len = query.shape[1]\n query, indices, q_cu_seqlens, q_max_seqlen_in_batch, _ = unpad_input(query, mask)\n key, _, k_cu_seqlens, k_max_seqlen_in_batch, _ = unpad_input(key, mask)\n value, _, _, _, _ = unpad_input(value, mask)\n x = flash_attn_varlen_func(\n query,\n key,\n value,\n q_cu_seqlens,\n k_cu_seqlens,\n q_max_seqlen_in_batch,\n k_max_seqlen_in_batch,\n )\n x = pad_input(x, indices, batch_size, total_seq_len)\n x = x.reshape(batch_size, -1, attn.heads * head_dim)\n else:\n x = flash_attn_func(query, key, value, dropout_p=0.0, causal=False)\n x = x.reshape(batch_size, -1, attn.heads * head_dim)\n\n x = x.to(query.dtype)\n\n # Split the attention outputs.\n x, c = (\n x[:, : residual.shape[1]],\n x[:, residual.shape[1] :],\n )\n\n # linear proj\n x = attn.to_out[0](x)\n # dropout\n x = attn.to_out[1](x)\n if not attn.context_pre_only:\n c = attn.to_out_c(c)\n\n if audio_mask is not None:\n x = x.masked_fill(~audio_mask.unsqueeze(-1), 0.0)\n if c_mask is not None:\n c = c.masked_fill(~c_mask.unsqueeze(-1), 0.0)\n\n return x, c\n\n\n# DiT Block\n\n\nclass DiTBlock(nn.Module):\n def __init__(\n self,\n dim,\n heads,\n dim_head,\n ff_mult=4,\n dropout=0.1,\n qk_norm=None,\n pe_attn_head=None,\n attn_backend=\"torch\", # \"torch\" or \"flash_attn\"\n attn_mask_enabled=True,\n ):\n super().__init__()\n\n self.attn_norm = AdaLayerNorm(dim)\n self.attn = Attention(\n processor=AttnProcessor(\n pe_attn_head=pe_attn_head,\n attn_backend=attn_backend,\n attn_mask_enabled=attn_mask_enabled,\n ),\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n dropout=dropout,\n qk_norm=qk_norm,\n )\n\n self.ff_norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate=\"tanh\")\n\n def forward(self, x, t, mask=None, rope=None): # x: noised input, t: time embedding\n # pre-norm & modulation for attention input\n norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)\n\n # attention\n attn_output = self.attn(x=norm, mask=mask, rope=rope)\n\n # process attention output for input x\n x = x + gate_msa.unsqueeze(1) * attn_output\n\n norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]\n ff_output = self.ff(norm)\n x = x + gate_mlp.unsqueeze(1) * ff_output\n\n return x\n\n\n# MMDiT Block https://arxiv.org/abs/2403.03206\n\n\nclass MMDiTBlock(nn.Module):\n r\"\"\"\n modified from diffusers/src/diffusers/models/attention.py\n\n notes.\n _c: context related. text, cond, etc. (left part in sd3 fig2.b)\n _x: noised input related. (right part)\n context_pre_only: last layer only do prenorm + modulation cuz no more ffn\n \"\"\"\n\n def __init__(\n self,\n dim,\n heads,\n dim_head,\n ff_mult=4,\n dropout=0.1,\n context_dim=None,\n context_pre_only=False,\n qk_norm=None,\n attn_backend=\"torch\",\n attn_mask_enabled=False,\n ):\n super().__init__()\n if context_dim is None:\n context_dim = dim\n self.context_pre_only = context_pre_only\n\n self.attn_norm_c = AdaLayerNorm_Final(context_dim) if context_pre_only else AdaLayerNorm(context_dim)\n self.attn_norm_x = AdaLayerNorm(dim)\n self.attn = Attention(\n processor=JointAttnProcessor(\n attn_backend=attn_backend,\n attn_mask_enabled=attn_mask_enabled,\n ),\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n dropout=dropout,\n context_dim=context_dim,\n context_pre_only=context_pre_only,\n qk_norm=qk_norm,\n )\n\n if not context_pre_only:\n self.ff_norm_c = nn.LayerNorm(context_dim, elementwise_affine=False, eps=1e-6)\n self.ff_c = FeedForward(dim=context_dim, mult=ff_mult, dropout=dropout, approximate=\"tanh\")\n else:\n self.ff_norm_c = None\n self.ff_c = None\n self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate=\"tanh\")\n\n def forward(\n self, x, c, t, mask=None, rope=None, c_rope=None, c_mask=None\n ): # x: noised input, c: context, t: time embedding\n # pre-norm & modulation for attention input\n if self.context_pre_only:\n norm_c = self.attn_norm_c(c, t)\n else:\n norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)\n norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)\n\n # attention\n x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask)\n\n # process attention output for context c\n if self.context_pre_only:\n c = None\n else: # if not last layer\n c = c + c_gate_msa.unsqueeze(1) * c_attn_output\n\n norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]\n c_ff_output = self.ff_c(norm_c)\n c = c + c_gate_mlp.unsqueeze(1) * c_ff_output\n\n # process attention output for input x\n x = x + x_gate_msa.unsqueeze(1) * x_attn_output\n\n norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]\n x_ff_output = self.ff_x(norm_x)\n x = x + x_gate_mlp.unsqueeze(1) * x_ff_output\n\n return c, x\n\n\n# time step conditioning embedding\n\n\nclass TimestepEmbedding(nn.Module):\n def __init__(self, dim, freq_embed_dim=256):\n super().__init__()\n self.time_embed = SinusPositionEmbedding(freq_embed_dim)\n self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))\n\n def forward(self, timestep: float[\"b\"]):\n time_hidden = self.time_embed(timestep)\n time_hidden = time_hidden.to(timestep.dtype)\n time = self.time_mlp(time_hidden) # b d\n return time\n" }, { "path": "src/f5_tts/model/trainer.py", "content": "from __future__ import annotations\n\nimport gc\nimport math\nimport os\n\nimport torch\nimport torchaudio\nimport wandb\nfrom accelerate import Accelerator\nfrom accelerate.utils import DistributedDataParallelKwargs\nfrom ema_pytorch import EMA\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import LinearLR, SequentialLR\nfrom torch.utils.data import DataLoader, Dataset, SequentialSampler\nfrom tqdm import tqdm\n\nfrom f5_tts.model import CFM\nfrom f5_tts.model.dataset import DynamicBatchSampler, collate_fn\nfrom f5_tts.model.utils import default, exists\n\n\n# trainer\n\n\nclass Trainer:\n def __init__(\n self,\n model: CFM,\n epochs,\n learning_rate,\n num_warmup_updates=20000,\n save_per_updates=1000,\n keep_last_n_checkpoints: int = -1, # -1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\n checkpoint_path=None,\n batch_size_per_gpu=32,\n batch_size_type: str = \"sample\",\n max_samples=32,\n grad_accumulation_steps=1,\n max_grad_norm=1.0,\n noise_scheduler: str | None = None,\n duration_predictor: torch.nn.Module | None = None,\n logger: str | None = \"wandb\", # \"wandb\" | \"tensorboard\" | None\n wandb_project=\"test_f5-tts\",\n wandb_run_name=\"test_run\",\n wandb_resume_id: str = None,\n log_samples: bool = False,\n last_per_updates=None,\n accelerate_kwargs: dict = dict(),\n ema_kwargs: dict = dict(),\n bnb_optimizer: bool = False,\n mel_spec_type: str = \"vocos\", # \"vocos\" | \"bigvgan\"\n is_local_vocoder: bool = False, # use local path vocoder\n local_vocoder_path: str = \"\", # local vocoder path\n model_cfg_dict: dict = dict(), # training config\n ):\n ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)\n\n if logger == \"wandb\" and not wandb.api.api_key:\n logger = None\n self.log_samples = log_samples\n\n self.accelerator = Accelerator(\n log_with=logger if logger == \"wandb\" else None,\n kwargs_handlers=[ddp_kwargs],\n gradient_accumulation_steps=grad_accumulation_steps,\n **accelerate_kwargs,\n )\n\n self.logger = logger\n if self.logger == \"wandb\":\n if exists(wandb_resume_id):\n init_kwargs = {\"wandb\": {\"resume\": \"allow\", \"name\": wandb_run_name, \"id\": wandb_resume_id}}\n else:\n init_kwargs = {\"wandb\": {\"resume\": \"allow\", \"name\": wandb_run_name}}\n\n if not model_cfg_dict:\n model_cfg_dict = {\n \"epochs\": epochs,\n \"learning_rate\": learning_rate,\n \"num_warmup_updates\": num_warmup_updates,\n \"batch_size_per_gpu\": batch_size_per_gpu,\n \"batch_size_type\": batch_size_type,\n \"max_samples\": max_samples,\n \"grad_accumulation_steps\": grad_accumulation_steps,\n \"max_grad_norm\": max_grad_norm,\n \"noise_scheduler\": noise_scheduler,\n \"bnb_optimizer\": bnb_optimizer,\n }\n model_cfg_dict[\"gpus\"] = self.accelerator.num_processes\n self.accelerator.init_trackers(\n project_name=wandb_project,\n init_kwargs=init_kwargs,\n config=model_cfg_dict,\n )\n\n elif self.logger == \"tensorboard\":\n from torch.utils.tensorboard import SummaryWriter\n\n self.writer = None\n if self.accelerator.is_main_process:\n self.writer = SummaryWriter(log_dir=f\"runs/{wandb_run_name}\")\n\n self.model = model\n\n if self.is_main:\n self.ema_model = EMA(model, include_online_model=False, **ema_kwargs)\n self.ema_model.to(self.accelerator.device)\n\n print(f\"Using logger: {logger}\")\n if grad_accumulation_steps > 1:\n print(\n \"Gradient accumulation checkpointing with per_updates now, old logic per_steps used with before f992c4e\"\n )\n\n self.epochs = epochs\n self.num_warmup_updates = num_warmup_updates\n self.save_per_updates = save_per_updates\n self.keep_last_n_checkpoints = keep_last_n_checkpoints\n self.last_per_updates = default(last_per_updates, save_per_updates)\n self.checkpoint_path = default(checkpoint_path, \"ckpts/test_f5-tts\")\n\n self.batch_size_per_gpu = batch_size_per_gpu\n self.batch_size_type = batch_size_type\n self.max_samples = max_samples\n self.grad_accumulation_steps = grad_accumulation_steps\n self.max_grad_norm = max_grad_norm\n\n # mel vocoder config\n self.vocoder_name = mel_spec_type\n self.is_local_vocoder = is_local_vocoder\n self.local_vocoder_path = local_vocoder_path\n\n self.noise_scheduler = noise_scheduler\n\n self.duration_predictor = duration_predictor\n\n if bnb_optimizer:\n import bitsandbytes as bnb\n\n self.optimizer = bnb.optim.AdamW8bit(model.parameters(), lr=learning_rate)\n else:\n self.optimizer = AdamW(model.parameters(), lr=learning_rate, fused=True)\n self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)\n\n @property\n def is_main(self):\n return self.accelerator.is_main_process\n\n def save_checkpoint(self, update, last=False):\n self.accelerator.wait_for_everyone()\n if self.is_main:\n checkpoint = dict(\n model_state_dict=self.accelerator.unwrap_model(self.model).state_dict(),\n optimizer_state_dict=self.optimizer.state_dict(),\n ema_model_state_dict=self.ema_model.state_dict(),\n scheduler_state_dict=self.scheduler.state_dict(),\n update=update,\n )\n if not os.path.exists(self.checkpoint_path):\n os.makedirs(self.checkpoint_path)\n if last:\n self.accelerator.save(checkpoint, f\"{self.checkpoint_path}/model_last.pt\")\n print(f\"Saved last checkpoint at update {update}\")\n else:\n if self.keep_last_n_checkpoints == 0:\n return\n self.accelerator.save(checkpoint, f\"{self.checkpoint_path}/model_{update}.pt\")\n if self.keep_last_n_checkpoints > 0:\n # Updated logic to exclude pretrained model from rotation\n checkpoints = [\n f\n for f in os.listdir(self.checkpoint_path)\n if f.startswith(\"model_\")\n and not f.startswith(\"pretrained_\") # Exclude pretrained models\n and f.endswith(\".pt\")\n and f != \"model_last.pt\"\n ]\n checkpoints.sort(key=lambda x: int(x.split(\"_\")[1].split(\".\")[0]))\n while len(checkpoints) > self.keep_last_n_checkpoints:\n oldest_checkpoint = checkpoints.pop(0)\n os.remove(os.path.join(self.checkpoint_path, oldest_checkpoint))\n print(f\"Removed old checkpoint: {oldest_checkpoint}\")\n\n def load_checkpoint(self):\n if (\n not exists(self.checkpoint_path)\n or not os.path.exists(self.checkpoint_path)\n or not any(filename.endswith((\".pt\", \".safetensors\")) for filename in os.listdir(self.checkpoint_path))\n ):\n return 0\n\n self.accelerator.wait_for_everyone()\n if \"model_last.pt\" in os.listdir(self.checkpoint_path):\n latest_checkpoint = \"model_last.pt\"\n else:\n # Updated to consider pretrained models for loading but prioritize training checkpoints\n all_checkpoints = [\n f\n for f in os.listdir(self.checkpoint_path)\n if (f.startswith(\"model_\") or f.startswith(\"pretrained_\")) and f.endswith((\".pt\", \".safetensors\"))\n ]\n\n # First try to find regular training checkpoints\n training_checkpoints = [f for f in all_checkpoints if f.startswith(\"model_\") and f != \"model_last.pt\"]\n if training_checkpoints:\n latest_checkpoint = sorted(\n training_checkpoints,\n key=lambda x: int(\"\".join(filter(str.isdigit, x))),\n )[-1]\n else:\n # If no training checkpoints, use pretrained model\n latest_checkpoint = next(f for f in all_checkpoints if f.startswith(\"pretrained_\"))\n\n if latest_checkpoint.endswith(\".safetensors\"): # always a pretrained checkpoint\n from safetensors.torch import load_file\n\n checkpoint = load_file(f\"{self.checkpoint_path}/{latest_checkpoint}\", device=\"cpu\")\n checkpoint = {\"ema_model_state_dict\": checkpoint}\n elif latest_checkpoint.endswith(\".pt\"):\n # checkpoint = torch.load(f\"{self.checkpoint_path}/{latest_checkpoint}\", map_location=self.accelerator.device) # rather use accelerator.load_state ಥ_ಥ\n checkpoint = torch.load(\n f\"{self.checkpoint_path}/{latest_checkpoint}\", weights_only=True, map_location=\"cpu\"\n )\n\n # patch for backward compatibility, 305e3ea\n for key in [\"ema_model.mel_spec.mel_stft.mel_scale.fb\", \"ema_model.mel_spec.mel_stft.spectrogram.window\"]:\n if key in checkpoint[\"ema_model_state_dict\"]:\n del checkpoint[\"ema_model_state_dict\"][key]\n\n if self.is_main:\n self.ema_model.load_state_dict(checkpoint[\"ema_model_state_dict\"])\n\n if \"update\" in checkpoint or \"step\" in checkpoint:\n # patch for backward compatibility, with before f992c4e\n if \"step\" in checkpoint:\n checkpoint[\"update\"] = checkpoint[\"step\"] // self.grad_accumulation_steps\n if self.grad_accumulation_steps > 1 and self.is_main:\n print(\n \"F5-TTS WARNING: Loading checkpoint saved with per_steps logic (before f992c4e), will convert to per_updates according to grad_accumulation_steps setting, may have unexpected behaviour.\"\n )\n # patch for backward compatibility, 305e3ea\n for key in [\"mel_spec.mel_stft.mel_scale.fb\", \"mel_spec.mel_stft.spectrogram.window\"]:\n if key in checkpoint[\"model_state_dict\"]:\n del checkpoint[\"model_state_dict\"][key]\n\n self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint[\"model_state_dict\"])\n self.optimizer.load_state_dict(checkpoint[\"optimizer_state_dict\"])\n if self.scheduler:\n self.scheduler.load_state_dict(checkpoint[\"scheduler_state_dict\"])\n update = checkpoint[\"update\"]\n else:\n checkpoint[\"model_state_dict\"] = {\n k.replace(\"ema_model.\", \"\"): v\n for k, v in checkpoint[\"ema_model_state_dict\"].items()\n if k not in [\"initted\", \"update\", \"step\"]\n }\n self.accelerator.unwrap_model(self.model).load_state_dict(checkpoint[\"model_state_dict\"])\n update = 0\n\n del checkpoint\n gc.collect()\n return update\n\n def train(self, train_dataset: Dataset, num_workers=16, resumable_with_seed: int = None):\n if self.log_samples:\n from f5_tts.infer.utils_infer import cfg_strength, load_vocoder, nfe_step, sway_sampling_coef\n\n vocoder = load_vocoder(\n vocoder_name=self.vocoder_name, is_local=self.is_local_vocoder, local_path=self.local_vocoder_path\n )\n target_sample_rate = self.accelerator.unwrap_model(self.model).mel_spec.target_sample_rate\n log_samples_path = f\"{self.checkpoint_path}/samples\"\n os.makedirs(log_samples_path, exist_ok=True)\n\n if exists(resumable_with_seed):\n generator = torch.Generator()\n generator.manual_seed(resumable_with_seed)\n else:\n generator = None\n\n if self.batch_size_type == \"sample\":\n train_dataloader = DataLoader(\n train_dataset,\n collate_fn=collate_fn,\n num_workers=num_workers,\n pin_memory=True,\n persistent_workers=True,\n batch_size=self.batch_size_per_gpu,\n shuffle=True,\n generator=generator,\n )\n elif self.batch_size_type == \"frame\":\n self.accelerator.even_batches = False\n sampler = SequentialSampler(train_dataset)\n batch_sampler = DynamicBatchSampler(\n sampler,\n self.batch_size_per_gpu,\n max_samples=self.max_samples,\n random_seed=resumable_with_seed, # This enables reproducible shuffling\n drop_residual=False,\n )\n train_dataloader = DataLoader(\n train_dataset,\n collate_fn=collate_fn,\n num_workers=num_workers,\n pin_memory=True,\n persistent_workers=True,\n batch_sampler=batch_sampler,\n )\n else:\n raise ValueError(f\"batch_size_type must be either 'sample' or 'frame', but received {self.batch_size_type}\")\n\n # accelerator.prepare() dispatches batches to devices;\n # which means the length of dataloader calculated before, should consider the number of devices\n warmup_updates = (\n self.num_warmup_updates * self.accelerator.num_processes\n ) # consider a fixed warmup steps while using accelerate multi-gpu ddp\n # otherwise by default with split_batches=False, warmup steps change with num_processes\n total_updates = math.ceil(len(train_dataloader) / self.grad_accumulation_steps) * self.epochs\n decay_updates = total_updates - warmup_updates\n warmup_scheduler = LinearLR(self.optimizer, start_factor=1e-8, end_factor=1.0, total_iters=warmup_updates)\n decay_scheduler = LinearLR(self.optimizer, start_factor=1.0, end_factor=1e-8, total_iters=decay_updates)\n self.scheduler = SequentialLR(\n self.optimizer, schedulers=[warmup_scheduler, decay_scheduler], milestones=[warmup_updates]\n )\n train_dataloader, self.scheduler = self.accelerator.prepare(\n train_dataloader, self.scheduler\n ) # actual multi_gpu updates = single_gpu updates / gpu nums\n start_update = self.load_checkpoint()\n global_update = start_update\n\n if exists(resumable_with_seed):\n orig_epoch_step = len(train_dataloader)\n start_step = start_update * self.grad_accumulation_steps\n skipped_epoch = int(start_step // orig_epoch_step)\n skipped_batch = start_step % orig_epoch_step\n skipped_dataloader = self.accelerator.skip_first_batches(train_dataloader, num_batches=skipped_batch)\n else:\n skipped_epoch = 0\n\n for epoch in range(skipped_epoch, self.epochs):\n self.model.train()\n if exists(resumable_with_seed) and epoch == skipped_epoch:\n progress_bar_initial = math.ceil(skipped_batch / self.grad_accumulation_steps)\n current_dataloader = skipped_dataloader\n else:\n progress_bar_initial = 0\n current_dataloader = train_dataloader\n\n # Set epoch for the batch sampler if it exists\n if hasattr(train_dataloader, \"batch_sampler\") and hasattr(train_dataloader.batch_sampler, \"set_epoch\"):\n train_dataloader.batch_sampler.set_epoch(epoch)\n\n progress_bar = tqdm(\n range(math.ceil(len(train_dataloader) / self.grad_accumulation_steps)),\n desc=f\"Epoch {epoch + 1}/{self.epochs}\",\n unit=\"update\",\n disable=not self.accelerator.is_local_main_process,\n initial=progress_bar_initial,\n )\n\n for batch in current_dataloader:\n with self.accelerator.accumulate(self.model):\n text_inputs = batch[\"text\"]\n mel_spec = batch[\"mel\"].permute(0, 2, 1)\n mel_lengths = batch[\"mel_lengths\"]\n\n # TODO. add duration predictor training\n if self.duration_predictor is not None and self.accelerator.is_local_main_process:\n dur_loss = self.duration_predictor(mel_spec, lens=batch.get(\"durations\"))\n self.accelerator.log({\"duration loss\": dur_loss.item()}, step=global_update)\n\n loss, cond, pred = self.model(\n mel_spec, text=text_inputs, lens=mel_lengths, noise_scheduler=self.noise_scheduler\n )\n self.accelerator.backward(loss)\n\n if self.max_grad_norm > 0 and self.accelerator.sync_gradients:\n self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)\n\n self.optimizer.step()\n self.scheduler.step()\n self.optimizer.zero_grad()\n\n if self.accelerator.sync_gradients:\n if self.is_main:\n self.ema_model.update()\n\n global_update += 1\n progress_bar.update(1)\n progress_bar.set_postfix(update=str(global_update), loss=loss.item())\n\n if self.accelerator.is_local_main_process:\n self.accelerator.log(\n {\"loss\": loss.item(), \"lr\": self.scheduler.get_last_lr()[0]}, step=global_update\n )\n if self.logger == \"tensorboard\" and self.accelerator.is_main_process:\n self.writer.add_scalar(\"loss\", loss.item(), global_update)\n self.writer.add_scalar(\"lr\", self.scheduler.get_last_lr()[0], global_update)\n\n if global_update % self.last_per_updates == 0 and self.accelerator.sync_gradients:\n self.save_checkpoint(global_update, last=True)\n\n if global_update % self.save_per_updates == 0 and self.accelerator.sync_gradients:\n self.save_checkpoint(global_update)\n\n if self.log_samples and self.accelerator.is_local_main_process:\n ref_audio_len = mel_lengths[0]\n infer_text = [\n text_inputs[0] + ([\" \"] if isinstance(text_inputs[0], list) else \" \") + text_inputs[0]\n ]\n with torch.inference_mode(), self.accelerator.autocast():\n generated, _ = self.accelerator.unwrap_model(self.model).sample(\n cond=mel_spec[0][:ref_audio_len].unsqueeze(0),\n text=infer_text,\n duration=ref_audio_len * 2,\n steps=nfe_step,\n cfg_strength=cfg_strength,\n sway_sampling_coef=sway_sampling_coef,\n )\n generated = generated.to(torch.float32)\n gen_mel_spec = generated[:, ref_audio_len:, :].permute(0, 2, 1).to(self.accelerator.device)\n ref_mel_spec = batch[\"mel\"][0, :, :ref_audio_len].unsqueeze(0)\n if self.vocoder_name == \"vocos\":\n gen_audio = vocoder.decode(gen_mel_spec).cpu()\n ref_audio = vocoder.decode(ref_mel_spec).cpu()\n elif self.vocoder_name == \"bigvgan\":\n gen_audio = vocoder(gen_mel_spec).squeeze(0).cpu()\n ref_audio = vocoder(ref_mel_spec).squeeze(0).cpu()\n\n torchaudio.save(\n f\"{log_samples_path}/update_{global_update}_gen.wav\", gen_audio, target_sample_rate\n )\n torchaudio.save(\n f\"{log_samples_path}/update_{global_update}_ref.wav\", ref_audio, target_sample_rate\n )\n self.model.train()\n\n self.save_checkpoint(global_update, last=True)\n\n self.accelerator.end_training()\n" }, { "path": "src/f5_tts/model/utils.py", "content": "# ruff: noqa: F722 F821\n\nfrom __future__ import annotations\n\nimport os\nimport random\nfrom collections import defaultdict\nfrom importlib.resources import files\n\nimport rjieba\nimport torch\nfrom pypinyin import Style, lazy_pinyin\nfrom torch.nn.utils.rnn import pad_sequence\n\n\n# seed everything\n\n\ndef seed_everything(seed=0):\n random.seed(seed)\n os.environ[\"PYTHONHASHSEED\"] = str(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n\n\n# helpers\n\n\ndef exists(v):\n return v is not None\n\n\ndef default(v, d):\n return v if exists(v) else d\n\n\ndef is_package_available(package_name: str) -> bool:\n try:\n import importlib\n\n package_exists = importlib.util.find_spec(package_name) is not None\n return package_exists\n except Exception:\n return False\n\n\n# tensor helpers\n\n\ndef lens_to_mask(t: int[\"b\"], length: int | None = None) -> bool[\"b n\"]:\n if not exists(length):\n length = t.amax()\n\n seq = torch.arange(length, device=t.device)\n return seq[None, :] < t[:, None]\n\n\ndef mask_from_start_end_indices(seq_len: int[\"b\"], start: int[\"b\"], end: int[\"b\"]):\n max_seq_len = seq_len.max().item()\n seq = torch.arange(max_seq_len, device=start.device).long()\n start_mask = seq[None, :] >= start[:, None]\n end_mask = seq[None, :] < end[:, None]\n return start_mask & end_mask\n\n\ndef mask_from_frac_lengths(seq_len: int[\"b\"], frac_lengths: float[\"b\"]):\n lengths = (frac_lengths * seq_len).long()\n max_start = seq_len - lengths\n\n rand = torch.rand_like(frac_lengths)\n start = (max_start * rand).long().clamp(min=0)\n end = start + lengths\n\n return mask_from_start_end_indices(seq_len, start, end)\n\n\ndef maybe_masked_mean(t: float[\"b n d\"], mask: bool[\"b n\"] = None) -> float[\"b d\"]:\n if not exists(mask):\n return t.mean(dim=1)\n\n t = torch.where(mask[:, :, None], t, torch.tensor(0.0, device=t.device))\n num = t.sum(dim=1)\n den = mask.float().sum(dim=1)\n\n return num / den.clamp(min=1.0)\n\n\n# simple utf-8 tokenizer, since paper went character based\ndef list_str_to_tensor(text: list[str], padding_value=-1) -> int[\"b nt\"]:\n list_tensors = [torch.tensor([*bytes(t, \"UTF-8\")]) for t in text] # ByT5 style\n text = pad_sequence(list_tensors, padding_value=padding_value, batch_first=True)\n return text\n\n\n# char tokenizer, based on custom dataset's extracted .txt file\ndef list_str_to_idx(\n text: list[str] | list[list[str]],\n vocab_char_map: dict[str, int], # {char: idx}\n padding_value=-1,\n) -> int[\"b nt\"]:\n list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text] # pinyin or char style\n text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)\n return text\n\n\n# Get tokenizer\n\n\ndef get_tokenizer(dataset_name, tokenizer: str = \"pinyin\"):\n \"\"\"\n tokenizer - \"pinyin\" do g2p for only chinese characters, need .txt vocab_file\n - \"char\" for char-wise tokenizer, need .txt vocab_file\n - \"byte\" for utf-8 tokenizer\n - \"custom\" if you're directly passing in a path to the vocab.txt you want to use\n vocab_size - if use \"pinyin\", all available pinyin types, common alphabets (also those with accent) and symbols\n - if use \"char\", derived from unfiltered character & symbol counts of custom dataset\n - if use \"byte\", set to 256 (unicode byte range)\n \"\"\"\n if tokenizer in [\"pinyin\", \"char\"]:\n tokenizer_path = os.path.join(files(\"f5_tts\").joinpath(\"../../data\"), f\"{dataset_name}_{tokenizer}/vocab.txt\")\n with open(tokenizer_path, \"r\", encoding=\"utf-8\") as f:\n vocab_char_map = {}\n for i, char in enumerate(f):\n vocab_char_map[char[:-1]] = i\n vocab_size = len(vocab_char_map)\n assert vocab_char_map[\" \"] == 0, \"make sure space is of idx 0 in vocab.txt, cuz 0 is used for unknown char\"\n\n elif tokenizer == \"byte\":\n vocab_char_map = None\n vocab_size = 256\n\n elif tokenizer == \"custom\":\n with open(dataset_name, \"r\", encoding=\"utf-8\") as f:\n vocab_char_map = {}\n for i, char in enumerate(f):\n vocab_char_map[char[:-1]] = i\n vocab_size = len(vocab_char_map)\n\n return vocab_char_map, vocab_size\n\n\n# convert char to pinyin\n\n\ndef convert_char_to_pinyin(text_list, polyphone=True):\n final_text_list = []\n custom_trans = str.maketrans(\n {\";\": \",\", \"“\": '\"', \"”\": '\"', \"‘\": \"'\", \"’\": \"'\"}\n ) # add custom trans here, to address oov\n\n def is_chinese(c):\n return (\n \"\\u3100\" <= c <= \"\\u9fff\" # common chinese characters\n )\n\n for text in text_list:\n char_list = []\n text = text.translate(custom_trans)\n for seg in rjieba.cut(text):\n seg_byte_len = len(bytes(seg, \"UTF-8\"))\n if seg_byte_len == len(seg): # if pure alphabets and symbols\n if char_list and seg_byte_len > 1 and char_list[-1] not in \" :'\\\"\":\n char_list.append(\" \")\n char_list.extend(seg)\n elif polyphone and seg_byte_len == 3 * len(seg): # if pure east asian characters\n seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)\n for i, c in enumerate(seg):\n if is_chinese(c):\n char_list.append(\" \")\n char_list.append(seg_[i])\n else: # if mixed characters, alphabets and symbols\n for c in seg:\n if ord(c) < 256:\n char_list.extend(c)\n elif is_chinese(c):\n char_list.append(\" \")\n char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))\n else:\n char_list.append(c)\n final_text_list.append(char_list)\n\n return final_text_list\n\n\n# filter func for dirty data with many repetitions\n\n\ndef repetition_found(text, length=2, tolerance=10):\n pattern_count = defaultdict(int)\n for i in range(len(text) - length + 1):\n pattern = text[i : i + length]\n pattern_count[pattern] += 1\n for pattern, count in pattern_count.items():\n if count > tolerance:\n return True\n return False\n\n\n# get the empirically pruned step for sampling\n\n\ndef get_epss_timesteps(n, device, dtype):\n dt = 1 / 32\n predefined_timesteps = {\n 5: [0, 2, 4, 8, 16, 32],\n 6: [0, 2, 4, 6, 8, 16, 32],\n 7: [0, 2, 4, 6, 8, 16, 24, 32],\n 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],\n 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],\n 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],\n }\n t = predefined_timesteps.get(n, [])\n if not t:\n return torch.linspace(0, 1, n + 1, device=device, dtype=dtype)\n return dt * torch.tensor(t, device=device, dtype=dtype)\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/.gitignore", "content": "# runtime/triton_trtllm related\nmodel.cache\nmodel_repo/\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/Dockerfile.server", "content": "FROM nvcr.io/nvidia/tritonserver:24.12-py3\nRUN pip install tritonclient[grpc] tensorrt-llm==0.16.0 torchaudio==2.5.1 rjieba pypinyin librosa vocos\nWORKDIR /workspace" }, { "path": "src/f5_tts/runtime/triton_trtllm/README.md", "content": "## Triton Inference Serving Best Practice for F5-TTS\n\n### Setup\n#### Option 1: Quick Start\n```sh\n# Directly launch the service using docker compose\nMODEL=F5TTS_v1_Base docker compose up\n```\n\n#### Option 2: Build from scratch\n```sh\n# Build the docker image\ndocker build . -f Dockerfile.server -t soar97/triton-f5-tts:24.12\n\n# Create Docker Container\nyour_mount_dir=/mnt:/mnt\ndocker run -it --name \"f5-server\" --gpus all --net host -v $your_mount_dir --shm-size=2g soar97/triton-f5-tts:24.12\n```\n\n### Build TensorRT-LLM Engines and Launch Server\nInside docker container, we would follow the official guide of TensorRT-LLM to build qwen and whisper TensorRT-LLM engines. See [here](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/whisper).\n```sh\n# F5TTS_v1_Base | F5TTS_Base | F5TTS_v1_Small | F5TTS_Small\nbash run.sh 0 4 F5TTS_v1_Base\n```\n> [!NOTE] \n> If use custom checkpoint, set `ckpt_file` and `vocab_file` in `run.sh`. \n> Remember to used matched model version (`F5TTS_v1_*` for v1, `F5TTS_*` for v0).\n> \n> If use checkpoint of different structure, see `scripts/convert_checkpoint.py`, and perform modification if necessary.\n\n> [!IMPORTANT] \n> If train or finetune with fp32, add `--dtype float32` flag when converting checkpoint in `run.sh` phase 1.\n\n### HTTP Client\n```sh\npython3 client_http.py\n```\n\n### Benchmarking\n#### Using Client-Server Mode\n```sh\n# bash run.sh 5 5 F5TTS_v1_Base\nnum_task=2\npython3 client_grpc.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name wenetspeech4tts\n```\n\n#### Using Offline TRT-LLM Mode\n```sh\n# bash run.sh 7 7 F5TTS_v1_Base\nbatch_size=1\nsplit_name=wenetspeech4tts\nbackend_type=trt\nlog_dir=./tests/benchmark_batch_size_${batch_size}_${split_name}_${backend_type}\nrm -r $log_dir\ntorchrun --nproc_per_node=1 \\\nbenchmark.py --output-dir $log_dir \\\n--batch-size $batch_size \\\n--enable-warmup \\\n--split-name $split_name \\\n--model-path $ckpt_file \\\n--vocab-file $vocab_file \\\n--vocoder-trt-engine-path $VOCODER_TRT_ENGINE_PATH \\\n--backend-type $backend_type \\\n--tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1\n```\n\n### Benchmark Results\nDecoding on a single L20 GPU, using 26 different prompt_audio & target_text pairs, 16 NFE.\n\n| Model | Concurrency | Avg Latency | RTF | Mode |\n|---------------------|----------------|-------------|--------|-----------------|\n| F5-TTS Base (Vocos) | 2 | 253 ms | 0.0394 | Client-Server |\n| F5-TTS Base (Vocos) | 1 (Batch_size) | - | 0.0402 | Offline TRT-LLM |\n| F5-TTS Base (Vocos) | 1 (Batch_size) | - | 0.1467 | Offline Pytorch |\n\n### Credits\n1. [Yuekai Zhang](https://github.com/yuekaizhang)\n2. [F5-TTS-TRTLLM](https://github.com/Bigfishering/f5-tts-trtllm)\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/benchmark.py", "content": "# Copyright (c) 2024 Tsinghua Univ. (authors: Xingchen Song)\n# 2025 (authors: Yuekai Zhang)\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Modified from https://github.com/xingchensong/S3Tokenizer/blob/main/s3tokenizer/cli.py\n\"\"\" Example Usage\ntorchrun --nproc_per_node=1 \\\nbenchmark.py --output-dir $log_dir \\\n--batch-size $batch_size \\\n--enable-warmup \\\n--split-name $split_name \\\n--model-path $CKPT_DIR/$model/model_1200000.pt \\\n--vocab-file $CKPT_DIR/$model/vocab.txt \\\n--vocoder-trt-engine-path $vocoder_trt_engine_path \\\n--backend-type $backend_type \\\n--tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1\n\"\"\"\n\nimport argparse\nimport importlib\nimport json\nimport os\nimport sys\nimport time\n\nimport datasets\nimport tensorrt as trt\nimport torch\nimport torch.distributed as dist\nimport torch.nn.functional as F\nimport torchaudio\nfrom datasets import load_dataset\nfrom huggingface_hub import hf_hub_download\nfrom tensorrt_llm._utils import trt_dtype_to_torch\nfrom tensorrt_llm.logger import logger\nfrom tensorrt_llm.runtime.session import Session, TensorInfo\nfrom torch.utils.data import DataLoader, DistributedSampler\nfrom tqdm import tqdm\nfrom vocos import Vocos\n\n\nsys.path.append(f\"{os.path.dirname(os.path.abspath(__file__))}/../../../../src/\")\n\nfrom f5_tts.eval.utils_eval import padded_mel_batch\nfrom f5_tts.model.modules import get_vocos_mel_spectrogram\nfrom f5_tts.model.utils import convert_char_to_pinyin, get_tokenizer, list_str_to_idx\n\n\nF5TTS = importlib.import_module(\"model_repo_f5_tts.f5_tts.1.f5_tts_trtllm\").F5TTS\n\ntorch.manual_seed(0)\n\n\ndef get_args():\n parser = argparse.ArgumentParser(description=\"extract speech code\")\n parser.add_argument(\n \"--split-name\",\n type=str,\n default=\"wenetspeech4tts\",\n choices=[\"wenetspeech4tts\", \"test_zh\", \"test_en\", \"test_hard\"],\n help=\"huggingface dataset split name\",\n )\n parser.add_argument(\"--output-dir\", required=True, type=str, help=\"dir to save result\")\n parser.add_argument(\n \"--vocab-file\",\n required=True,\n type=str,\n help=\"vocab file\",\n )\n parser.add_argument(\n \"--model-path\",\n required=True,\n type=str,\n help=\"model path, to load text embedding\",\n )\n parser.add_argument(\n \"--tllm-model-dir\",\n required=True,\n type=str,\n help=\"tllm model dir\",\n )\n parser.add_argument(\n \"--batch-size\",\n required=True,\n type=int,\n help=\"batch size (per-device) for inference\",\n )\n parser.add_argument(\"--num-workers\", type=int, default=0, help=\"workers for dataloader\")\n parser.add_argument(\"--prefetch\", type=int, default=None, help=\"prefetch for dataloader\")\n parser.add_argument(\n \"--vocoder\",\n default=\"vocos\",\n type=str,\n help=\"vocoder name\",\n )\n parser.add_argument(\n \"--vocoder-trt-engine-path\",\n default=None,\n type=str,\n help=\"vocoder trt engine path\",\n )\n parser.add_argument(\"--enable-warmup\", action=\"store_true\")\n parser.add_argument(\"--remove-input-padding\", action=\"store_true\")\n parser.add_argument(\"--use-perf\", action=\"store_true\", help=\"use nvtx to record performance\")\n parser.add_argument(\"--backend-type\", type=str, default=\"triton\", choices=[\"trt\", \"pytorch\"], help=\"backend type\")\n args = parser.parse_args()\n return args\n\n\ndef data_collator(batch, vocab_char_map, device=\"cuda\", use_perf=False):\n if use_perf:\n torch.cuda.nvtx.range_push(\"data_collator\")\n target_sample_rate = 24000\n target_rms = 0.1\n (\n ids,\n ref_rms_list,\n ref_mel_list,\n ref_mel_len_list,\n estimated_reference_target_mel_len,\n reference_target_texts_list,\n ) = (\n [],\n [],\n [],\n [],\n [],\n [],\n )\n for i, item in enumerate(batch):\n item_id, prompt_text, target_text = (\n item[\"id\"],\n item[\"prompt_text\"],\n item[\"target_text\"],\n )\n ids.append(item_id)\n reference_target_texts_list.append(prompt_text + target_text)\n\n ref_audio_org, ref_sr = (\n item[\"prompt_audio\"][\"array\"],\n item[\"prompt_audio\"][\"sampling_rate\"],\n )\n ref_audio_org = torch.from_numpy(ref_audio_org).unsqueeze(0).float()\n ref_rms = torch.sqrt(torch.mean(torch.square(ref_audio_org)))\n ref_rms_list.append(ref_rms)\n if ref_rms < target_rms:\n ref_audio_org = ref_audio_org * target_rms / ref_rms\n\n if ref_sr != target_sample_rate:\n resampler = torchaudio.transforms.Resample(ref_sr, target_sample_rate)\n ref_audio = resampler(ref_audio_org)\n else:\n ref_audio = ref_audio_org\n\n if use_perf:\n torch.cuda.nvtx.range_push(f\"mel_spectrogram {i}\")\n ref_audio = ref_audio.to(\"cuda\")\n ref_mel = get_vocos_mel_spectrogram(ref_audio).squeeze(0)\n if use_perf:\n torch.cuda.nvtx.range_pop()\n ref_mel_len = ref_mel.shape[-1]\n assert ref_mel.shape[0] == 100\n\n ref_mel_list.append(ref_mel)\n ref_mel_len_list.append(ref_mel_len)\n\n estimated_reference_target_mel_len.append(\n int(ref_mel_len * (1 + len(target_text.encode(\"utf-8\")) / len(prompt_text.encode(\"utf-8\"))))\n )\n\n ref_mel_batch = padded_mel_batch(ref_mel_list)\n ref_mel_len_batch = torch.LongTensor(ref_mel_len_list)\n\n pinyin_list = convert_char_to_pinyin(reference_target_texts_list, polyphone=True)\n text_pad_sequence = list_str_to_idx(pinyin_list, vocab_char_map)\n\n if use_perf:\n torch.cuda.nvtx.range_pop()\n return {\n \"ids\": ids,\n \"ref_rms_list\": ref_rms_list,\n \"ref_mel_batch\": ref_mel_batch,\n \"ref_mel_len_batch\": ref_mel_len_batch,\n \"text_pad_sequence\": text_pad_sequence,\n \"estimated_reference_target_mel_len\": estimated_reference_target_mel_len,\n }\n\n\ndef init_distributed():\n world_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n local_rank = int(os.environ.get(\"LOCAL_RANK\", 0))\n rank = int(os.environ.get(\"RANK\", 0))\n print(\n \"Inference on multiple gpus, this gpu {}\".format(local_rank)\n + \", rank {}, world_size {}\".format(rank, world_size)\n )\n torch.cuda.set_device(local_rank)\n # Initialize process group with explicit device IDs\n dist.init_process_group(\n \"nccl\",\n )\n return world_size, local_rank, rank\n\n\ndef load_vocoder(\n vocoder_name=\"vocos\", is_local=False, local_path=\"\", device=\"cuda\", hf_cache_dir=None, vocoder_trt_engine_path=None\n):\n if vocoder_name == \"vocos\":\n if vocoder_trt_engine_path is not None:\n vocoder = VocosTensorRT(engine_path=vocoder_trt_engine_path)\n else:\n # vocoder = Vocos.from_pretrained(\"charactr/vocos-mel-24khz\").to(device)\n if is_local:\n print(f\"Load vocos from local path {local_path}\")\n config_path = f\"{local_path}/config.yaml\"\n model_path = f\"{local_path}/pytorch_model.bin\"\n else:\n print(\"Download Vocos from huggingface charactr/vocos-mel-24khz\")\n repo_id = \"charactr/vocos-mel-24khz\"\n config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"config.yaml\")\n model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"pytorch_model.bin\")\n vocoder = Vocos.from_hparams(config_path)\n state_dict = torch.load(model_path, map_location=\"cpu\", weights_only=True)\n from vocos.feature_extractors import EncodecFeatures\n\n if isinstance(vocoder.feature_extractor, EncodecFeatures):\n encodec_parameters = {\n \"feature_extractor.encodec.\" + key: value\n for key, value in vocoder.feature_extractor.encodec.state_dict().items()\n }\n state_dict.update(encodec_parameters)\n vocoder.load_state_dict(state_dict)\n vocoder = vocoder.eval().to(device)\n elif vocoder_name == \"bigvgan\":\n raise NotImplementedError(\"BigVGAN is not implemented yet\")\n return vocoder\n\n\nclass VocosTensorRT:\n def __init__(self, engine_path=\"./vocos_vocoder.plan\", stream=None):\n TRT_LOGGER = trt.Logger(trt.Logger.WARNING)\n trt.init_libnvinfer_plugins(TRT_LOGGER, namespace=\"\")\n logger.info(f\"Loading vocoder engine from {engine_path}\")\n self.engine_path = engine_path\n with open(engine_path, \"rb\") as f:\n engine_buffer = f.read()\n self.session = Session.from_serialized_engine(engine_buffer)\n self.stream = stream if stream is not None else torch.cuda.current_stream().cuda_stream\n\n def decode(self, mels):\n mels = mels.contiguous()\n inputs = {\"mel\": mels}\n output_info = self.session.infer_shapes([TensorInfo(\"mel\", trt.DataType.FLOAT, mels.shape)])\n outputs = {\n t.name: torch.empty(tuple(t.shape), dtype=trt_dtype_to_torch(t.dtype), device=\"cuda\") for t in output_info\n }\n ok = self.session.run(inputs, outputs, self.stream)\n\n assert ok, \"Runtime execution failed for vae session\"\n\n samples = outputs[\"waveform\"]\n return samples\n\n\ndef main():\n args = get_args()\n os.makedirs(args.output_dir, exist_ok=True)\n\n assert torch.cuda.is_available()\n world_size, local_rank, rank = init_distributed()\n device = torch.device(f\"cuda:{local_rank}\")\n\n vocab_char_map, vocab_size = get_tokenizer(args.vocab_file, \"custom\")\n\n tllm_model_dir = args.tllm_model_dir\n with open(os.path.join(tllm_model_dir, \"config.json\")) as f:\n tllm_model_config = json.load(f)\n if args.backend_type == \"trt\":\n model = F5TTS(\n tllm_model_config,\n debug_mode=False,\n tllm_model_dir=tllm_model_dir,\n model_path=args.model_path,\n vocab_size=vocab_size,\n )\n elif args.backend_type == \"pytorch\":\n from f5_tts.infer.utils_infer import load_model\n from f5_tts.model import DiT\n\n pretrained_config = tllm_model_config[\"pretrained_config\"]\n pt_model_config = dict(\n dim=pretrained_config[\"hidden_size\"],\n depth=pretrained_config[\"num_hidden_layers\"],\n heads=pretrained_config[\"num_attention_heads\"],\n ff_mult=pretrained_config[\"ff_mult\"],\n text_dim=pretrained_config[\"text_dim\"],\n text_mask_padding=pretrained_config[\"text_mask_padding\"],\n conv_layers=pretrained_config[\"conv_layers\"],\n pe_attn_head=pretrained_config[\"pe_attn_head\"],\n # attn_backend=\"flash_attn\",\n # attn_mask_enabled=True,\n )\n model = load_model(DiT, pt_model_config, args.model_path)\n\n vocoder = load_vocoder(\n vocoder_name=args.vocoder, device=device, vocoder_trt_engine_path=args.vocoder_trt_engine_path\n )\n\n dataset = load_dataset(\n \"yuekai/seed_tts\",\n split=args.split_name,\n trust_remote_code=True,\n )\n\n def add_estimated_duration(example):\n prompt_audio_len = example[\"prompt_audio\"][\"array\"].shape[0]\n scale_factor = 1 + len(example[\"target_text\"]) / len(example[\"prompt_text\"])\n estimated_duration = prompt_audio_len * scale_factor\n example[\"estimated_duration\"] = estimated_duration / example[\"prompt_audio\"][\"sampling_rate\"]\n return example\n\n dataset = dataset.map(add_estimated_duration)\n dataset = dataset.sort(\"estimated_duration\", reverse=True)\n if args.use_perf:\n # dataset_list = [dataset.select(range(1)) for i in range(16)] # seq_len 1000\n dataset_list_short = [dataset.select([24]) for i in range(8)] # seq_len 719\n # dataset_list_long = [dataset.select([23]) for i in range(8)] # seq_len 2002\n # dataset = datasets.concatenate_datasets(dataset_list_short + dataset_list_long)\n dataset = datasets.concatenate_datasets(dataset_list_short)\n if world_size > 1:\n sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)\n else:\n # This would disable shuffling\n sampler = None\n\n dataloader = DataLoader(\n dataset,\n batch_size=args.batch_size,\n sampler=sampler,\n shuffle=False,\n num_workers=args.num_workers,\n prefetch_factor=args.prefetch,\n collate_fn=lambda x: data_collator(x, vocab_char_map, use_perf=args.use_perf),\n )\n\n total_steps = len(dataset)\n\n if args.enable_warmup:\n for batch in dataloader:\n ref_mels, ref_mel_lens = batch[\"ref_mel_batch\"].to(device), batch[\"ref_mel_len_batch\"].to(device)\n text_pad_seq = batch[\"text_pad_sequence\"].to(device)\n total_mel_lens = batch[\"estimated_reference_target_mel_len\"]\n cond_pad_seq = F.pad(ref_mels, (0, 0, 0, max(total_mel_lens) - ref_mels.shape[1], 0, 0))\n if args.backend_type == \"trt\":\n _ = model.sample(\n text_pad_seq,\n cond_pad_seq,\n ref_mel_lens,\n total_mel_lens,\n remove_input_padding=args.remove_input_padding,\n )\n elif args.backend_type == \"pytorch\":\n total_mel_lens = torch.tensor(total_mel_lens, device=device)\n with torch.inference_mode():\n generated, _ = model.sample(\n cond=ref_mels,\n text=text_pad_seq,\n duration=total_mel_lens,\n steps=32,\n cfg_strength=2.0,\n sway_sampling_coef=-1,\n )\n\n if rank == 0:\n progress_bar = tqdm(total=total_steps, desc=\"Processing\", unit=\"wavs\")\n\n decoding_time = 0\n vocoder_time = 0\n total_duration = 0\n if args.use_perf:\n torch.cuda.cudart().cudaProfilerStart()\n total_decoding_time = time.time()\n for batch in dataloader:\n if args.use_perf:\n torch.cuda.nvtx.range_push(\"data sample\")\n ref_mels, ref_mel_lens = batch[\"ref_mel_batch\"].to(device), batch[\"ref_mel_len_batch\"].to(device)\n text_pad_seq = batch[\"text_pad_sequence\"].to(device)\n total_mel_lens = batch[\"estimated_reference_target_mel_len\"]\n cond_pad_seq = F.pad(ref_mels, (0, 0, 0, max(total_mel_lens) - ref_mels.shape[1], 0, 0))\n if args.use_perf:\n torch.cuda.nvtx.range_pop()\n if args.backend_type == \"trt\":\n generated, cost_time = model.sample(\n text_pad_seq,\n cond_pad_seq,\n ref_mel_lens,\n total_mel_lens,\n remove_input_padding=args.remove_input_padding,\n use_perf=args.use_perf,\n )\n elif args.backend_type == \"pytorch\":\n total_mel_lens = torch.tensor(total_mel_lens, device=device)\n with torch.inference_mode():\n start_time = time.time()\n generated, _ = model.sample(\n cond=ref_mels,\n text=text_pad_seq,\n duration=total_mel_lens,\n lens=ref_mel_lens,\n steps=32,\n cfg_strength=2.0,\n sway_sampling_coef=-1,\n )\n cost_time = time.time() - start_time\n decoding_time += cost_time\n vocoder_start_time = time.time()\n target_rms = 0.1\n target_sample_rate = 24000\n for i, gen in enumerate(generated):\n gen = gen[ref_mel_lens[i] : total_mel_lens[i], :].unsqueeze(0)\n gen_mel_spec = gen.permute(0, 2, 1).to(torch.float32)\n if args.vocoder == \"vocos\":\n if args.use_perf:\n torch.cuda.nvtx.range_push(\"vocoder decode\")\n generated_wave = vocoder.decode(gen_mel_spec).cpu()\n if args.use_perf:\n torch.cuda.nvtx.range_pop()\n else:\n generated_wave = vocoder(gen_mel_spec).squeeze(0).cpu()\n\n if batch[\"ref_rms_list\"][i] < target_rms:\n generated_wave = generated_wave * batch[\"ref_rms_list\"][i] / target_rms\n\n utt = batch[\"ids\"][i]\n torchaudio.save(\n f\"{args.output_dir}/{utt}.wav\",\n generated_wave,\n target_sample_rate,\n )\n total_duration += generated_wave.shape[1] / target_sample_rate\n vocoder_time += time.time() - vocoder_start_time\n if rank == 0:\n progress_bar.update(world_size * len(batch[\"ids\"]))\n total_decoding_time = time.time() - total_decoding_time\n if rank == 0:\n progress_bar.close()\n rtf = total_decoding_time / total_duration\n s = f\"RTF: {rtf:.4f}\\n\"\n s += f\"total_duration: {total_duration:.3f} seconds\\n\"\n s += f\"({total_duration / 3600:.2f} hours)\\n\"\n s += f\"DiT time: {decoding_time:.3f} seconds ({decoding_time / 3600:.2f} hours)\\n\"\n s += f\"Vocoder time: {vocoder_time:.3f} seconds ({vocoder_time / 3600:.2f} hours)\\n\"\n s += f\"total decoding time: {total_decoding_time:.3f} seconds ({total_decoding_time / 3600:.2f} hours)\\n\"\n s += f\"batch size: {args.batch_size}\\n\"\n print(s)\n\n with open(f\"{args.output_dir}/rtf.txt\", \"w\") as f:\n f.write(s)\n\n dist.barrier()\n dist.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/client_grpc.py", "content": "#!/usr/bin/env python3\n# Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang)\n# 2023 Nvidia (authors: Yuekai Zhang)\n# 2023 Recurrent.ai (authors: Songtao Shi)\n# See LICENSE for clarification regarding multiple authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis script supports to load dataset from huggingface and sends it to the server\nfor decoding, in parallel.\n\nUsage:\nnum_task=2\n\n# For offline F5-TTS\npython3 client_grpc.py \\\n --server-addr localhost \\\n --model-name f5_tts \\\n --num-tasks $num_task \\\n --huggingface-dataset yuekai/seed_tts \\\n --split-name test_zh \\\n --log-dir ./log_concurrent_tasks_${num_task}\n\"\"\"\n\nimport argparse\nimport asyncio\nimport json\nimport os\nimport time\nimport types\nfrom pathlib import Path\n\nimport numpy as np\nimport soundfile as sf\nimport tritonclient\nimport tritonclient.grpc.aio as grpcclient\nfrom tritonclient.utils import np_to_triton_dtype\n\n\ndef write_triton_stats(stats, summary_file):\n with open(summary_file, \"w\") as summary_f:\n model_stats = stats[\"model_stats\"]\n # write a note, the log is from triton_client.get_inference_statistics(), to better human readability\n summary_f.write(\n \"The log is parsing from triton_client.get_inference_statistics(), to better human readability. \\n\"\n )\n summary_f.write(\"To learn more about the log, please refer to: \\n\")\n summary_f.write(\"1. https://github.com/triton-inference-server/server/blob/main/docs/user_guide/metrics.md \\n\")\n summary_f.write(\"2. https://github.com/triton-inference-server/server/issues/5374 \\n\\n\")\n summary_f.write(\n \"To better improve throughput, we always would like let requests wait in the queue for a while, and then execute them with a larger batch size. \\n\"\n )\n summary_f.write(\n \"However, there is a trade-off between the increased queue time and the increased batch size. \\n\"\n )\n summary_f.write(\n \"You may change 'max_queue_delay_microseconds' and 'preferred_batch_size' in the model configuration file to achieve this. \\n\"\n )\n summary_f.write(\n \"See https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md#delayed-batching for more details. \\n\\n\"\n )\n for model_state in model_stats:\n if \"last_inference\" not in model_state:\n continue\n summary_f.write(f\"model name is {model_state['name']} \\n\")\n model_inference_stats = model_state[\"inference_stats\"]\n total_queue_time_s = int(model_inference_stats[\"queue\"][\"ns\"]) / 1e9\n total_infer_time_s = int(model_inference_stats[\"compute_infer\"][\"ns\"]) / 1e9\n total_input_time_s = int(model_inference_stats[\"compute_input\"][\"ns\"]) / 1e9\n total_output_time_s = int(model_inference_stats[\"compute_output\"][\"ns\"]) / 1e9\n summary_f.write(\n f\"queue time {total_queue_time_s:<5.2f} s, compute infer time {total_infer_time_s:<5.2f} s, compute input time {total_input_time_s:<5.2f} s, compute output time {total_output_time_s:<5.2f} s \\n\" # noqa\n )\n model_batch_stats = model_state[\"batch_stats\"]\n for batch in model_batch_stats:\n batch_size = int(batch[\"batch_size\"])\n compute_input = batch[\"compute_input\"]\n compute_output = batch[\"compute_output\"]\n compute_infer = batch[\"compute_infer\"]\n batch_count = int(compute_infer[\"count\"])\n assert compute_infer[\"count\"] == compute_output[\"count\"] == compute_input[\"count\"]\n compute_infer_time_ms = int(compute_infer[\"ns\"]) / 1e6\n compute_input_time_ms = int(compute_input[\"ns\"]) / 1e6\n compute_output_time_ms = int(compute_output[\"ns\"]) / 1e6\n summary_f.write(\n f\"execuate inference with batch_size {batch_size:<2} total {batch_count:<5} times, total_infer_time {compute_infer_time_ms:<9.2f} ms, avg_infer_time {compute_infer_time_ms:<9.2f}/{batch_count:<5}={compute_infer_time_ms / batch_count:.2f} ms, avg_infer_time_per_sample {compute_infer_time_ms:<9.2f}/{batch_count:<5}/{batch_size}={compute_infer_time_ms / batch_count / batch_size:.2f} ms \\n\" # noqa\n )\n summary_f.write(\n f\"input {compute_input_time_ms:<9.2f} ms, avg {compute_input_time_ms / batch_count:.2f} ms, \" # noqa\n )\n summary_f.write(\n f\"output {compute_output_time_ms:<9.2f} ms, avg {compute_output_time_ms / batch_count:.2f} ms \\n\" # noqa\n )\n\n\ndef get_args():\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n\n parser.add_argument(\n \"--server-addr\",\n type=str,\n default=\"localhost\",\n help=\"Address of the server\",\n )\n\n parser.add_argument(\n \"--server-port\",\n type=int,\n default=8001,\n help=\"Grpc port of the triton server, default is 8001\",\n )\n\n parser.add_argument(\n \"--reference-audio\",\n type=str,\n default=None,\n help=\"Path to a single audio file. It can't be specified at the same time with --manifest-dir\",\n )\n\n parser.add_argument(\n \"--reference-text\",\n type=str,\n default=\"\",\n help=\"\",\n )\n\n parser.add_argument(\n \"--target-text\",\n type=str,\n default=\"\",\n help=\"\",\n )\n\n parser.add_argument(\n \"--huggingface-dataset\",\n type=str,\n default=\"yuekai/seed_tts\",\n help=\"dataset name in huggingface dataset hub\",\n )\n\n parser.add_argument(\n \"--split-name\",\n type=str,\n default=\"wenetspeech4tts\",\n choices=[\"wenetspeech4tts\", \"test_zh\", \"test_en\", \"test_hard\"],\n help=\"dataset split name, default is 'test'\",\n )\n\n parser.add_argument(\n \"--manifest-path\",\n type=str,\n default=None,\n help=\"Path to the manifest dir which includes wav.scp trans.txt files.\",\n )\n\n parser.add_argument(\n \"--model-name\",\n type=str,\n default=\"f5_tts\",\n help=\"triton model_repo module name to request\",\n )\n\n parser.add_argument(\n \"--num-tasks\",\n type=int,\n default=1,\n help=\"Number of concurrent tasks for sending\",\n )\n\n parser.add_argument(\n \"--log-interval\",\n type=int,\n default=5,\n help=\"Controls how frequently we print the log.\",\n )\n\n parser.add_argument(\n \"--compute-wer\",\n action=\"store_true\",\n default=False,\n help=\"\"\"True to compute WER.\n \"\"\",\n )\n\n parser.add_argument(\n \"--log-dir\",\n type=str,\n required=False,\n default=\"./tests/client_grpc\",\n help=\"log directory\",\n )\n\n parser.add_argument(\n \"--batch-size\",\n type=int,\n default=1,\n help=\"Inference batch_size per request for offline mode.\",\n )\n\n return parser.parse_args()\n\n\ndef load_audio(wav_path, target_sample_rate=24000):\n assert target_sample_rate == 24000, \"hard coding in server\"\n if isinstance(wav_path, dict):\n waveform = wav_path[\"array\"]\n sample_rate = wav_path[\"sampling_rate\"]\n else:\n waveform, sample_rate = sf.read(wav_path)\n if sample_rate != target_sample_rate:\n from scipy.signal import resample\n\n waveform = resample(waveform, int(len(waveform) * (target_sample_rate / sample_rate)))\n return waveform, target_sample_rate\n\n\nasync def send(\n manifest_item_list: list,\n name: str,\n triton_client: tritonclient.grpc.aio.InferenceServerClient,\n protocol_client: types.ModuleType,\n log_interval: int,\n model_name: str,\n padding_duration: int = None,\n audio_save_dir: str = \"./\",\n save_sample_rate: int = 24000,\n):\n total_duration = 0.0\n latency_data = []\n task_id = int(name[5:])\n\n print(f\"manifest_item_list: {manifest_item_list}\")\n for i, item in enumerate(manifest_item_list):\n if i % log_interval == 0:\n print(f\"{name}: {i}/{len(manifest_item_list)}\")\n waveform, sample_rate = load_audio(item[\"audio_filepath\"], target_sample_rate=24000)\n duration = len(waveform) / sample_rate\n lengths = np.array([[len(waveform)]], dtype=np.int32)\n\n reference_text, target_text = item[\"reference_text\"], item[\"target_text\"]\n\n estimated_target_duration = duration / len(reference_text) * len(target_text)\n\n if padding_duration:\n # padding to nearset 10 seconds\n samples = np.zeros(\n (\n 1,\n padding_duration\n * sample_rate\n * ((int(estimated_target_duration + duration) // padding_duration) + 1),\n ),\n dtype=np.float32,\n )\n\n samples[0, : len(waveform)] = waveform\n else:\n samples = waveform\n\n samples = samples.reshape(1, -1).astype(np.float32)\n\n inputs = [\n protocol_client.InferInput(\"reference_wav\", samples.shape, np_to_triton_dtype(samples.dtype)),\n protocol_client.InferInput(\"reference_wav_len\", lengths.shape, np_to_triton_dtype(lengths.dtype)),\n protocol_client.InferInput(\"reference_text\", [1, 1], \"BYTES\"),\n protocol_client.InferInput(\"target_text\", [1, 1], \"BYTES\"),\n ]\n inputs[0].set_data_from_numpy(samples)\n inputs[1].set_data_from_numpy(lengths)\n\n input_data_numpy = np.array([reference_text], dtype=object)\n input_data_numpy = input_data_numpy.reshape((1, 1))\n inputs[2].set_data_from_numpy(input_data_numpy)\n\n input_data_numpy = np.array([target_text], dtype=object)\n input_data_numpy = input_data_numpy.reshape((1, 1))\n inputs[3].set_data_from_numpy(input_data_numpy)\n\n outputs = [protocol_client.InferRequestedOutput(\"waveform\")]\n\n sequence_id = 100000000 + i + task_id * 10\n start = time.time()\n response = await triton_client.infer(model_name, inputs, request_id=str(sequence_id), outputs=outputs)\n\n audio = response.as_numpy(\"waveform\").reshape(-1)\n\n end = time.time() - start\n\n audio_save_path = os.path.join(audio_save_dir, f\"{item['target_audio_path']}.wav\")\n sf.write(audio_save_path, audio, save_sample_rate, \"PCM_16\")\n\n actual_duration = len(audio) / save_sample_rate\n latency_data.append((end, actual_duration))\n total_duration += actual_duration\n\n return total_duration, latency_data\n\n\ndef load_manifests(manifest_path):\n with open(manifest_path, \"r\") as f:\n manifest_list = []\n for line in f:\n assert len(line.strip().split(\"|\")) == 4\n utt, prompt_text, prompt_wav, gt_text = line.strip().split(\"|\")\n utt = Path(utt).stem\n # gt_wav = os.path.join(os.path.dirname(manifest_path), \"wavs\", utt + \".wav\")\n if not os.path.isabs(prompt_wav):\n prompt_wav = os.path.join(os.path.dirname(manifest_path), prompt_wav)\n manifest_list.append(\n {\n \"audio_filepath\": prompt_wav,\n \"reference_text\": prompt_text,\n \"target_text\": gt_text,\n \"target_audio_path\": utt,\n }\n )\n return manifest_list\n\n\ndef split_data(data, k):\n n = len(data)\n if n < k:\n print(f\"Warning: the length of the input list ({n}) is less than k ({k}). Setting k to {n}.\")\n k = n\n\n quotient = n // k\n remainder = n % k\n\n result = []\n start = 0\n for i in range(k):\n if i < remainder:\n end = start + quotient + 1\n else:\n end = start + quotient\n\n result.append(data[start:end])\n start = end\n\n return result\n\n\nasync def main():\n args = get_args()\n url = f\"{args.server_addr}:{args.server_port}\"\n\n triton_client = grpcclient.InferenceServerClient(url=url, verbose=False)\n protocol_client = grpcclient\n\n if args.reference_audio:\n args.num_tasks = 1\n args.log_interval = 1\n manifest_item_list = [\n {\n \"reference_text\": args.reference_text,\n \"target_text\": args.target_text,\n \"audio_filepath\": args.reference_audio,\n \"target_audio_path\": \"test\",\n }\n ]\n elif args.huggingface_dataset:\n import datasets\n\n dataset = datasets.load_dataset(\n args.huggingface_dataset,\n split=args.split_name,\n trust_remote_code=True,\n )\n manifest_item_list = []\n for i in range(len(dataset)):\n manifest_item_list.append(\n {\n \"audio_filepath\": dataset[i][\"prompt_audio\"],\n \"reference_text\": dataset[i][\"prompt_text\"],\n \"target_audio_path\": dataset[i][\"id\"],\n \"target_text\": dataset[i][\"target_text\"],\n }\n )\n else:\n manifest_item_list = load_manifests(args.manifest_path)\n\n args.num_tasks = min(args.num_tasks, len(manifest_item_list))\n manifest_item_list = split_data(manifest_item_list, args.num_tasks)\n\n os.makedirs(args.log_dir, exist_ok=True)\n tasks = []\n start_time = time.time()\n for i in range(args.num_tasks):\n task = asyncio.create_task(\n send(\n manifest_item_list[i],\n name=f\"task-{i}\",\n triton_client=triton_client,\n protocol_client=protocol_client,\n log_interval=args.log_interval,\n model_name=args.model_name,\n audio_save_dir=args.log_dir,\n padding_duration=1,\n save_sample_rate=24000,\n )\n )\n tasks.append(task)\n\n ans_list = await asyncio.gather(*tasks)\n\n end_time = time.time()\n elapsed = end_time - start_time\n\n total_duration = 0.0\n latency_data = []\n for ans in ans_list:\n total_duration += ans[0]\n latency_data += ans[1]\n\n rtf = elapsed / total_duration\n\n s = f\"RTF: {rtf:.4f}\\n\"\n s += f\"total_duration: {total_duration:.3f} seconds\\n\"\n s += f\"({total_duration / 3600:.2f} hours)\\n\"\n s += f\"processing time: {elapsed:.3f} seconds ({elapsed / 3600:.2f} hours)\\n\"\n\n latency_list = [chunk_end for (chunk_end, chunk_duration) in latency_data]\n latency_ms = sum(latency_list) / float(len(latency_list)) * 1000.0\n latency_variance = np.var(latency_list, dtype=np.float64) * 1000.0\n s += f\"latency_variance: {latency_variance:.2f}\\n\"\n s += f\"latency_50_percentile_ms: {np.percentile(latency_list, 50) * 1000.0:.2f}\\n\"\n s += f\"latency_90_percentile_ms: {np.percentile(latency_list, 90) * 1000.0:.2f}\\n\"\n s += f\"latency_95_percentile_ms: {np.percentile(latency_list, 95) * 1000.0:.2f}\\n\"\n s += f\"latency_99_percentile_ms: {np.percentile(latency_list, 99) * 1000.0:.2f}\\n\"\n s += f\"average_latency_ms: {latency_ms:.2f}\\n\"\n\n print(s)\n if args.manifest_path:\n name = Path(args.manifest_path).stem\n elif args.split_name:\n name = args.split_name\n with open(f\"{args.log_dir}/rtf-{name}.txt\", \"w\") as f:\n f.write(s)\n\n stats = await triton_client.get_inference_statistics(model_name=\"\", as_json=True)\n write_triton_stats(stats, f\"{args.log_dir}/stats_summary-{name}.txt\")\n\n metadata = await triton_client.get_model_config(model_name=args.model_name, as_json=True)\n with open(f\"{args.log_dir}/model_config-{name}.json\", \"w\") as f:\n json.dump(metadata, f, indent=4)\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/client_http.py", "content": "# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions\n# are met:\n# * Redistributions of source code must retain the above copyright\n# notice, this list of conditions and the following disclaimer.\n# * Redistributions in binary form must reproduce the above copyright\n# notice, this list of conditions and the following disclaimer in the\n# documentation and/or other materials provided with the distribution.\n# * Neither the name of NVIDIA CORPORATION nor the names of its\n# contributors may be used to endorse or promote products derived\n# from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\nimport argparse\nimport os\n\nimport numpy as np\nimport requests\nimport soundfile as sf\n\n\ndef get_args():\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n\n parser.add_argument(\n \"--server-url\",\n type=str,\n default=\"localhost:8000\",\n help=\"Address of the server\",\n )\n\n parser.add_argument(\n \"--reference-audio\",\n type=str,\n default=\"../../infer/examples/basic/basic_ref_en.wav\",\n help=\"Path to a single audio file. It can't be specified at the same time with --manifest-dir\",\n )\n\n parser.add_argument(\n \"--reference-text\",\n type=str,\n default=\"Some call me nature, others call me mother nature.\",\n help=\"\",\n )\n\n parser.add_argument(\n \"--target-text\",\n type=str,\n default=\"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.\",\n help=\"\",\n )\n\n parser.add_argument(\n \"--model-name\",\n type=str,\n default=\"f5_tts\",\n help=\"triton model_repo module name to request\",\n )\n\n parser.add_argument(\n \"--output-audio\",\n type=str,\n default=\"tests/client_http.wav\",\n help=\"Path to save the output audio\",\n )\n return parser.parse_args()\n\n\ndef prepare_request(\n waveform,\n reference_text,\n target_text,\n sample_rate=24000,\n audio_save_dir: str = \"./\",\n):\n assert len(waveform.shape) == 1, \"waveform should be 1D\"\n lengths = np.array([[len(waveform)]], dtype=np.int32)\n waveform = waveform.reshape(1, -1).astype(np.float32)\n\n data = {\n \"inputs\": [\n {\"name\": \"reference_wav\", \"shape\": waveform.shape, \"datatype\": \"FP32\", \"data\": waveform.tolist()},\n {\n \"name\": \"reference_wav_len\",\n \"shape\": lengths.shape,\n \"datatype\": \"INT32\",\n \"data\": lengths.tolist(),\n },\n {\"name\": \"reference_text\", \"shape\": [1, 1], \"datatype\": \"BYTES\", \"data\": [reference_text]},\n {\"name\": \"target_text\", \"shape\": [1, 1], \"datatype\": \"BYTES\", \"data\": [target_text]},\n ]\n }\n\n return data\n\n\ndef load_audio(wav_path, target_sample_rate=24000):\n assert target_sample_rate == 24000, \"hard coding in server\"\n if isinstance(wav_path, dict):\n waveform = wav_path[\"array\"]\n sample_rate = wav_path[\"sampling_rate\"]\n else:\n waveform, sample_rate = sf.read(wav_path)\n if sample_rate != target_sample_rate:\n from scipy.signal import resample\n\n waveform = resample(waveform, int(len(waveform) * (target_sample_rate / sample_rate)))\n return waveform, target_sample_rate\n\n\nif __name__ == \"__main__\":\n args = get_args()\n server_url = args.server_url\n if not server_url.startswith((\"http://\", \"https://\")):\n server_url = f\"http://{server_url}\"\n\n url = f\"{server_url}/v2/models/{args.model_name}/infer\"\n waveform, sr = load_audio(args.reference_audio)\n assert sr == 24000, \"sample rate hardcoded in server\"\n\n waveform = np.array(waveform, dtype=np.float32)\n data = prepare_request(waveform, args.reference_text, args.target_text)\n\n rsp = requests.post(\n url, headers={\"Content-Type\": \"application/json\"}, json=data, verify=False, params={\"request_id\": \"0\"}\n )\n result = rsp.json()\n audio = result[\"outputs\"][0][\"data\"]\n audio = np.array(audio, dtype=np.float32)\n os.makedirs(os.path.dirname(args.output_audio), exist_ok=True)\n sf.write(args.output_audio, audio, 24000, \"PCM_16\")\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/docker-compose.yml", "content": "services:\n tts:\n image: soar97/triton-f5-tts:24.12\n shm_size: '1gb'\n ports:\n - \"8000:8000\"\n - \"8001:8001\"\n - \"8002:8002\"\n environment:\n - PYTHONIOENCODING=utf-8\n - MODEL_ID=${MODEL_ID}\n deploy:\n resources:\n reservations:\n devices:\n - driver: nvidia\n device_ids: ['0']\n capabilities: [gpu]\n command: >\n /bin/bash -c \"pip install vocos && rm -rf F5-TTS && git clone https://github.com/SWivid/F5-TTS.git && cd F5-TTS/src/f5_tts/runtime/triton_trtllm/ && bash run.sh 0 4 $MODEL\"\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/f5_tts/1/f5_tts_trtllm.py", "content": "import math\nimport os\nimport time\nfrom functools import wraps\nfrom typing import List, Optional\n\nimport tensorrt as trt\nimport tensorrt_llm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch\nfrom tensorrt_llm.logger import logger\nfrom tensorrt_llm.runtime.session import Session\nfrom torch.nn.utils.rnn import pad_sequence\n\n\ndef remove_tensor_padding(input_tensor, input_tensor_lengths=None):\n # Audio tensor case: batch, seq_len, feature_len\n # position_ids case: batch, seq_len\n assert input_tensor_lengths is not None, \"input_tensor_lengths must be provided for 3D input_tensor\"\n\n # Initialize a list to collect valid sequences\n valid_sequences = []\n\n for i in range(input_tensor.shape[0]):\n valid_length = input_tensor_lengths[i]\n valid_sequences.append(input_tensor[i, :valid_length])\n\n # Concatenate all valid sequences along the batch dimension\n output_tensor = torch.cat(valid_sequences, dim=0).contiguous()\n return output_tensor\n\n\nclass TextEmbedding(nn.Module):\n def __init__(\n self, text_num_embeds, text_dim, mask_padding=True, conv_layers=0, conv_mult=2, precompute_max_pos=4096\n ):\n super().__init__()\n self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token\n self.mask_padding = mask_padding\n self.register_buffer(\"freqs_cis\", precompute_freqs_cis(text_dim, precompute_max_pos), persistent=False)\n self.text_blocks = nn.Sequential(*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)])\n\n def forward(self, text, seq_len, drop_text=False):\n text = text + 1\n text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens\n text = F.pad(text, (0, seq_len - text.shape[1]), value=0)\n if self.mask_padding:\n text_mask = text == 0\n\n if drop_text: # cfg for text\n text = torch.zeros_like(text)\n\n text = self.text_embed(text) # b n -> b n d\n text = text + self.freqs_cis[:seq_len, :]\n if self.mask_padding:\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n for block in self.text_blocks:\n text = block(text)\n text = text.masked_fill(text_mask.unsqueeze(-1).expand(-1, -1, text.size(-1)), 0.0)\n else:\n text = self.text_blocks(text)\n\n return text\n\n\nclass GRN(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.gamma = nn.Parameter(torch.zeros(1, 1, dim))\n self.beta = nn.Parameter(torch.zeros(1, 1, dim))\n\n def forward(self, x):\n Gx = torch.norm(x, p=2, dim=1, keepdim=True)\n Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)\n return self.gamma * (x * Nx) + self.beta + x\n\n\nclass ConvNeXtV2Block(nn.Module):\n def __init__(\n self,\n dim: int,\n intermediate_dim: int,\n dilation: int = 1,\n ):\n super().__init__()\n padding = (dilation * (7 - 1)) // 2\n self.dwconv = nn.Conv1d(\n dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation\n ) # depthwise conv\n self.norm = nn.LayerNorm(dim, eps=1e-6)\n self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers\n self.act = nn.GELU()\n self.grn = GRN(intermediate_dim)\n self.pwconv2 = nn.Linear(intermediate_dim, dim)\n\n def forward(self, x: torch.Tensor) -> torch.Tensor:\n residual = x\n x = x.transpose(1, 2) # b n d -> b d n\n x = self.dwconv(x)\n x = x.transpose(1, 2) # b d n -> b n d\n x = self.norm(x)\n x = self.pwconv1(x)\n x = self.act(x)\n x = self.grn(x)\n x = self.pwconv2(x)\n return residual + x\n\n\ndef precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, theta_rescale_factor=1.0):\n # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning\n # has some connection to NTK literature\n # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/\n # https://github.com/lucidrains/rotary-embedding-torch/blob/main/rotary_embedding_torch/rotary_embedding_torch.py\n theta *= theta_rescale_factor ** (dim / (dim - 2))\n freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))\n t = torch.arange(end, device=freqs.device) # type: ignore\n freqs = torch.outer(t, freqs).float() # type: ignore\n freqs_cos = torch.cos(freqs) # real part\n freqs_sin = torch.sin(freqs) # imaginary part\n return torch.cat([freqs_cos, freqs_sin], dim=-1)\n\n\ndef get_text_embed_dict(ckpt_path, use_ema=True):\n ckpt_type = ckpt_path.split(\".\")[-1]\n if ckpt_type == \"safetensors\":\n from safetensors.torch import load_file\n\n checkpoint = load_file(ckpt_path)\n else:\n checkpoint = torch.load(ckpt_path, map_location=\"cpu\", weights_only=True)\n\n if use_ema:\n if ckpt_type == \"safetensors\":\n checkpoint = {\"ema_model_state_dict\": checkpoint}\n checkpoint[\"model_state_dict\"] = {\n k.replace(\"ema_model.\", \"\"): v\n for k, v in checkpoint[\"ema_model_state_dict\"].items()\n if k not in [\"initted\", \"step\"]\n }\n else:\n if ckpt_type == \"safetensors\":\n checkpoint = {\"model_state_dict\": checkpoint}\n model_params = checkpoint[\"model_state_dict\"]\n\n text_embed_dict = {}\n for key in model_params.keys():\n # transformer.text_embed.text_embed.weight -> text_embed.weight\n if \"text_embed\" in key:\n text_embed_dict[key.replace(\"transformer.text_embed.\", \"\")] = model_params[key]\n return text_embed_dict\n\n\nclass F5TTS(object):\n def __init__(\n self,\n config,\n debug_mode=True,\n stream: Optional[torch.cuda.Stream] = None,\n tllm_model_dir: Optional[str] = None,\n model_path: Optional[str] = None,\n vocab_size: Optional[int] = None,\n ):\n self.dtype = config[\"pretrained_config\"][\"dtype\"]\n\n rank = tensorrt_llm.mpi_rank()\n world_size = config[\"pretrained_config\"][\"mapping\"][\"world_size\"]\n cp_size = config[\"pretrained_config\"][\"mapping\"][\"cp_size\"]\n tp_size = config[\"pretrained_config\"][\"mapping\"][\"tp_size\"]\n pp_size = config[\"pretrained_config\"][\"mapping\"][\"pp_size\"]\n assert pp_size == 1\n self.mapping = tensorrt_llm.Mapping(\n world_size=world_size, rank=rank, cp_size=cp_size, tp_size=tp_size, pp_size=1, gpus_per_node=1\n )\n\n local_rank = rank % self.mapping.gpus_per_node\n self.device = torch.device(f\"cuda:{local_rank}\")\n\n torch.cuda.set_device(self.device)\n\n self.stream = stream\n if self.stream is None:\n self.stream = torch.cuda.Stream(self.device)\n torch.cuda.set_stream(self.stream)\n\n engine_file = os.path.join(tllm_model_dir, f\"rank{rank}.engine\")\n logger.info(f\"Loading engine from {engine_file}\")\n with open(engine_file, \"rb\") as f:\n engine_buffer = f.read()\n\n assert engine_buffer is not None\n\n self.session = Session.from_serialized_engine(engine_buffer)\n\n self.debug_mode = debug_mode\n\n self.inputs = {}\n self.outputs = {}\n self.buffer_allocated = False\n\n expected_tensor_names = [\"noise\", \"cond\", \"time\", \"rope_cos\", \"rope_sin\", \"input_lengths\", \"denoised\"]\n\n found_tensor_names = [self.session.engine.get_tensor_name(i) for i in range(self.session.engine.num_io_tensors)]\n if not self.debug_mode and set(expected_tensor_names) != set(found_tensor_names):\n logger.error(\n f\"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}\"\n )\n logger.error(\n f\"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}\"\n )\n logger.error(f\"Expected tensor names: {expected_tensor_names}\")\n logger.error(f\"Found tensor names: {found_tensor_names}\")\n raise RuntimeError(\"Tensor names in engine are not the same as expected.\")\n if self.debug_mode:\n self.debug_tensors = list(set(found_tensor_names) - set(expected_tensor_names))\n\n self.max_mel_len = 4096\n self.text_embedding = TextEmbedding(\n text_num_embeds=vocab_size,\n text_dim=config[\"pretrained_config\"][\"text_dim\"],\n mask_padding=config[\"pretrained_config\"][\"text_mask_padding\"],\n conv_layers=config[\"pretrained_config\"][\"conv_layers\"],\n precompute_max_pos=self.max_mel_len,\n ).to(self.device)\n self.text_embedding.load_state_dict(get_text_embed_dict(model_path), strict=True)\n\n self.n_mel_channels = config[\"pretrained_config\"][\"mel_dim\"]\n self.head_dim = config[\"pretrained_config\"][\"dim_head\"]\n self.base_rescale_factor = 1.0\n self.interpolation_factor = 1.0\n base = 10000.0 * self.base_rescale_factor ** (self.head_dim / (self.head_dim - 2))\n inv_freq = 1.0 / (base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))\n freqs = torch.outer(torch.arange(self.max_mel_len, dtype=torch.float32), inv_freq) / self.interpolation_factor\n self.freqs = freqs.repeat_interleave(2, dim=-1).unsqueeze(0)\n self.rope_cos = self.freqs.cos().half()\n self.rope_sin = self.freqs.sin().half()\n\n self.nfe_steps = 32\n epss = {\n 5: [0, 2, 4, 8, 16, 32],\n 6: [0, 2, 4, 6, 8, 16, 32],\n 7: [0, 2, 4, 6, 8, 16, 24, 32],\n 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32],\n 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32],\n 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32],\n }\n t = 1 / 32 * torch.tensor(epss.get(self.nfe_steps, list(range(self.nfe_steps + 1))), dtype=torch.float32)\n time_step = 1 - torch.cos(torch.pi * t / 2)\n delta_t = torch.diff(time_step)\n\n freq_embed_dim = 256 # Warning: hard coding 256 here\n time_expand = torch.zeros((1, self.nfe_steps, freq_embed_dim), dtype=torch.float32)\n half_dim = freq_embed_dim // 2\n emb_factor = math.log(10000) / (half_dim - 1)\n emb_factor = 1000.0 * torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb_factor)\n for i in range(self.nfe_steps):\n emb = time_step[i] * emb_factor\n time_expand[:, i, :] = torch.cat((emb.sin(), emb.cos()), dim=-1)\n self.time_expand = time_expand.to(self.device)\n self.delta_t = torch.cat((delta_t, delta_t), dim=0).contiguous().to(self.device)\n\n def _tensor_dtype(self, name):\n # return torch dtype given tensor name for convenience\n dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name))\n return dtype\n\n def _setup(self, batch_size, seq_len):\n for i in range(self.session.engine.num_io_tensors):\n name = self.session.engine.get_tensor_name(i)\n if self.session.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:\n shape = list(self.session.engine.get_tensor_shape(name))\n shape[0] = batch_size\n shape[1] = seq_len\n self.outputs[name] = torch.empty(shape, dtype=self._tensor_dtype(name), device=self.device)\n\n self.buffer_allocated = True\n\n def cuda_stream_guard(func):\n \"\"\"Sync external stream and set current stream to the one bound to the session. Reset on exit.\"\"\"\n\n @wraps(func)\n def wrapper(self, *args, **kwargs):\n external_stream = torch.cuda.current_stream()\n if external_stream != self.stream:\n external_stream.synchronize()\n torch.cuda.set_stream(self.stream)\n ret = func(self, *args, **kwargs)\n if external_stream != self.stream:\n self.stream.synchronize()\n torch.cuda.set_stream(external_stream)\n return ret\n\n return wrapper\n\n @cuda_stream_guard\n def forward(\n self,\n noise: torch.Tensor,\n cond: torch.Tensor,\n time_expand: torch.Tensor,\n rope_cos: torch.Tensor,\n rope_sin: torch.Tensor,\n input_lengths: torch.Tensor,\n delta_t: torch.Tensor,\n use_perf: bool = False,\n ):\n if use_perf:\n torch.cuda.nvtx.range_push(\"flow matching\")\n cfg_strength = 2.0\n batch_size = noise.shape[0]\n half_batch = batch_size // 2\n noise_half = noise[:half_batch] # Store the initial half of noise\n\n input_type = str_dtype_to_torch(self.dtype)\n\n # Keep a copy of the initial tensors\n cond = cond.to(input_type)\n rope_cos = rope_cos.to(input_type)\n rope_sin = rope_sin.to(input_type)\n input_lengths = input_lengths.to(str_dtype_to_torch(\"int32\"))\n\n # Instead of iteratively updating noise within a single model context,\n # we'll do a single forward pass for each iteration with fresh context setup\n for i in range(self.nfe_steps):\n # Re-setup the buffers for clean execution\n self._setup(batch_size, noise.shape[1])\n if not self.buffer_allocated:\n raise RuntimeError(\"Buffer not allocated, please call setup first!\")\n\n # Re-create combined noises for this iteration\n current_noise = torch.cat([noise_half, noise_half], dim=0).to(input_type)\n\n # Get time step for this iteration\n current_time = time_expand[:, i].to(input_type)\n\n # Create fresh input dictionary for this iteration\n current_inputs = {\n \"noise\": current_noise,\n \"cond\": cond,\n \"time\": current_time,\n \"rope_cos\": rope_cos,\n \"rope_sin\": rope_sin,\n \"input_lengths\": input_lengths,\n }\n\n # Update inputs and set shapes\n self.inputs.clear() # Clear previous inputs\n self.inputs.update(**current_inputs)\n self.session.set_shapes(self.inputs)\n\n if use_perf:\n torch.cuda.nvtx.range_push(f\"execute {i}\")\n ok = self.session.run(self.inputs, self.outputs, self.stream.cuda_stream)\n assert ok, \"Failed to execute model\"\n # self.session.context.execute_async_v3(self.stream.cuda_stream)\n if use_perf:\n torch.cuda.nvtx.range_pop()\n # Process results\n t_scale = delta_t[i].unsqueeze(0).to(input_type)\n\n # Extract predictions\n pred_cond = self.outputs[\"denoised\"][:half_batch]\n pred_uncond = self.outputs[\"denoised\"][half_batch:]\n\n # Apply classifier-free guidance with safeguards\n guidance = pred_cond + (pred_cond - pred_uncond) * cfg_strength\n # Calculate update for noise\n noise_half = noise_half + guidance * t_scale\n if use_perf:\n torch.cuda.nvtx.range_pop()\n return noise_half\n\n def sample(\n self,\n text_pad_sequence: torch.Tensor,\n cond_pad_sequence: torch.Tensor,\n ref_mel_len_batch: torch.Tensor,\n estimated_reference_target_mel_len: List[int],\n remove_input_padding: bool = False,\n use_perf: bool = False,\n ):\n if use_perf:\n torch.cuda.nvtx.range_push(\"text embedding\")\n batch = text_pad_sequence.shape[0]\n max_seq_len = cond_pad_sequence.shape[1]\n\n # get text_embed one by one to avoid misalignment\n text_and_drop_embedding_list = []\n for i in range(batch):\n text_embedding_i = self.text_embedding(\n text_pad_sequence[i].unsqueeze(0).to(self.device),\n estimated_reference_target_mel_len[i],\n drop_text=False,\n )\n text_embedding_drop_i = self.text_embedding(\n text_pad_sequence[i].unsqueeze(0).to(self.device),\n estimated_reference_target_mel_len[i],\n drop_text=True,\n )\n text_and_drop_embedding_list.extend([text_embedding_i[0], text_embedding_drop_i[0]])\n\n # pad separately computed text_embed to form batch with max_seq_len\n text_and_drop_embedding = pad_sequence(\n text_and_drop_embedding_list,\n batch_first=True,\n padding_value=0,\n )\n text_embedding = text_and_drop_embedding[0::2]\n text_embedding_drop = text_and_drop_embedding[1::2]\n\n noise = torch.randn_like(cond_pad_sequence).to(self.device)\n rope_cos = self.rope_cos[:, :max_seq_len, :].float().repeat(batch, 1, 1)\n rope_sin = self.rope_sin[:, :max_seq_len, :].float().repeat(batch, 1, 1)\n\n cat_mel_text = torch.cat(\n (\n cond_pad_sequence,\n text_embedding,\n ),\n dim=-1,\n )\n cat_mel_text_drop = torch.cat(\n (\n torch.zeros((batch, max_seq_len, self.n_mel_channels), dtype=torch.float32).to(self.device),\n text_embedding_drop,\n ),\n dim=-1,\n )\n\n time_expand = self.time_expand.repeat(2 * batch, 1, 1).contiguous()\n\n # Convert estimated_reference_target_mel_len to tensor\n input_lengths = torch.tensor(estimated_reference_target_mel_len, dtype=torch.int32)\n\n # combine above along the batch dimension\n inputs = {\n \"noise\": torch.cat((noise, noise), dim=0).contiguous(),\n \"cond\": torch.cat((cat_mel_text, cat_mel_text_drop), dim=0).contiguous(),\n \"time_expand\": time_expand,\n \"rope_cos\": torch.cat((rope_cos, rope_cos), dim=0).contiguous(),\n \"rope_sin\": torch.cat((rope_sin, rope_sin), dim=0).contiguous(),\n \"input_lengths\": torch.cat((input_lengths, input_lengths), dim=0).contiguous(),\n \"delta_t\": self.delta_t,\n }\n if use_perf and remove_input_padding:\n torch.cuda.nvtx.range_push(\"remove input padding\")\n if remove_input_padding:\n max_seq_len = inputs[\"cond\"].shape[1]\n inputs[\"noise\"] = remove_tensor_padding(inputs[\"noise\"], inputs[\"input_lengths\"])\n inputs[\"cond\"] = remove_tensor_padding(inputs[\"cond\"], inputs[\"input_lengths\"])\n # for time_expand, convert from B,D to B,T,D by repeat\n inputs[\"time_expand\"] = inputs[\"time_expand\"].unsqueeze(1).repeat(1, max_seq_len, 1, 1)\n inputs[\"time_expand\"] = remove_tensor_padding(inputs[\"time_expand\"], inputs[\"input_lengths\"])\n inputs[\"rope_cos\"] = remove_tensor_padding(inputs[\"rope_cos\"], inputs[\"input_lengths\"])\n inputs[\"rope_sin\"] = remove_tensor_padding(inputs[\"rope_sin\"], inputs[\"input_lengths\"])\n if use_perf and remove_input_padding:\n torch.cuda.nvtx.range_pop()\n for key in inputs:\n inputs[key] = inputs[key].to(self.device)\n if use_perf:\n torch.cuda.nvtx.range_pop()\n start_time = time.time()\n denoised = self.forward(**inputs, use_perf=use_perf)\n cost_time = time.time() - start_time\n if use_perf and remove_input_padding:\n torch.cuda.nvtx.range_push(\"remove input padding output\")\n if remove_input_padding:\n denoised_list = []\n start_idx = 0\n for i in range(batch):\n denoised_list.append(denoised[start_idx : start_idx + inputs[\"input_lengths\"][i]])\n start_idx += inputs[\"input_lengths\"][i]\n if use_perf and remove_input_padding:\n torch.cuda.nvtx.range_pop()\n return denoised_list, cost_time\n return denoised, cost_time\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/f5_tts/1/model.py", "content": "# Copyright 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions\n# are met:\n# * Redistributions of source code must retain the above copyright\n# notice, this list of conditions and the following disclaimer.\n# * Redistributions in binary form must reproduce the above copyright\n# notice, this list of conditions and the following disclaimer in the\n# documentation and/or other materials provided with the distribution.\n# * Neither the name of NVIDIA CORPORATION nor the names of its\n# contributors may be used to endorse or promote products derived\n# from this software without specific prior written permission.\n#\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY\n# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\n# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR\n# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,\n# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,\n# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR\n# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY\n# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\nimport json\nimport os\n\nimport rjieba\nimport torch\nimport torchaudio\nimport triton_python_backend_utils as pb_utils\nfrom f5_tts_trtllm import F5TTS\nfrom pypinyin import Style, lazy_pinyin\nfrom torch.nn.utils.rnn import pad_sequence\nfrom torch.utils.dlpack import from_dlpack, to_dlpack\n\n\ndef get_tokenizer(vocab_file_path: str):\n \"\"\"\n tokenizer - \"pinyin\" do g2p for only chinese characters, need .txt vocab_file\n - \"char\" for char-wise tokenizer, need .txt vocab_file\n - \"byte\" for utf-8 tokenizer\n - \"custom\" if you're directly passing in a path to the vocab.txt you want to use\n vocab_size - if use \"pinyin\", all available pinyin types, common alphabets (also those with accent) and symbols\n - if use \"char\", derived from unfiltered character & symbol counts of custom dataset\n - if use \"byte\", set to 256 (unicode byte range)\n \"\"\"\n with open(vocab_file_path, \"r\", encoding=\"utf-8\") as f:\n vocab_char_map = {}\n for i, char in enumerate(f):\n vocab_char_map[char[:-1]] = i\n vocab_size = len(vocab_char_map)\n return vocab_char_map, vocab_size\n\n\ndef convert_char_to_pinyin(reference_target_texts_list, polyphone=True):\n final_reference_target_texts_list = []\n custom_trans = str.maketrans(\n {\";\": \",\", \"“\": '\"', \"”\": '\"', \"‘\": \"'\", \"’\": \"'\"}\n ) # add custom trans here, to address oov\n\n def is_chinese(c):\n return \"\\u3100\" <= c <= \"\\u9fff\" # common chinese characters\n\n for text in reference_target_texts_list:\n char_list = []\n text = text.translate(custom_trans)\n for seg in rjieba.cut(text):\n seg_byte_len = len(bytes(seg, \"UTF-8\"))\n if seg_byte_len == len(seg): # if pure alphabets and symbols\n if char_list and seg_byte_len > 1 and char_list[-1] not in \" :'\\\"\":\n char_list.append(\" \")\n char_list.extend(seg)\n elif polyphone and seg_byte_len == 3 * len(seg): # if pure east asian characters\n seg_ = lazy_pinyin(seg, style=Style.TONE3, tone_sandhi=True)\n for i, c in enumerate(seg):\n if is_chinese(c):\n char_list.append(\" \")\n char_list.append(seg_[i])\n else: # if mixed characters, alphabets and symbols\n for c in seg:\n if ord(c) < 256:\n char_list.extend(c)\n elif is_chinese(c):\n char_list.append(\" \")\n char_list.extend(lazy_pinyin(c, style=Style.TONE3, tone_sandhi=True))\n else:\n char_list.append(c)\n final_reference_target_texts_list.append(char_list)\n\n return final_reference_target_texts_list\n\n\ndef list_str_to_idx(\n text: list[str] | list[list[str]],\n vocab_char_map: dict[str, int], # {char: idx}\n padding_value=-1,\n): # noqa: F722\n list_idx_tensors = [torch.tensor([vocab_char_map.get(c, 0) for c in t]) for t in text] # pinyin or char style\n text = pad_sequence(list_idx_tensors, padding_value=padding_value, batch_first=True)\n return text\n\n\nclass TritonPythonModel:\n def initialize(self, args):\n self.use_perf = True\n self.device = torch.device(\"cuda\")\n self.target_audio_sample_rate = 24000\n self.target_rms = 0.1 # least rms when inference, normalize to if lower\n self.n_fft = 1024\n self.win_length = 1024\n self.hop_length = 256\n self.n_mel_channels = 100\n self.max_mel_len = 4096\n\n parameters = json.loads(args[\"model_config\"])[\"parameters\"]\n for key, value in parameters.items():\n parameters[key] = value[\"string_value\"]\n\n self.vocab_char_map, self.vocab_size = get_tokenizer(parameters[\"vocab_file\"])\n self.reference_sample_rate = int(parameters[\"reference_audio_sample_rate\"])\n self.resampler = torchaudio.transforms.Resample(self.reference_sample_rate, self.target_audio_sample_rate)\n\n self.tllm_model_dir = parameters[\"tllm_model_dir\"]\n config_file = os.path.join(self.tllm_model_dir, \"config.json\")\n with open(config_file) as f:\n config = json.load(f)\n self.model = F5TTS(\n config,\n debug_mode=False,\n tllm_model_dir=self.tllm_model_dir,\n model_path=parameters[\"model_path\"],\n vocab_size=self.vocab_size,\n )\n\n self.vocoder = parameters[\"vocoder\"]\n assert self.vocoder in [\"vocos\", \"bigvgan\"]\n if self.vocoder == \"vocos\":\n self.mel_stft = torchaudio.transforms.MelSpectrogram(\n sample_rate=self.target_audio_sample_rate,\n n_fft=self.n_fft,\n win_length=self.win_length,\n hop_length=self.hop_length,\n n_mels=self.n_mel_channels,\n power=1,\n center=True,\n normalized=False,\n norm=None,\n ).to(self.device)\n self.compute_mel_fn = self.get_vocos_mel_spectrogram\n elif self.vocoder == \"bigvgan\":\n self.compute_mel_fn = self.get_bigvgan_mel_spectrogram\n\n def get_vocos_mel_spectrogram(self, waveform):\n mel = self.mel_stft(waveform)\n mel = mel.clamp(min=1e-5).log()\n return mel.transpose(1, 2)\n\n def forward_vocoder(self, mel):\n mel = mel.to(torch.float32).contiguous().cpu()\n input_tensor_0 = pb_utils.Tensor.from_dlpack(\"mel\", to_dlpack(mel))\n\n inference_request = pb_utils.InferenceRequest(\n model_name=\"vocoder\", requested_output_names=[\"waveform\"], inputs=[input_tensor_0]\n )\n inference_response = inference_request.exec()\n if inference_response.has_error():\n raise pb_utils.TritonModelException(inference_response.error().message())\n else:\n waveform = pb_utils.get_output_tensor_by_name(inference_response, \"waveform\")\n waveform = torch.utils.dlpack.from_dlpack(waveform.to_dlpack()).cpu()\n\n return waveform\n\n def execute(self, requests):\n (\n reference_text_list,\n target_text_list,\n reference_target_texts_list,\n estimated_reference_target_mel_len,\n reference_mel_len,\n reference_rms_list,\n ) = [], [], [], [], [], []\n mel_features_list = []\n if self.use_perf:\n torch.cuda.nvtx.range_push(\"preprocess\")\n for request in requests:\n wav_tensor = pb_utils.get_input_tensor_by_name(request, \"reference_wav\")\n wav_lens = pb_utils.get_input_tensor_by_name(request, \"reference_wav_len\")\n\n reference_text = pb_utils.get_input_tensor_by_name(request, \"reference_text\").as_numpy()\n reference_text = reference_text[0][0].decode(\"utf-8\")\n reference_text_list.append(reference_text)\n target_text = pb_utils.get_input_tensor_by_name(request, \"target_text\").as_numpy()\n target_text = target_text[0][0].decode(\"utf-8\")\n target_text_list.append(target_text)\n\n text = reference_text + target_text\n reference_target_texts_list.append(text)\n\n wav = from_dlpack(wav_tensor.to_dlpack())\n wav_len = from_dlpack(wav_lens.to_dlpack())\n wav_len = wav_len.squeeze()\n assert wav.shape[0] == 1, \"Only support batch size 1 for now.\"\n wav = wav[:, :wav_len]\n\n ref_rms = torch.sqrt(torch.mean(torch.square(wav)))\n if ref_rms < self.target_rms:\n wav = wav * self.target_rms / ref_rms\n reference_rms_list.append(ref_rms)\n if self.reference_sample_rate != self.target_audio_sample_rate:\n wav = self.resampler(wav)\n wav = wav.to(self.device)\n if self.use_perf:\n torch.cuda.nvtx.range_push(\"compute_mel\")\n mel_features = self.compute_mel_fn(wav)\n if self.use_perf:\n torch.cuda.nvtx.range_pop()\n mel_features_list.append(mel_features)\n\n reference_mel_len.append(mel_features.shape[1])\n estimated_reference_target_mel_len.append(\n int(\n mel_features.shape[1] * (1 + len(target_text.encode(\"utf-8\")) / len(reference_text.encode(\"utf-8\")))\n )\n )\n\n max_seq_len = min(max(estimated_reference_target_mel_len), self.max_mel_len)\n\n batch = len(requests)\n mel_features = torch.zeros((batch, max_seq_len, self.n_mel_channels), dtype=torch.float32).to(self.device)\n for i, mel in enumerate(mel_features_list):\n mel_features[i, : mel.shape[1], :] = mel\n\n reference_mel_len_tensor = torch.LongTensor(reference_mel_len).to(self.device)\n\n pinyin_list = convert_char_to_pinyin(reference_target_texts_list, polyphone=True)\n text_pad_sequence = list_str_to_idx(pinyin_list, self.vocab_char_map)\n\n if self.use_perf:\n torch.cuda.nvtx.range_pop()\n\n denoised, cost_time = self.model.sample(\n text_pad_sequence,\n mel_features,\n reference_mel_len_tensor,\n estimated_reference_target_mel_len,\n remove_input_padding=False,\n use_perf=self.use_perf,\n )\n if self.use_perf:\n torch.cuda.nvtx.range_push(\"vocoder\")\n\n responses = []\n for i in range(batch):\n ref_mel_len = reference_mel_len[i]\n estimated_mel_len = estimated_reference_target_mel_len[i]\n denoised_one_item = denoised[i, ref_mel_len:estimated_mel_len, :].unsqueeze(0).transpose(1, 2)\n audio = self.forward_vocoder(denoised_one_item)\n if reference_rms_list[i] < self.target_rms:\n audio = audio * reference_rms_list[i] / self.target_rms\n\n audio = pb_utils.Tensor.from_dlpack(\"waveform\", to_dlpack(audio))\n inference_response = pb_utils.InferenceResponse(output_tensors=[audio])\n responses.append(inference_response)\n if self.use_perf:\n torch.cuda.nvtx.range_pop()\n return responses\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/f5_tts/config.pbtxt", "content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nname: \"f5_tts\"\nbackend: \"python\"\nmax_batch_size: 4\ndynamic_batching {\n max_queue_delay_microseconds: 1000\n}\nparameters [\n {\n key: \"vocab_file\"\n value: { string_value: \"${vocab}\"}\n },\n {\n key: \"model_path\", \n value: {string_value:\"${model}\"}\n },\n {\n key: \"tllm_model_dir\", \n value: {string_value:\"${trtllm}\"}\n },\n {\n key: \"reference_audio_sample_rate\", \n value: {string_value:\"24000\"}\n },\n {\n key: \"vocoder\", \n value: {string_value:\"${vocoder}\"}\n }\n]\n\ninput [\n {\n name: \"reference_wav\"\n data_type: TYPE_FP32\n dims: [-1]\n optional: True\n },\n {\n name: \"reference_wav_len\"\n data_type: TYPE_INT32\n dims: [1]\n optional: True\n },\n {\n name: \"reference_text\"\n data_type: TYPE_STRING\n dims: [1]\n },\n {\n name: \"target_text\"\n data_type: TYPE_STRING\n dims: [1]\n }\n]\noutput [\n {\n name: \"waveform\"\n data_type: TYPE_FP32\n dims: [ -1 ]\n }\n]\n\ninstance_group [\n {\n count: 1\n kind: KIND_GPU\n }\n]" }, { "path": "src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/vocoder/1/.gitkeep", "content": "" }, { "path": "src/f5_tts/runtime/triton_trtllm/model_repo_f5_tts/vocoder/config.pbtxt", "content": "name: \"vocoder\"\nbackend: \"tensorrt\"\ndefault_model_filename: \"vocoder.plan\"\nmax_batch_size: 4\n\ninput [\n {\n name: \"mel\"\n data_type: TYPE_FP32\n dims: [ 100, -1 ]\n }\n]\n\noutput [\n {\n name: \"waveform\"\n data_type: TYPE_FP32\n dims: [ -1 ]\n }\n]\n\ndynamic_batching {\n preferred_batch_size: [1, 2, 4]\n max_queue_delay_microseconds: 1\n}\n\ninstance_group [\n {\n count: 1\n kind: KIND_GPU \n }\n]" }, { "path": "src/f5_tts/runtime/triton_trtllm/patch/__init__.py", "content": "# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom .baichuan.model import BaichuanForCausalLM\nfrom .bert.model import (\n BertForQuestionAnswering,\n BertForSequenceClassification,\n BertModel,\n RobertaForQuestionAnswering,\n RobertaForSequenceClassification,\n RobertaModel,\n)\nfrom .bloom.model import BloomForCausalLM, BloomModel\nfrom .chatglm.config import ChatGLMConfig\nfrom .chatglm.model import ChatGLMForCausalLM, ChatGLMModel\nfrom .cogvlm.config import CogVLMConfig\nfrom .cogvlm.model import CogVLMForCausalLM\nfrom .commandr.model import CohereForCausalLM\nfrom .dbrx.config import DbrxConfig\nfrom .dbrx.model import DbrxForCausalLM\nfrom .deepseek_v1.model import DeepseekForCausalLM\nfrom .deepseek_v2.model import DeepseekV2ForCausalLM\nfrom .dit.model import DiT\nfrom .eagle.model import EagleForCausalLM\nfrom .enc_dec.model import DecoderModel, EncoderModel, WhisperEncoder\nfrom .f5tts.model import F5TTS\nfrom .falcon.config import FalconConfig\nfrom .falcon.model import FalconForCausalLM, FalconModel\nfrom .gemma.config import GEMMA2_ARCHITECTURE, GEMMA_ARCHITECTURE, GemmaConfig\nfrom .gemma.model import GemmaForCausalLM\nfrom .gpt.config import GPTConfig\nfrom .gpt.model import GPTForCausalLM, GPTModel\nfrom .gptj.config import GPTJConfig\nfrom .gptj.model import GPTJForCausalLM, GPTJModel\nfrom .gptneox.model import GPTNeoXForCausalLM, GPTNeoXModel\nfrom .grok.model import GrokForCausalLM\nfrom .llama.config import LLaMAConfig\nfrom .llama.model import LLaMAForCausalLM, LLaMAModel\nfrom .mamba.model import MambaForCausalLM\nfrom .medusa.config import MedusaConfig\nfrom .medusa.model import MedusaForCausalLm\nfrom .mllama.model import MLLaMAModel\nfrom .modeling_utils import PretrainedConfig, PretrainedModel, SpeculativeDecodingMode\nfrom .mpt.model import MPTForCausalLM, MPTModel\nfrom .nemotron_nas.model import DeciLMForCausalLM\nfrom .opt.model import OPTForCausalLM, OPTModel\nfrom .phi.model import PhiForCausalLM, PhiModel\nfrom .phi3.model import Phi3ForCausalLM, Phi3Model\nfrom .qwen.model import QWenForCausalLM\nfrom .recurrentgemma.model import RecurrentGemmaForCausalLM\nfrom .redrafter.model import ReDrafterForCausalLM\n\n\n__all__ = [\n \"BertModel\",\n \"BertForQuestionAnswering\",\n \"BertForSequenceClassification\",\n \"RobertaModel\",\n \"RobertaForQuestionAnswering\",\n \"RobertaForSequenceClassification\",\n \"BloomModel\",\n \"BloomForCausalLM\",\n \"DiT\",\n \"DeepseekForCausalLM\",\n \"FalconConfig\",\n \"DeepseekV2ForCausalLM\",\n \"FalconForCausalLM\",\n \"FalconModel\",\n \"GPTConfig\",\n \"GPTModel\",\n \"GPTForCausalLM\",\n \"OPTForCausalLM\",\n \"OPTModel\",\n \"LLaMAConfig\",\n \"LLaMAForCausalLM\",\n \"LLaMAModel\",\n \"MedusaConfig\",\n \"MedusaForCausalLm\",\n \"ReDrafterForCausalLM\",\n \"GPTJConfig\",\n \"GPTJModel\",\n \"GPTJForCausalLM\",\n \"GPTNeoXModel\",\n \"GPTNeoXForCausalLM\",\n \"PhiModel\",\n \"PhiConfig\",\n \"Phi3Model\",\n \"Phi3Config\",\n \"PhiForCausalLM\",\n \"Phi3ForCausalLM\",\n \"ChatGLMConfig\",\n \"ChatGLMForCausalLM\",\n \"ChatGLMModel\",\n \"BaichuanForCausalLM\",\n \"QWenConfigQWenForCausalLM\",\n \"QWenModel\",\n \"EncoderModel\",\n \"DecoderModel\",\n \"PretrainedConfig\",\n \"PretrainedModel\",\n \"WhisperEncoder\",\n \"MambaForCausalLM\",\n \"MambaConfig\",\n \"MPTForCausalLM\",\n \"MPTModel\",\n \"SkyworkForCausalLM\",\n \"GemmaConfig\",\n \"GemmaForCausalLM\",\n \"DbrxConfig\",\n \"DbrxForCausalLM\",\n \"RecurrentGemmaForCausalLM\",\n \"CogVLMConfig\",\n \"CogVLMForCausalLM\",\n \"EagleForCausalLM\",\n \"SpeculativeDecodingMode\",\n \"CohereForCausalLM\",\n \"MLLaMAModel\",\n \"F5TTS\",\n]\n\nMODEL_MAP = {\n \"GPT2LMHeadModel\": GPTForCausalLM,\n \"GPT2LMHeadCustomModel\": GPTForCausalLM,\n \"GPTBigCodeForCausalLM\": GPTForCausalLM,\n \"Starcoder2ForCausalLM\": GPTForCausalLM,\n \"FuyuForCausalLM\": GPTForCausalLM,\n \"Kosmos2ForConditionalGeneration\": GPTForCausalLM,\n \"JAISLMHeadModel\": GPTForCausalLM,\n \"GPTForCausalLM\": GPTForCausalLM,\n \"NemotronForCausalLM\": GPTForCausalLM,\n \"OPTForCausalLM\": OPTForCausalLM,\n \"BloomForCausalLM\": BloomForCausalLM,\n \"RWForCausalLM\": FalconForCausalLM,\n \"FalconForCausalLM\": FalconForCausalLM,\n \"PhiForCausalLM\": PhiForCausalLM,\n \"Phi3ForCausalLM\": Phi3ForCausalLM,\n \"Phi3VForCausalLM\": Phi3ForCausalLM,\n \"Phi3SmallForCausalLM\": Phi3ForCausalLM,\n \"PhiMoEForCausalLM\": Phi3ForCausalLM,\n \"MambaForCausalLM\": MambaForCausalLM,\n \"GPTNeoXForCausalLM\": GPTNeoXForCausalLM,\n \"GPTJForCausalLM\": GPTJForCausalLM,\n \"MPTForCausalLM\": MPTForCausalLM,\n \"GLMModel\": ChatGLMForCausalLM,\n \"ChatGLMModel\": ChatGLMForCausalLM,\n \"ChatGLMForCausalLM\": ChatGLMForCausalLM,\n \"LlamaForCausalLM\": LLaMAForCausalLM,\n \"ExaoneForCausalLM\": LLaMAForCausalLM,\n \"MistralForCausalLM\": LLaMAForCausalLM,\n \"MixtralForCausalLM\": LLaMAForCausalLM,\n \"ArcticForCausalLM\": LLaMAForCausalLM,\n \"Grok1ModelForCausalLM\": GrokForCausalLM,\n \"InternLMForCausalLM\": LLaMAForCausalLM,\n \"InternLM2ForCausalLM\": LLaMAForCausalLM,\n \"MedusaForCausalLM\": MedusaForCausalLm,\n \"ReDrafterForCausalLM\": ReDrafterForCausalLM,\n \"BaichuanForCausalLM\": BaichuanForCausalLM,\n \"BaiChuanForCausalLM\": BaichuanForCausalLM,\n \"SkyworkForCausalLM\": LLaMAForCausalLM,\n GEMMA_ARCHITECTURE: GemmaForCausalLM,\n GEMMA2_ARCHITECTURE: GemmaForCausalLM,\n \"QWenLMHeadModel\": QWenForCausalLM,\n \"QWenForCausalLM\": QWenForCausalLM,\n \"Qwen2ForCausalLM\": QWenForCausalLM,\n \"Qwen2MoeForCausalLM\": QWenForCausalLM,\n \"Qwen2ForSequenceClassification\": QWenForCausalLM,\n \"Qwen2VLForConditionalGeneration\": QWenForCausalLM,\n \"WhisperEncoder\": WhisperEncoder,\n \"EncoderModel\": EncoderModel,\n \"DecoderModel\": DecoderModel,\n \"DbrxForCausalLM\": DbrxForCausalLM,\n \"RecurrentGemmaForCausalLM\": RecurrentGemmaForCausalLM,\n \"CogVLMForCausalLM\": CogVLMForCausalLM,\n \"DiT\": DiT,\n \"DeepseekForCausalLM\": DeepseekForCausalLM,\n \"DeciLMForCausalLM\": DeciLMForCausalLM,\n \"DeepseekV2ForCausalLM\": DeepseekV2ForCausalLM,\n \"EagleForCausalLM\": EagleForCausalLM,\n \"CohereForCausalLM\": CohereForCausalLM,\n \"MllamaForConditionalGeneration\": MLLaMAModel,\n \"BertForQuestionAnswering\": BertForQuestionAnswering,\n \"BertForSequenceClassification\": BertForSequenceClassification,\n \"BertModel\": BertModel,\n \"RobertaModel\": RobertaModel,\n \"RobertaForQuestionAnswering\": RobertaForQuestionAnswering,\n \"RobertaForSequenceClassification\": RobertaForSequenceClassification,\n \"F5TTS\": F5TTS,\n}\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/patch/f5tts/model.py", "content": "from __future__ import annotations\n\nimport os\nimport sys\nfrom collections import OrderedDict\n\nimport numpy as np\nimport tensorrt as trt\nfrom tensorrt_llm._common import default_net\n\nfrom ..._utils import str_dtype_to_trt\nfrom ...functional import (\n Tensor,\n concat,\n constant,\n expand,\n shape,\n slice,\n unsqueeze,\n)\nfrom ...layers import Linear\nfrom ...module import Module, ModuleList\nfrom ...plugin import current_all_reduce_helper\nfrom ..modeling_utils import PretrainedConfig, PretrainedModel\nfrom .modules import AdaLayerNormZero_Final, ConvPositionEmbedding, DiTBlock, TimestepEmbedding\n\n\ncurrent_file_path = os.path.abspath(__file__)\nparent_dir = os.path.dirname(current_file_path)\nsys.path.append(parent_dir)\n\n\nclass InputEmbedding(Module):\n def __init__(self, mel_dim, text_dim, out_dim):\n super().__init__()\n self.proj = Linear(mel_dim * 2 + text_dim, out_dim)\n self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)\n\n def forward(self, x, cond, mask=None):\n x = self.proj(concat([x, cond], dim=-1))\n return self.conv_pos_embed(x, mask=mask) + x\n\n\nclass F5TTS(PretrainedModel):\n def __init__(self, config: PretrainedConfig):\n super().__init__(config)\n self.dtype = str_dtype_to_trt(config.dtype)\n\n self.time_embed = TimestepEmbedding(config.hidden_size)\n self.input_embed = InputEmbedding(config.mel_dim, config.text_dim, config.hidden_size)\n\n self.dim = config.hidden_size\n self.depth = config.num_hidden_layers\n self.transformer_blocks = ModuleList(\n [\n DiTBlock(\n dim=self.dim,\n heads=config.num_attention_heads,\n dim_head=config.dim_head,\n ff_mult=config.ff_mult,\n dropout=config.dropout,\n pe_attn_head=config.pe_attn_head,\n )\n for _ in range(self.depth)\n ]\n )\n\n self.norm_out = AdaLayerNormZero_Final(config.hidden_size) # final modulation\n self.proj_out = Linear(config.hidden_size, config.mel_dim)\n\n def forward(\n self,\n noise, # nosied input audio\n cond, # masked cond audio\n time, # time step\n rope_cos,\n rope_sin,\n input_lengths,\n scale=1.0,\n ):\n if default_net().plugin_config.remove_input_padding:\n mask = None\n else:\n N = shape(noise, 1)\n B = shape(noise, 0)\n seq_len_2d = concat([1, N])\n max_position_embeddings = 4096\n # create position ids\n position_ids_buffer = constant(np.expand_dims(np.arange(max_position_embeddings).astype(np.int32), 0))\n tmp_position_ids = slice(position_ids_buffer, starts=[0, 0], sizes=seq_len_2d)\n tmp_position_ids = expand(tmp_position_ids, concat([B, N])) # [B, N]\n tmp_input_lengths = unsqueeze(input_lengths, 1) # [B, 1]\n tmp_input_lengths = expand(tmp_input_lengths, concat([B, N])) # [B, N]\n mask = tmp_position_ids < tmp_input_lengths # [B, N]\n mask = mask.cast(\"int32\")\n\n t = self.time_embed(time)\n x = self.input_embed(noise, cond, mask=mask)\n for block in self.transformer_blocks:\n x = block(x, t, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask)\n denoise = self.proj_out(self.norm_out(x, t))\n denoise.mark_output(\"denoised\", self.dtype)\n return denoise\n\n def prepare_inputs(self, **kwargs):\n max_batch_size = kwargs[\"max_batch_size\"]\n batch_size_range = [2, 2, max_batch_size]\n mel_size = self.config.mel_dim\n max_seq_len = 3000 # 4096\n num_frames_range = [mel_size * 2, max_seq_len * 2, max_seq_len * max_batch_size]\n concat_feature_dim = mel_size + self.config.text_dim\n freq_embed_dim = 256 # Warning: hard coding 256 here\n head_dim = self.config.dim_head\n mapping = self.config.mapping\n if mapping.tp_size > 1:\n current_all_reduce_helper().set_workspace_tensor(mapping, 1)\n if default_net().plugin_config.remove_input_padding:\n noise = Tensor(\n name=\"noise\",\n dtype=self.dtype,\n shape=[-1, mel_size],\n dim_range=OrderedDict(\n [\n (\"num_frames\", [num_frames_range]),\n (\"n_mels\", [mel_size]),\n ]\n ),\n )\n cond = Tensor(\n name=\"cond\",\n dtype=self.dtype,\n shape=[-1, concat_feature_dim],\n dim_range=OrderedDict(\n [\n (\"num_frames\", [num_frames_range]),\n (\"embeded_length\", [concat_feature_dim]),\n ]\n ),\n )\n time = Tensor(\n name=\"time\",\n dtype=self.dtype,\n shape=[-1, freq_embed_dim],\n dim_range=OrderedDict(\n [\n (\"num_frames\", [num_frames_range]),\n (\"freq_dim\", [freq_embed_dim]),\n ]\n ),\n )\n rope_cos = Tensor(\n name=\"rope_cos\",\n dtype=self.dtype,\n shape=[-1, head_dim],\n dim_range=OrderedDict(\n [\n (\"num_frames\", [num_frames_range]),\n (\"head_dim\", [head_dim]),\n ]\n ),\n )\n rope_sin = Tensor(\n name=\"rope_sin\",\n dtype=self.dtype,\n shape=[-1, head_dim],\n dim_range=OrderedDict(\n [\n (\"num_frames\", [num_frames_range]),\n (\"head_dim\", [head_dim]),\n ]\n ),\n )\n\n else:\n noise = Tensor(\n name=\"noise\",\n dtype=self.dtype,\n shape=[-1, -1, mel_size],\n dim_range=OrderedDict(\n [\n (\"batch_size\", [batch_size_range]),\n (\"max_duratuion\", [[100, max_seq_len // 2, max_seq_len]]),\n (\"n_mels\", [mel_size]),\n ]\n ),\n )\n cond = Tensor(\n name=\"cond\",\n dtype=self.dtype,\n shape=[-1, -1, concat_feature_dim],\n dim_range=OrderedDict(\n [\n (\"batch_size\", [batch_size_range]),\n (\"max_duratuion\", [[100, max_seq_len // 2, max_seq_len]]),\n (\"embeded_length\", [concat_feature_dim]),\n ]\n ),\n )\n time = Tensor(\n name=\"time\",\n dtype=self.dtype,\n shape=[-1, freq_embed_dim],\n dim_range=OrderedDict(\n [\n (\"batch_size\", [batch_size_range]),\n (\"freq_dim\", [freq_embed_dim]),\n ]\n ),\n )\n rope_cos = Tensor(\n name=\"rope_cos\",\n dtype=self.dtype,\n shape=[-1, -1, head_dim],\n dim_range=OrderedDict(\n [\n (\"batch_size\", [batch_size_range]),\n (\"max_duratuion\", [[100, max_seq_len // 2, max_seq_len]]),\n (\"head_dim\", [head_dim]),\n ]\n ),\n )\n rope_sin = Tensor(\n name=\"rope_sin\",\n dtype=self.dtype,\n shape=[-1, -1, head_dim],\n dim_range=OrderedDict(\n [\n (\"batch_size\", [batch_size_range]),\n (\"max_duratuion\", [[100, max_seq_len // 2, max_seq_len]]),\n (\"head_dim\", [head_dim]),\n ]\n ),\n )\n input_lengths = Tensor(\n name=\"input_lengths\",\n dtype=trt.int32,\n shape=[-1],\n dim_range=OrderedDict([(\"batch_size\", [batch_size_range])]),\n )\n return {\n \"noise\": noise,\n \"cond\": cond,\n \"time\": time,\n \"rope_cos\": rope_cos,\n \"rope_sin\": rope_sin,\n \"input_lengths\": input_lengths,\n }\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/patch/f5tts/modules.py", "content": "from __future__ import annotations\n\nimport math\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom tensorrt_llm._common import default_net\n\nfrom ..._utils import str_dtype_to_trt, trt_dtype_to_np\nfrom ...functional import (\n Tensor,\n bert_attention,\n cast,\n chunk,\n concat,\n constant,\n expand_dims,\n expand_dims_like,\n expand_mask,\n gelu,\n matmul,\n permute,\n shape,\n silu,\n slice,\n softmax,\n squeeze,\n unsqueeze,\n view,\n)\nfrom ...layers import ColumnLinear, Conv1d, LayerNorm, Linear, Mish, RowLinear\nfrom ...module import Module\n\n\nclass FeedForward(Module):\n def __init__(self, dim, dim_out=None, mult=4, dropout=0.0):\n super().__init__()\n inner_dim = int(dim * mult)\n dim_out = dim_out if dim_out is not None else dim\n\n self.project_in = Linear(dim, inner_dim)\n self.ff = Linear(inner_dim, dim_out)\n\n def forward(self, x):\n return self.ff(gelu(self.project_in(x)))\n\n\nclass AdaLayerNormZero(Module):\n def __init__(self, dim):\n super().__init__()\n\n self.linear = Linear(dim, dim * 6)\n self.norm = LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n\n def forward(self, x, emb=None):\n emb = self.linear(silu(emb))\n shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = chunk(emb, 6, dim=1)\n x = self.norm(x)\n ones = constant(np.ones(1, dtype=np.float32)).cast(x.dtype)\n if default_net().plugin_config.remove_input_padding:\n x = x * (ones + scale_msa) + shift_msa\n else:\n x = x * (ones + unsqueeze(scale_msa, 1)) + unsqueeze(shift_msa, 1)\n return x, gate_msa, shift_mlp, scale_mlp, gate_mlp\n\n\nclass AdaLayerNormZero_Final(Module):\n def __init__(self, dim):\n super().__init__()\n\n self.linear = Linear(dim, dim * 2)\n\n self.norm = LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n\n def forward(self, x, emb):\n emb = self.linear(silu(emb))\n scale, shift = chunk(emb, 2, dim=1)\n ones = constant(np.ones(1, dtype=np.float32)).cast(x.dtype)\n if default_net().plugin_config.remove_input_padding:\n x = self.norm(x) * (ones + scale) + shift\n else:\n x = self.norm(x) * unsqueeze((ones + scale), 1)\n x = x + unsqueeze(shift, 1)\n return x\n\n\nclass ConvPositionEmbedding(Module):\n def __init__(self, dim, kernel_size=31, groups=16):\n super().__init__()\n assert kernel_size % 2 != 0\n self.conv1d1 = Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2)\n self.conv1d2 = Conv1d(dim, dim, kernel_size, groups=groups, padding=kernel_size // 2)\n self.mish = Mish()\n\n def forward(self, x, mask=None):\n if default_net().plugin_config.remove_input_padding:\n x = unsqueeze(x, 0)\n if mask is not None:\n mask = mask.view(concat([shape(mask, 0), 1, shape(mask, 1)])) # [B 1 N]\n mask = expand_dims_like(mask, x) # [B D N]\n mask = cast(mask, x.dtype)\n x = permute(x, [0, 2, 1]) # [B D N]\n\n if mask is not None:\n x = self.mish(self.conv1d2(self.mish(self.conv1d1(x * mask) * mask)) * mask)\n else:\n x = self.mish(self.conv1d2(self.mish(self.conv1d1(x))))\n\n x = permute(x, [0, 2, 1]) # [B N D]\n if default_net().plugin_config.remove_input_padding:\n x = squeeze(x, 0)\n return x\n\n\nclass Attention(Module):\n def __init__(\n self,\n processor: AttnProcessor,\n dim: int,\n heads: int = 16,\n dim_head: int = 64,\n dropout: float = 0.0,\n context_dim: Optional[int] = None, # if not None -> joint attention\n context_pre_only=None,\n ):\n super().__init__()\n\n if not hasattr(F, \"scaled_dot_product_attention\"):\n raise ImportError(\"Attention equires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.\")\n\n self.processor = processor\n\n self.dim = dim # hidden_size\n self.heads = heads\n self.inner_dim = dim_head * heads\n self.dropout = dropout\n self.attention_head_size = dim_head\n self.context_dim = context_dim\n self.context_pre_only = context_pre_only\n self.tp_size = 1\n self.num_attention_heads = heads // self.tp_size\n self.num_attention_kv_heads = heads // self.tp_size # 8\n self.dtype = str_dtype_to_trt(\"float32\")\n self.attention_hidden_size = self.attention_head_size * self.num_attention_heads\n self.to_q = ColumnLinear(\n dim,\n self.tp_size * self.num_attention_heads * self.attention_head_size,\n bias=True,\n dtype=self.dtype,\n tp_group=None,\n tp_size=self.tp_size,\n )\n self.to_k = ColumnLinear(\n dim,\n self.tp_size * self.num_attention_heads * self.attention_head_size,\n bias=True,\n dtype=self.dtype,\n tp_group=None,\n tp_size=self.tp_size,\n )\n self.to_v = ColumnLinear(\n dim,\n self.tp_size * self.num_attention_heads * self.attention_head_size,\n bias=True,\n dtype=self.dtype,\n tp_group=None,\n tp_size=self.tp_size,\n )\n\n if self.context_dim is not None:\n self.to_k_c = Linear(context_dim, self.inner_dim)\n self.to_v_c = Linear(context_dim, self.inner_dim)\n if self.context_pre_only is not None:\n self.to_q_c = Linear(context_dim, self.inner_dim)\n\n self.to_out = RowLinear(\n self.tp_size * self.num_attention_heads * self.attention_head_size,\n dim,\n bias=True,\n dtype=self.dtype,\n tp_group=None,\n tp_size=self.tp_size,\n )\n\n if self.context_pre_only is not None and not self.context_pre_only:\n self.to_out_c = Linear(self.inner_dim, dim)\n\n def forward(\n self,\n x, # noised input x\n rope_cos,\n rope_sin,\n input_lengths,\n mask=None,\n c=None, # context c\n scale=1.0,\n rope=None,\n c_rope=None, # rotary position embedding for c\n ) -> torch.Tensor:\n if c is not None:\n return self.processor(self, x, c=c, input_lengths=input_lengths, scale=scale, rope=rope, c_rope=c_rope)\n else:\n return self.processor(\n self, x, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale\n )\n\n\ndef rotate_every_two_3dim(tensor: Tensor) -> Tensor:\n shape_tensor = concat(\n [shape(tensor, i) / 2 if i == (tensor.ndim() - 1) else shape(tensor, i) for i in range(tensor.ndim())]\n )\n if default_net().plugin_config.remove_input_padding:\n assert tensor.ndim() == 2\n x1 = slice(tensor, [0, 0], shape_tensor, [1, 2])\n x2 = slice(tensor, [0, 1], shape_tensor, [1, 2])\n x1 = expand_dims(x1, 2)\n x2 = expand_dims(x2, 2)\n zero = constant(np.ascontiguousarray(np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype))))\n x2 = zero - x2\n x = concat([x2, x1], 2)\n out = view(x, concat([shape(x, 0), shape(x, 1) * 2]))\n else:\n assert tensor.ndim() == 3\n\n x1 = slice(tensor, [0, 0, 0], shape_tensor, [1, 1, 2])\n x2 = slice(tensor, [0, 0, 1], shape_tensor, [1, 1, 2])\n x1 = expand_dims(x1, 3)\n x2 = expand_dims(x2, 3)\n zero = constant(np.ascontiguousarray(np.zeros([1], dtype=trt_dtype_to_np(tensor.dtype))))\n x2 = zero - x2\n x = concat([x2, x1], 3)\n out = view(x, concat([shape(x, 0), shape(x, 1), shape(x, 2) * 2]))\n\n return out\n\n\ndef apply_rotary_pos_emb_3dim(x, rope_cos, rope_sin, pe_attn_head):\n full_dim = x.size(-1)\n head_dim = rope_cos.size(-1) # attn head dim, e.g. 64\n if pe_attn_head is None:\n pe_attn_head = full_dim // head_dim\n rotated_dim = head_dim * pe_attn_head\n\n rotated_and_unrotated_list = []\n\n if default_net().plugin_config.remove_input_padding: # for [N, D] input\n new_t_shape = concat([shape(x, 0), head_dim]) # (2, -1, 64)\n\n for i in range(pe_attn_head):\n x_slice_i = slice(x, [0, i * 64], new_t_shape, [1, 1])\n x_rotated_i = x_slice_i * rope_cos + rotate_every_two_3dim(x_slice_i) * rope_sin\n rotated_and_unrotated_list.append(x_rotated_i)\n\n new_t_unrotated_shape = concat([shape(x, 0), full_dim - rotated_dim]) # (2, -1, 1024 - 64 * pe_attn_head)\n x_unrotated = slice(x, concat([0, rotated_dim]), new_t_unrotated_shape, [1, 1])\n rotated_and_unrotated_list.append(x_unrotated)\n\n else: # for [B, N, D] input\n new_t_shape = concat([shape(x, 0), shape(x, 1), head_dim]) # (2, -1, 64)\n\n for i in range(pe_attn_head):\n x_slice_i = slice(x, [0, 0, i * 64], new_t_shape, [1, 1, 1])\n x_rotated_i = x_slice_i * rope_cos + rotate_every_two_3dim(x_slice_i) * rope_sin\n rotated_and_unrotated_list.append(x_rotated_i)\n\n new_t_unrotated_shape = concat(\n [shape(x, 0), shape(x, 1), full_dim - rotated_dim]\n ) # (2, -1, 1024 - 64 * pe_attn_head)\n x_unrotated = slice(x, concat([0, 0, rotated_dim]), new_t_unrotated_shape, [1, 1, 1])\n rotated_and_unrotated_list.append(x_unrotated)\n\n out = concat(rotated_and_unrotated_list, dim=-1)\n\n return out\n\n\nclass AttnProcessor:\n def __init__(\n self,\n pe_attn_head: Optional[int] = None, # number of attention head to apply rope, None for all\n ):\n self.pe_attn_head = pe_attn_head\n\n def __call__(\n self,\n attn,\n x, # noised input x\n rope_cos,\n rope_sin,\n input_lengths,\n scale=1.0,\n rope=None,\n mask=None,\n ) -> torch.FloatTensor:\n query = attn.to_q(x)\n key = attn.to_k(x)\n value = attn.to_v(x)\n # k,v,q all (2,1226,1024)\n query = apply_rotary_pos_emb_3dim(query, rope_cos, rope_sin, self.pe_attn_head)\n key = apply_rotary_pos_emb_3dim(key, rope_cos, rope_sin, self.pe_attn_head)\n\n # attention\n inner_dim = key.shape[-1]\n norm_factor = math.sqrt(attn.attention_head_size)\n q_scaling = 1.0 / norm_factor\n if default_net().plugin_config.remove_input_padding:\n mask = None\n\n if default_net().plugin_config.bert_attention_plugin:\n qkv = concat([query, key, value], dim=-1)\n # TRT plugin mode\n assert input_lengths is not None\n if default_net().plugin_config.remove_input_padding:\n qkv = qkv.view(concat([-1, 3 * inner_dim]))\n max_input_length = constant(\n np.zeros(\n [\n 2048,\n ],\n dtype=np.int32,\n )\n )\n else:\n max_input_length = None\n context = bert_attention(\n qkv,\n input_lengths,\n attn.num_attention_heads,\n attn.attention_head_size,\n q_scaling=q_scaling,\n max_input_length=max_input_length,\n )\n else:\n assert not default_net().plugin_config.remove_input_padding\n\n def transpose_for_scores(x):\n new_x_shape = concat([shape(x, 0), shape(x, 1), attn.num_attention_heads, attn.attention_head_size])\n\n y = x.view(new_x_shape)\n y = y.transpose(1, 2)\n return y\n\n def transpose_for_scores_k(x):\n new_x_shape = concat([shape(x, 0), shape(x, 1), attn.num_attention_heads, attn.attention_head_size])\n\n y = x.view(new_x_shape)\n y = y.permute([0, 2, 3, 1])\n return y\n\n query = transpose_for_scores(query)\n key = transpose_for_scores_k(key)\n value = transpose_for_scores(value)\n\n attention_scores = matmul(query, key, use_fp32_acc=False)\n\n if mask is not None:\n attention_mask = expand_mask(mask, shape(query, 2))\n attention_mask = cast(attention_mask, attention_scores.dtype)\n attention_scores = attention_scores + attention_mask\n\n attention_probs = softmax(attention_scores, dim=-1)\n\n context = matmul(attention_probs, value, use_fp32_acc=False).transpose(1, 2)\n context = context.view(concat([shape(context, 0), shape(context, 1), attn.attention_hidden_size]))\n context = attn.to_out(context)\n if mask is not None:\n mask = mask.view(concat([shape(mask, 0), shape(mask, 1), 1]))\n mask = expand_dims_like(mask, context)\n mask = cast(mask, context.dtype)\n context = context * mask\n return context\n\n\n# DiT Block\nclass DiTBlock(Module):\n def __init__(self, dim, heads, dim_head, ff_mult=2, dropout=0.1, pe_attn_head=None):\n super().__init__()\n\n self.attn_norm = AdaLayerNormZero(dim)\n self.attn = Attention(\n processor=AttnProcessor(pe_attn_head=pe_attn_head),\n dim=dim,\n heads=heads,\n dim_head=dim_head,\n dropout=dropout,\n )\n\n self.ff_norm = LayerNorm(dim, elementwise_affine=False, eps=1e-6)\n self.ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout)\n\n def forward(\n self, x, t, rope_cos, rope_sin, input_lengths, scale=1.0, rope=ModuleNotFoundError, mask=None\n ): # x: noised input, t: time embedding\n # pre-norm & modulation for attention input\n norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, emb=t)\n # attention\n # norm ----> (2,1226,1024)\n attn_output = self.attn(\n x=norm, rope_cos=rope_cos, rope_sin=rope_sin, input_lengths=input_lengths, scale=scale, mask=mask\n )\n # process attention output for input x\n if default_net().plugin_config.remove_input_padding:\n x = x + gate_msa * attn_output\n else:\n x = x + unsqueeze(gate_msa, 1) * attn_output\n ones = constant(np.ones(1, dtype=np.float32)).cast(x.dtype)\n if default_net().plugin_config.remove_input_padding:\n norm = self.ff_norm(x) * (ones + scale_mlp) + shift_mlp\n else:\n norm = self.ff_norm(x) * (ones + unsqueeze(scale_mlp, 1)) + unsqueeze(shift_mlp, 1)\n # norm = self.ff_norm(x) * (ones + scale_mlp) + shift_mlp\n ff_output = self.ff(norm)\n if default_net().plugin_config.remove_input_padding:\n x = x + gate_mlp * ff_output\n else:\n x = x + unsqueeze(gate_mlp, 1) * ff_output\n\n return x\n\n\nclass TimestepEmbedding(Module):\n def __init__(self, dim, freq_embed_dim=256, dtype=None):\n super().__init__()\n # self.time_embed = SinusPositionEmbedding(freq_embed_dim)\n self.mlp1 = Linear(freq_embed_dim, dim, bias=True, dtype=dtype)\n self.mlp2 = Linear(dim, dim, bias=True, dtype=dtype)\n\n def forward(self, timestep):\n t_freq = self.mlp1(timestep)\n t_freq = silu(t_freq)\n t_emb = self.mlp2(t_freq)\n return t_emb\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/run.sh", "content": "stage=$1\nstop_stage=$2\nmodel=$3 # F5TTS_v1_Base | F5TTS_Base | F5TTS_v1_Small | F5TTS_Small\nif [ -z \"$model\" ]; then\n model=F5TTS_v1_Base\nfi\necho \"Start stage: $stage, Stop stage: $stop_stage, Model: $model\"\nexport CUDA_VISIBLE_DEVICES=0\n\nCKPT_DIR=../../../../ckpts\nTRTLLM_CKPT_DIR=$CKPT_DIR/$model/trtllm_ckpt\nTRTLLM_ENGINE_DIR=$CKPT_DIR/$model/trtllm_engine\n\nVOCODER_ONNX_PATH=$CKPT_DIR/vocos_vocoder.onnx\nVOCODER_TRT_ENGINE_PATH=$CKPT_DIR/vocos_vocoder.plan\nMODEL_REPO=./model_repo\n\nif [ $stage -le 0 ] && [ $stop_stage -ge 0 ]; then\n echo \"Downloading F5-TTS from huggingface\"\n huggingface-cli download SWivid/F5-TTS $model/model_*.* $model/vocab.txt --local-dir $CKPT_DIR\nfi\n\nckpt_file=$(ls $CKPT_DIR/$model/model_*.* 2>/dev/null | sort -V | tail -1) # default select latest update\nvocab_file=$CKPT_DIR/$model/vocab.txt\n\nif [ $stage -le 1 ] && [ $stop_stage -ge 1 ]; then\n echo \"Converting checkpoint\"\n python3 scripts/convert_checkpoint.py \\\n --pytorch_ckpt $ckpt_file \\\n --output_dir $TRTLLM_CKPT_DIR --model_name $model\n python_package_path=/usr/local/lib/python3.12/dist-packages\n cp -r patch/* $python_package_path/tensorrt_llm/models\n trtllm-build --checkpoint_dir $TRTLLM_CKPT_DIR \\\n --max_batch_size 8 \\\n --output_dir $TRTLLM_ENGINE_DIR --remove_input_padding disable\nfi\n\nif [ $stage -le 2 ] && [ $stop_stage -ge 2 ]; then\n echo \"Exporting vocos vocoder\"\n python3 scripts/export_vocoder_to_onnx.py --vocoder vocos --output-path $VOCODER_ONNX_PATH\n bash scripts/export_vocos_trt.sh $VOCODER_ONNX_PATH $VOCODER_TRT_ENGINE_PATH\nfi\n\nif [ $stage -le 3 ] && [ $stop_stage -ge 3 ]; then\n echo \"Building triton server\"\n rm -r $MODEL_REPO\n cp -r ./model_repo_f5_tts $MODEL_REPO\n python3 scripts/fill_template.py -i $MODEL_REPO/f5_tts/config.pbtxt vocab:$vocab_file,model:$ckpt_file,trtllm:$TRTLLM_ENGINE_DIR,vocoder:vocos\n cp $VOCODER_TRT_ENGINE_PATH $MODEL_REPO/vocoder/1/vocoder.plan\nfi\n\nif [ $stage -le 4 ] && [ $stop_stage -ge 4 ]; then\n echo \"Starting triton server\"\n tritonserver --model-repository=$MODEL_REPO\nfi\n\nif [ $stage -le 5 ] && [ $stop_stage -ge 5 ]; then\n echo \"Testing triton server\"\n num_task=1\n split_name=wenetspeech4tts\n log_dir=./tests/client_grpc_${model}_concurrent_${num_task}_${split_name}\n rm -r $log_dir\n python3 client_grpc.py --num-tasks $num_task --huggingface-dataset yuekai/seed_tts --split-name $split_name --log-dir $log_dir\nfi\n\nif [ $stage -le 6 ] && [ $stop_stage -ge 6 ]; then\n echo \"Testing http client\"\n audio=../../infer/examples/basic/basic_ref_en.wav\n reference_text=\"Some call me nature, others call me mother nature.\"\n target_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.\"\n python3 client_http.py --reference-audio $audio --reference-text \"$reference_text\" --target-text \"$target_text\" --output-audio \"./tests/client_http_$model.wav\"\nfi\n\nif [ $stage -le 7 ] && [ $stop_stage -ge 7 ]; then\n echo \"TRT-LLM: offline decoding benchmark test\"\n batch_size=2\n split_name=wenetspeech4tts\n backend_type=trt\n log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type}\n rm -r $log_dir\n torchrun --nproc_per_node=1 \\\n benchmark.py --output-dir $log_dir \\\n --batch-size $batch_size \\\n --enable-warmup \\\n --split-name $split_name \\\n --model-path $ckpt_file \\\n --vocab-file $vocab_file \\\n --vocoder-trt-engine-path $VOCODER_TRT_ENGINE_PATH \\\n --backend-type $backend_type \\\n --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1\nfi\n\nif [ $stage -le 8 ] && [ $stop_stage -ge 8 ]; then\n echo \"Native Pytorch: offline decoding benchmark test\"\n if ! python3 -c \"import f5_tts\" &> /dev/null; then\n pip install -e ../../../../\n fi\n batch_size=1 # set attn_mask_enabled=True if batching in actual use case\n split_name=wenetspeech4tts\n backend_type=pytorch\n log_dir=./tests/benchmark_${model}_batch_size_${batch_size}_${split_name}_${backend_type}\n rm -r $log_dir\n torchrun --nproc_per_node=1 \\\n benchmark.py --output-dir $log_dir \\\n --batch-size $batch_size \\\n --split-name $split_name \\\n --enable-warmup \\\n --model-path $ckpt_file \\\n --vocab-file $vocab_file \\\n --backend-type $backend_type \\\n --tllm-model-dir $TRTLLM_ENGINE_DIR || exit 1\nfi" }, { "path": "src/f5_tts/runtime/triton_trtllm/scripts/conv_stft.py", "content": "# Modified from https://github.com/echocatzh/conv-stft/blob/master/conv_stft/conv_stft.py\n\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# MIT License\n\n# Copyright (c) 2020 Shimin Zhang\n\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\nimport torch as th\nimport torch.nn.functional as F\nfrom scipy.signal import check_COLA, get_window\n\n\nsupport_clp_op = None\nif th.__version__ >= \"1.7.0\":\n from torch.fft import rfft as fft\n\n support_clp_op = True\nelse:\n from torch import rfft as fft\n\n\nclass STFT(th.nn.Module):\n def __init__(\n self,\n win_len=1024,\n win_hop=512,\n fft_len=1024,\n enframe_mode=\"continue\",\n win_type=\"hann\",\n win_sqrt=False,\n pad_center=True,\n ):\n \"\"\"\n Implement of STFT using 1D convolution and 1D transpose convolutions.\n Implement of framing the signal in 2 ways, `break` and `continue`.\n `break` method is a kaldi-like framing.\n `continue` method is a librosa-like framing.\n\n More information about `perfect reconstruction`:\n 1. https://ww2.mathworks.cn/help/signal/ref/stft.html\n 2. https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.get_window.html\n\n Args:\n win_len (int): Number of points in one frame. Defaults to 1024.\n win_hop (int): Number of framing stride. Defaults to 512.\n fft_len (int): Number of DFT points. Defaults to 1024.\n enframe_mode (str, optional): `break` and `continue`. Defaults to 'continue'.\n win_type (str, optional): The type of window to create. Defaults to 'hann'.\n win_sqrt (bool, optional): using square root window. Defaults to True.\n pad_center (bool, optional): `perfect reconstruction` opts. Defaults to True.\n \"\"\"\n super(STFT, self).__init__()\n assert enframe_mode in [\"break\", \"continue\"]\n assert fft_len >= win_len\n self.win_len = win_len\n self.win_hop = win_hop\n self.fft_len = fft_len\n self.mode = enframe_mode\n self.win_type = win_type\n self.win_sqrt = win_sqrt\n self.pad_center = pad_center\n self.pad_amount = self.fft_len // 2\n\n en_k, fft_k, ifft_k, ola_k = self.__init_kernel__()\n self.register_buffer(\"en_k\", en_k)\n self.register_buffer(\"fft_k\", fft_k)\n self.register_buffer(\"ifft_k\", ifft_k)\n self.register_buffer(\"ola_k\", ola_k)\n\n def __init_kernel__(self):\n \"\"\"\n Generate enframe_kernel, fft_kernel, ifft_kernel and overlap-add kernel.\n ** enframe_kernel: Using conv1d layer and identity matrix.\n ** fft_kernel: Using linear layer for matrix multiplication. In fact,\n enframe_kernel and fft_kernel can be combined, But for the sake of\n readability, I took the two apart.\n ** ifft_kernel, pinv of fft_kernel.\n ** overlap-add kernel, just like enframe_kernel, but transposed.\n\n Returns:\n tuple: four kernels.\n \"\"\"\n enframed_kernel = th.eye(self.fft_len)[:, None, :]\n if support_clp_op:\n tmp = fft(th.eye(self.fft_len))\n fft_kernel = th.stack([tmp.real, tmp.imag], dim=2)\n else:\n fft_kernel = fft(th.eye(self.fft_len), 1)\n if self.mode == \"break\":\n enframed_kernel = th.eye(self.win_len)[:, None, :]\n fft_kernel = fft_kernel[: self.win_len]\n fft_kernel = th.cat((fft_kernel[:, :, 0], fft_kernel[:, :, 1]), dim=1)\n ifft_kernel = th.pinverse(fft_kernel)[:, None, :]\n window = get_window(self.win_type, self.win_len)\n\n self.perfect_reconstruct = check_COLA(window, self.win_len, self.win_len - self.win_hop)\n window = th.FloatTensor(window)\n if self.mode == \"continue\":\n left_pad = (self.fft_len - self.win_len) // 2\n right_pad = left_pad + (self.fft_len - self.win_len) % 2\n window = F.pad(window, (left_pad, right_pad))\n if self.win_sqrt:\n self.padded_window = window\n window = th.sqrt(window)\n else:\n self.padded_window = window**2\n\n fft_kernel = fft_kernel.T * window\n ifft_kernel = ifft_kernel * window\n ola_kernel = th.eye(self.fft_len)[: self.win_len, None, :]\n if self.mode == \"continue\":\n ola_kernel = th.eye(self.fft_len)[:, None, : self.fft_len]\n return enframed_kernel, fft_kernel, ifft_kernel, ola_kernel\n\n def is_perfect(self):\n \"\"\"\n Whether the parameters win_len, win_hop and win_sqrt\n obey constants overlap-add(COLA)\n\n Returns:\n bool: Return true if parameters obey COLA.\n \"\"\"\n return self.perfect_reconstruct and self.pad_center\n\n def transform(self, inputs, return_type=\"complex\"):\n \"\"\"Take input data (audio) to STFT domain.\n\n Args:\n inputs (tensor): Tensor of floats, with shape (num_batch, num_samples)\n return_type (str, optional): return (mag, phase) when `magphase`,\n return (real, imag) when `realimag` and complex(real, imag) when `complex`.\n Defaults to 'complex'.\n\n Returns:\n tuple: (mag, phase) when `magphase`, return (real, imag) when\n `realimag`. Defaults to 'complex', each elements with shape\n [num_batch, num_frequencies, num_frames]\n \"\"\"\n assert return_type in [\"magphase\", \"realimag\", \"complex\"]\n if inputs.dim() == 2:\n inputs = th.unsqueeze(inputs, 1)\n self.num_samples = inputs.size(-1)\n if self.pad_center:\n inputs = F.pad(inputs, (self.pad_amount, self.pad_amount), mode=\"reflect\")\n enframe_inputs = F.conv1d(inputs, self.en_k, stride=self.win_hop)\n outputs = th.transpose(enframe_inputs, 1, 2)\n outputs = F.linear(outputs, self.fft_k)\n outputs = th.transpose(outputs, 1, 2)\n dim = self.fft_len // 2 + 1\n real = outputs[:, :dim, :]\n imag = outputs[:, dim:, :]\n if return_type == \"realimag\":\n return real, imag\n elif return_type == \"complex\":\n assert support_clp_op\n return th.complex(real, imag)\n else:\n mags = th.sqrt(real**2 + imag**2)\n phase = th.atan2(imag, real)\n return mags, phase\n\n def inverse(self, input1, input2=None, input_type=\"magphase\"):\n \"\"\"Call the inverse STFT (iSTFT), given tensors produced\n by the `transform` function.\n\n Args:\n input1 (tensors): Magnitude/Real-part of STFT with shape\n [num_batch, num_frequencies, num_frames]\n input2 (tensors): Phase/Imag-part of STFT with shape\n [num_batch, num_frequencies, num_frames]\n input_type (str, optional): Mathematical meaning of input tensor's.\n Defaults to 'magphase'.\n\n Returns:\n tensors: Reconstructed audio given magnitude and phase. Of\n shape [num_batch, num_samples]\n \"\"\"\n assert input_type in [\"magphase\", \"realimag\"]\n if input_type == \"realimag\":\n real, imag = None, None\n if support_clp_op and th.is_complex(input1):\n real, imag = input1.real, input1.imag\n else:\n real, imag = input1, input2\n else:\n real = input1 * th.cos(input2)\n imag = input1 * th.sin(input2)\n inputs = th.cat([real, imag], dim=1)\n outputs = F.conv_transpose1d(inputs, self.ifft_k, stride=self.win_hop)\n t = (self.padded_window[None, :, None]).repeat(1, 1, inputs.size(-1))\n t = t.to(inputs.device)\n coff = F.conv_transpose1d(t, self.ola_k, stride=self.win_hop)\n\n num_frames = input1.size(-1)\n num_samples = num_frames * self.win_hop\n\n rm_start, rm_end = self.pad_amount, self.pad_amount + num_samples\n\n outputs = outputs[..., rm_start:rm_end]\n coff = coff[..., rm_start:rm_end]\n coffidx = th.where(coff > 1e-8)\n outputs[coffidx] = outputs[coffidx] / (coff[coffidx])\n return outputs.squeeze(dim=1)\n\n def forward(self, inputs):\n \"\"\"Take input data (audio) to STFT domain and then back to audio.\n\n Args:\n inputs (tensor): Tensor of floats, with shape [num_batch, num_samples]\n\n Returns:\n tensor: Reconstructed audio given magnitude and phase.\n Of shape [num_batch, num_samples]\n \"\"\"\n mag, phase = self.transform(inputs)\n rec_wav = self.inverse(mag, phase)\n return rec_wav\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/scripts/convert_checkpoint.py", "content": "import argparse\nimport json\nimport os\nimport re\nimport time\nimport traceback\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\nimport safetensors.torch\nimport torch\nfrom tensorrt_llm import str_dtype_to_torch\nfrom tensorrt_llm.mapping import Mapping\nfrom tensorrt_llm.models.convert_utils import split, split_matrix_tp\n\n\ndef split_q_tp(v, n_head, n_hidden, tensor_parallel, rank):\n split_v = split(v, tensor_parallel, rank, dim=1)\n return split_v.contiguous()\n\n\ndef split_q_bias_tp(v, n_head, n_hidden, tensor_parallel, rank):\n split_v = split(v, tensor_parallel, rank, dim=0)\n return split_v.contiguous()\n\n\ndef parse_arguments():\n parser = argparse.ArgumentParser()\n parser.add_argument(\"--pytorch_ckpt\", type=str, default=\"./ckpts/model_last.pt\")\n parser.add_argument(\n \"--output_dir\", type=str, default=\"./tllm_checkpoint\", help=\"The path to save the TensorRT-LLM checkpoint\"\n )\n parser.add_argument(\"--tp_size\", type=int, default=1, help=\"N-way tensor parallelism size\")\n parser.add_argument(\"--cp_size\", type=int, default=1, help=\"Context parallelism size\")\n parser.add_argument(\"--pp_size\", type=int, default=1, help=\"N-way pipeline parallelism size\")\n parser.add_argument(\"--dtype\", type=str, default=\"float16\", choices=[\"float32\", \"bfloat16\", \"float16\"])\n parser.add_argument(\"--fp8_linear\", action=\"store_true\", help=\"Whether use FP8 for linear layers\")\n parser.add_argument(\n \"--workers\", type=int, default=1, help=\"The number of workers for converting checkpoint in parallel\"\n )\n parser.add_argument(\n \"--model_name\",\n type=str,\n default=\"F5TTS_Custom\",\n choices=[\n \"F5TTS_v1_Base\",\n \"F5TTS_Base\",\n \"F5TTS_v1_Small\",\n \"F5TTS_Small\",\n ], # if set, overwrite the below hyperparams\n )\n parser.add_argument(\"--hidden_size\", type=int, default=1024, help=\"The hidden size of DiT\")\n parser.add_argument(\"--depth\", type=int, default=22, help=\"The number of DiTBlock layers\")\n parser.add_argument(\"--num_heads\", type=int, default=16, help=\"The number of heads of attention module\")\n parser.add_argument(\"--dim_head\", type=int, default=64, help=\"The dimension of attention head\")\n parser.add_argument(\"--ff_mult\", type=int, default=2, help=\"The FFN intermediate dimension multiplier\")\n parser.add_argument(\"--text_dim\", type=int, default=512, help=\"The output dimension of text encoder\")\n parser.add_argument(\n \"--text_mask_padding\",\n type=lambda x: x.lower() == \"true\",\n choices=[True, False],\n default=True,\n help=\"Whether apply padding mask for conv layers in text encoder\",\n )\n parser.add_argument(\"--conv_layers\", type=int, default=4, help=\"The number of conv layers of text encoder\")\n parser.add_argument(\"--pe_attn_head\", type=int, default=None, help=\"The number of attn head that apply pos emb\")\n args = parser.parse_args()\n\n # overwrite if --model_name ordered\n if args.model_name == \"F5TTS_v1_Base\":\n args.hidden_size = 1024\n args.depth = 22\n args.num_heads = 16\n args.dim_head = 64\n args.ff_mult = 2\n args.text_dim = 512\n args.text_mask_padding = True\n args.conv_layers = 4\n args.pe_attn_head = None\n elif args.model_name == \"F5TTS_Base\":\n args.hidden_size = 1024\n args.depth = 22\n args.num_heads = 16\n args.dim_head = 64\n args.ff_mult = 2\n args.text_dim = 512\n args.text_mask_padding = False\n args.conv_layers = 4\n args.pe_attn_head = 1\n elif args.model_name == \"F5TTS_v1_Small\":\n args.hidden_size = 768\n args.depth = 18\n args.num_heads = 12\n args.dim_head = 64\n args.ff_mult = 2\n args.text_dim = 512\n args.text_mask_padding = True\n args.conv_layers = 4\n args.pe_attn_head = None\n elif args.model_name == \"F5TTS_Small\":\n args.hidden_size = 768\n args.depth = 18\n args.num_heads = 12\n args.dim_head = 64\n args.ff_mult = 2\n args.text_dim = 512\n args.text_mask_padding = False\n args.conv_layers = 4\n args.pe_attn_head = 1\n\n return args\n\n\ndef convert_pytorch_dit_to_trtllm_weight(args, mapping, dtype=\"float32\", use_ema=True):\n weights = {}\n tik = time.time()\n torch_dtype = str_dtype_to_torch(dtype)\n tensor_parallel = mapping.tp_size\n\n ckpt_path = args.pytorch_ckpt\n ckpt_type = ckpt_path.split(\".\")[-1]\n if ckpt_type == \"safetensors\":\n from safetensors.torch import load_file\n\n model_params = load_file(ckpt_path)\n else:\n ckpt = torch.load(ckpt_path, map_location=\"cpu\", weights_only=True)\n model_params = ckpt[\"ema_model_state_dict\"] if use_ema else ckpt[\"model_state_dict\"]\n\n prefix = \"ema_model.transformer.\" if use_ema else \"transformer.\"\n if any(k.startswith(prefix) for k in model_params.keys()):\n model_params = {\n key[len(prefix) :] if key.startswith(prefix) else key: value\n for key, value in model_params.items()\n if key.startswith(prefix)\n }\n\n pytorch_to_trtllm_name = {\n r\"^time_embed\\.time_mlp\\.0\\.(weight|bias)$\": r\"time_embed.mlp1.\\1\",\n r\"^time_embed\\.time_mlp\\.2\\.(weight|bias)$\": r\"time_embed.mlp2.\\1\",\n r\"^input_embed\\.conv_pos_embed\\.conv1d\\.0\\.(weight|bias)$\": r\"input_embed.conv_pos_embed.conv1d1.\\1\",\n r\"^input_embed\\.conv_pos_embed\\.conv1d\\.2\\.(weight|bias)$\": r\"input_embed.conv_pos_embed.conv1d2.\\1\",\n r\"^transformer_blocks\\.(\\d+)\\.attn\\.to_out\\.0\\.(weight|bias)$\": r\"transformer_blocks.\\1.attn.to_out.\\2\",\n r\"^transformer_blocks\\.(\\d+)\\.ff\\.ff\\.0\\.0\\.(weight|bias)$\": r\"transformer_blocks.\\1.ff.project_in.\\2\",\n r\"^transformer_blocks\\.(\\d+)\\.ff\\.ff\\.2\\.(weight|bias)$\": r\"transformer_blocks.\\1.ff.ff.\\2\",\n }\n\n def get_trtllm_name(pytorch_name):\n for pytorch_name_pattern, trtllm_name_replacement in pytorch_to_trtllm_name.items():\n trtllm_name_if_matched = re.sub(pytorch_name_pattern, trtllm_name_replacement, pytorch_name)\n if trtllm_name_if_matched != pytorch_name:\n return trtllm_name_if_matched\n return pytorch_name\n\n weights = dict()\n for name, param in model_params.items():\n if name == \"input_embed.conv_pos_embed.conv1d.0.weight\" or name == \"input_embed.conv_pos_embed.conv1d.2.weight\":\n weights[get_trtllm_name(name)] = param.contiguous().to(torch_dtype).unsqueeze(-1)\n else:\n weights[get_trtllm_name(name)] = param.contiguous().to(torch_dtype)\n\n assert len(weights) == len(model_params)\n\n # new_prefix = \"f5_transformer.\"\n new_prefix = \"\"\n weights = {new_prefix + key: value for key, value in weights.items()}\n import math\n\n scale_factor = math.pow(64, -0.25)\n for k, v in weights.items():\n if re.match(\"^transformer_blocks.*.attn.to_k.weight$\", k):\n weights[k] *= scale_factor\n weights[k] = split_q_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n\n elif re.match(\"^transformer_blocks.*.attn.to_k.bias$\", k):\n weights[k] *= scale_factor\n weights[k] = split_q_bias_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n\n elif re.match(\"^transformer_blocks.*.attn.to_q.weight$\", k):\n weights[k] = split_q_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n weights[k] *= scale_factor\n\n elif re.match(\"^transformer_blocks.*.attn.to_q.bias$\", k):\n weights[k] = split_q_bias_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n weights[k] *= scale_factor\n\n elif re.match(\"^transformer_blocks.*.attn.to_v.weight$\", k):\n weights[k] = split_q_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n\n elif re.match(\"^transformer_blocks.*.attn.to_v.bias$\", k):\n weights[k] = split_q_bias_tp(v, args.num_heads, args.hidden_size, tensor_parallel, mapping.tp_rank)\n\n elif re.match(\"^transformer_blocks.*.attn.to_out.weight$\", k):\n weights[k] = split_matrix_tp(v, tensor_parallel, mapping.tp_rank, dim=1)\n\n tok = time.time()\n t = time.strftime(\"%H:%M:%S\", time.gmtime(tok - tik))\n print(f\"Weights loaded. Total time: {t}\")\n return weights\n\n\ndef save_config(args):\n if not os.path.exists(args.output_dir):\n os.makedirs(args.output_dir)\n config = {\n \"architecture\": \"F5TTS\", # set the same as in ../patch/__init__.py\n \"dtype\": args.dtype,\n \"hidden_size\": args.hidden_size,\n \"num_hidden_layers\": args.depth,\n \"num_attention_heads\": args.num_heads,\n \"dim_head\": args.dim_head,\n \"dropout\": 0.0, # inference-only\n \"ff_mult\": args.ff_mult,\n \"mel_dim\": 100,\n \"text_dim\": args.text_dim,\n \"text_mask_padding\": args.text_mask_padding,\n \"conv_layers\": args.conv_layers,\n \"pe_attn_head\": args.pe_attn_head,\n \"mapping\": {\n \"world_size\": args.cp_size * args.tp_size * args.pp_size,\n \"cp_size\": args.cp_size,\n \"tp_size\": args.tp_size,\n \"pp_size\": args.pp_size,\n },\n }\n if args.fp8_linear:\n config[\"quantization\"] = {\n \"quant_algo\": \"FP8\",\n # TODO: add support for exclude modules.\n # \"exclude_modules\": \"*final_layer*\",\n }\n\n with open(os.path.join(args.output_dir, \"config.json\"), \"w\") as f:\n json.dump(config, f, indent=4)\n\n\ndef covert_and_save(args, rank):\n if rank == 0:\n save_config(args)\n\n mapping = Mapping(\n world_size=args.cp_size * args.tp_size * args.pp_size,\n rank=rank,\n cp_size=args.cp_size,\n tp_size=args.tp_size,\n pp_size=args.pp_size,\n )\n\n weights = convert_pytorch_dit_to_trtllm_weight(args, mapping, dtype=args.dtype)\n\n safetensors.torch.save_file(weights, os.path.join(args.output_dir, f\"rank{rank}.safetensors\"))\n\n\ndef execute(workers, func, args):\n if workers == 1:\n for rank, f in enumerate(func):\n f(args, rank)\n else:\n with ThreadPoolExecutor(max_workers=workers) as p:\n futures = [p.submit(f, args, rank) for rank, f in enumerate(func)]\n exceptions = []\n for future in as_completed(futures):\n try:\n future.result()\n except Exception as e:\n traceback.print_exc()\n exceptions.append(e)\n assert len(exceptions) == 0, \"Checkpoint conversion failed, please check error log.\"\n\n\ndef main():\n args = parse_arguments()\n world_size = args.cp_size * args.tp_size * args.pp_size\n\n assert args.pp_size == 1, \"PP is not supported yet.\"\n\n tik = time.time()\n if args.pytorch_ckpt is None:\n return\n print(\"Start execute\")\n execute(args.workers, [covert_and_save] * world_size, args)\n\n tok = time.time()\n t = time.strftime(\"%H:%M:%S\", time.gmtime(tok - tik))\n print(f\"Total time of converting checkpoints: {t}\")\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/scripts/export_vocoder_to_onnx.py", "content": "# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\n\nimport torch\nimport torch.nn as nn\nfrom conv_stft import STFT\nfrom huggingface_hub import hf_hub_download\nfrom vocos import Vocos\n\n\nopset_version = 17\n\n\ndef get_args():\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n parser.add_argument(\n \"--vocoder\",\n type=str,\n default=\"vocos\",\n choices=[\"vocos\", \"bigvgan\"],\n help=\"Vocoder to export\",\n )\n parser.add_argument(\n \"--output-path\",\n type=str,\n default=\"./vocos_vocoder.onnx\",\n help=\"Output path\",\n )\n return parser.parse_args()\n\n\nclass ISTFTHead(nn.Module):\n def __init__(self, n_fft: int, hop_length: int):\n super().__init__()\n self.out = None\n self.stft = STFT(fft_len=n_fft, win_hop=hop_length, win_len=n_fft)\n\n def forward(self, x: torch.Tensor):\n x = self.out(x).transpose(1, 2)\n mag, p = x.chunk(2, dim=1)\n mag = torch.exp(mag)\n mag = torch.clip(mag, max=1e2)\n real = mag * torch.cos(p)\n imag = mag * torch.sin(p)\n audio = self.stft.inverse(input1=real, input2=imag, input_type=\"realimag\")\n return audio\n\n\nclass VocosVocoder(nn.Module):\n def __init__(self, vocos_vocoder):\n super(VocosVocoder, self).__init__()\n self.vocos_vocoder = vocos_vocoder\n istft_head_out = self.vocos_vocoder.head.out\n n_fft = self.vocos_vocoder.head.istft.n_fft\n hop_length = self.vocos_vocoder.head.istft.hop_length\n istft_head_for_export = ISTFTHead(n_fft, hop_length)\n istft_head_for_export.out = istft_head_out\n self.vocos_vocoder.head = istft_head_for_export\n\n def forward(self, mel):\n waveform = self.vocos_vocoder.decode(mel)\n return waveform\n\n\ndef export_VocosVocoder(vocos_vocoder, output_path, verbose):\n vocos_vocoder = VocosVocoder(vocos_vocoder).cuda()\n vocos_vocoder.eval()\n\n dummy_batch_size = 8\n dummy_input_length = 500\n\n dummy_mel = torch.randn(dummy_batch_size, 100, dummy_input_length).cuda()\n\n with torch.no_grad():\n dummy_waveform = vocos_vocoder(mel=dummy_mel)\n print(dummy_waveform.shape)\n\n dummy_input = dummy_mel\n\n torch.onnx.export(\n vocos_vocoder,\n dummy_input,\n output_path,\n opset_version=opset_version,\n do_constant_folding=True,\n input_names=[\"mel\"],\n output_names=[\"waveform\"],\n dynamic_axes={\n \"mel\": {0: \"batch_size\", 2: \"input_length\"},\n \"waveform\": {0: \"batch_size\", 1: \"output_length\"},\n },\n verbose=verbose,\n )\n\n print(\"Exported to {}\".format(output_path))\n\n\ndef load_vocoder(vocoder_name=\"vocos\", is_local=False, local_path=\"\", device=\"cpu\", hf_cache_dir=None):\n if vocoder_name == \"vocos\":\n # vocoder = Vocos.from_pretrained(\"charactr/vocos-mel-24khz\").to(device)\n if is_local:\n print(f\"Load vocos from local path {local_path}\")\n config_path = f\"{local_path}/config.yaml\"\n model_path = f\"{local_path}/pytorch_model.bin\"\n else:\n print(\"Download Vocos from huggingface charactr/vocos-mel-24khz\")\n repo_id = \"charactr/vocos-mel-24khz\"\n config_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"config.yaml\")\n model_path = hf_hub_download(repo_id=repo_id, cache_dir=hf_cache_dir, filename=\"pytorch_model.bin\")\n vocoder = Vocos.from_hparams(config_path)\n state_dict = torch.load(model_path, map_location=\"cpu\", weights_only=True)\n vocoder.load_state_dict(state_dict)\n vocoder = vocoder.eval().to(device)\n elif vocoder_name == \"bigvgan\":\n raise NotImplementedError(\"BigVGAN is not supported yet\")\n vocoder.remove_weight_norm()\n vocoder = vocoder.eval().to(device)\n return vocoder\n\n\nif __name__ == \"__main__\":\n args = get_args()\n vocoder = load_vocoder(vocoder_name=args.vocoder, device=\"cpu\", hf_cache_dir=None)\n if args.vocoder == \"vocos\":\n export_VocosVocoder(vocoder, args.output_path, verbose=False)\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/scripts/export_vocos_trt.sh", "content": "#!/bin/bash\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Manual installation of TensorRT, in case not using NVIDIA NGC:\n# https://docs.nvidia.com/deeplearning/tensorrt/latest/installing-tensorrt/installing.html#downloading-tensorrt\nTRTEXEC=\"/usr/src/tensorrt/bin/trtexec\"\n\nONNX_PATH=$1\nENGINE_PATH=$2\necho \"ONNX_PATH: $ONNX_PATH\"\necho \"ENGINE_PATH: $ENGINE_PATH\"\nPRECISION=\"fp32\"\n\n\nMIN_BATCH_SIZE=1\nOPT_BATCH_SIZE=1\nMAX_BATCH_SIZE=8\n\nMIN_INPUT_LENGTH=1\nOPT_INPUT_LENGTH=1000\nMAX_INPUT_LENGTH=3000 # 4096\n\nMEL_MIN_SHAPE=\"${MIN_BATCH_SIZE}x100x${MIN_INPUT_LENGTH}\"\nMEL_OPT_SHAPE=\"${OPT_BATCH_SIZE}x100x${OPT_INPUT_LENGTH}\"\nMEL_MAX_SHAPE=\"${MAX_BATCH_SIZE}x100x${MAX_INPUT_LENGTH}\"\n\n${TRTEXEC} \\\n --minShapes=\"mel:${MEL_MIN_SHAPE}\" \\\n --optShapes=\"mel:${MEL_OPT_SHAPE}\" \\\n --maxShapes=\"mel:${MEL_MAX_SHAPE}\" \\\n --onnx=${ONNX_PATH} \\\n --saveEngine=${ENGINE_PATH}\n" }, { "path": "src/f5_tts/runtime/triton_trtllm/scripts/fill_template.py", "content": "#! /usr/bin/env python3\nfrom argparse import ArgumentParser\nfrom string import Template\n\n\ndef main(file_path, substitutions, in_place, participant_ids):\n with open(file_path) as f:\n pbtxt = Template(f.read())\n\n sub_dict = {\"max_queue_size\": 0}\n sub_dict[\"participant_ids\"] = participant_ids\n for sub in substitutions.split(\",\"):\n key, value = sub.split(\":\")\n sub_dict[key] = value\n\n pbtxt = pbtxt.safe_substitute(sub_dict)\n\n if in_place:\n with open(file_path, \"w\") as f:\n f.write(pbtxt)\n else:\n print(pbtxt)\n\n\nif __name__ == \"__main__\":\n parser = ArgumentParser()\n parser.add_argument(\"file_path\", help=\"path of the .pbtxt to modify\")\n parser.add_argument(\n \"substitutions\",\n help=\"substitutions to perform, in the format variable_name_1:value_1,variable_name_2:value_2...\",\n )\n parser.add_argument(\"--in_place\", \"-i\", action=\"store_true\", help=\"do the operation in-place\")\n parser.add_argument(\"--participant_ids\", help=\"Participant IDs for the model\", default=\"\")\n args = parser.parse_args()\n\n main(**vars(args))\n" }, { "path": "src/f5_tts/scripts/count_max_epoch.py", "content": "\"\"\"ADAPTIVE BATCH SIZE\"\"\"\n\nprint(\"Adaptive batch size: using grouping batch sampler, frames_per_gpu fixed fed in\")\nprint(\" -> least padding, gather wavs with accumulated frames in a batch\\n\")\n\n# data\ntotal_hours = 95282\nmel_hop_length = 256\nmel_sampling_rate = 24000\n\n# target\nwanted_max_updates = 1200000\n\n# train params\ngpus = 8\nframes_per_gpu = 38400 # 8 * 38400 = 307200\ngrad_accum = 1\n\n# intermediate\nmini_batch_frames = frames_per_gpu * grad_accum * gpus\nmini_batch_hours = mini_batch_frames * mel_hop_length / mel_sampling_rate / 3600\nupdates_per_epoch = total_hours / mini_batch_hours\n# steps_per_epoch = updates_per_epoch * grad_accum\n\n# result\nepochs = wanted_max_updates / updates_per_epoch\nprint(f\"epochs should be set to: {epochs:.0f} ({epochs / grad_accum:.1f} x gd_acum {grad_accum})\")\nprint(f\"progress_bar should show approx. 0/{updates_per_epoch:.0f} updates\")\n# print(f\" or approx. 0/{steps_per_epoch:.0f} steps\")\n\n# others\nprint(f\"total {total_hours:.0f} hours\")\nprint(f\"mini-batch of {mini_batch_frames:.0f} frames, {mini_batch_hours:.2f} hours per mini-batch\")\n" }, { "path": "src/f5_tts/scripts/count_max_epoch_precise.py", "content": "import math\n\nfrom torch.utils.data import SequentialSampler\n\nfrom f5_tts.model.dataset import DynamicBatchSampler, load_dataset\n\n\ntrain_dataset = load_dataset(\"Emilia_ZH_EN\", \"pinyin\")\nsampler = SequentialSampler(train_dataset)\n\ngpus = 8\nbatch_size_per_gpu = 38400\nmax_samples_per_gpu = 64\nmax_updates = 1250000\n\nbatch_sampler = DynamicBatchSampler(\n sampler,\n batch_size_per_gpu,\n max_samples=max_samples_per_gpu,\n random_seed=666,\n drop_residual=False,\n)\nupdates_per_epoch = int(len(batch_sampler) / gpus)\n\nprint(\n f\"One epoch has {updates_per_epoch} updates if gpus={gpus}, with \"\n f\"batch_size_per_gpu={batch_size_per_gpu} (frames) & \"\n f\"max_samples_per_gpu={max_samples_per_gpu}.\"\n)\nprint(f\"If gpus={gpus}, for max_updates={max_updates} should set epoch={math.ceil(max_updates / updates_per_epoch)}.\")\n" }, { "path": "src/f5_tts/scripts/count_params_gflops.py", "content": "import os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport thop\nimport torch\n\nfrom f5_tts.model import CFM, DiT\n\n\n\"\"\" ~155M \"\"\"\n# transformer = UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4)\n# transformer = UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4, text_dim = 512, conv_layers = 4)\n# transformer = DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2)\n# transformer = DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4)\n# transformer = DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4, long_skip_connection = True)\n# transformer = MMDiT(dim = 512, depth = 16, heads = 16, ff_mult = 2)\n\n\"\"\" ~335M \"\"\"\n# FLOPs: 622.1 G, Params: 333.2 M\n# transformer = UNetT(dim = 1024, depth = 24, heads = 16, ff_mult = 4)\n# FLOPs: 363.4 G, Params: 335.8 M\ntransformer = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)\n\n\nmodel = CFM(transformer=transformer)\ntarget_sample_rate = 24000\nn_mel_channels = 100\nhop_length = 256\nduration = 20\nframe_length = int(duration * target_sample_rate / hop_length)\ntext_length = 150\n\nflops, params = thop.profile(\n model, inputs=(torch.randn(1, frame_length, n_mel_channels), torch.zeros(1, text_length, dtype=torch.long))\n)\nprint(f\"FLOPs: {flops / 1e9} G\")\nprint(f\"Params: {params / 1e6} M\")\n" }, { "path": "src/f5_tts/socket_client.py", "content": "import asyncio\nimport logging\nimport socket\nimport time\n\nimport numpy as np\nimport pyaudio\n\n\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\n\nasync def listen_to_F5TTS(text, server_ip=\"localhost\", server_port=9998):\n client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n await asyncio.get_event_loop().run_in_executor(None, client_socket.connect, (server_ip, int(server_port)))\n\n start_time = time.time()\n first_chunk_time = None\n\n async def play_audio_stream():\n nonlocal first_chunk_time\n p = pyaudio.PyAudio()\n stream = p.open(format=pyaudio.paFloat32, channels=1, rate=24000, output=True, frames_per_buffer=2048)\n\n try:\n while True:\n data = await asyncio.get_event_loop().run_in_executor(None, client_socket.recv, 8192)\n if not data:\n break\n if data == b\"END\":\n logger.info(\"End of audio received.\")\n break\n\n audio_array = np.frombuffer(data, dtype=np.float32)\n stream.write(audio_array.tobytes())\n\n if first_chunk_time is None:\n first_chunk_time = time.time()\n\n finally:\n stream.stop_stream()\n stream.close()\n p.terminate()\n\n logger.info(f\"Total time taken: {time.time() - start_time:.4f} seconds\")\n\n try:\n data_to_send = f\"{text}\".encode(\"utf-8\")\n await asyncio.get_event_loop().run_in_executor(None, client_socket.sendall, data_to_send)\n await play_audio_stream()\n\n except Exception as e:\n logger.error(f\"Error in listen_to_F5TTS: {e}\")\n\n finally:\n client_socket.close()\n\n\nif __name__ == \"__main__\":\n 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\"\n\n asyncio.run(listen_to_F5TTS(text_to_send))\n" }, { "path": "src/f5_tts/socket_server.py", "content": "import argparse\nimport gc\nimport logging\nimport queue\nimport socket\nimport struct\nimport threading\nimport traceback\nimport wave\nfrom importlib.resources import files\n\nimport numpy as np\nimport torch\nimport torchaudio\nfrom huggingface_hub import hf_hub_download\nfrom hydra.utils import get_class\nfrom omegaconf import OmegaConf\n\nfrom f5_tts.infer.utils_infer import (\n chunk_text,\n infer_batch_process,\n load_model,\n load_vocoder,\n preprocess_ref_audio_text,\n)\n\n\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\n\nclass AudioFileWriterThread(threading.Thread):\n \"\"\"Threaded file writer to avoid blocking the TTS streaming process.\"\"\"\n\n def __init__(self, output_file, sampling_rate):\n super().__init__()\n self.output_file = output_file\n self.sampling_rate = sampling_rate\n self.queue = queue.Queue()\n self.stop_event = threading.Event()\n self.audio_data = []\n\n def run(self):\n \"\"\"Process queued audio data and write it to a file.\"\"\"\n logger.info(\"AudioFileWriterThread started.\")\n with wave.open(self.output_file, \"wb\") as wf:\n wf.setnchannels(1)\n wf.setsampwidth(2)\n wf.setframerate(self.sampling_rate)\n\n while not self.stop_event.is_set() or not self.queue.empty():\n try:\n chunk = self.queue.get(timeout=0.1)\n if chunk is not None:\n chunk = np.int16(chunk * 32767)\n self.audio_data.append(chunk)\n wf.writeframes(chunk.tobytes())\n except queue.Empty:\n continue\n\n def add_chunk(self, chunk):\n \"\"\"Add a new chunk to the queue.\"\"\"\n self.queue.put(chunk)\n\n def stop(self):\n \"\"\"Stop writing and ensure all queued data is written.\"\"\"\n self.stop_event.set()\n self.join()\n logger.info(\"Audio writing completed.\")\n\n\nclass TTSStreamingProcessor:\n def __init__(self, model, ckpt_file, vocab_file, ref_audio, ref_text, device=None, dtype=torch.float32):\n self.device = device or (\n \"cuda\"\n if torch.cuda.is_available()\n else \"xpu\"\n if torch.xpu.is_available()\n else \"mps\"\n if torch.backends.mps.is_available()\n else \"cpu\"\n )\n model_cfg = OmegaConf.load(str(files(\"f5_tts\").joinpath(f\"configs/{model}.yaml\")))\n self.model_cls = get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\n self.model_arc = model_cfg.model.arch\n self.mel_spec_type = model_cfg.model.mel_spec.mel_spec_type\n self.sampling_rate = model_cfg.model.mel_spec.target_sample_rate\n\n self.model = self.load_ema_model(ckpt_file, vocab_file, dtype)\n self.vocoder = self.load_vocoder_model()\n\n self.update_reference(ref_audio, ref_text)\n self._warm_up()\n self.file_writer_thread = None\n self.first_package = True\n\n def load_ema_model(self, ckpt_file, vocab_file, dtype):\n return load_model(\n self.model_cls,\n self.model_arc,\n ckpt_path=ckpt_file,\n mel_spec_type=self.mel_spec_type,\n vocab_file=vocab_file,\n ode_method=\"euler\",\n use_ema=True,\n device=self.device,\n ).to(self.device, dtype=dtype)\n\n def load_vocoder_model(self):\n return load_vocoder(vocoder_name=self.mel_spec_type, is_local=False, local_path=None, device=self.device)\n\n def update_reference(self, ref_audio, ref_text):\n self.ref_audio, self.ref_text = preprocess_ref_audio_text(ref_audio, ref_text)\n self.audio, self.sr = torchaudio.load(self.ref_audio)\n\n ref_audio_duration = self.audio.shape[-1] / self.sr\n ref_text_byte_len = len(self.ref_text.encode(\"utf-8\"))\n self.max_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration))\n self.few_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 2)\n self.min_chars = int(ref_text_byte_len / (ref_audio_duration) * (25 - ref_audio_duration) / 4)\n\n def _warm_up(self):\n logger.info(\"Warming up the model...\")\n gen_text = \"Warm-up text for the model.\"\n for _ in infer_batch_process(\n (self.audio, self.sr),\n self.ref_text,\n [gen_text],\n self.model,\n self.vocoder,\n progress=None,\n device=self.device,\n streaming=True,\n ):\n pass\n logger.info(\"Warm-up completed.\")\n\n def generate_stream(self, text, conn):\n text_batches = chunk_text(text, max_chars=self.max_chars)\n if self.first_package:\n text_batches = chunk_text(text_batches[0], max_chars=self.few_chars) + text_batches[1:]\n text_batches = chunk_text(text_batches[0], max_chars=self.min_chars) + text_batches[1:]\n self.first_package = False\n\n audio_stream = infer_batch_process(\n (self.audio, self.sr),\n self.ref_text,\n text_batches,\n self.model,\n self.vocoder,\n progress=None,\n device=self.device,\n streaming=True,\n chunk_size=2048,\n )\n\n # Reset the file writer thread\n if self.file_writer_thread is not None:\n self.file_writer_thread.stop()\n self.file_writer_thread = AudioFileWriterThread(\"output.wav\", self.sampling_rate)\n self.file_writer_thread.start()\n\n for audio_chunk, _ in audio_stream:\n if len(audio_chunk) > 0:\n logger.info(f\"Generated audio chunk of size: {len(audio_chunk)}\")\n\n # Send audio chunk via socket\n conn.sendall(struct.pack(f\"{len(audio_chunk)}f\", *audio_chunk))\n\n # Write to file asynchronously\n self.file_writer_thread.add_chunk(audio_chunk)\n\n logger.info(\"Finished sending audio stream.\")\n conn.sendall(b\"END\") # Send end signal\n\n # Ensure all audio data is written before exiting\n self.file_writer_thread.stop()\n\n\ndef handle_client(conn, processor):\n try:\n with conn:\n conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)\n while True:\n data = conn.recv(1024)\n if not data:\n processor.first_package = True\n break\n data_str = data.decode(\"utf-8\").strip()\n logger.info(f\"Received text: {data_str}\")\n\n try:\n processor.generate_stream(data_str, conn)\n except Exception as inner_e:\n logger.error(f\"Error during processing: {inner_e}\")\n traceback.print_exc()\n break\n except Exception as e:\n logger.error(f\"Error handling client: {e}\")\n traceback.print_exc()\n\n\ndef start_server(host, port, processor):\n with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:\n s.bind((host, port))\n s.listen()\n logger.info(f\"Server started on {host}:{port}\")\n while True:\n conn, addr = s.accept()\n logger.info(f\"Connected by {addr}\")\n handle_client(conn, processor)\n\n\nif __name__ == \"__main__\":\n parser = argparse.ArgumentParser()\n\n parser.add_argument(\"--host\", default=\"0.0.0.0\")\n parser.add_argument(\"--port\", default=9998)\n\n parser.add_argument(\n \"--model\",\n default=\"F5TTS_v1_Base\",\n help=\"The model name, e.g. F5TTS_v1_Base\",\n )\n parser.add_argument(\n \"--ckpt_file\",\n default=str(hf_hub_download(repo_id=\"SWivid/F5-TTS\", filename=\"F5TTS_v1_Base/model_1250000.safetensors\")),\n help=\"Path to the model checkpoint file\",\n )\n parser.add_argument(\n \"--vocab_file\",\n default=\"\",\n help=\"Path to the vocab file if customized\",\n )\n\n parser.add_argument(\n \"--ref_audio\",\n default=str(files(\"f5_tts\").joinpath(\"infer/examples/basic/basic_ref_en.wav\")),\n help=\"Reference audio to provide model with speaker characteristics\",\n )\n parser.add_argument(\n \"--ref_text\",\n default=\"\",\n help=\"Reference audio subtitle, leave empty to auto-transcribe\",\n )\n\n parser.add_argument(\"--device\", default=None, help=\"Device to run the model on\")\n parser.add_argument(\"--dtype\", default=torch.float32, help=\"Data type to use for model inference\")\n\n args = parser.parse_args()\n\n try:\n # Initialize the processor with the model and vocoder\n processor = TTSStreamingProcessor(\n model=args.model,\n ckpt_file=args.ckpt_file,\n vocab_file=args.vocab_file,\n ref_audio=args.ref_audio,\n ref_text=args.ref_text,\n device=args.device,\n dtype=args.dtype,\n )\n\n # Start the server\n start_server(args.host, args.port, processor)\n\n except KeyboardInterrupt:\n gc.collect()\n" }, { "path": "src/f5_tts/train/README.md", "content": "# Training\n\nCheck your FFmpeg installation:\n```bash\nffmpeg -version\n```\nIf not found, install it first (or skip assuming you know of other backends available).\n\n## Prepare Dataset\n\nExample data processing scripts, and you may tailor your own one along with a Dataset class in `src/f5_tts/model/dataset.py`.\n\n### 1. Some specific Datasets preparing scripts\nDownload corresponding dataset first, and fill in the path in scripts.\n\n```bash\n# Prepare the Emilia dataset\npython src/f5_tts/train/datasets/prepare_emilia.py\n\n# Prepare the Wenetspeech4TTS dataset\npython src/f5_tts/train/datasets/prepare_wenetspeech4tts.py\n\n# Prepare the LibriTTS dataset\npython src/f5_tts/train/datasets/prepare_libritts.py\n\n# Prepare the LJSpeech dataset\npython src/f5_tts/train/datasets/prepare_ljspeech.py\n```\n\n### 2. Create custom dataset with CSV\nPrepare a CSV with two columns using a required header: `audio_file|text`. Audio paths must be absolute.\nUse guidance see [#57 here](https://github.com/SWivid/F5-TTS/discussions/57#discussioncomment-10959029).\n\n```bash\npython src/f5_tts/train/datasets/prepare_csv_wavs.py /path/to/metadata.csv /path/to/output\n```\n\n## Training & Finetuning\n\nOnce your datasets are prepared, you can start the training process.\n\n### 1. Training script used for pretrained model\n\n```bash\n# setup accelerate config, e.g. use multi-gpu ddp, fp16\n# will be to: ~/.cache/huggingface/accelerate/default_config.yaml \naccelerate config\n\n# .yaml files are under src/f5_tts/configs directory\naccelerate launch src/f5_tts/train/train.py --config-name F5TTS_v1_Base.yaml\n\n# possible to overwrite accelerate and hydra config\naccelerate launch --mixed_precision=fp16 src/f5_tts/train/train.py --config-name F5TTS_v1_Base.yaml ++datasets.batch_size_per_gpu=19200\n```\n\n### 2. Finetuning practice\nDiscussion board for Finetuning [#57](https://github.com/SWivid/F5-TTS/discussions/57).\n\nGradio UI training/finetuning with `src/f5_tts/train/finetune_gradio.py` see [#143](https://github.com/SWivid/F5-TTS/discussions/143).\n\nIf want to finetune with a variant version e.g. *F5TTS_v1_Base_no_zero_init*, manually download pretrained checkpoint from model weight repository and fill in the path correspondingly on web interface.\n\nIf use tensorboard as logger, install it first with `pip install tensorboard`.\n\nThe `use_ema = True` might be harmful for early-stage finetuned checkpoints (which goes just few updates, thus ema weights still dominated by pretrained ones), try turn it off with finetune gradio option or `load_model(..., use_ema=False)`, see if offer better results.\n\n### 3. W&B Logging\n\nThe `wandb/` dir will be created under path you run training/finetuning scripts.\n\nBy default, the training script does NOT use logging (assuming you didn't manually log in using `wandb login`).\n\nTo turn on wandb logging, you can either:\n\n1. Manually login with `wandb login`: Learn more [here](https://docs.wandb.ai/ref/cli/wandb-login)\n2. Automatically login programmatically by setting an environment variable: Get an API KEY at https://wandb.ai/authorize and set the environment variable as follows:\n\nOn Mac & Linux:\n\n```\nexport WANDB_API_KEY=\n```\n\nOn Windows:\n\n```\nset WANDB_API_KEY=\n```\nMoreover, if you couldn't access W&B and want to log metrics offline, you can set the environment variable as follows:\n\n```\nexport WANDB_MODE=offline\n```\n" }, { "path": "src/f5_tts/train/datasets/prepare_csv_wavs.py", "content": "\"\"\"\nUsage:\n python prepare_csv_wavs.py /path/to/metadata.csv /output/dataset/path [--pretrain] [--workers N]\n\nCSV format (header required, \"|\" delimiter):\n audio_file|text\n /path/to/wavs/audio_0001.wav|Yo! Hello? Hello?\n /path/to/wavs/audio_0002.wav|Hi, how are you doing today? I want to go shopping and buy me some lemons.\n\nNotes:\n - audio_file must be an absolute path.\n\"\"\"\n\nimport concurrent.futures\nimport multiprocessing\nimport os\nimport shutil\nimport signal\nimport subprocess\nimport sys\nfrom contextlib import contextmanager\n\n\nsys.path.append(os.getcwd())\n\nimport argparse\nimport csv\nimport json\nfrom importlib.resources import files\nfrom pathlib import Path\n\nimport soundfile as sf\nimport torchaudio\nfrom datasets.arrow_writer import ArrowWriter\nfrom tqdm import tqdm\n\nfrom f5_tts.model.utils import convert_char_to_pinyin\n\n\nPRETRAINED_VOCAB_PATH = files(\"f5_tts\").joinpath(\"../../data/Emilia_ZH_EN_pinyin/vocab.txt\")\n\n# Configuration constants\nBATCH_SIZE = 100 # Batch size for text conversion\nMAX_WORKERS = max(1, multiprocessing.cpu_count() - 1) # Leave one CPU free\nTHREAD_NAME_PREFIX = \"AudioProcessor\"\nCHUNK_SIZE = 100 # Number of files to process per worker batch\nexecutor = None # Global executor for cleanup\n\n\ndef is_csv_wavs_format(input_path):\n fpath = Path(input_path).expanduser()\n return fpath.is_file() and fpath.suffix.lower() == \".csv\"\n\n\n@contextmanager\ndef graceful_exit():\n \"\"\"Context manager for graceful shutdown on signals\"\"\"\n\n def signal_handler(signum, frame):\n print(\"\\nReceived signal to terminate. Cleaning up...\")\n if executor is not None:\n print(\"Shutting down executor...\")\n executor.shutdown(wait=False, cancel_futures=True)\n sys.exit(1)\n\n # Set up signal handlers\n signal.signal(signal.SIGINT, signal_handler)\n signal.signal(signal.SIGTERM, signal_handler)\n\n try:\n yield\n finally:\n if executor is not None:\n executor.shutdown(wait=False)\n\n\ndef process_audio_file(audio_path, text, polyphone):\n \"\"\"Process a single audio file by checking its existence and extracting duration.\"\"\"\n if not Path(audio_path).exists():\n print(f\"audio {audio_path} not found, skipping\")\n return None\n try:\n audio_duration = get_audio_duration(audio_path)\n if audio_duration <= 0:\n raise ValueError(f\"Duration {audio_duration} is non-positive.\")\n return (audio_path, text, audio_duration)\n except Exception as e:\n print(f\"Warning: Failed to process {audio_path} due to error: {e}. Skipping corrupt file.\")\n return None\n\n\ndef batch_convert_texts(texts, polyphone, batch_size=BATCH_SIZE):\n \"\"\"Convert a list of texts to pinyin in batches.\"\"\"\n converted_texts = []\n for i in tqdm(\n range(0, len(texts), batch_size),\n total=(len(texts) + batch_size - 1) // batch_size,\n desc=\"Converting texts to pinyin\",\n ):\n batch = texts[i : i + batch_size]\n converted_batch = convert_char_to_pinyin(batch, polyphone=polyphone)\n converted_texts.extend(converted_batch)\n return converted_texts\n\n\ndef prepare_csv_wavs_dir(input_path, num_workers=None):\n global executor\n if not is_csv_wavs_format(input_path):\n raise ValueError(f\"input must be a .csv file: {input_path}\")\n audio_path_text_pairs = read_audio_text_pairs(Path(input_path).expanduser().as_posix())\n\n polyphone = True\n total_files = len(audio_path_text_pairs)\n if total_files == 0:\n raise RuntimeError(\"No valid rows found in CSV.\")\n\n # Use provided worker count or calculate optimal number\n worker_count = num_workers if num_workers is not None else min(MAX_WORKERS, total_files)\n print(f\"\\nProcessing {total_files} audio files using {worker_count} workers...\")\n\n with graceful_exit():\n # Initialize thread pool with optimized settings\n with concurrent.futures.ThreadPoolExecutor(\n max_workers=worker_count, thread_name_prefix=THREAD_NAME_PREFIX\n ) as exec:\n executor = exec\n results = []\n\n # Process files in chunks for better efficiency\n for i in range(0, len(audio_path_text_pairs), CHUNK_SIZE):\n chunk = audio_path_text_pairs[i : i + CHUNK_SIZE]\n # Submit futures in order\n chunk_futures = [executor.submit(process_audio_file, pair[0], pair[1], polyphone) for pair in chunk]\n\n # Iterate over futures in the original submission order to preserve ordering\n for future in tqdm(\n chunk_futures,\n total=len(chunk),\n desc=f\"Processing chunk {i // CHUNK_SIZE + 1}/{(total_files + CHUNK_SIZE - 1) // CHUNK_SIZE}\",\n ):\n try:\n result = future.result()\n if result is not None:\n results.append(result)\n except Exception as e:\n print(f\"Error processing file: {e}\")\n\n executor = None\n\n # Filter out failed results\n processed = [res for res in results if res is not None]\n if not processed:\n raise RuntimeError(\"No valid audio files were processed!\")\n\n # Batch process text conversion\n raw_texts = [item[1] for item in processed]\n converted_texts = batch_convert_texts(raw_texts, polyphone, batch_size=BATCH_SIZE)\n\n # Prepare final results\n sub_result = []\n durations = []\n vocab_set = set()\n\n for (audio_path, _, duration), conv_text in zip(processed, converted_texts):\n sub_result.append({\"audio_path\": audio_path, \"text\": conv_text, \"duration\": duration})\n durations.append(duration)\n vocab_set.update(list(conv_text))\n\n return sub_result, durations, vocab_set\n\n\ndef get_audio_duration(audio_path, timeout=5):\n \"\"\"Get the duration of an audio file in seconds with fallbacks.\"\"\"\n try:\n return sf.info(audio_path).duration\n except Exception as e:\n print(f\"Warning: soundfile failed for {audio_path} with error: {e}. Falling back to ffprobe.\")\n\n try:\n cmd = [\n \"ffprobe\",\n \"-v\",\n \"error\",\n \"-show_entries\",\n \"format=duration\",\n \"-of\",\n \"default=noprint_wrappers=1:nokey=1\",\n audio_path,\n ]\n result = subprocess.run(\n cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True, timeout=timeout\n )\n duration_str = result.stdout.strip()\n if duration_str:\n return float(duration_str)\n raise ValueError(\"Empty duration string from ffprobe.\")\n except (subprocess.TimeoutExpired, subprocess.SubprocessError, ValueError) as e:\n print(f\"Warning: ffprobe failed for {audio_path} with error: {e}. Falling back to torchaudio.info.\")\n\n try:\n info = torchaudio.info(audio_path)\n if info.sample_rate > 0:\n return info.num_frames / info.sample_rate\n raise ValueError(\"Invalid sample_rate from torchaudio.info.\")\n except Exception as e:\n raise RuntimeError(f\"failed to get duration for {audio_path}: {e}\")\n\n\ndef read_audio_text_pairs(csv_file_path):\n audio_text_pairs = []\n\n csv_path = Path(csv_file_path).expanduser().absolute()\n with open(csv_path.as_posix(), mode=\"r\", newline=\"\", encoding=\"utf-8-sig\") as csvfile:\n reader = csv.reader(csvfile, delimiter=\"|\")\n header = next(reader, None)\n if header is None:\n return audio_text_pairs\n if len(header) < 2 or header[0].strip() != \"audio_file\" or header[1].strip() != \"text\":\n raise ValueError(\"CSV header must be: audio_file|text\")\n for row_idx, row in enumerate(reader, start=2):\n if len(row) < 2:\n continue\n audio_file = row[0].strip()\n text = row[1].strip()\n if not audio_file:\n continue\n audio_path = Path(audio_file).expanduser()\n if not audio_path.is_absolute():\n raise ValueError(f\"audio_file must be an absolute path (row {row_idx}): {audio_file}\")\n audio_text_pairs.append((audio_path.as_posix(), text))\n\n return audio_text_pairs\n\n\ndef save_prepped_dataset(out_dir, result, duration_list, text_vocab_set, is_finetune):\n out_dir = Path(out_dir)\n out_dir.mkdir(exist_ok=True, parents=True)\n print(f\"\\nSaving to {out_dir} ...\")\n\n raw_arrow_path = out_dir / \"raw.arrow\"\n with ArrowWriter(path=raw_arrow_path.as_posix()) as writer:\n for line in tqdm(result, desc=\"Writing to raw.arrow ...\"):\n writer.write(line)\n writer.finalize()\n\n # Save durations to JSON\n dur_json_path = out_dir / \"duration.json\"\n with open(dur_json_path.as_posix(), \"w\", encoding=\"utf-8\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n # Handle vocab file - write only once based on finetune flag\n voca_out_path = out_dir / \"vocab.txt\"\n if is_finetune:\n file_vocab_finetune = PRETRAINED_VOCAB_PATH.as_posix()\n shutil.copy2(file_vocab_finetune, voca_out_path)\n else:\n with open(voca_out_path.as_posix(), \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n\n dataset_name = out_dir.stem\n print(f\"\\nFor {dataset_name}, sample count: {len(result)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\")\n print(f\"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours\")\n\n\ndef prepare_and_save_set(inp_dir, out_dir, is_finetune: bool = True, num_workers: int = None):\n if is_finetune:\n assert PRETRAINED_VOCAB_PATH.exists(), f\"pretrained vocab.txt not found: {PRETRAINED_VOCAB_PATH}\"\n sub_result, durations, vocab_set = prepare_csv_wavs_dir(inp_dir, num_workers=num_workers)\n save_prepped_dataset(out_dir, sub_result, durations, vocab_set, is_finetune)\n\n\ndef get_args():\n parser = argparse.ArgumentParser(description=\"Prepare and save dataset.\")\n parser.add_argument(\n \"inp_dir\",\n type=str,\n help=\"Input CSV with header 'audio_file|text' and absolute wav paths.\",\n )\n parser.add_argument(\"out_dir\", type=str, help=\"Output directory to save the prepared data.\")\n parser.add_argument(\"--pretrain\", action=\"store_true\", help=\"Enable for new pretrain, otherwise is a fine-tune\")\n parser.add_argument(\"--workers\", type=int, help=f\"Number of worker threads (default: {MAX_WORKERS})\")\n return parser.parse_args()\n\n\ndef cli():\n try:\n args = get_args()\n prepare_and_save_set(args.inp_dir, args.out_dir, is_finetune=not args.pretrain, num_workers=args.workers)\n except KeyboardInterrupt:\n print(\"\\nOperation cancelled by user. Cleaning up...\")\n if executor is not None:\n executor.shutdown(wait=False, cancel_futures=True)\n sys.exit(1)\n\n\nif __name__ == \"__main__\":\n cli()\n" }, { "path": "src/f5_tts/train/datasets/prepare_emilia.py", "content": "# Emilia Dataset: https://huggingface.co/datasets/amphion/Emilia-Dataset/tree/fc71e07\n# if use updated new version, i.e. WebDataset, feel free to modify / draft your own script\n\n# generate audio text map for Emilia ZH & EN\n# evaluate for vocab size\n\nimport os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport json\nfrom concurrent.futures import ProcessPoolExecutor\nfrom importlib.resources import files\nfrom pathlib import Path\n\nfrom datasets.arrow_writer import ArrowWriter\nfrom tqdm import tqdm\n\nfrom f5_tts.model.utils import convert_char_to_pinyin, repetition_found\n\n\nout_zh = {\n \"ZH_B00041_S06226\",\n \"ZH_B00042_S09204\",\n \"ZH_B00065_S09430\",\n \"ZH_B00065_S09431\",\n \"ZH_B00066_S09327\",\n \"ZH_B00066_S09328\",\n}\nzh_filters = [\"い\", \"て\"]\n# seems synthesized audios, or heavily code-switched\nout_en = {\n \"EN_B00013_S00913\",\n \"EN_B00042_S00120\",\n \"EN_B00055_S04111\",\n \"EN_B00061_S00693\",\n \"EN_B00061_S01494\",\n \"EN_B00061_S03375\",\n \"EN_B00059_S00092\",\n \"EN_B00111_S04300\",\n \"EN_B00100_S03759\",\n \"EN_B00087_S03811\",\n \"EN_B00059_S00950\",\n \"EN_B00089_S00946\",\n \"EN_B00078_S05127\",\n \"EN_B00070_S04089\",\n \"EN_B00074_S09659\",\n \"EN_B00061_S06983\",\n \"EN_B00061_S07060\",\n \"EN_B00059_S08397\",\n \"EN_B00082_S06192\",\n \"EN_B00091_S01238\",\n \"EN_B00089_S07349\",\n \"EN_B00070_S04343\",\n \"EN_B00061_S02400\",\n \"EN_B00076_S01262\",\n \"EN_B00068_S06467\",\n \"EN_B00076_S02943\",\n \"EN_B00064_S05954\",\n \"EN_B00061_S05386\",\n \"EN_B00066_S06544\",\n \"EN_B00076_S06944\",\n \"EN_B00072_S08620\",\n \"EN_B00076_S07135\",\n \"EN_B00076_S09127\",\n \"EN_B00065_S00497\",\n \"EN_B00059_S06227\",\n \"EN_B00063_S02859\",\n \"EN_B00075_S01547\",\n \"EN_B00061_S08286\",\n \"EN_B00079_S02901\",\n \"EN_B00092_S03643\",\n \"EN_B00096_S08653\",\n \"EN_B00063_S04297\",\n \"EN_B00063_S04614\",\n \"EN_B00079_S04698\",\n \"EN_B00104_S01666\",\n \"EN_B00061_S09504\",\n \"EN_B00061_S09694\",\n \"EN_B00065_S05444\",\n \"EN_B00063_S06860\",\n \"EN_B00065_S05725\",\n \"EN_B00069_S07628\",\n \"EN_B00083_S03875\",\n \"EN_B00071_S07665\",\n \"EN_B00071_S07665\",\n \"EN_B00062_S04187\",\n \"EN_B00065_S09873\",\n \"EN_B00065_S09922\",\n \"EN_B00084_S02463\",\n \"EN_B00067_S05066\",\n \"EN_B00106_S08060\",\n \"EN_B00073_S06399\",\n \"EN_B00073_S09236\",\n \"EN_B00087_S00432\",\n \"EN_B00085_S05618\",\n \"EN_B00064_S01262\",\n \"EN_B00072_S01739\",\n \"EN_B00059_S03913\",\n \"EN_B00069_S04036\",\n \"EN_B00067_S05623\",\n \"EN_B00060_S05389\",\n \"EN_B00060_S07290\",\n \"EN_B00062_S08995\",\n}\nen_filters = [\"ا\", \"い\", \"て\"]\n\n\ndef deal_with_audio_dir(audio_dir):\n audio_jsonl = audio_dir.with_suffix(\".jsonl\")\n sub_result, durations = [], []\n vocab_set = set()\n bad_case_zh = 0\n bad_case_en = 0\n with open(audio_jsonl, \"r\") as f:\n lines = f.readlines()\n for line in tqdm(lines, desc=f\"{audio_jsonl.stem}\"):\n obj = json.loads(line)\n text = obj[\"text\"]\n if obj[\"language\"] == \"zh\":\n if obj[\"wav\"].split(\"/\")[1] in out_zh or any(f in text for f in zh_filters) or repetition_found(text):\n bad_case_zh += 1\n continue\n else:\n text = text.translate(\n str.maketrans({\",\": \"，\", \"!\": \"！\", \"?\": \"？\"})\n ) # not \"。\" cuz much code-switched\n if obj[\"language\"] == \"en\":\n if (\n obj[\"wav\"].split(\"/\")[1] in out_en\n or any(f in text for f in en_filters)\n or repetition_found(text, length=4)\n ):\n bad_case_en += 1\n continue\n if tokenizer == \"pinyin\":\n text = convert_char_to_pinyin([text], polyphone=polyphone)[0]\n duration = obj[\"duration\"]\n sub_result.append({\"audio_path\": str(audio_dir.parent / obj[\"wav\"]), \"text\": text, \"duration\": duration})\n durations.append(duration)\n vocab_set.update(list(text))\n return sub_result, durations, vocab_set, bad_case_zh, bad_case_en\n\n\ndef main():\n assert tokenizer in [\"pinyin\", \"char\"]\n result = []\n duration_list = []\n text_vocab_set = set()\n total_bad_case_zh = 0\n total_bad_case_en = 0\n\n # process raw data\n executor = ProcessPoolExecutor(max_workers=max_workers)\n futures = []\n for lang in langs:\n dataset_path = Path(os.path.join(dataset_dir, lang))\n [\n futures.append(executor.submit(deal_with_audio_dir, audio_dir))\n for audio_dir in dataset_path.iterdir()\n if audio_dir.is_dir()\n ]\n for futures in tqdm(futures, total=len(futures)):\n sub_result, durations, vocab_set, bad_case_zh, bad_case_en = futures.result()\n result.extend(sub_result)\n duration_list.extend(durations)\n text_vocab_set.update(vocab_set)\n total_bad_case_zh += bad_case_zh\n total_bad_case_en += bad_case_en\n executor.shutdown()\n\n # save preprocessed dataset to disk\n if not os.path.exists(f\"{save_dir}\"):\n os.makedirs(f\"{save_dir}\")\n print(f\"\\nSaving to {save_dir} ...\")\n\n # dataset = Dataset.from_dict({\"audio_path\": audio_path_list, \"text\": text_list, \"duration\": duration_list}) # oom\n # dataset.save_to_disk(f\"{save_dir}/raw\", max_shard_size=\"2GB\")\n with ArrowWriter(path=f\"{save_dir}/raw.arrow\") as writer:\n for line in tqdm(result, desc=\"Writing to raw.arrow ...\"):\n writer.write(line)\n writer.finalize()\n\n # dup a json separately saving duration in case for DynamicBatchSampler ease\n with open(f\"{save_dir}/duration.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n # vocab map, i.e. tokenizer\n # add alphabets and symbols (optional, if plan to ft on de/fr etc.)\n # if tokenizer == \"pinyin\":\n # text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])\n with open(f\"{save_dir}/vocab.txt\", \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n\n print(f\"\\nFor {dataset_name}, sample count: {len(result)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\")\n print(f\"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours\")\n if \"ZH\" in langs:\n print(f\"Bad zh transcription case: {total_bad_case_zh}\")\n if \"EN\" in langs:\n print(f\"Bad en transcription case: {total_bad_case_en}\\n\")\n\n\nif __name__ == \"__main__\":\n max_workers = 32\n\n tokenizer = \"pinyin\" # \"pinyin\" | \"char\"\n polyphone = True\n\n langs = [\"ZH\", \"EN\"]\n dataset_dir = \"/Emilia_Dataset/raw\"\n dataset_name = f\"Emilia_{'_'.join(langs)}_{tokenizer}\"\n save_dir = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/data/{dataset_name}\"\n print(f\"\\nPrepare for {dataset_name}, will save to {save_dir}\\n\")\n\n main()\n\n # Emilia ZH & EN\n # samples count 37837916 (after removal)\n # pinyin vocab size 2543 (polyphone)\n # total duration 95281.87 (hours)\n # bad zh asr cnt 230435 (samples)\n # bad eh asr cnt 37217 (samples)\n\n # vocab size may be slightly different due to rjieba tokenizer and pypinyin (e.g. way of polyphoneme)\n # please be careful if using pretrained model, make sure the vocab.txt is same\n" }, { "path": "src/f5_tts/train/datasets/prepare_emilia_v2.py", "content": "# put in src/f5_tts/train/datasets/prepare_emilia_v2.py\n# prepares Emilia dataset with the new format w/ Emilia-YODAS\n\nimport json\nimport os\nfrom concurrent.futures import ProcessPoolExecutor\nfrom importlib.resources import files\nfrom pathlib import Path\n\nfrom datasets.arrow_writer import ArrowWriter\nfrom tqdm import tqdm\n\nfrom f5_tts.model.utils import repetition_found\n\n\n# Define filters for exclusion\nout_en = set()\nen_filters = [\"ا\", \"い\", \"て\"]\n\n\ndef process_audio_directory(audio_dir):\n sub_result, durations, vocab_set = [], [], set()\n bad_case_en = 0\n\n for file in audio_dir.iterdir():\n if file.suffix == \".json\":\n with open(file, \"r\") as f:\n obj = json.load(f)\n text = obj[\"text\"]\n if any(f in text for f in en_filters) or repetition_found(text, length=4):\n bad_case_en += 1\n continue\n\n duration = obj[\"duration\"]\n audio_file = file.with_suffix(\".mp3\")\n if audio_file.exists():\n sub_result.append({\"audio_path\": str(audio_file), \"text\": text, \"duration\": duration})\n durations.append(duration)\n vocab_set.update(list(text))\n\n return sub_result, durations, vocab_set, bad_case_en\n\n\ndef main():\n assert tokenizer in [\"pinyin\", \"char\"]\n result, duration_list, text_vocab_set = [], [], set()\n total_bad_case_en = 0\n\n executor = ProcessPoolExecutor(max_workers=max_workers)\n futures = []\n dataset_path = Path(dataset_dir)\n for sub_dir in dataset_path.iterdir():\n if sub_dir.is_dir():\n futures.append(executor.submit(process_audio_directory, sub_dir))\n\n for future in tqdm(futures, total=len(futures)):\n sub_result, durations, vocab_set, bad_case_en = future.result()\n result.extend(sub_result)\n duration_list.extend(durations)\n text_vocab_set.update(vocab_set)\n total_bad_case_en += bad_case_en\n\n executor.shutdown()\n\n if not os.path.exists(f\"{save_dir}\"):\n os.makedirs(f\"{save_dir}\")\n\n with ArrowWriter(path=f\"{save_dir}/raw.arrow\") as writer:\n for line in tqdm(result, desc=\"Writing to raw.arrow ...\"):\n writer.write(line)\n writer.finalize()\n\n with open(f\"{save_dir}/duration.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n with open(f\"{save_dir}/vocab.txt\", \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n\n print(f\"For {dataset_name}, sample count: {len(result)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\")\n print(f\"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours\")\n print(f\"Bad en transcription case: {total_bad_case_en}\\n\")\n\n\nif __name__ == \"__main__\":\n max_workers = 32\n tokenizer = \"char\"\n dataset_dir = \"/home/ubuntu/emilia-dataset/Emilia-YODAS/EN\"\n dataset_name = f\"Emilia_EN_{tokenizer}\"\n # save_dir = os.path.expanduser(f\"~/F5-TTS/data/{dataset_name}\")\n save_dir = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/data/{dataset_name}\"\n\n print(f\"Prepare for {dataset_name}, will save to {save_dir}\\n\")\n main()\n" }, { "path": "src/f5_tts/train/datasets/prepare_libritts.py", "content": "import os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport json\nfrom concurrent.futures import ProcessPoolExecutor\nfrom importlib.resources import files\nfrom pathlib import Path\n\nimport soundfile as sf\nfrom datasets.arrow_writer import ArrowWriter\nfrom tqdm import tqdm\n\n\ndef deal_with_audio_dir(audio_dir):\n sub_result, durations = [], []\n vocab_set = set()\n audio_lists = list(audio_dir.rglob(\"*.wav\"))\n\n for line in audio_lists:\n text_path = line.with_suffix(\".normalized.txt\")\n text = open(text_path, \"r\").read().strip()\n duration = sf.info(line).duration\n if duration < 0.4 or duration > 30:\n continue\n sub_result.append({\"audio_path\": str(line), \"text\": text, \"duration\": duration})\n durations.append(duration)\n vocab_set.update(list(text))\n return sub_result, durations, vocab_set\n\n\ndef main():\n result = []\n duration_list = []\n text_vocab_set = set()\n\n # process raw data\n executor = ProcessPoolExecutor(max_workers=max_workers)\n futures = []\n\n for subset in tqdm(SUB_SET):\n dataset_path = Path(os.path.join(dataset_dir, subset))\n [\n futures.append(executor.submit(deal_with_audio_dir, audio_dir))\n for audio_dir in dataset_path.iterdir()\n if audio_dir.is_dir()\n ]\n for future in tqdm(futures, total=len(futures)):\n sub_result, durations, vocab_set = future.result()\n result.extend(sub_result)\n duration_list.extend(durations)\n text_vocab_set.update(vocab_set)\n executor.shutdown()\n\n # save preprocessed dataset to disk\n if not os.path.exists(f\"{save_dir}\"):\n os.makedirs(f\"{save_dir}\")\n print(f\"\\nSaving to {save_dir} ...\")\n\n with ArrowWriter(path=f\"{save_dir}/raw.arrow\") as writer:\n for line in tqdm(result, desc=\"Writing to raw.arrow ...\"):\n writer.write(line)\n writer.finalize()\n\n # dup a json separately saving duration in case for DynamicBatchSampler ease\n with open(f\"{save_dir}/duration.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n # vocab map, i.e. tokenizer\n with open(f\"{save_dir}/vocab.txt\", \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n\n print(f\"\\nFor {dataset_name}, sample count: {len(result)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\")\n print(f\"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours\")\n\n\nif __name__ == \"__main__\":\n max_workers = 36\n\n tokenizer = \"char\" # \"pinyin\" | \"char\"\n\n SUB_SET = [\"train-clean-100\", \"train-clean-360\", \"train-other-500\"]\n dataset_dir = \"/LibriTTS\"\n dataset_name = f\"LibriTTS_{'_'.join(SUB_SET)}_{tokenizer}\".replace(\"train-clean-\", \"\").replace(\"train-other-\", \"\")\n save_dir = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/data/{dataset_name}\"\n print(f\"\\nPrepare for {dataset_name}, will save to {save_dir}\\n\")\n main()\n\n # For LibriTTS_100_360_500_char, sample count: 354218\n # For LibriTTS_100_360_500_char, vocab size is: 78\n # For LibriTTS_100_360_500_char, total 554.09 hours\n" }, { "path": "src/f5_tts/train/datasets/prepare_ljspeech.py", "content": "import os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport json\nfrom importlib.resources import files\nfrom pathlib import Path\n\nimport soundfile as sf\nfrom datasets.arrow_writer import ArrowWriter\nfrom tqdm import tqdm\n\n\ndef main():\n result = []\n duration_list = []\n text_vocab_set = set()\n\n with open(meta_info, \"r\") as f:\n lines = f.readlines()\n for line in tqdm(lines):\n uttr, text, norm_text = line.split(\"|\")\n norm_text = norm_text.strip()\n wav_path = Path(dataset_dir) / \"wavs\" / f\"{uttr}.wav\"\n duration = sf.info(wav_path).duration\n if duration < 0.4 or duration > 30:\n continue\n result.append({\"audio_path\": str(wav_path), \"text\": norm_text, \"duration\": duration})\n duration_list.append(duration)\n text_vocab_set.update(list(norm_text))\n\n # save preprocessed dataset to disk\n if not os.path.exists(f\"{save_dir}\"):\n os.makedirs(f\"{save_dir}\")\n print(f\"\\nSaving to {save_dir} ...\")\n\n with ArrowWriter(path=f\"{save_dir}/raw.arrow\") as writer:\n for line in tqdm(result, desc=\"Writing to raw.arrow ...\"):\n writer.write(line)\n writer.finalize()\n\n # dup a json separately saving duration in case for DynamicBatchSampler ease\n with open(f\"{save_dir}/duration.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n # vocab map, i.e. tokenizer\n # add alphabets and symbols (optional, if plan to ft on de/fr etc.)\n with open(f\"{save_dir}/vocab.txt\", \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n\n print(f\"\\nFor {dataset_name}, sample count: {len(result)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\")\n print(f\"For {dataset_name}, total {sum(duration_list) / 3600:.2f} hours\")\n\n\nif __name__ == \"__main__\":\n tokenizer = \"char\" # \"pinyin\" | \"char\"\n\n dataset_dir = \"/LJSpeech-1.1\"\n dataset_name = f\"LJSpeech_{tokenizer}\"\n meta_info = os.path.join(dataset_dir, \"metadata.csv\")\n save_dir = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/data/{dataset_name}\"\n print(f\"\\nPrepare for {dataset_name}, will save to {save_dir}\\n\")\n\n main()\n" }, { "path": "src/f5_tts/train/datasets/prepare_wenetspeech4tts.py", "content": "# generate audio text map for WenetSpeech4TTS\n# evaluate for vocab size\n\nimport os\nimport sys\n\n\nsys.path.append(os.getcwd())\n\nimport json\nfrom concurrent.futures import ProcessPoolExecutor\nfrom importlib.resources import files\n\nimport torchaudio\nfrom datasets import Dataset\nfrom tqdm import tqdm\n\nfrom f5_tts.model.utils import convert_char_to_pinyin\n\n\ndef deal_with_sub_path_files(dataset_path, sub_path):\n print(f\"Dealing with: {sub_path}\")\n\n text_dir = os.path.join(dataset_path, sub_path, \"txts\")\n audio_dir = os.path.join(dataset_path, sub_path, \"wavs\")\n text_files = os.listdir(text_dir)\n\n audio_paths, texts, durations = [], [], []\n for text_file in tqdm(text_files):\n with open(os.path.join(text_dir, text_file), \"r\", encoding=\"utf-8\") as file:\n first_line = file.readline().split(\"\\t\")\n audio_nm = first_line[0]\n audio_path = os.path.join(audio_dir, audio_nm + \".wav\")\n text = first_line[1].strip()\n\n audio_paths.append(audio_path)\n\n if tokenizer == \"pinyin\":\n texts.extend(convert_char_to_pinyin([text], polyphone=polyphone))\n elif tokenizer == \"char\":\n texts.append(text)\n\n audio, sample_rate = torchaudio.load(audio_path)\n durations.append(audio.shape[-1] / sample_rate)\n\n return audio_paths, texts, durations\n\n\ndef main():\n assert tokenizer in [\"pinyin\", \"char\"]\n\n audio_path_list, text_list, duration_list = [], [], []\n\n executor = ProcessPoolExecutor(max_workers=max_workers)\n futures = []\n for dataset_path in dataset_paths:\n sub_items = os.listdir(dataset_path)\n sub_paths = [item for item in sub_items if os.path.isdir(os.path.join(dataset_path, item))]\n for sub_path in sub_paths:\n futures.append(executor.submit(deal_with_sub_path_files, dataset_path, sub_path))\n for future in tqdm(futures, total=len(futures)):\n audio_paths, texts, durations = future.result()\n audio_path_list.extend(audio_paths)\n text_list.extend(texts)\n duration_list.extend(durations)\n executor.shutdown()\n\n if not os.path.exists(\"data\"):\n os.makedirs(\"data\")\n\n print(f\"\\nSaving to {save_dir} ...\")\n dataset = Dataset.from_dict({\"audio_path\": audio_path_list, \"text\": text_list, \"duration\": duration_list})\n dataset.save_to_disk(f\"{save_dir}/raw\", max_shard_size=\"2GB\") # arrow format\n\n with open(f\"{save_dir}/duration.json\", \"w\", encoding=\"utf-8\") as f:\n json.dump(\n {\"duration\": duration_list}, f, ensure_ascii=False\n ) # dup a json separately saving duration in case for DynamicBatchSampler ease\n\n print(\"\\nEvaluating vocab size (all characters and symbols / all phonemes) ...\")\n text_vocab_set = set()\n for text in tqdm(text_list):\n text_vocab_set.update(list(text))\n\n # add alphabets and symbols (optional, if plan to ft on de/fr etc.)\n if tokenizer == \"pinyin\":\n text_vocab_set.update([chr(i) for i in range(32, 127)] + [chr(i) for i in range(192, 256)])\n\n with open(f\"{save_dir}/vocab.txt\", \"w\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n print(f\"\\nFor {dataset_name}, sample count: {len(text_list)}\")\n print(f\"For {dataset_name}, vocab size is: {len(text_vocab_set)}\\n\")\n\n\nif __name__ == \"__main__\":\n max_workers = 32\n\n tokenizer = \"pinyin\" # \"pinyin\" | \"char\"\n polyphone = True\n dataset_choice = 1 # 1: Premium, 2: Standard, 3: Basic\n\n dataset_name = (\n [\"WenetSpeech4TTS_Premium\", \"WenetSpeech4TTS_Standard\", \"WenetSpeech4TTS_Basic\"][dataset_choice - 1]\n + \"_\"\n + tokenizer\n )\n dataset_paths = [\n \"/WenetSpeech4TTS/Basic\",\n \"/WenetSpeech4TTS/Standard\",\n \"/WenetSpeech4TTS/Premium\",\n ][-dataset_choice:]\n save_dir = str(files(\"f5_tts\").joinpath(\"../../\")) + f\"/data/{dataset_name}\"\n print(f\"\\nChoose Dataset: {dataset_name}, will save to {save_dir}\\n\")\n\n main()\n\n # Results (if adding alphabets with accents and symbols):\n # WenetSpeech4TTS Basic Standard Premium\n # samples count 3932473 1941220 407494\n # pinyin vocab size 1349 1348 1344 (no polyphone)\n # - - 1459 (polyphone)\n # char vocab size 5264 5219 5042\n\n # vocab size may be slightly different due to rjieba tokenizer and pypinyin (e.g. way of polyphoneme)\n # please be careful if using pretrained model, make sure the vocab.txt is same\n" }, { "path": "src/f5_tts/train/finetune_cli.py", "content": "import argparse\nimport os\nimport shutil\nfrom importlib.resources import files\n\nfrom cached_path import cached_path\n\nfrom f5_tts.model import CFM, DiT, Trainer, UNetT\nfrom f5_tts.model.dataset import load_dataset\nfrom f5_tts.model.utils import get_tokenizer\n\n\n# -------------------------- Dataset Settings --------------------------- #\ntarget_sample_rate = 24000\nn_mel_channels = 100\nhop_length = 256\nwin_length = 1024\nn_fft = 1024\nmel_spec_type = \"vocos\" # 'vocos' or 'bigvgan'\n\n\n# -------------------------- Argument Parsing --------------------------- #\ndef parse_args():\n parser = argparse.ArgumentParser(description=\"Train CFM Model\")\n\n parser.add_argument(\n \"--exp_name\",\n type=str,\n default=\"F5TTS_v1_Base\",\n choices=[\"F5TTS_v1_Base\", \"F5TTS_Base\", \"E2TTS_Base\"],\n help=\"Experiment name\",\n )\n parser.add_argument(\"--dataset_name\", type=str, default=\"Emilia_ZH_EN\", help=\"Name of the dataset to use\")\n parser.add_argument(\"--learning_rate\", type=float, default=1e-5, help=\"Learning rate for training\")\n parser.add_argument(\"--batch_size_per_gpu\", type=int, default=3200, help=\"Batch size per GPU\")\n parser.add_argument(\n \"--batch_size_type\", type=str, default=\"frame\", choices=[\"frame\", \"sample\"], help=\"Batch size type\"\n )\n parser.add_argument(\"--max_samples\", type=int, default=64, help=\"Max sequences per batch\")\n parser.add_argument(\"--grad_accumulation_steps\", type=int, default=1, help=\"Gradient accumulation steps\")\n parser.add_argument(\"--max_grad_norm\", type=float, default=1.0, help=\"Max gradient norm for clipping\")\n parser.add_argument(\"--epochs\", type=int, default=100, help=\"Number of training epochs\")\n parser.add_argument(\"--num_warmup_updates\", type=int, default=20000, help=\"Warmup updates\")\n parser.add_argument(\"--save_per_updates\", type=int, default=50000, help=\"Save checkpoint every N updates\")\n parser.add_argument(\n \"--keep_last_n_checkpoints\",\n type=int,\n default=-1,\n help=\"-1 to keep all, 0 to not save intermediate, > 0 to keep last N checkpoints\",\n )\n parser.add_argument(\"--last_per_updates\", type=int, default=5000, help=\"Save last checkpoint every N updates\")\n parser.add_argument(\"--finetune\", action=\"store_true\", help=\"Use Finetune\")\n parser.add_argument(\"--pretrain\", type=str, default=None, help=\"the path to the checkpoint\")\n parser.add_argument(\n \"--tokenizer\", type=str, default=\"pinyin\", choices=[\"pinyin\", \"char\", \"custom\"], help=\"Tokenizer type\"\n )\n parser.add_argument(\n \"--tokenizer_path\",\n type=str,\n default=None,\n help=\"Path to custom tokenizer vocab file (only used if tokenizer = 'custom')\",\n )\n parser.add_argument(\n \"--log_samples\",\n action=\"store_true\",\n help=\"Log inferenced samples per ckpt save updates\",\n )\n parser.add_argument(\"--logger\", type=str, default=None, choices=[None, \"wandb\", \"tensorboard\"], help=\"logger\")\n parser.add_argument(\n \"--bnb_optimizer\",\n action=\"store_true\",\n help=\"Use 8-bit Adam optimizer from bitsandbytes\",\n )\n\n return parser.parse_args()\n\n\n# -------------------------- Training Settings -------------------------- #\n\n\ndef main():\n args = parse_args()\n\n checkpoint_path = str(files(\"f5_tts\").joinpath(f\"../../ckpts/{args.dataset_name}\"))\n\n # Model parameters based on experiment name\n\n if args.exp_name == \"F5TTS_v1_Base\":\n wandb_resume_id = None\n model_cls = DiT\n model_cfg = dict(\n dim=1024,\n depth=22,\n heads=16,\n ff_mult=2,\n text_dim=512,\n conv_layers=4,\n )\n if args.finetune:\n if args.pretrain is None:\n ckpt_path = str(cached_path(\"hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors\"))\n else:\n ckpt_path = args.pretrain\n\n elif args.exp_name == \"F5TTS_Base\":\n wandb_resume_id = None\n model_cls = DiT\n model_cfg = dict(\n dim=1024,\n depth=22,\n heads=16,\n ff_mult=2,\n text_dim=512,\n text_mask_padding=False,\n conv_layers=4,\n pe_attn_head=1,\n )\n if args.finetune:\n if args.pretrain is None:\n ckpt_path = str(cached_path(\"hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt\"))\n else:\n ckpt_path = args.pretrain\n\n elif args.exp_name == \"E2TTS_Base\":\n wandb_resume_id = None\n model_cls = UNetT\n model_cfg = dict(\n dim=1024,\n depth=24,\n heads=16,\n ff_mult=4,\n text_mask_padding=False,\n pe_attn_head=1,\n )\n if args.finetune:\n if args.pretrain is None:\n ckpt_path = str(cached_path(\"hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt\"))\n else:\n ckpt_path = args.pretrain\n\n if args.finetune:\n if not os.path.isdir(checkpoint_path):\n os.makedirs(checkpoint_path, exist_ok=True)\n\n file_checkpoint = os.path.basename(ckpt_path)\n if not file_checkpoint.startswith(\"pretrained_\"): # Change: Add 'pretrained_' prefix to copied model\n file_checkpoint = \"pretrained_\" + file_checkpoint\n file_checkpoint = os.path.join(checkpoint_path, file_checkpoint)\n if not os.path.isfile(file_checkpoint):\n shutil.copy2(ckpt_path, file_checkpoint)\n print(\"copy checkpoint for finetune\")\n\n # Use the tokenizer and tokenizer_path provided in the command line arguments\n\n tokenizer = args.tokenizer\n if tokenizer == \"custom\":\n if not args.tokenizer_path:\n raise ValueError(\"Custom tokenizer selected, but no tokenizer_path provided.\")\n tokenizer_path = args.tokenizer_path\n else:\n tokenizer_path = args.dataset_name\n\n vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)\n\n print(\"\\nvocab : \", vocab_size)\n print(\"\\nvocoder : \", mel_spec_type)\n\n mel_spec_kwargs = dict(\n n_fft=n_fft,\n hop_length=hop_length,\n win_length=win_length,\n n_mel_channels=n_mel_channels,\n target_sample_rate=target_sample_rate,\n mel_spec_type=mel_spec_type,\n )\n\n model = CFM(\n transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels),\n mel_spec_kwargs=mel_spec_kwargs,\n vocab_char_map=vocab_char_map,\n )\n\n trainer = Trainer(\n model,\n args.epochs,\n args.learning_rate,\n num_warmup_updates=args.num_warmup_updates,\n save_per_updates=args.save_per_updates,\n keep_last_n_checkpoints=args.keep_last_n_checkpoints,\n checkpoint_path=checkpoint_path,\n batch_size_per_gpu=args.batch_size_per_gpu,\n batch_size_type=args.batch_size_type,\n max_samples=args.max_samples,\n grad_accumulation_steps=args.grad_accumulation_steps,\n max_grad_norm=args.max_grad_norm,\n logger=args.logger,\n wandb_project=args.dataset_name,\n wandb_run_name=args.exp_name,\n wandb_resume_id=wandb_resume_id,\n log_samples=args.log_samples,\n last_per_updates=args.last_per_updates,\n bnb_optimizer=args.bnb_optimizer,\n )\n\n train_dataset = load_dataset(args.dataset_name, tokenizer, mel_spec_kwargs=mel_spec_kwargs)\n\n trainer.train(\n train_dataset,\n resumable_with_seed=666, # seed for shuffling dataset\n )\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/train/finetune_gradio.py", "content": "import gc\nimport json\nimport os\nimport platform\nimport queue\nimport random\nimport re\nimport shutil\nimport signal\nimport subprocess\nimport sys\nimport tempfile\nimport threading\nimport time\nfrom glob import glob\nfrom importlib.resources import files\n\nimport click\nimport gradio as gr\nimport librosa\nimport numpy as np\nimport psutil\nimport torch\nimport torchaudio\nfrom cached_path import cached_path\nfrom datasets import Dataset as Dataset_\nfrom datasets.arrow_writer import ArrowWriter\nfrom safetensors.torch import load_file, save_file\nfrom scipy.io import wavfile\n\nfrom f5_tts.api import F5TTS\nfrom f5_tts.infer.utils_infer import transcribe\nfrom f5_tts.model.utils import convert_char_to_pinyin\n\n\ntraining_process = None\nsystem = platform.system()\npython_executable = sys.executable or \"python\"\ntts_api = None\nlast_checkpoint = \"\"\nlast_device = \"\"\nlast_ema = None\n\n\npath_data = str(files(\"f5_tts\").joinpath(\"../../data\"))\npath_project_ckpts = str(files(\"f5_tts\").joinpath(\"../../ckpts\"))\nfile_train = str(files(\"f5_tts\").joinpath(\"train/finetune_cli.py\"))\n\ndevice = (\n \"cuda\"\n if torch.cuda.is_available()\n else \"xpu\"\n if torch.xpu.is_available()\n else \"mps\"\n if torch.backends.mps.is_available()\n else \"cpu\"\n)\n\n\n# Save settings from a JSON file\ndef save_settings(\n project_name,\n exp_name,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n grad_accumulation_steps,\n max_grad_norm,\n epochs,\n num_warmup_updates,\n save_per_updates,\n keep_last_n_checkpoints,\n last_per_updates,\n finetune,\n file_checkpoint_train,\n tokenizer_type,\n tokenizer_file,\n mixed_precision,\n logger,\n ch_8bit_adam,\n):\n path_project = os.path.join(path_project_ckpts, project_name)\n os.makedirs(path_project, exist_ok=True)\n file_setting = os.path.join(path_project, \"setting.json\")\n\n settings = {\n \"exp_name\": exp_name,\n \"learning_rate\": learning_rate,\n \"batch_size_per_gpu\": batch_size_per_gpu,\n \"batch_size_type\": batch_size_type,\n \"max_samples\": max_samples,\n \"grad_accumulation_steps\": grad_accumulation_steps,\n \"max_grad_norm\": max_grad_norm,\n \"epochs\": epochs,\n \"num_warmup_updates\": num_warmup_updates,\n \"save_per_updates\": save_per_updates,\n \"keep_last_n_checkpoints\": keep_last_n_checkpoints,\n \"last_per_updates\": last_per_updates,\n \"finetune\": finetune,\n \"file_checkpoint_train\": file_checkpoint_train,\n \"tokenizer_type\": tokenizer_type,\n \"tokenizer_file\": tokenizer_file,\n \"mixed_precision\": mixed_precision,\n \"logger\": logger,\n \"bnb_optimizer\": ch_8bit_adam,\n }\n with open(file_setting, \"w\") as f:\n json.dump(settings, f, indent=4)\n return \"Settings saved!\"\n\n\n# Load settings from a JSON file\ndef load_settings(project_name):\n project_name = project_name.replace(\"_pinyin\", \"\").replace(\"_char\", \"\")\n path_project = os.path.join(path_project_ckpts, project_name)\n file_setting = os.path.join(path_project, \"setting.json\")\n\n # Default settings\n default_settings = {\n \"exp_name\": \"F5TTS_v1_Base\",\n \"learning_rate\": 1e-5,\n \"batch_size_per_gpu\": 3200,\n \"batch_size_type\": \"frame\",\n \"max_samples\": 64,\n \"grad_accumulation_steps\": 1,\n \"max_grad_norm\": 1.0,\n \"epochs\": 100,\n \"num_warmup_updates\": 100,\n \"save_per_updates\": 500,\n \"keep_last_n_checkpoints\": -1,\n \"last_per_updates\": 100,\n \"finetune\": True,\n \"file_checkpoint_train\": \"\",\n \"tokenizer_type\": \"pinyin\",\n \"tokenizer_file\": \"\",\n \"mixed_precision\": \"fp16\",\n \"logger\": \"none\",\n \"bnb_optimizer\": False,\n }\n if device == \"mps\":\n default_settings[\"mixed_precision\"] = \"none\"\n\n # Load settings from file if it exists\n if os.path.isfile(file_setting):\n with open(file_setting, \"r\") as f:\n file_settings = json.load(f)\n default_settings.update(file_settings)\n\n # Return as a tuple in the correct order\n return (\n default_settings[\"exp_name\"],\n default_settings[\"learning_rate\"],\n default_settings[\"batch_size_per_gpu\"],\n default_settings[\"batch_size_type\"],\n default_settings[\"max_samples\"],\n default_settings[\"grad_accumulation_steps\"],\n default_settings[\"max_grad_norm\"],\n default_settings[\"epochs\"],\n default_settings[\"num_warmup_updates\"],\n default_settings[\"save_per_updates\"],\n default_settings[\"keep_last_n_checkpoints\"],\n default_settings[\"last_per_updates\"],\n default_settings[\"finetune\"],\n default_settings[\"file_checkpoint_train\"],\n default_settings[\"tokenizer_type\"],\n default_settings[\"tokenizer_file\"],\n default_settings[\"mixed_precision\"],\n default_settings[\"logger\"],\n default_settings[\"bnb_optimizer\"],\n )\n\n\n# Load metadata\ndef get_audio_duration(audio_path):\n \"\"\"Calculate the duration mono of an audio file.\"\"\"\n audio, sample_rate = torchaudio.load(audio_path)\n return audio.shape[1] / sample_rate\n\n\nclass Slicer: # https://github.com/RVC-Boss/GPT-SoVITS/blob/main/tools/slicer2.py\n def __init__(\n self,\n sr: int,\n threshold: float = -40.0,\n min_length: int = 20000, # 20 seconds\n min_interval: int = 300,\n hop_size: int = 20,\n max_sil_kept: int = 2000,\n ):\n if not min_length >= min_interval >= hop_size:\n raise ValueError(\"The following condition must be satisfied: min_length >= min_interval >= hop_size\")\n if not max_sil_kept >= hop_size:\n raise ValueError(\"The following condition must be satisfied: max_sil_kept >= hop_size\")\n min_interval = sr * min_interval / 1000\n self.threshold = 10 ** (threshold / 20.0)\n self.hop_size = round(sr * hop_size / 1000)\n self.win_size = min(round(min_interval), 4 * self.hop_size)\n self.min_length = round(sr * min_length / 1000 / self.hop_size)\n self.min_interval = round(min_interval / self.hop_size)\n self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)\n\n def _apply_slice(self, waveform, begin, end):\n if len(waveform.shape) > 1:\n return waveform[:, begin * self.hop_size : min(waveform.shape[1], end * self.hop_size)]\n else:\n return waveform[begin * self.hop_size : min(waveform.shape[0], end * self.hop_size)]\n\n # @timeit\n def slice(self, waveform):\n if len(waveform.shape) > 1:\n samples = waveform.mean(axis=0)\n else:\n samples = waveform\n if samples.shape[0] <= self.min_length:\n return [waveform]\n rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)\n sil_tags = []\n silence_start = None\n clip_start = 0\n for i, rms in enumerate(rms_list):\n # Keep looping while frame is silent.\n if rms < self.threshold:\n # Record start of silent frames.\n if silence_start is None:\n silence_start = i\n continue\n # Keep looping while frame is not silent and silence start has not been recorded.\n if silence_start is None:\n continue\n # Clear recorded silence start if interval is not enough or clip is too short\n is_leading_silence = silence_start == 0 and i > self.max_sil_kept\n need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length\n if not is_leading_silence and not need_slice_middle:\n silence_start = None\n continue\n # Need slicing. Record the range of silent frames to be removed.\n if i - silence_start <= self.max_sil_kept:\n pos = rms_list[silence_start : i + 1].argmin() + silence_start\n if silence_start == 0:\n sil_tags.append((0, pos))\n else:\n sil_tags.append((pos, pos))\n clip_start = pos\n elif i - silence_start <= self.max_sil_kept * 2:\n pos = rms_list[i - self.max_sil_kept : silence_start + self.max_sil_kept + 1].argmin()\n pos += i - self.max_sil_kept\n pos_l = rms_list[silence_start : silence_start + self.max_sil_kept + 1].argmin() + silence_start\n pos_r = rms_list[i - self.max_sil_kept : i + 1].argmin() + i - self.max_sil_kept\n if silence_start == 0:\n sil_tags.append((0, pos_r))\n clip_start = pos_r\n else:\n sil_tags.append((min(pos_l, pos), max(pos_r, pos)))\n clip_start = max(pos_r, pos)\n else:\n pos_l = rms_list[silence_start : silence_start + self.max_sil_kept + 1].argmin() + silence_start\n pos_r = rms_list[i - self.max_sil_kept : i + 1].argmin() + i - self.max_sil_kept\n if silence_start == 0:\n sil_tags.append((0, pos_r))\n else:\n sil_tags.append((pos_l, pos_r))\n clip_start = pos_r\n silence_start = None\n # Deal with trailing silence.\n total_frames = rms_list.shape[0]\n if silence_start is not None and total_frames - silence_start >= self.min_interval:\n silence_end = min(total_frames, silence_start + self.max_sil_kept)\n pos = rms_list[silence_start : silence_end + 1].argmin() + silence_start\n sil_tags.append((pos, total_frames + 1))\n # Apply and return slices: [chunk, start, end]\n if len(sil_tags) == 0:\n return [[waveform, 0, int(total_frames * self.hop_size)]]\n else:\n chunks = []\n if sil_tags[0][0] > 0:\n chunks.append([self._apply_slice(waveform, 0, sil_tags[0][0]), 0, int(sil_tags[0][0] * self.hop_size)])\n for i in range(len(sil_tags) - 1):\n chunks.append(\n [\n self._apply_slice(waveform, sil_tags[i][1], sil_tags[i + 1][0]),\n int(sil_tags[i][1] * self.hop_size),\n int(sil_tags[i + 1][0] * self.hop_size),\n ]\n )\n if sil_tags[-1][1] < total_frames:\n chunks.append(\n [\n self._apply_slice(waveform, sil_tags[-1][1], total_frames),\n int(sil_tags[-1][1] * self.hop_size),\n int(total_frames * self.hop_size),\n ]\n )\n return chunks\n\n\n# terminal\ndef terminate_process_tree(pid, including_parent=True):\n try:\n parent = psutil.Process(pid)\n except psutil.NoSuchProcess:\n # Process already terminated\n return\n\n children = parent.children(recursive=True)\n for child in children:\n try:\n os.kill(child.pid, signal.SIGTERM) # or signal.SIGKILL\n except OSError:\n pass\n if including_parent:\n try:\n os.kill(parent.pid, signal.SIGTERM) # or signal.SIGKILL\n except OSError:\n pass\n\n\ndef terminate_process(pid):\n if system == \"Windows\":\n cmd = f\"taskkill /t /f /pid {pid}\"\n os.system(cmd)\n else:\n terminate_process_tree(pid)\n\n\ndef start_training(\n dataset_name,\n exp_name,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n grad_accumulation_steps,\n max_grad_norm,\n epochs,\n num_warmup_updates,\n save_per_updates,\n keep_last_n_checkpoints,\n last_per_updates,\n finetune,\n file_checkpoint_train,\n tokenizer_type,\n tokenizer_file,\n mixed_precision,\n stream,\n logger,\n ch_8bit_adam,\n):\n global training_process, tts_api, stop_signal\n\n if tts_api is not None:\n if tts_api is not None:\n del tts_api\n\n gc.collect()\n torch.cuda.empty_cache()\n tts_api = None\n\n path_project = os.path.join(path_data, dataset_name)\n\n if not os.path.isdir(path_project):\n yield (\n f\"There is not project with name {dataset_name}\",\n gr.update(interactive=True),\n gr.update(interactive=False),\n )\n return\n\n file_raw = os.path.join(path_project, \"raw.arrow\")\n if not os.path.isfile(file_raw):\n yield f\"There is no file {file_raw}\", gr.update(interactive=True), gr.update(interactive=False)\n return\n\n # Check if a training process is already running\n if training_process is not None:\n return \"Train run already!\", gr.update(interactive=False), gr.update(interactive=True)\n\n yield \"start train\", gr.update(interactive=False), gr.update(interactive=False)\n\n # Command to run the training script with the specified arguments\n\n if tokenizer_file == \"\":\n if dataset_name.endswith(\"_pinyin\"):\n tokenizer_type = \"pinyin\"\n elif dataset_name.endswith(\"_char\"):\n tokenizer_type = \"char\"\n else:\n tokenizer_type = \"custom\"\n\n dataset_name = dataset_name.replace(\"_pinyin\", \"\").replace(\"_char\", \"\")\n\n if mixed_precision != \"none\":\n fp16 = f\"--mixed_precision={mixed_precision}\"\n else:\n fp16 = \"\"\n\n cmd = (\n f'accelerate launch {fp16} \"{file_train}\" --exp_name {exp_name}'\n f\" --learning_rate {learning_rate}\"\n f\" --batch_size_per_gpu {batch_size_per_gpu}\"\n f\" --batch_size_type {batch_size_type}\"\n f\" --max_samples {max_samples}\"\n f\" --grad_accumulation_steps {grad_accumulation_steps}\"\n f\" --max_grad_norm {max_grad_norm}\"\n f\" --epochs {epochs}\"\n f\" --num_warmup_updates {num_warmup_updates}\"\n f\" --save_per_updates {save_per_updates}\"\n f\" --keep_last_n_checkpoints {keep_last_n_checkpoints}\"\n f\" --last_per_updates {last_per_updates}\"\n f\" --dataset_name {dataset_name}\"\n )\n\n if finetune:\n cmd += \" --finetune\"\n\n if file_checkpoint_train != \"\":\n cmd += f' --pretrain \"{file_checkpoint_train}\"'\n\n if tokenizer_file != \"\":\n cmd += f\" --tokenizer_path {tokenizer_file}\"\n\n cmd += f\" --tokenizer {tokenizer_type}\"\n\n if logger != \"none\":\n cmd += f\" --logger {logger}\"\n\n cmd += \" --log_samples\"\n\n if ch_8bit_adam:\n cmd += \" --bnb_optimizer\"\n\n print(\"run command : \\n\" + cmd + \"\\n\")\n\n save_settings(\n dataset_name,\n exp_name,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n grad_accumulation_steps,\n max_grad_norm,\n epochs,\n num_warmup_updates,\n save_per_updates,\n keep_last_n_checkpoints,\n last_per_updates,\n finetune,\n file_checkpoint_train,\n tokenizer_type,\n tokenizer_file,\n mixed_precision,\n logger,\n ch_8bit_adam,\n )\n\n try:\n if not stream:\n # Start the training process\n training_process = subprocess.Popen(cmd, shell=True)\n\n time.sleep(5)\n yield \"train start\", gr.update(interactive=False), gr.update(interactive=True)\n\n # Wait for the training process to finish\n training_process.wait()\n else:\n\n def stream_output(pipe, output_queue):\n try:\n for line in iter(pipe.readline, \"\"):\n output_queue.put(line)\n except Exception as e:\n output_queue.put(f\"Error reading pipe: {str(e)}\")\n finally:\n pipe.close()\n\n env = os.environ.copy()\n env[\"PYTHONUNBUFFERED\"] = \"1\"\n\n training_process = subprocess.Popen(\n cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, env=env\n )\n yield \"Training started ...\", gr.update(interactive=False), gr.update(interactive=True)\n\n stdout_queue = queue.Queue()\n stderr_queue = queue.Queue()\n\n stdout_thread = threading.Thread(target=stream_output, args=(training_process.stdout, stdout_queue))\n stderr_thread = threading.Thread(target=stream_output, args=(training_process.stderr, stderr_queue))\n stdout_thread.daemon = True\n stderr_thread.daemon = True\n stdout_thread.start()\n stderr_thread.start()\n stop_signal = False\n while True:\n if stop_signal:\n training_process.terminate()\n time.sleep(0.5)\n if training_process.poll() is None:\n training_process.kill()\n yield \"Training stopped by user.\", gr.update(interactive=True), gr.update(interactive=False)\n break\n\n process_status = training_process.poll()\n\n # Handle stdout\n try:\n while True:\n output = stdout_queue.get_nowait()\n print(output, end=\"\")\n match = re.search(\n r\"Epoch (\\d+)/(\\d+):\\s+(\\d+)%\\|.*\\[(\\d+:\\d+)<.*?loss=(\\d+\\.\\d+), update=(\\d+)\", output\n )\n if match:\n current_epoch = match.group(1)\n total_epochs = match.group(2)\n percent_complete = match.group(3)\n elapsed_time = match.group(4)\n loss = match.group(5)\n current_update = match.group(6)\n message = (\n f\"Epoch: {current_epoch}/{total_epochs}, \"\n f\"Progress: {percent_complete}%, \"\n f\"Elapsed Time: {elapsed_time}, \"\n f\"Loss: {loss}, \"\n f\"Update: {current_update}\"\n )\n yield message, gr.update(interactive=False), gr.update(interactive=True)\n elif output.strip():\n yield output, gr.update(interactive=False), gr.update(interactive=True)\n except queue.Empty:\n pass\n\n # Handle stderr\n try:\n while True:\n error_output = stderr_queue.get_nowait()\n print(error_output, end=\"\")\n if error_output.strip():\n yield f\"{error_output.strip()}\", gr.update(interactive=False), gr.update(interactive=True)\n except queue.Empty:\n pass\n\n if process_status is not None and stdout_queue.empty() and stderr_queue.empty():\n if process_status != 0:\n yield (\n f\"Process crashed with exit code {process_status}!\",\n gr.update(interactive=False),\n gr.update(interactive=True),\n )\n else:\n yield (\n \"Training complete or paused ...\",\n gr.update(interactive=False),\n gr.update(interactive=True),\n )\n break\n\n # Small sleep to prevent CPU thrashing\n time.sleep(0.1)\n\n # Clean up\n training_process.stdout.close()\n training_process.stderr.close()\n training_process.wait()\n\n time.sleep(1)\n\n if training_process is None:\n text_info = \"Train stopped !\"\n else:\n text_info = \"Train complete at end !\"\n\n except Exception as e: # Catch all exceptions\n # Ensure that we reset the training process variable in case of an error\n text_info = f\"An error occurred: {str(e)}\"\n\n training_process = None\n\n yield text_info, gr.update(interactive=True), gr.update(interactive=False)\n\n\ndef stop_training():\n global training_process, stop_signal\n\n if training_process is None:\n return \"Train not running !\", gr.update(interactive=True), gr.update(interactive=False)\n terminate_process_tree(training_process.pid)\n # training_process = None\n stop_signal = True\n return \"Train stopped !\", gr.update(interactive=True), gr.update(interactive=False)\n\n\ndef get_list_projects():\n project_list = []\n for folder in os.listdir(path_data):\n path_folder = os.path.join(path_data, folder)\n if not os.path.isdir(path_folder):\n continue\n folder = folder.lower()\n if folder == \"emilia_zh_en_pinyin\":\n continue\n project_list.append(folder)\n\n projects_selelect = None if not project_list else project_list[-1]\n\n return project_list, projects_selelect\n\n\ndef create_data_project(name, tokenizer_type):\n name += \"_\" + tokenizer_type\n os.makedirs(os.path.join(path_data, name), exist_ok=True)\n os.makedirs(os.path.join(path_data, name, \"dataset\"), exist_ok=True)\n project_list, projects_selelect = get_list_projects()\n return gr.update(choices=project_list, value=name)\n\n\ndef transcribe_all(name_project, audio_files, language, user=False, progress=gr.Progress()):\n path_project = os.path.join(path_data, name_project)\n path_dataset = os.path.join(path_project, \"dataset\")\n path_project_wavs = os.path.join(path_project, \"wavs\")\n file_metadata = os.path.join(path_project, \"metadata.csv\")\n\n if not user:\n if audio_files is None:\n return \"You need to load an audio file.\"\n\n if os.path.isdir(path_project_wavs):\n shutil.rmtree(path_project_wavs)\n\n if os.path.isfile(file_metadata):\n os.remove(file_metadata)\n\n os.makedirs(path_project_wavs, exist_ok=True)\n\n if user:\n file_audios = [\n file\n for format in (\"*.wav\", \"*.ogg\", \"*.opus\", \"*.mp3\", \"*.flac\")\n for file in glob(os.path.join(path_dataset, format))\n ]\n if file_audios == []:\n return \"No audio file was found in the dataset.\"\n else:\n file_audios = audio_files\n\n alpha = 0.5\n _max = 1.0\n slicer = Slicer(24000)\n\n num = 0\n error_num = 0\n data = \"\"\n for file_audio in progress.tqdm(file_audios, desc=\"transcribe files\", total=len((file_audios))):\n audio, _ = librosa.load(file_audio, sr=24000, mono=True)\n\n list_slicer = slicer.slice(audio)\n for chunk, start, end in progress.tqdm(list_slicer, total=len(list_slicer), desc=\"slicer files\"):\n name_segment = os.path.join(f\"segment_{num}\")\n file_segment = os.path.join(path_project_wavs, f\"{name_segment}.wav\")\n\n tmp_max = np.abs(chunk).max()\n if tmp_max > 1:\n chunk /= tmp_max\n chunk = (chunk / tmp_max * (_max * alpha)) + (1 - alpha) * chunk\n wavfile.write(file_segment, 24000, (chunk * 32767).astype(np.int16))\n\n try:\n text = transcribe(file_segment, language)\n text = text.strip()\n\n data += f\"{name_segment}|{text}\\n\"\n\n num += 1\n except: # noqa: E722\n error_num += 1\n\n with open(file_metadata, \"w\", encoding=\"utf-8-sig\") as f:\n f.write(data)\n\n if error_num != []:\n error_text = f\"\\nerror files : {error_num}\"\n else:\n error_text = \"\"\n\n return f\"transcribe complete samples : {num}\\npath : {path_project_wavs}{error_text}\"\n\n\ndef format_seconds_to_hms(seconds):\n hours = int(seconds / 3600)\n minutes = int((seconds % 3600) / 60)\n seconds = seconds % 60\n return \"{:02d}:{:02d}:{:02d}\".format(hours, minutes, int(seconds))\n\n\ndef get_correct_audio_path(\n audio_input,\n base_path=\"wavs\",\n supported_formats=(\"wav\", \"mp3\", \"aac\", \"flac\", \"m4a\", \"alac\", \"ogg\", \"aiff\", \"wma\", \"amr\"),\n):\n file_audio = None\n\n # Helper function to check if file has a supported extension\n def has_supported_extension(file_name):\n return any(file_name.endswith(f\".{ext}\") for ext in supported_formats)\n\n # Case 1: If it's a full path with a valid extension, use it directly\n if os.path.isabs(audio_input) and has_supported_extension(audio_input):\n file_audio = audio_input\n\n # Case 2: If it has a supported extension but is not a full path\n elif has_supported_extension(audio_input) and not os.path.isabs(audio_input):\n file_audio = os.path.join(base_path, audio_input)\n\n # Case 3: If only the name is given (no extension and not a full path)\n elif not has_supported_extension(audio_input) and not os.path.isabs(audio_input):\n for ext in supported_formats:\n potential_file = os.path.join(base_path, f\"{audio_input}.{ext}\")\n if os.path.exists(potential_file):\n file_audio = potential_file\n break\n else:\n file_audio = os.path.join(base_path, f\"{audio_input}.{supported_formats[0]}\")\n return file_audio\n\n\ndef create_metadata(name_project, ch_tokenizer, progress=gr.Progress()):\n path_project = os.path.join(path_data, name_project)\n path_project_wavs = os.path.join(path_project, \"wavs\")\n file_metadata = os.path.join(path_project, \"metadata.csv\")\n file_raw = os.path.join(path_project, \"raw.arrow\")\n file_duration = os.path.join(path_project, \"duration.json\")\n file_vocab = os.path.join(path_project, \"vocab.txt\")\n\n if not os.path.isfile(file_metadata):\n return \"The file was not found in \" + file_metadata, \"\"\n\n with open(file_metadata, \"r\", encoding=\"utf-8-sig\") as f:\n data = f.read()\n\n audio_path_list = []\n text_list = []\n duration_list = []\n\n count = data.split(\"\\n\")\n lenght = 0\n result = []\n error_files = []\n text_vocab_set = set()\n for line in progress.tqdm(data.split(\"\\n\"), total=count):\n sp_line = line.split(\"|\")\n if len(sp_line) != 2:\n continue\n name_audio, text = sp_line[:2]\n\n file_audio = get_correct_audio_path(name_audio, path_project_wavs)\n\n if not os.path.isfile(file_audio):\n error_files.append([file_audio, \"error path\"])\n continue\n\n try:\n duration = get_audio_duration(file_audio)\n except Exception as e:\n error_files.append([file_audio, \"duration\"])\n print(f\"Error processing {file_audio}: {e}\")\n continue\n\n if duration < 1 or duration > 30:\n if duration > 30:\n error_files.append([file_audio, \"duration > 30 sec\"])\n if duration < 1:\n error_files.append([file_audio, \"duration < 1 sec \"])\n continue\n if len(text) < 3:\n error_files.append([file_audio, \"very short text length 3\"])\n continue\n\n text = text.strip()\n text = convert_char_to_pinyin([text], polyphone=True)[0]\n\n audio_path_list.append(file_audio)\n duration_list.append(duration)\n text_list.append(text)\n\n result.append({\"audio_path\": file_audio, \"text\": text, \"duration\": duration})\n if ch_tokenizer:\n text_vocab_set.update(list(text))\n\n lenght += duration\n\n if duration_list == []:\n return f\"Error: No audio files found in the specified path : {path_project_wavs}\", \"\"\n\n min_second = round(min(duration_list), 2)\n max_second = round(max(duration_list), 2)\n\n with ArrowWriter(path=file_raw) as writer:\n for line in progress.tqdm(result, total=len(result), desc=\"prepare data\"):\n writer.write(line)\n writer.finalize()\n\n with open(file_duration, \"w\") as f:\n json.dump({\"duration\": duration_list}, f, ensure_ascii=False)\n\n new_vocal = \"\"\n if not ch_tokenizer:\n if not os.path.isfile(file_vocab):\n file_vocab_finetune = os.path.join(path_data, \"Emilia_ZH_EN_pinyin/vocab.txt\")\n if not os.path.isfile(file_vocab_finetune):\n return \"Error: Vocabulary file 'Emilia_ZH_EN_pinyin' not found!\", \"\"\n shutil.copy2(file_vocab_finetune, file_vocab)\n\n with open(file_vocab, \"r\", encoding=\"utf-8-sig\") as f:\n vocab_char_map = {}\n for i, char in enumerate(f):\n vocab_char_map[char[:-1]] = i\n vocab_size = len(vocab_char_map)\n\n else:\n with open(file_vocab, \"w\", encoding=\"utf-8-sig\") as f:\n for vocab in sorted(text_vocab_set):\n f.write(vocab + \"\\n\")\n new_vocal += vocab + \"\\n\"\n vocab_size = len(text_vocab_set)\n\n if error_files != []:\n error_text = \"\\n\".join([\" = \".join(item) for item in error_files])\n else:\n error_text = \"\"\n\n return (\n f\"prepare complete \\nsamples : {len(text_list)}\\ntime data : {format_seconds_to_hms(lenght)}\\nmin sec : {min_second}\\nmax sec : {max_second}\\nfile_arrow : {file_raw}\\nvocab : {vocab_size}\\n{error_text}\",\n new_vocal,\n )\n\n\ndef check_user(value):\n return gr.update(visible=not value), gr.update(visible=value)\n\n\ndef calculate_train(\n name_project,\n epochs,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n num_warmup_updates,\n finetune,\n):\n path_project = os.path.join(path_data, name_project)\n file_duration = os.path.join(path_project, \"duration.json\")\n\n hop_length = 256\n sampling_rate = 24000\n\n if not os.path.isfile(file_duration):\n return (\n epochs,\n learning_rate,\n batch_size_per_gpu,\n max_samples,\n num_warmup_updates,\n \"project not found !\",\n )\n\n with open(file_duration, \"r\") as file:\n data = json.load(file)\n\n duration_list = data[\"duration\"]\n max_sample_length = max(duration_list) * sampling_rate / hop_length\n total_samples = len(duration_list)\n total_duration = sum(duration_list)\n\n if torch.cuda.is_available():\n gpu_count = torch.cuda.device_count()\n total_memory = 0\n for i in range(gpu_count):\n gpu_properties = torch.cuda.get_device_properties(i)\n total_memory += gpu_properties.total_memory / (1024**3) # in GB\n elif torch.xpu.is_available():\n gpu_count = torch.xpu.device_count()\n total_memory = 0\n for i in range(gpu_count):\n gpu_properties = torch.xpu.get_device_properties(i)\n total_memory += gpu_properties.total_memory / (1024**3)\n elif torch.backends.mps.is_available():\n gpu_count = 1\n total_memory = psutil.virtual_memory().available / (1024**3)\n\n avg_gpu_memory = total_memory / gpu_count\n\n # rough estimate of batch size\n if batch_size_type == \"frame\":\n batch_size_per_gpu = max(int(38400 * (avg_gpu_memory - 5) / 75), int(max_sample_length))\n elif batch_size_type == \"sample\":\n batch_size_per_gpu = int(200 / (total_duration / total_samples))\n\n if total_samples < 64:\n max_samples = int(total_samples * 0.25)\n\n num_warmup_updates = max(num_warmup_updates, int(total_samples * 0.05))\n\n # take 1.2M updates as the maximum\n max_updates = 1200000\n\n if batch_size_type == \"frame\":\n mini_batch_duration = batch_size_per_gpu * gpu_count * hop_length / sampling_rate\n updates_per_epoch = total_duration / mini_batch_duration\n elif batch_size_type == \"sample\":\n updates_per_epoch = total_samples / batch_size_per_gpu / gpu_count\n\n epochs = int(max_updates / updates_per_epoch)\n\n if finetune:\n learning_rate = 1e-5\n else:\n learning_rate = 7.5e-5\n\n return (\n epochs,\n learning_rate,\n batch_size_per_gpu,\n max_samples,\n num_warmup_updates,\n total_samples,\n )\n\n\ndef prune_checkpoint(checkpoint_path: str, new_checkpoint_path: str, save_ema: bool, safetensors: bool) -> str:\n try:\n checkpoint = torch.load(checkpoint_path, weights_only=True)\n print(\"Original Checkpoint Keys:\", checkpoint.keys())\n\n to_retain = \"ema_model_state_dict\" if save_ema else \"model_state_dict\"\n try:\n model_state_dict_to_retain = checkpoint[to_retain]\n except KeyError:\n return f\"{to_retain} not found in the checkpoint.\"\n\n if safetensors:\n new_checkpoint_path = new_checkpoint_path.replace(\".pt\", \".safetensors\")\n save_file(model_state_dict_to_retain, new_checkpoint_path)\n else:\n new_checkpoint_path = new_checkpoint_path.replace(\".safetensors\", \".pt\")\n new_checkpoint = {\"ema_model_state_dict\": model_state_dict_to_retain}\n torch.save(new_checkpoint, new_checkpoint_path)\n\n return f\"New checkpoint saved at: {new_checkpoint_path}\"\n\n except Exception as e:\n return f\"An error occurred: {e}\"\n\n\ndef expand_model_embeddings(ckpt_path, new_ckpt_path, num_new_tokens=42):\n seed = 666\n random.seed(seed)\n os.environ[\"PYTHONHASHSEED\"] = str(seed)\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.benchmark = False\n\n if ckpt_path.endswith(\".safetensors\"):\n ckpt = load_file(ckpt_path, device=\"cpu\")\n ckpt = {\"ema_model_state_dict\": ckpt}\n elif ckpt_path.endswith(\".pt\"):\n ckpt = torch.load(ckpt_path, map_location=\"cpu\")\n\n ema_sd = ckpt.get(\"ema_model_state_dict\", {})\n embed_key_ema = \"ema_model.transformer.text_embed.text_embed.weight\"\n old_embed_ema = ema_sd[embed_key_ema]\n\n vocab_old = old_embed_ema.size(0)\n embed_dim = old_embed_ema.size(1)\n vocab_new = vocab_old + num_new_tokens\n\n def expand_embeddings(old_embeddings):\n new_embeddings = torch.zeros((vocab_new, embed_dim))\n new_embeddings[:vocab_old] = old_embeddings\n new_embeddings[vocab_old:] = torch.randn((num_new_tokens, embed_dim))\n return new_embeddings\n\n ema_sd[embed_key_ema] = expand_embeddings(ema_sd[embed_key_ema])\n\n if new_ckpt_path.endswith(\".safetensors\"):\n save_file(ema_sd, new_ckpt_path)\n elif new_ckpt_path.endswith(\".pt\"):\n torch.save(ckpt, new_ckpt_path)\n\n return vocab_new\n\n\ndef vocab_count(text):\n return str(len(text.split(\",\")))\n\n\ndef vocab_extend(project_name, symbols, model_type):\n if symbols == \"\":\n return \"Symbols empty!\"\n\n name_project = project_name\n path_project = os.path.join(path_data, name_project)\n file_vocab_project = os.path.join(path_project, \"vocab.txt\")\n\n file_vocab = os.path.join(path_data, \"Emilia_ZH_EN_pinyin/vocab.txt\")\n if not os.path.isfile(file_vocab):\n return f\"the file {file_vocab} not found !\"\n\n symbols = symbols.split(\",\")\n if symbols == []:\n return \"Symbols to extend not found.\"\n\n with open(file_vocab, \"r\", encoding=\"utf-8-sig\") as f:\n data = f.read()\n vocab = data.split(\"\\n\")\n vocab_check = set(vocab)\n\n miss_symbols = []\n for item in symbols:\n item = item.replace(\" \", \"\")\n if item in vocab_check:\n continue\n miss_symbols.append(item)\n\n if miss_symbols == []:\n return \"Symbols are okay no need to extend.\"\n\n size_vocab = len(vocab)\n vocab.pop()\n for item in miss_symbols:\n vocab.append(item)\n\n vocab.append(\"\")\n\n with open(file_vocab_project, \"w\", encoding=\"utf-8\") as f:\n f.write(\"\\n\".join(vocab))\n\n if model_type == \"F5TTS_v1_Base\":\n ckpt_path = str(cached_path(\"hf://SWivid/F5-TTS/F5TTS_v1_Base/model_1250000.safetensors\"))\n elif model_type == \"F5TTS_Base\":\n ckpt_path = str(cached_path(\"hf://SWivid/F5-TTS/F5TTS_Base/model_1200000.pt\"))\n elif model_type == \"E2TTS_Base\":\n ckpt_path = str(cached_path(\"hf://SWivid/E2-TTS/E2TTS_Base/model_1200000.pt\"))\n\n vocab_size_new = len(miss_symbols)\n\n dataset_name = name_project.replace(\"_pinyin\", \"\").replace(\"_char\", \"\")\n new_ckpt_path = os.path.join(path_project_ckpts, dataset_name)\n os.makedirs(new_ckpt_path, exist_ok=True)\n\n # Add pretrained_ prefix to model when copying for consistency with finetune_cli.py\n new_ckpt_file = os.path.join(new_ckpt_path, \"pretrained_\" + os.path.basename(ckpt_path))\n\n size = expand_model_embeddings(ckpt_path, new_ckpt_file, num_new_tokens=vocab_size_new)\n\n vocab_new = \"\\n\".join(miss_symbols)\n return f\"vocab old size : {size_vocab}\\nvocab new size : {size}\\nvocab add : {vocab_size_new}\\nnew symbols :\\n{vocab_new}\"\n\n\ndef vocab_check(project_name, tokenizer_type):\n name_project = project_name\n path_project = os.path.join(path_data, name_project)\n\n file_metadata = os.path.join(path_project, \"metadata.csv\")\n\n file_vocab = os.path.join(path_data, \"Emilia_ZH_EN_pinyin/vocab.txt\")\n if not os.path.isfile(file_vocab):\n return f\"the file {file_vocab} not found !\", \"\"\n\n with open(file_vocab, \"r\", encoding=\"utf-8-sig\") as f:\n data = f.read()\n vocab = data.split(\"\\n\")\n vocab = set(vocab)\n\n if not os.path.isfile(file_metadata):\n return f\"the file {file_metadata} not found !\", \"\"\n\n with open(file_metadata, \"r\", encoding=\"utf-8-sig\") as f:\n data = f.read()\n\n miss_symbols = []\n miss_symbols_keep = {}\n for item in data.split(\"\\n\"):\n sp = item.split(\"|\")\n if len(sp) != 2:\n continue\n\n text = sp[1].strip()\n if tokenizer_type == \"pinyin\":\n text = convert_char_to_pinyin([text], polyphone=True)[0]\n\n for t in text:\n if t not in vocab and t not in miss_symbols_keep:\n miss_symbols.append(t)\n miss_symbols_keep[t] = t\n\n if miss_symbols == []:\n vocab_miss = \"\"\n info = \"You can train using your language !\"\n else:\n vocab_miss = \",\".join(miss_symbols)\n info = f\"The following {len(miss_symbols)} symbols are missing in your language\\n\\n\"\n\n return info, vocab_miss\n\n\ndef get_random_sample_prepare(project_name):\n name_project = project_name\n path_project = os.path.join(path_data, name_project)\n file_arrow = os.path.join(path_project, \"raw.arrow\")\n if not os.path.isfile(file_arrow):\n return \"\", None\n dataset = Dataset_.from_file(file_arrow)\n random_sample = dataset.shuffle(seed=random.randint(0, 1000)).select([0])\n text = \"[\" + \" , \".join([\"' \" + t + \" '\" for t in random_sample[\"text\"][0]]) + \"]\"\n audio_path = random_sample[\"audio_path\"][0]\n return text, audio_path\n\n\ndef get_random_sample_transcribe(project_name):\n name_project = project_name\n path_project = os.path.join(path_data, name_project)\n file_metadata = os.path.join(path_project, \"metadata.csv\")\n if not os.path.isfile(file_metadata):\n return \"\", None\n\n data = \"\"\n with open(file_metadata, \"r\", encoding=\"utf-8-sig\") as f:\n data = f.read()\n\n list_data = []\n for item in data.split(\"\\n\"):\n sp = item.split(\"|\")\n if len(sp) != 2:\n continue\n\n # fixed audio when it is absolute\n file_audio = get_correct_audio_path(sp[0], os.path.join(path_project, \"wavs\"))\n list_data.append([file_audio, sp[1]])\n\n if list_data == []:\n return \"\", None\n\n random_item = random.choice(list_data)\n\n return random_item[1], random_item[0]\n\n\ndef get_random_sample_infer(project_name):\n text, audio = get_random_sample_transcribe(project_name)\n return (\n text,\n text,\n audio,\n )\n\n\ndef infer(\n project, file_checkpoint, exp_name, ref_text, ref_audio, gen_text, nfe_step, use_ema, speed, seed, remove_silence\n):\n global last_checkpoint, last_device, tts_api, last_ema\n\n if not os.path.isfile(file_checkpoint):\n return None, \"checkpoint not found!\"\n\n if training_process is not None:\n device_test = \"cpu\"\n else:\n device_test = None\n\n if last_checkpoint != file_checkpoint or last_device != device_test or last_ema != use_ema or tts_api is None:\n if last_checkpoint != file_checkpoint:\n last_checkpoint = file_checkpoint\n\n if last_device != device_test:\n last_device = device_test\n\n if last_ema != use_ema:\n last_ema = use_ema\n\n vocab_file = os.path.join(path_data, project, \"vocab.txt\")\n\n tts_api = F5TTS(\n model=exp_name, ckpt_file=file_checkpoint, vocab_file=vocab_file, device=device_test, use_ema=use_ema\n )\n\n print(\"update >> \", device_test, file_checkpoint, use_ema)\n\n if seed == -1: # -1 used for random\n seed = None\n\n with tempfile.NamedTemporaryFile(delete=False, suffix=\".wav\") as f:\n tts_api.infer(\n ref_file=ref_audio,\n ref_text=ref_text.strip(),\n gen_text=gen_text.strip(),\n nfe_step=nfe_step,\n speed=speed,\n remove_silence=remove_silence,\n file_wave=f.name,\n seed=seed,\n )\n return f.name, tts_api.device, str(tts_api.seed)\n\n\ndef check_finetune(finetune):\n return gr.update(interactive=finetune), gr.update(interactive=finetune), gr.update(interactive=finetune)\n\n\ndef get_checkpoints_project(project_name, is_gradio=True):\n if project_name is None:\n return [], \"\"\n project_name = project_name.replace(\"_pinyin\", \"\").replace(\"_char\", \"\")\n\n if os.path.isdir(path_project_ckpts):\n files_checkpoints = glob(os.path.join(path_project_ckpts, project_name, \"*.pt\"))\n # Separate pretrained and regular checkpoints\n pretrained_checkpoints = [f for f in files_checkpoints if \"pretrained_\" in os.path.basename(f)]\n regular_checkpoints = [\n f\n for f in files_checkpoints\n if \"pretrained_\" not in os.path.basename(f) and \"model_last.pt\" not in os.path.basename(f)\n ]\n last_checkpoint = [f for f in files_checkpoints if \"model_last.pt\" in os.path.basename(f)]\n\n # Sort regular checkpoints by number\n regular_checkpoints = sorted(\n regular_checkpoints, key=lambda x: int(os.path.basename(x).split(\"_\")[1].split(\".\")[0])\n )\n\n # Combine in order: pretrained, regular, last\n files_checkpoints = pretrained_checkpoints + regular_checkpoints + last_checkpoint\n else:\n files_checkpoints = []\n\n selelect_checkpoint = None if not files_checkpoints else files_checkpoints[0]\n\n if is_gradio:\n return gr.update(choices=files_checkpoints, value=selelect_checkpoint)\n\n return files_checkpoints, selelect_checkpoint\n\n\ndef get_audio_project(project_name, is_gradio=True):\n if project_name is None:\n return [], \"\"\n project_name = project_name.replace(\"_pinyin\", \"\").replace(\"_char\", \"\")\n\n if os.path.isdir(path_project_ckpts):\n files_audios = glob(os.path.join(path_project_ckpts, project_name, \"samples\", \"*.wav\"))\n files_audios = sorted(files_audios, key=lambda x: int(os.path.basename(x).split(\"_\")[1].split(\".\")[0]))\n\n files_audios = [item.replace(\"_gen.wav\", \"\") for item in files_audios if item.endswith(\"_gen.wav\")]\n else:\n files_audios = []\n\n selelect_checkpoint = None if not files_audios else files_audios[0]\n\n if is_gradio:\n return gr.update(choices=files_audios, value=selelect_checkpoint)\n\n return files_audios, selelect_checkpoint\n\n\ndef get_gpu_stats():\n gpu_stats = \"\"\n\n if torch.cuda.is_available():\n gpu_count = torch.cuda.device_count()\n for i in range(gpu_count):\n gpu_name = torch.cuda.get_device_name(i)\n gpu_properties = torch.cuda.get_device_properties(i)\n total_memory = gpu_properties.total_memory / (1024**3) # in GB\n allocated_memory = torch.cuda.memory_allocated(i) / (1024**2) # in MB\n reserved_memory = torch.cuda.memory_reserved(i) / (1024**2) # in MB\n\n gpu_stats += (\n f\"GPU {i} Name: {gpu_name}\\n\"\n f\"Total GPU memory (GPU {i}): {total_memory:.2f} GB\\n\"\n f\"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\\n\"\n f\"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\\n\\n\"\n )\n elif torch.xpu.is_available():\n gpu_count = torch.xpu.device_count()\n for i in range(gpu_count):\n gpu_name = torch.xpu.get_device_name(i)\n gpu_properties = torch.xpu.get_device_properties(i)\n total_memory = gpu_properties.total_memory / (1024**3) # in GB\n allocated_memory = torch.xpu.memory_allocated(i) / (1024**2) # in MB\n reserved_memory = torch.xpu.memory_reserved(i) / (1024**2) # in MB\n\n gpu_stats += (\n f\"GPU {i} Name: {gpu_name}\\n\"\n f\"Total GPU memory (GPU {i}): {total_memory:.2f} GB\\n\"\n f\"Allocated GPU memory (GPU {i}): {allocated_memory:.2f} MB\\n\"\n f\"Reserved GPU memory (GPU {i}): {reserved_memory:.2f} MB\\n\\n\"\n )\n elif torch.backends.mps.is_available():\n gpu_count = 1\n gpu_stats += \"MPS GPU\\n\"\n total_memory = psutil.virtual_memory().total / (\n 1024**3\n ) # Total system memory (MPS doesn't have its own memory)\n allocated_memory = 0\n reserved_memory = 0\n\n gpu_stats += (\n f\"Total system memory: {total_memory:.2f} GB\\n\"\n f\"Allocated GPU memory (MPS): {allocated_memory:.2f} MB\\n\"\n f\"Reserved GPU memory (MPS): {reserved_memory:.2f} MB\\n\"\n )\n\n else:\n gpu_stats = \"No GPU available\"\n\n return gpu_stats\n\n\ndef get_cpu_stats():\n cpu_usage = psutil.cpu_percent(interval=1)\n memory_info = psutil.virtual_memory()\n memory_used = memory_info.used / (1024**2)\n memory_total = memory_info.total / (1024**2)\n memory_percent = memory_info.percent\n\n pid = os.getpid()\n process = psutil.Process(pid)\n nice_value = process.nice()\n\n cpu_stats = (\n f\"CPU Usage: {cpu_usage:.2f}%\\n\"\n f\"System Memory: {memory_used:.2f} MB used / {memory_total:.2f} MB total ({memory_percent}% used)\\n\"\n f\"Process Priority (Nice value): {nice_value}\"\n )\n\n return cpu_stats\n\n\ndef get_combined_stats():\n gpu_stats = get_gpu_stats()\n cpu_stats = get_cpu_stats()\n combined_stats = f\"### GPU Stats\\n{gpu_stats}\\n\\n### CPU Stats\\n{cpu_stats}\"\n return combined_stats\n\n\ndef get_audio_select(file_sample):\n select_audio_ref = file_sample\n select_audio_gen = file_sample\n\n if file_sample is not None:\n select_audio_ref += \"_ref.wav\"\n select_audio_gen += \"_gen.wav\"\n\n return select_audio_ref, select_audio_gen\n\n\nwith gr.Blocks() as app:\n gr.Markdown(\n \"\"\"\n# F5 TTS Automatic Finetune\n\nThis is a local web UI for F5 TTS finetuning support. This app supports the following TTS models:\n\n* [F5-TTS](https://arxiv.org/abs/2410.06885) (A Fairytaler that Fakes Fluent and Faithful Speech with Flow Matching)\n* [E2 TTS](https://arxiv.org/abs/2406.18009) (Embarrassingly Easy Fully Non-Autoregressive Zero-Shot TTS)\n\nThe pretrained checkpoints support English and Chinese.\n\nFor tutorial and updates check here (https://github.com/SWivid/F5-TTS/discussions/143)\n\"\"\"\n )\n\n with gr.Row():\n projects, projects_selelect = get_list_projects()\n tokenizer_type = gr.Radio(label=\"Tokenizer Type\", choices=[\"pinyin\", \"char\", \"custom\"], value=\"pinyin\")\n project_name = gr.Textbox(label=\"Project Name\", value=\"my_speak\")\n bt_create = gr.Button(\"Create a New Project\")\n\n with gr.Row():\n cm_project = gr.Dropdown(\n choices=projects, value=projects_selelect, label=\"Project\", allow_custom_value=True, scale=6\n )\n ch_refresh_project = gr.Button(\"Refresh\", scale=1)\n\n bt_create.click(fn=create_data_project, inputs=[project_name, tokenizer_type], outputs=[cm_project])\n\n with gr.Tabs():\n with gr.TabItem(\"Transcribe Data\"):\n gr.Markdown(\"\"\"```plaintext \nSkip this step if you have your dataset, metadata.csv, and a folder wavs with all the audio files. \n```\"\"\")\n\n ch_manual = gr.Checkbox(label=\"Audio from Path\", value=False)\n\n mark_info_transcribe = gr.Markdown(\n \"\"\"```plaintext \n Place your 'wavs' folder and 'metadata.csv' file in the '{your_project_name}' directory. \n \n my_speak/\n │\n └── dataset/\n ├── audio1.wav\n └── audio2.wav\n ...\n ```\"\"\",\n visible=False,\n )\n\n audio_speaker = gr.File(label=\"Voice\", type=\"filepath\", file_count=\"multiple\")\n txt_lang = gr.Textbox(label=\"Language\", value=\"English\")\n bt_transcribe = bt_create = gr.Button(\"Transcribe\")\n txt_info_transcribe = gr.Textbox(label=\"Info\", value=\"\")\n bt_transcribe.click(\n fn=transcribe_all,\n inputs=[cm_project, audio_speaker, txt_lang, ch_manual],\n outputs=[txt_info_transcribe],\n )\n ch_manual.change(fn=check_user, inputs=[ch_manual], outputs=[audio_speaker, mark_info_transcribe])\n\n random_sample_transcribe = gr.Button(\"Random Sample\")\n\n with gr.Row():\n random_text_transcribe = gr.Textbox(label=\"Text\")\n random_audio_transcribe = gr.Audio(label=\"Audio\", type=\"filepath\")\n\n random_sample_transcribe.click(\n fn=get_random_sample_transcribe,\n inputs=[cm_project],\n outputs=[random_text_transcribe, random_audio_transcribe],\n )\n\n with gr.TabItem(\"Vocab Check\"):\n gr.Markdown(\"\"\"```plaintext \nCheck the vocabulary for fine-tuning Emilia_ZH_EN to ensure all symbols are included. For fine-tuning a new language.\n```\"\"\")\n\n check_button = gr.Button(\"Check Vocab\")\n txt_info_check = gr.Textbox(label=\"Info\", value=\"\")\n\n gr.Markdown(\"\"\"```plaintext \nUsing the extended model, you can finetune to a new language that is missing symbols in the vocab. This creates a new model with a new vocabulary size and saves it in your ckpts/project folder.\n```\"\"\")\n\n exp_name_extend = gr.Radio(\n label=\"Model\", choices=[\"F5TTS_v1_Base\", \"F5TTS_Base\", \"E2TTS_Base\"], value=\"F5TTS_v1_Base\"\n )\n\n with gr.Row():\n txt_extend = gr.Textbox(\n label=\"Symbols\",\n value=\"\",\n placeholder=\"To add new symbols, make sure to use ',' for each symbol\",\n scale=6,\n )\n txt_count_symbol = gr.Textbox(label=\"New Vocab Size\", value=\"\", scale=1)\n\n extend_button = gr.Button(\"Extend\")\n txt_info_extend = gr.Textbox(label=\"Info\", value=\"\")\n\n txt_extend.change(vocab_count, inputs=[txt_extend], outputs=[txt_count_symbol])\n check_button.click(\n fn=vocab_check, inputs=[cm_project, tokenizer_type], outputs=[txt_info_check, txt_extend]\n )\n extend_button.click(\n fn=vocab_extend, inputs=[cm_project, txt_extend, exp_name_extend], outputs=[txt_info_extend]\n )\n\n with gr.TabItem(\"Prepare Data\"):\n gr.Markdown(\"\"\"```plaintext \nSkip this step if you have your dataset, raw.arrow, duration.json, and vocab.txt\n```\"\"\")\n\n gr.Markdown(\n \"\"\"```plaintext \n Place all your \"wavs\" folder and your \"metadata.csv\" file in your project name directory.\n\n Supported audio formats: \"wav\", \"mp3\", \"aac\", \"flac\", \"m4a\", \"alac\", \"ogg\", \"aiff\", \"wma\", \"amr\"\n\n Example wav format: \n my_speak/\n │\n ├── wavs/\n │ ├── audio1.wav\n │ └── audio2.wav\n | ...\n │\n └── metadata.csv\n \n File format metadata.csv:\n\n audio1|text1 or audio1.wav|text1 or your_path/audio1.wav|text1 \n audio2|text1 or audio2.wav|text1 or your_path/audio2.wav|text1 \n ...\n\n ```\"\"\"\n )\n ch_tokenizern = gr.Checkbox(label=\"Create Vocabulary\", value=False, visible=False)\n\n bt_prepare = bt_create = gr.Button(\"Prepare\")\n txt_info_prepare = gr.Textbox(label=\"Info\", value=\"\")\n txt_vocab_prepare = gr.Textbox(label=\"Vocab\", value=\"\")\n\n bt_prepare.click(\n fn=create_metadata, inputs=[cm_project, ch_tokenizern], outputs=[txt_info_prepare, txt_vocab_prepare]\n )\n\n random_sample_prepare = gr.Button(\"Random Sample\")\n\n with gr.Row():\n random_text_prepare = gr.Textbox(label=\"Tokenizer\")\n random_audio_prepare = gr.Audio(label=\"Audio\", type=\"filepath\")\n\n random_sample_prepare.click(\n fn=get_random_sample_prepare, inputs=[cm_project], outputs=[random_text_prepare, random_audio_prepare]\n )\n\n with gr.TabItem(\"Train Model\"):\n gr.Markdown(\"\"\"```plaintext \nThe auto-setting is still experimental. Set a large value of epoch if not sure; and keep last N checkpoints if limited disk space.\nIf you encounter a memory error, try reducing the batch size per GPU to a smaller number.\n```\"\"\")\n with gr.Row():\n exp_name = gr.Radio(label=\"Model\", choices=[\"F5TTS_v1_Base\", \"F5TTS_Base\", \"E2TTS_Base\"])\n tokenizer_file = gr.Textbox(label=\"Tokenizer File\")\n file_checkpoint_train = gr.Textbox(label=\"Path to the Pretrained Checkpoint\")\n\n with gr.Row():\n ch_finetune = bt_create = gr.Checkbox(label=\"Finetune\")\n lb_samples = gr.Label(label=\"Samples\")\n bt_calculate = bt_create = gr.Button(\"Auto Settings\")\n\n with gr.Row():\n epochs = gr.Number(label=\"Epochs\")\n learning_rate = gr.Number(label=\"Learning Rate\", step=0.5e-5)\n max_grad_norm = gr.Number(label=\"Max Gradient Norm\")\n num_warmup_updates = gr.Number(label=\"Warmup Updates\")\n\n with gr.Row():\n batch_size_type = gr.Radio(\n label=\"Batch Size Type\",\n choices=[\"frame\", \"sample\"],\n info=\"frame is calculated as seconds * sampling_rate / hop_length\",\n )\n batch_size_per_gpu = gr.Number(label=\"Batch Size per GPU\", info=\"N frames or N samples\")\n grad_accumulation_steps = gr.Number(\n label=\"Gradient Accumulation Steps\", info=\"Effective batch size is multiplied by this value\"\n )\n max_samples = gr.Number(label=\"Max Samples\", info=\"Maximum number of samples per single GPU batch\")\n\n with gr.Row():\n save_per_updates = gr.Number(\n label=\"Save per Updates\",\n info=\"Save intermediate checkpoints every N updates\",\n minimum=10,\n )\n keep_last_n_checkpoints = gr.Number(\n label=\"Keep Last N Checkpoints\",\n step=1,\n precision=0,\n info=\"-1 to keep all, 0 to not save intermediate, > 0 to keep last N\",\n minimum=-1,\n )\n last_per_updates = gr.Number(\n label=\"Last per Updates\",\n info=\"Save latest checkpoint with suffix _last.pt every N updates\",\n minimum=10,\n )\n gr.Radio(label=\"\") # placeholder\n\n with gr.Row():\n ch_8bit_adam = gr.Checkbox(label=\"Use 8-bit Adam optimizer\")\n mixed_precision = gr.Radio(label=\"Mixed Precision\", choices=[\"none\", \"fp16\", \"bf16\"])\n cd_logger = gr.Radio(label=\"Logger\", choices=[\"none\", \"wandb\", \"tensorboard\"])\n with gr.Column():\n start_button = gr.Button(\"Start Training\")\n stop_button = gr.Button(\"Stop Training\", interactive=False)\n\n if projects_selelect is not None:\n (\n exp_name_value,\n learning_rate_value,\n batch_size_per_gpu_value,\n batch_size_type_value,\n max_samples_value,\n grad_accumulation_steps_value,\n max_grad_norm_value,\n epochs_value,\n num_warmup_updates_value,\n save_per_updates_value,\n keep_last_n_checkpoints_value,\n last_per_updates_value,\n finetune_value,\n file_checkpoint_train_value,\n tokenizer_type_value,\n tokenizer_file_value,\n mixed_precision_value,\n logger_value,\n bnb_optimizer_value,\n ) = load_settings(projects_selelect)\n\n # Assigning values to the respective components\n exp_name.value = exp_name_value\n learning_rate.value = learning_rate_value\n batch_size_per_gpu.value = batch_size_per_gpu_value\n batch_size_type.value = batch_size_type_value\n max_samples.value = max_samples_value\n grad_accumulation_steps.value = grad_accumulation_steps_value\n max_grad_norm.value = max_grad_norm_value\n epochs.value = epochs_value\n num_warmup_updates.value = num_warmup_updates_value\n save_per_updates.value = save_per_updates_value\n keep_last_n_checkpoints.value = keep_last_n_checkpoints_value\n last_per_updates.value = last_per_updates_value\n ch_finetune.value = finetune_value\n file_checkpoint_train.value = file_checkpoint_train_value\n tokenizer_type.value = tokenizer_type_value\n tokenizer_file.value = tokenizer_file_value\n mixed_precision.value = mixed_precision_value\n cd_logger.value = logger_value\n ch_8bit_adam.value = bnb_optimizer_value\n\n ch_stream = gr.Checkbox(label=\"Stream Output Experiment\", value=True)\n txt_info_train = gr.Textbox(label=\"Info\", value=\"\")\n\n list_audios, select_audio = get_audio_project(projects_selelect, False)\n\n select_audio_ref = select_audio\n select_audio_gen = select_audio\n\n if select_audio is not None:\n select_audio_ref += \"_ref.wav\"\n select_audio_gen += \"_gen.wav\"\n\n with gr.Row():\n ch_list_audio = gr.Dropdown(\n choices=list_audios,\n value=select_audio,\n label=\"Audios\",\n allow_custom_value=True,\n scale=6,\n interactive=True,\n )\n bt_stream_audio = gr.Button(\"Refresh\", scale=1)\n bt_stream_audio.click(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio])\n cm_project.change(fn=get_audio_project, inputs=[cm_project], outputs=[ch_list_audio])\n\n with gr.Row():\n audio_ref_stream = gr.Audio(label=\"Original\", type=\"filepath\", value=select_audio_ref)\n audio_gen_stream = gr.Audio(label=\"Generate\", type=\"filepath\", value=select_audio_gen)\n\n ch_list_audio.change(\n fn=get_audio_select,\n inputs=[ch_list_audio],\n outputs=[audio_ref_stream, audio_gen_stream],\n )\n\n start_button.click(\n fn=start_training,\n inputs=[\n cm_project,\n exp_name,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n grad_accumulation_steps,\n max_grad_norm,\n epochs,\n num_warmup_updates,\n save_per_updates,\n keep_last_n_checkpoints,\n last_per_updates,\n ch_finetune,\n file_checkpoint_train,\n tokenizer_type,\n tokenizer_file,\n mixed_precision,\n ch_stream,\n cd_logger,\n ch_8bit_adam,\n ],\n outputs=[txt_info_train, start_button, stop_button],\n )\n stop_button.click(fn=stop_training, outputs=[txt_info_train, start_button, stop_button])\n\n bt_calculate.click(\n fn=calculate_train,\n inputs=[\n cm_project,\n epochs,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n num_warmup_updates,\n ch_finetune,\n ],\n outputs=[\n epochs,\n learning_rate,\n batch_size_per_gpu,\n max_samples,\n num_warmup_updates,\n lb_samples,\n ],\n )\n\n ch_finetune.change(\n check_finetune, inputs=[ch_finetune], outputs=[file_checkpoint_train, tokenizer_file, tokenizer_type]\n )\n\n def setup_load_settings():\n output_components = [\n exp_name,\n learning_rate,\n batch_size_per_gpu,\n batch_size_type,\n max_samples,\n grad_accumulation_steps,\n max_grad_norm,\n epochs,\n num_warmup_updates,\n save_per_updates,\n keep_last_n_checkpoints,\n last_per_updates,\n ch_finetune,\n file_checkpoint_train,\n tokenizer_type,\n tokenizer_file,\n mixed_precision,\n cd_logger,\n ch_8bit_adam,\n ]\n return output_components\n\n outputs = setup_load_settings()\n\n cm_project.change(\n fn=load_settings,\n inputs=[cm_project],\n outputs=outputs,\n )\n\n ch_refresh_project.click(\n fn=load_settings,\n inputs=[cm_project],\n outputs=outputs,\n )\n\n with gr.TabItem(\"Test Model\"):\n gr.Markdown(\"\"\"```plaintext \nCheck the use_ema setting (True or False) for your model to see what works best for you. Set seed to -1 for random.\n```\"\"\")\n exp_name = gr.Radio(\n label=\"Model\", choices=[\"F5TTS_v1_Base\", \"F5TTS_Base\", \"E2TTS_Base\"], value=\"F5TTS_v1_Base\"\n )\n list_checkpoints, checkpoint_select = get_checkpoints_project(projects_selelect, False)\n\n with gr.Row():\n nfe_step = gr.Number(label=\"NFE Step\", value=32)\n speed = gr.Slider(label=\"Speed\", value=1.0, minimum=0.3, maximum=2.0, step=0.1)\n seed = gr.Number(label=\"Random Seed\", value=-1, minimum=-1)\n remove_silence = gr.Checkbox(label=\"Remove Silence\")\n\n with gr.Row():\n ch_use_ema = gr.Checkbox(\n label=\"Use EMA\", value=True, info=\"Turn off at early stage might offer better results\"\n )\n cm_checkpoint = gr.Dropdown(\n choices=list_checkpoints, value=checkpoint_select, label=\"Checkpoints\", allow_custom_value=True\n )\n bt_checkpoint_refresh = gr.Button(\"Refresh\")\n\n random_sample_infer = gr.Button(\"Random Sample\")\n\n ref_text = gr.Textbox(label=\"Reference Text\")\n ref_audio = gr.Audio(label=\"Reference Audio\", type=\"filepath\")\n gen_text = gr.Textbox(label=\"Text to Generate\")\n\n random_sample_infer.click(\n fn=get_random_sample_infer, inputs=[cm_project], outputs=[ref_text, gen_text, ref_audio]\n )\n\n with gr.Row():\n txt_info_gpu = gr.Textbox(\"\", label=\"Inference on Device :\")\n seed_info = gr.Textbox(label=\"Used Random Seed :\")\n check_button_infer = gr.Button(\"Inference\")\n\n gen_audio = gr.Audio(label=\"Generated Audio\", type=\"filepath\")\n\n check_button_infer.click(\n fn=infer,\n inputs=[\n cm_project,\n cm_checkpoint,\n exp_name,\n ref_text,\n ref_audio,\n gen_text,\n nfe_step,\n ch_use_ema,\n speed,\n seed,\n remove_silence,\n ],\n outputs=[gen_audio, txt_info_gpu, seed_info],\n )\n\n bt_checkpoint_refresh.click(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint])\n cm_project.change(fn=get_checkpoints_project, inputs=[cm_project], outputs=[cm_checkpoint])\n\n with gr.TabItem(\"Prune Checkpoint\"):\n gr.Markdown(\"\"\"```plaintext \nReduce the Base model size from 5GB to 1.3GB. The new checkpoint file prunes out optimizer and etc., can be used for inference or finetuning afterward, but not able to resume pretraining.\n```\"\"\")\n txt_path_checkpoint = gr.Textbox(label=\"Path to Checkpoint:\")\n txt_path_checkpoint_small = gr.Textbox(label=\"Path to Output:\")\n with gr.Row():\n ch_save_ema = gr.Checkbox(label=\"Save EMA checkpoint\", value=True)\n ch_safetensors = gr.Checkbox(label=\"Save with safetensors format\", value=True)\n txt_info_reduse = gr.Textbox(label=\"Info\", value=\"\")\n reduse_button = gr.Button(\"Prune\")\n reduse_button.click(\n fn=prune_checkpoint,\n inputs=[txt_path_checkpoint, txt_path_checkpoint_small, ch_save_ema, ch_safetensors],\n outputs=[txt_info_reduse],\n )\n\n with gr.TabItem(\"System Info\"):\n output_box = gr.Textbox(label=\"GPU and CPU Information\", lines=20)\n\n def update_stats():\n return get_combined_stats()\n\n update_button = gr.Button(\"Update Stats\")\n update_button.click(fn=update_stats, outputs=output_box)\n\n def auto_update():\n yield gr.update(value=update_stats())\n\n gr.update(fn=auto_update, inputs=[], outputs=output_box)\n\n\n@click.command()\n@click.option(\"--port\", \"-p\", default=None, type=int, help=\"Port to run the app on\")\n@click.option(\"--host\", \"-H\", default=None, help=\"Host to run the app on\")\n@click.option(\n \"--share\",\n \"-s\",\n default=False,\n is_flag=True,\n help=\"Share the app via Gradio share link\",\n)\n@click.option(\"--api\", \"-a\", default=True, is_flag=True, help=\"Allow API access\")\ndef main(port, host, share, api):\n global app\n print(\"Starting app...\")\n app.queue(api_open=api).launch(server_name=host, server_port=port, share=share, show_api=api)\n\n\nif __name__ == \"__main__\":\n main()\n" }, { "path": "src/f5_tts/train/train.py", "content": "# training script.\n\nimport os\nfrom importlib.resources import files\n\nimport hydra\nfrom omegaconf import OmegaConf\n\nfrom f5_tts.model import CFM, Trainer\nfrom f5_tts.model.dataset import load_dataset\nfrom f5_tts.model.utils import get_tokenizer\n\n\nos.chdir(str(files(\"f5_tts\").joinpath(\"../..\"))) # change working directory to root of project (local editable)\n\n\n@hydra.main(version_base=\"1.3\", config_path=str(files(\"f5_tts\").joinpath(\"configs\")), config_name=None)\ndef main(model_cfg):\n model_cls = hydra.utils.get_class(f\"f5_tts.model.{model_cfg.model.backbone}\")\n model_arc = model_cfg.model.arch\n tokenizer = model_cfg.model.tokenizer\n mel_spec_type = model_cfg.model.mel_spec.mel_spec_type\n\n wandb_project = model_cfg.ckpts.get(\"wandb_project\", \"CFM-TTS\")\n wandb_run_name = model_cfg.ckpts.get(\n \"wandb_run_name\",\n f\"{model_cfg.model.name}_{mel_spec_type}_{model_cfg.model.tokenizer}_{model_cfg.datasets.name}\",\n )\n wandb_resume_id = model_cfg.ckpts.get(\"wandb_resume_id\", None)\n\n # set text tokenizer\n if tokenizer != \"custom\":\n tokenizer_path = model_cfg.datasets.name\n else:\n tokenizer_path = model_cfg.model.tokenizer_path\n vocab_char_map, vocab_size = get_tokenizer(tokenizer_path, tokenizer)\n\n # set model\n model = CFM(\n transformer=model_cls(**model_arc, text_num_embeds=vocab_size, mel_dim=model_cfg.model.mel_spec.n_mel_channels),\n mel_spec_kwargs=model_cfg.model.mel_spec,\n vocab_char_map=vocab_char_map,\n )\n\n # init trainer\n trainer = Trainer(\n model,\n epochs=model_cfg.optim.epochs,\n learning_rate=model_cfg.optim.learning_rate,\n num_warmup_updates=model_cfg.optim.num_warmup_updates,\n save_per_updates=model_cfg.ckpts.save_per_updates,\n keep_last_n_checkpoints=model_cfg.ckpts.keep_last_n_checkpoints,\n checkpoint_path=str(files(\"f5_tts\").joinpath(f\"../../{model_cfg.ckpts.save_dir}\")),\n batch_size_per_gpu=model_cfg.datasets.batch_size_per_gpu,\n batch_size_type=model_cfg.datasets.batch_size_type,\n max_samples=model_cfg.datasets.max_samples,\n grad_accumulation_steps=model_cfg.optim.grad_accumulation_steps,\n max_grad_norm=model_cfg.optim.max_grad_norm,\n logger=model_cfg.ckpts.logger,\n wandb_project=wandb_project,\n wandb_run_name=wandb_run_name,\n wandb_resume_id=wandb_resume_id,\n last_per_updates=model_cfg.ckpts.last_per_updates,\n log_samples=model_cfg.ckpts.log_samples,\n bnb_optimizer=model_cfg.optim.bnb_optimizer,\n mel_spec_type=mel_spec_type,\n is_local_vocoder=model_cfg.model.vocoder.is_local,\n local_vocoder_path=model_cfg.model.vocoder.local_path,\n model_cfg_dict=OmegaConf.to_container(model_cfg, resolve=True),\n )\n\n train_dataset = load_dataset(model_cfg.datasets.name, tokenizer, mel_spec_kwargs=model_cfg.model.mel_spec)\n trainer.train(\n train_dataset,\n num_workers=model_cfg.datasets.num_workers,\n resumable_with_seed=666, # seed for shuffling dataset\n )\n\n\nif __name__ == \"__main__\":\n main()\n" } ]