[ { "path": ".github/workflows/publish_action.yml", "content": "name: Publish to Comfy registry\non:\n workflow_dispatch:\n push:\n branches:\n - master\n - main\n paths:\n - \"pyproject.toml\"\n\njobs:\n publish-node:\n name: Publish Custom Node to registry\n runs-on: ubuntu-latest\n steps:\n - name: Check out code\n uses: actions/checkout@v4\n - name: Publish Custom Node\n uses: Comfy-Org/publish-node-action@main\n with:\n ## Add your own personal access token to your Github Repository secrets and reference it here.\n personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}\n" }, { "path": ".gitignore", "content": "# 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# UV\n# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n#uv.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\n# Ruff stuff:\n.ruff_cache/\n\n# PyPI configuration file\n.pypirc\n" }, { "path": "LICENSE", "content": "\n Apache License\n Version 2.0, January 2004\n http://www.apache.org/licenses/\n\n TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n 1. Definitions.\n\n \"License\" shall mean the terms and conditions for use, reproduction,\n and distribution as defined by Sections 1 through 9 of this document.\n\n \"Licensor\" shall mean the copyright owner or entity authorized by\n the copyright owner that is granting the License.\n\n \"Legal Entity\" shall mean the union of the acting entity and all\n other entities that control, are controlled by, or are under common\n control with that entity. For the purposes of this definition,\n \"control\" means (i) the power, direct or indirect, to cause the\n direction or management of such entity, whether by contract or\n otherwise, or (ii) ownership of fifty percent (50%) or more of the\n outstanding shares, or (iii) beneficial ownership of such entity.\n\n \"You\" (or \"Your\") shall mean an individual or Legal Entity\n exercising permissions granted by this License.\n\n \"Source\" form shall mean the preferred form for making modifications,\n including but not limited to software source code, documentation\n source, and configuration files.\n\n \"Object\" form shall mean any form resulting from mechanical\n transformation or translation of a Source form, including but\n not limited to compiled object code, generated documentation,\n and conversions to other media types.\n\n \"Work\" shall mean the work of authorship, whether in Source or\n Object form, made available under the License, as indicated by a\n copyright notice that is included in or attached to the work\n (an example is provided in the Appendix below).\n\n \"Derivative Works\" shall mean any work, whether in Source or Object\n form, that is based on (or derived from) the Work and for which the\n editorial revisions, annotations, elaborations, or other modifications\n represent, as a whole, an original work of authorship. For the purposes\n of this License, Derivative Works shall not include works that remain\n separable from, or merely link (or bind by name) to the interfaces of,\n the Work and Derivative Works thereof.\n\n \"Contribution\" shall mean any work of authorship, including\n the original version of the Work and any modifications or additions\n to that Work or Derivative Works thereof, that is intentionally\n submitted to Licensor for inclusion in the Work by the copyright owner\n or by an individual or Legal Entity authorized to submit on behalf of\n the copyright owner. For the purposes of this definition, \"submitted\"\n means any form of electronic, verbal, or written communication sent\n to the Licensor or its representatives, including but not limited to\n communication on electronic mailing lists, source code control systems,\n and issue tracking systems that are managed by, or on behalf of, the\n Licensor for the purpose of discussing and improving the Work, but\n excluding communication that is conspicuously marked or otherwise\n designated in writing by the copyright owner as \"Not a Contribution.\"\n\n \"Contributor\" shall mean Licensor and any individual or Legal Entity\n on behalf of whom a Contribution has been received by Licensor and\n subsequently incorporated within the Work.\n\n 2. Grant of Copyright License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n copyright license to reproduce, prepare Derivative Works of,\n publicly display, publicly perform, sublicense, and distribute the\n Work and such Derivative Works in Source or Object form.\n\n 3. Grant of Patent License. Subject to the terms and conditions of\n this License, each Contributor hereby grants to You a perpetual,\n worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n (except as stated in this section) patent license to make, have made,\n use, offer to sell, sell, import, and otherwise transfer the Work,\n where such license applies only to those patent claims licensable\n by such Contributor that are necessarily infringed by their\n Contribution(s) alone or by combination of their Contribution(s)\n with the Work to which such Contribution(s) was submitted. If You\n institute patent litigation against any entity (including a\n cross-claim or counterclaim in a lawsuit) alleging that the Work\n or a Contribution incorporated within the Work constitutes direct\n or contributory patent infringement, then any patent licenses\n granted to You under this License for that Work shall terminate\n as of the date such litigation is filed.\n\n 4. Redistribution. You may reproduce and distribute copies of the\n Work or Derivative Works thereof in any medium, with or without\n modifications, and in Source or Object form, provided that You\n meet the following conditions:\n\n (a) You must give any other recipients of the Work or\n Derivative Works a copy of this License; and\n\n (b) You must cause any modified files to carry prominent notices\n stating that You changed the files; and\n\n (c) You must retain, in the Source form of any Derivative Works\n that You distribute, all copyright, patent, trademark, and\n attribution notices from the Source form of the Work,\n excluding those notices that do not pertain to any part of\n the Derivative Works; and\n\n (d) If the Work includes a \"NOTICE\" text file as part of its\n distribution, then any Derivative Works that You distribute must\n include a readable copy of the attribution notices contained\n within such NOTICE file, excluding those notices that do not\n pertain to any part of the Derivative Works, in at least one\n of the following places: within a NOTICE text file distributed\n as part of the Derivative Works; within the Source form or\n documentation, if provided along with the Derivative Works; or,\n within a display generated by the Derivative Works, if and\n wherever such third-party notices normally appear. The contents\n of the NOTICE file are for informational purposes only and\n do not modify the License. You may add Your own attribution\n notices within Derivative Works that You distribute, alongside\n or as an addendum to the NOTICE text from the Work, provided\n that such additional attribution notices cannot be construed\n as modifying the License.\n\n You may add Your own copyright statement to Your modifications and\n may provide additional or different license terms and conditions\n for use, reproduction, or distribution of Your modifications, or\n for any such Derivative Works as a whole, provided Your use,\n reproduction, and distribution of the Work otherwise complies with\n the conditions stated in this License.\n\n 5. Submission of Contributions. Unless You explicitly state otherwise,\n any Contribution intentionally submitted for inclusion in the Work\n by You to the Licensor shall be under the terms and conditions of\n this License, without any additional terms or conditions.\n Notwithstanding the above, nothing herein shall supersede or modify\n the terms of any separate license agreement you may have executed\n with Licensor regarding such Contributions.\n\n 6. Trademarks. This License does not grant permission to use the trade\n names, trademarks, service marks, or product names of the Licensor,\n except as required for reasonable and customary use in describing the\n origin of the Work and reproducing the content of the NOTICE file.\n\n 7. Disclaimer of Warranty. Unless required by applicable law or\n agreed to in writing, Licensor provides the Work (and each\n Contributor provides its Contributions) on an \"AS IS\" BASIS,\n WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n implied, including, without limitation, any warranties or conditions\n of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n PARTICULAR PURPOSE. You are solely responsible for determining the\n appropriateness of using or redistributing the Work and assume any\n risks associated with Your exercise of permissions under this License.\n\n 8. Limitation of Liability. In no event and under no legal theory,\n whether in tort (including negligence), contract, or otherwise,\n unless required by applicable law (such as deliberate and grossly\n negligent acts) or agreed to in writing, shall any Contributor be\n liable to You for damages, including any direct, indirect, special,\n incidental, or consequential damages of any character arising as a\n result of this License or out of the use or inability to use the\n Work (including but not limited to damages for loss of goodwill,\n work stoppage, computer failure or malfunction, or any and all\n other commercial damages or losses), even if such Contributor\n has been advised of the possibility of such damages.\n\n 9. Accepting Warranty or Additional Liability. While redistributing\n the Work or Derivative Works thereof, You may choose to offer,\n and charge a fee for, acceptance of support, warranty, indemnity,\n or other liability obligations and/or rights consistent with this\n License. However, in accepting such obligations, You may act only\n on Your own behalf and on Your sole responsibility, not on behalf\n of any other Contributor, and only if You agree to indemnify,\n defend, and hold each Contributor harmless for any liability\n incurred by, or claims asserted against, such Contributor by reason\n of your accepting any such warranty or additional liability.\n\n END OF TERMS AND CONDITIONS\n\n APPENDIX: How to apply the Apache License to your work.\n\n To apply the Apache License to your work, attach the following\n boilerplate notice, with the fields enclosed by brackets \"[]\"\n replaced with your own identifying information. (Don't include\n the brackets!) The text should be enclosed in the appropriate\n comment syntax for the file format. We also recommend that a\n file or class name and description of purpose be included on the\n same \"printed page\" as the copyright notice for easier\n identification within third-party archives.\n\n Copyright [2025] ByteDance\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." }, { "path": "README-CN.md", "content": "[中文](README-CN.md) | [English](README.md) \r\n\r\n# ComfyUI 的 MegaTTS3 声音克隆节点\r\n\r\n声音克隆质量非常高, 支持中英文, 并可跨语言克隆. **支持自定义音色!!! 超长文本!!! 双人对话!!! Windows 正常安装 pynini, 不再是阉割版 TTS!!!**.\r\n\r\n## 📣 更新\r\n\r\n[2025-06-07]⚒️: v2.0.0. **支持自定义音色, 支持超长文本, 支持双人对话, Windows 正常安装 pynini, 不再是阉割版 TTS!**.\r\n\r\n```\r\n[S1] MegaTTS 真开源版本来了，效果666\r\n[S2] 晕 xuan4 是一种 gan3 觉\r\n[S1] 我爱你！I love you!“我爱你”的英语是“I love you”\r\n[S2] 2.5平方电线,共465篇，约315万字\r\n[S1] 2002年的第一场雪，下在了2003年\r\n```\r\n\r\nhttps://github.com/user-attachments/assets/b734e6bd-9303-4311-b3a4-618241ca6535\r\n\r\n[2025-04-28]⚒️: 新增预览音色节点, 先预览音色, 满意再进行克隆. 感谢 @chenpipi0807 的 idea😍. 可在 `speakers` 文件夹下分门别类建更多文件夹.\r\n\r\n[2025-04-06]⚒️: 发布 v1.0.0.\r\n\r\n## 使用\r\n\r\n- 单人克隆(超长文本用空行隔开):\r\n\r\n![image](https://github.com/billwuhao/ComfyUI_MegaTTS3/blob/main/images/2025-04-06_13-52-57.png)\r\n\r\n- 双人对话:\r\n\r\n![image](https://github.com/billwuhao/ComfyUI_MegaTTS3/blob/main/images/2025-04-06_14-49-12.png)\r\n\r\n## 安装\r\n\r\n- **Windows 先安装以下依赖**:\r\n\r\n[pynini-windows-wheels](https://github.com/billwuhao/pynini-windows-wheels/releases/tag/v2.1.6.post1) 下载相应 python 版本的 pynini 轮子.\r\n\r\n示例:\r\n```\r\nD:\\AIGC\\python\\py310\\python.exe -m pip install pynini-2.1.6.post1-cp3xx-cp3xx-win_amd64.whl\r\nD:\\AIGC\\python\\py310\\python.exe -m pip install importlib_resources\r\nD:\\AIGC\\python\\py310\\python.exe -m pip install WeTextProcessing>=1.0.4 --no-deps\r\n```\r\n\r\n- **然后正常进行下列安装**:\r\n```\r\ncd ComfyUI/custom_nodes\r\ngit clone https://github.com/billwuhao/ComfyUI_MegaTTS3.git\r\ncd ComfyUI_MegaTTS3\r\npip install -r requirements.txt\r\n\r\n# python_embeded\r\n./python_embeded/python.exe -m pip install -r requirements.txt\r\n```\r\n\r\n## 模型下载\r\n\r\n- 模型和音色需要手动下载放到 `ComfyUI\\models\\TTS` 路径下:\r\n\r\n[MegaTTS3](https://huggingface.co/ByteDance/MegaTTS3/tree/main) 整个文件夹全部下载放到 `TTS` 文件夹下.\r\n\r\n- **VAE 编码模型, 加微信公众号获取, 放到 `TTS\\MegaTTS3\\wavvae` 文件夹下, 即可自定义音色而无需 `.npy` 文件**:\r\n\r\n\"\"\r\n\r\n - [Google 云盘](https://drive.google.com/drive/folders/1p9GNdNJqeK_94lIJW8lew_G3EazU-9Wx?usp=sharing)\r\n\r\n- 请将音频放到 `TTS\\speakers` 目录下. 我将会把所有 TTS 节点的说话者音频全部统一放到 `ComfyUI\\models\\TTS\\speakers` 路径下, 这些节点包括 `IndexTTS, CSM, Dia, KokoroTTS, MegaTTS, QuteTTS, SparkTTS, StepAudioTTS` 等.\r\n\r\n结构如下:\r\n\r\n```\r\n.\r\n│ .gitattributes\r\n│ config.json\r\n│ README.md\r\n│\r\n├─aligner_lm\r\n│ config.yaml\r\n│ model_only_last.ckpt\r\n│\r\n├─diffusion_transformer\r\n│ config.yaml\r\n│ model_only_last.ckpt\r\n│\r\n├─duration_lm\r\n│ config.yaml\r\n│ model_only_last.ckpt\r\n│\r\n├─g2p\r\n│ added_tokens.json\r\n│ config.json\r\n│ generation_config.json\r\n│ latest\r\n│ merges.txt\r\n│ model.safetensors\r\n│ special_tokens_map.json\r\n│ tokenizer.json\r\n│ tokenizer_config.json\r\n│ trainer_state.json\r\n│ vocab.json\r\n│\r\n└─wavvae\r\n config.yaml\r\n decoder.ckpt\r\n model_only_last.ckpt\r\n```\r\n\r\n## 鸣谢\r\n\r\n- [MegaTTS3](https://github.com/bytedance/MegaTTS3)\r\n\r\n## 打赏\r\n\r\n您的赞赏是我最大的动力! 感谢您支持我一杯咖啡!\r\n\r\n\"\"\r\n" }, { "path": "README.md", "content": "[中文](README-CN.md) | [English](README.md)\n\n# MegaTTS3 Voice Cloning Nodes for ComfyUI\n\nHigh-quality voice cloning, supporting both Chinese and English, with cross-lingual cloning capabilities. **Supports custom voice cloning!!! Extra-long text!!! Two-person dialogue!!! Full pynini installation on Windows, no more stripped-down TTS!!!**.\n\n## 📣 Updates\n\n[2025-06-07]⚒️: v2.0.0. **Supports custom voice cloning, extra-long text, two-person dialogue, and full pynini installation on Windows, no more stripped-down TTS!**.\n\n```\n[S1] MegaTTS 真开源版本来了，效果666\n[S2] 晕 xuan4 是一种 gan3 觉\n[S1] 我爱你！I love you!“我爱你”的英语是“I love you”\n[S2] 2.5平方电线,共465篇，约315万字\n[S1] 2002年的第一场雪，下在了2003年\n```\n\nhttps://github.com/user-attachments/assets/b734e6bd-9303-4311-b3a4-618241ca6535\n\n[2025-04-28]⚒️: Added a voice preview node. Preview the voice first, then clone if you're satisfied. Thanks to @chenpipi0807 for the idea😍. You can create categorized subfolders within the `speakers` folder.\n\n[2025-04-06]⚒️: Released v1.0.0.\n\n## Usage\n\n- Single-person cloning (separate long text with blank lines):\n\n![image](https://github.com/billwuhao/ComfyUI_MegaTTS3/blob/main/images/2025-04-06_13-52-57.png)\n\n- Two-person dialogue:\n\n![image](https://github.com/billwuhao/ComfyUI_MegaTTS3/blob/main/images/2025-04-06_14-49-12.png)\n\n## Installation\n\n- **For Windows, install the following dependencies first**:\n\n[pynini-windows-wheels](https://github.com/billwuhao/pynini-windows-wheels/releases/tag/v2.1.6.post1) Download the pynini wheel file corresponding to your Python version.\n\nExample:\n```\nD:\\AIGC\\python\\py310\\python.exe -m pip install pynini-2.1.6.post1-cp3xx-cp3xx-win_amd64.whl\nD:\\AIGC\\python\\py310\\python.exe -m pip install importlib_resources\nD:\\AIGC\\python\\py310\\python.exe -m pip install WeTextProcessing>=1.0.4 --no-deps\n```\n\n- **Then, proceed with the normal installation**:\n```\ncd ComfyUI/custom_nodes\ngit clone https://github.com/billwuhao/ComfyUI_MegaTTS3.git\ncd ComfyUI_MegaTTS3\npip install -r requirements.txt\n\n# For python_embeded\n./python_embeded/python.exe -m pip install -r requirements.txt\n```\n\n## Model Download\n\n- Models and voices need to be downloaded manually and placed in the `ComfyUI\\models\\TTS` directory:\n\n[MegaTTS3](https://huggingface.co/ByteDance/MegaTTS3/tree/main) Download the entire folder and place it in the `TTS` directory.\n\n- **For the VAE encoder model, which enables custom voice cloning without `.npy` files, please follow our WeChat Official Account to obtain it. Place it in the `TTS\\MegaTTS3\\wavvae` folder**:\n\n\"\"\n\n - [Google Cloud Drive](https://drive.google.com/drive/folders/1p9GNdNJqeK_94lIJW8lew_G3EazU-9Wx?usp=sharing)\n\n- Please place the audio in the `TTS\\speakers` directory. I will unify all speaker audios for TTS nodes into the `ComfyUI\\models\\TTS\\speakers` path. These nodes include `IndexTTS, CSM, Dia, KokoroTTS, MegaTTS, QuteTTS, SparkTTS, StepAudioTTS`, etc.\n\nThe structure is as follows:\n\n```\n.\n│ .gitattributes\n│ config.json\n│ README.md\n│\n├─aligner_lm\n│ config.yaml\n│ model_only_last.ckpt\n│\n├─diffusion_transformer\n│ config.yaml\n│ model_only_last.ckpt\n│\n├─duration_lm\n│ config.yaml\n│ model_only_last.ckpt\n│\n├─g2p\n│ added_tokens.json\n│ config.json\n│ generation_config.json\n│ latest\n│ merges.txt\n│ model.safetensors\n│ special_tokens_map.json\n│ tokenizer.json\n│ tokenizer_config.json\n│ trainer_state.json\n│ vocab.json\n│\n└─wavvae\n config.yaml\n decoder.ckpt\n model_only_last.ckpt\n```\n\n\n## Credits\n\n- [MegaTTS3](https://github.com/bytedance/MegaTTS3)\n\n## Donation\n\nYour appreciation is my greatest motivation! Thank you for supporting me with a cup of coffee!\n\n\"\"\n" }, { "path": "__init__.py", "content": "from .megatts3node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS\r\n\r\n__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']" }, { "path": "megatts3node.py", "content": "import json\nimport os\nimport librosa\nimport numpy as np\nimport torch\nimport torchaudio\nfrom typing import List, Union, Optional\nfrom tn.chinese.normalizer import Normalizer as ZhNormalizer\nfrom tn.english.normalizer import Normalizer as EnNormalizer\nfrom langdetect import detect as classify_language\nimport pyloudnorm as pyln\nimport folder_paths\nimport gc\nimport re\nimport sys\ncurrent_dir = os.path.dirname(os.path.abspath(__file__))\nif current_dir not in sys.path:\n sys.path.append(current_dir)\n\nfrom tts.modules.ar_dur.commons.nar_tts_modules import LengthRegulator\nfrom tts.frontend_function import g2p, align, make_dur_prompt, dur_pred, prepare_inputs_for_dit\nfrom tts.utils.audio_utils.io import convert_to_wav_bytes, combine_audio_segments\nfrom tts.utils.commons.ckpt_utils import load_ckpt\nfrom tts.utils.commons.hparams import set_hparams, hparams\nfrom tts.utils.text_utils.text_encoder import TokenTextEncoder\nfrom tts.utils.text_utils.split_text import chunk_text_chinese, chunk_text_english, chunk_text_chinesev2\nfrom tts.utils.commons.hparams import hparams, set_hparams\n\n\nmodels_dir = folder_paths.models_dir\nmodel_path = os.path.join(models_dir, \"TTS\")\nspeakers_dir = os.path.join(model_path, \"speakers\")\ncache_dir = folder_paths.get_temp_directory()\n\ndef get_all_files(\n root_dir: str,\n return_type: str = \"list\",\n extensions: Optional[List[str]] = None,\n exclude_dirs: Optional[List[str]] = None,\n relative_path: bool = False\n) -> Union[List[str], dict]:\n \"\"\"\n 递归获取目录下所有文件路径\n \n :param root_dir: 要遍历的根目录\n :param return_type: 返回类型 - \"list\"(列表) 或 \"dict\"(按目录分组)\n :param extensions: 可选的文件扩展名过滤列表 (如 ['.py', '.txt'])\n :param exclude_dirs: 要排除的目录名列表 (如 ['__pycache__', '.git'])\n :param relative_path: 是否返回相对路径 (相对于root_dir)\n :return: 文件路径列表或字典\n \"\"\"\n file_paths = []\n file_dict = {}\n \n # 规范化目录路径\n root_dir = os.path.normpath(root_dir)\n \n for dirpath, dirnames, filenames in os.walk(root_dir):\n # 处理排除目录\n if exclude_dirs:\n dirnames[:] = [d for d in dirnames if d not in exclude_dirs]\n \n current_files = []\n for filename in filenames:\n # 扩展名过滤\n if extensions:\n if not any(filename.lower().endswith(ext.lower()) for ext in extensions):\n continue\n \n # 构建完整路径\n full_path = os.path.join(dirpath, filename)\n \n # 处理相对路径\n if relative_path:\n full_path = os.path.relpath(full_path, root_dir)\n \n current_files.append(full_path)\n \n if return_type == \"dict\":\n # 使用相对路径或绝对路径作为键\n dict_key = os.path.relpath(dirpath, root_dir) if relative_path else dirpath\n if current_files:\n file_dict[dict_key] = current_files\n else:\n file_paths.extend(current_files)\n \n return file_dict if return_type == \"dict\" else file_paths\n\ndef get_speakers():\n if not os.path.exists(speakers_dir):\n os.makedirs(speakers_dir, exist_ok=True)\n return []\n speakers = get_all_files(speakers_dir, extensions=[\".wav\", \".mp3\", \".flac\", \".mp4\", \".WAV\", \".MP3\", \".FLAC\", \".MP4\"], relative_path=True)\n return speakers\n\n\nclass MegaTTS3DiTInfer():\n def __init__(\n self,\n device=None,\n ckpt_root=os.path.join(model_path, \"MegaTTS3\"),\n dit_exp_name='diffusion_transformer',\n frontend_exp_name='aligner_lm',\n wavvae_exp_name='wavvae',\n dur_ckpt_path='duration_lm',\n g2p_exp_name='g2p',\n precision=torch.float16,\n **kwargs\n ):\n self.sr = 24000\n self.fm = 8\n if device is None:\n device = 'cuda' if torch.cuda.is_available() else 'cpu'\n self.device = device\n self.precision = precision\n\n # build models\n self.dit_exp_name = os.path.join(ckpt_root, dit_exp_name)\n self.frontend_exp_name = os.path.join(ckpt_root, frontend_exp_name)\n self.wavvae_exp_name = os.path.join(ckpt_root, wavvae_exp_name)\n self.dur_exp_name = os.path.join(ckpt_root, dur_ckpt_path)\n self.g2p_exp_name = os.path.join(ckpt_root, g2p_exp_name)\n self.build_model(self.device)\n\n # init text normalizer\n self.zh_normalizer = ZhNormalizer(overwrite_cache=False, remove_erhua=False, remove_interjections=False)\n self.en_normalizer = EnNormalizer(overwrite_cache=False)\n\n # loudness meter\n self.loudness_meter = pyln.Meter(self.sr)\n \n self.ph_ref = None\n self.tone_ref = None\n self.mel2ph_ref = None\n self.vae_latent = None\n self.ctx_dur_tokens = None\n self.incremental_state_dur_prompt = None\n\n self.audio_bytes = None\n \n def clean(self):\n import gc\n self.dur_model = None\n self.dit= None\n self.g2p_model = None\n self.wavvae_en = None\n self.wavvae_de = None\n self.aligner_lm = None\n\n self.audio_bytes = None\n self.ph_ref = None\n self.tone_ref = None\n self.mel2ph_ref = None\n self.vae_latent = None\n self.ctx_dur_tokens = None\n self.incremental_state_dur_prompt = None\n\n gc.collect()\n torch.cuda.empty_cache()\n\n def build_model(self, device):\n set_hparams(exp_name=self.dit_exp_name, print_hparams=False)\n\n ''' Load Dict '''\n current_dir = os.path.dirname(os.path.abspath(__file__))\n ling_dict = json.load(open(f\"{current_dir}/tts/utils/text_utils/dict.json\", encoding='utf-8-sig'))\n self.ling_dict = {k: TokenTextEncoder(None, vocab_list=ling_dict[k], replace_oov='') for k in ['phone', 'tone']}\n self.token_encoder = token_encoder = self.ling_dict['phone']\n ph_dict_size = len(token_encoder)\n\n ''' Load Duration LM '''\n from tts.modules.ar_dur.ar_dur_predictor import ARDurPredictor\n hp_dur_model = self.hp_dur_model = set_hparams(f'{self.dur_exp_name}/config.yaml', global_hparams=False)\n hp_dur_model['frames_multiple'] = hparams['frames_multiple']\n self.dur_model = ARDurPredictor(\n hp_dur_model, hp_dur_model['dur_txt_hs'], hp_dur_model['dur_model_hidden_size'],\n hp_dur_model['dur_model_layers'], ph_dict_size,\n hp_dur_model['dur_code_size'],\n use_rot_embed=hp_dur_model.get('use_rot_embed', False))\n self.length_regulator = LengthRegulator()\n load_ckpt(self.dur_model, f'{self.dur_exp_name}', 'dur_model')\n self.dur_model.eval()\n self.dur_model.to(device)\n\n ''' Load Diffusion Transformer '''\n from tts.modules.llm_dit.dit import Diffusion\n self.dit = Diffusion()\n load_ckpt(self.dit, f'{self.dit_exp_name}', 'dit', strict=False)\n self.dit.eval()\n self.dit.to(device)\n self.cfg_mask_token_phone = 302 - 1\n self.cfg_mask_token_tone = 32 - 1\n\n ''' Load Frontend LM '''\n from tts.modules.aligner.whisper_small import Whisper\n self.aligner_lm = Whisper()\n load_ckpt(self.aligner_lm, f'{self.frontend_exp_name}', 'model')\n self.aligner_lm.eval()\n self.aligner_lm.to(device)\n self.kv_cache = None\n self.hooks = None\n\n ''' Load G2P LM'''\n from transformers import AutoTokenizer, AutoModelForCausalLM\n g2p_tokenizer = AutoTokenizer.from_pretrained(self.g2p_exp_name, padding_side=\"right\")\n g2p_tokenizer.padding_side = \"right\"\n self.g2p_model = AutoModelForCausalLM.from_pretrained(self.g2p_exp_name).eval().to(device)\n self.g2p_tokenizer = g2p_tokenizer\n self.speech_start_idx = g2p_tokenizer.encode('')[0]\n\n ''' Wav VAE '''\n self.hp_wavvae = hp_wavvae = set_hparams(f'{self.wavvae_exp_name}/config.yaml', global_hparams=False)\n from tts.modules.wavvae.decoder.wavvae_v3 import WavVAE_V3\n\n self.wavvae_en = WavVAE_V3(hparams=hp_wavvae)\n self.wavvae_de = WavVAE_V3(hparams=hp_wavvae)\n\n if os.path.exists(f'{self.wavvae_exp_name}/model_only_last.ckpt'):\n load_ckpt(self.wavvae_en, f'{self.wavvae_exp_name}/model_only_last.ckpt', 'model_gen', strict=True)\n self.has_vae_encoder = True\n self.wavvae_en.eval()\n self.wavvae_en.to(device)\n else:\n load_ckpt(self.wavvae_de, f'{self.wavvae_exp_name}/decoder.ckpt', 'model_gen', strict=False)\n self.has_vae_encoder = False\n self.wavvae_de.eval()\n self.wavvae_de.to(device)\n\n self.vae_stride = hp_wavvae.get('vae_stride', 4)\n self.hop_size = hp_wavvae.get('hop_size', 4)\n \n def preprocess(self, audio_bytes, latent_file=None, topk_dur=1, **kwargs):\n if self.audio_bytes != audio_bytes:\n self.audio_bytes = audio_bytes\n wav_bytes = convert_to_wav_bytes(audio_bytes)\n\n ''' Load wav '''\n wav, _ = librosa.core.load(wav_bytes, sr=self.sr)\n # Pad wav if necessary\n ws = hparams['win_size']\n if len(wav) % ws < ws - 1:\n wav = np.pad(wav, (0, ws - 1 - (len(wav) % ws)), mode='constant', constant_values=0.0).astype(np.float32)\n wav = np.pad(wav, (0, 12000), mode='constant', constant_values=0.0).astype(np.float32)\n self.loudness_prompt = self.loudness_meter.integrated_loudness(wav.astype(float))\n\n ''' obtain alignments with aligner_lm '''\n ph_ref, tone_ref, mel2ph_ref = align(self, wav)\n\n self.kv_cache = None\n self.hooks = None\n\n with torch.inference_mode():\n ''' Forward WaveVAE to obtain: prompt latent '''\n if self.has_vae_encoder:\n if latent_file is None:\n wav = torch.FloatTensor(wav)[None].to(self.device)\n vae_latent = self.wavvae_en.encode_latent(wav)\n else:\n vae_latent = torch.from_numpy(np.load(latent_file)).to(self.device)\n vae_latent = vae_latent[:, :mel2ph_ref.size(1)//4]\n else:\n assert latent_file is not None, \"WaveVAE encode model does not exist, an npy file must be provided!!!\"\n vae_latent = torch.from_numpy(np.load(latent_file)).to(self.device)\n vae_latent = vae_latent[:, :mel2ph_ref.size(1)//4]\n \n ''' Duration Prompting '''\n self.dur_model.hparams[\"infer_top_k\"] = topk_dur if topk_dur > 1 else None\n incremental_state_dur_prompt, ctx_dur_tokens = make_dur_prompt(self, mel2ph_ref, ph_ref, tone_ref)\n\n self.ph_ref = ph_ref.to(self.device)\n self.tone_ref = tone_ref.to(self.device)\n self.mel2ph_ref = mel2ph_ref.to(self.device)\n self.vae_latent = vae_latent.to(self.device)\n self.ctx_dur_tokens = ctx_dur_tokens.to(self.device)\n self.incremental_state_dur_prompt = incremental_state_dur_prompt\n\n def forward(self, texts, time_step, p_w, t_w, dur_disturb=0.1, dur_alpha=1.0, **kwargs):\n\n with torch.inference_mode():\n ''' Generating '''\n waveforms = []\n for input_text in texts:\n wav_pred_ = []\n language_type = classify_language(input_text)\n if language_type == 'en':\n input_text = self.en_normalizer.normalize(input_text)\n text_segs = chunk_text_english(input_text, max_chars=130)\n else:\n input_text = self.zh_normalizer.normalize(input_text)\n text_segs = chunk_text_chinesev2(input_text, limit=60)\n\n for seg_i, text in enumerate(text_segs):\n ''' G2P '''\n ph_pred, tone_pred = g2p(self, text)\n\n ''' Duration Prediction '''\n mel2ph_pred = dur_pred(self, self.ctx_dur_tokens, self.incremental_state_dur_prompt, ph_pred, tone_pred, seg_i, dur_disturb, dur_alpha, is_first=seg_i==0, is_final=seg_i==len(text_segs)-1)\n \n inputs = prepare_inputs_for_dit(self, self.mel2ph_ref, mel2ph_pred, self.ph_ref, self.tone_ref, ph_pred, tone_pred, self.vae_latent)\n # Speech dit inference\n with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):\n x = self.dit.inference(inputs, timesteps=time_step, seq_cfg_w=[p_w, t_w]).float()\n \n # WavVAE decode\n x[:, :self.vae_latent.size(1)] = self.vae_latent\n if self.has_vae_encoder:\n wav_pred = self.wavvae_en.decode(x)[0,0].to(torch.float32)\n else:\n wav_pred = self.wavvae_de.decode(x)[0,0].to(torch.float32)\n \n ''' Post-processing '''\n # Trim prompt wav\n wav_pred = wav_pred[self.vae_latent.size(1)*self.vae_stride*self.hop_size:].cpu().numpy()\n # Norm generated wav to prompt wav's level\n meter = pyln.Meter(self.sr) # create BS.1770 meter\n loudness_pred = self.loudness_meter.integrated_loudness(wav_pred.astype(float))\n wav_pred = pyln.normalize.loudness(wav_pred, loudness_pred, self.loudness_prompt)\n if np.abs(wav_pred).max() >= 1:\n wav_pred = wav_pred / np.abs(wav_pred).max() * 0.95\n\n # Apply hamming window\n wav_pred_.append(wav_pred)\n\n gc.collect()\n torch.cuda.empty_cache()\n\n wav_pred = combine_audio_segments(wav_pred_, sr=self.sr).astype(np.float32)\n waveform = torch.tensor(wav_pred)\n waveforms.append(waveform.cpu())\n\n return torch.cat(waveforms, dim=0), self.sr\n\n\nclass MegaTTS3SpeakersPreview:\n @classmethod\n def INPUT_TYPES(s):\n speakers = get_speakers()\n return {\n \"required\": {\"speaker\":(speakers,),},}\n\n RETURN_TYPES = (\"AUDIO\", \"STRING\", )\n RETURN_NAMES = (\"audio\", \"npy_file\", )\n FUNCTION = \"preview\"\n CATEGORY = \"🎤MW/MW-MegaTTS3\"\n\n def preview(self, speaker):\n wav_path = os.path.join(speakers_dir, speaker)\n latent_file = wav_path.rsplit('.', 1)[0] + '.npy'\n if not os.path.exists(latent_file):\n latent_file = \"\"\n\n waveform, sample_rate = torchaudio.load(wav_path)\n waveform = waveform.unsqueeze(0)\n output_audio = {\n \"waveform\": waveform,\n \"sample_rate\": sample_rate\n }\n return (output_audio, latent_file)\n\n\ndef cache_audio_tensor(\n cache_dir,\n audio_tensor: torch.Tensor,\n sample_rate: int,\n filename_prefix: str = \"cached_audio_\",\n audio_format: Optional[str] = \".wav\"\n) -> str:\n import tempfile\n try:\n with tempfile.NamedTemporaryFile(\n prefix=filename_prefix,\n suffix=audio_format,\n dir=cache_dir,\n delete=False \n ) as tmp_file:\n temp_filepath = tmp_file.name\n \n torchaudio.save(temp_filepath, audio_tensor, sample_rate)\n\n return temp_filepath\n except Exception as e:\n raise Exception(f\"Error caching audio tensor: {e}\")\n\ndef statistical_compare(tensor1, tensor2):\n \"\"\"通过统计特征快速比较\"\"\"\n stats1 = {\n 'mean': tensor1.mean(),\n 'std': tensor1.std(),\n 'max': tensor1.max(),\n 'min': tensor1.min()\n }\n stats2 = {\n 'mean': tensor2.mean(),\n 'std': tensor2.std(),\n 'max': tensor2.max(),\n 'min': tensor2.min()\n }\n return all(torch.allclose(stats1[k], stats2[k], rtol=1e-3) for k in stats1)\n\n\nINFER_INS_CACHE = None\nclass MegaTTS3Run:\n def __init__(self):\n self.resource_context = None\n self.audio_tensor = None\n self.audio_prompt = None\n\n @classmethod\n def INPUT_TYPES(s):\n return {\n \"required\": {\n \"audio\": (\"AUDIO\",),\n \"text\": (\"STRING\", {\"forceInput\": True}),\n \"time_step\": (\"INT\", {\"default\": 32, \"min\": 1,}),\n \"p_w\": (\"FLOAT\", {\"default\":1.6, \"min\": 0.1,}),\n \"t_w\": (\"FLOAT\", {\"default\": 2.5, \"min\": 0.1,}),\n \"unload_model\": (\"BOOLEAN\", {\"default\": True}),\n },\n \"optional\": {\n \"dialogue_audio_s2\":(\"AUDIO\",),\n \"audio_npy_file\": (\"STRING\", {\"forceInput\": True, \"tooltip\": \"No `npy_file` will use VAE to encode audio. 不提供 .npy 文件, 将使用 WaveVAE 编码音频\"}),\n \"audio_s2_npy_file\": (\"STRING\", {\"forceInput\": True, \"tooltip\": \"No `npy_file` will use VAE to encode audio. 不提供 .npy 文件, 将使用 WaveVAE 编码音频\"}),\n }\n }\n\n RETURN_TYPES = (\"AUDIO\",)\n RETURN_NAMES = (\"audio\",)\n FUNCTION = \"clone\"\n CATEGORY = \"🎤MW/MW-MegaTTS3\"\n\n def clone(self, audio, text, time_step, p_w, t_w, unload_model, audio_npy_file=None, dialogue_audio_s2=None, audio_s2_npy_file=None):\n if not os.path.exists(os.path.join(model_path, \"MegaTTS3\", 'wavvae', 'model_only_last.ckpt')):\n print(\"WaveVAE encode model does not exist, an npy file must be provided!!!\")\n waveform = audio[\"waveform\"].squeeze(0)\n\n global INFER_INS_CACHE\n if INFER_INS_CACHE is None:\n INFER_INS_CACHE = MegaTTS3DiTInfer()\n \n latent_file = audio_npy_file if audio_npy_file else None\n try:\n import gc\n if dialogue_audio_s2 is None:\n # 只有音频改变时, 才重新预处理\n if self.audio_tensor is None or self.audio_prompt is None or statistical_compare(self.audio_tensor, waveform) == False:\n self.audio_tensor = waveform\n self.audio_prompt = cache_audio_tensor(cache_dir, waveform, audio[\"sample_rate\"])\n\n texts = [i.strip() for i in re.split(r'\\n\\s*\\n', text.strip()) if i.strip()]\n with open(self.audio_prompt, 'rb') as file:\n file_content = file.read()\n INFER_INS_CACHE.preprocess(file_content, latent_file=latent_file)\n\n del file_content\n gc.collect()\n torch.cuda.empty_cache()\n\n waveform, sr = INFER_INS_CACHE.forward(texts=texts, time_step=time_step, p_w=p_w, t_w=t_w)\n gc.collect()\n torch.cuda.empty_cache()\n else:\n latent_file_2 = audio_s2_npy_file if audio_s2_npy_file else None\n audio_1 = cache_audio_tensor(cache_dir, waveform, audio[\"sample_rate\"])\n audio_2 = cache_audio_tensor(cache_dir, dialogue_audio_s2[\"waveform\"].squeeze(0), dialogue_audio_s2[\"sample_rate\"])\n with open(audio_1, 'rb') as file:\n file_content_1 = file.read()\n with open(audio_2, 'rb') as file:\n file_content_2 = file.read()\n\n gc.collect()\n torch.cuda.empty_cache()\n \n ress = []\n for t, a, n in self.get_speaker_text_audio(text, audio_1, audio_2):\n texts = [i.strip() for i in re.split(r'\\n\\s*\\n', t.strip()) if i.strip()]\n if a == audio_1:\n INFER_INS_CACHE.preprocess(file_content_1, latent_file=latent_file)\n res_sub, sr = INFER_INS_CACHE.forward(texts=texts, time_step=time_step, p_w=p_w, t_w=t_w)\n ress.append([res_sub, n])\n else:\n INFER_INS_CACHE.preprocess(file_content_2, latent_file=latent_file_2)\n res_sub, sr = INFER_INS_CACHE.forward(texts=texts, time_step=time_step, p_w=p_w, t_w=t_w)\n ress.append([res_sub, n])\n\n del file_content_1\n del file_content_2\n gc.collect()\n torch.cuda.empty_cache()\n waveform = torch.cat(list(zip(*sorted(ress, key=lambda x: x[1])))[0], dim=0)\n\n except Exception as e:\n if unload_model:\n import gc\n INFER_INS_CACHE.clean()\n INFER_INS_CACHE = None\n self.resource_context = None\n gc.collect()\n torch.cuda.empty_cache()\n raise e\n\n if unload_model:\n import gc\n INFER_INS_CACHE.clean()\n INFER_INS_CACHE = None\n self.resource_context = None\n gc.collect()\n torch.cuda.empty_cache()\n\n return ({\"waveform\": waveform.unsqueeze(0).unsqueeze(0), \"sample_rate\": sr},)\n\n def get_speaker_text_audio(self, text, audio_1, audio_2):\n pattern = r'(\\[s?S?1\\]|\\[s?S?2\\])\\s*([\\s\\S]*?)(?=\\[s?S?[12]\\]|$)'\n matches = re.findall(pattern, text)\n if len(matches) == 0:\n raise ValueError(\"No speaker tags found in the text: [S1]... [S2]...\")\n labels = []\n contents = []\n audios = []\n\n for label, content in matches:\n labels.append(label)\n contents.append(content)\n \n audios = [\n audio_1 if i.lower() == '[s1]' else audio_2 for i in labels\n ]\n\n return sorted(zip(contents, audios, range(len(contents))), key=lambda x: x[1])\n \n\nclass MultiLinePromptMG:\n @classmethod\n def INPUT_TYPES(cls):\n \n return {\n \"required\": {\n \"multi_line_prompt\": (\"STRING\", {\n \"multiline\": True, \n \"default\": \"\"}),\n },\n }\n\n CATEGORY = \"🎤MW/MW-MegaTTS3\"\n RETURN_TYPES = (\"STRING\",)\n RETURN_NAMES = (\"text\",)\n FUNCTION = \"promptgen\"\n \n def promptgen(self, multi_line_prompt: str):\n return (multi_line_prompt.strip(),)\n\n\nNODE_CLASS_MAPPINGS = {\n \"MegaTTS3SpeakersPreview\": MegaTTS3SpeakersPreview,\n \"MegaTTS3Run\": MegaTTS3Run,\n \"MultiLinePromptMG\": MultiLinePromptMG,\n}\n\nNODE_DISPLAY_NAME_MAPPINGS = {\n \"MegaTTS3SpeakersPreview\": \"MegaTTS3 Speakers Preview\",\n \"MegaTTS3Run\": \"MegaTTS3 Run\",\n \"MultiLinePromptMG\": \"Multi Line Text\",\n}" }, { "path": "pyproject.toml", "content": "[project]\r\nname = \"megatts3-mw\"\r\ndescription = \"Lightweight and Efficient, 🎧Ultra High-Quality Voice Cloning, Chinese and English.\"\r\nversion = \"2.0.0\"\r\nlicense = {file = \"LICENSE\"}\r\ndependencies = [\"setproctitle\", \"attrdict\", \"librosa\", \"pydub\", \"pyloudnorm\", \"x-transformers\", \"torchdiffeq\", \"openai-whisper>=20240930\"]\r\n\r\n[project.urls]\r\nRepository = \"https://github.com/billwuhao/ComfyUI_MegaTTS3\"\r\n# Used by Comfy Registry https://comfyregistry.org\r\n\r\n[tool.comfy]\r\nPublisherId = \"mw\"\r\nDisplayName = \"MW-ComfyUI_MegaTTS3\"\r\nIcon = \"\"\r\n" }, { "path": "requirements.txt", "content": "setproctitle\nattrdict\nlibrosa\npyloudnorm\nx-transformers\ntorchdiffeq\nopenai-whisper>=20240930\nlangdetect\npynini==2.1.6; platform_system!=\"Windows\"\nWeTextProcessing>=1.0.3; platform_system!=\"Windows\"\n" }, { "path": "tts/frontend_function.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\nimport torch.nn.functional as F\nimport whisper\nimport librosa\nfrom copy import deepcopy\nfrom tts.utils.text_utils.ph_tone_convert import split_ph_timestamp, split_ph\nfrom tts.utils.audio_utils.align import mel2token_to_dur\n\n''' Graphme to phoneme function '''\ndef g2p(self, text_inp):\n # prepare inputs\n txt_token = self.g2p_tokenizer('' + text_inp + '')['input_ids']\n input_ids = torch.LongTensor([txt_token+[145+self.speech_start_idx]]).to(self.device)\n\n # model forward\n with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):\n outputs = self.g2p_model.generate(input_ids, max_new_tokens=256, do_sample=True, top_k=1, eos_token_id=800+1+self.speech_start_idx)\n \n # process outputs\n ph_tokens = outputs[:, len(txt_token):-1]-self.speech_start_idx\n ph_pred, tone_pred = split_ph(ph_tokens[0])\n ph_pred, tone_pred = ph_pred[None, :].to(self.device), tone_pred[None, :].to(self.device)\n return ph_pred, tone_pred\n\n''' Get phoneme2mel align of prompt speech '''\ndef align(self, wav):\n with torch.inference_mode():\n whisper_wav = librosa.resample(wav, orig_sr=self.sr, target_sr=16000)\n mel = torch.FloatTensor(whisper.log_mel_spectrogram(whisper_wav).T).to(self.device)[None].transpose(1,2)\n prompt_max_frame = mel.size(2) // self.fm * self.fm\n mel = mel[:, :, :prompt_max_frame]\n token = torch.LongTensor([[798]]).to(self.device)\n audio_features = self.aligner_lm.embed_audio(mel)\n for i in range(768):\n with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):\n logits = self.aligner_lm.logits(token, audio_features, None)\n token_pred = torch.argmax(F.softmax(logits[:, -1], dim=-1), 1)[None]\n token = torch.cat([token, token_pred], dim=1)\n if token_pred[0] == 799:\n break\n alignment_tokens = token\n \n ph_ref, tone_ref, dur_ref, _ = split_ph_timestamp(deepcopy(alignment_tokens)[0, 1:-1])\n ph_ref = torch.Tensor(ph_ref)[None].to(self.device)\n tone_ref = torch.Tensor(tone_ref)[None].to(self.device)\n if dur_ref.sum() < prompt_max_frame:\n dur_ref[-1] += prompt_max_frame - dur_ref.sum()\n elif dur_ref.sum() > prompt_max_frame:\n len_diff = dur_ref.sum() - prompt_max_frame\n while True:\n for i in range(len(dur_ref)):\n dur_ref[i] -= 1\n len_diff -= 1\n if len_diff == 0:\n break\n if len_diff == 0:\n break\n mel2ph_ref = self.length_regulator(dur_ref[None]).to(self.device)\n mel2ph_ref = mel2ph_ref[:, :mel2ph_ref.size(1)//self.fm*self.fm]\n return ph_ref, tone_ref, mel2ph_ref\n\n''' Duration Prompting '''\ndef make_dur_prompt(self, mel2ph_ref, ph_ref, tone_ref):\n dur_tokens_2d_ = mel2token_to_dur(mel2ph_ref, ph_ref.shape[1]).clamp(\n max=self.hp_dur_model['dur_code_size'] - 1) + 1\n\n ctx_dur_tokens = dur_tokens_2d_.clone().flatten(0, 1).to(self.device)\n txt_tokens_flat_ = ph_ref.flatten(0, 1)\n ctx_dur_tokens = ctx_dur_tokens[txt_tokens_flat_ > 0][None]\n\n last_dur_pos_prompt = ctx_dur_tokens.shape[1]\n dur_spk_pos_ids_flat = range(0, last_dur_pos_prompt)\n dur_spk_pos_ids_flat = torch.LongTensor([dur_spk_pos_ids_flat]).to(self.device)\n with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):\n _, incremental_state_dur_prompt = self.dur_model.infer(\n ph_ref, {'tone': tone_ref}, None, None, None,\n ctx_vqcodes=ctx_dur_tokens, spk_pos_ids_flat=dur_spk_pos_ids_flat, return_state=True)\n return incremental_state_dur_prompt, ctx_dur_tokens\n\n''' Duration Prediction '''\ndef dur_pred(self, ctx_dur_tokens, incremental_state_dur_prompt, ph_pred, tone_pred, seg_i, dur_disturb, dur_alpha, is_first, is_final):\n last_dur_token = ctx_dur_tokens[:, -1:]\n last_dur_pos_prompt = ctx_dur_tokens.shape[1]\n incremental_state_dur = deepcopy(incremental_state_dur_prompt)\n txt_len = ph_pred.shape[1]\n dur_spk_pos_ids_flat = range(last_dur_pos_prompt, last_dur_pos_prompt + txt_len)\n dur_spk_pos_ids_flat = torch.LongTensor([dur_spk_pos_ids_flat]).to(self.device)\n last_dur_pos_prompt = last_dur_pos_prompt + txt_len\n\n with torch.cuda.amp.autocast(dtype=self.precision, enabled=True):\n dur_pred = self.dur_model.infer(\n ph_pred, {'tone': tone_pred}, None, None, None,\n incremental_state=incremental_state_dur,\n first_decoder_inp=last_dur_token,\n spk_pos_ids_flat=dur_spk_pos_ids_flat,\n )\n\n dur_pred = dur_pred - 1\n dur_pred = dur_pred.clamp(0, self.hp_dur_model['dur_code_size'] - 1)\n # if is_final:\n # dur_pred[:, -1] = dur_pred[:, -1].clamp(64, 128)\n # else:\n # dur_pred[:, -1] = dur_pred[:, -1].clamp(48, 128)\n # if seg_i > 0:\n # dur_pred[:, 0] = 0\n # ['。', '！', '？', 'sil']\n # for sil_token in [148, 153, 166, 145]:\n # dur_pred[ph_pred==sil_token].clamp_min(32)\n # # ['，', '；'] \n # for sil_token in [163, 165]:\n # dur_pred[ph_pred==sil_token].clamp_min(16)\n if not is_final:\n # add 0.32ms for crossfade\n dur_pred[:, -1] = dur_pred[:, -1] + 32\n else:\n dur_pred[:, -1] = dur_pred[:, -1].clamp(64, 128)\n\n ''' DiT target speech generation '''\n dur_disturb_choice = (torch.rand_like(dur_pred.float()) > 0.5).float()\n dur_disturb_r = 1 + torch.rand_like(dur_pred.float()) * dur_disturb\n dur_pred = dur_pred * dur_disturb_r * dur_disturb_choice + \\\n dur_pred / dur_disturb_r * (1 - dur_disturb_choice)\n dur_pred = torch.round(dur_pred * dur_alpha).clamp(0, 127)\n # ['。', '！', '？', 'sil']\n for sil_token in [148, 153, 166, 145]:\n dur_pred[ph_pred==sil_token] = dur_pred[ph_pred==sil_token].clamp_min(64)\n # ['，', '；'] \n for sil_token in [163, 165]:\n dur_pred[ph_pred==sil_token] = dur_pred[ph_pred==sil_token].clamp_min(32)\n if is_first:\n dur_pred[:, 0] = 8\n \n dur_sum = dur_pred.sum()\n npad = self.fm - dur_sum % self.fm\n if npad < self.fm:\n dur_pred[:, -1] += npad\n mel2ph_pred = self.length_regulator(dur_pred).to(self.device)\n return mel2ph_pred\n\ndef prepare_inputs_for_dit(self, mel2ph_ref, mel2ph_pred, ph_ref, tone_ref, ph_pred, tone_pred, vae_latent):\n # Prepare duration token \n mel2ph_pred = torch.cat((mel2ph_ref, mel2ph_pred+ph_ref.size(1)), dim=1)\n mel2ph_pred = mel2ph_pred[:, :mel2ph_pred.size(1)//self.fm*self.fm].repeat(3, 1)\n # Prepare phone and tone token\n ph_pred = torch.cat((ph_ref, ph_pred), dim=1)\n tone_pred = torch.cat((tone_ref, tone_pred), dim=1)\n # Disable the English tone (set them to 3)\"\"\"\n en_tone_idx = ~((tone_pred == 4) | ( (11 <= tone_pred) & (tone_pred <= 15)) | (tone_pred == 0))\n tone_pred[en_tone_idx] = 3\n \n # Prepare cfg inputs\n ph_seq = torch.cat([ph_pred, ph_pred, torch.full(ph_pred.size(), self.cfg_mask_token_phone, device=self.device)], 0)\n tone_seq = torch.cat([tone_pred, tone_pred, torch.full(tone_pred.size(), self.cfg_mask_token_tone, device=self.device)], 0)\n target_size = mel2ph_pred.size(1)//self.vae_stride\n vae_latent_ = vae_latent.repeat(3, 1, 1)\n ctx_mask = torch.ones_like(vae_latent_[:, :, 0:1])\n vae_latent_ = F.pad(vae_latent_, (0, 0, 0, target_size - vae_latent.size(1)), mode='constant', value=0)\n vae_latent_[1:] = 0.0\n ctx_mask = F.pad(ctx_mask, (0, 0, 0, target_size - vae_latent.size(1)), mode='constant', value=0)\n\n return {\n 'phone': ph_seq,\n 'tone': tone_seq,\n \"lat_ctx\": vae_latent_ * ctx_mask,\n \"ctx_mask\": ctx_mask,\n \"dur\": mel2ph_pred,\n }" }, { "path": "tts/modules/aligner/whisper_small.py", "content": "# MIT License\n\n# Copyright (c) 2022 OpenAI\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\n# Copyright (c) [2022] [OpenAI] \n# Copyright (c) [2025] [Ziyue Jiang] \n# SPDX-License-Identifier: MIT\n# This file has been modified by Ziyue Jiang on 2025/03/19\n# Original file was released under MIT, with the full license text # available at https://github.com/openai/whisper/blob/v20240930/LICENSE.\n# This modified file is released under the same license.\n\nfrom contextlib import contextmanager\nfrom typing import Dict, Iterable, Optional, Tuple\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom torch import Tensor, nn\n\nfrom torch.nn.functional import scaled_dot_product_attention\nSDPA_AVAILABLE = True\n\n\nclass LayerNorm(nn.LayerNorm):\n def forward(self, x: Tensor) -> Tensor:\n return super().forward(x.float()).type(x.dtype)\n\n\nclass Linear(nn.Linear):\n def forward(self, x: Tensor) -> Tensor:\n return F.linear(\n x,\n self.weight.to(x.dtype),\n None if self.bias is None else self.bias.to(x.dtype),\n )\n\n\nclass Conv1d(nn.Conv1d):\n def _conv_forward(\n self, x: Tensor, weight: Tensor, bias: Optional[Tensor]\n ) -> Tensor:\n return super()._conv_forward(\n x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)\n )\n\n\ndef sinusoids(length, channels, max_timescale=10000):\n \"\"\"Returns sinusoids for positional embedding\"\"\"\n assert channels % 2 == 0\n log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)\n inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))\n scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]\n return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)\n\n\n@contextmanager\ndef disable_sdpa():\n prev_state = MultiHeadAttention.use_sdpa\n try:\n MultiHeadAttention.use_sdpa = False\n yield\n finally:\n MultiHeadAttention.use_sdpa = prev_state\n\n\nclass MultiHeadAttention(nn.Module):\n use_sdpa = True\n\n def __init__(self, n_state: int, n_head: int):\n super().__init__()\n self.n_head = n_head\n self.query = Linear(n_state, n_state)\n self.key = Linear(n_state, n_state, bias=False)\n self.value = Linear(n_state, n_state)\n self.out = Linear(n_state, n_state)\n\n def forward(\n self,\n x: Tensor,\n xa: Optional[Tensor] = None,\n mask: Optional[Tensor] = None,\n kv_cache: Optional[dict] = None,\n casual: Optional[bool] = None\n ):\n q = self.query(x)\n\n if kv_cache is None or xa is None or self.key not in kv_cache:\n # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;\n # otherwise, perform key/value projections for self- or cross-attention as usual.\n k = self.key(x if xa is None else xa)\n v = self.value(x if xa is None else xa)\n else:\n # for cross-attention, calculate keys and values once and reuse in subsequent calls.\n k = kv_cache[self.key]\n v = kv_cache[self.value]\n\n wv = self.qkv_attention(q, k, v, mask, casual)\n return self.out(wv)\n\n def qkv_attention(\n self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None, casual: Optional[bool] = None\n ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:\n n_batch, n_ctx, n_state = q.shape\n scale = (n_state // self.n_head) ** -0.25\n q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)\n k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)\n v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)\n\n a = scaled_dot_product_attention(\n q, k, v, is_causal=casual and n_ctx > 1, attn_mask=mask[:, None, None, :] if mask is not None else None\n )\n out = a.permute(0, 2, 1, 3).flatten(start_dim=2)\n return out\n\n\nclass ResidualAttentionBlock(nn.Module):\n def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):\n super().__init__()\n\n self.attn = MultiHeadAttention(n_state, n_head)\n self.attn_ln = LayerNorm(n_state)\n\n self.cross_attn = (\n MultiHeadAttention(n_state, n_head) if cross_attention else None\n )\n self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None\n\n n_mlp = n_state * 4\n self.mlp = nn.Sequential(\n Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state)\n )\n self.mlp_ln = LayerNorm(n_state)\n\n def forward(\n self,\n x: Tensor,\n xa: Optional[Tensor] = None,\n mask: Optional[Tensor] = None,\n kv_cache: Optional[dict] = None,\n casual: Optional[bool] = None,\n ):\n x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache, casual=casual)\n if self.cross_attn:\n # TODO: Cross attention mask\n x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache, casual=False)\n x = x + self.mlp(self.mlp_ln(x))\n return x\n\n\nclass AudioEncoder(nn.Module):\n def __init__(\n self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int\n ):\n super().__init__()\n self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)\n self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)\n self.register_buffer(\"positional_embedding\", sinusoids(n_ctx, n_state))\n\n self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(\n [ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]\n )\n self.ln_post = LayerNorm(n_state)\n\n def forward(self, x: Tensor, attn_mask: Tensor):\n \"\"\"\n x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)\n the mel spectrogram of the audio\n \"\"\"\n x = F.gelu(self.conv1(x))\n x = F.gelu(self.conv2(x))\n x = x.permute(0, 2, 1)\n\n # assert x.shape[1:] == self.positional_embedding.shape, \"incorrect audio shape\"\n x = (x + self.positional_embedding[:x.size(1)]).to(x.dtype)\n\n for block in self.blocks:\n x = block(x, mask=attn_mask, casual=False)\n\n x = self.ln_post(x)\n return x\n\n\nclass TextDecoder(nn.Module):\n def __init__(\n self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int\n ):\n super().__init__()\n\n self.token_embedding = nn.Embedding(n_vocab, n_state)\n self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state))\n\n self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(\n [\n ResidualAttentionBlock(n_state, n_head, cross_attention=True)\n for _ in range(n_layer)\n ]\n )\n self.ln = LayerNorm(n_state)\n\n self.out_proj = nn.Linear(n_state, n_vocab)\n\n def forward(self, x: Tensor, attn_mask: Tensor, xa: Tensor, kv_cache: Optional[dict] = None):\n \"\"\"\n x : torch.LongTensor, shape = (batch_size, <= n_ctx)\n the text tokens\n xa : torch.Tensor, shape = (batch_size, n_audio_ctx, n_audio_state)\n the encoded audio features to be attended on\n \"\"\"\n offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0\n x = (\n self.token_embedding(x)\n + self.positional_embedding[offset : offset + x.shape[-1]]\n )\n x = x.to(xa.dtype)\n\n for block in self.blocks:\n x = block(x, xa, mask=attn_mask, kv_cache=kv_cache, casual=True)\n\n x = self.ln(x)\n # logits = (\n # x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)\n # ).float()\n logits = self.out_proj(x)\n\n return logits\n\n\nclass Whisper(nn.Module):\n def __init__(self):\n super().__init__()\n self.n_vocab = 6800\n self.n_text_layer = 6\n self.n_text_head = 8\n self.n_text_ctx = 2048\n\n self.encoder = AudioEncoder(\n n_mels=80, n_ctx=3000, n_state=512, n_head=8, n_layer=6,\n )\n self.decoder = TextDecoder(\n n_vocab=6800, n_ctx=2048, n_state=512, n_head=8, n_layer=6,\n )\n\n def embed_audio(self, mel: torch.Tensor):\n return self.encoder(mel, None)\n\n def logits(self, tokens, audio_features, kv_cache=None):\n return self.decoder(tokens, None, audio_features, kv_cache=kv_cache)\n\n def forward(\n self, mel, mel_len, token, token_len\n ) -> Dict[str, torch.Tensor]:\n attn_mask_enc = self.sequence_mask(mel_len//2, device=mel.device) > 0\n attn_mask_dec = self.sequence_mask(token_len, device=mel.device) > 0\n return self.decoder(token, attn_mask_dec, self.encoder(mel, attn_mask_enc))\n\n @property\n def device(self):\n return next(self.parameters()).device\n\n def install_kv_cache_hooks(self, cache: Optional[dict] = None):\n \"\"\"\n The `MultiHeadAttention` module optionally accepts `kv_cache` which stores the key and value\n tensors calculated for the previous positions. This method returns a dictionary that stores\n all caches, and the necessary hooks for the key and value projection modules that save the\n intermediate tensors to be reused during later calculations.\n\n Returns\n -------\n cache : Dict[nn.Module, torch.Tensor]\n A dictionary object mapping the key/value projection modules to its cache\n hooks : List[RemovableHandle]\n List of PyTorch RemovableHandle objects to stop the hooks to be called\n \"\"\"\n cache = {**cache} if cache is not None else {}\n hooks = []\n\n def save_to_cache(module, _, output):\n if module not in cache or output.shape[1] > self.n_text_ctx:\n # save as-is, for the first token or cross attention\n cache[module] = output\n else:\n cache[module] = torch.cat([cache[module], output], dim=1).detach()\n return cache[module]\n\n def install_hooks(layer: nn.Module):\n if isinstance(layer, MultiHeadAttention):\n hooks.append(layer.key.register_forward_hook(save_to_cache))\n hooks.append(layer.value.register_forward_hook(save_to_cache))\n\n self.decoder.apply(install_hooks)\n return cache, hooks\n \n def sequence_mask(self, seq_lens, max_len=None, device='cpu'):\n b = seq_lens.shape[0]\n if max_len is None:\n max_len = seq_lens.max()\n mask = torch.arange(max_len).unsqueeze(0).to(device) # [1, t]\n mask = mask < (seq_lens.unsqueeze(1)) # [1, t] + [b, 1] = [b, t]\n mask = mask.float()\n return mask\n" }, { "path": "tts/modules/ar_dur/ar_dur_predictor.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 random\nfrom copy import deepcopy\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch.nn import Linear\nfrom tqdm import tqdm\n\nfrom tts.modules.ar_dur.commons.layers import Embedding, LayerNorm\nfrom tts.modules.ar_dur.commons.nar_tts_modules import PosEmb\nfrom tts.modules.ar_dur.commons.rot_transformer import RotTransformerDecoderLayer\nfrom tts.modules.ar_dur.commons.transformer import SinusoidalPositionalEmbedding\nfrom tts.modules.ar_dur.commons.rel_transformer import RelTransformerEncoder\n\nFS_ENCODERS = {\n 'rel_fft': lambda hp, dict_size: RelTransformerEncoder(\n dict_size, hp['hidden_size'], hp['hidden_size'],\n hp['ffn_hidden_size'], hp['num_heads'], hp['enc_layers'],\n hp['enc_ffn_kernel_size'], hp['dropout'], prenet=hp['enc_prenet'], pre_ln=hp['enc_pre_ln']),\n}\n\ndef fill_with_neg_inf2(t):\n \"\"\"FP16-compatible function that fills a tensor with -inf.\"\"\"\n return t.float().fill_(-1e8).type_as(t)\n\ndef expand_states(h, mel2token):\n h = F.pad(h, [0, 0, 1, 0])\n mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]])\n h = torch.gather(h, 1, mel2token_) # [B, T, H]\n return h\n\n\nclass CodePredictor(nn.Module):\n def __init__(self, hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size):\n super().__init__()\n self.hparams = deepcopy(hparams)\n self.hparams['hidden_size'] = hidden_size\n self.hidden_size = hidden_size\n char_dict_size = hparams.get('char_dict_size', 4000)\n if not hparams.get('lm_use_enc'):\n self.encoder = nn.Embedding(dict_size, self.hidden_size, padding_idx=0)\n if hparams.get('mega_use_char', True):\n self.char_encoder = nn.Embedding(char_dict_size,\n self.hidden_size, padding_idx=0)\n else:\n self.encoder = FS_ENCODERS[self.hparams['encoder_type']](self.hparams, dict_size)\n if hparams.get('mega_use_char', True):\n self.char_encoder = FS_ENCODERS[self.hparams['encoder_type']](self.hparams, char_dict_size)\n if hparams['use_ph_pos_embed']:\n self.ph_pos_embed = PosEmb(self.hidden_size)\n\n self.char_empty_embed = nn.Embedding(1, self.hidden_size)\n if hparams.get('use_bert_input'):\n self.bert_input_proj = nn.Linear(768, self.hidden_size)\n self.ling_label_embed_layers = nn.ModuleDict()\n for k, s in zip(hparams['ling_labels'], hparams['ling_label_dict_size']):\n self.ling_label_embed_layers[k] = Embedding(s + 3, self.hidden_size, padding_idx=0)\n\n self.dec_hidden_size = dec_hidden_size\n self.enc_proj = nn.Linear(self.hidden_size, dec_hidden_size)\n self.code_emb = Embedding(code_size + 2, dec_hidden_size, 0)\n self.use_pos_embed = hparams.get('use_pos_embed', False)\n if self.use_pos_embed:\n self.embed_positions = SinusoidalPositionalEmbedding(dec_hidden_size, 0, init_size=1024)\n self.use_post_ln = hparams.get('use_post_ln', False)\n self.layers = None\n if not self.use_post_ln:\n self.layer_norm = LayerNorm(dec_hidden_size)\n self.code_size = code_size\n self.project_out_dim = Linear(dec_hidden_size, code_size + 1, bias=True)\n\n def forward_ling_encoder(\n self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed, spk_id, spk_embed, mels_timbre):\n ph_tokens = txt_tokens\n hparams = self.hparams\n ph_nonpadding = (ph_tokens > 0).float()[:, :, None] # [B, T_phone, 1]\n x_spk = self.forward_style_embed(spk_embed, spk_id, mels_timbre)\n\n # enc_ph\n if not hparams.get('lm_use_enc'):\n x_ph = self.encoder(ph_tokens)\n x_ph = x_ph + sum(\n [self.ling_label_embed_layers[k](ling_feas[k]) for k in hparams['ling_labels']]) \\\n if len(hparams['ling_labels']) > 0 else 0\n x_ph = x_ph + x_spk\n else:\n # enc_ph\n ph_enc_oembed = sum(\n [self.ling_label_embed_layers[k](ling_feas[k]) for k in hparams['ling_labels']]) \\\n if len(hparams['ling_labels']) > 0 else 0\n ph_enc_oembed = ph_enc_oembed + self.ph_pos_embed(\n torch.arange(0, ph_tokens.shape[1])[None,].to(ph_tokens.device))\n ph_enc_oembed = ph_enc_oembed + x_spk\n ph_enc_oembed = ph_enc_oembed * ph_nonpadding\n x_ph = self.encoder(ph_tokens, other_embeds=ph_enc_oembed)\n\n # enc_char\n if char_tokens is not None and ph2char is not None:\n char_nonpadding = (char_tokens > 0).float()[:, :, None]\n x_char = self.char_encoder(char_tokens)\n empty_char = (ph2char > 100000).long()\n ph2char = ph2char * (1 - empty_char)\n x_char_phlevel = \\\n expand_states(x_char * char_nonpadding, ph2char) \\\n * (1 - empty_char)[..., None] + \\\n self.char_empty_embed(torch.zeros_like(ph_tokens)) * empty_char[..., None]\n else:\n x_char_phlevel = 0\n # x_ling\n x_ling = x_ph + x_char_phlevel\n x_ling = x_ling * ph_nonpadding\n x_ling = self.enc_proj(x_ling)\n return x_ling\n\n def sample_one_step(self, vq_pred):\n hparams = self.hparams\n if hparams.get('infer_top_k'):\n top_k = hparams.get('infer_top_k')\n temperature = hparams.get('infer_temperature', 1)\n vq_pred = vq_pred[:, -1] / temperature\n # optionally crop the logits to only the top k options\n if top_k is not None:\n v, _ = torch.topk(vq_pred, min(top_k, vq_pred.size(-1)))\n vq_pred[vq_pred < v[:, [-1]]] = -float('Inf')\n # apply softmax to convert logits to (normalized) probabilities\n probs = F.softmax(vq_pred, dim=-1)\n # sample from the distribution\n vq_pred = torch.multinomial(probs, num_samples=1)\n else:\n vq_pred = torch.argmax(F.softmax(vq_pred[:, -1], dim=-1), 1)\n return vq_pred\n\n def forward_style_embed(self, spk_embed=None, spk_id=None, mel_ref=None):\n # add spk embed\n style_embed = 0\n if self.hparams['use_spk_embed']:\n style_embed = style_embed + self.spk_embed_proj(spk_embed)[:, None, :]\n if self.hparams['use_spk_id']:\n style_embed = style_embed + self.spk_id_proj(spk_id)[:, None, :]\n if self.hparams['use_spk_enc']:\n style_embed = style_embed + self.spk_enc(mel_ref)[:, None, :]\n return style_embed\n\n def buffered_future_mask(self, tensor):\n dim = tensor.size(0)\n if (\n not hasattr(self, '_future_mask')\n or self._future_mask is None\n or self._future_mask.device != tensor.device\n or self._future_mask.size(0) < dim\n ):\n self._future_mask = torch.triu(fill_with_neg_inf2(tensor.new(dim, dim)), 1)\n return self._future_mask[:dim, :dim]\n\n\nclass ARDurPredictor(CodePredictor):\n def __init__(self, hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size, use_rot_embed=True,\n op_version=1):\n super().__init__(hparams, hidden_size, dec_hidden_size, lm_num_layers, dict_size, code_size)\n self.use_rot_embed = use_rot_embed\n bias = hparams.get('lm_bias', True)\n if self.use_rot_embed:\n self.layers = nn.ModuleList([])\n self.layers.extend([\n RotTransformerDecoderLayer(\n dec_hidden_size, 0.0, kernel_size=1, ffn_hidden_size=dec_hidden_size * 4,\n post_ln=self.use_post_ln, op_version=op_version, bias=bias)\n for _ in range(lm_num_layers)\n ])\n if hparams['dur_model_type'] == 'ar_mse':\n self.project_out_dim = nn.Sequential(torch.nn.Linear(dec_hidden_size, 1), nn.Softplus())\n else:\n self.project_out_dim = torch.nn.Linear(dec_hidden_size, code_size + 1)\n\n def forward(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,\n prev_code, spk_id=None, spk_embed=None, mels_timbre=None, mel2ph=None,\n incremental_state=None, x_ling=None, attn_mask=None, spk_pos_ids_flat=None,\n prompt_length=None, cache_size=20, streaming=False):\n x = self.code_emb(prev_code)\n if x_ling is None:\n x_ling = self.forward_ling_encoder(\n txt_tokens, ling_feas, char_tokens, ph2char, bert_embed, spk_id, spk_embed, mels_timbre)\n x_ling = x_ling.flatten(0, 1)\n txt_tokens = txt_tokens.flatten(0, 1)\n x_ling = x_ling[txt_tokens > 0][None]\n\n # run decoder\n self_attn_padding_mask = None\n if self.use_pos_embed:\n positions = self.embed_positions(\n prev_code,\n incremental_state=incremental_state\n )\n if incremental_state is not None:\n x_ling = x_ling[:, x.shape[1] - 1:x.shape[1]]\n if spk_pos_ids_flat is not None:\n spk_pos_ids_flat = spk_pos_ids_flat[:, x.shape[1] - 1:x.shape[1]]\n x = x[:, -1:]\n if self.use_pos_embed:\n positions = positions[:, -1:]\n if streaming:\n # Shift Pos: query pos is min(cache_size, idx)\n spk_pos_ids_flat = torch.min(torch.LongTensor([prompt_length + cache_size]).to(x.device),\n spk_pos_ids_flat)\n\n # # B x T x C -> T x B x C\n if self.use_pos_embed:\n x = x + positions\n x_ling = x_ling[:, :self.hparams['max_tokens']].contiguous()\n T = min(self.hparams.get('max_tokens_per_item', 1e9), x_ling.shape[1])\n x_ling = x_ling.reshape(-1, T, x_ling.shape[-1])\n x = x + x_ling\n x = x.transpose(0, 1)\n\n for idx, layer in enumerate(self.layers):\n if incremental_state is None:\n self_attn_mask = self.buffered_future_mask(x)\n if attn_mask is not None:\n self_attn_mask = self_attn_mask + (1 - attn_mask.float()) * -1e8\n self_attn_mask = self_attn_mask.clamp_min(-1e8)\n else:\n self_attn_mask = None\n\n x, attn_weights = layer(\n x,\n incremental_state=incremental_state,\n self_attn_mask=self_attn_mask,\n self_attn_padding_mask=self_attn_padding_mask,\n spk_pos_ids_flat=spk_pos_ids_flat\n )\n\n if streaming and incremental_state != {}:\n for k, v in incremental_state.items():\n if 'attn_state' in k:\n prev_key, prev_value = incremental_state[k]['prev_key'], incremental_state[k]['prev_value']\n cur_length = prev_key.shape[2]\n if cur_length - prompt_length > cache_size:\n prev_key = torch.cat((prev_key[:, :, :prompt_length], prev_key[:, :, -cache_size:]), dim=2)\n prev_value = torch.cat((prev_value[:, :, :prompt_length], prev_value[:, :, -cache_size:]),\n dim=2)\n incremental_state[k]['prev_key'], incremental_state[k]['prev_value'] = prev_key, prev_value\n\n if not self.use_post_ln:\n x = self.layer_norm(x)\n # T x B x C -> B x T x C\n x = x.transpose(0, 1)\n x = self.project_out_dim(x)\n return x\n\n def infer(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,\n spk_id=None, spk_embed=None, mels_timbre=None,\n incremental_state=None, ctx_vqcodes=None, spk_pos_ids_flat=None, return_state=False,\n first_step_min=0, return_probs=False, first_decoder_inp=None, dur_disturb=0.0, **kwargs):\n if incremental_state is None:\n incremental_state = {}\n x_ling = self.forward_ling_encoder(\n txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,\n spk_id, spk_embed, mels_timbre)\n x_ling = x_ling.flatten(0, 1)\n txt_tokens_ori = txt_tokens\n txt_tokens_withpad = txt_tokens = txt_tokens.flatten(0, 1)\n x_ling = x_ling[txt_tokens > 0][None]\n txt_tokens = txt_tokens[txt_tokens > 0][None]\n\n decoded = torch.zeros_like(txt_tokens)\n decoded = F.pad(decoded, [1, 0], value=self.code_size + 1)\n if incremental_state != {}:\n if first_decoder_inp is None:\n assert ctx_vqcodes is not None\n decoded[:, :ctx_vqcodes.shape[1]] = ctx_vqcodes\n ctx_vqcodes = None\n else:\n decoded[:, :1] = first_decoder_inp\n probs = []\n for step in range(decoded.shape[1] - 1):\n vq_pred = self(txt_tokens, None, None, None, None,\n decoded[:, :step + 1], None, None, None,\n incremental_state=incremental_state, x_ling=x_ling,\n spk_pos_ids_flat=spk_pos_ids_flat, **kwargs)\n probs.append(vq_pred.cpu())\n if ctx_vqcodes is None or step >= ctx_vqcodes.shape[1]:\n if self.hparams['dur_model_type'] == 'ar_mse':\n d = vq_pred[:, -1, 0]\n if dur_disturb > 0 and step >= 1:\n if random.random() > 0.5:\n d = d * (1 + random.random() * dur_disturb)\n else:\n d = d / (1 + random.random() * dur_disturb)\n d = torch.clamp_max(d, self.code_size - 1)\n vq_pred = torch.round(d).long()\n else:\n vq_pred = self.sample_one_step(vq_pred)\n decoded[:, step + 1] = torch.clamp_min(vq_pred, 1)\n if step == 0:\n decoded[:, step + 1] = torch.clamp_min(vq_pred, first_step_min)\n else:\n decoded[:, step + 1] = ctx_vqcodes[:, step]\n decoded = decoded[:, 1:]\n decoded_2d = torch.zeros_like(txt_tokens_ori)\n decoded_2d.flatten(0, 1)[txt_tokens_withpad > 0] = decoded\n if return_state:\n return decoded_2d, incremental_state\n if return_probs:\n return decoded_2d, torch.cat(probs, 1)\n return decoded_2d\n\n def streaming_infer(self, txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,\n spk_id=None, spk_embed=None, mels_timbre=None,\n incremental_state=None, ctx_vqcodes=None, spk_pos_ids_flat=None, return_state=False,\n **kwargs):\n if incremental_state is None:\n incremental_state = {}\n x_ling = self.forward_ling_encoder(\n txt_tokens, ling_feas, char_tokens, ph2char, bert_embed,\n spk_id, spk_embed, mels_timbre)\n x_ling = x_ling.flatten(0, 1)\n txt_tokens_ori = txt_tokens\n txt_tokens_withpad = txt_tokens = txt_tokens.flatten(0, 1)\n x_ling = x_ling[txt_tokens > 0][None]\n txt_tokens = txt_tokens[txt_tokens > 0][None]\n\n vq_decoded = torch.zeros_like(txt_tokens)\n vq_decoded = F.pad(vq_decoded, [1, 0], value=self.code_size + 1)\n if incremental_state != {}:\n assert ctx_vqcodes is not None\n vq_decoded[:, :ctx_vqcodes.shape[1]] = ctx_vqcodes\n ctx_vqcodes = None\n prompt_length = list(incremental_state.items())[0][1]['prev_key'].shape[2]\n for step in tqdm(range(vq_decoded.shape[1] - 1), desc='AR Duration Predictor inference...'):\n vq_pred = self(txt_tokens, None, None, None, None,\n vq_decoded[:, :step + 1], None, None, None,\n incremental_state=incremental_state, x_ling=x_ling,\n spk_pos_ids_flat=spk_pos_ids_flat, prompt_length=prompt_length, streaming=True, **kwargs)\n if ctx_vqcodes is None or step >= ctx_vqcodes.shape[1]:\n if self.hparams['dur_model_type'] == 'ar_mse':\n vq_pred = torch.round(vq_pred[:, -1, 0]).long()\n else:\n vq_pred = self.sample_one_step(vq_pred)\n vq_decoded[:, step + 1] = vq_pred\n else:\n vq_decoded[:, step + 1] = ctx_vqcodes[:, step]\n vq_decoded = vq_decoded[:, 1:]\n vq_decoded_2d = torch.zeros_like(txt_tokens_ori)\n vq_decoded_2d.flatten(0, 1)[txt_tokens_withpad > 0] = vq_decoded\n if return_state:\n return vq_decoded_2d, incremental_state\n return vq_decoded_2d" }, { "path": "tts/modules/ar_dur/commons/layers.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\nfrom torch import nn\n\n\nclass LayerNorm(torch.nn.LayerNorm):\n \"\"\"Layer normalization module.\n :param int nout: output dim size\n :param int dim: dimension to be normalized\n \"\"\"\n\n def __init__(self, nout, dim=-1, eps=1e-5):\n \"\"\"Construct an LayerNorm object.\"\"\"\n super(LayerNorm, self).__init__(nout, eps=eps)\n self.dim = dim\n\n def forward(self, x):\n \"\"\"Apply layer normalization.\n :param torch.Tensor x: input tensor\n :return: layer normalized tensor\n :rtype torch.Tensor\n \"\"\"\n if self.dim == -1:\n return super(LayerNorm, self).forward(x)\n return super(LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)\n\n\nclass Reshape(nn.Module):\n def __init__(self, *args):\n super(Reshape, self).__init__()\n self.shape = args\n\n def forward(self, x):\n return x.view(self.shape)\n\n\nclass Permute(nn.Module):\n def __init__(self, *args):\n super(Permute, self).__init__()\n self.args = args\n\n def forward(self, x):\n return x.permute(self.args)\n\n\ndef Embedding(num_embeddings, embedding_dim, padding_idx=None):\n m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)\n nn.init.normal_(m.weight, mean=0, std=embedding_dim ** -0.5)\n if padding_idx is not None:\n nn.init.constant_(m.weight[padding_idx], 0)\n return m\n" }, { "path": "tts/modules/ar_dur/commons/nar_tts_modules.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 math\n\nimport torch\nfrom torch import nn\n\nimport torch.nn.functional as F\n\n\nclass LengthRegulator(torch.nn.Module):\n def __init__(self, pad_value=0.0):\n super(LengthRegulator, self).__init__()\n self.pad_value = pad_value\n\n def forward(self, dur, dur_padding=None, alpha=1.0):\n \"\"\"\n Example (no batch dim version):\n 1. dur = [2,2,3]\n 2. token_idx = [[1],[2],[3]], dur_cumsum = [2,4,7], dur_cumsum_prev = [0,2,4]\n 3. token_mask = [[1,1,0,0,0,0,0],\n [0,0,1,1,0,0,0],\n [0,0,0,0,1,1,1]]\n 4. token_idx * token_mask = [[1,1,0,0,0,0,0],\n [0,0,2,2,0,0,0],\n [0,0,0,0,3,3,3]]\n 5. (token_idx * token_mask).sum(0) = [1,1,2,2,3,3,3]\n\n :param dur: Batch of durations of each frame (B, T_txt)\n :param dur_padding: Batch of padding of each frame (B, T_txt)\n :param alpha: duration rescale coefficient\n :return:\n mel2ph (B, T_speech)\n assert alpha > 0\n \"\"\"\n dur = torch.round(dur.float() * alpha).long()\n if dur_padding is not None:\n dur = dur * (1 - dur_padding.long())\n token_idx = torch.arange(1, dur.shape[1] + 1)[None, :, None].to(dur.device)\n dur_cumsum = torch.cumsum(dur, 1)\n dur_cumsum_prev = F.pad(dur_cumsum, [1, -1], mode='constant', value=0)\n\n pos_idx = torch.arange(dur.sum(-1).max())[None, None].to(dur.device)\n token_mask = (pos_idx >= dur_cumsum_prev[:, :, None]) & (pos_idx < dur_cumsum[:, :, None])\n mel2token = (token_idx * token_mask.long()).sum(1)\n return mel2token\n\n\nclass PosEmb(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.dim = dim\n half_dim = self.dim // 2\n emb = math.log(10000) / (half_dim - 1)\n emb = torch.exp(torch.arange(half_dim) * -emb)\n self.emb = emb # TODO\n\n def forward(self, x):\n emb = x[:, :, None] * self.emb[None, None, :].to(x.device)\n emb = torch.cat((emb.sin(), emb.cos()), dim=-1)\n return emb\n" }, { "path": "tts/modules/ar_dur/commons/rel_transformer.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 math\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom tts.modules.ar_dur.commons.layers import Embedding\n\n\ndef convert_pad_shape(pad_shape):\n l = pad_shape[::-1]\n pad_shape = [item for sublist in l for item in sublist]\n return pad_shape\n\n\ndef shift_1d(x):\n x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1]\n return x\n\n\ndef sequence_mask(length, max_length=None):\n if max_length is None:\n max_length = length.max()\n x = torch.arange(max_length, dtype=length.dtype, device=length.device)\n return x.unsqueeze(0) < length.unsqueeze(1)\n\n\nclass Encoder(nn.Module):\n def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0.,\n window_size=None, block_length=None, pre_ln=False, **kwargs):\n super().__init__()\n self.hidden_channels = hidden_channels\n self.filter_channels = filter_channels\n self.n_heads = n_heads\n self.n_layers = n_layers\n self.kernel_size = kernel_size\n self.p_dropout = p_dropout\n self.window_size = window_size\n self.block_length = block_length\n self.pre_ln = pre_ln\n\n self.drop = nn.Dropout(p_dropout)\n self.attn_layers = nn.ModuleList()\n self.norm_layers_1 = nn.ModuleList()\n self.ffn_layers = nn.ModuleList()\n self.norm_layers_2 = nn.ModuleList()\n for i in range(self.n_layers):\n self.attn_layers.append(\n MultiHeadAttention(hidden_channels, hidden_channels, n_heads, window_size=window_size,\n p_dropout=p_dropout, block_length=block_length))\n self.norm_layers_1.append(LayerNorm(hidden_channels))\n self.ffn_layers.append(\n FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))\n self.norm_layers_2.append(LayerNorm(hidden_channels))\n if pre_ln:\n self.last_ln = LayerNorm(hidden_channels)\n\n def forward(self, x, x_mask, attn_mask=1):\n if isinstance(attn_mask, torch.Tensor):\n attn_mask = attn_mask[:, None]\n attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) * attn_mask\n for i in range(self.n_layers):\n x = x * x_mask\n x_ = x\n if self.pre_ln:\n x = self.norm_layers_1[i](x)\n y = self.attn_layers[i](x, x, attn_mask)\n y = self.drop(y)\n x = x_ + y\n if not self.pre_ln:\n x = self.norm_layers_1[i](x)\n\n x_ = x\n if self.pre_ln:\n x = self.norm_layers_2[i](x)\n y = self.ffn_layers[i](x, x_mask)\n y = self.drop(y)\n x = x_ + y\n if not self.pre_ln:\n x = self.norm_layers_2[i](x)\n if self.pre_ln:\n x = self.last_ln(x)\n x = x * x_mask\n return x\n\n\nclass MultiHeadAttention(nn.Module):\n def __init__(self, channels, out_channels, n_heads, window_size=None, heads_share=True, p_dropout=0.,\n block_length=None, proximal_bias=False, proximal_init=False):\n super().__init__()\n assert channels % n_heads == 0\n\n self.channels = channels\n self.out_channels = out_channels\n self.n_heads = n_heads\n self.window_size = window_size\n self.heads_share = heads_share\n self.block_length = block_length\n self.proximal_bias = proximal_bias\n self.p_dropout = p_dropout\n self.attn = None\n\n self.k_channels = channels // n_heads\n self.conv_q = nn.Conv1d(channels, channels, 1)\n self.conv_k = nn.Conv1d(channels, channels, 1)\n self.conv_v = nn.Conv1d(channels, channels, 1)\n if window_size is not None:\n n_heads_rel = 1 if heads_share else n_heads\n rel_stddev = self.k_channels ** -0.5\n self.emb_rel_k = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)\n self.emb_rel_v = nn.Parameter(torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) * rel_stddev)\n self.conv_o = nn.Conv1d(channels, out_channels, 1)\n self.drop = nn.Dropout(p_dropout)\n\n nn.init.xavier_uniform_(self.conv_q.weight)\n nn.init.xavier_uniform_(self.conv_k.weight)\n if proximal_init:\n self.conv_k.weight.data.copy_(self.conv_q.weight.data)\n self.conv_k.bias.data.copy_(self.conv_q.bias.data)\n nn.init.xavier_uniform_(self.conv_v.weight)\n\n def forward(self, x, c, attn_mask=None):\n q = self.conv_q(x)\n k = self.conv_k(c)\n v = self.conv_v(c)\n\n x, self.attn = self.attention(q, k, v, mask=attn_mask)\n\n x = self.conv_o(x)\n return x\n\n def attention(self, query, key, value, mask=None):\n # reshape [b, d, t] -> [b, n_h, t, d_k]\n b, d, t_s, t_t = (*key.size(), query.size(2))\n query = query.view(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)\n key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)\n value = value.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)\n\n scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.k_channels)\n if self.window_size is not None:\n assert t_s == t_t, \"Relative attention is only available for self-attention.\"\n key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)\n rel_logits = self._matmul_with_relative_keys(query, key_relative_embeddings)\n rel_logits = self._relative_position_to_absolute_position(rel_logits)\n scores_local = rel_logits / math.sqrt(self.k_channels)\n scores = scores + scores_local\n if self.proximal_bias:\n assert t_s == t_t, \"Proximal bias is only available for self-attention.\"\n scores = scores + self._attention_bias_proximal(t_s).to(device=scores.device, dtype=scores.dtype)\n if mask is not None:\n scores = scores.masked_fill(mask == 0, -1e4)\n if self.block_length is not None:\n block_mask = torch.ones_like(scores).triu(-self.block_length).tril(self.block_length)\n scores = scores * block_mask + -1e4 * (1 - block_mask)\n p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s]\n p_attn = self.drop(p_attn)\n output = torch.matmul(p_attn, value)\n if self.window_size is not None:\n relative_weights = self._absolute_position_to_relative_position(p_attn)\n value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)\n output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)\n output = output.transpose(2, 3).contiguous().view(b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t]\n return output, p_attn\n\n def _matmul_with_relative_values(self, x, y):\n \"\"\"\n x: [b, h, l, m]\n y: [h or 1, m, d]\n ret: [b, h, l, d]\n \"\"\"\n ret = torch.matmul(x, y.unsqueeze(0))\n return ret\n\n def _matmul_with_relative_keys(self, x, y):\n \"\"\"\n x: [b, h, l, d]\n y: [h or 1, m, d]\n ret: [b, h, l, m]\n \"\"\"\n ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1))\n return ret\n\n def _get_relative_embeddings(self, relative_embeddings, length):\n max_relative_position = 2 * self.window_size + 1\n # Pad first before slice to avoid using cond ops.\n pad_length = max(length - (self.window_size + 1), 0)\n slice_start_position = max((self.window_size + 1) - length, 0)\n slice_end_position = slice_start_position + 2 * length - 1\n if pad_length > 0:\n padded_relative_embeddings = F.pad(\n relative_embeddings,\n convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))\n else:\n padded_relative_embeddings = relative_embeddings\n used_relative_embeddings = padded_relative_embeddings[:, slice_start_position:slice_end_position]\n return used_relative_embeddings\n\n def _relative_position_to_absolute_position(self, x):\n \"\"\"\n x: [b, h, l, 2*l-1]\n ret: [b, h, l, l]\n \"\"\"\n batch, heads, length, _ = x.size()\n # Concat columns of pad to shift from relative to absolute indexing.\n x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, 1]]))\n\n # Concat extra elements so to add up to shape (len+1, 2*len-1).\n x_flat = x.view([batch, heads, length * 2 * length])\n x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [0, length - 1]]))\n\n # Reshape and slice out the padded elements.\n x_final = x_flat.view([batch, heads, length + 1, 2 * length - 1])[:, :, :length, length - 1:]\n return x_final\n\n def _absolute_position_to_relative_position(self, x):\n \"\"\"\n x: [b, h, l, l]\n ret: [b, h, l, 2*l-1]\n \"\"\"\n batch, heads, length, _ = x.size()\n # padd along column\n x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length - 1]]))\n x_flat = x.view([batch, heads, -1])\n # add 0's in the beginning that will skew the elements after reshape\n x_flat = F.pad(x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]]))\n x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:]\n return x_final\n\n def _attention_bias_proximal(self, length):\n \"\"\"Bias for self-attention to encourage attention to close positions.\n Args:\n length: an integer scalar.\n Returns:\n a Tensor with shape [1, 1, length, length]\n \"\"\"\n r = torch.arange(length, dtype=torch.float32)\n diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1)\n return torch.unsqueeze(torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0)\n\n\nclass FFN(nn.Module):\n def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None):\n super().__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.filter_channels = filter_channels\n self.kernel_size = kernel_size\n self.p_dropout = p_dropout\n self.activation = activation\n\n self.conv_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size // 2)\n self.conv_2 = nn.Conv1d(filter_channels, out_channels, 1)\n self.drop = nn.Dropout(p_dropout)\n\n def forward(self, x, x_mask):\n x = self.conv_1(x * x_mask)\n if self.activation == \"gelu\":\n x = x * torch.sigmoid(1.702 * x)\n else:\n x = torch.relu(x)\n x = self.drop(x)\n x = self.conv_2(x * x_mask)\n return x * x_mask\n\n\nclass LayerNorm(nn.Module):\n def __init__(self, channels, eps=1e-4):\n super().__init__()\n self.channels = channels\n self.eps = eps\n\n self.gamma = nn.Parameter(torch.ones(channels))\n self.beta = nn.Parameter(torch.zeros(channels))\n\n def forward(self, x):\n n_dims = len(x.shape)\n mean = torch.mean(x, 1, keepdim=True)\n variance = torch.mean((x - mean) ** 2, 1, keepdim=True)\n\n x = (x - mean) * torch.rsqrt(variance + self.eps)\n\n shape = [1, -1] + [1] * (n_dims - 2)\n x = x * self.gamma.view(*shape) + self.beta.view(*shape)\n return x\n\n\nclass ConvReluNorm(nn.Module):\n def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):\n super().__init__()\n self.in_channels = in_channels\n self.hidden_channels = hidden_channels\n self.out_channels = out_channels\n self.kernel_size = kernel_size\n self.n_layers = n_layers\n self.p_dropout = p_dropout\n assert n_layers > 1, \"Number of layers should be larger than 0.\"\n\n self.conv_layers = nn.ModuleList()\n self.norm_layers = nn.ModuleList()\n self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2))\n self.norm_layers.append(LayerNorm(hidden_channels))\n self.relu_drop = nn.Sequential(\n nn.ReLU(),\n nn.Dropout(p_dropout))\n for _ in range(n_layers - 1):\n self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2))\n self.norm_layers.append(LayerNorm(hidden_channels))\n self.proj = nn.Conv1d(hidden_channels, out_channels, 1)\n self.proj.weight.data.zero_()\n self.proj.bias.data.zero_()\n\n def forward(self, x, x_mask):\n x_org = x\n for i in range(self.n_layers):\n x = self.conv_layers[i](x * x_mask)\n x = self.norm_layers[i](x)\n x = self.relu_drop(x)\n x = x_org + self.proj(x)\n return x * x_mask\n\n\nclass RelTransformerEncoder(nn.Module):\n def __init__(self,\n n_vocab,\n out_channels,\n hidden_channels,\n filter_channels,\n n_heads,\n n_layers,\n kernel_size,\n p_dropout=0.0,\n window_size=4,\n block_length=None,\n in_channels=None,\n prenet=True,\n pre_ln=True,\n ):\n\n super().__init__()\n\n self.n_vocab = n_vocab\n self.out_channels = out_channels\n self.hidden_channels = hidden_channels\n self.filter_channels = filter_channels\n self.n_heads = n_heads\n self.n_layers = n_layers\n self.kernel_size = kernel_size\n self.p_dropout = p_dropout\n self.window_size = window_size\n self.block_length = block_length\n self.prenet = prenet\n if n_vocab > 0:\n self.emb = Embedding(n_vocab, hidden_channels, padding_idx=0)\n\n if prenet:\n if in_channels is None:\n in_channels = hidden_channels\n self.pre = ConvReluNorm(in_channels, in_channels, in_channels,\n kernel_size=5, n_layers=3, p_dropout=0)\n if in_channels is not None and in_channels != hidden_channels:\n self.encoder_inp_proj = nn.Conv1d(in_channels, hidden_channels, 1)\n self.encoder = Encoder(\n hidden_channels,\n filter_channels,\n n_heads,\n n_layers,\n kernel_size,\n p_dropout,\n window_size=window_size,\n block_length=block_length,\n pre_ln=pre_ln,\n )\n\n def forward(self, x, x_mask=None, other_embeds=0, attn_mask=1):\n if self.n_vocab > 0:\n x_lengths = (x > 0).long().sum(-1)\n x = self.emb(x) * math.sqrt(self.hidden_channels) # [b, t, h]\n else:\n x_lengths = (x.abs().sum(-1) > 0).long().sum(-1)\n x = x + other_embeds\n x = torch.transpose(x, 1, -1) # [b, h, t]\n x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)\n\n if self.prenet:\n x = self.pre(x, x_mask)\n self.prenet_out = x.transpose(1, 2)\n if hasattr(self, 'encoder_inp_proj'):\n x = self.encoder_inp_proj(x) * x_mask\n x = self.encoder(x, x_mask, attn_mask)\n return x.transpose(1, 2)\n" }, { "path": "tts/modules/ar_dur/commons/rot_transformer.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 math\nimport torch\nfrom typing import Optional, Tuple\nfrom torch import nn\nfrom torch.nn import Parameter, Linear\nfrom tts.modules.ar_dur.commons.layers import LayerNorm, Embedding\nfrom tts.modules.ar_dur.commons.transformer import TransformerFFNLayer, MultiheadAttention\nfrom tts.modules.ar_dur.commons.seq_utils import get_incremental_state, set_incremental_state, softmax, make_positions\nimport torch.nn.functional as F\n\nDEFAULT_MAX_SOURCE_POSITIONS = 3000\nDEFAULT_MAX_TARGET_POSITIONS = 3000\n\n\nclass SinusoidalPositionalEmbedding(nn.Module):\n \"\"\"This module produces sinusoidal positional embeddings of any length.\n\n Padding symbols are ignored.\n \"\"\"\n\n def __init__(self, embedding_dim, padding_idx, init_size=1024):\n super().__init__()\n self.embedding_dim = embedding_dim\n self.padding_idx = padding_idx\n self.weights = SinusoidalPositionalEmbedding.get_embedding(\n init_size,\n embedding_dim,\n padding_idx,\n )\n self.register_buffer('_float_tensor', torch.FloatTensor(1))\n\n @staticmethod\n def get_embedding(num_embeddings, embedding_dim, padding_idx=None):\n \"\"\"Build sinusoidal embeddings.\n\n This matches the implementation in tensor2tensor, but differs slightly\n from the description in Section 3.5 of \"Attention Is All You Need\".\n \"\"\"\n half_dim = embedding_dim // 2\n emb = math.log(10000) / (half_dim - 1)\n emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)\n emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)\n emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)\n if embedding_dim % 2 == 1:\n # zero pad\n emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)\n if padding_idx is not None:\n emb[padding_idx, :] = 0\n return emb\n\n def forward(self, input, incremental_state=None, timestep=None, positions=None, **kwargs):\n \"\"\"Input is expected to be of size [bsz x seqlen].\"\"\"\n bsz, seq_len = input.shape[:2]\n max_pos = self.padding_idx + 1 + seq_len\n if self.weights is None or max_pos > self.weights.size(0):\n # recompute/expand embeddings if needed\n self.weights = SinusoidalPositionalEmbedding.get_embedding(\n max_pos,\n self.embedding_dim,\n self.padding_idx,\n )\n self.weights = self.weights.to(self._float_tensor)\n\n if incremental_state is not None:\n # positions is the same for every token when decoding a single step\n pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len\n return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)\n\n positions = make_positions(input, self.padding_idx) if positions is None else positions\n return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()\n\n def max_positions(self):\n \"\"\"Maximum number of supported positions.\"\"\"\n return int(1e5) # an arbitrary large number\n\n\nclass RotaryEmbeddings(nn.Module):\n cos: torch.Tensor\n sin: torch.Tensor\n theta: torch.Tensor\n\n def __init__(\n self,\n width: int,\n *,\n seq_len: int = 40000,\n base: int = 10000,\n device: Optional[torch.device] = None,\n ):\n \"\"\"Rotary embeddings (Su et al., 2021) layer. The rotary embedding\n will be precomputed for up to 'seq _len' positions. The embedding\n will be recomputed when a longer sequence is found in the input.\n\n :param width:\n Rotary embedding dimensionality, must be even.\n :param seq_len:\n Number of positons to initially precompute.\n :param base:\n The base used for Θ_i, determines the cycle length of the\n embeddings.\n :param device: Device on which the module is to be initialized.\n \"\"\"\n super().__init__()\n\n if width % 2:\n raise ValueError(f\"Width of rotary embeddings must be even, was: {width}\")\n\n # Ignore allocations on the meta device as we don't persist our buffer,\n # i.e., we don't expect the backing tensor to be replaced with pretrained weights.\n if device is not None and device.type == \"meta\":\n device = None\n # Θ_i = 10000^(-2(i-1)/d)\n theta = torch.pow(\n base, -torch.arange(0, width, 2, dtype=torch.float, device=device) / width\n )\n self.register_buffer(\"theta\", theta, persistent=False)\n\n self._create_rotary_embed(width=width, length=seq_len)\n\n def _create_rotary_embed(self, *, width: int, length: int):\n # mΘ\n position = torch.arange(length, device=self.theta.device).unsqueeze(1)\n m_theta = position * self.theta.unsqueeze(0)\n\n # We apply both sin and cos twice (see Eq 15, 34), but the ordering\n # is changed for compatibility with most common implementations.\n m_theta = torch.cat([m_theta, m_theta], dim=-1)\n\n re_cos = m_theta.cos().view([length, width])\n re_sin = m_theta.sin().view([length, width])\n\n self.register_buffer(\"cos\", re_cos, persistent=False)\n self.register_buffer(\"sin\", re_sin, persistent=False)\n\n def _rotate(self, input: torch.Tensor):\n \"\"\"Rotate the input tensor by half of its innermost width.\n\n input (Tensor): array to rotate.\n RETURNS (Tensor): rotated array.\n\n Shapes:\n input - (..., width)\n output - (..., width)\n \"\"\"\n half_idx = input.shape[-1] // 2\n input_1 = -input[..., half_idx:]\n input_2 = input[..., :half_idx]\n return torch.cat([input_1, input_2], dim=-1)\n\n def forward(self, input: torch.Tensor, *, positions: Optional[torch.Tensor] = None):\n \"\"\"\n Apply rotary embeddings to an array.\n\n :param input: Array to apply the rotary embeddings to.\n :param positions: positions of the inputs. If no positions are\n provided, they are assumed to be [0, seq_len).\n :return: Array with the rotary embeddings applied.\n\n Shapes:\n input - (batch_size, num_heads, seq_len, width_per_head)\n positions - (batch_size, seq_len)\n output - (batch_size, num_heads, seq_len, width_per_head)\n \"\"\"\n batch_size, _, seq_len, width = input.shape\n\n if positions is None:\n # Fastpath: positions from [0..seq_len), avoid indexing.\n if self.cos.size(-2) < seq_len:\n self._create_rotary_embed(width=width, length=seq_len)\n rot_cos = self.cos[:seq_len, :].view(1, 1, seq_len, width)\n rot_sin = self.sin[:seq_len, :].view(1, 1, seq_len, width)\n else:\n max_len = int(positions.max()) + 1\n if self.cos.size(-2) < max_len:\n self._create_rotary_embed(width=width, length=max_len)\n\n # Flatten positions to index cos/sin arrays, then unflatten.\n #\n # Example shapes:\n #\n # positions_flat - (batch_size * seq_len)\n # self.cos - (max_len, width)\n # rot_cos - (batch_size, seq_len, width)\n positions_flat = positions.view(-1)\n rot_cos = self.cos[positions_flat].view(batch_size, 1, seq_len, width)\n rot_sin = self.sin[positions_flat].view(batch_size, 1, seq_len, width)\n\n # Eq 34 with ordering changed for compatibility.\n return rot_cos * input + rot_sin * self._rotate(input)\n\n\nclass RotMultiheadAttention(MultiheadAttention):\n def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=True,\n add_bias_kv=False, add_zero_attn=False, self_attention=False,\n encoder_decoder_attention=False):\n super().__init__(embed_dim, num_heads, kdim=kdim, vdim=vdim, dropout=dropout, bias=bias,\n add_bias_kv=add_bias_kv, add_zero_attn=add_zero_attn, self_attention=self_attention,\n encoder_decoder_attention=encoder_decoder_attention)\n self.rotary_embeds = RotaryEmbeddings(width=embed_dim // num_heads)\n\n def forward(\n self,\n query, key, value,\n spk_pos_ids_flat=None,\n key_padding_mask=None,\n incremental_state=None,\n need_weights=True,\n static_kv=False,\n attn_mask=None,\n before_softmax=False,\n need_head_weights=False,\n enc_dec_attn_constraint_mask=None,\n reset_attn_weight=None\n ):\n \"\"\"Input shape: Time x Batch x Channel\n\n Args:\n key_padding_mask (ByteTensor, optional): mask to exclude\n keys that are pads, of shape `(batch, src_len)`, where\n padding elements are indicated by 1s.\n need_weights (bool, optional): return the attention weights,\n averaged over heads (default: False).\n attn_mask (ByteTensor, optional): typically used to\n implement causal attention, where the mask prevents the\n attention from looking forward in time (default: None).\n before_softmax (bool, optional): return the raw attention\n weights and values before the attention softmax.\n need_head_weights (bool, optional): return the attention\n weights for each head. Implies *need_weights*. Default:\n return the average attention weights over all heads.\n \"\"\"\n if need_head_weights:\n need_weights = True\n\n tgt_len, bsz, embed_dim = query.size()\n assert embed_dim == self.embed_dim\n assert list(query.size()) == [tgt_len, bsz, embed_dim]\n\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_key' in saved_state:\n # previous time steps are cached - no need to recompute\n # key and value if they are static\n if static_kv:\n assert self.encoder_decoder_attention and not self.self_attention\n key = value = None\n else:\n saved_state = None\n\n if self.self_attention:\n # self-attention\n q, k, v = self.in_proj_qkv(query)\n elif self.encoder_decoder_attention:\n # encoder-decoder attention\n q = self.in_proj_q(query)\n if key is None:\n assert value is None\n k = v = None\n else:\n k = self.in_proj_k(key)\n v = self.in_proj_v(key)\n else:\n q = self.in_proj_q(query)\n k = self.in_proj_k(key)\n v = self.in_proj_v(value)\n q = q * self.scaling\n\n if self.bias_k is not None:\n assert self.bias_v is not None\n k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])\n v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])\n if attn_mask is not None:\n attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)\n if key_padding_mask is not None:\n key_padding_mask = torch.cat(\n [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)], dim=1)\n\n q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if k is not None:\n k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if v is not None:\n v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n\n # Apply rot embedding and store incremental_state\n q = self.rotary_embeds(q[None, :], positions=spk_pos_ids_flat)[0]\n if saved_state is not None:\n # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)\n if 'prev_key' in saved_state:\n prev_key = saved_state['prev_key'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n k = prev_key\n else:\n k = torch.cat((prev_key, k), dim=1)\n if 'prev_value' in saved_state:\n prev_value = saved_state['prev_value'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n v = prev_value\n else:\n v = torch.cat((prev_value, v), dim=1)\n saved_state['prev_key'], saved_state['prev_value'] = k.view(bsz, self.num_heads, -1, self.head_dim), v.view(\n bsz, self.num_heads, -1, self.head_dim)\n self._set_input_buffer(incremental_state, saved_state)\n if incremental_state is not None:\n key_pos = torch.arange(k.shape[-2], device=q.device).unsqueeze(0)\n else:\n key_pos = spk_pos_ids_flat\n k = self.rotary_embeds(k[None, :], positions=key_pos)[0]\n\n src_len = k.size(1)\n\n # This is part of a workaround to get around fork/join parallelism\n # not supporting Optional types.\n if key_padding_mask is not None and key_padding_mask.shape == torch.Size([]):\n key_padding_mask = None\n\n if key_padding_mask is not None:\n assert key_padding_mask.size(0) == bsz\n assert key_padding_mask.size(1) == src_len\n\n if self.add_zero_attn:\n src_len += 1\n k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)\n v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)\n if attn_mask is not None:\n attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)\n if key_padding_mask is not None:\n key_padding_mask = torch.cat(\n [key_padding_mask, torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask)], dim=1)\n\n attn_weights = torch.bmm(q, k.transpose(1, 2))\n attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)\n assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]\n\n if attn_mask is not None:\n if len(attn_mask.shape) == 2:\n attn_mask = attn_mask.unsqueeze(0)\n elif len(attn_mask.shape) == 3:\n attn_mask = attn_mask[:, None].repeat([1, self.num_heads, 1, 1]).reshape(\n bsz * self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights + attn_mask\n\n if enc_dec_attn_constraint_mask is not None: # bs x head x L_kv\n attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights.masked_fill(\n enc_dec_attn_constraint_mask.unsqueeze(2).bool(),\n -1e8,\n )\n attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)\n\n if key_padding_mask is not None:\n # don't attend to padding symbols\n attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights.masked_fill(\n key_padding_mask.unsqueeze(1).unsqueeze(2),\n -1e8,\n )\n attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)\n\n attn_logits = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n\n if before_softmax:\n return attn_weights, v\n\n attn_weights_float = softmax(attn_weights, dim=-1)\n attn_weights = attn_weights_float.type_as(attn_weights)\n attn_probs = F.dropout(attn_weights_float.type_as(attn_weights), p=self.dropout, training=self.training)\n\n if reset_attn_weight is not None:\n if reset_attn_weight:\n self.last_attn_probs = attn_probs.detach()\n else:\n assert self.last_attn_probs is not None\n attn_probs = self.last_attn_probs\n attn = torch.bmm(attn_probs, v)\n assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]\n attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)\n attn = self.out_proj(attn)\n\n if need_weights:\n attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)\n if not need_head_weights:\n # average attention weights over heads\n attn_weights = attn_weights.mean(dim=0)\n else:\n attn_weights = None\n\n return attn, (attn_weights, attn_logits)\n\n\nclass RotMultiheadAttention2(MultiheadAttention):\n def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=True,\n add_bias_kv=False, add_zero_attn=False, self_attention=False,\n encoder_decoder_attention=False):\n super().__init__(embed_dim, num_heads, kdim=kdim, vdim=vdim, dropout=dropout, bias=bias,\n add_bias_kv=add_bias_kv, add_zero_attn=add_zero_attn, self_attention=self_attention,\n encoder_decoder_attention=encoder_decoder_attention)\n self.rotary_embeds = RotaryEmbeddings(width=embed_dim // num_heads)\n\n def forward(\n self,\n query, key, value,\n spk_pos_ids_flat=None,\n key_padding_mask=None,\n incremental_state=None,\n need_weights=True,\n static_kv=False,\n attn_mask=None,\n before_softmax=False,\n need_head_weights=False,\n enc_dec_attn_constraint_mask=None,\n reset_attn_weight=None\n ):\n \"\"\"Input shape: Time x Batch x Channel\n\n Args:\n key_padding_mask (ByteTensor, optional): mask to exclude\n keys that are pads, of shape `(batch, src_len)`, where\n padding elements are indicated by 1s.\n need_weights (bool, optional): return the attention weights,\n averaged over heads (default: False).\n attn_mask (ByteTensor, optional): typically used to\n implement causal attention, where the mask prevents the\n attention from looking forward in time (default: None).\n before_softmax (bool, optional): return the raw attention\n weights and values before the attention softmax.\n need_head_weights (bool, optional): return the attention\n weights for each head. Implies *need_weights*. Default:\n return the average attention weights over all heads.\n \"\"\"\n if need_head_weights:\n need_weights = True\n\n tgt_len, bsz, embed_dim = query.size()\n assert embed_dim == self.embed_dim\n assert list(query.size()) == [tgt_len, bsz, embed_dim]\n\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_key' in saved_state:\n # previous time steps are cached - no need to recompute\n # key and value if they are static\n if static_kv:\n assert self.encoder_decoder_attention and not self.self_attention\n key = value = None\n else:\n saved_state = None\n\n if self.self_attention:\n # self-attention\n q, k, v = self.in_proj_qkv(query)\n elif self.encoder_decoder_attention:\n # encoder-decoder attention\n q = self.in_proj_q(query)\n if key is None:\n assert value is None\n k = v = None\n else:\n k = self.in_proj_k(key)\n v = self.in_proj_v(key)\n else:\n q = self.in_proj_q(query)\n k = self.in_proj_k(key)\n v = self.in_proj_v(value)\n\n if self.bias_k is not None:\n assert self.bias_v is not None\n k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])\n v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])\n if attn_mask is not None:\n attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)\n if key_padding_mask is not None:\n key_padding_mask = torch.cat(\n [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)], dim=1)\n\n q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if k is not None:\n k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if v is not None:\n v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n\n # Apply rot embedding and store incremental_state\n q = self.rotary_embeds(q[None, :], positions=spk_pos_ids_flat)[0]\n if saved_state is not None:\n # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)\n if 'prev_key' in saved_state:\n prev_key = saved_state['prev_key'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n k = prev_key\n else:\n k = torch.cat((prev_key, k), dim=1)\n if 'prev_value' in saved_state:\n prev_value = saved_state['prev_value'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n v = prev_value\n else:\n v = torch.cat((prev_value, v), dim=1)\n saved_state['prev_key'], saved_state['prev_value'] = k.view(bsz, self.num_heads, -1, self.head_dim), v.view(\n bsz, self.num_heads, -1, self.head_dim)\n self._set_input_buffer(incremental_state, saved_state)\n key_pos = torch.arange(k.shape[-2], device=q.device).unsqueeze(0)\n k = self.rotary_embeds(k[None, :], positions=key_pos)[0]\n\n src_len = k.size(1)\n\n # This is part of a workaround to get around fork/join parallelism\n # not supporting Optional types.\n if key_padding_mask is not None and key_padding_mask.shape == torch.Size([]):\n key_padding_mask = None\n\n if key_padding_mask is not None:\n assert key_padding_mask.size(0) == bsz\n assert key_padding_mask.size(1) == src_len\n\n if attn_mask is not None:\n if len(attn_mask.shape) == 2:\n attn_mask = attn_mask.unsqueeze(0)\n elif len(attn_mask.shape) == 3:\n attn_mask = attn_mask[:, None].repeat([1, self.num_heads, 1, 1]).reshape(\n bsz * self.num_heads, tgt_len, src_len)\n attn = torch.nn.functional.scaled_dot_product_attention(\n q, k, v, attn_mask=attn_mask, dropout_p=0, is_causal=False)\n assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]\n attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)\n attn_logits = None\n attn_weights = None\n return attn, (attn_weights, attn_logits)\n\n\nclass RotDecSALayer(nn.Module):\n def __init__(self, c, num_heads, dropout, attention_dropout=0.1, relu_dropout=0.1,\n kernel_size=9, ffn_hidden_size=1024, act='gelu', post_ln=False, bias=True):\n super().__init__()\n self.c = c\n self.dropout = dropout\n self.layer_norm1 = LayerNorm(c)\n self.self_attn = RotMultiheadAttention(\n c, num_heads, self_attention=True, dropout=attention_dropout, bias=False\n )\n self.layer_norm2 = LayerNorm(c)\n self.ffn = TransformerFFNLayer(\n c, ffn_hidden_size, padding='LEFT', kernel_size=kernel_size,\n dropout=relu_dropout, act=act, bias=bias)\n self.post_ln = post_ln\n\n def forward(\n self,\n x,\n encoder_out=None,\n encoder_padding_mask=None,\n incremental_state=None,\n self_attn_mask=None,\n self_attn_padding_mask=None,\n attn_out=None,\n reset_attn_weight=None,\n spk_pos_ids_flat=None,\n **kwargs,\n ):\n layer_norm_training = kwargs.get('layer_norm_training', None)\n if layer_norm_training is not None:\n self.layer_norm1.training = layer_norm_training\n self.layer_norm2.training = layer_norm_training\n residual = x\n if not self.post_ln:\n x = self.layer_norm1(x)\n\n x, (attn_weights, _) = self.self_attn(\n query=x,\n key=x,\n value=x,\n key_padding_mask=self_attn_padding_mask,\n incremental_state=incremental_state,\n attn_mask=self_attn_mask,\n spk_pos_ids_flat=spk_pos_ids_flat\n )\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n if self.post_ln:\n x = self.layer_norm1(x)\n\n residual = x\n if not self.post_ln:\n x = self.layer_norm2(x)\n x = self.ffn(x, incremental_state=incremental_state)\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n if self.post_ln:\n x = self.layer_norm2(x)\n return x, attn_weights\n\n def clear_buffer(self, input, encoder_out=None, encoder_padding_mask=None, incremental_state=None):\n self.encoder_attn.clear_buffer(incremental_state)\n self.ffn.clear_buffer(incremental_state)\n\n def set_buffer(self, name, tensor, incremental_state):\n return set_incremental_state(self, incremental_state, name, tensor)\n\n\nclass RotDecSALayer2(RotDecSALayer):\n def __init__(self, c, num_heads, dropout, attention_dropout=0.1, relu_dropout=0.1, kernel_size=9,\n ffn_hidden_size=1024, act='gelu', post_ln=False):\n super().__init__(c, num_heads, dropout, attention_dropout, relu_dropout, kernel_size, ffn_hidden_size, act,\n post_ln)\n self.self_attn = RotMultiheadAttention2(\n c, num_heads, self_attention=True, dropout=attention_dropout, bias=False\n )\n\n\nclass RotTransformerDecoderLayer(nn.Module):\n def __init__(self, hidden_size, dropout, kernel_size=9, num_heads=8, ffn_hidden_size=1024, post_ln=False,\n op_version=1, bias=True):\n super().__init__()\n self.hidden_size = hidden_size\n self.dropout = dropout\n self.num_heads = num_heads\n if op_version == 1:\n self.op = RotDecSALayer(\n hidden_size, num_heads, dropout=dropout,\n attention_dropout=0.0, relu_dropout=dropout,\n kernel_size=kernel_size, ffn_hidden_size=ffn_hidden_size,\n post_ln=post_ln, bias=bias)\n else:\n self.op = RotDecSALayer2(\n hidden_size, num_heads, dropout=dropout,\n attention_dropout=0.0, relu_dropout=dropout,\n kernel_size=kernel_size, ffn_hidden_size=ffn_hidden_size,\n post_ln=post_ln)\n\n def forward(self, x, **kwargs):\n return self.op(x, **kwargs)\n\n def clear_buffer(self, *args):\n return self.op.clear_buffer(*args)\n\n def set_buffer(self, *args):\n return self.op.set_buffer(*args)\n" }, { "path": "tts/modules/ar_dur/commons/seq_utils.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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\nfrom collections import defaultdict\nimport torch\nimport torch.nn.functional as F\n\n\ndef make_positions(tensor, padding_idx):\n \"\"\"Replace non-padding symbols with their position numbers.\n\n Position numbers begin at padding_idx+1. Padding symbols are ignored.\n \"\"\"\n # The series of casts and type-conversions here are carefully\n # balanced to both work with ONNX export and XLA. In particular XLA\n # prefers ints, cumsum defaults to output longs, and ONNX doesn't know\n # how to handle the dtype kwarg in cumsum.\n mask = tensor.ne(padding_idx).int()\n return (\n torch.cumsum(mask, dim=1).type_as(mask) * mask\n ).long() + padding_idx\n\n\ndef softmax(x, dim):\n return F.softmax(x, dim=dim, dtype=torch.float32)\n\n\ndef sequence_mask(lengths, maxlen=None, dtype=torch.bool):\n if maxlen is None:\n maxlen = lengths.max()\n mask = ~(torch.ones((len(lengths), maxlen)).to(lengths.device).cumsum(dim=1).t() > lengths).t()\n mask.type(dtype)\n return mask\n\n\ndef weights_nonzero_speech(target):\n # target : B x T x mel\n # Assign weight 1.0 to all labels except for padding (id=0).\n dim = target.size(-1)\n return target.abs().sum(-1, keepdim=True).ne(0).float().repeat(1, 1, dim)\n\n\nINCREMENTAL_STATE_INSTANCE_ID = defaultdict(lambda: 0)\n\n\ndef _get_full_incremental_state_key(module_instance, key):\n module_name = module_instance.__class__.__name__\n\n # assign a unique ID to each module instance, so that incremental state is\n # not shared across module instances\n if not hasattr(module_instance, '_instance_id'):\n INCREMENTAL_STATE_INSTANCE_ID[module_name] += 1\n module_instance._instance_id = INCREMENTAL_STATE_INSTANCE_ID[module_name]\n\n return '{}.{}.{}'.format(module_name, module_instance._instance_id, key)\n\n\ndef get_incremental_state(module, incremental_state, key):\n \"\"\"Helper for getting incremental state for an nn.Module.\"\"\"\n full_key = _get_full_incremental_state_key(module, key)\n if incremental_state is None or full_key not in incremental_state:\n return None\n return incremental_state[full_key]\n\n\ndef set_incremental_state(module, incremental_state, key, value):\n \"\"\"Helper for setting incremental state for an nn.Module.\"\"\"\n if incremental_state is not None:\n full_key = _get_full_incremental_state_key(module, key)\n incremental_state[full_key] = value\n\n\ndef fill_with_neg_inf(t):\n \"\"\"FP16-compatible function that fills a tensor with -inf.\"\"\"\n return t.float().fill_(float('-inf')).type_as(t)\n\n\ndef fill_with_neg_inf2(t):\n \"\"\"FP16-compatible function that fills a tensor with -inf.\"\"\"\n return t.float().fill_(-1e8).type_as(t)\n\n\ndef select_attn(attn_logits, type='best'):\n \"\"\"\n\n :param attn_logits: [n_layers, B, n_head, T_sp, T_txt]\n :return:\n \"\"\"\n encdec_attn = torch.stack(attn_logits, 0).transpose(1, 2)\n # [n_layers * n_head, B, T_sp, T_txt]\n encdec_attn = (encdec_attn.reshape([-1, *encdec_attn.shape[2:]])).softmax(-1)\n if type == 'best':\n indices = encdec_attn.max(-1).values.sum(-1).argmax(0)\n encdec_attn = encdec_attn.gather(\n 0, indices[None, :, None, None].repeat(1, 1, encdec_attn.size(-2), encdec_attn.size(-1)))[0]\n return encdec_attn\n elif type == 'mean':\n return encdec_attn.mean(0)\n\n\ndef make_pad_mask(lengths, xs=None, length_dim=-1):\n \"\"\"Make mask tensor containing indices of padded part.\n Args:\n lengths (LongTensor or List): Batch of lengths (B,).\n xs (Tensor, optional): The reference tensor.\n If set, masks will be the same shape as this tensor.\n length_dim (int, optional): Dimension indicator of the above tensor.\n See the example.\n Returns:\n Tensor: Mask tensor containing indices of padded part.\n dtype=torch.uint8 in PyTorch 1.2-\n dtype=torch.bool in PyTorch 1.2+ (including 1.2)\n Examples:\n With only lengths.\n >>> lengths = [5, 3, 2]\n >>> make_non_pad_mask(lengths)\n masks = [[0, 0, 0, 0 ,0],\n [0, 0, 0, 1, 1],\n [0, 0, 1, 1, 1]]\n With the reference tensor.\n >>> xs = torch.zeros((3, 2, 4))\n >>> make_pad_mask(lengths, xs)\n tensor([[[0, 0, 0, 0],\n [0, 0, 0, 0]],\n [[0, 0, 0, 1],\n [0, 0, 0, 1]],\n [[0, 0, 1, 1],\n [0, 0, 1, 1]]], dtype=torch.uint8)\n >>> xs = torch.zeros((3, 2, 6))\n >>> make_pad_mask(lengths, xs)\n tensor([[[0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1]],\n [[0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1]],\n [[0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)\n With the reference tensor and dimension indicator.\n >>> xs = torch.zeros((3, 6, 6))\n >>> make_pad_mask(lengths, xs, 1)\n tensor([[[0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1]],\n [[0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1]],\n [[0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1]]], dtype=torch.uint8)\n >>> make_pad_mask(lengths, xs, 2)\n tensor([[[0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1],\n [0, 0, 0, 0, 0, 1]],\n [[0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1],\n [0, 0, 0, 1, 1, 1]],\n [[0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1],\n [0, 0, 1, 1, 1, 1]]], dtype=torch.uint8)\n \"\"\"\n if length_dim == 0:\n raise ValueError(\"length_dim cannot be 0: {}\".format(length_dim))\n\n if not isinstance(lengths, list):\n lengths = lengths.tolist()\n bs = int(len(lengths))\n if xs is None:\n maxlen = int(max(lengths))\n else:\n maxlen = xs.size(length_dim)\n\n seq_range = torch.arange(0, maxlen, dtype=torch.int64)\n seq_range_expand = seq_range.unsqueeze(0).expand(bs, maxlen)\n seq_length_expand = seq_range_expand.new(lengths).unsqueeze(-1)\n mask = seq_range_expand >= seq_length_expand\n\n if xs is not None:\n assert xs.size(0) == bs, (xs.size(0), bs)\n\n if length_dim < 0:\n length_dim = xs.dim() + length_dim\n # ind = (:, None, ..., None, :, , None, ..., None)\n ind = tuple(\n slice(None) if i in (0, length_dim) else None for i in range(xs.dim())\n )\n mask = mask[ind].expand_as(xs).to(xs.device)\n return mask\n\n\ndef make_non_pad_mask(lengths, xs=None, length_dim=-1):\n \"\"\"Make mask tensor containing indices of non-padded part.\n Args:\n lengths (LongTensor or List): Batch of lengths (B,).\n xs (Tensor, optional): The reference tensor.\n If set, masks will be the same shape as this tensor.\n length_dim (int, optional): Dimension indicator of the above tensor.\n See the example.\n Returns:\n ByteTensor: mask tensor containing indices of padded part.\n dtype=torch.uint8 in PyTorch 1.2-\n dtype=torch.bool in PyTorch 1.2+ (including 1.2)\n Examples:\n With only lengths.\n >>> lengths = [5, 3, 2]\n >>> make_non_pad_mask(lengths)\n masks = [[1, 1, 1, 1 ,1],\n [1, 1, 1, 0, 0],\n [1, 1, 0, 0, 0]]\n With the reference tensor.\n >>> xs = torch.zeros((3, 2, 4))\n >>> make_non_pad_mask(lengths, xs)\n tensor([[[1, 1, 1, 1],\n [1, 1, 1, 1]],\n [[1, 1, 1, 0],\n [1, 1, 1, 0]],\n [[1, 1, 0, 0],\n [1, 1, 0, 0]]], dtype=torch.uint8)\n >>> xs = torch.zeros((3, 2, 6))\n >>> make_non_pad_mask(lengths, xs)\n tensor([[[1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0]],\n [[1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0]],\n [[1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)\n With the reference tensor and dimension indicator.\n >>> xs = torch.zeros((3, 6, 6))\n >>> make_non_pad_mask(lengths, xs, 1)\n tensor([[[1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [0, 0, 0, 0, 0, 0]],\n [[1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0]],\n [[1, 1, 1, 1, 1, 1],\n [1, 1, 1, 1, 1, 1],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0],\n [0, 0, 0, 0, 0, 0]]], dtype=torch.uint8)\n >>> make_non_pad_mask(lengths, xs, 2)\n tensor([[[1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0],\n [1, 1, 1, 1, 1, 0]],\n [[1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0],\n [1, 1, 1, 0, 0, 0]],\n [[1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0],\n [1, 1, 0, 0, 0, 0]]], dtype=torch.uint8)\n \"\"\"\n return ~make_pad_mask(lengths, xs, length_dim)\n\n\ndef get_mask_from_lengths(lengths):\n max_len = torch.max(lengths).item()\n ids = torch.arange(0, max_len).to(lengths.device)\n mask = (ids < lengths.unsqueeze(1)).bool()\n return mask\n\n\ndef group_hidden_by_segs(h, seg_ids, max_len):\n \"\"\"\n\n :param h: [B, T, H]\n :param seg_ids: [B, T]\n :return: h_ph: [B, T_ph, H]\n \"\"\"\n B, T, H = h.shape\n h_gby_segs = h.new_zeros([B, max_len + 1, H]).scatter_add_(1, seg_ids[:, :, None].repeat([1, 1, H]), h)\n all_ones = h.new_ones(h.shape[:2])\n cnt_gby_segs = h.new_zeros([B, max_len + 1]).scatter_add_(1, seg_ids, all_ones).contiguous()\n h_gby_segs = h_gby_segs[:, 1:]\n cnt_gby_segs = cnt_gby_segs[:, 1:]\n h_gby_segs = h_gby_segs / torch.clamp(cnt_gby_segs[:, :, None], min=1)\n return h_gby_segs, cnt_gby_segs\n\ndef expand_by_repeat_times(source_encoding, lengths):\n \"\"\"\n source_encoding: [T, C]\n lengths, list of int, [T,], how many times each token should repeat\n return:\n expanded_encoding: [T_expand, C]\n \"\"\"\n hid_dim = source_encoding.shape[1]\n out2source = []\n for i, length in enumerate(lengths):\n out2source += [i for _ in range(length)]\n out2source = torch.LongTensor(out2source).to(source_encoding.device)\n out2source_ = out2source[:, None].repeat([1, hid_dim])\n expanded_encoding = torch.gather(source_encoding, 0, out2source_) # [B, T, H]\n return expanded_encoding\n\n\ndef expand_word2ph(word_encoding, ph2word):\n word_encoding = F.pad(word_encoding,[0,0,1,0])\n ph2word_ = ph2word[:, :, None].repeat([1, 1, word_encoding.shape[-1]])\n out = torch.gather(word_encoding, 1, ph2word_) # [B, T, H]\n return out\n" }, { "path": "tts/modules/ar_dur/commons/transformer.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 math\nimport torch\nfrom torch import nn\nfrom torch.nn import Parameter, Linear\nfrom tts.modules.ar_dur.commons.layers import LayerNorm, Embedding\nfrom tts.modules.ar_dur.commons.seq_utils import get_incremental_state, set_incremental_state, softmax, make_positions\nimport torch.nn.functional as F\n\nDEFAULT_MAX_SOURCE_POSITIONS = 3000\nDEFAULT_MAX_TARGET_POSITIONS = 3000\n\n\nclass SinusoidalPositionalEmbedding(nn.Module):\n \"\"\"This module produces sinusoidal positional embeddings of any length.\n\n Padding symbols are ignored.\n \"\"\"\n\n def __init__(self, embedding_dim, padding_idx, init_size=1024):\n super().__init__()\n self.embedding_dim = embedding_dim\n self.padding_idx = padding_idx\n self.weights = SinusoidalPositionalEmbedding.get_embedding(\n init_size,\n embedding_dim,\n padding_idx,\n )\n self.register_buffer('_float_tensor', torch.FloatTensor(1))\n\n @staticmethod\n def get_embedding(num_embeddings, embedding_dim, padding_idx=None):\n \"\"\"Build sinusoidal embeddings.\n\n This matches the implementation in tensor2tensor, but differs slightly\n from the description in Section 3.5 of \"Attention Is All You Need\".\n \"\"\"\n half_dim = embedding_dim // 2\n emb = math.log(10000) / (half_dim - 1)\n emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)\n emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)\n emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)\n if embedding_dim % 2 == 1:\n # zero pad\n emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)\n if padding_idx is not None:\n emb[padding_idx, :] = 0\n return emb\n\n def forward(self, input, incremental_state=None, timestep=None, positions=None, **kwargs):\n \"\"\"Input is expected to be of size [bsz x seqlen].\"\"\"\n bsz, seq_len = input.shape[:2]\n max_pos = self.padding_idx + 1 + seq_len\n if self.weights is None or max_pos > self.weights.size(0):\n # recompute/expand embeddings if needed\n self.weights = SinusoidalPositionalEmbedding.get_embedding(\n max_pos,\n self.embedding_dim,\n self.padding_idx,\n )\n self.weights = self.weights.to(self._float_tensor)\n\n if incremental_state is not None:\n # positions is the same for every token when decoding a single step\n pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len\n return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)\n\n positions = make_positions(input, self.padding_idx) if positions is None else positions\n return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()\n\n def max_positions(self):\n \"\"\"Maximum number of supported positions.\"\"\"\n return int(1e5) # an arbitrary large number\n\n\nclass TransformerFFNLayer(nn.Module):\n def __init__(self, hidden_size, filter_size, padding=\"SAME\", kernel_size=1, dropout=0., act='gelu', bias=True):\n super().__init__()\n self.kernel_size = kernel_size\n self.dropout = dropout\n self.act = act\n if padding == 'SAME':\n self.ffn_1 = nn.Conv1d(hidden_size, filter_size, kernel_size,\n padding=kernel_size // 2, bias=bias)\n elif padding == 'LEFT':\n self.ffn_1 = nn.Sequential(\n nn.ConstantPad1d((kernel_size - 1, 0), 0.0),\n nn.Conv1d(hidden_size, filter_size, kernel_size, bias=bias)\n )\n self.ffn_2 = Linear(filter_size, hidden_size, bias=bias)\n\n def forward(self, x, incremental_state=None):\n # x: T x B x C\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_input' in saved_state:\n prev_input = saved_state['prev_input']\n x = torch.cat((prev_input, x), dim=0)\n x = x[-self.kernel_size:]\n saved_state['prev_input'] = x\n self._set_input_buffer(incremental_state, saved_state)\n\n x = self.ffn_1(x.permute(1, 2, 0)).permute(2, 0, 1)\n x = x * self.kernel_size ** -0.5\n\n if incremental_state is not None:\n x = x[-1:]\n if self.act == 'gelu':\n x = F.gelu(x)\n if self.act == 'relu':\n x = F.relu(x)\n x = F.dropout(x, self.dropout, training=self.training)\n x = self.ffn_2(x)\n return x\n\n def _get_input_buffer(self, incremental_state):\n return get_incremental_state(\n self,\n incremental_state,\n 'f',\n ) or {}\n\n def _set_input_buffer(self, incremental_state, buffer):\n set_incremental_state(\n self,\n incremental_state,\n 'f',\n buffer,\n )\n\n def clear_buffer(self, incremental_state):\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_input' in saved_state:\n del saved_state['prev_input']\n self._set_input_buffer(incremental_state, saved_state)\n\n\nclass MultiheadAttention(nn.Module):\n def __init__(self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0., bias=True,\n add_bias_kv=False, add_zero_attn=False, self_attention=False,\n encoder_decoder_attention=False):\n super().__init__()\n self.embed_dim = embed_dim\n self.kdim = kdim if kdim is not None else embed_dim\n self.vdim = vdim if vdim is not None else embed_dim\n self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim\n\n self.num_heads = num_heads\n self.dropout = dropout\n self.head_dim = embed_dim // num_heads\n assert self.head_dim * num_heads == self.embed_dim, \"embed_dim must be divisible by num_heads\"\n self.scaling = self.head_dim ** -0.5\n\n self.self_attention = self_attention\n self.encoder_decoder_attention = encoder_decoder_attention\n\n assert not self.self_attention or self.qkv_same_dim, 'Self-attention requires query, key and ' \\\n 'value to be of the same size'\n\n if self.qkv_same_dim:\n self.in_proj_weight = Parameter(torch.Tensor(3 * embed_dim, embed_dim))\n else:\n self.k_proj_weight = Parameter(torch.Tensor(embed_dim, self.kdim))\n self.v_proj_weight = Parameter(torch.Tensor(embed_dim, self.vdim))\n self.q_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))\n\n if bias:\n self.in_proj_bias = Parameter(torch.Tensor(3 * embed_dim))\n else:\n self.register_parameter('in_proj_bias', None)\n\n self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)\n\n if add_bias_kv:\n self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))\n self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))\n else:\n self.bias_k = self.bias_v = None\n\n self.add_zero_attn = add_zero_attn\n\n self.reset_parameters()\n\n self.enable_torch_version = False\n self.last_attn_probs = None\n\n def reset_parameters(self):\n if self.qkv_same_dim:\n nn.init.xavier_uniform_(self.in_proj_weight)\n else:\n nn.init.xavier_uniform_(self.k_proj_weight)\n nn.init.xavier_uniform_(self.v_proj_weight)\n nn.init.xavier_uniform_(self.q_proj_weight)\n\n nn.init.xavier_uniform_(self.out_proj.weight)\n if self.in_proj_bias is not None:\n nn.init.constant_(self.in_proj_bias, 0.)\n nn.init.constant_(self.out_proj.bias, 0.)\n if self.bias_k is not None:\n nn.init.xavier_normal_(self.bias_k)\n if self.bias_v is not None:\n nn.init.xavier_normal_(self.bias_v)\n\n def forward(\n self,\n query, key, value,\n key_padding_mask=None,\n incremental_state=None,\n need_weights=True,\n static_kv=False,\n attn_mask=None,\n before_softmax=False,\n need_head_weights=False,\n enc_dec_attn_constraint_mask=None,\n reset_attn_weight=None\n ):\n \"\"\"Input shape: Time x Batch x Channel\n\n Args:\n key_padding_mask (ByteTensor, optional): mask to exclude\n keys that are pads, of shape `(batch, src_len)`, where\n padding elements are indicated by 1s.\n need_weights (bool, optional): return the attention weights,\n averaged over heads (default: False).\n attn_mask (ByteTensor, optional): typically used to\n implement causal attention, where the mask prevents the\n attention from looking forward in time (default: None).\n before_softmax (bool, optional): return the raw attention\n weights and values before the attention softmax.\n need_head_weights (bool, optional): return the attention\n weights for each head. Implies *need_weights*. Default:\n return the average attention weights over all heads.\n \"\"\"\n if need_head_weights:\n need_weights = True\n\n tgt_len, bsz, embed_dim = query.size()\n assert embed_dim == self.embed_dim\n assert list(query.size()) == [tgt_len, bsz, embed_dim]\n\n if self.enable_torch_version and incremental_state is None and not static_kv and reset_attn_weight is None:\n if self.qkv_same_dim:\n return F.multi_head_attention_forward(query, key, value,\n self.embed_dim, self.num_heads,\n self.in_proj_weight,\n self.in_proj_bias, self.bias_k, self.bias_v,\n self.add_zero_attn, self.dropout,\n self.out_proj.weight, self.out_proj.bias,\n self.training, key_padding_mask, need_weights,\n attn_mask)\n else:\n return F.multi_head_attention_forward(query, key, value,\n self.embed_dim, self.num_heads,\n torch.empty([0]),\n self.in_proj_bias, self.bias_k, self.bias_v,\n self.add_zero_attn, self.dropout,\n self.out_proj.weight, self.out_proj.bias,\n self.training, key_padding_mask, need_weights,\n attn_mask, use_separate_proj_weight=True,\n q_proj_weight=self.q_proj_weight,\n k_proj_weight=self.k_proj_weight,\n v_proj_weight=self.v_proj_weight)\n\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_key' in saved_state:\n # previous time steps are cached - no need to recompute\n # key and value if they are static\n if static_kv:\n assert self.encoder_decoder_attention and not self.self_attention\n key = value = None\n else:\n saved_state = None\n\n if self.self_attention:\n # self-attention\n q, k, v = self.in_proj_qkv(query)\n elif self.encoder_decoder_attention:\n # encoder-decoder attention\n q = self.in_proj_q(query)\n if key is None:\n assert value is None\n k = v = None\n else:\n k = self.in_proj_k(key)\n v = self.in_proj_v(key)\n\n else:\n q = self.in_proj_q(query)\n k = self.in_proj_k(key)\n v = self.in_proj_v(value)\n q = q * self.scaling\n\n if self.bias_k is not None:\n assert self.bias_v is not None\n k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])\n v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])\n if attn_mask is not None:\n attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)\n if key_padding_mask is not None:\n key_padding_mask = torch.cat(\n [key_padding_mask, key_padding_mask.new_zeros(key_padding_mask.size(0), 1)], dim=1)\n\n q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if k is not None:\n k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n if v is not None:\n v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)\n\n if saved_state is not None:\n # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)\n if 'prev_key' in saved_state:\n prev_key = saved_state['prev_key'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n k = prev_key\n else:\n k = torch.cat((prev_key, k), dim=1)\n if 'prev_value' in saved_state:\n prev_value = saved_state['prev_value'].view(bsz * self.num_heads, -1, self.head_dim)\n if static_kv:\n v = prev_value\n else:\n v = torch.cat((prev_value, v), dim=1)\n if 'prev_key_padding_mask' in saved_state and saved_state['prev_key_padding_mask'] is not None:\n prev_key_padding_mask = saved_state['prev_key_padding_mask']\n if static_kv:\n key_padding_mask = prev_key_padding_mask\n else:\n key_padding_mask = torch.cat((prev_key_padding_mask, key_padding_mask), dim=1)\n\n saved_state['prev_key'] = k.view(bsz, self.num_heads, -1, self.head_dim)\n saved_state['prev_value'] = v.view(bsz, self.num_heads, -1, self.head_dim)\n saved_state['prev_key_padding_mask'] = key_padding_mask\n\n self._set_input_buffer(incremental_state, saved_state)\n\n src_len = k.size(1)\n\n # This is part of a workaround to get around fork/join parallelism\n # not supporting Optional types.\n if key_padding_mask is not None and key_padding_mask.shape == torch.Size([]):\n key_padding_mask = None\n\n if key_padding_mask is not None:\n assert key_padding_mask.size(0) == bsz\n assert key_padding_mask.size(1) == src_len\n\n if self.add_zero_attn:\n src_len += 1\n k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)\n v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)\n if attn_mask is not None:\n attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1)\n if key_padding_mask is not None:\n key_padding_mask = torch.cat(\n [key_padding_mask, torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask)], dim=1)\n\n attn_weights = torch.bmm(q, k.transpose(1, 2))\n attn_weights = self.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)\n\n assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]\n\n if attn_mask is not None:\n if len(attn_mask.shape) == 2:\n attn_mask = attn_mask.unsqueeze(0)\n elif len(attn_mask.shape) == 3:\n attn_mask = attn_mask[:, None].repeat([1, self.num_heads, 1, 1]).reshape(\n bsz * self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights + attn_mask\n\n if enc_dec_attn_constraint_mask is not None: # bs x head x L_kv\n attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights.masked_fill(\n enc_dec_attn_constraint_mask.unsqueeze(2).bool(),\n -1e8,\n )\n attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)\n\n if key_padding_mask is not None:\n # don't attend to padding symbols\n attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n attn_weights = attn_weights.masked_fill(\n key_padding_mask.unsqueeze(1).unsqueeze(2),\n -1e8,\n )\n attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)\n\n attn_logits = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)\n\n if before_softmax:\n return attn_weights, v\n\n attn_weights_float = softmax(attn_weights, dim=-1)\n attn_weights = attn_weights_float.type_as(attn_weights)\n attn_probs = F.dropout(attn_weights_float.type_as(attn_weights), p=self.dropout, training=self.training)\n\n if reset_attn_weight is not None:\n if reset_attn_weight:\n self.last_attn_probs = attn_probs.detach()\n else:\n assert self.last_attn_probs is not None\n attn_probs = self.last_attn_probs\n attn = torch.bmm(attn_probs, v)\n assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]\n attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)\n attn = self.out_proj(attn)\n\n if need_weights:\n attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0)\n if not need_head_weights:\n # average attention weights over heads\n attn_weights = attn_weights.mean(dim=0)\n else:\n attn_weights = None\n\n return attn, (attn_weights, attn_logits)\n\n def in_proj_qkv(self, query):\n return self._in_proj(query).chunk(3, dim=-1)\n\n def in_proj_q(self, query):\n if self.qkv_same_dim:\n return self._in_proj(query, end=self.embed_dim)\n else:\n bias = self.in_proj_bias\n if bias is not None:\n bias = bias[:self.embed_dim]\n return F.linear(query, self.q_proj_weight, bias)\n\n def in_proj_k(self, key):\n if self.qkv_same_dim:\n return self._in_proj(key, start=self.embed_dim, end=2 * self.embed_dim)\n else:\n weight = self.k_proj_weight\n bias = self.in_proj_bias\n if bias is not None:\n bias = bias[self.embed_dim:2 * self.embed_dim]\n return F.linear(key, weight, bias)\n\n def in_proj_v(self, value):\n if self.qkv_same_dim:\n return self._in_proj(value, start=2 * self.embed_dim)\n else:\n weight = self.v_proj_weight\n bias = self.in_proj_bias\n if bias is not None:\n bias = bias[2 * self.embed_dim:]\n return F.linear(value, weight, bias)\n\n def _in_proj(self, input, start=0, end=None):\n weight = self.in_proj_weight\n bias = self.in_proj_bias\n weight = weight[start:end, :]\n if bias is not None:\n bias = bias[start:end]\n return F.linear(input, weight, bias)\n\n def _get_input_buffer(self, incremental_state):\n return get_incremental_state(\n self,\n incremental_state,\n 'attn_state',\n ) or {}\n\n def _set_input_buffer(self, incremental_state, buffer):\n set_incremental_state(\n self,\n incremental_state,\n 'attn_state',\n buffer,\n )\n\n def apply_sparse_mask(self, attn_weights, tgt_len, src_len, bsz):\n return attn_weights\n\n def clear_buffer(self, incremental_state=None):\n if incremental_state is not None:\n saved_state = self._get_input_buffer(incremental_state)\n if 'prev_key' in saved_state:\n del saved_state['prev_key']\n if 'prev_value' in saved_state:\n del saved_state['prev_value']\n self._set_input_buffer(incremental_state, saved_state)\n\n\nclass EncSALayer(nn.Module):\n def __init__(self, c, num_heads, dropout, attention_dropout=0.1,\n relu_dropout=0.1, kernel_size=9, padding='SAME', act='gelu',\n ffn_hidden_size=1024):\n super().__init__()\n self.c = c\n self.dropout = dropout\n self.num_heads = num_heads\n if num_heads > 0:\n self.layer_norm1 = LayerNorm(c)\n self.self_attn = MultiheadAttention(\n self.c, num_heads, self_attention=True, dropout=attention_dropout, bias=False)\n self.layer_norm2 = LayerNorm(c)\n self.ffn = TransformerFFNLayer(\n c, ffn_hidden_size, kernel_size=kernel_size, dropout=relu_dropout, padding=padding, act=act)\n\n def forward(self, x, encoder_padding_mask=None, **kwargs):\n layer_norm_training = kwargs.get('layer_norm_training', None)\n if layer_norm_training is not None:\n self.layer_norm1.training = layer_norm_training\n self.layer_norm2.training = layer_norm_training\n if self.num_heads > 0:\n residual = x\n x = self.layer_norm1(x)\n x, _, = self.self_attn(\n query=x,\n key=x,\n value=x,\n key_padding_mask=encoder_padding_mask\n )\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None]\n\n residual = x\n x = self.layer_norm2(x)\n x = self.ffn(x)\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None]\n return x\n\n\nclass DecSALayer(nn.Module):\n def __init__(self, c, num_heads, dropout, attention_dropout=0.1, relu_dropout=0.1,\n kernel_size=9, ffn_hidden_size=1024, act='gelu', post_ln=False):\n super().__init__()\n self.c = c\n self.dropout = dropout\n self.layer_norm1 = LayerNorm(c)\n self.self_attn = MultiheadAttention(\n c, num_heads, self_attention=True, dropout=attention_dropout, bias=False\n )\n self.layer_norm2 = LayerNorm(c)\n self.encoder_attn = MultiheadAttention(\n c, num_heads, encoder_decoder_attention=True, dropout=attention_dropout, bias=False,\n )\n self.layer_norm3 = LayerNorm(c)\n self.ffn = TransformerFFNLayer(\n c, ffn_hidden_size, padding='LEFT', kernel_size=kernel_size, dropout=relu_dropout, act=act)\n self.post_ln = post_ln\n\n def forward(\n self,\n x,\n encoder_out=None,\n encoder_padding_mask=None,\n incremental_state=None,\n self_attn_mask=None,\n self_attn_padding_mask=None,\n attn_out=None,\n reset_attn_weight=None,\n **kwargs,\n ):\n layer_norm_training = kwargs.get('layer_norm_training', None)\n if layer_norm_training is not None:\n self.layer_norm1.training = layer_norm_training\n self.layer_norm2.training = layer_norm_training\n self.layer_norm3.training = layer_norm_training\n residual = x\n if not self.post_ln:\n x = self.layer_norm1(x)\n x, _ = self.self_attn(\n query=x,\n key=x,\n value=x,\n key_padding_mask=self_attn_padding_mask,\n incremental_state=incremental_state,\n attn_mask=self_attn_mask\n )\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n if self.post_ln:\n x = self.layer_norm1(x)\n\n attn_logits = None\n if encoder_out is not None or attn_out is not None:\n residual = x\n if not self.post_ln:\n x = self.layer_norm2(x)\n if encoder_out is not None:\n x, attn = self.encoder_attn(\n query=x,\n key=encoder_out,\n value=encoder_out,\n key_padding_mask=encoder_padding_mask,\n incremental_state=incremental_state,\n static_kv=True,\n enc_dec_attn_constraint_mask=get_incremental_state(self, incremental_state,\n 'enc_dec_attn_constraint_mask'),\n reset_attn_weight=reset_attn_weight\n )\n attn_logits = attn[1]\n elif attn_out is not None:\n x = self.encoder_attn.in_proj_v(attn_out)\n if encoder_out is not None or attn_out is not None:\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n if self.post_ln:\n x = self.layer_norm2(x)\n\n residual = x\n if not self.post_ln:\n x = self.layer_norm3(x)\n x = self.ffn(x, incremental_state=incremental_state)\n x = F.dropout(x, self.dropout, training=self.training)\n x = residual + x\n if self.post_ln:\n x = self.layer_norm3(x)\n return x, attn_logits\n\n def clear_buffer(self, input, encoder_out=None, encoder_padding_mask=None, incremental_state=None):\n self.encoder_attn.clear_buffer(incremental_state)\n self.ffn.clear_buffer(incremental_state)\n\n def set_buffer(self, name, tensor, incremental_state):\n return set_incremental_state(self, incremental_state, name, tensor)\n\n\nclass TransformerEncoderLayer(nn.Module):\n def __init__(self, hidden_size, dropout, kernel_size=9, num_heads=2, ffn_hidden_size=1024):\n super().__init__()\n self.hidden_size = hidden_size\n self.dropout = dropout\n self.num_heads = num_heads\n self.op = EncSALayer(\n hidden_size, num_heads, dropout=dropout,\n attention_dropout=0.0, relu_dropout=dropout,\n kernel_size=kernel_size, ffn_hidden_size=ffn_hidden_size)\n\n def forward(self, x, **kwargs):\n return self.op(x, **kwargs)\n\n\nclass TransformerDecoderLayer(nn.Module):\n def __init__(self, hidden_size, dropout, kernel_size=9, num_heads=2, ffn_hidden_size=1024, post_ln=False):\n super().__init__()\n self.hidden_size = hidden_size\n self.dropout = dropout\n self.num_heads = num_heads\n self.op = DecSALayer(\n hidden_size, num_heads, dropout=dropout,\n attention_dropout=0.0, relu_dropout=dropout,\n kernel_size=kernel_size, ffn_hidden_size=ffn_hidden_size,\n post_ln=post_ln)\n\n def forward(self, x, **kwargs):\n return self.op(x, **kwargs)\n\n def clear_buffer(self, *args):\n return self.op.clear_buffer(*args)\n\n def set_buffer(self, *args):\n return self.op.set_buffer(*args)\n\n\nclass FFTBlocks(nn.Module):\n def __init__(self, hidden_size, num_layers, ffn_kernel_size=9, dropout=0.0,\n num_heads=2, use_pos_embed=True, use_last_norm=True,\n use_pos_embed_alpha=True, ffn_hidden_size=1024):\n super().__init__()\n self.num_layers = num_layers\n embed_dim = self.hidden_size = hidden_size\n self.dropout = dropout\n self.use_pos_embed = use_pos_embed\n self.use_last_norm = use_last_norm\n if use_pos_embed:\n self.max_source_positions = DEFAULT_MAX_TARGET_POSITIONS\n self.padding_idx = 0\n self.pos_embed_alpha = nn.Parameter(torch.Tensor([1])) if use_pos_embed_alpha else 1\n self.embed_positions = SinusoidalPositionalEmbedding(\n embed_dim, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS,\n )\n\n self.layers = nn.ModuleList([])\n self.layers.extend([\n TransformerEncoderLayer(self.hidden_size, self.dropout,\n kernel_size=ffn_kernel_size, num_heads=num_heads,\n ffn_hidden_size=ffn_hidden_size)\n for _ in range(self.num_layers)\n ])\n if self.use_last_norm:\n self.layer_norm = nn.LayerNorm(embed_dim)\n else:\n self.layer_norm = None\n\n def forward(self, x, padding_mask=None, attn_mask=None, return_hiddens=False):\n \"\"\"\n :param x: [B, T, C]\n :param padding_mask: [B, T]\n :return: [B, T, C] or [L, B, T, C]\n \"\"\"\n padding_mask = x.abs().sum(-1).eq(0).data if padding_mask is None else padding_mask\n nonpadding_mask_TB = 1 - padding_mask.transpose(0, 1).float()[:, :, None] # [T, B, 1]\n if self.use_pos_embed:\n positions = self.pos_embed_alpha * self.embed_positions(x[..., 0])\n x = x + positions\n x = F.dropout(x, p=self.dropout, training=self.training)\n # B x T x C -> T x B x C\n x = x.transpose(0, 1) * nonpadding_mask_TB\n hiddens = []\n for layer in self.layers:\n x = layer(x, encoder_padding_mask=padding_mask, attn_mask=attn_mask) * nonpadding_mask_TB\n hiddens.append(x)\n if self.use_last_norm:\n x = self.layer_norm(x) * nonpadding_mask_TB\n if return_hiddens:\n x = torch.stack(hiddens, 0) # [L, T, B, C]\n x = x.transpose(1, 2) # [L, B, T, C]\n else:\n x = x.transpose(0, 1) # [B, T, C]\n return x\n\n\nclass FastSpeechEncoder(FFTBlocks):\n def __init__(self, dict_size, hidden_size=256, num_layers=4, kernel_size=9,\n dropout=0.0, num_heads=2, ffn_hidden_size=1024):\n super().__init__(hidden_size, num_layers, kernel_size, num_heads=num_heads,\n use_pos_embed=False, dropout=dropout, ffn_hidden_size=ffn_hidden_size)\n self.embed_tokens = Embedding(dict_size, hidden_size, 0)\n self.embed_scale = math.sqrt(hidden_size)\n self.padding_idx = 0\n self.embed_positions = SinusoidalPositionalEmbedding(\n hidden_size, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS,\n )\n\n def forward(self, txt_tokens, attn_mask=None, other_embeds=0):\n \"\"\"\n\n :param txt_tokens: [B, T]\n :return: {\n 'encoder_out': [B x T x C]\n }\n \"\"\"\n encoder_padding_mask = txt_tokens.eq(self.padding_idx).data\n x = self.forward_embedding(txt_tokens) + other_embeds # [B, T, H]\n if self.num_layers > 0:\n x = super(FastSpeechEncoder, self).forward(x, encoder_padding_mask, attn_mask=attn_mask)\n return x\n\n def forward_embedding(self, txt_tokens):\n # embed tokens and positions\n x = self.embed_scale * self.embed_tokens(txt_tokens)\n if self.use_pos_embed:\n positions = self.embed_positions(txt_tokens)\n x = x + positions\n x = F.dropout(x, p=self.dropout, training=self.training)\n return x\n" }, { "path": "tts/modules/llm_dit/cfm.py", "content": "# MIT License\n\n# Copyright (c) 2023 Alexander Tong\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\n# Copyright (c) [2023] [Alexander Tong] \n# Copyright (c) [2025] [Ziyue Jiang] \n# SPDX-License-Identifier: MIT\n# This file has been modified by Ziyue Jiang on 2025/03/19\n# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.\n# This modified file is released under the same license.\n\nimport math\nimport torch\nfrom typing import Union\nfrom torch.distributions import LogisticNormal\n\n\nclass LogitNormalTrainingTimesteps:\n def __init__(self, T=1000.0, loc=0.0, scale=1.0):\n assert T > 0\n self.T = T\n self.dist = LogisticNormal(loc, scale)\n\n def sample(self, size, device):\n t = self.dist.sample(size)[..., 0].to(device)\n return t\n \n\ndef pad_t_like_x(t, x):\n \"\"\"Function to reshape the time vector t by the number of dimensions of x.\n\n Parameters\n ----------\n x : Tensor, shape (bs, *dim)\n represents the source minibatch\n t : FloatTensor, shape (bs)\n\n Returns\n -------\n t : Tensor, shape (bs, number of x dimensions)\n\n Example\n -------\n x: Tensor (bs, C, W, H)\n t: Vector (bs)\n pad_t_like_x(t, x): Tensor (bs, 1, 1, 1)\n \"\"\"\n if isinstance(t, (float, int)):\n return t\n return t.reshape(-1, *([1] * (x.dim() - 1)))\n\n\nclass ConditionalFlowMatcher:\n \"\"\"Base class for conditional flow matching methods. This class implements the independent\n conditional flow matching methods from [1] and serves as a parent class for all other flow\n matching methods.\n\n It implements:\n - Drawing data from gaussian probability path N(t * x1 + (1 - t) * x0, sigma) function\n - conditional flow matching ut(x1|x0) = x1 - x0\n - score function $\\nabla log p_t(x|x0, x1)$\n \"\"\"\n\n def __init__(self, sigma: Union[float, int] = 0.0):\n r\"\"\"Initialize the ConditionalFlowMatcher class. It requires the hyper-parameter $\\sigma$.\n\n Parameters\n ----------\n sigma : Union[float, int]\n \"\"\"\n self.sigma = sigma\n self.time_sampler = LogitNormalTrainingTimesteps()\n\n def compute_mu_t(self, x0, x1, t):\n \"\"\"\n Compute the mean of the probability path N(t * x1 + (1 - t) * x0, sigma), see (Eq.14) [1].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n t : FloatTensor, shape (bs)\n\n Returns\n -------\n mean mu_t: t * x1 + (1 - t) * x0\n\n References\n ----------\n [1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.\n \"\"\"\n t = pad_t_like_x(t, x0)\n return t * x1 + (1 - t) * x0\n\n def compute_sigma_t(self, t):\n \"\"\"\n Compute the standard deviation of the probability path N(t * x1 + (1 - t) * x0, sigma), see (Eq.14) [1].\n\n Parameters\n ----------\n t : FloatTensor, shape (bs)\n\n Returns\n -------\n standard deviation sigma\n\n References\n ----------\n [1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.\n \"\"\"\n del t\n return self.sigma\n\n def sample_xt(self, x0, x1, t, epsilon):\n \"\"\"\n Draw a sample from the probability path N(t * x1 + (1 - t) * x0, sigma), see (Eq.14) [1].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n t : FloatTensor, shape (bs)\n epsilon : Tensor, shape (bs, *dim)\n noise sample from N(0, 1)\n\n Returns\n -------\n xt : Tensor, shape (bs, *dim)\n\n References\n ----------\n [1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.\n \"\"\"\n mu_t = self.compute_mu_t(x0, x1, t)\n sigma_t = self.compute_sigma_t(t)\n sigma_t = pad_t_like_x(sigma_t, x0)\n return mu_t + sigma_t * epsilon\n\n def compute_conditional_flow(self, x0, x1, t, xt):\n \"\"\"\n Compute the conditional vector field ut(x1|x0) = x1 - x0, see Eq.(15) [1].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n t : FloatTensor, shape (bs)\n xt : Tensor, shape (bs, *dim)\n represents the samples drawn from probability path pt\n\n Returns\n -------\n ut : conditional vector field ut(x1|x0) = x1 - x0\n\n References\n ----------\n [1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.\n \"\"\"\n del t, xt\n return x1 - x0\n\n def sample_noise_like(self, x):\n return torch.randn_like(x)\n\n def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False):\n \"\"\"\n Compute the sample xt (drawn from N(t * x1 + (1 - t) * x0, sigma))\n and the conditional vector field ut(x1|x0) = x1 - x0, see Eq.(15) [1].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n (optionally) t : Tensor, shape (bs)\n represents the time levels\n if None, drawn from uniform [0,1]\n return_noise : bool\n return the noise sample epsilon\n\n\n Returns\n -------\n t : FloatTensor, shape (bs)\n xt : Tensor, shape (bs, *dim)\n represents the samples drawn from probability path pt\n ut : conditional vector field ut(x1|x0) = x1 - x0\n (optionally) eps: Tensor, shape (bs, *dim) such that xt = mu_t + sigma_t * epsilon\n\n References\n ----------\n [1] Improving and Generalizing Flow-Based Generative Models with minibatch optimal transport, Preprint, Tong et al.\n \"\"\"\n if t is None:\n # t = torch.rand(x0.shape[0]).type_as(x0)\n t = self.time_sampler.sample([x0.shape[0]], x0.device).type_as(x0)\n\n assert len(t) == x0.shape[0], \"t has to have batch size dimension\"\n\n eps = self.sample_noise_like(x0)\n xt = self.sample_xt(x0, x1, t, eps)\n ut = self.compute_conditional_flow(x0, x1, t, xt)\n if return_noise:\n return t, xt, ut, eps\n else:\n return t, xt, ut\n\n def compute_lambda(self, t):\n \"\"\"Compute the lambda function, see Eq.(23) [3].\n\n Parameters\n ----------\n t : FloatTensor, shape (bs)\n\n Returns\n -------\n lambda : score weighting function\n\n References\n ----------\n [4] Simulation-free Schrodinger bridges via score and flow matching, Preprint, Tong et al.\n \"\"\"\n sigma_t = self.compute_sigma_t(t)\n return 2 * sigma_t / (self.sigma**2 + 1e-8)\n\n\nclass VariancePreservingConditionalFlowMatcher(ConditionalFlowMatcher):\n \"\"\"Albergo et al. 2023 trigonometric interpolants class. This class inherits the\n ConditionalFlowMatcher and override the compute_mu_t and compute_conditional_flow functions in\n order to compute [3]'s trigonometric interpolants.\n\n [3] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, Albergo et al.\n \"\"\"\n\n def compute_mu_t(self, x0, x1, t):\n r\"\"\"Compute the mean of the probability path (Eq.5) from [3].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n t : FloatTensor, shape (bs)\n\n Returns\n -------\n mean mu_t: cos(pi t/2)x0 + sin(pi t/2)x1\n\n References\n ----------\n [3] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, Albergo et al.\n \"\"\"\n t = pad_t_like_x(t, x0)\n return torch.cos(math.pi / 2 * t) * x0 + torch.sin(math.pi / 2 * t) * x1\n\n def compute_conditional_flow(self, x0, x1, t, xt):\n r\"\"\"Compute the conditional vector field similar to [3].\n\n ut(x1|x0) = pi/2 (cos(pi*t/2) x1 - sin(pi*t/2) x0),\n see Eq.(21) [3].\n\n Parameters\n ----------\n x0 : Tensor, shape (bs, *dim)\n represents the source minibatch\n x1 : Tensor, shape (bs, *dim)\n represents the target minibatch\n t : FloatTensor, shape (bs)\n xt : Tensor, shape (bs, *dim)\n represents the samples drawn from probability path pt\n\n Returns\n -------\n ut : conditional vector field\n ut(x1|x0) = pi/2 (cos(pi*t/2) x1 - sin(\\pi*t/2) x0)\n\n References\n ----------\n [3] Stochastic Interpolants: A Unifying Framework for Flows and Diffusions, Albergo et al.\n \"\"\"\n del xt\n t = pad_t_like_x(t, x0)\n return math.pi / 2 * (torch.cos(math.pi / 2 * t) * x1 - torch.sin(math.pi / 2 * t) * x0)\n" }, { "path": "tts/modules/llm_dit/dit.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\nfrom torch import nn\n\nfrom tts.modules.llm_dit.cfm import ConditionalFlowMatcher\nfrom tts.modules.ar_dur.commons.layers import Embedding\nfrom tts.modules.ar_dur.commons.nar_tts_modules import PosEmb\nfrom tts.modules.ar_dur.commons.rel_transformer import RelTransformerEncoder\nfrom tts.modules.ar_dur.ar_dur_predictor import expand_states\nfrom tts.modules.llm_dit.transformer import Transformer\nfrom tts.modules.llm_dit.time_embedding import TimestepEmbedding\n\n\nclass Diffusion(nn.Module):\n def __init__(self):\n super().__init__()\n # Hparams\n # cond dim\n self.local_cond_dim = 512\n self.ctx_mask_dim = 16\n self.in_channels = 32\n self.out_channels = 32\n # LLM\n self.encoder_dim = 1024\n self.encoder_n_layers = 24\n self.encoder_n_heads = 16\n self.max_seq_len = 16384\n self.multiple_of = 256\n\n self.ctx_mask_proj = nn.Linear(1, self.ctx_mask_dim)\n self.local_cond_project = nn.Linear(\n self.out_channels + self.ctx_mask_dim, self.local_cond_dim)\n\n self.encoder = Transformer(self.encoder_n_layers, self.encoder_dim, self.encoder_n_heads, self.max_seq_len)\n\n self.x_prenet = nn.Linear(self.in_channels, self.encoder_dim)\n self.prenet = nn.Linear(self.local_cond_dim, self.encoder_dim)\n self.postnet = nn.Linear(self.encoder_dim, self.out_channels)\n \n self.flow_matcher = ConditionalFlowMatcher(sigma=0.0)\n # The implementation of TimestepEmbedding is a modified version from F5-TTS (https://github.com/SWivid/F5-TTS), \n # which is licensed under the MIT License.\n self.f5_time_embed = TimestepEmbedding(self.encoder_dim)\n\n # text encoder\n self.ph_encoder = RelTransformerEncoder(\n 302, self.encoder_dim, self.encoder_dim,\n self.encoder_dim * 2, 4, 6,\n 3, 0.0, prenet=True, pre_ln=True)\n self.tone_embed = Embedding(32, self.encoder_dim, padding_idx=0)\n self.ph_pos_embed = PosEmb(self.encoder_dim)\n self.ling_pre_net = torch.nn.Sequential(*[\n torch.nn.Conv1d(self.encoder_dim, self.encoder_dim, kernel_size=s * 2, stride=s, padding=s // 2)\n for i, s in enumerate([2, 2])\n ])\n \n def forward(self, inputs, sigmas=None, x_noisy=None):\n ctx_mask = inputs['ctx_mask']\n ctx_feature = inputs['lat_ctx'] * ctx_mask\n\n \"\"\" local conditioning (prompt_latent + spk_embed) \"\"\"\n ctx_mask_emb = self.ctx_mask_proj(ctx_mask)\n # ctx_feature = ctx_feature * (1 - inputs[\"spk_cfg_mask\"][:, :, None])\n local_cond = torch.cat([ctx_feature, ctx_mask_emb], dim=-1)\n local_cond = self.local_cond_project(local_cond)\n\n \"\"\" diffusion target latent \"\"\"\n x = inputs['lat']\n \n # Here, x is x1 in CFM\n x0 = torch.randn_like(x)\n t, xt, ut = self.flow_matcher.sample_location_and_conditional_flow(x0, x)\n \n # define noisy_input and target\n t = t.bfloat16()\n x_noisy = (xt * (1 - ctx_mask)).bfloat16()\n target = ut\n\n # concat condition.\n x_ling = self.forward_ling_encoder(inputs[\"phone\"], inputs[\"tone\"])\n x_ling = self.ling_pre_net(expand_states(x_ling, inputs['mel2ph']).transpose(1, 2)).transpose(1, 2)\n x_noisy = self.x_prenet(x_noisy) + self.prenet(local_cond) + x_ling\n encoder_out = self.encoder(x_noisy, self.f5_time_embed(t), attn_mask=inputs[\"text_mel_mask\"], do_checkpoint=False)\n pred = self.postnet(encoder_out)\n\n return pred, target\n \n def forward_ling_encoder(self, txt_tokens, tone_tokens):\n ph_tokens = txt_tokens\n ph_nonpadding = (ph_tokens > 0).float()[:, :, None] # [B, T_phone, 1]\n\n # enc_ph\n ph_enc_oembed = self.tone_embed(tone_tokens)\n ph_enc_oembed = ph_enc_oembed + self.ph_pos_embed(\n torch.arange(0, ph_tokens.shape[1])[None,].to(ph_tokens.device))\n ph_enc_oembed = ph_enc_oembed\n ph_enc_oembed = ph_enc_oembed * ph_nonpadding\n x_ling = self.ph_encoder(ph_tokens, other_embeds=ph_enc_oembed) * ph_nonpadding\n return x_ling\n\n def _forward(self, x, local_cond, x_ling, timesteps, ctx_mask, dur=None, seq_cfg_w=[1.0,1.0]):\n \"\"\" When we use torchdiffeq, we need to include the CFG process inside _forward() \"\"\"\n x = x * (1 - ctx_mask)\n x = self.x_prenet(x) + self.prenet(local_cond) + x_ling\n pred_v = self.encoder(x, self.f5_time_embed(timesteps), attn_mask=torch.ones((x.size(0), x.size(1)), device=x.device))\n pred = self.postnet(pred_v)\n\n \"\"\" Perform multi-cond CFG \"\"\"\n cond_spk_txt, cond_txt, uncond = pred.chunk(3)\n pred = uncond + seq_cfg_w[0] * (cond_txt - uncond) + seq_cfg_w[1] * (cond_spk_txt - cond_txt)\n return pred\n\n @torch.no_grad()\n def inference(self, inputs, timesteps=20, seq_cfg_w=[1.0, 1.0], **kwargs):\n # txt embedding\n x_ling = self.forward_ling_encoder(inputs[\"phone\"], inputs[\"tone\"])\n x_ling = self.ling_pre_net(expand_states(x_ling, inputs['dur']).transpose(1, 2)).transpose(1, 2)\n\n # speaker embedding\n ctx_feature = inputs['lat_ctx']\n ctx_feature[1:, :, :] = 0 # prefix spk cfg\n ctx_mask_emb = self.ctx_mask_proj(inputs['ctx_mask'])\n\n # local conditioning.\n local_cond = torch.cat([ctx_feature, ctx_mask_emb], dim=-1)\n local_cond = self.local_cond_project(local_cond)\n \n ''' Euler ODE solver '''\n bsz, device, frm_len = (local_cond.size(0), local_cond.device, local_cond.size(1))\n # Sway sampling from F5-TTS (https://github.com/SWivid/F5-TTS), \n # which is licensed under the MIT License.\n sway_sampling_coef = -1.0\n t_schedule = torch.linspace(0, 1, timesteps + 1, device=device, dtype=x_ling.dtype)\n if sway_sampling_coef is not None:\n t_schedule = t_schedule + sway_sampling_coef * (torch.cos(torch.pi / 2 * t_schedule) - 1 + t_schedule)\n \n # AMO sampling implementation for \"AMO Sampler: Enhancing Text Rendering with Overshooting\" (https://arxiv.org/pdf/2411.19415)\n def amo_sampling(z_t, t, t_next, v):\n # Upcast to avoid precision issues when computing prev_sample\n z_t = z_t.to(torch.float32)\n\n # Constant definition in Algorithm 1\n s = t_next\n c = 3\n\n # Line 7 in Algorithm 1\n o = min(t_next + c * (t_next - t), 1)\n pred_z_o = z_t + (o - t) * v\n\n # Line 11 in Algorithm 1\n a = s / o\n b = ((1 - s) ** 2 - (a * (1 - o)) ** 2) ** 0.5\n noise_i = torch.randn(size=z_t.shape, device=z_t.device)\n z_t_next = a * pred_z_o + b * noise_i\n return z_t_next.to(v.dtype)\n\n x = torch.randn([1, frm_len, self.out_channels], device=device)\n for step_index in range(timesteps):\n x = x.to(torch.float32)\n sigma = t_schedule[step_index].to(x_ling.dtype)\n sigma_next = t_schedule[step_index + 1]\n model_out = self._forward(torch.cat([x] * bsz), local_cond, x_ling, timesteps=sigma.unsqueeze(0), ctx_mask=inputs['ctx_mask'], dur=inputs['dur'], seq_cfg_w=seq_cfg_w)\n x = amo_sampling(x, sigma, sigma_next, model_out)\n # Cast sample back to model compatible dtype\n x = x.to(model_out.dtype)\n \n return x\n" }, { "path": "tts/modules/llm_dit/time_embedding.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 math\nimport torch\nfrom torch import nn\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\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): # noqa: F821\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" }, { "path": "tts/modules/llm_dit/transformer.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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\nfrom typing import Any, Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\n\n\ndef precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):\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_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64\n return freqs_cis\n\n\ndef reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):\n ndim = x.ndim\n assert 0 <= 1 < ndim\n assert freqs_cis.shape == (x.shape[1], x.shape[-1])\n shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]\n return freqs_cis.view(*shape)\n\n\ndef apply_rotary_emb(\n xq: torch.Tensor,\n xk: torch.Tensor,\n freqs_cis: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))\n xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))\n freqs_cis = reshape_for_broadcast(freqs_cis, xq_)\n xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)\n xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)\n return xq_out.type_as(xq), xk_out.type_as(xk)\n\n\nclass AdaLNZero(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.silu = nn.SiLU()\n self.linear = nn.Linear(dim, dim * 6)\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 x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]\n return x, gate_msa, shift_mlp, scale_mlp, gate_mlp\n\n\nclass AdaLNZero_Out(nn.Module):\n def __init__(self, dim):\n super().__init__()\n self.silu = nn.SiLU()\n self.linear = nn.Linear(dim, dim * 2)\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 x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]\n return x\n\n\nclass Attention(nn.Module):\n def __init__(self, encoder_dim, encoder_n_heads, max_seq_len):\n super().__init__()\n self.encoder_n_kv_heads = encoder_n_heads\n model_parallel_size = 1\n self.n_local_heads = encoder_n_heads // model_parallel_size\n self.n_local_kv_heads = self.encoder_n_kv_heads // model_parallel_size\n self.n_rep = self.n_local_heads // self.n_local_kv_heads\n self.head_dim = encoder_dim // encoder_n_heads\n\n self.wq = nn.Linear(\n encoder_dim,\n encoder_n_heads * self.head_dim,\n )\n self.wk = nn.Linear(\n encoder_dim,\n self.encoder_n_kv_heads * self.head_dim,\n )\n self.wv = nn.Linear(\n encoder_dim,\n self.encoder_n_kv_heads * self.head_dim,\n )\n self.wo = nn.Linear(\n encoder_n_heads * self.head_dim,\n encoder_dim,\n )\n\n def forward(\n self,\n x: torch.Tensor,\n start_pos: int,\n freqs_cis: torch.Tensor,\n mask: Optional[torch.Tensor],\n ):\n bsz, seqlen, _ = x.shape\n xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)\n xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)\n xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)\n xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)\n\n xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)\n xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim)\n keys = xk.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)\n values = xv.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)\n\n output = F.scaled_dot_product_attention(xq, keys, values, mask[:, None, None, :], is_causal=False)\n output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)\n return self.wo(output)\n\n\nclass FeedForward(nn.Module):\n def __init__(\n self,\n dim: int,\n hidden_dim: int,\n multiple_of: int,\n ffn_dim_multiplier: Optional[float],\n ):\n super().__init__()\n if ffn_dim_multiplier is not None:\n hidden_dim = int(ffn_dim_multiplier * hidden_dim)\n hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)\n\n self.w1 = nn.Linear(\n dim, hidden_dim\n )\n self.w2 = nn.Linear(\n hidden_dim, dim\n )\n\n def forward(self, x):\n return self.w2(F.silu(self.w1(x)))\n\n\nclass TransformerBlock(nn.Module):\n def __init__(self, encoder_dim, encoder_n_heads, max_seq_len):\n super().__init__()\n self.encoder_n_heads = encoder_n_heads\n self.encoder_dim = encoder_dim\n self.head_dim = encoder_dim // encoder_n_heads\n self.attention = Attention(encoder_dim, encoder_n_heads, max_seq_len)\n self.feed_forward = FeedForward(\n dim=encoder_dim,\n hidden_dim=2 * encoder_dim,\n multiple_of=256,\n ffn_dim_multiplier=None,\n )\n self.attention_norm = AdaLNZero(encoder_dim)\n self.ffn_norm = nn.LayerNorm(encoder_dim, elementwise_affine=False, eps=1e-6)\n\n def forward(\n self,\n x: torch.Tensor,\n t: torch.Tensor,\n start_pos: int,\n freqs_cis: torch.Tensor,\n mask: Optional[torch.Tensor],\n ):\n \"\"\"\n Perform a forward pass through the TransformerBlock.\n\n Args:\n x (torch.Tensor): Input tensor.\n start_pos (int): Starting position for attention caching.\n freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies.\n mask (torch.Tensor, optional): Masking tensor for attention. Defaults to None.\n\n Returns:\n torch.Tensor: Output tensor after applying attention and feedforward layers.\n\n \"\"\"\n # pre-norm & modulation for attention input\n norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attention_norm(x, emb=t)\n\n # attention\n attn_output = self.attention(norm, start_pos, freqs_cis, mask=mask)\n\n # process attention output for input x\n h = x + gate_msa.unsqueeze(1) * attn_output\n\n norm = self.ffn_norm(h) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]\n ff_output = self.feed_forward(norm)\n out = h + gate_mlp.unsqueeze(1) * ff_output\n\n return out\n\n\nclass Transformer(nn.Module):\n def __init__(self, encoder_n_layers, encoder_dim, encoder_n_heads, max_seq_len):\n super().__init__()\n # Decoder\n self.layers = torch.nn.ModuleList()\n for _ in range(encoder_n_layers):\n self.layers.append(TransformerBlock(encoder_dim, encoder_n_heads, max_seq_len))\n\n self.norm = AdaLNZero_Out(encoder_dim)\n self.out_proj = nn.Linear(encoder_dim, encoder_dim)\n\n # Rope embedding\n freqs_cis = precompute_freqs_cis(\n encoder_dim // encoder_n_heads, max_seq_len\n )\n self.register_buffer(\"freqs_cis\", torch.view_as_real(freqs_cis), persistent=False)\n \n def forward(self, x, t, attn_mask, start_pos=0):\n freqs_cis = torch.view_as_complex(self.freqs_cis.float())[start_pos: start_pos + x.size(1)]\n for i, layer in enumerate(self.layers):\n x = layer(x, t, start_pos, freqs_cis, attn_mask)\n x = self.norm(x, t)\n x = self.out_proj(x)\n return x" }, { "path": "tts/modules/wavvae/decoder/diag_gaussian.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\nimport numpy as np\n\nclass DiagonalGaussianDistribution(object):\n def __init__(self, parameters: torch.Tensor, deterministic: bool = False):\n self.parameters = parameters\n self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)\n self.logvar = torch.clamp(self.logvar, -30.0, 20.0)\n self.deterministic = deterministic\n self.std = torch.exp(0.5 * self.logvar)\n self.var = torch.exp(self.logvar)\n if self.deterministic:\n self.var = self.std = torch.zeros_like(\n self.mean, device=self.parameters.device, dtype=self.parameters.dtype\n )\n\n def sample(self, generator=None) -> torch.Tensor:\n # make sure sample is on the same device as the parameters and has same dtype\n sample = torch.randn(\n self.mean.shape,\n generator=generator,\n device=self.parameters.device,\n dtype=self.parameters.dtype,\n )\n x = self.mean + self.std * sample\n return x\n\n def kl(self, other: \"DiagonalGaussianDistribution\" = None) -> torch.Tensor:\n if self.deterministic:\n return torch.Tensor([0.0])\n else:\n if other is None:\n return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar\n else:\n return 0.5 * (\n torch.pow(self.mean - other.mean, 2) / other.var\n + self.var / other.var\n - 1.0\n - self.logvar\n + other.logvar\n )\n\n def nll(self, sample, dims) -> torch.Tensor:\n if self.deterministic:\n return torch.Tensor([0.0])\n logtwopi = np.log(2.0 * np.pi)\n return 0.5 * torch.sum(\n logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,\n dim=dims,\n )\n\n def mode(self) -> torch.Tensor:\n return self.mean" }, { "path": "tts/modules/wavvae/decoder/hifigan_modules.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch.nn as nn\nimport torch.nn.functional as F\nimport torch\nimport torch.utils.data\nfrom librosa.filters import mel as librosa_mel_fn\nfrom torch.nn.utils import weight_norm, remove_weight_norm\nfrom torch.nn import Conv1d\nimport numpy as np\n\n\ndef init_weights(m, mean=0.0, std=0.01):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n m.weight.data.normal_(mean, std)\n\n\ndef get_padding(kernel_size, dilation=1):\n return int((kernel_size*dilation - dilation)/2)\n\n\nclass Upsample(nn.Module):\n def __init__(self, mult, r):\n super(Upsample, self).__init__()\n self.r = r\n self.upsample = nn.Sequential(nn.Upsample(mode=\"nearest\", scale_factor=r),\n nn.LeakyReLU(0.2),\n nn.ReflectionPad1d(3),\n nn.utils.weight_norm(nn.Conv1d(mult, mult // 2, kernel_size=7, stride=1))\n )\n r_kernel = r if r >= 5 else 5\n self.trans_upsample = nn.Sequential(nn.LeakyReLU(0.2),\n nn.utils.weight_norm(nn.ConvTranspose1d(mult, mult // 2,\n kernel_size=r_kernel * 2, stride=r,\n padding=r_kernel - r // 2,\n output_padding=r % 2)\n ))\n\n def forward(self, x):\n x = torch.sin(x) + x\n out1 = self.upsample(x)\n out2 = self.trans_upsample(x)\n return out1 + out2\n\n\nclass Downsample(nn.Module):\n def __init__(self, mult, r):\n super(Downsample, self).__init__()\n self.r = r\n r_kernel = r if r >= 5 else 5\n self.trans_downsample = nn.Sequential(nn.LeakyReLU(0.2),\n nn.utils.weight_norm(nn.Conv1d(mult, mult * 2,\n kernel_size=r_kernel * 2, stride=r,\n padding=r_kernel - r // 2)\n ))\n\n def forward(self, x):\n out = self.trans_downsample(x)\n return out\n\n\ndef weights_init(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n m.weight.data.normal_(0.0, 0.02)\n elif classname.find(\"BatchNorm2d\") != -1:\n m.weight.data.normal_(1.0, 0.02)\n m.bias.data.fill_(0)\n\n\ndef weights_zero_init(m):\n classname = m.__class__.__name__\n if classname.find(\"Conv\") != -1:\n m.weight.data.fill_(0.0)\n m.bias.data.fill_(0.0)\n\n\ndef WNConv1d(*args, **kwargs):\n return weight_norm(nn.Conv1d(*args, **kwargs))\n\n\ndef WNConvTranspose1d(*args, **kwargs):\n return weight_norm(nn.ConvTranspose1d(*args, **kwargs))\n\n\nclass Audio2Mel(nn.Module):\n def __init__(\n self,\n hop_length=300,\n sampling_rate=24000,\n n_mel_channels=80,\n mel_fmin=0.,\n mel_fmax=None,\n frame_size=0.05,\n device='cpu'\n ):\n super().__init__()\n ##############################################\n # FFT Parameters #\n ##############################################\n\n self.n_fft = int(np.power(2., np.ceil(np.log(sampling_rate * frame_size) / np.log(2))))\n window = torch.hann_window(int(sampling_rate * frame_size)).float()\n mel_basis = librosa_mel_fn(\n sampling_rate, self.n_fft, n_mel_channels, mel_fmin, mel_fmax\n ) # Mel filter (by librosa)\n mel_basis = torch.from_numpy(mel_basis).float()\n self.register_buffer(\"mel_basis\", mel_basis)\n self.register_buffer(\"window\", window)\n\n self.hop_length = hop_length\n self.win_length = int(sampling_rate * frame_size)\n self.sampling_rate = sampling_rate\n self.n_mel_channels = n_mel_channels\n\n def forward(self, audio):\n fft = torch.stft(\n audio.squeeze(1),\n n_fft=self.n_fft,\n hop_length=self.hop_length,\n win_length=self.win_length,\n window=self.window,\n center=True,\n )\n real_part, imag_part = fft.unbind(-1)\n magnitude = torch.sqrt(torch.clamp(real_part ** 2 + imag_part ** 2, min=1e-5))\n mel_output = torch.matmul(self.mel_basis, magnitude)\n\n log_mel_spec = 20 * torch.log10(torch.clamp(mel_output, min=1e-5)) - 20\n norm_mel = (log_mel_spec + 115.) / 115.\n mel_comp = torch.clamp(norm_mel * 8. - 4., -4., 4.)\n\n return mel_comp\n\n\nclass ResnetBlock(nn.Module):\n def __init__(self, dim, dilation=1, dim_in=None):\n super().__init__()\n if dim_in is None:\n dim_in = dim\n\n self.block = nn.Sequential(\n nn.LeakyReLU(0.2),\n nn.ReflectionPad1d(dilation),\n WNConv1d(dim_in, dim, kernel_size=3, dilation=dilation),\n nn.LeakyReLU(0.2),\n WNConv1d(dim, dim, kernel_size=1),\n )\n self.shortcut = WNConv1d(dim_in, dim, kernel_size=1)\n\n def forward(self, x):\n return self.shortcut(x) + self.block(x)\n\n\n'''\n参照hifigan（https://arxiv.org/pdf/2010.05646.pdf）v2结构\n多尺度主要是kernel_size不同，3组并行卷积模块，每个卷积模块内部采用不同的串行dilation size，且中间交叉正常无dilation卷积层\n'''\n\n\nclass ResBlockMRFV2(torch.nn.Module):\n def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):\n super(ResBlockMRFV2, self).__init__()\n self.convs1 = nn.ModuleList([\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],\n padding=get_padding(kernel_size, dilation[0]))),\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],\n padding=get_padding(kernel_size, dilation[1]))),\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],\n padding=get_padding(kernel_size, dilation[2])))\n ])\n self.convs1.apply(init_weights)\n\n self.convs2 = nn.ModuleList([\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,\n padding=get_padding(kernel_size, 1))),\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,\n padding=get_padding(kernel_size, 1))),\n weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,\n padding=get_padding(kernel_size, 1)))\n ])\n self.convs2.apply(init_weights)\n\n def forward(self, x):\n for c1, c2 in zip(self.convs1, self.convs2):\n xt = F.leaky_relu(x, 0.2)\n xt = c1(xt)\n xt = F.leaky_relu(xt, 0.2)\n xt = c2(xt)\n x = xt + x\n return x\n\n def remove_weight_norm(self):\n for l in self.convs1:\n remove_weight_norm(l)\n for l in self.convs2:\n remove_weight_norm(l)\n\n\nclass ResBlockMRFV2Inter(torch.nn.Module):\n def __init__(self, channels, kernel_size=3):\n super(ResBlockMRFV2Inter, self).__init__()\n self.block1 = ResBlockMRFV2(channels)\n self.block2 = ResBlockMRFV2(channels, 7)\n self.block3 = ResBlockMRFV2(channels, 11)\n\n def forward(self, x):\n xs = self.block1(x)\n xs += self.block2(x)\n xs += self.block3(x)\n x = xs / 3\n return x\n\n\nclass Generator(nn.Module):\n def __init__(self, input_size_, ngf, n_residual_layers, num_band, args, ratios=[5, 5, 4, 3], onnx_export=False,\n device='cpu'):\n super().__init__()\n self.hop_length = args.frame_shift\n self.args = args\n self.onnx_export = onnx_export\n\n # ------------- Define upsample layers ----------------\n mult = int(2 ** len(ratios))\n model_up = []\n input_size = input_size_\n model_up += [\n nn.ReflectionPad1d(3),\n WNConv1d(input_size, mult * ngf, kernel_size=7, padding=0),\n ]\n\n # Upsample to raw audio scale\n for i, r in enumerate(ratios):\n model_up += [Upsample(mult * ngf, r)]\n model_up += [ResBlockMRFV2Inter(mult * ngf // 2)]\n mult //= 2\n\n model_up += [\n nn.LeakyReLU(0.2),\n nn.ReflectionPad1d(3),\n WNConv1d(ngf, num_band, kernel_size=7, padding=0),\n nn.Tanh(),\n ]\n if not args.use_tanh:\n model_up[-1] = nn.Conv1d(num_band, num_band, 1)\n model_up[-2].apply(weights_zero_init)\n\n self.model_up = nn.Sequential(*model_up)\n\n self.apply(weights_init)\n\n def forward(self, mel, step=None):\n # mel input: (batch_size, seq_num, 80)\n if self.onnx_export:\n mel = mel.transpose(1, 2)\n # on onnx, for engineering, mel input: (batch_size, 80, seq_num)\n\n # Between Down and up\n x = mel\n\n # Upsample pipline\n cnt_after_upsample = 0\n\n for i, m in enumerate(self.model_up):\n x = m(x)\n\n if type(m) == Upsample:\n cnt_after_upsample += 1\n\n return x" }, { "path": "tts/modules/wavvae/decoder/seanet_encoder.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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\nfrom typing import List\n\nimport torch\nfrom torch import nn\nfrom tts.modules.wavvae.encoder.common_modules.seanet import SEANetEncoder\n\nclass Encoder(nn.Module):\n def __init__(\n self,\n dowmsamples: List[int] = [6, 5, 5, 4, 2],\n ):\n super().__init__()\n\n # breakpoint()\n self.frame_rate = 25 # not use\n self.encoder = SEANetEncoder(causal=False, n_residual_layers=1, norm='weight_norm', pad_mode='reflect', lstm=2,\n dimension=512, channels=1, n_filters=32, ratios=dowmsamples, activation='ELU',\n kernel_size=7, residual_kernel_size=3, last_kernel_size=7, dilation_base=2,\n true_skip=False, compress=2)\n\n def forward(self, audio: torch.Tensor):\n audio = audio.unsqueeze(1) # audio(16,24000)\n emb = self.encoder(audio)\n return emb\n" }, { "path": "tts/modules/wavvae/decoder/wavvae_v3.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\nfrom tts.modules.wavvae.decoder.seanet_encoder import Encoder\nfrom tts.modules.wavvae.decoder.diag_gaussian import DiagonalGaussianDistribution\nfrom tts.modules.wavvae.decoder.hifigan_modules import Generator, Upsample\n\n\nclass WavVAE_V3(nn.Module):\n def __init__(self, hparams=None):\n super().__init__()\n self.encoder = Encoder(dowmsamples=[6, 5, 4, 4, 2])\n self.proj_to_z = nn.Linear(512, 64)\n self.proj_to_decoder = nn.Linear(32, 320)\n\n config_path = hparams['melgan_config']\n args = argparse.Namespace()\n args.__dict__.update(config_path)\n self.latent_upsampler = Upsample(320, 4)\n self.decoder = Generator(\n input_size_=160, ngf=128, n_residual_layers=4,\n num_band=1, args=args, ratios=[5,4,4,3])\n\n ''' encode waveform into 25 hz latent representation '''\n def encode_latent(self, audio):\n posterior = self.encode(audio)\n latent = posterior.sample().permute(0, 2, 1) # (b,t,latent_channel)\n return latent\n\n def encode(self, audio):\n x = self.encoder(audio).permute(0, 2, 1)\n x = self.proj_to_z(x).permute(0, 2, 1)\n poseterior = DiagonalGaussianDistribution(x)\n return poseterior\n\n def decode(self, latent):\n latent = self.proj_to_decoder(latent).permute(0, 2, 1)\n return self.decoder(self.latent_upsampler(latent))\n\n def forward(self, audio):\n posterior = self.encode(audio)\n latent = posterior.sample().permute(0, 2, 1) # (b, t, latent_channel)\n recon_wav = self.decode(latent)\n return recon_wav, posterior" }, { "path": "tts/modules/wavvae/encoder/common_modules/conv.py", "content": "# MIT License\n\n# Copyright (c) Meta Platforms, Inc. and affiliates.\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\n# Copyright (c) [2023] [Meta Platforms, Inc. and affiliates.] \n# Copyright (c) [2025] [Ziyue Jiang] \n# SPDX-License-Identifier: MIT\n# This file has been modified by Ziyue Jiang on 2025/03/19\n# Original file was released under MIT, with the full license text # available at https://github.com/facebookresearch/encodec/blob/gh-pages/LICENSE.\n# This modified file is released under the same license.\n\n\"\"\"Convolutional layers wrappers and utilities.\"\"\"\n\nimport math\nimport typing as tp\nimport warnings\nimport einops\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.nn.utils import spectral_norm, weight_norm\n\n\nCONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',\n 'time_layer_norm', 'layer_norm', 'time_group_norm'])\n\n\ndef apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:\n assert norm in CONV_NORMALIZATIONS\n if norm == 'weight_norm':\n return weight_norm(module)\n elif norm == 'spectral_norm':\n return spectral_norm(module)\n else:\n return module\n\n\ndef get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:\n assert norm in CONV_NORMALIZATIONS\n if norm == 'layer_norm':\n assert isinstance(module, nn.modules.conv._ConvNd)\n return ConvLayerNorm(module.out_channels, **norm_kwargs)\n elif norm == 'time_group_norm':\n if causal:\n raise ValueError(\"GroupNorm doesn't support causal evaluation.\")\n assert isinstance(module, nn.modules.conv._ConvNd)\n return nn.GroupNorm(1, module.out_channels, **norm_kwargs)\n else:\n return nn.Identity()\n\n\ndef get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,\n padding_total: int = 0) -> int:\n length = x.shape[-1]\n n_frames = (length - kernel_size + padding_total) / stride + 1\n ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)\n return ideal_length - length\n\n\ndef pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):\n length = x.shape[-1]\n padding_left, padding_right = paddings\n assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)\n if mode == 'reflect':\n max_pad = max(padding_left, padding_right)\n extra_pad = 0\n if length <= max_pad:\n extra_pad = max_pad - length + 1\n x = F.pad(x, (0, extra_pad))\n padded = F.pad(x, paddings, mode, value)\n end = padded.shape[-1] - extra_pad\n return padded[..., :end]\n else:\n return F.pad(x, paddings, mode, value)\n\n\nclass ConvLayerNorm(nn.LayerNorm):\n def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):\n super().__init__(normalized_shape, **kwargs)\n\n def forward(self, x):\n x = einops.rearrange(x, 'b ... t -> b t ...')\n x = super().forward(x)\n x = einops.rearrange(x, 'b t ... -> b ... t')\n return\n\n\nclass NormConv1d(nn.Module):\n def __init__(self, *args, causal: bool = False, norm: str = 'none',\n norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):\n super().__init__()\n self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)\n self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)\n self.norm_type = norm\n\n def forward(self, x):\n x = self.conv(x)\n x = self.norm(x)\n return x\n\n\nclass SConv1d(nn.Module):\n def __init__(self, in_channels: int, out_channels: int,\n kernel_size: int, stride: int = 1, dilation: int = 1,\n groups: int = 1, bias: bool = True, causal: bool = False,\n norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},\n pad_mode: str = 'reflect'):\n super().__init__()\n # warn user on unusual setup between dilation and stride\n if stride > 1 and dilation > 1:\n warnings.warn('SConv1d has been initialized with stride > 1 and dilation > 1'\n f' (kernel_size={kernel_size} stride={stride}, dilation={dilation}).')\n self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,\n dilation=dilation, groups=groups, bias=bias, causal=causal,\n norm=norm, norm_kwargs=norm_kwargs)\n self.causal = causal\n self.pad_mode = pad_mode\n\n def forward(self, x):\n B, C, T = x.shape\n kernel_size = self.conv.conv.kernel_size[0]\n stride = self.conv.conv.stride[0]\n dilation = self.conv.conv.dilation[0]\n kernel_size = (kernel_size - 1) * dilation + 1 # effective kernel size with dilations\n padding_total = kernel_size - stride\n extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)\n if self.causal:\n # Left padding for causal\n x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)\n else:\n # Asymmetric padding required for odd strides\n padding_right = padding_total // 2\n padding_left = padding_total - padding_right\n x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)\n return self.conv(x)" }, { "path": "tts/modules/wavvae/encoder/common_modules/lstm.py", "content": "# MIT License\n\n# Copyright (c) Meta Platforms, Inc. and affiliates.\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\n# Copyright (c) [2023] [Meta Platforms, Inc. and affiliates.] \n# Copyright (c) [2025] [Ziyue Jiang] \n# SPDX-License-Identifier: MIT\n# This file has been modified by Ziyue Jiang on 2025/03/19\n# Original file was released under MIT, with the full license text # available at https://github.com/facebookresearch/encodec/blob/gh-pages/LICENSE.\n# This modified file is released under the same license.\n\n\"\"\"LSTM layers module.\"\"\"\nfrom torch import nn\n\n\nclass SLSTM(nn.Module):\n \"\"\"\n LSTM without worrying about the hidden state, nor the layout of the data.\n Expects input as convolutional layout.\n \"\"\"\n def __init__(self, dimension: int, num_layers: int = 2, skip: bool = True):\n super().__init__()\n self.skip = skip\n self.lstm = nn.LSTM(dimension, dimension, num_layers)\n\n # 修改transpose顺序\n def forward(self, x):\n x1 = x.permute(2, 0, 1)\n y, _ = self.lstm(x1)\n y = y.permute(1, 2, 0)\n if self.skip:\n y = y + x\n return y\n" }, { "path": "tts/modules/wavvae/encoder/common_modules/seanet.py", "content": "# MIT License\n\n# Copyright (c) Meta Platforms, Inc. and affiliates.\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\n# Copyright (c) [2023] [Meta Platforms, Inc. and affiliates.] \n# Copyright (c) [2025] [Ziyue Jiang] \n# SPDX-License-Identifier: MIT\n# This file has been modified by Ziyue Jiang on 2025/03/19\n# Original file was released under MIT, with the full license text # available at https://github.com/facebookresearch/encodec/blob/gh-pages/LICENSE.\n# This modified file is released under the same license.\n\n\"\"\"Encodec SEANet-based encoder and decoder implementation.\"\"\"\n\nimport typing as tp\n\nimport numpy as np\nimport torch.nn as nn\n\nfrom .conv import SConv1d\nfrom .lstm import SLSTM\n\n\nclass SEANetResnetBlock(nn.Module):\n def __init__(self, dim: int, kernel_sizes: tp.List[int] = [3, 1], dilations: tp.List[int] = [1, 1],\n activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},\n norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, causal: bool = False,\n pad_mode: str = 'reflect', compress: int = 2, true_skip: bool = True):\n super().__init__()\n assert len(kernel_sizes) == len(dilations), 'Number of kernel sizes should match number of dilations'\n act = getattr(nn, activation)\n hidden = dim // compress\n block = []\n for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):\n in_chs = dim if i == 0 else hidden\n out_chs = dim if i == len(kernel_sizes) - 1 else hidden\n block += [\n act(**activation_params),\n SConv1d(in_chs, out_chs, kernel_size=kernel_size, dilation=dilation,\n norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode),\n ]\n self.block = nn.Sequential(*block)\n self.shortcut: nn.Module\n if true_skip:\n self.shortcut = nn.Identity()\n else:\n self.shortcut = SConv1d(dim, dim, kernel_size=1, norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode)\n\n def forward(self, x):\n return self.shortcut(x) + self.block(x)\n\n\nclass SEANetEncoder(nn.Module):\n def __init__(self, channels: int = 1, dimension: int = 128, n_filters: int = 32, n_residual_layers: int = 1,\n ratios: tp.List[int] = [8, 5, 4, 2], activation: str = 'ELU', activation_params: dict = {'alpha': 1.0},\n norm: str = 'weight_norm', norm_params: tp.Dict[str, tp.Any] = {}, kernel_size: int = 7,\n last_kernel_size: int = 7, residual_kernel_size: int = 3, dilation_base: int = 2, causal: bool = False,\n pad_mode: str = 'reflect', true_skip: bool = False, compress: int = 2, lstm: int = 2):\n super().__init__()\n self.channels = channels\n self.dimension = dimension\n self.n_filters = n_filters\n self.ratios = list(reversed(ratios))\n del ratios\n self.n_residual_layers = n_residual_layers\n self.hop_length = np.prod(self.ratios)\n\n act = getattr(nn, activation)\n mult = 1\n model: tp.List[nn.Module] = [\n SConv1d(channels, mult * n_filters, kernel_size, norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode)\n ]\n # Downsample to raw audio scale\n for i, ratio in enumerate(self.ratios):\n # Add residual layers\n for j in range(n_residual_layers):\n model += [\n SEANetResnetBlock(mult * n_filters, kernel_sizes=[residual_kernel_size, 1],\n dilations=[dilation_base ** j, 1],\n norm=norm, norm_params=norm_params,\n activation=activation, activation_params=activation_params,\n causal=causal, pad_mode=pad_mode, compress=compress, true_skip=true_skip)]\n\n # Add downsampling layers\n model += [\n act(**activation_params),\n SConv1d(mult * n_filters, mult * n_filters * 2,\n kernel_size=ratio * 2, stride=ratio,\n norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode),\n ]\n mult *= 2\n\n if lstm:\n model += [SLSTM(mult * n_filters, num_layers=lstm)]\n\n model += [\n act(**activation_params),\n SConv1d(mult * n_filters, dimension, last_kernel_size, norm=norm, norm_kwargs=norm_params,\n causal=causal, pad_mode=pad_mode)\n ]\n\n self.model = nn.Sequential(*model)\n\n def forward(self, x):\n return self.model(x)" }, { "path": "tts/utils/audio_utils/align.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\n\ndef mel2token_to_dur(mel2token, T_txt=None, max_dur=None):\n is_torch = isinstance(mel2token, torch.Tensor)\n has_batch_dim = True\n if not is_torch:\n mel2token = torch.LongTensor(mel2token)\n if T_txt is None:\n T_txt = mel2token.max()\n if len(mel2token.shape) == 1:\n mel2token = mel2token[None, ...]\n has_batch_dim = False\n B, _ = mel2token.shape\n dur = mel2token.new_zeros(B, T_txt + 1).scatter_add(1, mel2token, torch.ones_like(mel2token))\n dur = dur[:, 1:]\n if max_dur is not None:\n dur = dur.clamp(max=max_dur)\n if not is_torch:\n dur = dur.numpy()\n if not has_batch_dim:\n dur = dur[0]\n return dur\n" }, { "path": "tts/utils/audio_utils/io.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 io\nimport os\nimport subprocess\n\nimport numpy as np\nfrom scipy.io import wavfile\nimport pyloudnorm as pyln\nfrom pydub import AudioSegment\n\n\ndef to_wav_bytes(wav, sr, norm=False):\n wav = wav.astype(float)\n if norm:\n meter = pyln.Meter(sr) # create BS.1770 meter\n loudness = meter.integrated_loudness(wav)\n wav = pyln.normalize.loudness(wav, loudness, -18.0)\n if np.abs(wav).max() >= 1:\n wav = wav / np.abs(wav).max() * 0.95\n wav = wav * 32767\n bytes_io = io.BytesIO()\n wavfile.write(bytes_io, sr, wav.astype(np.int16))\n return bytes_io.getvalue()\n\n\ndef save_wav(wav_bytes, path):\n with open(path[:-4] + '.wav', 'wb') as file:\n file.write(wav_bytes)\n if path[-4:] == '.mp3':\n to_mp3(path[:-4])\n\n\ndef to_mp3(out_path):\n if out_path[-4:] == '.wav':\n out_path = out_path[:-4]\n subprocess.check_call(\n f'ffmpeg -threads 1 -loglevel error -i \"{out_path}.wav\" -vn -b:a 192k -y -hide_banner -async 1 \"{out_path}.mp3\"',\n shell=True, stdin=subprocess.PIPE)\n subprocess.check_call(f'rm -f \"{out_path}.wav\"', shell=True)\n\n\ndef convert_to_wav(wav_path):\n # Check if the file exists\n if not os.path.exists(wav_path):\n print(f\"The file '{wav_path}' does not exist.\")\n return\n\n # Check if the file already has a .wav extension\n if not wav_path.endswith(\".wav\"):\n # Define the output path with a .wav extension\n out_path = os.path.splitext(wav_path)[0] + \".wav\"\n\n # Load the audio file using pydub and convert it to WAV\n audio = AudioSegment.from_file(wav_path)\n audio.export(out_path, format=\"wav\")\n\n print(f\"Converted '{wav_path}' to '{out_path}'\")\n\n\ndef convert_to_wav_bytes(audio_binary):\n # Load the audio binary using pydub and convert it to WAV\n audio = AudioSegment.from_file(io.BytesIO(audio_binary))\n wav_bytes = io.BytesIO()\n audio.export(wav_bytes, format=\"wav\")\n wav_bytes.seek(0)\n return wav_bytes\n\n\n''' Smoothly combine audio segments using crossfade transitions.\" '''\ndef combine_audio_segments(segments, crossfade_duration=0.16, sr=24000):\n window_length = int(sr * crossfade_duration)\n hanning_window = np.hanning(2 * window_length)\n # Combine\n for i, segment in enumerate(segments):\n if i == 0:\n combined_audio = segment\n else:\n overlap = combined_audio[-window_length:] * hanning_window[window_length:] + segment[:window_length] * hanning_window[:window_length]\n combined_audio = np.concatenate(\n [combined_audio[:-window_length], overlap, segment[window_length:]]\n )\n return combined_audio" }, { "path": "tts/utils/audio_utils/plot.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 matplotlib\n\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\n\nLINE_COLORS = ['w', 'r', 'orange', 'k', 'cyan', 'm', 'b', 'lime', 'g', 'brown', 'navy']\n\n\ndef spec_to_figure(spec, vmin=None, vmax=None, title='', f0s=None, dur_info=None, figsize=(12, 6)):\n if isinstance(spec, torch.Tensor):\n spec = spec.cpu().numpy()\n H = spec.shape[1] // 2\n fig = plt.figure(figsize=figsize)\n plt.title(title)\n plt.pcolor(spec.T, vmin=vmin, vmax=vmax)\n\n if dur_info is not None:\n assert isinstance(dur_info, dict)\n txt = dur_info['txt']\n dur_gt = dur_info['dur_gt']\n if isinstance(dur_gt, torch.Tensor):\n dur_gt = dur_gt.cpu().numpy()\n dur_gt = np.cumsum(dur_gt).astype(int)\n for i in range(len(dur_gt)):\n shift = (i % 8) + 1\n plt.text(dur_gt[i], shift * 4, txt[i])\n plt.vlines(dur_gt[i], 0, H // 2, colors='b') # blue is gt\n plt.xlim(0, dur_gt[-1])\n if 'dur_pred' in dur_info:\n dur_pred = dur_info['dur_pred']\n if isinstance(dur_pred, torch.Tensor):\n dur_pred = dur_pred.cpu().numpy()\n dur_pred = np.cumsum(dur_pred).astype(int)\n for i in range(len(dur_pred)):\n shift = (i % 8) + 1\n plt.text(dur_pred[i], H + shift * 4, txt[i])\n plt.vlines(dur_pred[i], H, H * 1.5, colors='r') # red is pred\n plt.xlim(0, max(dur_gt[-1], dur_pred[-1]))\n if f0s is not None:\n ax = plt.gca()\n ax2 = ax.twinx()\n # ax.set_xticks()\n\n if not isinstance(f0s, dict):\n f0s = {'f0': f0s}\n for i, (k, f0) in enumerate(f0s.items()):\n if f0 is not None:\n if isinstance(f0, torch.Tensor):\n f0 = f0.cpu().numpy()\n ax2.plot(\n np.arange(len(f0)) + 0.5, f0, label=k, c=LINE_COLORS[i], linewidth=1, alpha=0.5)\n ax2.set_ylim(0, 1000)\n ax2.legend()\n return fig\n\n\ndef align_to_figure(align, dur_info):\n if isinstance(align, torch.Tensor):\n align = align.cpu().numpy()\n H = align.shape[1]\n fig = plt.figure(figsize=(12, 6))\n plt.pcolor(align.T, vmin=0, vmax=1)\n if dur_info is not None:\n assert isinstance(dur_info, dict)\n txt = dur_info['txt']\n dur_gt = dur_info['dur_gt']\n if isinstance(dur_gt, torch.Tensor):\n dur_gt = dur_gt.cpu().numpy()\n dur_gt = np.cumsum(dur_gt).astype(int) // 2\n for i in range(len(dur_gt)):\n plt.text(dur_gt[i], i, txt[i], color='red')\n plt.vlines(dur_gt[i], 0, H, colors='b') # blue is gt\n # plt.xlim(0, dur_gt[-1])\n return fig\n" }, { "path": "tts/utils/commons/ckpt_utils.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 contextlib\nimport glob\nimport os\nimport re\nimport subprocess\nimport traceback\n\nimport torch\nfrom torch.nn.parallel import DistributedDataParallel\nimport torch.distributed as dist\n\n\n@contextlib.contextmanager\ndef dist_load(path):\n if not dist.is_initialized() or dist.get_world_size() == 1 or os.path.realpath(path).startswith('/dev/shm'):\n yield path\n else:\n from tts.utils.commons.hparams import hparams\n from tts.utils.commons.trainer import LOCAL_RANK\n tmpdir = '/dev/shm'\n assert len(os.path.basename(path)) > 0\n shm_ckpt_path = f'{tmpdir}/{hparams[\"exp_name\"]}/{os.path.basename(path)}'\n if LOCAL_RANK == 0:\n subprocess.check_call(\n f'mkdir -p {os.path.dirname(shm_ckpt_path)}; '\n f'cp -Lr {path} {shm_ckpt_path}', shell=True)\n dist.barrier()\n yield shm_ckpt_path\n dist.barrier()\n if LOCAL_RANK == 0:\n subprocess.check_call(f'rm -rf {shm_ckpt_path}', shell=True)\n\n\ndef torch_load_dist(path, map_location='cpu'):\n with dist_load(path) as tmp_path:\n checkpoint = torch.load(tmp_path, map_location=map_location)\n return checkpoint\n\n\ndef get_last_checkpoint(work_dir, steps=None):\n checkpoint = None\n last_ckpt_path = None\n ckpt_paths = get_all_ckpts(work_dir, steps)\n if len(ckpt_paths) > 0:\n last_ckpt_path = ckpt_paths[0]\n checkpoint = torch_load_dist(last_ckpt_path, map_location='cpu')\n return checkpoint, last_ckpt_path\n\n\ndef get_all_ckpts(work_dir, steps=None):\n if steps is None or steps == 0:\n ckpt_path_pattern = f'{work_dir}/model_ckpt_steps_*.ckpt'\n else:\n ckpt_path_pattern = f'{work_dir}/model_ckpt_steps_{steps}.ckpt'\n return sorted(glob.glob(ckpt_path_pattern),\n key=lambda x: -int(re.findall('.*steps\\_(\\d+)\\.ckpt', x)[0]))\n\n\ndef load_ckpt(cur_model, ckpt_base_dir, model_name='model', force=True, strict=True,\n silent=False, load_opt=False, opts=None, steps=None, checkpoint=None, ckpt_path='', delete_unmatch=True):\n if checkpoint is None:\n if os.path.isfile(ckpt_base_dir):\n base_dir = os.path.dirname(ckpt_base_dir)\n ckpt_path = ckpt_base_dir\n checkpoint = torch_load_dist(ckpt_base_dir, map_location='cpu')\n else:\n base_dir = ckpt_base_dir\n if load_opt:\n checkpoint, ckpt_path = get_last_checkpoint(ckpt_base_dir, steps)\n else:\n ckpt_path = f'{ckpt_base_dir}/model_only_last.ckpt'\n if os.path.exists(ckpt_path):\n checkpoint = torch_load_dist(ckpt_path, map_location='cpu')\n else:\n checkpoint, ckpt_path = get_last_checkpoint(ckpt_base_dir, steps)\n if checkpoint is not None:\n state_dict_all = {\n k.replace('module.', '').replace('_orig_mod.', ''): v for k, v in checkpoint[\"state_dict\"].items()}\n if not isinstance(cur_model, list):\n cur_models = [cur_model]\n model_names = [model_name]\n else:\n cur_models = cur_model\n model_names = model_name\n for model_name, cur_model in zip(model_names, cur_models):\n if isinstance(cur_model, DistributedDataParallel):\n cur_model = cur_model.module\n device = next(cur_model.parameters()).device\n if '.' not in model_name:\n state_dict = state_dict_all[model_name]\n else:\n base_model_name = model_name.split('.')[0]\n rest_model_name = model_name[len(base_model_name) + 1:]\n state_dict = {\n k[len(rest_model_name) + 1:]: v for k, v in state_dict_all[base_model_name].items()\n if k.startswith(f'{rest_model_name}.')}\n state_dict = {k.replace('module.', '').replace('_orig_mod.', ''): v for k, v in state_dict.items()}\n if not strict and delete_unmatch:\n try:\n cur_model.load_state_dict(state_dict, strict=True)\n if not silent:\n print(f\"| loaded '{model_name}' from '{ckpt_path}' with strict=True.\")\n except:\n cur_model_state_dict = cur_model.state_dict()\n cur_model_state_dict = {k.replace('module.', '').replace('_orig_mod.', ''): v for k, v in\n cur_model_state_dict.items()}\n unmatched_keys = []\n for key, param in state_dict.items():\n if key in cur_model_state_dict:\n new_param = cur_model_state_dict[key]\n if new_param.shape != param.shape:\n unmatched_keys.append(key)\n print(\"| Unmatched keys: \", key, \"cur model: \", new_param.shape,\n \"ckpt model: \", param.shape)\n for key in unmatched_keys:\n del state_dict[key]\n load_results = cur_model.load_state_dict(state_dict, strict=strict)\n cur_model.to(device)\n if not silent:\n print(f\"| loaded '{model_name}' from '{ckpt_path}'.\")\n missing_keys, unexpected_keys = load_results.missing_keys, load_results.unexpected_keys\n print(f\"| Missing keys: {len(missing_keys)}, Unexpected keys: {len(unexpected_keys)}\")\n if load_opt:\n optimizer_states = checkpoint['optimizer_states']\n assert len(opts) == len(optimizer_states)\n for optimizer, opt_state in zip(opts, optimizer_states):\n opt_state = {k.replace('_orig_mod.', ''): v for k, v in opt_state.items()}\n if optimizer is None:\n return\n try:\n optimizer.load_state_dict(opt_state)\n for i, state in enumerate(optimizer.state.values()):\n for k, v in state.items():\n if isinstance(v, torch.Tensor):\n state[k] = v.to(device)\n except ValueError:\n print(f\"| WARMING: optimizer {optimizer} parameters not match !!!\")\n return checkpoint.get('global_step', 0)\n else:\n e_msg = f\"| ckpt not found in {base_dir}.\"\n if force:\n assert False, e_msg\n else:\n print(e_msg)\n\n\ndef load_with_size_mismatch(model, state_dict, prefix=\"\"):\n current_model_dict = model.state_dict()\n cm_keys = current_model_dict.keys()\n mismatch_keys = {k.replace(prefix, \"\") for k, v in state_dict.items() if k.replace(prefix, \"\") in cm_keys and v.size() != current_model_dict[k.replace(prefix, \"\")].size()}\n new_state_dict = {k.replace(prefix, \"\"): v for k, v in state_dict.items() if k.replace(prefix, \"\") in cm_keys and v.size() == current_model_dict[k.replace(prefix, \"\")].size()}\n missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)\n print(f\"| mismatch keys: \", mismatch_keys)\n if len(missing_keys) > 0:\n print(f\"| missing_keys in dit: {missing_keys}\")\n if len(unexpected_keys) > 0:\n print(f\"| unexpected_keys in dit: {unexpected_keys}\")\n" }, { "path": "tts/utils/commons/hparams.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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\nimport json\nimport os\nimport re\nimport ast\nimport yaml\n\nglobal_print_hparams = True\nhparams = {}\n\n\nclass Args:\n def __init__(self, **kwargs):\n for k, v in kwargs.items():\n self.__setattr__(k, v)\n\n\ndef override_config(old_config: dict, new_config: dict):\n if new_config.get('__replace', False):\n old_config.clear()\n for k, v in new_config.items():\n if isinstance(v, dict) and k in old_config:\n override_config(old_config[k], new_config[k])\n else:\n old_config[k] = v\n\n\ndef traverse_dict(d, func, ctx):\n for k in list(d.keys()):\n v = d[k]\n if isinstance(v, dict):\n traverse_dict(v, func, ctx)\n else:\n d[k] = func(v, ctx)\n\n\ndef parse_config(v, context=None):\n if context is None:\n context = {}\n\n if isinstance(v, str):\n if v.startswith('^'):\n return load_config(v[1:], [], set())\n\n match = re.match(r\"\\${(.*)}\", v)\n if match:\n expression = match.group(1)\n return ast.literal_eval(expression, {}, context)\n return v\n\n\ndef remove_meta_key(d):\n for k in list(d.keys()):\n v = d[k]\n if isinstance(v, dict):\n remove_meta_key(v)\n else:\n if k[:2] == '__':\n del d[k]\n\n\ndef load_config(config_fn, config_chains, loaded_configs):\n # deep first inheritance and avoid the second visit of one node\n if not os.path.exists(config_fn):\n print(f\"| WARN: {config_fn} not exist.\", )\n return {}\n with open(config_fn) as f:\n hparams_ = yaml.safe_load(f)\n loaded_configs.add(config_fn)\n\n if 'base_config' in hparams_:\n ret_hparams = {}\n if not isinstance(hparams_['base_config'], list):\n hparams_['base_config'] = [hparams_['base_config']]\n for c in hparams_['base_config']:\n if c.startswith('.'):\n c = f'{os.path.dirname(config_fn)}/{c}'\n c = os.path.normpath(c)\n if c not in loaded_configs:\n override_config(ret_hparams, load_config(c, config_chains, loaded_configs))\n override_config(ret_hparams, hparams_)\n else:\n ret_hparams = hparams_\n\n config_chains.append(config_fn)\n return ret_hparams\n\n\ndef set_hparams(config='', exp_name='', hparams_str='', print_hparams=True, global_hparams=True):\n if config == '' and exp_name == '':\n parser = argparse.ArgumentParser(description='')\n parser.add_argument('--config', type=str, default='',\n help='location of the data corpus')\n parser.add_argument('--exp_name', type=str, default='', help='exp_name')\n parser.add_argument('-hp', '--hparams', type=str, default='',\n help='location of the data corpus')\n parser.add_argument('--infer', action='store_true', help='infer')\n parser.add_argument('--validate', action='store_true', help='validate')\n parser.add_argument('--reset', action='store_true', help='reset hparams')\n parser.add_argument('--remove', action='store_true', help='remove old ckpt')\n parser.add_argument('--debug', action='store_true', help='debug')\n parser.add_argument('--start_rank', type=int, default=-1,\n help='the start rank id for DDP, keep 0 when single-machine multi-GPU')\n parser.add_argument('--world_size', type=int, default=-1,\n help='the total number of GPU used across all machines, keep -1 for single-machine multi-GPU')\n parser.add_argument('--init_method', type=str, default='tcp', help='method to init ddp, use tcp or file')\n parser.add_argument('--master_addr', type=str, default='', help='')\n parser.add_argument('--ddp_dir', type=str, default='', help='')\n\n args, unknown = parser.parse_known_args()\n if print_hparams:\n print(\"| set_hparams Unknow hparams: \", unknown)\n else:\n args = Args(config=config, exp_name=exp_name, hparams=hparams_str,\n infer=False, validate=False, reset=False, debug=False, remove=False,\n start_rank=-1, world_size=-1, init_method='tcp', ddp_dir='', master_addr='')\n global hparams\n assert args.config != '' or args.exp_name != ''\n if args.config != '':\n assert os.path.exists(args.config), f\"{args.config} not exists\"\n\n saved_hparams = {}\n args_work_dir = ''\n if args.exp_name != '':\n args_work_dir = f'{args.exp_name}'\n ckpt_config_path = f'{args_work_dir}/config.yaml'\n if os.path.exists(ckpt_config_path):\n with open(ckpt_config_path) as f:\n saved_hparams_ = yaml.safe_load(f)\n if saved_hparams_ is not None:\n saved_hparams.update(saved_hparams_)\n hparams_ = {}\n config_chains = []\n if args.config != '':\n hparams_.update(load_config(args.config, config_chains, set()))\n if len(config_chains) > 1 and print_hparams:\n print('| Hparams chains: ', config_chains)\n if not args.reset:\n hparams_.update(saved_hparams)\n traverse_dict(hparams_, parse_config, hparams_)\n hparams_['work_dir'] = args_work_dir\n\n # Support config overriding in command line. Support list type config overriding.\n # Examples: --hparams=\"a=1,b.c=2,d=[1 1 1]\"\n if args.hparams != \"\":\n for new_hparam in args.hparams.split(\",\"):\n k, v = new_hparam.split(\"=\")\n v = v.strip(\"\\'\\\" \")\n config_node = hparams_\n for k_ in k.split(\".\")[:-1]:\n config_node = config_node[k_]\n k = k.split(\".\")[-1]\n if k in config_node:\n if v in ['True', 'False'] or type(config_node[k]) in [bool, list, dict]:\n if type(config_node[k]) == list:\n v = v.replace(\" \", \",\").replace('^', \"\\\"\")\n if '|' in v:\n tp = type(config_node[k][0]) if len(config_node[k]) else str\n config_node[k] = [tp(x) for x in v.split(\"|\") if x != '']\n continue\n config_node[k] = ast.literal_eval(v)\n else:\n config_node[k] = type(config_node[k])(v)\n else:\n config_node[k] = v\n try:\n config_node[k] = float(v)\n except:\n pass\n try:\n config_node[k] = int(v)\n except:\n pass\n if v.lower() in ['false', 'true']:\n config_node[k] = v.lower() == 'true'\n\n if args_work_dir != '' and not args.infer:\n os.makedirs(hparams_['work_dir'], exist_ok=True)\n\n hparams_['infer'] = args.infer\n hparams_['debug'] = args.debug\n hparams_['validate'] = args.validate\n hparams_['exp_name'] = args.exp_name\n\n hparams_['start_rank'] = args.start_rank # useful for multi-machine training\n hparams_['world_size'] = args.world_size\n hparams_['init_method'] = args.init_method\n hparams_['ddp_dir'] = args.ddp_dir\n hparams_['master_addr'] = args.master_addr\n\n remove_meta_key(hparams_)\n global global_print_hparams\n if global_hparams:\n hparams.clear()\n hparams.update(hparams_)\n if print_hparams and global_print_hparams and global_hparams:\n print('| Hparams: ', json.dumps(hparams_, indent=2, sort_keys=True))\n # for i, (k, v) in enumerate(sorted(hparams_.items())):\n # print(f\"\\033[;33;m{k}\\033[0m: {v}, \", end=\"\\n\" if i % 5 == 4 else \"\")\n global_print_hparams = False\n return hparams_" }, { "path": "tts/utils/text_utils/dict.json", "content": "{\"phone\": [\"C0a\", \"C0ai\", \"C0air\", \"C0an\", \"C0ang\", \"C0angr\", \"C0anr\", \"C0ao\", \"C0aor\", \"C0ar\", \"C0b\", \"C0c\", \"C0ch\", \"C0d\", \"C0e\", \"C0ei\", \"C0eir\", \"C0en\", \"C0eng\", \"C0engr\", \"C0enr\", \"C0er\", \"C0f\", \"C0g\", \"C0h\", \"C0i\", \"C0ia\", \"C0ian\", \"C0iang\", \"C0iangr\", \"C0ianr\", \"C0iao\", \"C0iaor\", \"C0iar\", \"C0ie\", \"C0ier\", \"C0ii\", \"C0iii\", \"C0iiir\", \"C0iir\", \"C0in\", \"C0ing\", \"C0ingr\", \"C0inr\", \"C0io\", \"C0iong\", \"C0iongr\", \"C0iou\", \"C0iour\", \"C0ir\", \"C0j\", \"C0k\", \"C0l\", \"C0m\", \"C0n\", \"C0ng\", \"C0o\", \"C0ong\", \"C0ongr\", \"C0or\", \"C0ou\", \"C0our\", \"C0p\", \"C0q\", \"C0r\", \"C0s\", \"C0sh\", \"C0t\", \"C0u\", \"C0ua\", \"C0uai\", \"C0uair\", \"C0uan\", \"C0uang\", \"C0uangr\", \"C0uanr\", \"C0uar\", \"C0uei\", \"C0ueir\", \"C0uen\", \"C0ueng\", \"C0uengr\", \"C0uenr\", \"C0uo\", \"C0uor\", \"C0ur\", \"C0v\", \"C0van\", \"C0vanr\", \"C0ve\", \"C0ver\", \"C0vn\", \"C0vnr\", \"C0vr\", \"C0x\", \"C0z\", \"C0zh\", \"C0＿\", \"E0aa\", \"E0ae\", \"E0ah\", \"E0ao\", \"E0aw\", \"E0ax\", \"E0ay\", \"E0b\", \"E0ch\", \"E0d\", \"E0dh\", \"E0eh\", \"E0ehr\", \"E0er\", \"E0ey\", \"E0f\", \"E0g\", \"E0hh\", \"E0ih\", \"E0iy\", \"E0iyr\", \"E0jh\", \"E0k\", \"E0l\", \"E0m\", \"E0n\", \"E0ng\", \"E0oh\", \"E0ow\", \"E0oy\", \"E0p\", \"E0r\", \"E0s\", \"E0sh\", \"E0t\", \"E0th\", \"E0uh\", \"E0uw\", \"E0uwr\", \"E0v\", \"E0w\", \"E0y\", \"E0z\", \"E0zh\", \"sil\", \"…\", \"、\", \"。\", \"《\", \"》\", \"【\", \"】\", \"！\", \"＂\", \"＃\", \"＄\", \"％\", \"＇\", \"＇＇\", \"（\", \"）\", \"＊\", \"，\", \"：\", \"；\", \"？\", \"＼\", \"＾\", \"＿\", \"｀\", \"｛\", \"｝\", \"～\"], \"tone\": [\"0\", \"1\", \"10\", \"11\", \"12\", \"13\", \"15\", \"17\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\", \"9\"], \"wordCategory\": [\"0\", \"B\", \"E\", \"M\", \"S\"], \"prosody\": [\"0\", \"1\", \"2\", \"3\", \"4\"], \"focus\": [\"0\", \"1\"], \"intonation\": [\"0\", \"1\", \"2\"], \"phraseAccent\": [\"0\", \"H-\", \"L-\"], \"boundaryTone\": [\"0\", \"H%\", \"L%\"], \"accentType\": [\"!H*\", \"0\", \"H*\", \"L*\", \"L*+H\", \"L+H*\"]}\n" }, { "path": "tts/utils/text_utils/ph_tone_convert.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 torch\nimport torch.nn.functional as F\n\ndef map_phone_to_tokendict(item, pad_bos_eos=True):\n # Merge Chinese phone and tone (Original dict ends at 173, i.e., ph_dict_size=173). 146~173 is punctuations.\n phone = item['txt_token'].clone()\n merged_phone = item['txt_token'].clone()\n tone_tmp = item['tone'].clone()\n # In tone_dict, tone_1 is 4, tone_2 is 11, tone_3 is 12, tone_4 is 13, tone_5 is 14, tone_6 is 15\n tone_tmp[tone_tmp==4] = 1\n tone_tmp[tone_tmp==11] = 2\n tone_tmp[tone_tmp==12] = 3\n tone_tmp[tone_tmp==13] = 4\n tone_tmp[tone_tmp==14] = 5\n tone_tmp[tone_tmp==15] = 6\n # Chinese phones lie in 3~100 in the phone_dict, we map them to 200~788\n ch_phone_idx = (phone >= 3) & (phone <= 100)\n merged_phone[ch_phone_idx] = (merged_phone[ch_phone_idx] - 3) * 6 + 200 + tone_tmp[ch_phone_idx]\n\n if pad_bos_eos:\n merged_phone = F.pad(merged_phone, (1, 0), mode='constant', value=798)\n merged_phone = F.pad(merged_phone, (0, 1), mode='constant', value=799)\n return merged_phone\n \ndef split_ph_timestamp(ph_timestamp):\n ''' Input: ph_timestamp, shape [T] '''\n\n # Map the timestamp of each phone back to its original frame-level lengths\n ph_timestamp[ph_timestamp >= 800] -= 800\n\n ph_list = []\n tone_list = []\n dur_list = []\n cur_timestamp = 0\n for idx, item in enumerate(ph_timestamp):\n if idx % 2 == 0:\n # Map Chinese phones back to its original phone_dict\n if (200 <= item <= 788):\n ph = (item - 200 - 1) // 6 + 3\n tone = (item - 200 - 1) % 6 + 1\n if tone == 1:\n tone = 4\n else:\n tone = tone + 9\n # Set English tone to '3'\n else:\n ph = item\n tone = 3\n ph_list.append(ph)\n tone_list.append(tone)\n else:\n dur_list.append((item - cur_timestamp))\n cur_timestamp = item\n assert len(ph_list) == len(dur_list), f\"{len(ph_list)}, {len(dur_list)}\"\n ph_seq, tone_seq, dur_seq = torch.LongTensor(ph_list), torch.LongTensor(tone_list), torch.LongTensor(dur_list)\n return ph_seq, tone_seq, dur_seq, ph_timestamp[-1]\n \ndef split_ph(ph_seq):\n ''' Input: ph_timestamp, shape [T] '''\n ph_list = []\n tone_list = []\n for idx, item in enumerate(ph_seq):\n # Map Chinese phones back to its original phone_dict\n if (200 <= item <= 788):\n ph = (item - 200 - 1) // 6 + 3\n tone = (item - 200 - 1) % 6 + 1\n if tone == 1:\n tone = 4\n else:\n tone = tone + 9\n # Set English tone to '3'\n else:\n ph = item\n tone = 3\n ph_list.append(ph)\n tone_list.append(tone)\n \n assert len(ph_list) == len(tone_list)\n ph_seq, tone_seq = torch.LongTensor(ph_list), torch.LongTensor(tone_list)\n return ph_seq, tone_seq" }, { "path": "tts/utils/text_utils/split_text.py", "content": "# -*- coding: utf-8 -*-\n# Copyright 2025 ByteDance and/or its affiliates.\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 re\n\ndef chunk_text_chinese(text, limit=60):\n # 中文字符匹配\n chinese_pattern = re.compile(r'[\\u4e00-\\u9fff]')\n # 标点符号匹配\n punctuation = r\"，。！？；：,.!?;:\"\n \n result = [] # 存储断句结果\n current_chunk = [] # 当前片段\n chinese_count = 0 # 中文字符计数\n\n i = 0\n while i < len(text):\n char = text[i]\n current_chunk.append(char)\n if chinese_pattern.match(char):\n chinese_count += 1\n \n if chinese_count >= limit: # 达到限制字符数\n # 从当前位置往前找最近的标点符号\n for j in range(len(current_chunk) - 1, -1, -1):\n if current_chunk[j] in punctuation:\n result.append(''.join(current_chunk[:j + 1]))\n current_chunk = current_chunk[j + 1:]\n chinese_count = sum(1 for c in current_chunk if chinese_pattern.match(c))\n break\n else:\n # 如果前面没有标点符号，则继续找后面的标点符号\n for k in range(i + 1, len(text)):\n if text[k] in punctuation:\n result.append(''.join(current_chunk)+text[i+1:k+1])\n current_chunk = []\n chinese_count = 0\n i = k\n break\n i+=1\n\n # 添加最后剩余的部分\n if current_chunk:\n result.append(''.join(current_chunk))\n\n return result\n\ndef chunk_text_english(text, max_chars=130):\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\ndef chunk_text_chinesev2(text, limit=60, look_ahead_limit=30):\n \"\"\"\n 将中文文本分成多个块，优先确保每个块以句号、感叹号或问号结尾，\n 其次考虑逗号等其他标点符号，避免在无标点处断句\n \n 参数:\n text: 要分块的文本\n limit: 每个块的中文字符数限制\n look_ahead_limit: 向后查找的最大字符数限制\n \n 返回:\n 分块后的文本列表\n \"\"\"\n # 中文字符匹配\n chinese_pattern = re.compile(r'[\\u4e00-\\u9fff]')\n \n # 分级定义标点符号（优先级从高到低）\n primary_end_marks = \"。.！!？?\" # 首选：句号、感叹号、问号\n secondary_end_marks = \"，,；;：\" # 次选：逗号、分号、冒号\n tertiary_end_marks = \"、…—-~～\" # 再次：顿号、省略号、破折号等\n \n result = [] # 存储断句结果\n current_chunk = [] # 当前片段\n chinese_count = 0 # 中文字符计数\n\n i = 0\n while i < len(text):\n char = text[i]\n current_chunk.append(char)\n\n if chinese_pattern.match(char):\n chinese_count += 1\n \n if chinese_count >= limit: # 达到字符数限制，需要寻找断句点\n found_end = False\n \n # 依次尝试不同优先级的断句策略\n \n # 1. 向后查找首选标点\n for k in range(1, min(look_ahead_limit, len(text) - i)):\n next_char = text[i + k]\n if next_char in primary_end_marks:\n result.append(''.join(current_chunk) + text[i+1:i+k+1])\n current_chunk = []\n chinese_count = 0\n i = i + k\n found_end = True\n break\n \n if not found_end:\n # 2. 向前查找首选标点\n for j in range(len(current_chunk) - 1, -1, -1):\n if current_chunk[j] in primary_end_marks:\n result.append(''.join(current_chunk[:j + 1]))\n current_chunk = current_chunk[j + 1:]\n chinese_count = sum(1 for c in current_chunk if chinese_pattern.match(c))\n found_end = True\n break\n \n if not found_end:\n # 3. 向后查找次选标点\n for k in range(1, min(look_ahead_limit, len(text) - i)):\n next_char = text[i + k]\n if next_char in secondary_end_marks:\n result.append(''.join(current_chunk) + text[i+1:i+k+1])\n current_chunk = []\n chinese_count = 0\n i = i + k\n found_end = True\n break\n \n if not found_end:\n # 4. 向前查找次选标点\n for j in range(len(current_chunk) - 1, -1, -1):\n if current_chunk[j] in secondary_end_marks:\n result.append(''.join(current_chunk[:j + 1]))\n current_chunk = current_chunk[j + 1:]\n chinese_count = sum(1 for c in current_chunk if chinese_pattern.match(c))\n found_end = True\n break\n \n if not found_end:\n # 5. 向后查找三级标点\n for k in range(1, min(look_ahead_limit, len(text) - i)):\n next_char = text[i + k]\n if next_char in tertiary_end_marks:\n result.append(''.join(current_chunk) + text[i+1:i+k+1])\n current_chunk = []\n chinese_count = 0\n i = i + k\n found_end = True\n break\n \n if not found_end:\n # 6. 向前查找三级标点\n for j in range(len(current_chunk) - 1, -1, -1):\n if current_chunk[j] in tertiary_end_marks:\n result.append(''.join(current_chunk[:j + 1]))\n current_chunk = current_chunk[j + 1:]\n chinese_count = sum(1 for c in current_chunk if chinese_pattern.match(c))\n found_end = True\n break\n \n if not found_end:\n # 万不得已，在此处断句（这种情况很少见，因为汉语文本中通常会有标点）\n result.append(''.join(current_chunk))\n current_chunk = []\n chinese_count = 0\n \n i += 1\n\n # 添加最后剩余的部分\n if current_chunk:\n result.append(''.join(current_chunk))\n\n # 英文标点替换为中文标点\n punctuation_map = {\n '.': '。',\n ',': '，',\n '!': '！',\n '?': '？',\n ';': '；',\n ':': '：'\n }\n \n for i in range(len(result)):\n for eng_punc, cn_punc in punctuation_map.items():\n result[i] = result[i].replace(eng_punc, cn_punc)\n\n return result\n\nif __name__ == '__main__':\n print(chunk_text_chinese(\"哇塞！家人们，你们太好运了。我居然发现了一个宝藏零食大礼包，简直适合所有人的口味！有香辣的，让你舌尖跳舞；有盐焗的，咸香可口；还有五香的，香气四溢。就连怀孕的姐妹都吃得津津有味！整整三十包啊！什么手撕蟹柳、辣子鸡、嫩豆干、手撕素肉、鹌鹑蛋、小肉枣肠、猪肉腐、魔芋、魔芋丝等等，应有尽有。香辣土豆爽辣过瘾，各种素肉嚼劲十足，鹌鹑蛋营养美味，真的太多太多啦，...家人们，现在价格太划算了，赶紧下单。\"))\n print(chunk_text_english(\"Washington CNN When President Donald Trump declared in the House Chamber this week that executives at the nation’s top automakers were “so excited” about their prospects amid his new tariff regime, it did not entirely reflect the conversation he’d held with them earlier that day.\"))\n text = \"欢迎收听《TED Talks Daily》，在这里，我们每天为您带来新思想，激发您的好奇心。我是您的主持人，Elise Hugh。当我们去看医生时，医生会评估我们的身体健康状况，检查我们的生命体征，可能还会关注我们的胆固醇水平，确保我们整体处于健康状态。医生可能还会通过一系列问题来检查我们的心理健康。然而，人际交往专家Casley Killam指出，我们在理解健康时忽略了一个关键指标，那就是我们的社会健康。在2024年的演讲中，她解释了为什么人际关系如此重要，以及忽视它可能带来的代价。几年前，我认识的一位女士，我们暂且称她为Maya，在短时间内经历了许多重大变化。她结婚了，和丈夫因工作搬到了一个陌生的城市，在那里她谁也不认识。她开始了一份在家办公的新工作，同时还要应对父亲新确诊的痴呆症。为了应对这些变化带来的压力，Maya加倍关注自己的身心健康。她几乎每天都锻炼，吃健康的食物，每周去看一次心理医生。这些措施确实有帮助，她的身体变得更加强壮，心理也更具韧性，但效果有限。她仍然感到困扰，经常在半夜失眠，白天感到注意力不集中，缺乏动力。Maya做了医生通常建议我们做的所有事情来保持身心健康，但似乎还缺少些什么。如果我告诉你，Maya所缺少的东西，也是全球数十亿人所缺少的，甚至可能也是你所缺少的呢？如果我告诉你，缺乏它会削弱我们为保持健康所做的其他努力，甚至可能缩短你的寿命呢？我研究这个问题已经超过十年，我发现，我们传统上对健康的理解是不完整的。通过将健康主要视为身体和心理的健康，我们忽略了我认为是我们这个时代最大的挑战和机遇——社会健康。身体健康关乎我们的身体，心理健康关乎我们的思想，而社会健康则关乎我们的人际关系。如果你以前没有听说过这个词，那是因为它还没有进入主流词汇，但它同样重要。Maya在她的新家还没有归属感。她不再亲自见到她的家人、朋友或同事，她经常一连几周只和丈夫共度时光。她的故事告诉我们，如果我们只照顾身体和心理，而不关注人际关系，我们就无法完全健康，无法真正茁壮成长。与Maya类似，全球有数亿人连续几周不与任何朋友或家人交谈。全球范围内，有四分之一的人感到孤独。20%的成年人觉得他们没有任何人可以求助。想想看，你遇到的每五个人中，可能有一个人觉得自己孤立无援。这不仅令人心碎，也是一场公共卫生危机。\"\n for res in chunk_text_chinesev2(text):\n print(res)" }, { "path": "tts/utils/text_utils/text_encoder.py", "content": "# Copyright 2025 ByteDance and/or its affiliates.\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 json\nimport re\nimport six\nfrom six.moves import range # pylint: disable=redefined-builtin\n\nPAD = \"\"\nEOS = \"\"\nUNK = \"\"\nSEG = \"|\"\nPUNCS = '!,.?;:'\nRESERVED_TOKENS = [PAD, EOS, UNK]\nNUM_RESERVED_TOKENS = len(RESERVED_TOKENS)\nPAD_ID = RESERVED_TOKENS.index(PAD) # Normally 0\nEOS_ID = RESERVED_TOKENS.index(EOS) # Normally 1\nUNK_ID = RESERVED_TOKENS.index(UNK) # Normally 2\n\nif six.PY2:\n RESERVED_TOKENS_BYTES = RESERVED_TOKENS\nelse:\n RESERVED_TOKENS_BYTES = [bytes(PAD, \"ascii\"), bytes(EOS, \"ascii\")]\n\n# Regular expression for unescaping token strings.\n# '\\u' is converted to '_'\n# '\\\\' is converted to '\\'\n# '\\213;' is converted to unichr(213)\n_UNESCAPE_REGEX = re.compile(r\"\\\\u|\\\\\\\\|\\\\([0-9]+);\")\n_ESCAPE_CHARS = set(u\"\\\\_u;0123456789\")\n\n\ndef strip_ids(ids, ids_to_strip):\n \"\"\"Strip ids_to_strip from the end ids.\"\"\"\n ids = list(ids)\n while ids and ids[-1] in ids_to_strip:\n ids.pop()\n return ids\n\n\nclass TextEncoder(object):\n \"\"\"Base class for converting from ints to/from human readable strings.\"\"\"\n\n def __init__(self, num_reserved_ids=NUM_RESERVED_TOKENS):\n self._num_reserved_ids = num_reserved_ids\n\n @property\n def num_reserved_ids(self):\n return self._num_reserved_ids\n\n def encode(self, s):\n \"\"\"Transform a human-readable string into a sequence of int ids.\n\n The ids should be in the range [num_reserved_ids, vocab_size). Ids [0,\n num_reserved_ids) are reserved.\n\n EOS is not appended.\n\n Args:\n s: human-readable string to be converted.\n\n Returns:\n ids: list of integers\n \"\"\"\n return [int(w) + self._num_reserved_ids for w in s.split()]\n\n def decode(self, ids, strip_extraneous=False):\n \"\"\"Transform a sequence of int ids into a human-readable string.\n\n EOS is not expected in ids.\n\n Args:\n ids: list of integers to be converted.\n strip_extraneous: bool, whether to strip off extraneous tokens\n (EOS and PAD).\n\n Returns:\n s: human-readable string.\n \"\"\"\n if strip_extraneous:\n ids = strip_ids(ids, list(range(self._num_reserved_ids or 0)))\n return \" \".join(self.decode_list(ids))\n\n def decode_list(self, ids):\n \"\"\"Transform a sequence of int ids into a their string versions.\n\n This method supports transforming individual input/output ids to their\n string versions so that sequence to/from text conversions can be visualized\n in a human readable format.\n\n Args:\n ids: list of integers to be converted.\n\n Returns:\n strs: list of human-readable string.\n \"\"\"\n decoded_ids = []\n for id_ in ids:\n if 0 <= id_ < self._num_reserved_ids:\n decoded_ids.append(RESERVED_TOKENS[int(id_)])\n else:\n decoded_ids.append(id_ - self._num_reserved_ids)\n return [str(d) for d in decoded_ids]\n\n @property\n def vocab_size(self):\n raise NotImplementedError()\n\n\nclass TokenTextEncoder(TextEncoder):\n \"\"\"Encoder based on a user-supplied vocabulary (file or list).\"\"\"\n\n def __init__(self,\n vocab_filename,\n reverse=False,\n vocab_list=None,\n replace_oov=None,\n num_reserved_ids=NUM_RESERVED_TOKENS):\n \"\"\"Initialize from a file or list, one token per line.\n\n Handling of reserved tokens works as follows:\n - When initializing from a list, we add reserved tokens to the vocab.\n - When initializing from a file, we do not add reserved tokens to the vocab.\n - When saving vocab files, we save reserved tokens to the file.\n\n Args:\n vocab_filename: If not None, the full filename to read vocab from. If this\n is not None, then vocab_list should be None.\n reverse: Boolean indicating if tokens should be reversed during encoding\n and decoding.\n vocab_list: If not None, a list of elements of the vocabulary. If this is\n not None, then vocab_filename should be None.\n replace_oov: If not None, every out-of-vocabulary token seen when\n encoding will be replaced by this string (which must be in vocab).\n num_reserved_ids: Number of IDs to save for reserved tokens like .\n \"\"\"\n super(TokenTextEncoder, self).__init__(num_reserved_ids=num_reserved_ids)\n self._reverse = reverse\n self._replace_oov = replace_oov\n if vocab_filename:\n self._init_vocab_from_file(vocab_filename)\n else:\n assert vocab_list is not None\n self._init_vocab_from_list(vocab_list)\n self.pad_index = self.token_to_id[PAD]\n self.eos_index = self.token_to_id[EOS]\n self.unk_index = self.token_to_id[UNK]\n self.seg_index = self.token_to_id[SEG] if SEG in self.token_to_id else self.eos_index\n\n def encode(self, s):\n \"\"\"Converts a space-separated string of tokens to a list of ids.\"\"\"\n if isinstance(s, str):\n sentence = s\n tokens = sentence.strip().split()\n else:\n tokens = s\n if self._replace_oov is not None:\n tokens = [t if t in self.token_to_id else self._replace_oov\n for t in tokens]\n ret = [self.token_to_id[tok] for tok in tokens]\n return ret[::-1] if self._reverse else ret\n\n def decode(self, ids, strip_eos=False, strip_padding=False):\n if strip_padding and self.pad() in list(ids):\n pad_pos = list(ids).index(self.pad())\n ids = ids[:pad_pos]\n if strip_eos and self.eos() in list(ids):\n eos_pos = list(ids).index(self.eos())\n ids = ids[:eos_pos]\n return \" \".join(self.decode_list(ids))\n\n def decode_list(self, ids):\n seq = reversed(ids) if self._reverse else ids\n return [self._safe_id_to_token(i) for i in seq]\n\n @property\n def vocab_size(self):\n return len(self.id_to_token)\n\n def __len__(self):\n return self.vocab_size\n\n def _safe_id_to_token(self, idx):\n return self.id_to_token.get(idx, \"ID_%d\" % idx)\n\n def _init_vocab_from_file(self, filename):\n \"\"\"Load vocab from a file.\n\n Args:\n filename: The file to load vocabulary from.\n \"\"\"\n with open(filename) as f:\n tokens = [token.strip() for token in f.readlines()]\n\n def token_gen():\n for token in tokens:\n yield token\n\n self._init_vocab(token_gen(), add_reserved_tokens=False)\n\n def _init_vocab_from_list(self, vocab_list):\n \"\"\"Initialize tokens from a list of tokens.\n\n It is ok if reserved tokens appear in the vocab list. They will be\n removed. The set of tokens in vocab_list should be unique.\n\n Args:\n vocab_list: A list of tokens.\n \"\"\"\n\n def token_gen():\n for token in vocab_list:\n if token not in RESERVED_TOKENS:\n yield token\n\n self._init_vocab(token_gen())\n\n def _init_vocab(self, token_generator, add_reserved_tokens=True):\n \"\"\"Initialize vocabulary with tokens from token_generator.\"\"\"\n\n self.id_to_token = {}\n non_reserved_start_index = 0\n\n if add_reserved_tokens:\n self.id_to_token.update(enumerate(RESERVED_TOKENS))\n non_reserved_start_index = len(RESERVED_TOKENS)\n\n self.id_to_token.update(\n enumerate(token_generator, start=non_reserved_start_index))\n\n # _token_to_id is the reverse of _id_to_token\n self.token_to_id = dict((v, k) for k, v in six.iteritems(self.id_to_token))\n\n def pad(self):\n return self.pad_index\n\n def eos(self):\n return self.eos_index\n\n def unk(self):\n return self.unk_index\n\n def seg(self):\n return self.seg_index\n\n def store_to_file(self, filename):\n \"\"\"Write vocab file to disk.\n\n Vocab files have one token per line. The file ends in a newline. Reserved\n tokens are written to the vocab file as well.\n\n Args:\n filename: Full path of the file to store the vocab to.\n \"\"\"\n with open(filename, \"w\") as f:\n for i in range(len(self.id_to_token)):\n f.write(self.id_to_token[i] + \"\\n\")\n\n def sil_phonemes(self):\n return [p for p in self.id_to_token.values() if is_sil_phoneme(p)]\n\n\ndef build_token_encoder(token_list_file):\n token_list = json.load(open(token_list_file))\n return TokenTextEncoder(None, vocab_list=token_list, replace_oov='')\n\n\ndef is_sil_phoneme(p):\n return p == '' or not p[0].isalpha() or p == 'sil' or p == 'sp' or p == 'XX'\n" }, { "path": "workflow-examples/单人语音.json", "content": "{\"id\":\"f4285961-b2b3-477d-8027-364fea281241\",\"revision\":0,\"last_node_id\":27,\"last_link_id\":74,\"nodes\":[{\"id\":15,\"type\":\"PreviewAudio\",\"pos\":[1494.0950927734375,-45.68428039550781],\"size\":[270,88],\"flags\":{},\"order\":3,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":73},{\"localized_name\":\"audioUI\",\"name\":\"audioUI\",\"type\":\"AUDIO_UI\",\"widget\":{\"name\":\"audioUI\"},\"link\":null}],\"outputs\":[],\"properties\":{\"Node name for S&R\":\"PreviewAudio\"},\"widgets_values\":[]},{\"id\":11,\"type\":\"MegaTTS3SpeakersPreview\",\"pos\":[814.6781616210938,-40.99158477783203],\"size\":[315,78],\"flags\":{},\"order\":0,\"mode\":0,\"inputs\":[{\"localized_name\":\"speaker\",\"name\":\"speaker\",\"type\":\"COMBO\",\"widget\":{\"name\":\"speaker\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"links\":[71]},{\"localized_name\":\"npy_file\",\"name\":\"npy_file\",\"type\":\"STRING\",\"links\":[]}],\"properties\":{\"Node name for S&R\":\"MegaTTS3SpeakersPreview\"},\"widgets_values\":[\"中文女-已知风格\\\\御姐配音.wav\"]},{\"id\":13,\"type\":\"MultiLinePromptMG\",\"pos\":[822.5565795898438,91.6104965209961],\"size\":[292.10406494140625,190.9091796875],\"flags\":{},\"order\":1,\"mode\":0,\"inputs\":[{\"localized_name\":\"multi_line_prompt\",\"name\":\"multi_line_prompt\",\"type\":\"STRING\",\"widget\":{\"name\":\"multi_line_prompt\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"text\",\"name\":\"text\",\"type\":\"STRING\",\"links\":[72]}],\"properties\":{\"Node name for S&R\":\"MultiLinePromptMG\"},\"widgets_values\":[\"MegaTTS 真开源版本来了，效果666, 我爱你！I love you!“我爱你”的英语是“I love you”\\n\\n2.5平方电线,共465篇，约315万字, 2002年的第一场雪，下在了2003年\\n\\nHigh-quality voice cloning, supports Chinese and English, and can perform cross-lingual cloning\\n\\n亲爱的,你最好了,不要生气了嘛,人家真的不是故意的啦,原谅我吧,好不好嘛\"]},{\"id\":19,\"type\":\"SaveAudioMP3\",\"pos\":[1500.0428466796875,108.09651184082031],\"size\":[270,136],\"flags\":{},\"order\":4,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":74},{\"localized_name\":\"filename_prefix\",\"name\":\"filename_prefix\",\"type\":\"STRING\",\"widget\":{\"name\":\"filename_prefix\"},\"link\":null},{\"localized_name\":\"quality\",\"name\":\"quality\",\"type\":\"COMBO\",\"widget\":{\"name\":\"quality\"},\"link\":null},{\"localized_name\":\"audioUI\",\"name\":\"audioUI\",\"type\":\"AUDIO_UI\",\"widget\":{\"name\":\"audioUI\"},\"link\":null}],\"outputs\":[],\"properties\":{\"Node name for S&R\":\"SaveAudioMP3\"},\"widgets_values\":[\"单人\",\"V0\"]},{\"id\":27,\"type\":\"MegaTTS3Run\",\"pos\":[1150.0458984375,-12.683919906616211],\"size\":[315,210],\"flags\":{},\"order\":2,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":71},{\"localized_name\":\"text\",\"name\":\"text\",\"type\":\"STRING\",\"link\":72},{\"localized_name\":\"dialogue_audio_s2\",\"name\":\"dialogue_audio_s2\",\"shape\":7,\"type\":\"AUDIO\",\"link\":null},{\"localized_name\":\"audio_npy_file\",\"name\":\"audio_npy_file\",\"shape\":7,\"type\":\"STRING\",\"link\":null},{\"localized_name\":\"audio_s2_npy_file\",\"name\":\"audio_s2_npy_file\",\"shape\":7,\"type\":\"STRING\",\"link\":null},{\"localized_name\":\"time_step\",\"name\":\"time_step\",\"type\":\"INT\",\"widget\":{\"name\":\"time_step\"},\"link\":null},{\"localized_name\":\"p_w\",\"name\":\"p_w\",\"type\":\"FLOAT\",\"widget\":{\"name\":\"p_w\"},\"link\":null},{\"localized_name\":\"t_w\",\"name\":\"t_w\",\"type\":\"FLOAT\",\"widget\":{\"name\":\"t_w\"},\"link\":null},{\"localized_name\":\"unload_model\",\"name\":\"unload_model\",\"type\":\"BOOLEAN\",\"widget\":{\"name\":\"unload_model\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"links\":[73,74]}],\"title\":\"Mega TTS3 Run\",\"properties\":{\"Node name for S&R\":\"MegaTTS3Run\"},\"widgets_values\":[32,1.6,2.5,false]}],\"links\":[[71,11,0,27,0,\"AUDIO\"],[72,13,0,27,1,\"STRING\"],[73,27,0,15,0,\"AUDIO\"],[74,27,0,19,0,\"AUDIO\"]],\"groups\":[],\"config\":{},\"extra\":{\"ds\":{\"scale\":1.1000000000000003,\"offset\":[-508.3255535295901,162.01651221912473]},\"frontendVersion\":\"1.16.9\",\"VHS_latentpreview\":false,\"VHS_latentpreviewrate\":0,\"VHS_MetadataImage\":true,\"VHS_KeepIntermediate\":true},\"version\":0.4}" }, { "path": "workflow-examples/双人会话.json", "content": "{\"id\":\"eb4d7439-617f-4e18-beba-af841ebc9d1a\",\"revision\":0,\"last_node_id\":22,\"last_link_id\":54,\"nodes\":[{\"id\":11,\"type\":\"MegaTTS3SpeakersPreview\",\"pos\":[812.1293334960938,-43.54039764404297],\"size\":[315,78],\"flags\":{},\"order\":0,\"mode\":0,\"inputs\":[{\"localized_name\":\"speaker\",\"name\":\"speaker\",\"type\":\"COMBO\",\"widget\":{\"name\":\"speaker\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"links\":[48]},{\"localized_name\":\"npy_file\",\"name\":\"npy_file\",\"type\":\"STRING\",\"links\":[50]}],\"properties\":{\"Node name for S&R\":\"MegaTTS3SpeakersPreview\"},\"widgets_values\":[\"中文未分男女\\\\磁性-中音-中-00001.wav\"]},{\"id\":17,\"type\":\"MegaTTS3SpeakersPreview\",\"pos\":[810.3323974609375,284.3142395019531],\"size\":[315,78],\"flags\":{},\"order\":1,\"mode\":0,\"inputs\":[{\"localized_name\":\"speaker\",\"name\":\"speaker\",\"type\":\"COMBO\",\"widget\":{\"name\":\"speaker\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"links\":[53]},{\"localized_name\":\"npy_file\",\"name\":\"npy_file\",\"type\":\"STRING\",\"links\":[54]}],\"properties\":{\"Node name for S&R\":\"MegaTTS3SpeakersPreview\"},\"widgets_values\":[\"中文女-已知风格\\\\御姐配音.wav\"]},{\"id\":13,\"type\":\"MultiLinePromptMG\",\"pos\":[821.70703125,91.6104965209961],\"size\":[287.00628662109375,133.1352081298828],\"flags\":{},\"order\":2,\"mode\":0,\"inputs\":[{\"localized_name\":\"multi_line_prompt\",\"name\":\"multi_line_prompt\",\"type\":\"STRING\",\"widget\":{\"name\":\"multi_line_prompt\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"text\",\"name\":\"text\",\"type\":\"STRING\",\"links\":[49]}],\"properties\":{\"Node name for S&R\":\"MultiLinePromptMG\"},\"widgets_values\":[\"[S1] MegaTTS 真开源版本来了，效果666\\n[S2] 晕 xuan4 是一种 gan3 觉\\n[S1] 我爱你！I love you!“我爱你”的英语是“I love you”\\n[S2] 2.5平方电线,共465篇，约315万字\\n[S1] 2002年的第一场雪，下在了2003年\"]},{\"id\":15,\"type\":\"PreviewAudio\",\"pos\":[1496.6439208984375,-0.654519259929657],\"size\":[270,88],\"flags\":{},\"order\":4,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":51},{\"localized_name\":\"audioUI\",\"name\":\"audioUI\",\"type\":\"AUDIO_UI\",\"widget\":{\"name\":\"audioUI\"},\"link\":null}],\"outputs\":[],\"properties\":{\"Node name for S&R\":\"PreviewAudio\"},\"widgets_values\":[]},{\"id\":19,\"type\":\"SaveAudioMP3\",\"pos\":[1507.689453125,157.37432861328125],\"size\":[270,136],\"flags\":{},\"order\":5,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":52},{\"localized_name\":\"filename_prefix\",\"name\":\"filename_prefix\",\"type\":\"STRING\",\"widget\":{\"name\":\"filename_prefix\"},\"link\":null},{\"localized_name\":\"quality\",\"name\":\"quality\",\"type\":\"COMBO\",\"widget\":{\"name\":\"quality\"},\"link\":null},{\"localized_name\":\"audioUI\",\"name\":\"audioUI\",\"type\":\"AUDIO_UI\",\"widget\":{\"name\":\"audioUI\"},\"link\":null}],\"outputs\":[],\"properties\":{\"Node name for S&R\":\"SaveAudioMP3\"},\"widgets_values\":[\"双人对话\",\"V0\"]},{\"id\":22,\"type\":\"MegaTTS3Run\",\"pos\":[1157.6923828125,-2.4885244369506836],\"size\":[315,210],\"flags\":{},\"order\":3,\"mode\":0,\"inputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"link\":48},{\"localized_name\":\"text\",\"name\":\"text\",\"type\":\"STRING\",\"link\":49},{\"localized_name\":\"dialogue_audio_s2\",\"name\":\"dialogue_audio_s2\",\"shape\":7,\"type\":\"AUDIO\",\"link\":53},{\"localized_name\":\"audio_npy_file\",\"name\":\"audio_npy_file\",\"shape\":7,\"type\":\"STRING\",\"link\":50},{\"localized_name\":\"audio_s2_npy_file\",\"name\":\"audio_s2_npy_file\",\"shape\":7,\"type\":\"STRING\",\"link\":54},{\"localized_name\":\"time_step\",\"name\":\"time_step\",\"type\":\"INT\",\"widget\":{\"name\":\"time_step\"},\"link\":null},{\"localized_name\":\"p_w\",\"name\":\"p_w\",\"type\":\"FLOAT\",\"widget\":{\"name\":\"p_w\"},\"link\":null},{\"localized_name\":\"t_w\",\"name\":\"t_w\",\"type\":\"FLOAT\",\"widget\":{\"name\":\"t_w\"},\"link\":null},{\"localized_name\":\"unload_model\",\"name\":\"unload_model\",\"type\":\"BOOLEAN\",\"widget\":{\"name\":\"unload_model\"},\"link\":null}],\"outputs\":[{\"localized_name\":\"audio\",\"name\":\"audio\",\"type\":\"AUDIO\",\"links\":[51,52]}],\"title\":\"Mega TTS3 Run\",\"properties\":{\"Node name for S&R\":\"MegaTTS3Run\"},\"widgets_values\":[32,1.6,2.5,false]}],\"links\":[[48,11,0,22,0,\"AUDIO\"],[49,13,0,22,1,\"STRING\"],[50,11,1,22,3,\"STRING\"],[51,22,0,15,0,\"AUDIO\"],[52,22,0,19,0,\"AUDIO\"],[53,17,0,22,2,\"AUDIO\"],[54,17,1,22,4,\"STRING\"]],\"groups\":[],\"config\":{},\"extra\":{\"ds\":{\"scale\":1.1000000000000003,\"offset\":[-476.88972922206165,158.6180469758719]},\"frontendVersion\":\"1.16.9\",\"VHS_latentpreview\":false,\"VHS_latentpreviewrate\":0,\"VHS_MetadataImage\":true,\"VHS_KeepIntermediate\":true},\"version\":0.4}" } ]