From b51a0820b3e8104c69bca246d00ccb281554e512 Mon Sep 17 00:00:00 2001 From: Yatharth Sharma Date: Fri, 23 Jan 2026 15:12:06 -0500 Subject: [PATCH] Initial commit --- .gitignore | 195 ++ LICENSE | 201 ++ README.md | 23 + pyproject.toml | 21 + requirements.txt | 23 + zipvoice/__init__.py | 7 + zipvoice/bin/compute_fbank.py | 272 +++ zipvoice/bin/generate_averaged_model.py | 229 +++ zipvoice/bin/infer_zipvoice.py | 899 +++++++++ zipvoice/bin/infer_zipvoice_dialog.py | 1286 +++++++++++++ zipvoice/bin/infer_zipvoice_onnx.py | 924 +++++++++ zipvoice/bin/onnx_export.py | 429 +++++ zipvoice/bin/prepare_dataset.py | 274 +++ zipvoice/bin/prepare_tokens.py | 103 + zipvoice/bin/tensorrt_export.py | 382 ++++ zipvoice/bin/train_zipvoice.py | 1130 +++++++++++ zipvoice/bin/train_zipvoice_dialog.py | 980 ++++++++++ zipvoice/bin/train_zipvoice_dialog_stereo.py | 963 ++++++++++ zipvoice/bin/train_zipvoice_distill.py | 1158 ++++++++++++ zipvoice/luxvoice.py | 52 + zipvoice/modeling_utils.py | 157 ++ zipvoice/models/modules/scaling.py | 1590 ++++++++++++++++ zipvoice/models/modules/solver.py | 277 +++ zipvoice/models/modules/zipformer.py | 1680 +++++++++++++++++ .../models/modules/zipformer_two_stream.py | 264 +++ zipvoice/models/zipvoice.py | 534 ++++++ zipvoice/models/zipvoice_dialog.py | 358 ++++ zipvoice/models/zipvoice_distill.py | 94 + zipvoice/onnx_modeling.py | 206 ++ zipvoice/tokenizer/normalizer.py | 170 ++ zipvoice/tokenizer/tokenizer.py | 648 +++++++ zipvoice/utils/checkpoint.py | 570 ++++++ zipvoice/utils/common.py | 670 +++++++ zipvoice/utils/diagnostics.py | 723 +++++++ zipvoice/utils/feature.py | 120 ++ zipvoice/utils/hooks.py | 111 ++ zipvoice/utils/infer.py | 414 ++++ zipvoice/utils/lr_scheduler.py | 245 +++ zipvoice/utils/optim.py | 868 +++++++++ zipvoice/utils/scaling_converter.py | 105 ++ zipvoice/utils/tensorrt.py | 143 ++ 41 files changed, 19498 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 requirements.txt create mode 100644 zipvoice/__init__.py create mode 100755 zipvoice/bin/compute_fbank.py create mode 100644 zipvoice/bin/generate_averaged_model.py create mode 100644 zipvoice/bin/infer_zipvoice.py create mode 100644 zipvoice/bin/infer_zipvoice_dialog.py create mode 100644 zipvoice/bin/infer_zipvoice_onnx.py create mode 100644 zipvoice/bin/onnx_export.py create mode 100644 zipvoice/bin/prepare_dataset.py create mode 100644 zipvoice/bin/prepare_tokens.py create mode 100644 zipvoice/bin/tensorrt_export.py create mode 100644 zipvoice/bin/train_zipvoice.py create mode 100644 zipvoice/bin/train_zipvoice_dialog.py create mode 100644 zipvoice/bin/train_zipvoice_dialog_stereo.py create mode 100644 zipvoice/bin/train_zipvoice_distill.py create mode 100644 zipvoice/luxvoice.py create mode 100644 zipvoice/modeling_utils.py create mode 100644 zipvoice/models/modules/scaling.py create mode 100644 zipvoice/models/modules/solver.py create mode 100644 zipvoice/models/modules/zipformer.py create mode 100644 zipvoice/models/modules/zipformer_two_stream.py create mode 100644 zipvoice/models/zipvoice.py create mode 100644 zipvoice/models/zipvoice_dialog.py create mode 100644 zipvoice/models/zipvoice_distill.py create mode 100644 zipvoice/onnx_modeling.py create mode 100644 zipvoice/tokenizer/normalizer.py create mode 100644 zipvoice/tokenizer/tokenizer.py create mode 100644 zipvoice/utils/checkpoint.py create mode 100644 zipvoice/utils/common.py create mode 100644 zipvoice/utils/diagnostics.py create mode 100644 zipvoice/utils/feature.py create mode 100644 zipvoice/utils/hooks.py create mode 100644 zipvoice/utils/infer.py create mode 100644 zipvoice/utils/lr_scheduler.py create mode 100644 zipvoice/utils/optim.py create mode 100644 zipvoice/utils/scaling_converter.py create mode 100644 zipvoice/utils/tensorrt.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d9d1331 --- /dev/null +++ b/.gitignore @@ -0,0 +1,195 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST +results/ + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the enitre vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Cursor +# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to +# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data +# refer to https://docs.cursor.com/context/ignore-files +.cursorignore +.cursorindexingignore \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..6a86d63 --- /dev/null +++ b/README.md @@ -0,0 +1,23 @@ +Simple installation: +``` +git clone https://github.com/ysharma3501/LuxTTS.git +cd LuxTTS +pip install -r requirements.txt +``` + +Usage: +```python +from zipvoice.luxtts import LuxTTS +lux_tts = LuxTTS('YatharthS/LuxTTS', device='cuda') ## change device to cpu for cpu usage +``` + +Infer: +```python +text = "Hey, what's up loser? I think you should shut up? You have NO dignity either way, ugh!" +prompt_audio = '/kaggle/input/voices/ElevenLabs_2025-11-02T22_31_30_Jessica_pre_sp100_s35_sb80_v3.mp3' + +encoded_prompt = lux_tts.encode_prompt(prompt_audio, rms=0.001) +final_wav = lux_tts.generate_speech(text, encoded_prompt, num_steps=4, t_shift=0.9) + +display(Audio(final_wav, rate=48000)) +``` diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..992bb74 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,21 @@ +[project] +name = "Zipvoice" +version = "0.0.11" +description = "Linacodec is a highly compressive and rapid audio tokenizer for audio models." +readme = "README.md" +requires-python = ">=3.10" +dependencies = [ + "huggingface_hub", +] + +[project.optional-dependencies] +train = [ + "einops", +] + +[build-system] +requires = ["uv_build>=0.9.6,<0.10.0"] +build-backend = "uv_build" + +[tool.ruff] +line-length = 120 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..f008ac6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,23 @@ +--find-links https://k2-fsa.github.io/icefall/piper_phonemize.html + +torch +torchaudio +numpy +lhotse +huggingface_hub +safetensors +tensorboard +vocos +pydub +onnxruntime +git+https://github.com/ysharma3501/LinaCodec.git + +# Normalization +cn2an +inflect + +# Tokenization +jieba +piper_phonemize +pypinyin +setuptools<81 diff --git a/zipvoice/__init__.py b/zipvoice/__init__.py new file mode 100644 index 0000000..4e221d7 --- /dev/null +++ b/zipvoice/__init__.py @@ -0,0 +1,7 @@ +import warnings + +warnings.filterwarnings( + "ignore", + category=UserWarning, + message="pkg_resources is deprecated as an API.*", +) diff --git a/zipvoice/bin/compute_fbank.py b/zipvoice/bin/compute_fbank.py new file mode 100755 index 0000000..b707bbb --- /dev/null +++ b/zipvoice/bin/compute_fbank.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +# Copyright 2024-2025 Xiaomi Corp. (authors: Wei Kang +# Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Usage: + python3 -m zipvoice.bin.compute_fbank \ + --source-dir data/manifests \ + --dest-dir data/fbank \ + --dataset libritts \ + --subset dev-other \ + --sampling-rate 24000 \ + --num-jobs 20 + +The input would be data/manifests/libritts-cuts_dev-other.jsonl.gz or + (libritts_supervisions_dev-other.jsonl.gz and librittsrecordings_dev-other.jsonl.gz) + +The output would be data/fbank/libritts-cuts_dev-other.jsonl.gz +""" + + +import argparse +import logging +from concurrent.futures import ProcessPoolExecutor as Pool +from pathlib import Path + +import lhotse +import torch +from lhotse import CutSet, LilcomChunkyWriter, load_manifest_lazy + +from zipvoice.utils.common import str2bool +from zipvoice.utils.feature import VocosFbank + +# Torch's multithreaded behavior needs to be disabled or +# it wastes a lot of CPU and slow things down. +# Do this outside of main() in case it needs to take effect +# even when we are not invoking the main (e.g. when spawning subprocesses). +torch.set_num_threads(1) +torch.set_num_interop_threads(1) + +lhotse.set_audio_duration_mismatch_tolerance(0.1) + + +def get_args(): + parser = argparse.ArgumentParser() + + parser.add_argument( + "--sampling-rate", + type=int, + default=24000, + help="The target sampling rate, the audio will be resampled to it.", + ) + + parser.add_argument( + "--type", + type=str, + default="vocos", + help="fbank type", + ) + + parser.add_argument( + "--dataset", + type=str, + help="Dataset name.", + ) + + parser.add_argument( + "--subset", + type=str, + help="The subset of the dataset.", + ) + + parser.add_argument( + "--source-dir", + type=str, + default="data/manifests", + help="The source directory of manifest files.", + ) + + parser.add_argument( + "--dest-dir", + type=str, + default="data/fbank", + help="The destination directory of manifest files.", + ) + + parser.add_argument( + "--split-cuts", + type=str2bool, + default=False, + help="Whether to use splited cuts.", + ) + + parser.add_argument( + "--split-begin", + type=int, + help="Start idx of splited cuts.", + ) + + parser.add_argument( + "--split-end", + type=int, + help="End idx of splited cuts.", + ) + + parser.add_argument( + "--batch-duration", + type=int, + default=1000, + help="The batch duration when computing the features.", + ) + + parser.add_argument( + "--num-jobs", + type=int, + default=20, + help="The number of extractor workers.", + ) + + return parser.parse_args() + + +def compute_fbank_split_single(params, idx): + logging.info( + f"Computing features for {idx}-th split of " + f"{params.dataset} dataset {params.subset} subset" + ) + lhotse.set_audio_duration_mismatch_tolerance(0.1) # for emilia + src_dir = Path(params.source_dir) + output_dir = Path(params.dest_dir) + + if not src_dir.exists(): + logging.error(f"{src_dir} not exists") + return + + if not output_dir.exists(): + output_dir.mkdir(parents=True, exist_ok=True) + + num_digits = 8 + if params.type == "vocos": + extractor = VocosFbank() + else: + raise NotImplementedError(f"{params.type} is not supported") + + prefix = params.dataset + subset = params.subset + suffix = "jsonl.gz" + + idx = f"{idx}".zfill(num_digits) + cuts_filename = f"{prefix}_cuts_{subset}.{idx}.{suffix}" + + if (src_dir / cuts_filename).is_file(): + logging.info(f"Loading manifests {src_dir / cuts_filename}") + cut_set = load_manifest_lazy(src_dir / cuts_filename) + else: + logging.warning(f"Raw {cuts_filename} not exists, skipping") + return + + cut_set = cut_set.resample(params.sampling_rate) + + if (output_dir / cuts_filename).is_file(): + logging.info(f"{cuts_filename} already exists - skipping.") + return + + logging.info(f"Processing {subset}.{idx} of {prefix}") + + cut_set = cut_set.compute_and_store_features_batch( + extractor=extractor, + storage_path=f"{output_dir}/{prefix}_feats_{subset}_{idx}", + num_workers=4, + batch_duration=params.batch_duration, + storage_type=LilcomChunkyWriter, + overwrite=True, + ) + logging.info(f"Saving file to {output_dir / cuts_filename}") + cut_set.to_file(output_dir / cuts_filename) + + +def compute_fbank_split(params): + if params.split_end < params.split_begin: + logging.warning( + f"Split begin should be smaller than split end, given " + f"{params.split_begin} -> {params.split_end}." + ) + + with Pool(max_workers=params.num_jobs) as pool: + futures = [ + pool.submit(compute_fbank_split_single, params, i) + for i in range(params.split_begin, params.split_end) + ] + for f in futures: + f.result() + f.done() + + +def compute_fbank(params): + logging.info( + f"Computing features for {params.dataset} dataset {params.subset} subset" + ) + src_dir = Path(params.source_dir) + output_dir = Path(params.dest_dir) + num_jobs = params.num_jobs + if not output_dir.exists(): + output_dir.mkdir(parents=True, exist_ok=True) + + prefix = params.dataset + subset = params.subset + suffix = "jsonl.gz" + + cut_set_name = f"{prefix}_cuts_{subset}.{suffix}" + + if (src_dir / cut_set_name).is_file(): + logging.info(f"Loading manifests {src_dir / cut_set_name}") + cut_set = load_manifest_lazy(src_dir / cut_set_name) + else: + recordings = load_manifest_lazy( + src_dir / f"{prefix}_recordings_{subset}.{suffix}" + ) + supervisions = load_manifest_lazy( + src_dir / f"{prefix}_supervisions_{subset}.{suffix}" + ) + cut_set = CutSet.from_manifests( + recordings=recordings, + supervisions=supervisions, + ) + + cut_set = cut_set.resample(params.sampling_rate) + if params.type == "vocos": + extractor = VocosFbank() + else: + raise NotImplementedError(f"{params.type} is not supported") + + cuts_filename = f"{prefix}_cuts_{subset}.{suffix}" + if (output_dir / cuts_filename).is_file(): + logging.info(f"{prefix} {subset} already exists - skipping.") + return + logging.info(f"Processing {subset} of {prefix}") + + cut_set = cut_set.compute_and_store_features( + extractor=extractor, + storage_path=f"{output_dir}/{prefix}_feats_{subset}", + num_jobs=num_jobs, + storage_type=LilcomChunkyWriter, + ) + logging.info(f"Saving file to {output_dir / cuts_filename}") + cut_set.to_file(output_dir / cuts_filename) + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + args = get_args() + logging.info(vars(args)) + if args.split_cuts: + compute_fbank_split(params=args) + else: + compute_fbank(params=args) + logging.info("Done!") diff --git a/zipvoice/bin/generate_averaged_model.py b/zipvoice/bin/generate_averaged_model.py new file mode 100644 index 0000000..7ff432b --- /dev/null +++ b/zipvoice/bin/generate_averaged_model.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 +# +# Copyright 2021-2022 Xiaomi Corporation +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Usage: +This script loads checkpoints and averages them. + +python3 -m zipvoice.bin.generate_averaged_model \ + --epoch 11 \ + --avg 4 \ + --model-name zipvoice \ + --exp-dir exp/zipvoice + +It will generate a file `epoch-11-avg-14.pt` in the given `exp_dir`. +You can later load it by `torch.load("epoch-11-avg-4.pt")`. +""" + +import argparse +import json +import logging +from pathlib import Path + +import torch + +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.models.zipvoice_dialog import ZipVoiceDialog, ZipVoiceDialogStereo +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import SimpleTokenizer +from zipvoice.utils.checkpoint import ( + average_checkpoints_with_averaged_model, + find_checkpoints, +) +from zipvoice.utils.common import AttributeDict + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--epoch", + type=int, + default=11, + help="""It specifies the checkpoint to use for decoding. + Note: Epoch counts from 1. + You can specify --avg to use more checkpoints for model averaging.""", + ) + + parser.add_argument( + "--iter", + type=int, + default=0, + help="""If positive, --epoch is ignored and it + will use the checkpoint exp_dir/checkpoint-iter.pt. + You can specify --avg to use more checkpoints for model averaging. + """, + ) + + parser.add_argument( + "--avg", + type=int, + default=4, + help="Number of checkpoints to average. Automatically select " + "consecutive checkpoints before the checkpoint specified by " + "'--epoch' or --iter", + ) + + parser.add_argument( + "--exp-dir", + type=str, + default="exp/zipvoice", + help="The experiment dir", + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice", + choices=[ + "zipvoice", + "zipvoice_distill", + "zipvoice_dialog", + "zipvoice_dialog_stereo", + ], + help="The model type to be averaged. ", + ) + + return parser + + +@torch.no_grad() +def main(): + parser = get_parser() + args = parser.parse_args() + params = AttributeDict() + params.update(vars(args)) + params.exp_dir = Path(params.exp_dir) + + with open(params.exp_dir / "model.json", "r") as f: + model_config = json.load(f) + + # Any tokenizer can be used here. + # Use SimpleTokenizer for simplicity. + tokenizer = SimpleTokenizer(token_file=params.exp_dir / "tokens.txt") + if params.model_name in ["zipvoice", "zipvoice_distill"]: + tokenizer_config = { + "vocab_size": tokenizer.vocab_size, + "pad_id": tokenizer.pad_id, + } + elif params.model_name in ["zipvoice_dialog", "zipvoice_dialog_stereo"]: + tokenizer_config = { + "vocab_size": tokenizer.vocab_size, + "pad_id": tokenizer.pad_id, + "spk_a_id": tokenizer.spk_a_id, + "spk_b_id": tokenizer.spk_b_id, + } + + params.suffix = f"epoch-{params.epoch}-avg-{params.avg}" + + logging.info("Script started") + + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + logging.info("About to create model") + if params.model_name == "zipvoice": + model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + elif params.model_name == "zipvoice_distill": + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + elif params.model_name == "zipvoice_dialog": + model = ZipVoiceDialog( + **model_config["model"], + **tokenizer_config, + ) + elif params.model_name == "zipvoice_dialog_stereo": + model = ZipVoiceDialogStereo( + **model_config["model"], + **tokenizer_config, + ) + else: + raise ValueError(f"Unknown model name: {params.model_name}") + + if params.iter > 0: + filenames = find_checkpoints(params.exp_dir, iteration=-params.iter)[ + : params.avg + 1 + ] + if len(filenames) == 0: + raise ValueError( + f"No checkpoints found for" f" --iter {params.iter}, --avg {params.avg}" + ) + elif len(filenames) < params.avg + 1: + raise ValueError( + f"Not enough checkpoints ({len(filenames)}) found for" + f" --iter {params.iter}, --avg {params.avg}" + ) + filename_start = filenames[-1] + filename_end = filenames[0] + logging.info( + "Calculating the averaged model over iteration checkpoints" + f" from {filename_start} (excluded) to {filename_end}" + ) + model.to(params.device) + model.load_state_dict( + average_checkpoints_with_averaged_model( + filename_start=filename_start, + filename_end=filename_end, + device=params.device, + ), + strict=True, + ) + else: + assert params.avg > 0, params.avg + start = params.epoch - params.avg + assert start >= 1, start + filename_start = f"{params.exp_dir}/epoch-{start}.pt" + filename_end = f"{params.exp_dir}/epoch-{params.epoch}.pt" + logging.info( + f"Calculating the averaged model over epoch range from " + f"{start} (excluded) to {params.epoch}" + ) + model.to(params.device) + model.load_state_dict( + average_checkpoints_with_averaged_model( + filename_start=filename_start, + filename_end=filename_end, + device=params.device, + ), + strict=True, + ) + if params.iter > 0: + filename = params.exp_dir / f"iter-{params.iter}-avg-{params.avg}.pt" + else: + filename = params.exp_dir / f"epoch-{params.epoch}-avg-{params.avg}.pt" + + logging.info(f"Saving the averaged checkpoint to {filename}") + torch.save({"model": model.state_dict()}, filename) + + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of model parameters: {num_param}") + + logging.info("Done!") + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + main() diff --git a/zipvoice/bin/infer_zipvoice.py b/zipvoice/bin/infer_zipvoice.py new file mode 100644 index 0000000..a741036 --- /dev/null +++ b/zipvoice/bin/infer_zipvoice.py @@ -0,0 +1,899 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script generates speech with our pre-trained ZipVoice or + ZipVoice-Distill models. If no local model is specified, + Required files will be automatically downloaded from HuggingFace. + +Usage: + +Note: If you having trouble connecting to HuggingFace, + try switching endpoint to mirror site: +export HF_ENDPOINT=https://hf-mirror.com + +(1) Inference of a single sentence: + +python3 -m zipvoice.bin.infer_zipvoice \ + --model-name zipvoice \ + --prompt-wav prompt.wav \ + --prompt-text "I am a prompt." \ + --text "I am a sentence." \ + --res-wav-path result.wav + +(2) Inference of a list of sentences: + +python3 -m zipvoice.bin.infer_zipvoice \ + --model-name zipvoice \ + --test-list test.tsv \ + --res-dir results + +`--model-name` can be `zipvoice` or `zipvoice_distill`, + which are the models before and after distillation, respectively. + +Each line of `test.tsv` is in the format of + `{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}`. + + +(3) Inference with TensorRT: + +python3 -m zipvoice.bin.infer_zipvoice \ + --model-name zipvoice_distill \ + --prompt-wav prompt.wav \ + --prompt-text "I am a prompt." \ + --text "I am a sentence." \ + --res-wav-path result.wav \ + --trt-engine-path models/zipvoice_distill_onnx_trt/fm_decoder.fp16.plan +""" + +import argparse +import datetime as dt +import json +import logging +import os +from pathlib import Path +from typing import Optional + +import numpy as np +import safetensors.torch +import torch +import torchaudio +from huggingface_hub import hf_hub_download +from lhotse.utils import fix_random_seed +from vocos import Vocos + +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import ( + EmiliaTokenizer, + EspeakTokenizer, + LibriTTSTokenizer, + SimpleTokenizer, +) +from zipvoice.utils.checkpoint import load_checkpoint +from zipvoice.utils.common import AttributeDict, str2bool +from zipvoice.utils.feature import VocosFbank +from zipvoice.utils.infer import ( + add_punctuation, + batchify_tokens, + chunk_tokens_punctuation, + cross_fade_concat, + load_prompt_wav, + remove_silence, + rms_norm, +) +from zipvoice.utils.tensorrt import load_trt + +HUGGINGFACE_REPO = "k2-fsa/ZipVoice" +MODEL_DIR = { + "zipvoice": "zipvoice", + "zipvoice_distill": "zipvoice_distill", +} + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice", + choices=["zipvoice", "zipvoice_distill"], + help="The model used for inference", + ) + + parser.add_argument( + "--model-dir", + type=str, + default=None, + help="The model directory that contains model checkpoint, configuration " + "file model.json, and tokens file tokens.txt. Will download pre-trained " + "checkpoint from huggingface if not specified.", + ) + + parser.add_argument( + "--checkpoint-name", + type=str, + default="model.pt", + help="The name of model checkpoint.", + ) + + parser.add_argument( + "--vocoder-path", + type=str, + default=None, + help="The vocoder checkpoint. " + "Will download pre-trained vocoder from huggingface if not specified.", + ) + + parser.add_argument( + "--tokenizer", + type=str, + default="emilia", + choices=["emilia", "libritts", "espeak", "simple"], + help="Tokenizer type.", + ) + + parser.add_argument( + "--lang", + type=str, + default="en-us", + help="Language identifier, used when tokenizer type is espeak. see" + "https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md", + ) + + parser.add_argument( + "--test-list", + type=str, + default=None, + help="The list of prompt speech, prompt_transcription, " + "and text to synthesizein the format of " + "'{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}'.", + ) + + parser.add_argument( + "--prompt-wav", + type=str, + default=None, + help="The prompt wav to mimic", + ) + + parser.add_argument( + "--prompt-text", + type=str, + default=None, + help="The transcription of the prompt wav", + ) + + parser.add_argument( + "--text", + type=str, + default=None, + help="The text to synthesize", + ) + + parser.add_argument( + "--res-dir", + type=str, + default="results", + help=""" + Path name of the generated wavs dir, + used when test-list is not None + """, + ) + + parser.add_argument( + "--res-wav-path", + type=str, + default="result.wav", + help=""" + Path name of the generated wav path, + used when test-list is None + """, + ) + + parser.add_argument( + "--guidance-scale", + type=float, + default=None, + help="The scale of classifier-free guidance during inference.", + ) + + parser.add_argument( + "--num-step", + type=int, + default=None, + help="The number of sampling steps.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--speed", + type=float, + default=1.0, + help="Control speech speed, 1.0 means normal, >1.0 means speed up", + ) + + parser.add_argument( + "--t-shift", + type=float, + default=0.5, + help="Shift t to smaller ones if t_shift < 1.0", + ) + + parser.add_argument( + "--target-rms", + type=float, + default=0.1, + help="Target speech normalization rms value, set to 0 to disable normalization", + ) + + parser.add_argument( + "--seed", + type=int, + default=666, + help="Random seed", + ) + + parser.add_argument( + "--num-thread", + type=int, + default=1, + help="Number of threads to use for PyTorch on CPU.", + ) + + parser.add_argument( + "--raw-evaluation", + type=str2bool, + default=False, + help="Whether to use the 'raw' evaluation mode where provided " + "prompts and text are fed to the model without pre-processing", + ) + + parser.add_argument( + "--max-duration", + type=float, + default=100, + help="Maximum duration (seconds) in a single batch, including " + "durations of the prompt and generated wavs. You can reduce it " + "if it causes CUDA OOM.", + ) + + parser.add_argument( + "--remove-long-sil", + type=str2bool, + default=False, + help="Whether to remove long silences in the middle of the generated " + "speech (edge silences will be removed by default).", + ) + + parser.add_argument( + "--trt-engine-path", + type=str, + default=None, + help="The path to the TensorRT engine file.", + ) + return parser + + +def get_vocoder(vocos_local_path: Optional[str] = None): + if vocos_local_path: + vocoder = Vocos.from_hparams(f"{vocos_local_path}/config.yaml") + state_dict = torch.load( + f"{vocos_local_path}/pytorch_model.bin", + weights_only=True, + map_location="cpu", + ) + vocoder.load_state_dict(state_dict) + else: + vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz") + return vocoder + + +def generate_sentence_raw_evaluation( + save_path: str, + prompt_text: str, + prompt_wav: str, + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function directly feed the prompt_text, prompt_wav and text to the model. + It is not efficient and can have poor results for some inappropriate inputs. + (e.g., prompt wav contains long silence, text to be generated is too long) + This function can be used to evaluate the "raw" performance of the model. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (str): Path to the prompt wav file. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (EmiliaTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and process prompt wav + prompt_wav = load_prompt_wav(prompt_wav, sampling_rate=sampling_rate) + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + prompt_features_lens = torch.tensor([prompt_features.size(1)], device=device) + + # Convert text to tokens + tokens = tokenizer.texts_to_token_ids([text]) + prompt_tokens = tokenizer.texts_to_token_ids([prompt_text]) + + # Start timing + start_t = dt.datetime.now() + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + prompt_features_lens=prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + + # Start vocoder processing + start_vocoder_t = dt.datetime.now() + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + + # Calculate processing times and real-time factors + t = (dt.datetime.now() - start_t).total_seconds() + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + torchaudio.save(save_path, wav.cpu(), sample_rate=sampling_rate) + + return metrics + + +def generate_sentence( + save_path: str, + prompt_text: str, + prompt_wav: str, + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + max_duration: float = 100, + remove_long_sil: bool = False, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function will do the following to improve the generation quality: + 1. chunk the text according to punctuations. + 2. process chunked texts in batches. + 3. remove long silences in the prompt audio. + 4. add punctuation to the end of prompt text and text if there is not. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (str): Path to the prompt wav file. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (EmiliaTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + max_duration (float, optional): The maximum duration to process in each + batch. Used to control memory consumption when generating long audios. + remove_long_sil (bool, optional): Whether to remove long silences in the + middle of the generated speech (edge silences will be removed by default). + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and process prompt wav + prompt_wav = load_prompt_wav(prompt_wav, sampling_rate=sampling_rate) + + # Remove edge and long silences in the prompt wav. + # Add 0.2s trailing silence to avoid leaking prompt to generated speech. + prompt_wav = remove_silence( + prompt_wav, sampling_rate, only_edge=False, trail_sil=200 + ) + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + prompt_duration = prompt_wav.shape[-1] / sampling_rate + + if prompt_duration > 20: + logging.warning( + f"Given prompt wav is too long ({prompt_duration}s). " + f"Please provide a shorter one (1-3 seconds is recommended)." + ) + elif prompt_duration > 10: + logging.warning( + f"Given prompt wav is long ({prompt_duration}s). " + f"It will lead to slower inference speed and possibly worse speech quality." + ) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + + # Add punctuation in the end if there is not + text = add_punctuation(text) + prompt_text = add_punctuation(prompt_text) + + # Tokenize text (str tokens), punctuations will be preserved. + tokens_str = tokenizer.texts_to_tokens([text])[0] + prompt_tokens_str = tokenizer.texts_to_tokens([prompt_text])[0] + + # chunk text so that each len(prompt wav + generated wav) is around 25 seconds. + token_duration = (prompt_wav.shape[-1] / sampling_rate) / ( + len(prompt_tokens_str) * speed + ) + max_tokens = int((25 - prompt_duration) / token_duration) + chunked_tokens_str = chunk_tokens_punctuation(tokens_str, max_tokens=max_tokens) + + # Tokenize text (int tokens) + chunked_tokens = tokenizer.tokens_to_token_ids(chunked_tokens_str) + prompt_tokens = tokenizer.tokens_to_token_ids([prompt_tokens_str]) + + # Batchify chunked texts for faster processing + tokens_batches, chunked_index = batchify_tokens( + chunked_tokens, max_duration, prompt_duration, token_duration + ) + + # Start predicting features + chunked_features = [] + start_t = dt.datetime.now() + + for batch_tokens in tokens_batches: + batch_prompt_tokens = prompt_tokens * len(batch_tokens) + + batch_prompt_features = prompt_features.repeat(len(batch_tokens), 1, 1) + batch_prompt_features_lens = torch.full( + (len(batch_tokens),), prompt_features.size(1), device=device + ) + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=batch_tokens, + prompt_tokens=batch_prompt_tokens, + prompt_features=batch_prompt_features, + prompt_features_lens=batch_prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + chunked_features.append((pred_features, pred_features_lens)) + + # Start vocoder processing + chunked_wavs = [] + start_vocoder_t = dt.datetime.now() + + for pred_features, pred_features_lens in chunked_features: + batch_wav = [] + for i in range(pred_features.size(0)): + + wav = ( + vocoder.decode(pred_features[i][None, :, : pred_features_lens[i]]) + .squeeze(1) + .clamp(-1, 1) + ) + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + batch_wav.append(wav) + chunked_wavs.extend(batch_wav) + + # Finish model generation + t = (dt.datetime.now() - start_t).total_seconds() + + # Merge chunked wavs + indexed_chunked_wavs = [ + (index, wav) for index, wav in zip(chunked_index, chunked_wavs) + ] + sequential_indexed_chunked_wavs = sorted(indexed_chunked_wavs, key=lambda x: x[0]) + sequential_chunked_wavs = [ + sequential_indexed_chunked_wavs[i][1] + for i in range(len(sequential_indexed_chunked_wavs)) + ] + final_wav = cross_fade_concat( + sequential_chunked_wavs, fade_duration=0.1, sample_rate=sampling_rate + ) + final_wav = remove_silence( + final_wav, sampling_rate, only_edge=(not remove_long_sil), trail_sil=0 + ) + + # Calculate processing time metrics + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = final_wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + torchaudio.save(save_path, final_wav.cpu(), sample_rate=sampling_rate) + return metrics + + +def generate_list( + res_dir: str, + test_list: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + raw_evaluation: bool = False, + max_duration: float = 100, + remove_long_sil: bool = False, +): + total_t = [] + total_t_no_vocoder = [] + total_t_vocoder = [] + total_wav_seconds = [] + + with open(test_list, "r") as fr: + lines = fr.readlines() + + for i, line in enumerate(lines): + wav_name, prompt_text, prompt_wav, text = line.strip().split("\t") + save_path = f"{res_dir}/{wav_name}.wav" + + common_params = { + "save_path": save_path, + "prompt_text": prompt_text, + "prompt_wav": prompt_wav, + "text": text, + "model": model, + "vocoder": vocoder, + "tokenizer": tokenizer, + "feature_extractor": feature_extractor, + "device": device, + "num_step": num_step, + "guidance_scale": guidance_scale, + "speed": speed, + "t_shift": t_shift, + "target_rms": target_rms, + "feat_scale": feat_scale, + "sampling_rate": sampling_rate, + } + + if raw_evaluation: + metrics = generate_sentence_raw_evaluation(**common_params) + else: + metrics = generate_sentence( + **common_params, + max_duration=max_duration, + remove_long_sil=remove_long_sil, + ) + logging.info(f"[Sentence: {i}] Saved to: {save_path}") + logging.info(f"[Sentence: {i}] RTF: {metrics['rtf']:.4f}") + total_t.append(metrics["t"]) + total_t_no_vocoder.append(metrics["t_no_vocoder"]) + total_t_vocoder.append(metrics["t_vocoder"]) + total_wav_seconds.append(metrics["wav_seconds"]) + + logging.info(f"Average RTF: {np.sum(total_t) / np.sum(total_wav_seconds):.4f}") + logging.info( + f"Average RTF w/o vocoder: " + f"{np.sum(total_t_no_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + logging.info( + f"Average RTF vocoder: " + f"{np.sum(total_t_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + + +@torch.inference_mode() +def main(): + parser = get_parser() + args = parser.parse_args() + + torch.set_num_threads(args.num_thread) + torch.set_num_interop_threads(args.num_thread) + + params = AttributeDict() + params.update(vars(args)) + fix_random_seed(params.seed) + + model_defaults = { + "zipvoice": { + "num_step": 16, + "guidance_scale": 1.0, + }, + "zipvoice_distill": { + "num_step": 8, + "guidance_scale": 3.0, + }, + } + + model_specific_defaults = model_defaults.get(params.model_name, {}) + + for param, value in model_specific_defaults.items(): + if getattr(params, param) is None: + setattr(params, param, value) + logging.info(f"Setting {param} to default value: {value}") + + assert (params.test_list is not None) ^ ( + (params.prompt_wav and params.prompt_text and params.text) is not None + ), ( + "For inference, please provide prompts and text with either '--test-list'" + " or '--prompt-wav, --prompt-text and --text'." + ) + + if params.model_dir is not None: + params.model_dir = Path(params.model_dir) + if not params.model_dir.is_dir(): + raise FileNotFoundError(f"{params.model_dir} does not exist") + for filename in [params.checkpoint_name, "model.json", "tokens.txt"]: + if not (params.model_dir / filename).is_file(): + raise FileNotFoundError(f"{params.model_dir / filename} does not exist") + model_ckpt = params.model_dir / params.checkpoint_name + model_config = params.model_dir / "model.json" + token_file = params.model_dir / "tokens.txt" + logging.info( + f"Using {params.model_name} in local model dir {params.model_dir}, " + f"checkpoint {params.checkpoint_name}" + ) + else: + logging.info(f"Using pretrained {params.model_name} model from the Huggingface") + model_ckpt = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/model.pt" + ) + model_config = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/model.json" + ) + + token_file = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/tokens.txt" + ) + + if params.tokenizer == "emilia": + tokenizer = EmiliaTokenizer(token_file=token_file) + elif params.tokenizer == "libritts": + tokenizer = LibriTTSTokenizer(token_file=token_file) + elif params.tokenizer == "espeak": + tokenizer = EspeakTokenizer(token_file=token_file, lang=params.lang) + else: + assert params.tokenizer == "simple" + tokenizer = SimpleTokenizer(token_file=token_file) + + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + + with open(model_config, "r") as f: + model_config = json.load(f) + + if params.model_name == "zipvoice": + model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + else: + assert params.model_name == "zipvoice_distill" + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + + if str(model_ckpt).endswith(".safetensors"): + safetensors.torch.load_model(model, model_ckpt) + elif str(model_ckpt).endswith(".pt"): + load_checkpoint(filename=model_ckpt, model=model, strict=True) + else: + raise NotImplementedError(f"Unsupported model checkpoint format: {model_ckpt}") + + if torch.cuda.is_available(): + params.device = torch.device("cuda", 0) + elif torch.backends.mps.is_available(): + params.device = torch.device("mps") + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + model = model.to(params.device) + model.eval() + + if params.trt_engine_path: + load_trt(model, params.trt_engine_path) + + vocoder = get_vocoder(params.vocoder_path) + vocoder = vocoder.to(params.device) + vocoder.eval() + + if model_config["feature"]["type"] == "vocos": + feature_extractor = VocosFbank() + else: + raise NotImplementedError( + f"Unsupported feature type: {model_config['feature']['type']}" + ) + params.sampling_rate = model_config["feature"]["sampling_rate"] + + logging.info("Start generating...") + if params.test_list: + res_dir = params.res_dir + os.makedirs(res_dir, exist_ok=True) + generate_list( + res_dir=params.res_dir, + test_list=params.test_list, + model=model, + vocoder=vocoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + device=params.device, + num_step=params.num_step, + guidance_scale=params.guidance_scale, + speed=params.speed, + t_shift=params.t_shift, + target_rms=params.target_rms, + feat_scale=params.feat_scale, + sampling_rate=params.sampling_rate, + raw_evaluation=params.raw_evaluation, + max_duration=params.max_duration, + remove_long_sil=params.remove_long_sil, + ) + else: + assert ( + not params.raw_evaluation + ), "Raw evaluation is only valid with --test-list" + generate_sentence( + save_path=params.res_wav_path, + prompt_text=params.prompt_text, + prompt_wav=params.prompt_wav, + text=params.text, + model=model, + vocoder=vocoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + device=params.device, + num_step=params.num_step, + guidance_scale=params.guidance_scale, + speed=params.speed, + t_shift=params.t_shift, + target_rms=params.target_rms, + feat_scale=params.feat_scale, + sampling_rate=params.sampling_rate, + max_duration=params.max_duration, + remove_long_sil=params.remove_long_sil, + ) + logging.info(f"Saved to: {params.res_wav_path}") + logging.info("Done") + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + main() diff --git a/zipvoice/bin/infer_zipvoice_dialog.py b/zipvoice/bin/infer_zipvoice_dialog.py new file mode 100644 index 0000000..9503408 --- /dev/null +++ b/zipvoice/bin/infer_zipvoice_dialog.py @@ -0,0 +1,1286 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script generates speech with our pre-trained ZipVoice-Dialog or + ZipVoice-Dialog-Stereo models. If no local model is specified, + Required files will be automatically downloaded from HuggingFace. + +Usage: + +Note: If you having trouble connecting to HuggingFace, + try switching endpoint to mirror site: +export HF_ENDPOINT=https://hf-mirror.com + +python3 -m zipvoice.bin.infer_zipvoice_dialog \ + --model-name zipvoice_dialog \ + --test-list test.tsv \ + --res-dir results + +`--model-name` can be `zipvoice_dialog` or `zipvoice_dialog_stereo`, + which generate mono and stereo dialogues, respectively. + +Each line of `test.tsv` is in the format of merged conversation: + '{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}' + or splited conversation: + '{wav_name}\t{spk1_prompt_transcription}\t{spk2_prompt_transcription} + \t{spk1_prompt_wav}\t{spk2_prompt_wav}\t{text}' +""" + +import argparse +import datetime as dt +import json +import logging +import os +from pathlib import Path +from typing import List, Optional, Union + +import numpy as np +import safetensors.torch +import torch +import torchaudio +from huggingface_hub import hf_hub_download +from lhotse.utils import fix_random_seed +from vocos import Vocos + +from zipvoice.models.zipvoice_dialog import ZipVoiceDialog, ZipVoiceDialogStereo +from zipvoice.tokenizer.tokenizer import DialogTokenizer +from zipvoice.utils.checkpoint import load_checkpoint +from zipvoice.utils.common import AttributeDict, str2bool +from zipvoice.utils.feature import VocosFbank +from zipvoice.utils.infer import ( + add_punctuation, + batchify_tokens, + chunk_tokens_dialog, + cross_fade_concat, + load_prompt_wav, + remove_silence, + rms_norm, +) + +HUGGINGFACE_REPO = "k2-fsa/ZipVoice" +MODEL_DIR = { + "zipvoice_dialog": "zipvoice_dialog", + "zipvoice_dialog_stereo": "zipvoice_dialog_stereo", +} + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice_dialog", + choices=["zipvoice_dialog", "zipvoice_dialog_stereo"], + help="The model used for inference", + ) + + parser.add_argument( + "--model-dir", + type=str, + default=None, + help="The model directory that contains model checkpoint, configuration " + "file model.json, and tokens file tokens.txt. Will download pre-trained " + "checkpoint from huggingface if not specified.", + ) + + parser.add_argument( + "--checkpoint-name", + type=str, + default="model.pt", + help="The name of model checkpoint.", + ) + + parser.add_argument( + "--vocoder-path", + type=str, + default=None, + help="The vocoder checkpoint. " + "Will download pre-trained vocoder from huggingface if not specified.", + ) + + parser.add_argument( + "--test-list", + type=str, + default=None, + help="The list of prompt speech, prompt_transcription, " + "and text to synthesizein the format of merged conversation: " + "'{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}' " + "or splited conversation: " + "'{wav_name}\t{spk1_prompt_transcription}\t{spk2_prompt_transcription}" + "\t{spk1_prompt_wav}\t{spk2_prompt_wav}\t{text}'.", + ) + + parser.add_argument( + "--res-dir", + type=str, + default="results", + help=""" + Path name of the generated wavs dir, + used when test-list is not None + """, + ) + + parser.add_argument( + "--guidance-scale", + type=float, + default=1.5, + help="The scale of classifier-free guidance during inference.", + ) + + parser.add_argument( + "--num-step", + type=int, + default=16, + help="The number of sampling steps.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--speed", + type=float, + default=1.0, + help="Control speech speed, 1.0 means normal, >1.0 means speed up", + ) + + parser.add_argument( + "--t-shift", + type=float, + default=0.5, + help="Shift t to smaller ones if t_shift < 1.0", + ) + + parser.add_argument( + "--target-rms", + type=float, + default=0.1, + help="Target speech normalization rms value, set to 0 to disable normalization", + ) + + parser.add_argument( + "--seed", + type=int, + default=666, + help="Random seed", + ) + + parser.add_argument( + "--silence-wav", + type=str, + default="assets/silence.wav", + help="Path of the silence wav file, used in two-channel generation " + "with single-channel prompts", + ) + + parser.add_argument( + "--num-thread", + type=int, + default=1, + help="Number of threads to use for PyTorch on CPU.", + ) + + parser.add_argument( + "--raw-evaluation", + type=str2bool, + default=False, + help="Whether to use the 'raw' evaluation mode where provided " + "prompts and text are fed to the model without pre-processing", + ) + + parser.add_argument( + "--max-duration", + type=float, + default=100, + help="Maximum duration (seconds) in a single batch, including " + "durations of the prompt and generated wavs. You can reduce it " + "if it causes CUDA OOM.", + ) + + parser.add_argument( + "--remove-long-sil", + type=str2bool, + default=False, + help="Whether to remove long silences in the middle of the generated " + "speech (edge silences will be removed by default).", + ) + + return parser + + +def get_vocoder(vocos_local_path: Optional[str] = None): + if vocos_local_path: + vocoder = Vocos.from_hparams(f"{vocos_local_path}/config.yaml") + state_dict = torch.load( + f"{vocos_local_path}/pytorch_model.bin", + weights_only=True, + map_location="cpu", + ) + vocoder.load_state_dict(state_dict) + else: + vocoder = Vocos.from_pretrained("charactr/vocos-mel-24khz") + return vocoder + + +def generate_sentence_raw_evaluation( + save_path: str, + prompt_text: str, + prompt_wav: Union[str, List[str]], + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: DialogTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function directly feed the prompt_text, prompt_wav and text to the model. + It is not efficient and can have poor results for some inappropriate inputs. + (e.g., prompt wav contains long silence, text to be generated is too long) + This function can be used to evaluate the "raw" performance of the model. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (Union[str, List[str]]): Path to the prompt wav file, can be + one or two wav files, which corresponding to a merged conversational + speech or two seperate speaker's speech. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (DialogTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and preprocess prompt wav + if isinstance(prompt_wav, str): + prompt_wav = [ + prompt_wav, + ] + else: + assert len(prompt_wav) == 2 and isinstance(prompt_wav[0], str) + + loaded_prompt_wavs = prompt_wav + for i in range(len(prompt_wav)): + loaded_prompt_wavs[i] = load_prompt_wav( + loaded_prompt_wavs[i], sampling_rate=sampling_rate + ) + if loaded_prompt_wavs[i].size(0) != 1: + loaded_prompt_wavs[i] = loaded_prompt_wavs[i].mean(0, keepdim=True) + + if len(loaded_prompt_wavs) == 1: + prompt_wav = loaded_prompt_wavs[0] + else: + prompt_wav = torch.cat(loaded_prompt_wavs, dim=1) + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + prompt_features_lens = torch.tensor([prompt_features.size(1)], device=device) + + # Convert text to tokens + tokens = tokenizer.texts_to_token_ids([text]) + prompt_tokens = tokenizer.texts_to_token_ids([prompt_text]) + + # Start timing + start_t = dt.datetime.now() + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + prompt_features_lens=prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + + # Start vocoder processing + start_vocoder_t = dt.datetime.now() + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + + # Calculate processing times and real-time factors + t = (dt.datetime.now() - start_t).total_seconds() + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + torchaudio.save(save_path, wav.cpu(), sample_rate=sampling_rate) + + return metrics + + +def generate_sentence( + save_path: str, + prompt_text: str, + prompt_wav: Union[str, List[str]], + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: DialogTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + max_duration: float = 100, + remove_long_sil: bool = False, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function will do the following to improve the generation quality: + 1. chunk the text according to speaker-turn symbol [S1]. + 2. process chunked texts in batches. + 3. remove long silences in the prompt audio. + 4. add punctuation to the end of prompt text and text if there is not. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (Union[str, List[str]]): Path to the prompt wav file, can be + one or two wav files, which corresponding to a merged conversational + speech or two seperate speaker's speech. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (DialogTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + max_duration (float, optional): The maximum duration to process in each + batch. Used to control memory consumption when generating long audios. + remove_long_sil (bool, optional): Whether to remove long silences in the + middle of the generated speech (edge silences will be removed by default). + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and preprocess prompt wav + if isinstance(prompt_wav, str): + prompt_wav = [ + prompt_wav, + ] + else: + assert len(prompt_wav) == 2 and isinstance(prompt_wav[0], str) + + loaded_prompt_wavs = prompt_wav + for i in range(len(prompt_wav)): + loaded_prompt_wavs[i] = load_prompt_wav( + loaded_prompt_wavs[i], sampling_rate=sampling_rate + ) + if loaded_prompt_wavs[i].size(0) != 1: + loaded_prompt_wavs[i] = loaded_prompt_wavs[i].mean(0, keepdim=True) + + if len(loaded_prompt_wavs) == 1: + prompt_wav = loaded_prompt_wavs[0] + else: + prompt_wav = torch.cat(loaded_prompt_wavs, dim=1) + + # Remove edge and long silences in the prompt wav. + # Add 0.2s trailing silence to avoid leaking prompt to generated speech. + prompt_wav = remove_silence( + prompt_wav, sampling_rate, only_edge=False, trail_sil=200 + ) + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + prompt_duration = prompt_wav.shape[-1] / sampling_rate + + if prompt_duration > 40: + logging.warning( + f"Given prompt wav is too long ({prompt_duration}s). " + f"Please provide a shorter one (prompt shorter than 10 " + f"seconds is recommended)." + ) + elif prompt_duration > 20: + logging.warning( + f"Given prompt wav is long ({prompt_duration}s). " + f"It will lead to slower inference speed and possibly worse speech quality." + ) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + + # Add punctuation in the end if there is not + text = add_punctuation(text) + prompt_text = add_punctuation(prompt_text) + + # Tokenize text (str tokens), punctuations will be preserved. + tokens_str = tokenizer.texts_to_tokens([text])[0] + prompt_tokens_str = tokenizer.texts_to_tokens([prompt_text])[0] + + # chunk text so that each len(prompt wav + generated wav) is around 40 seconds. + token_duration = (prompt_wav.shape[-1] / sampling_rate) / ( + len(prompt_tokens_str) * speed + ) + max_tokens = int((40 - prompt_duration) / token_duration) + chunked_tokens_str = chunk_tokens_dialog(tokens_str, max_tokens=max_tokens) + + # Tokenize text (int tokens) + chunked_tokens = tokenizer.tokens_to_token_ids(chunked_tokens_str) + prompt_tokens = tokenizer.tokens_to_token_ids([prompt_tokens_str]) + + # Batchify chunked texts for faster processing + tokens_batches, chunked_index = batchify_tokens( + chunked_tokens, max_duration, prompt_duration, token_duration + ) + + # Start predicting features + chunked_features = [] + start_t = dt.datetime.now() + + for batch_tokens in tokens_batches: + batch_prompt_tokens = prompt_tokens * len(batch_tokens) + + batch_prompt_features = prompt_features.repeat(len(batch_tokens), 1, 1) + batch_prompt_features_lens = torch.full( + (len(batch_tokens),), prompt_features.size(1), device=device + ) + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=batch_tokens, + prompt_tokens=batch_prompt_tokens, + prompt_features=batch_prompt_features, + prompt_features_lens=batch_prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + chunked_features.append((pred_features, pred_features_lens)) + + # Start vocoder processing + chunked_wavs = [] + start_vocoder_t = dt.datetime.now() + + for pred_features, pred_features_lens in chunked_features: + batch_wav = [] + for i in range(pred_features.size(0)): + + wav = ( + vocoder.decode(pred_features[i][None, :, : pred_features_lens[i]]) + .squeeze(1) + .clamp(-1, 1) + ) + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + batch_wav.append(wav) + chunked_wavs.extend(batch_wav) + + # Finish model generation + t = (dt.datetime.now() - start_t).total_seconds() + + # Merge chunked wavs + indexed_chunked_wavs = [ + (index, wav) for index, wav in zip(chunked_index, chunked_wavs) + ] + sequential_indexed_chunked_wavs = sorted(indexed_chunked_wavs, key=lambda x: x[0]) + sequential_chunked_wavs = [ + sequential_indexed_chunked_wavs[i][1] + for i in range(len(sequential_indexed_chunked_wavs)) + ] + final_wav = cross_fade_concat( + sequential_chunked_wavs, fade_duration=0.1, sample_rate=sampling_rate + ) + final_wav = remove_silence( + final_wav, sampling_rate, only_edge=(not remove_long_sil), trail_sil=0 + ) + + # Calculate processing time metrics + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = final_wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + torchaudio.save(save_path, final_wav.cpu(), sample_rate=sampling_rate) + return metrics + + +def generate_sentence_stereo_raw_evaluation( + save_path: str, + prompt_text: str, + prompt_wav: Union[str, List[str]], + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: DialogTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + silence_wav: Optional[str] = None, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function directly feed the prompt_text, prompt_wav and text to the model. + It is not efficient and can have poor results for some inappropriate inputs. + (e.g., prompt wav contains long silence, text to be generated is too long) + This function can be used to evaluate the "raw" performance of the model. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (Union[str, List[str]]): Path to the prompt wav file, can be + one or two wav files, which corresponding to a merged conversational + speech or two seperate speaker's speech. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (DialogTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + silence_wav (str): Path of the silence wav file, used in two-channel + generation with single-channel prompts + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and preprocess prompt wav + if isinstance(prompt_wav, str): + prompt_wav = [ + prompt_wav, + ] + else: + assert len(prompt_wav) == 2 and isinstance(prompt_wav[0], str) + + loaded_prompt_wavs = prompt_wav + for i in range(len(prompt_wav)): + loaded_prompt_wavs[i] = load_prompt_wav( + loaded_prompt_wavs[i], sampling_rate=sampling_rate + ) + + if len(loaded_prompt_wavs) == 1: + assert ( + loaded_prompt_wavs[0].size(0) == 2 + ), "Merged prompt wav must be stereo for stereo dialogue generation" + prompt_wav = loaded_prompt_wavs[0] + + else: + assert len(loaded_prompt_wavs) == 2 + if loaded_prompt_wavs[0].size(0) == 2: + prompt_wav = torch.cat(loaded_prompt_wavs, dim=1) + else: + assert loaded_prompt_wavs[0].size(0) == 1 + silence_wav, silence_sampling_rate = torchaudio.load(silence_wav) + assert silence_sampling_rate == sampling_rate + prompt_wav = silence_wav[ + :, : loaded_prompt_wavs[0].size(1) + loaded_prompt_wavs[1].size(1) + ] + prompt_wav[0, : loaded_prompt_wavs[0].size(1)] = loaded_prompt_wavs[0] + prompt_wav[1, loaded_prompt_wavs[0].size(1) :] = loaded_prompt_wavs[1] + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + prompt_features_lens = torch.tensor([prompt_features.size(1)], device=device) + + # Convert text to tokens + tokens = tokenizer.texts_to_token_ids([text]) + prompt_tokens = tokenizer.texts_to_token_ids([prompt_text]) + + # Start timing + start_t = dt.datetime.now() + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + prompt_features_lens=prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + + # Start vocoder processing + start_vocoder_t = dt.datetime.now() + feat_dim = pred_features.size(1) // 2 + wav_left = vocoder.decode(pred_features[:, :feat_dim]).squeeze(1).clamp(-1, 1) + wav_right = ( + vocoder.decode(pred_features[:, feat_dim : feat_dim * 2]) + .squeeze(1) + .clamp(-1, 1) + ) + + wav = torch.cat([wav_left, wav_right], dim=0) + + # Calculate processing times and real-time factors + t = (dt.datetime.now() - start_t).total_seconds() + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + torchaudio.save(save_path, wav.cpu(), sample_rate=sampling_rate) + + return metrics + + +def generate_sentence_stereo( + save_path: str, + prompt_text: str, + prompt_wav: Union[str, List[str]], + text: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: DialogTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + silence_wav: Optional[str] = None, + max_duration: float = 100, + remove_long_sil: bool = False, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function will do the following to improve the generation quality: + 1. chunk the text according to speaker-turn symbol [S1]. + 2. process chunked texts in batches. + 3. remove long silences in the prompt audio. + 4. add punctuation to the end of prompt text and text if there is not. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (Union[str, List[str]]): Path to the prompt wav file, can be + one or two wav files, which corresponding to a merged conversational + speech or two seperate speaker's speech. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (DialogTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + device (torch.device): The device on which computations are performed. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + silence_wav (str): Path of the silence wav file, used in two-channel + generation with single-channel prompts + max_duration (float, optional): The maximum duration to process in each + batch. Used to control memory consumption when generating long audios. + remove_long_sil (bool, optional): Whether to remove long silences in the + middle of the generated speech (edge silences will be removed by default). + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and preprocess prompt wav + if isinstance(prompt_wav, str): + prompt_wav = [ + prompt_wav, + ] + else: + assert len(prompt_wav) == 2 and isinstance(prompt_wav[0], str) + + loaded_prompt_wavs = prompt_wav + for i in range(len(prompt_wav)): + loaded_prompt_wavs[i] = load_prompt_wav( + loaded_prompt_wavs[i], sampling_rate=sampling_rate + ) + + if len(loaded_prompt_wavs) == 1: + assert ( + loaded_prompt_wavs[0].size(0) == 2 + ), "Merged prompt wav must be stereo for stereo dialogue generation" + prompt_wav = loaded_prompt_wavs[0] + + else: + assert len(loaded_prompt_wavs) == 2 + if loaded_prompt_wavs[0].size(0) == 2: + prompt_wav = torch.cat(loaded_prompt_wavs, dim=1) + else: + assert loaded_prompt_wavs[0].size(0) == 1 + silence_wav, silence_sampling_rate = torchaudio.load(silence_wav) + assert silence_sampling_rate == sampling_rate + prompt_wav = silence_wav[ + :, : loaded_prompt_wavs[0].size(1) + loaded_prompt_wavs[1].size(1) + ] + prompt_wav[0, : loaded_prompt_wavs[0].size(1)] = loaded_prompt_wavs[0] + prompt_wav[1, loaded_prompt_wavs[0].size(1) :] = loaded_prompt_wavs[1] + + # Remove edge and long silences in the prompt wav. + # Add 0.2s trailing silence to avoid leaking prompt to generated speech. + prompt_wav = remove_silence( + prompt_wav, sampling_rate, only_edge=False, trail_sil=200 + ) + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + prompt_duration = prompt_wav.shape[-1] / sampling_rate + + if prompt_duration > 40: + logging.warning( + f"Given prompt wav is too long ({prompt_duration}s). " + f"Please provide a shorter one (prompt shorter than 10 " + f"seconds is recommended)." + ) + elif prompt_duration > 20: + logging.warning( + f"Given prompt wav is long ({prompt_duration}s). " + f"It will lead to slower inference speed and possibly worse speech quality." + ) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=sampling_rate + ).to(device) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + + # Add punctuation in the end if there is not + text = add_punctuation(text) + prompt_text = add_punctuation(prompt_text) + + # Tokenize text (str tokens), punctuations will be preserved. + tokens_str = tokenizer.texts_to_tokens([text])[0] + prompt_tokens_str = tokenizer.texts_to_tokens([prompt_text])[0] + + # chunk text so that each len(prompt wav + generated wav) is around 40 seconds. + token_duration = (prompt_wav.shape[-1] / sampling_rate) / ( + len(prompt_tokens_str) * speed + ) + max_tokens = int((40 - prompt_duration) / token_duration) + chunked_tokens_str = chunk_tokens_dialog(tokens_str, max_tokens=max_tokens) + + # Tokenize text (int tokens) + chunked_tokens = tokenizer.tokens_to_token_ids(chunked_tokens_str) + prompt_tokens = tokenizer.tokens_to_token_ids([prompt_tokens_str]) + + # Batchify chunked texts for faster processing + tokens_batches, chunked_index = batchify_tokens( + chunked_tokens, max_duration, prompt_duration, token_duration + ) + + # Start predicting features + chunked_features = [] + start_t = dt.datetime.now() + + for batch_tokens in tokens_batches: + batch_prompt_tokens = prompt_tokens * len(batch_tokens) + + batch_prompt_features = prompt_features.repeat(len(batch_tokens), 1, 1) + batch_prompt_features_lens = torch.full( + (len(batch_tokens),), prompt_features.size(1), device=device + ) + + # Generate features + ( + pred_features, + pred_features_lens, + pred_prompt_features, + pred_prompt_features_lens, + ) = model.sample( + tokens=batch_tokens, + prompt_tokens=batch_prompt_tokens, + prompt_features=batch_prompt_features, + prompt_features_lens=batch_prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration="predict", + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + chunked_features.append((pred_features, pred_features_lens)) + + # Start vocoder processing + chunked_wavs = [] + start_vocoder_t = dt.datetime.now() + + for pred_features, pred_features_lens in chunked_features: + batch_wav = [] + for i in range(pred_features.size(0)): + + feat_dim = pred_features.size(1) // 2 + wav_left = ( + vocoder.decode( + pred_features[i][None, :feat_dim, : pred_features_lens[i]] + ) + .squeeze(1) + .clamp(-1, 1) + ) + wav_right = ( + vocoder.decode( + pred_features[i][ + None, feat_dim : feat_dim * 2, : pred_features_lens[i] + ] + ) + .squeeze(1) + .clamp(-1, 1) + ) + wav = torch.cat([wav_left, wav_right], dim=0) + + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + batch_wav.append(wav) + chunked_wavs.extend(batch_wav) + + # Finish model generation + t = (dt.datetime.now() - start_t).total_seconds() + + # Merge chunked wavs + indexed_chunked_wavs = [ + (index, wav) for index, wav in zip(chunked_index, chunked_wavs) + ] + sequential_indexed_chunked_wavs = sorted(indexed_chunked_wavs, key=lambda x: x[0]) + sequential_chunked_wavs = [ + sequential_indexed_chunked_wavs[i][1] + for i in range(len(sequential_indexed_chunked_wavs)) + ] + final_wav = cross_fade_concat( + sequential_chunked_wavs, fade_duration=0.1, sample_rate=sampling_rate + ) + final_wav = remove_silence( + final_wav, sampling_rate, only_edge=(not remove_long_sil), trail_sil=0 + ) + + # Calculate processing time metrics + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = final_wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + torchaudio.save(save_path, final_wav.cpu(), sample_rate=sampling_rate) + return metrics + + +def generate_list( + model_name: str, + res_dir: str, + test_list: str, + model: torch.nn.Module, + vocoder: torch.nn.Module, + tokenizer: DialogTokenizer, + feature_extractor: VocosFbank, + device: torch.device, + num_step: int = 16, + guidance_scale: float = 1.5, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + silence_wav: Optional[str] = None, + raw_evaluation: bool = False, + max_duration: float = 100, + remove_long_sil: bool = False, +): + total_t = [] + total_t_no_vocoder = [] + total_t_vocoder = [] + total_wav_seconds = [] + + with open(test_list, "r") as fr: + lines = fr.readlines() + + for i, line in enumerate(lines): + items = line.strip().split("\t") + if len(items) == 6: + ( + wav_name, + prompt_text_1, + prompt_text_2, + prompt_wav_1, + prompt_wav_2, + text, + ) = items + prompt_text = f"[S1]{prompt_text_1}[S2]{prompt_text_2}" + prompt_wav = [prompt_wav_1, prompt_wav_2] + elif len(items) == 4: + wav_name, prompt_text, prompt_wav, text = items + else: + raise ValueError(f"Invalid line: {line}") + assert text.startswith("[S1]") + + save_path = f"{res_dir}/{wav_name}.wav" + + common_params = { + "save_path": save_path, + "prompt_text": prompt_text, + "prompt_wav": prompt_wav, + "text": text, + "model": model, + "vocoder": vocoder, + "tokenizer": tokenizer, + "feature_extractor": feature_extractor, + "device": device, + "num_step": num_step, + "guidance_scale": guidance_scale, + "speed": speed, + "t_shift": t_shift, + "target_rms": target_rms, + "feat_scale": feat_scale, + "sampling_rate": sampling_rate, + } + + if model_name == "zipvoice_dialog": + if raw_evaluation: + metrics = generate_sentence_raw_evaluation(**common_params) + else: + metrics = generate_sentence( + **common_params, + max_duration=max_duration, + remove_long_sil=remove_long_sil, + ) + else: + assert model_name == "zipvoice_dialog_stereo" + if raw_evaluation: + metrics = generate_sentence_stereo_raw_evaluation( + **common_params, + silence_wav=silence_wav, + ) + else: + metrics = generate_sentence_stereo( + **common_params, + silence_wav=silence_wav, + max_duration=max_duration, + remove_long_sil=remove_long_sil, + ) + logging.info(f"[Sentence: {i}] Saved to: {save_path}") + logging.info(f"[Sentence: {i}] RTF: {metrics['rtf']:.4f}") + total_t.append(metrics["t"]) + total_t_no_vocoder.append(metrics["t_no_vocoder"]) + total_t_vocoder.append(metrics["t_vocoder"]) + total_wav_seconds.append(metrics["wav_seconds"]) + + logging.info(f"Average RTF: {np.sum(total_t) / np.sum(total_wav_seconds):.4f}") + logging.info( + f"Average RTF w/o vocoder: " + f"{np.sum(total_t_no_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + logging.info( + f"Average RTF vocoder: " + f"{np.sum(total_t_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + + +@torch.inference_mode() +def main(): + parser = get_parser() + args = parser.parse_args() + + torch.set_num_threads(args.num_thread) + torch.set_num_interop_threads(args.num_thread) + + params = AttributeDict() + params.update(vars(args)) + fix_random_seed(params.seed) + + assert ( + params.test_list is not None + ), "For inference, please provide prompts and text with '--test-list'" + + if params.model_dir is not None: + params.model_dir = Path(params.model_dir) + if not params.model_dir.is_dir(): + raise FileNotFoundError(f"{params.model_dir} does not exist") + for filename in [params.checkpoint_name, "model.json", "tokens.txt"]: + if not (params.model_dir / filename).is_file(): + raise FileNotFoundError(f"{params.model_dir / filename} does not exist") + model_ckpt = params.model_dir / params.checkpoint_name + model_config = params.model_dir / "model.json" + token_file = params.model_dir / "tokens.txt" + logging.info( + f"Using {params.model_name} in local model dir {params.model_dir}, " + f"checkpoint {params.checkpoint_name}" + ) + else: + logging.info(f"Using pretrained {params.model_name} model from the Huggingface") + model_ckpt = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/model.pt" + ) + model_config = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/model.json" + ) + + token_file = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/tokens.txt" + ) + + tokenizer = DialogTokenizer(token_file=token_file) + + tokenizer_config = { + "vocab_size": tokenizer.vocab_size, + "pad_id": tokenizer.pad_id, + "spk_a_id": tokenizer.spk_a_id, + "spk_b_id": tokenizer.spk_b_id, + } + + with open(model_config, "r") as f: + model_config = json.load(f) + + if params.model_name == "zipvoice_dialog": + model = ZipVoiceDialog( + **model_config["model"], + **tokenizer_config, + ) + else: + assert params.model_name == "zipvoice_dialog_stereo" + model = ZipVoiceDialogStereo( + **model_config["model"], + **tokenizer_config, + ) + + if str(model_ckpt).endswith(".safetensors"): + safetensors.torch.load_model(model, model_ckpt) + elif str(model_ckpt).endswith(".pt"): + load_checkpoint(filename=model_ckpt, model=model, strict=True) + else: + raise NotImplementedError(f"Unsupported model checkpoint format: {model_ckpt}") + + if torch.cuda.is_available(): + params.device = torch.device("cuda", 0) + elif torch.backends.mps.is_available(): + params.device = torch.device("mps") + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + model = model.to(params.device) + model.eval() + + vocoder = get_vocoder(params.vocoder_path) + vocoder = vocoder.to(params.device) + vocoder.eval() + + if model_config["feature"]["type"] == "vocos": + if params.model_name == "zipvoice_dialog": + num_channels = 1 + else: + assert params.model_name == "zipvoice_dialog_stereo" + num_channels = 2 + feature_extractor = VocosFbank(num_channels=num_channels) + else: + raise NotImplementedError( + f"Unsupported feature type: {model_config['feature']['type']}" + ) + params.sampling_rate = model_config["feature"]["sampling_rate"] + + logging.info("Start generating...") + os.makedirs(params.res_dir, exist_ok=True) + generate_list( + model_name=params.model_name, + res_dir=params.res_dir, + test_list=params.test_list, + model=model, + vocoder=vocoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + device=params.device, + num_step=params.num_step, + guidance_scale=params.guidance_scale, + speed=params.speed, + t_shift=params.t_shift, + target_rms=params.target_rms, + feat_scale=params.feat_scale, + sampling_rate=params.sampling_rate, + silence_wav=params.silence_wav, + raw_evaluation=params.raw_evaluation, + max_duration=params.max_duration, + remove_long_sil=params.remove_long_sil, + ) + logging.info("Done") + + +if __name__ == "__main__": + + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + main() diff --git a/zipvoice/bin/infer_zipvoice_onnx.py b/zipvoice/bin/infer_zipvoice_onnx.py new file mode 100644 index 0000000..6852535 --- /dev/null +++ b/zipvoice/bin/infer_zipvoice_onnx.py @@ -0,0 +1,924 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu, +# Zengwei Yao) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script generates speech with our pre-trained ZipVoice or ZipVoice-Distill + ONNX models. If no local model is specified, + Required files will be automatically downloaded from HuggingFace. + +Usage: + +Note: If you having trouble connecting to HuggingFace, + try switching endpoint to mirror site: +export HF_ENDPOINT=https://hf-mirror.com + +(1) Inference of a single sentence: + +python3 -m zipvoice.bin.infer_zipvoice_onnx \ + --onnx-int8 False \ + --model-name zipvoice \ + --prompt-wav prompt.wav \ + --prompt-text "I am a prompt." \ + --text "I am a sentence." \ + --res-wav-path result.wav + +(2) Inference of a list of sentences: +python3 -m zipvoice.bin.infer_zipvoice_onnx \ + --onnx-int8 False \ + --model-name zipvoice \ + --test-list test.tsv \ + --res-dir results + +`--model-name` can be `zipvoice` or `zipvoice_distill`, + which are the models before and after distillation, respectively. + +Each line of `test.tsv` is in the format of + `{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}`. + +Set `--onnx-int8 True` to use int8 quantizated ONNX model. +Quantizated model has faster but lower quality. +""" + +import argparse +import datetime as dt +import json +import logging +import os +from pathlib import Path +from typing import List, Tuple + +import numpy as np +import onnxruntime as ort +import torch +import torchaudio +from huggingface_hub import hf_hub_download +from lhotse.utils import fix_random_seed +from torch import Tensor, nn + +from zipvoice.bin.infer_zipvoice import get_vocoder +from zipvoice.models.modules.solver import get_time_steps +from zipvoice.tokenizer.tokenizer import ( + EmiliaTokenizer, + EspeakTokenizer, + LibriTTSTokenizer, + SimpleTokenizer, +) +from zipvoice.utils.common import AttributeDict, str2bool +from zipvoice.utils.feature import VocosFbank +from zipvoice.utils.infer import ( + add_punctuation, + chunk_tokens_punctuation, + cross_fade_concat, + load_prompt_wav, + remove_silence, + rms_norm, +) + +HUGGINGFACE_REPO = "k2-fsa/ZipVoice" +MODEL_DIR = { + "zipvoice": "zipvoice", + "zipvoice_distill": "zipvoice_distill", +} + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--onnx-int8", + type=str2bool, + default=False, + help="Whether to use the int8 model", + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice", + choices=["zipvoice", "zipvoice_distill"], + help="The model used for inference", + ) + + parser.add_argument( + "--model-dir", + type=str, + default=None, + help="The path to the local onnx model. " + "Will download pre-trained checkpoint from huggingface if not specified.", + ) + + parser.add_argument( + "--vocoder-path", + type=str, + default=None, + help="The vocoder checkpoint. " + "Will download pre-trained vocoder from huggingface if not specified.", + ) + + parser.add_argument( + "--tokenizer", + type=str, + default="emilia", + choices=["emilia", "libritts", "espeak", "simple"], + help="Tokenizer type.", + ) + + parser.add_argument( + "--lang", + type=str, + default="en-us", + help="Language identifier, used when tokenizer type is espeak. see" + "https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md", + ) + + parser.add_argument( + "--test-list", + type=str, + default=None, + help="The list of prompt speech, prompt_transcription, " + "and text to synthesizein the format of " + "'{wav_name}\t{prompt_transcription}\t{prompt_wav}\t{text}'.", + ) + + parser.add_argument( + "--prompt-wav", + type=str, + default=None, + help="The prompt wav to mimic", + ) + + parser.add_argument( + "--prompt-text", + type=str, + default=None, + help="The transcription of the prompt wav", + ) + + parser.add_argument( + "--text", + type=str, + default=None, + help="The text to synthesize", + ) + + parser.add_argument( + "--res-dir", + type=str, + default="results", + help=""" + Path name of the generated wavs dir, + used when test-list is not None + """, + ) + + parser.add_argument( + "--res-wav-path", + type=str, + default="result.wav", + help=""" + Path name of the generated wav path, + used when test-list is None + """, + ) + + parser.add_argument( + "--guidance-scale", + type=float, + default=None, + help="The scale of classifier-free guidance during inference.", + ) + + parser.add_argument( + "--num-step", + type=int, + default=None, + help="The number of sampling steps.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--speed", + type=float, + default=1.0, + help="Control speech speed, 1.0 means normal, >1.0 means speed up", + ) + + parser.add_argument( + "--t-shift", + type=float, + default=0.5, + help="Shift t to smaller ones if t_shift < 1.0", + ) + + parser.add_argument( + "--target-rms", + type=float, + default=0.1, + help="Target speech normalization rms value, set to 0 to disable normalization", + ) + + parser.add_argument( + "--seed", + type=int, + default=666, + help="Random seed", + ) + + parser.add_argument( + "--num-thread", + type=int, + default=1, + help="Number of threads to use for ONNX Runtime and PyTorch.", + ) + + parser.add_argument( + "--raw-evaluation", + type=str2bool, + default=False, + help="Whether to use the 'raw' evaluation mode where provided " + "prompts and text are fed to the model without pre-processing", + ) + + parser.add_argument( + "--remove-long-sil", + type=str2bool, + default=False, + help="Whether to remove long silences in the middle of the generated " + "speech (edge silences will be removed by default).", + ) + return parser + + +class OnnxModel: + def __init__( + self, + text_encoder_path: str, + fm_decoder_path: str, + num_thread: int = 1, + ): + session_opts = ort.SessionOptions() + session_opts.inter_op_num_threads = num_thread + session_opts.intra_op_num_threads = num_thread + + self.session_opts = session_opts + + self.init_text_encoder(text_encoder_path) + self.init_fm_decoder(fm_decoder_path) + + def init_text_encoder(self, model_path: str): + self.text_encoder = ort.InferenceSession( + model_path, + sess_options=self.session_opts, + providers=["CPUExecutionProvider"], + ) + + def init_fm_decoder(self, model_path: str): + self.fm_decoder = ort.InferenceSession( + model_path, + sess_options=self.session_opts, + providers=["CPUExecutionProvider"], + ) + meta = self.fm_decoder.get_modelmeta().custom_metadata_map + self.feat_dim = int(meta["feat_dim"]) + + def run_text_encoder( + self, + tokens: Tensor, + prompt_tokens: Tensor, + prompt_features_len: Tensor, + speed: Tensor, + ) -> Tuple[Tensor, Tensor]: + out = self.text_encoder.run( + [ + self.text_encoder.get_outputs()[0].name, + ], + { + self.text_encoder.get_inputs()[0].name: tokens.numpy(), + self.text_encoder.get_inputs()[1].name: prompt_tokens.numpy(), + self.text_encoder.get_inputs()[2].name: prompt_features_len.numpy(), + self.text_encoder.get_inputs()[3].name: speed.numpy(), + }, + ) + return torch.from_numpy(out[0]) + + def run_fm_decoder( + self, + t: Tensor, + x: Tensor, + text_condition: Tensor, + speech_condition: torch.Tensor, + guidance_scale: Tensor, + ) -> Tensor: + out = self.fm_decoder.run( + [ + self.fm_decoder.get_outputs()[0].name, + ], + { + self.fm_decoder.get_inputs()[0].name: t.numpy(), + self.fm_decoder.get_inputs()[1].name: x.numpy(), + self.fm_decoder.get_inputs()[2].name: text_condition.numpy(), + self.fm_decoder.get_inputs()[3].name: speech_condition.numpy(), + self.fm_decoder.get_inputs()[4].name: guidance_scale.numpy(), + }, + ) + return torch.from_numpy(out[0]) + + +def sample( + model: OnnxModel, + tokens: List[List[int]], + prompt_tokens: List[List[int]], + prompt_features: Tensor, + speed: float = 1.0, + t_shift: float = 0.5, + guidance_scale: float = 1.0, + num_step: int = 16, +) -> torch.Tensor: + """ + Generate acoustic features, given text tokens, prompts feature and prompt + transcription's text tokens. + + Args: + tokens: a list of list of text tokens. + prompt_tokens: a list of list of prompt tokens. + prompt_features: the prompt feature with the shape + (batch_size, seq_len, feat_dim). + speed : speed control. + t_shift: time shift. + guidance_scale: the guidance scale for classifier-free guidance. + num_step: the number of steps to use in the ODE solver. + """ + # Run text encoder + assert len(tokens) == len(prompt_tokens) == 1 + tokens = torch.tensor(tokens, dtype=torch.int64) + prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.int64) + prompt_features_len = torch.tensor(prompt_features.size(1), dtype=torch.int64) + speed = torch.tensor(speed, dtype=torch.float32) + + text_condition = model.run_text_encoder( + tokens, prompt_tokens, prompt_features_len, speed + ) + + batch_size, num_frames, _ = text_condition.shape + assert batch_size == 1 + feat_dim = model.feat_dim + + # Run flow matching model + timesteps = get_time_steps( + t_start=0.0, + t_end=1.0, + num_step=num_step, + t_shift=t_shift, + ) + x = torch.randn(batch_size, num_frames, feat_dim) + speech_condition = torch.nn.functional.pad( + prompt_features, (0, 0, 0, num_frames - prompt_features.shape[1]) + ) # (B, T, F) + guidance_scale = torch.tensor(guidance_scale, dtype=torch.float32) + + for step in range(num_step): + v = model.run_fm_decoder( + t=timesteps[step], + x=x, + text_condition=text_condition, + speech_condition=speech_condition, + guidance_scale=guidance_scale, + ) + x = x + v * (timesteps[step + 1] - timesteps[step]) + + x = x[:, prompt_features_len.item() :, :] + return x + + +# Copied from zipvoice/bin/infer_zipvoice.py, but call an external sample function +def generate_sentence_raw_evaluation( + save_path: str, + prompt_text: str, + prompt_wav: str, + text: str, + model: OnnxModel, + vocoder: nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function directly feed the prompt_text, prompt_wav and text to the model. + It is not efficient and can have poor results for some inappropriate inputs. + (e.g., prompt wav contains long silence, text to be generated is too long) + This function can be used to evaluate the "raw" performance of the model. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (str): Path to the prompt wav file. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (EmiliaTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and process prompt wav + prompt_wav = load_prompt_wav(prompt_wav, sampling_rate=sampling_rate) + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract(prompt_wav, sampling_rate=sampling_rate) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + + # Convert text to tokens + tokens = tokenizer.texts_to_token_ids([text]) + prompt_tokens = tokenizer.texts_to_token_ids([prompt_text]) + + # Start timing + start_t = dt.datetime.now() + + # Generate features + pred_features = sample( + model=model, + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + speed=speed, + t_shift=t_shift, + guidance_scale=guidance_scale, + num_step=num_step, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + + # Start vocoder processing + start_vocoder_t = dt.datetime.now() + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + + # Calculate processing times and real-time factors + t = (dt.datetime.now() - start_t).total_seconds() + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + torchaudio.save(save_path, wav.cpu(), sample_rate=sampling_rate) + + return metrics + + +def generate_sentence( + save_path: str, + prompt_text: str, + prompt_wav: str, + text: str, + model: OnnxModel, + vocoder: nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + remove_long_sil: bool = False, +): + """ + Generate waveform of a text based on a given prompt waveform and its transcription, + this function will do the following to improve the generation quality: + 1. chunk the text according to punctuations. + 2. process chunked texts sequentially. + 3. remove long silences in the prompt audio. + 4. add punctuation to the end of prompt text and text if there is not. + + Args: + save_path (str): Path to save the generated wav. + prompt_text (str): Transcription of the prompt wav. + prompt_wav (str): Path to the prompt wav file. + text (str): Text to be synthesized into a waveform. + model (torch.nn.Module): The model used for generation. + vocoder (torch.nn.Module): The vocoder used to convert features to waveforms. + tokenizer (EmiliaTokenizer): The tokenizer used to convert text to tokens. + feature_extractor (VocosFbank): The feature extractor used to + extract acoustic features. + num_step (int, optional): Number of steps for decoding. Defaults to 16. + guidance_scale (float, optional): Scale for classifier-free guidance. + Defaults to 1.0. + speed (float, optional): Speed control. Defaults to 1.0. + t_shift (float, optional): Time shift. Defaults to 0.5. + target_rms (float, optional): Target RMS for waveform normalization. + Defaults to 0.1. + feat_scale (float, optional): Scale for features. + Defaults to 0.1. + sampling_rate (int, optional): Sampling rate for the waveform. + Defaults to 24000. + remove_long_sil (bool, optional): Whether to remove long silences in the + middle of the generated speech (edge silences will be removed by default). + Returns: + metrics (dict): Dictionary containing time and real-time + factor metrics for processing. + """ + + # Load and process prompt wav + prompt_wav = load_prompt_wav(prompt_wav, sampling_rate=sampling_rate) + + # Remove edge and long silences in the prompt wav. + # Add 0.2s trailing silence to avoid leaking prompt to generated speech. + prompt_wav = remove_silence( + prompt_wav, sampling_rate, only_edge=False, trail_sil=200 + ) + + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + prompt_duration = prompt_wav.shape[-1] / sampling_rate + + if prompt_duration > 20: + logging.warning( + f"Given prompt wav is too long ({prompt_duration}s). " + f"Please provide a shorter one (1-3 seconds is recommended)." + ) + elif prompt_duration > 10: + logging.warning( + f"Given prompt wav is long ({prompt_duration}s). " + f"It will lead to slower inference speed and possibly worse speech quality." + ) + + # Extract features from prompt wav + prompt_features = feature_extractor.extract(prompt_wav, sampling_rate=sampling_rate) + + prompt_features = prompt_features.unsqueeze(0) * feat_scale + + # Add punctuation in the end if there is not + text = add_punctuation(text) + prompt_text = add_punctuation(prompt_text) + + # Tokenize text (str tokens), punctuations will be preserved. + tokens_str = tokenizer.texts_to_tokens([text])[0] + prompt_tokens_str = tokenizer.texts_to_tokens([prompt_text])[0] + + # chunk text so that each len(prompt wav + generated wav) is around 25 seconds. + token_duration = (prompt_wav.shape[-1] / sampling_rate) / ( + len(prompt_tokens_str) * speed + ) + max_tokens = int((25 - prompt_duration) / token_duration) + chunked_tokens_str = chunk_tokens_punctuation(tokens_str, max_tokens=max_tokens) + print(len(chunked_tokens_str)) + print(chunked_tokens_str) + + # Tokenize text (int tokens) + chunked_tokens = tokenizer.tokens_to_token_ids(chunked_tokens_str) + prompt_tokens = tokenizer.tokens_to_token_ids([prompt_tokens_str]) + + # Start predicting features + chunked_features = [] + start_t = dt.datetime.now() + for tokens in chunked_tokens: + + # Generate features + pred_features = sample( + model=model, + tokens=[tokens], + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + speed=speed, + t_shift=t_shift, + guidance_scale=guidance_scale, + num_step=num_step, + ) + + # Postprocess predicted features + pred_features = pred_features.permute(0, 2, 1) / feat_scale # (B, C, T) + chunked_features.append(pred_features) + + # Start vocoder processing + chunked_wavs = [] + start_vocoder_t = dt.datetime.now() + + for pred_features in chunked_features: + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + # Adjust wav volume if necessary + if prompt_rms < target_rms: + wav = wav * prompt_rms / target_rms + chunked_wavs.append(wav) + + # Finish model generation + t = (dt.datetime.now() - start_t).total_seconds() + + # Merge chunked wavs + final_wav = cross_fade_concat( + chunked_wavs, fade_duration=0.1, sample_rate=sampling_rate + ) + final_wav = remove_silence( + final_wav, sampling_rate, only_edge=(not remove_long_sil), trail_sil=0 + ) + + # Calculate processing time metrics + t_no_vocoder = (start_vocoder_t - start_t).total_seconds() + t_vocoder = (dt.datetime.now() - start_vocoder_t).total_seconds() + wav_seconds = final_wav.shape[-1] / sampling_rate + rtf = t / wav_seconds + rtf_no_vocoder = t_no_vocoder / wav_seconds + rtf_vocoder = t_vocoder / wav_seconds + metrics = { + "t": t, + "t_no_vocoder": t_no_vocoder, + "t_vocoder": t_vocoder, + "wav_seconds": wav_seconds, + "rtf": rtf, + "rtf_no_vocoder": rtf_no_vocoder, + "rtf_vocoder": rtf_vocoder, + } + + torchaudio.save(save_path, final_wav.cpu(), sample_rate=sampling_rate) + return metrics + + +def generate_list( + res_dir: str, + test_list: str, + model: OnnxModel, + vocoder: nn.Module, + tokenizer: EmiliaTokenizer, + feature_extractor: VocosFbank, + num_step: int = 16, + guidance_scale: float = 1.0, + speed: float = 1.0, + t_shift: float = 0.5, + target_rms: float = 0.1, + feat_scale: float = 0.1, + sampling_rate: int = 24000, + raw_evaluation: bool = False, + remove_long_sil: bool = False, +): + total_t = [] + total_t_no_vocoder = [] + total_t_vocoder = [] + total_wav_seconds = [] + + with open(test_list, "r") as fr: + lines = fr.readlines() + + for i, line in enumerate(lines): + wav_name, prompt_text, prompt_wav, text = line.strip().split("\t") + save_path = f"{res_dir}/{wav_name}.wav" + + common_params = { + "save_path": save_path, + "prompt_text": prompt_text, + "prompt_wav": prompt_wav, + "text": text, + "model": model, + "vocoder": vocoder, + "tokenizer": tokenizer, + "feature_extractor": feature_extractor, + "num_step": num_step, + "guidance_scale": guidance_scale, + "speed": speed, + "t_shift": t_shift, + "target_rms": target_rms, + "feat_scale": feat_scale, + "sampling_rate": sampling_rate, + } + + if raw_evaluation: + metrics = generate_sentence_raw_evaluation(**common_params) + else: + metrics = generate_sentence( + **common_params, + remove_long_sil=remove_long_sil, + ) + logging.info(f"[Sentence: {i}] Saved to: {save_path}") + logging.info(f"[Sentence: {i}] RTF: {metrics['rtf']:.4f}") + total_t.append(metrics["t"]) + total_t_no_vocoder.append(metrics["t_no_vocoder"]) + total_t_vocoder.append(metrics["t_vocoder"]) + total_wav_seconds.append(metrics["wav_seconds"]) + + logging.info(f"Average RTF: {np.sum(total_t) / np.sum(total_wav_seconds):.4f}") + logging.info( + f"Average RTF w/o vocoder: " + f"{np.sum(total_t_no_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + logging.info( + f"Average RTF vocoder: " + f"{np.sum(total_t_vocoder) / np.sum(total_wav_seconds):.4f}" + ) + + +@torch.inference_mode() +def main(): + parser = get_parser() + args = parser.parse_args() + + torch.set_num_threads(args.num_thread) + torch.set_num_interop_threads(args.num_thread) + + params = AttributeDict() + params.update(vars(args)) + fix_random_seed(params.seed) + + model_defaults = { + "zipvoice": { + "num_step": 16, + "guidance_scale": 1.0, + }, + "zipvoice_distill": { + "num_step": 8, + "guidance_scale": 3.0, + }, + } + + model_specific_defaults = model_defaults.get(params.model_name, {}) + + for param, value in model_specific_defaults.items(): + if getattr(params, param) is None: + setattr(params, param, value) + logging.info(f"Setting {param} to default value: {value}") + + assert (params.test_list is not None) ^ ( + (params.prompt_wav and params.prompt_text and params.text) is not None + ), ( + "For inference, please provide prompts and text with either '--test-list'" + " or '--prompt-wav, --prompt-text and --text'." + ) + + if params.onnx_int8: + text_encoder_name = "text_encoder_int8.onnx" + fm_decoder_name = "fm_decoder_int8.onnx" + else: + text_encoder_name = "text_encoder.onnx" + fm_decoder_name = "fm_decoder.onnx" + + if params.model_dir is not None: + params.model_dir = Path(params.model_dir) + if not params.model_dir.is_dir(): + raise FileNotFoundError(f"{params.model_dir} does not exist") + + for filename in [ + text_encoder_name, + fm_decoder_name, + "model.json", + "tokens.txt", + ]: + if not (params.model_dir / filename).is_file(): + raise FileNotFoundError(f"{params.model_dir / filename} does not exist") + text_encoder_path = params.model_dir / text_encoder_name + fm_decoder_path = params.model_dir / fm_decoder_name + model_config = params.model_dir / "model.json" + token_file = params.model_dir / "tokens.txt" + logging.info(f"Using local model dir {params.model_dir}.") + else: + logging.info("Using pretrained model from the Huggingface") + text_encoder_path = hf_hub_download( + HUGGINGFACE_REPO, + filename=f"{MODEL_DIR[params.model_name]}/{text_encoder_name}", + ) + fm_decoder_path = hf_hub_download( + HUGGINGFACE_REPO, + filename=f"{MODEL_DIR[params.model_name]}/{fm_decoder_name}", + ) + model_config = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/model.json" + ) + + token_file = hf_hub_download( + HUGGINGFACE_REPO, filename=f"{MODEL_DIR[params.model_name]}/tokens.txt" + ) + + if params.tokenizer == "emilia": + tokenizer = EmiliaTokenizer(token_file=token_file) + elif params.tokenizer == "libritts": + tokenizer = LibriTTSTokenizer(token_file=token_file) + elif params.tokenizer == "espeak": + tokenizer = EspeakTokenizer(token_file=token_file, lang=params.lang) + else: + assert params.tokenizer == "simple" + tokenizer = SimpleTokenizer(token_file=token_file) + + with open(model_config, "r") as f: + model_config = json.load(f) + + model = OnnxModel(text_encoder_path, fm_decoder_path, num_thread=args.num_thread) + + vocoder = get_vocoder(params.vocoder_path) + vocoder.eval() + + if model_config["feature"]["type"] == "vocos": + feature_extractor = VocosFbank() + else: + raise NotImplementedError( + f"Unsupported feature type: {model_config['feature']['type']}" + ) + params.sampling_rate = model_config["feature"]["sampling_rate"] + + logging.info("Start generating...") + if params.test_list: + os.makedirs(params.res_dir, exist_ok=True) + generate_list( + res_dir=params.res_dir, + test_list=params.test_list, + model=model, + vocoder=vocoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + num_step=params.num_step, + guidance_scale=params.guidance_scale, + speed=params.speed, + t_shift=params.t_shift, + target_rms=params.target_rms, + feat_scale=params.feat_scale, + sampling_rate=params.sampling_rate, + raw_evaluation=params.raw_evaluation, + remove_long_sil=params.remove_long_sil, + ) + else: + assert ( + not params.raw_evaluation + ), "Raw evaluation is only valid with --test-list" + generate_sentence( + save_path=params.res_wav_path, + prompt_text=params.prompt_text, + prompt_wav=params.prompt_wav, + text=params.text, + model=model, + vocoder=vocoder, + tokenizer=tokenizer, + feature_extractor=feature_extractor, + num_step=params.num_step, + guidance_scale=params.guidance_scale, + speed=params.speed, + t_shift=params.t_shift, + target_rms=params.target_rms, + feat_scale=params.feat_scale, + sampling_rate=params.sampling_rate, + remove_long_sil=params.remove_long_sil, + ) + logging.info(f"Saved to: {params.res_wav_path}") + logging.info("Done") + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + main() diff --git a/zipvoice/bin/onnx_export.py b/zipvoice/bin/onnx_export.py new file mode 100644 index 0000000..64b5aae --- /dev/null +++ b/zipvoice/bin/onnx_export.py @@ -0,0 +1,429 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Zengwei Yao) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script exports a pre-trained ZipVoice or ZipVoice-Distill model from PyTorch to +ONNX. + +Usage: + +python3 -m zipvoice.bin.onnx_export \ + --model-name zipvoice \ + --model-dir exp/zipvoice \ + --checkpoint-name epoch-11-avg-4.pt \ + --onnx-model-dir exp/zipvoice + +`--model-name` can be `zipvoice` or `zipvoice_distill`, + which are the models before and after distillation, respectively. +""" + + +import argparse +import json +import logging +from pathlib import Path +from typing import Dict + +import onnx +import safetensors.torch +import torch +from onnxruntime.quantization import QuantType, quantize_dynamic +from torch import Tensor, nn + +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import SimpleTokenizer +from zipvoice.utils.checkpoint import load_checkpoint +from zipvoice.utils.common import AttributeDict +from zipvoice.utils.scaling_converter import convert_scaled_to_non_scaled + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--onnx-model-dir", + type=str, + default="exp", + help="Dir to the exported models", + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice", + choices=["zipvoice", "zipvoice_distill"], + help="The model used for inference", + ) + + parser.add_argument( + "--model-dir", + type=str, + default=None, + help="The model directory that contains model checkpoint, configuration " + "file model.json, and tokens file tokens.txt. Will download pre-trained " + "checkpoint from huggingface if not specified.", + ) + + parser.add_argument( + "--checkpoint-name", + type=str, + default="model.pt", + help="The name of model checkpoint.", + ) + + return parser + + +def add_meta_data(filename: str, meta_data: Dict[str, str]): + """Add meta data to an ONNX model. It is changed in-place. + + Args: + filename: + Filename of the ONNX model to be changed. + meta_data: + Key-value pairs. + """ + model = onnx.load(filename) + for key, value in meta_data.items(): + meta = model.metadata_props.add() + meta.key = key + meta.value = value + + onnx.save(model, filename) + + +class OnnxTextModel(nn.Module): + def __init__(self, model: nn.Module): + """A wrapper for ZipVoice text encoder.""" + super().__init__() + self.embed = model.embed + self.text_encoder = model.text_encoder + self.pad_id = model.pad_id + + def forward( + self, + tokens: Tensor, + prompt_tokens: Tensor, + prompt_features_len: Tensor, + speed: Tensor, + ) -> Tensor: + cat_tokens = torch.cat([prompt_tokens, tokens], dim=1) + cat_tokens = nn.functional.pad(cat_tokens, (0, 1), value=self.pad_id) + tokens_len = cat_tokens.shape[1] - 1 + padding_mask = (torch.arange(tokens_len + 1) == tokens_len).unsqueeze(0) + + embed = self.embed(cat_tokens) + embed = self.text_encoder(x=embed, t=None, padding_mask=padding_mask) + + features_len = torch.ceil( + (prompt_features_len / prompt_tokens.shape[1] * tokens_len / speed) + ).to(dtype=torch.int64) + + token_dur = torch.div(features_len, tokens_len, rounding_mode="floor").to( + dtype=torch.int64 + ) + + # If you pass a scalar tensor, ONNX may infer the shape as [1] (rank-1 tensor), + # but sometimes expects an actual scalar (rank-0). + # When exporting, ONNX may generate a model where Concat expects inputs of the + # same rank, but receives [1] and []. + # In PyTorch, this is usually fine. In ONNX Runtime (C++), this causes the error like + # "Ranks of input data are different, cannot concatenate them. expected rank: 1 got: 2" + # If you use x.item(), ONNX loses the dynamic link and the input mismatch error can happen at inference. + # use reshape(()) to convert a rank-1 tensor to a rank-0 tensor. + + token_dur = token_dur.reshape(()) + features_len = features_len.reshape(()) + + text_condition = embed[:, :-1, :].unsqueeze(2).expand(-1, -1, token_dur, -1) + text_condition = text_condition.reshape(embed.shape[0], -1, embed.shape[2]) + + text_condition = torch.cat( + [ + text_condition, + embed[:, -1:, :].expand(-1, features_len - text_condition.shape[1], -1), + ], + dim=1, + ) + + return text_condition + + +class OnnxFlowMatchingModel(nn.Module): + def __init__(self, model: nn.Module, distill: bool = False): + """A wrapper for ZipVoice flow-matching decoder.""" + super().__init__() + self.distill = distill + self.fm_decoder = model.fm_decoder + self.model_func = getattr(model, "forward_fm_decoder") + self.feat_dim = model.feat_dim + + def forward( + self, + t: Tensor, + x: Tensor, + text_condition: Tensor, + speech_condition: torch.Tensor, + guidance_scale: Tensor, + ) -> Tensor: + if self.distill: + return self.model_func( + t=t, + xt=x, + text_condition=text_condition, + speech_condition=speech_condition, + guidance_scale=guidance_scale, + ) + else: + x = x.repeat(2, 1, 1) + text_condition = torch.cat( + [torch.zeros_like(text_condition), text_condition], dim=0 + ) + speech_condition = torch.cat( + [ + torch.where( + t > 0.5, torch.zeros_like(speech_condition), speech_condition + ), + speech_condition, + ], + dim=0, + ) + guidance_scale = torch.where(t > 0.5, guidance_scale, guidance_scale * 2.0) + data_uncond, data_cond = self.model_func( + t=t, + xt=x, + text_condition=text_condition, + speech_condition=speech_condition, + ).chunk(2, dim=0) + v = (1 + guidance_scale) * data_cond - guidance_scale * data_uncond + return v + + +def export_text_encoder( + model: OnnxTextModel, + filename: str, + opset_version: int = 13, +) -> None: + """Export the text encoder model to ONNX format. + + Args: + model: + The input model + filename: + The filename to save the exported ONNX model. + opset_version: + The opset version to use. + """ + tokens = torch.tensor([[2, 3, 4, 5]], dtype=torch.int64) + prompt_tokens = torch.tensor([[0, 1]], dtype=torch.int64) + prompt_features_len = torch.tensor(10, dtype=torch.int64) + speed = torch.tensor(1.0, dtype=torch.float32) + + model = torch.jit.trace(model, (tokens, prompt_tokens, prompt_features_len, speed)) + + torch.onnx.export( + model, + (tokens, prompt_tokens, prompt_features_len, speed), + filename, + verbose=False, + opset_version=opset_version, + input_names=["tokens", "prompt_tokens", "prompt_features_len", "speed"], + output_names=["text_condition"], + dynamic_axes={ + "tokens": {0: "N", 1: "T"}, + "prompt_tokens": {0: "N", 1: "T"}, + "text_condition": {0: "N", 1: "T"}, + }, + ) + + meta_data = { + "version": "1", + "model_author": "k2-fsa", + "comment": "ZipVoice text encoder", + "use_espeak": "1", + "use_pinyin": "1", + } + logging.info(f"meta_data: {meta_data}") + add_meta_data(filename=filename, meta_data=meta_data) + + logging.info(f"Exported to {filename}") + + +def export_fm_decoder( + model: OnnxFlowMatchingModel, + filename: str, + opset_version: int = 13, +) -> None: + """Export the flow matching decoder model to ONNX format. + + Args: + model: + The input model + filename: + The filename to save the exported ONNX model. + opset_version: + The opset version to use. + """ + feat_dim = model.feat_dim + seq_len = 200 + t = torch.tensor(0.5, dtype=torch.float32) + x = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + text_condition = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + speech_condition = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + guidance_scale = torch.tensor(1.0, dtype=torch.float32) + + model = torch.jit.trace( + model, (t, x, text_condition, speech_condition, guidance_scale) + ) + + torch.onnx.export( + model, + (t, x, text_condition, speech_condition, guidance_scale), + filename, + verbose=False, + opset_version=opset_version, + input_names=["t", "x", "text_condition", "speech_condition", "guidance_scale"], + output_names=["v"], + dynamic_axes={ + "x": {0: "N", 1: "T"}, + "text_condition": {0: "N", 1: "T"}, + "speech_condition": {0: "N", 1: "T"}, + "v": {0: "N", 1: "T"}, + }, + ) + + meta_data = { + "version": "1", + "model_author": "k2-fsa", + "comment": "ZipVoice flow-matching decoder", + "feat_dim": str(feat_dim), + "sample_rate": "24000", + "n_fft": "1024", + "hop_length": "256", + "window_length": "1024", + "num_mels": "100", + } + logging.info(f"meta_data: {meta_data}") + add_meta_data(filename=filename, meta_data=meta_data) + + logging.info(f"Exported to {filename}") + + +@torch.no_grad() +def main(): + parser = get_parser() + args = parser.parse_args() + + params = AttributeDict() + params.update(vars(args)) + + params.model_dir = Path(params.model_dir) + if not params.model_dir.is_dir(): + raise FileNotFoundError(f"{params.model_dir} does not exist") + for filename in [params.checkpoint_name, "model.json", "tokens.txt"]: + if not (params.model_dir / filename).is_file(): + raise FileNotFoundError(f"{params.model_dir / filename} does not exist") + model_ckpt = params.model_dir / params.checkpoint_name + model_config = params.model_dir / "model.json" + token_file = params.model_dir / "tokens.txt" + + logging.info(f"Loading model from {params.model_dir}") + + tokenizer = SimpleTokenizer(token_file) + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + + with open(model_config, "r") as f: + model_config = json.load(f) + + if params.model_name == "zipvoice": + model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + distill = False + else: + assert params.model_name == "zipvoice_distill" + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + distill = True + + if str(model_ckpt).endswith(".safetensors"): + safetensors.torch.load_model(model, model_ckpt) + elif str(model_ckpt).endswith(".pt"): + load_checkpoint(filename=model_ckpt, model=model, strict=True) + else: + raise NotImplementedError(f"Unsupported model checkpoint format: {model_ckpt}") + + device = torch.device("cpu") + model = model.to(device) + model.eval() + + convert_scaled_to_non_scaled(model, inplace=True, is_onnx=True) + + logging.info("Exporting model") + onnx_model_dir = Path(params.onnx_model_dir) + onnx_model_dir.mkdir(parents=True, exist_ok=True) + opset_version = 13 + + text_encoder = OnnxTextModel(model=model) + text_encoder_file = onnx_model_dir / "text_encoder.onnx" + export_text_encoder( + model=text_encoder, + filename=text_encoder_file, + opset_version=opset_version, + ) + + fm_decoder = OnnxFlowMatchingModel(model=model, distill=distill) + fm_decoder_file = onnx_model_dir / "fm_decoder.onnx" + export_fm_decoder( + model=fm_decoder, + filename=fm_decoder_file, + opset_version=opset_version, + ) + + logging.info("Generate int8 quantization models") + + text_encoder_int8_file = onnx_model_dir / "text_encoder_int8.onnx" + quantize_dynamic( + model_input=text_encoder_file, + model_output=text_encoder_int8_file, + op_types_to_quantize=["MatMul"], + weight_type=QuantType.QInt8, + ) + + fm_decoder_int8_file = onnx_model_dir / "fm_decoder_int8.onnx" + quantize_dynamic( + model_input=fm_decoder_file, + model_output=fm_decoder_int8_file, + op_types_to_quantize=["MatMul"], + weight_type=QuantType.QInt8, + ) + + logging.info("Done!") + + +if __name__ == "__main__": + + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + main() diff --git a/zipvoice/bin/prepare_dataset.py b/zipvoice/bin/prepare_dataset.py new file mode 100644 index 0000000..e3a2b28 --- /dev/null +++ b/zipvoice/bin/prepare_dataset.py @@ -0,0 +1,274 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script generates lhotse manifest files from TSV files for custom datasets. + +Each line of the TSV files should be in one of the following formats: +1. "{uniq_id}\t{text}\t{wav_path}" if the text corresponds to the full wav", +2. "{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} if text corresponds + to part of the wav. The start_time and end_time specify the start and end + times of the text within the wav, which should be in seconds. + +Note: {uniq_id} must be unique for each line. + +Usage: + +Suppose you have two TSV files: "custom_train.tsv" and "custom_dev.tsv", +where "custom" is your dataset name, "train"/"dev" are used for training and +validation respectively. + +(1) Prepare the training data + +python3 -m zipvoice.bin.prepare_dataset \ + --tsv-path data/raw/custom_train.tsv \ + --prefix "custom" \ + --subset "train" \ + --num-jobs 20 \ + --output-dir "data/manifests" + +The output file would be "data/manifests/custom_cuts_train.jsonl.gz". + +(2) Prepare the validation data + +python3 -m zipvoice.bin.prepare_dataset \ + --tsv-path data/raw/custom_dev.tsv \ + --prefix "custom" \ + --subset "dev" \ + --num-jobs 1 \ + --output-dir "data/manifests" + +The output file would be "data/manifests/custom_cuts_dev.jsonl.gz". + +""" + +import argparse +import logging +import re +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import List, Optional, Tuple + +from lhotse import CutSet, validate_recordings_and_supervisions +from lhotse.audio import Recording, RecordingSet +from lhotse.qa import fix_manifests +from lhotse.supervision import SupervisionSegment, SupervisionSet +from lhotse.utils import Pathlike +from tqdm.auto import tqdm + + +def get_args(): + parser = argparse.ArgumentParser() + + parser.add_argument( + "--tsv-path", + type=str, + help="The path of the tsv file. Each line should be in the format: " + "{uniq_id}\t{text}\t{wav_path}\t{start_time}\t{end_time} " + "if text corresponds to part of the wav or {uniq_id}\t{text}\t{wav_path} " + "if the text corresponds to the full wav", + ) + parser.add_argument( + "--prefix", + type=str, + default="custom", + help="Prefix of the output manifest file.", + ) + + parser.add_argument( + "--subset", + type=str, + default="train", + help="Subset name manifest file, typically train or dev.", + ) + + parser.add_argument( + "--num-jobs", + type=int, + default=20, + help="Number of jobs to processing.", + ) + + parser.add_argument( + "--output-dir", + type=str, + default="data/manifests", + help="The destination directory of manifest files.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=24000, + help="The target sampling rate.", + ) + return parser.parse_args() + + +def _parse_recording( + wav_path: str, +) -> Tuple[Recording, str]: + """ + :param wav_path: Path to the audio file + :return: a tuple of "recording" and "recording_id" + """ + + recording_id = wav_path.replace("/", "_").replace(".", "_") + recording = Recording.from_file(path=wav_path, recording_id=recording_id) + + return recording, recording_id + + +def _parse_supervision( + supervision: List, recording_dict: dict +) -> Optional[SupervisionSegment]: + """ + :param line: A line from the TSV file + :param recording_dict: Dictionary mapping recording IDs to Recording objects + :return: A SupervisionSegment object + """ + + uniq_id, text, wav_path, start, end = supervision + try: + recording_id = wav_path.replace("/", "_").replace(".", "_") + + recording = recording_dict[recording_id] + duration = end - start if end is not None else recording.duration + assert duration <= recording.duration, f"Duration {duration} is greater than " + f"recording duration {recording.duration}" + + text = re.sub("_", " ", text) # "_" is treated as padding symbol + text = re.sub(r"\s+", " ", text) # remove extra whitespace + + return SupervisionSegment( + id=f"{uniq_id}", + recording_id=recording.id, + start=start, + duration=duration, + channel=recording.channel_ids, + text=text.strip(), + ) + except Exception as e: + logging.warning(f"Error processing line: {e}") + return None + + +def prepare_dataset( + tsv_path: Pathlike, + prefix: str, + subset: str, + sampling_rate: int, + num_jobs: int, + output_dir: Pathlike, +): + """ + Returns the manifests which consist of the Recordings and Supervisions + + :param tsv_path: Path to the TSV file + :param output_dir: Path where to write the manifests + :param num_jobs: Number of processes for parallel processing + :return: The CutSet containing the data + """ + logging.info(f"Preparing {prefix} dataset {subset} subset.") + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + file_name = f"{prefix}_cuts_{subset}.jsonl.gz" + if (output_dir / file_name).is_file(): + logging.info(f"{file_name} exists, skipping.") + return + + # Step 1: Read all unique recording paths + recordings_path_set = set() + supervision_list = list() + with open(tsv_path, "r") as fr: + for line in fr: + items = line.strip().split("\t") + if len(items) == 3: + uniq_id, text, wav_path = items + start, end = 0, None + elif len(items) == 5: + uniq_id, text, wav_path, start, end = items + start, end = float(start), float(end) + else: + raise ValueError( + f"Invalid line format: {line}," + "requries to be 3 columns or 5 columns" + ) + recordings_path_set.add(wav_path) + supervision_list.append((uniq_id, text, wav_path, start, end)) + + logging.info("Starting to process recordings...") + # Step 2: Process recordings + futures = [] + recording_dict = {} + with ThreadPoolExecutor(max_workers=num_jobs) as ex: + for wav_path in tqdm(recordings_path_set, desc="Submitting jobs"): + futures.append(ex.submit(_parse_recording, wav_path)) + + for future in tqdm(futures, desc="Processing recordings"): + try: + recording, recording_id = future.result() + recording_dict[recording_id] = recording + except Exception as e: + logging.warning( + f"Error processing recording {recording_id} with error: {e}" + ) + + recording_set = RecordingSet.from_recordings(recording_dict.values()) + + logging.info("Starting to process supervisions...") + # Step 3: Process supervisions + supervisions = [] + for supervision in tqdm(supervision_list, desc="Processing supervisions"): + seg = _parse_supervision(supervision, recording_dict) + if seg is not None: + supervisions.append(seg) + + logging.info("Processing Cuts...") + + # Step 4: Create and validate manifests + supervision_set = SupervisionSet.from_segments(supervisions) + + recording_set, supervision_set = fix_manifests(recording_set, supervision_set) + validate_recordings_and_supervisions(recording_set, supervision_set) + + cut_set = CutSet.from_manifests( + recordings=recording_set, supervisions=supervision_set + ) + cut_set = cut_set.sort_by_recording_id() + cut_set = cut_set.resample(sampling_rate) + cut_set = cut_set.trim_to_supervisions(keep_overlapping=False) + + logging.info(f"Saving file to {output_dir / file_name}") + # Step 5: Write manifests to disk + cut_set.to_file(output_dir / file_name) + logging.info("Done!") + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + args = get_args() + + prepare_dataset( + tsv_path=args.tsv_path, + prefix=args.prefix, + subset=args.subset, + sampling_rate=args.sampling_rate, + num_jobs=args.num_jobs, + output_dir=args.output_dir, + ) diff --git a/zipvoice/bin/prepare_tokens.py b/zipvoice/bin/prepare_tokens.py new file mode 100644 index 0000000..2847e86 --- /dev/null +++ b/zipvoice/bin/prepare_tokens.py @@ -0,0 +1,103 @@ +""" +This file reads the texts in given manifest and save the new cuts with prepared tokens. +""" + +import argparse +import logging +from functools import partial +from pathlib import Path + +from lhotse import load_manifest, split_parallelize_combine + +from zipvoice.tokenizer.tokenizer import add_tokens + + +def get_args(): + parser = argparse.ArgumentParser() + + parser.add_argument( + "--input-file", + type=str, + help="Input manifest without tokens", + ) + + parser.add_argument( + "--output-file", + type=str, + help="Output manifest with tokens.", + ) + + parser.add_argument( + "--num-jobs", + type=int, + default=20, + help="Number of jobs to run in parallel.", + ) + + parser.add_argument( + "--tokenizer", + type=str, + default="emilia", + choices=["emilia", "espeak", "dialog", "libritts", "simple"], + help="The destination directory of manifest files.", + ) + + parser.add_argument( + "--lang", + type=str, + default="en-us", + help="Language identifier, used when tokenizer type is espeak. see" + "https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md", + ) + + return parser.parse_args() + + +def prepare_tokens( + input_file: Path, + output_file: Path, + num_jobs: int, + tokenizer: str, + lang: str = "en-us", +): + logging.info(f"Processing {input_file}") + if output_file.is_file(): + logging.info(f"{output_file} exists, skipping.") + return + logging.info(f"loading manifest from {input_file}") + cut_set = load_manifest(input_file) + + _add_tokens = partial(add_tokens, tokenizer=tokenizer, lang=lang) + + logging.info("Adding tokens") + + cut_set = split_parallelize_combine( + num_jobs=num_jobs, manifest=cut_set, fn=_add_tokens + ) + + logging.info(f"Saving file to {output_file}") + cut_set.to_file(output_file) + + +if __name__ == "__main__": + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + + args = get_args() + input_file = Path(args.input_file) + output_file = Path(args.output_file) + num_jobs = args.num_jobs + tokenizer = args.tokenizer + lang = args.lang + + output_file.parent.mkdir(parents=True, exist_ok=True) + + prepare_tokens( + input_file=input_file, + output_file=output_file, + num_jobs=num_jobs, + tokenizer=tokenizer, + lang=lang, + ) + + logging.info("Done!") diff --git a/zipvoice/bin/tensorrt_export.py b/zipvoice/bin/tensorrt_export.py new file mode 100644 index 0000000..dab3d87 --- /dev/null +++ b/zipvoice/bin/tensorrt_export.py @@ -0,0 +1,382 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Zengwei Yao) +# Copyright 2025 Nvidia Corp. (authors: Yuekai Zhang) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script exports a pre-trained ZipVoice or ZipVoice-Distill model from PyTorch to +ONNX. + +Usage: + +python3 -m zipvoice.bin.tensorrt_export \ + --model-name zipvoice_distill \ + --model-dir models/zipvoice_distill \ + --checkpoint-name model.pt \ + --trt-engine-file-name fm_decoder.fp16.max_batch_4.plan \ + --tensorrt-model-dir models/zipvoice_distill_trt || exit 1 + +`--model-name` can be `zipvoice` or `zipvoice_distill`, + which are the models before and after distillation, respectively. +""" + + +import argparse +import json +import logging +from pathlib import Path +from typing import Dict +import math + +import safetensors.torch +import torch +from torch import Tensor, nn + +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import SimpleTokenizer +from zipvoice.utils.checkpoint import load_checkpoint +from zipvoice.utils.common import AttributeDict +from zipvoice.utils.scaling_converter import convert_scaled_to_non_scaled +from zipvoice.models.modules.zipformer import CompactRelPositionalEncoding + +# Monkey-patching CompactRelPositionalEncoding.extend_pe +def extend_pe(self, x: Tensor, left_context_len: int = 0) -> None: + """Reset the positional encodings.""" + T = x.size(0) + left_context_len + + # if self.pe is not None: + # # self.pe contains both positive and negative parts + # # the length of self.pe is 2 * input_len - 1 + # if self.pe.size(0) >= T * 2 - 1: + # self.pe = self.pe.to(dtype=x.dtype, device=x.device) + # return + + # if T == 4, x would contain [ -3, -2, 1, 0, 1, 2, 3 ] + x = torch.arange(-(T - 1), T, device=x.device).to(torch.float32).unsqueeze(1) + + freqs = 1 + torch.arange(self.embed_dim // 2, device=x.device) + + # `compression_length` this is arbitrary/heuristic, if it is larger we have more + # resolution for small time offsets but less resolution for large time offsets. + compression_length = self.embed_dim**0.5 + # x_compressed, like X, goes from -infinity to infinity as T goes from -infinity + # to infinity; but it does so more slowly than T for large absolute values of T. + # The formula is chosen so that d(x_compressed )/dx is 1 around x == 0, which is + # important. + x_compressed = ( + compression_length + * x.sign() + * ((x.abs() + compression_length).log() - math.log(compression_length)) + ) + + # if self.length_factor == 1.0, then length_scale is chosen so that the + # FFT can exactly separate points close to the origin (T == 0). So this + # part of the formulation is not really heuristic. + # But empirically, for ASR at least, length_factor > 1.0 seems to work better. + length_scale = self.length_factor * self.embed_dim / (2.0 * math.pi) + + # note for machine implementations: if atan is not available, we can use: + # x.sign() * ((1 / (x.abs() + 1)) - 1) * (-math.pi/2) + # check on wolframalpha.com: plot(sign(x) * (1 / ( abs(x) + 1) - 1 ) * -pi/2 , + # atan(x)) + x_atan = (x_compressed / length_scale).atan() # results between -pi and pi + + cosines = (x_atan * freqs).cos() + sines = (x_atan * freqs).sin() + + pe = torch.zeros(x.shape[0], self.embed_dim, device=x.device) + pe[:, 0::2] = cosines + pe[:, 1::2] = sines + pe[:, -1] = 1.0 # for bias. + + self.pe = pe.to(dtype=x.dtype) + + +CompactRelPositionalEncoding.extend_pe = extend_pe + + +def get_trt_kwargs_dynamic_batch( + min_batch_size: int = 1, + opt_batch_size: int = 2, + max_batch_size: int = 4, +) -> Dict: + """Get keyword arguments for TensorRT with dynamic batch size.""" + feat_dim = 300 + min_seq_len = 100 + opt_seq_len = 200 + max_seq_len = 3000 + min_shape = [(min_batch_size, min_seq_len, feat_dim), (min_batch_size,), (min_batch_size, min_seq_len), (min_batch_size,)] + opt_shape = [(opt_batch_size, opt_seq_len, feat_dim), (opt_batch_size,), (opt_batch_size, opt_seq_len), (opt_batch_size,)] + max_shape = [(max_batch_size, max_seq_len, feat_dim), (max_batch_size,), (max_batch_size, max_seq_len), (max_batch_size,)] + input_names = ["x", "t", "padding_mask", "guidance_scale"] + return { + "min_shape": min_shape, + "opt_shape": opt_shape, + "max_shape": max_shape, + "input_names": input_names, + } + + +def convert_onnx_to_trt( + trt_model: str, trt_kwargs: Dict, onnx_model: str, dtype: torch.dtype = torch.float16 +): + """ + Convert an ONNX model to a TensorRT engine. + + Args: + trt_model (str): The path to save the TensorRT engine. + trt_kwargs (Dict): Keyword arguments for TensorRT. + onnx_model (str): The path to the ONNX model. + dtype (torch.dtype, optional): The data type to use. Defaults to torch.float16. + """ + logging.info("Converting onnx to trt...") + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + logger = trt.Logger(trt.Logger.INFO) + builder = trt.Builder(logger) + network = builder.create_network(network_flags) + parser = trt.OnnxParser(network, logger) + config = builder.create_builder_config() + # config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 32) # 4GB + if dtype == torch.float16: + config.set_flag(trt.BuilderFlag.FP16) + + profile = builder.create_optimization_profile() + # load onnx model + with open(onnx_model, "rb") as f: + if not parser.parse(f.read()): + for error in range(parser.num_errors): + print(parser.get_error(error)) + raise ValueError('failed to parse {}'.format(onnx_model)) + # set input shapes + for i in range(len(trt_kwargs['input_names'])): + profile.set_shape(trt_kwargs['input_names'][i], trt_kwargs['min_shape'][i], trt_kwargs['opt_shape'][i], trt_kwargs['max_shape'][i]) + if dtype == torch.float16: + tensor_dtype = trt.DataType.HALF + elif dtype == torch.bfloat16: + tensor_dtype = trt.DataType.BF16 + elif dtype == torch.float32: + tensor_dtype = trt.DataType.FLOAT + else: + raise ValueError('invalid dtype {}'.format(dtype)) + # set input and output data type + for i in range(network.num_inputs): + input_tensor = network.get_input(i) + input_tensor.dtype = tensor_dtype + for i in range(network.num_outputs): + output_tensor = network.get_output(i) + output_tensor.dtype = tensor_dtype + config.add_optimization_profile(profile) + engine_bytes = builder.build_serialized_network(network, config) + # save trt engine + with open(trt_model, "wb") as f: + f.write(engine_bytes) + logging.info("Succesfully convert onnx to trt...") + + +def get_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--tensorrt-model-dir", + type=str, + default="exp", + help="Dir to the exported models", + ) + + parser.add_argument( + "--model-name", + type=str, + default="zipvoice", + choices=["zipvoice", "zipvoice_distill"], + help="The model used for inference", + ) + + parser.add_argument( + "--model-dir", + type=str, + default=None, + help="The model directory that contains model checkpoint, configuration " + "file model.json, and tokens file tokens.txt. Will download pre-trained " + "checkpoint from huggingface if not specified.", + ) + + parser.add_argument( + "--checkpoint-name", + type=str, + default="model.pt", + help="The name of model checkpoint.", + ) + + parser.add_argument( + "--trt-engine-file-name", + type=str, + default=None, + help="The name of TensorRT engine file.", + ) + + parser.add_argument( + "--max-batch-size", + type=int, + default=4, + help="The maximum batch size to use for TensorRT.", + ) + + return parser + +def export_onnx_fm_decoder( + model: torch.nn.Module, + filename: str, + opset_version: int = 18, + distill: bool = False, +) -> None: + """Export the flow matching decoder model to ONNX format. + + Args: + model: + The input model + filename: + The filename to save the exported ONNX model. + opset_version: + The opset version to use. + """ + + + feat_dim, seq_len = model.feat_dim, 200 + + t = torch.tensor(0.5, dtype=torch.float32).unsqueeze(0) + guidance_scale = torch.tensor(1.0, dtype=torch.float32).unsqueeze(0) + padding_mask = torch.zeros(1, seq_len, dtype=torch.bool) + x = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + text_condition = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + speech_condition = torch.randn(1, seq_len, feat_dim, dtype=torch.float32) + xt= torch.cat([x, text_condition, speech_condition], dim=2) + xt = xt.repeat(2, 1, 1) + t = t.repeat(2) + padding_mask = padding_mask.repeat(2, 1) + guidance_scale = guidance_scale.repeat(2) + + inputs_tensors = [xt, t, padding_mask] + input_names = ['x', 't', 'padding_mask'] + dynamic_axes = { + 'x': {0: 'N', 1: 'T'}, + 't': {0: 'N'}, + 'padding_mask': {0: 'N', 1: 'T'}, + } + if distill: + inputs_tensors.append(guidance_scale) + input_names.append('guidance_scale') + dynamic_axes['guidance_scale'] = {0: 'N'} + estimator = model.fm_decoder + estimator = torch.jit.trace(estimator, inputs_tensors) + torch.onnx.export( + estimator, + inputs_tensors, + filename, + opset_version=opset_version, + input_names=input_names, + output_names=['v'], + dynamic_axes=dynamic_axes, + dynamo=False, + ) + logging.info(f"Exported to {filename}") + + +@torch.no_grad() +def main(): + parser = get_parser() + args = parser.parse_args() + + params = AttributeDict() + params.update(vars(args)) + + params.model_dir = Path(params.model_dir) + if not params.model_dir.is_dir(): + raise FileNotFoundError(f"{params.model_dir} does not exist") + for filename in [params.checkpoint_name, "model.json", "tokens.txt"]: + if not (params.model_dir / filename).is_file(): + raise FileNotFoundError(f"{params.model_dir / filename} does not exist") + model_ckpt = params.model_dir / params.checkpoint_name + model_config = params.model_dir / "model.json" + token_file = params.model_dir / "tokens.txt" + + logging.info(f"Loading model from {params.model_dir}") + + tokenizer = SimpleTokenizer(token_file) + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + + with open(model_config, "r") as f: + model_config = json.load(f) + + if params.model_name == "zipvoice": + model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + distill = False + else: + assert params.model_name == "zipvoice_distill" + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + distill = True + + if str(model_ckpt).endswith(".safetensors"): + safetensors.torch.load_model(model, model_ckpt) + elif str(model_ckpt).endswith(".pt"): + load_checkpoint(filename=model_ckpt, model=model, strict=True) + else: + raise NotImplementedError(f"Unsupported model checkpoint format: {model_ckpt}") + + device = torch.device("cpu") + model = model.to(device) + model.eval() + + convert_scaled_to_non_scaled(model, inplace=True, is_onnx=True) + + logging.info("Exporting model") + tensorrt_model_dir = Path(params.tensorrt_model_dir) + tensorrt_model_dir.mkdir(parents=True, exist_ok=True) + opset_version = 18 + + + fm_decoder_onnx_file = tensorrt_model_dir / "fm_decoder.onnx" + + export_onnx_fm_decoder( + model=model, + filename=fm_decoder_onnx_file, + opset_version=opset_version, + distill=distill, + ) + + logging.info("Exported to TensorRT model") + + trt_engine_file = f'{str(tensorrt_model_dir)}/{params.trt_engine_file_name}' + trt_kwargs = get_trt_kwargs_dynamic_batch(min_batch_size=1, opt_batch_size=2, max_batch_size=params.max_batch_size) + convert_onnx_to_trt(trt_engine_file, trt_kwargs, fm_decoder_onnx_file, dtype=torch.float16) + + logging.info("Done!") + + +if __name__ == "__main__": + + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + logging.basicConfig(format=formatter, level=logging.INFO, force=True) + import tensorrt as trt + main() diff --git a/zipvoice/bin/train_zipvoice.py b/zipvoice/bin/train_zipvoice.py new file mode 100644 index 0000000..59d3ace --- /dev/null +++ b/zipvoice/bin/train_zipvoice.py @@ -0,0 +1,1130 @@ +#!/usr/bin/env python3 +# Copyright 2024-2025 Xiaomi Corp. (authors: Wei Kang, +# Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script trains a ZipVoice model with the flow-matching loss. + +Usage: + +python3 -m zipvoice.bin.train_zipvoice \ + --world-size 8 \ + --use-fp16 1 \ + --num-epochs 11 \ + --max-duration 500 \ + --lr-hours 30000 \ + --model-config conf/zipvoice_base.json \ + --tokenizer emilia \ + --token-file "data/tokens_emilia.txt" \ + --dataset emilia \ + --manifest-dir data/fbank \ + --exp-dir exp/zipvoice +""" + +import argparse +import copy +import json +import logging +import os +from functools import partial +from pathlib import Path +from shutil import copyfile +from typing import List, Optional, Tuple, Union + +import torch +import torch.multiprocessing as mp +import torch.nn as nn +from lhotse.cut import Cut, CutSet +from lhotse.utils import fix_random_seed +from torch import Tensor +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.utils.tensorboard import SummaryWriter + +import zipvoice.utils.diagnostics as diagnostics +from zipvoice.dataset.datamodule import TtsDataModule +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.tokenizer.tokenizer import ( + EmiliaTokenizer, + EspeakTokenizer, + LibriTTSTokenizer, + SimpleTokenizer, +) +from zipvoice.utils.checkpoint import ( + load_checkpoint, + remove_checkpoints, + resume_checkpoint, + save_checkpoint, + save_checkpoint_with_global_batch_idx, + update_averaged_model, +) +from zipvoice.utils.common import ( + AttributeDict, + GradScaler, + MetricsTracker, + cleanup_dist, + create_grad_scaler, + get_adjusted_batch_count, + get_env_info, + get_parameter_groups_with_lrs, + prepare_input, + set_batch_count, + setup_dist, + setup_logger, + str2bool, + torch_autocast, +) +from zipvoice.utils.hooks import register_inf_check_hooks +from zipvoice.utils.lr_scheduler import Eden, FixedLRScheduler, LRScheduler +from zipvoice.utils.optim import ScaledAdam + +LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, LRScheduler] + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12356, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + parser.add_argument( + "--num-epochs", + type=int, + default=11, + help="Number of epochs to train.", + ) + + parser.add_argument( + "--num-iters", + type=int, + default=0, + help="Number of iter to train, will ignore num_epochs if > 0.", + ) + + parser.add_argument( + "--start-epoch", + type=int, + default=1, + help="""Resume training from this epoch. It should be positive. + If larger than 1, it will load checkpoint from + exp-dir/epoch-{start_epoch-1}.pt + """, + ) + + parser.add_argument( + "--checkpoint", + type=str, + default=None, + help="""Checkpoints of pre-trained models, will load it if not None + """, + ) + + parser.add_argument( + "--exp-dir", + type=str, + default="exp/zipvoice", + help="""The experiment dir. + It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + """, + ) + + parser.add_argument( + "--base-lr", type=float, default=0.02, help="The base learning rate." + ) + + parser.add_argument( + "--lr-batches", + type=float, + default=7500, + help="""Number of steps that affects how rapidly the learning rate + decreases. We suggest not to change this.""", + ) + + parser.add_argument( + "--lr-epochs", + type=float, + default=10, + help="""Number of epochs that affects how rapidly the learning rate decreases. + """, + ) + + parser.add_argument( + "--lr-hours", + type=float, + default=0, + help="""If positive, --epoch is ignored and it specifies the number of hours + that affects how rapidly the learning rate decreases. + """, + ) + + parser.add_argument( + "--ref-duration", + type=float, + default=50, + help="""Reference batch duration for purposes of adjusting batch counts for" + setting various schedules inside the model". + """, + ) + + parser.add_argument( + "--finetune", + type=str2bool, + default=False, + help="Whether to use the fine-tuning mode, will used a fixed learning rate " + "schedule and skip the large dropout phase.", + ) + + parser.add_argument( + "--seed", + type=int, + default=42, + help="The seed for random generators intended for reproducibility", + ) + + parser.add_argument( + "--print-diagnostics", + type=str2bool, + default=False, + help="Accumulate stats on activations, print them and exit.", + ) + + parser.add_argument( + "--scan-oom", + type=str2bool, + default=False, + help="Scan pessimistic batches to see whether they cause OOMs.", + ) + + parser.add_argument( + "--inf-check", + type=str2bool, + default=False, + help="Add hooks to check for infinite module outputs and gradients.", + ) + + parser.add_argument( + "--save-every-n", + type=int, + default=5000, + help="""Save checkpoint after processing this number of batches" + periodically. We save checkpoint to exp-dir/ whenever + params.batch_idx_train % save_every_n == 0. The checkpoint filename + has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt' + Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the + end of each epoch where `xxx` is the epoch number counting from 1. + """, + ) + + parser.add_argument( + "--valid-by-epoch", + type=str2bool, + default=False, + help="""Whether to validate after each epoch. If False, will validate + after every save_every_n iterations. + """, + ) + + parser.add_argument( + "--keep-last-k", + type=int, + default=30, + help="""Only keep this number of checkpoints on disk. + For instance, if it is 3, there are only 3 checkpoints + in the exp-dir with filenames `checkpoint-xxx.pt`. + It does not affect checkpoints with name `epoch-xxx.pt`. + """, + ) + + parser.add_argument( + "--average-period", + type=int, + default=200, + help="""Update the averaged model, namely `model_avg`, after processing + this number of batches. `model_avg` is a separate version of model, + in which each floating-point parameter is the average of all the + parameters from the start of training. Each time we take the average, + we do: `model_avg = model * (average_period / batch_idx_train) + + model_avg * ((batch_idx_train - average_period) / batch_idx_train)`. + """, + ) + + parser.add_argument( + "--use-fp16", + type=str2bool, + default=True, + help="Whether to use half precision training.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--condition-drop-ratio", + type=float, + default=0.2, + help="The drop rate of text condition during training.", + ) + + parser.add_argument( + "--dataset", + type=str, + default="emilia", + choices=["emilia", "libritts", "custom"], + help="The used training dataset", + ) + + parser.add_argument( + "--train-manifest", + type=str, + help="Path of the training manifest", + ) + + parser.add_argument( + "--dev-manifest", + type=str, + help="Path of the validation manifest", + ) + + parser.add_argument( + "--min-len", + type=float, + default=1.0, + help="The minimum audio length used for training", + ) + + parser.add_argument( + "--max-len", + type=float, + default=30.0, + help="The maximum audio length used for training", + ) + + parser.add_argument( + "--model-config", + type=str, + default="conf/zipvoice_base.json", + help="The model configuration file.", + ) + + parser.add_argument( + "--tokenizer", + type=str, + default="emilia", + choices=["emilia", "libritts", "espeak", "simple"], + help="Tokenizer type.", + ) + + parser.add_argument( + "--lang", + type=str, + default="en-us", + help="Language identifier, used when tokenizer type is espeak. see" + "https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md", + ) + + parser.add_argument( + "--token-file", + type=str, + default="data/tokens_emilia.txt", + help="The file that contains information that maps tokens to ids," + "which is a text file with '{token}\t{token_id}' per line.", + ) + + return parser + + +def get_params() -> AttributeDict: + """Return a dict containing training parameters. + + All training related parameters that are not passed from the commandline + are saved in the variable `params`. + + Commandline options are merged into `params` after they are parsed, so + you can also access them via `params`. + + Explanation of options saved in `params`: + + - best_train_loss: Best training loss so far. It is used to select + the model that has the lowest training loss. It is + updated during the training. + + - best_valid_loss: Best validation loss so far. It is used to select + the model that has the lowest validation loss. It is + updated during the training. + + - best_train_epoch: It is the epoch that has the best training loss. + + - best_valid_epoch: It is the epoch that has the best validation loss. + + - batch_idx_train: Used to writing statistics to tensorboard. It + contains number of batches trained so far across + epochs. + + - log_interval: Print training loss if batch_idx % log_interval` is 0 + + - reset_interval: Reset statistics if batch_idx % reset_interval is 0 + + - env_info: A dict containing information about the environment. + + """ + params = AttributeDict( + { + "best_train_loss": float("inf"), + "best_valid_loss": float("inf"), + "best_train_epoch": -1, + "best_valid_epoch": -1, + "batch_idx_train": 0, + "log_interval": 50, + "reset_interval": 200, + "env_info": get_env_info(), + } + ) + + return params + + +def compute_fbank_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + features: Tensor, + features_lens: Tensor, + tokens: List[List[int]], + is_training: bool, +) -> Tuple[Tensor, MetricsTracker]: + """ + Compute loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. + features: + The target acoustic feature. + features_lens: + The number of frames of each utterance. + tokens: + Input tokens that representing the transcripts. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + """ + + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batch_size, num_frames, _ = features.shape + + noise = torch.randn_like(features) # (B, T, F) + + # Sampling t from uniform distribution + if is_training: + t = torch.rand(batch_size, 1, 1, device=device) + else: + t = ( + (torch.arange(batch_size, device=device) / batch_size) + .unsqueeze(1) + .unsqueeze(2) + ) + with torch.set_grad_enabled(is_training): + + loss = model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=noise, + t=t, + condition_drop_ratio=params.condition_drop_ratio, + ) + + assert loss.requires_grad == is_training + info = MetricsTracker() + num_frames = features_lens.sum().item() + info["frames"] = num_frames + info["loss"] = loss.detach().cpu().item() * num_frames + + return loss, info + + +def train_one_epoch( + params: AttributeDict, + model: Union[nn.Module, DDP], + optimizer: Optimizer, + scheduler: LRSchedulerType, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + scaler: GradScaler, + model_avg: Optional[nn.Module] = None, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, + rank: int = 0, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches or every epochs. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + optimizer: + The optimizer. + scheduler: + The learning rate scheduler, we call step() every epoch. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + scaler: + The scaler used for mix precision training. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + rank: + The rank of the node in DDP training. If no DDP is used, it should + be set to 0. + """ + model.train() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to track the stats over iterations in one epoch + tot_loss = MetricsTracker() + + saved_bad_model = False + + def save_bad_model(suffix: str = ""): + save_checkpoint( + filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt", + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=0, + ) + + for batch_idx, batch in enumerate(train_dl): + + if batch_idx % 10 == 0: + if params.finetune: + set_batch_count(model, get_adjusted_batch_count(params) + 100000) + else: + set_batch_count(model, get_adjusted_batch_count(params)) + + if ( + params.valid_by_epoch and batch_idx == 0 and not params.print_diagnostics + ) or ( + not params.valid_by_epoch + and params.batch_idx_train % params.valid_interval == 0 + and not params.print_diagnostics + ): + logging.info("Computing validation loss") + valid_info = compute_validation_loss( + params=params, + model=model, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, global_batch_idx: {params.batch_idx_train}," + f" validation: {valid_info}" + ) + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + if tb_writer is not None: + valid_info.write_summary( + tb_writer, "train/valid_", params.batch_idx_train + ) + + params.batch_idx_train += 1 + + batch_size = len(batch["text"]) + + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + + tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info + + scaler.scale(loss).backward() + + scheduler.step_batch(params.batch_idx_train) + # Use the number of hours of speech to adjust the learning rate + if params.lr_hours > 0: + scheduler.step_epoch( + params.batch_idx_train + * params.max_duration + * params.world_size + / 3600 + ) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + except Exception as e: + logging.info(f"Caught exception : {e}.") + save_bad_model() + raise + + if params.print_diagnostics and batch_idx == 5: + return + + if ( + rank == 0 + and params.batch_idx_train > 0 + and params.batch_idx_train % params.average_period == 0 + ): + update_averaged_model( + params=params, + model_cur=model, + model_avg=model_avg, + ) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.save_every_n == 0 + ): + save_checkpoint_with_global_batch_idx( + out_dir=params.exp_dir, + global_batch_idx=params.batch_idx_train, + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + remove_checkpoints( + out_dir=params.exp_dir, + topk=params.keep_last_k, + rank=rank, + ) + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + if params.batch_idx_train % 100 == 0 and params.use_fp16: + # If the grad scale was less than 1, try increasing it. The _growth_interval + # of the grad scaler is configurable, but we can't configure it to have + # different behavior depending on the current grad scale. + cur_grad_scale = scaler._scale.item() + + if cur_grad_scale < 1024.0 or ( + cur_grad_scale < 4096.0 and params.batch_idx_train % 400 == 0 + ): + scaler.update(cur_grad_scale * 2.0) + if cur_grad_scale < 0.01: + if not saved_bad_model: + save_bad_model(suffix="-first-warning") + saved_bad_model = True + logging.warning(f"Grad scale is small: {cur_grad_scale}") + if cur_grad_scale < 1.0e-05: + save_bad_model() + raise RuntimeError( + f"grad_scale is too small, exiting: {cur_grad_scale}" + ) + + if params.batch_idx_train % params.log_interval == 0: + cur_lr = max(scheduler.get_last_lr()) + cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0 + + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"global_batch_idx: {params.batch_idx_train}, " + f"batch size: {batch_size}, " + f"loss[{loss_info}], tot_loss[{tot_loss}], " + f"cur_lr: {cur_lr:.2e}, " + + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "") + ) + + if tb_writer is not None: + tb_writer.add_scalar( + "train/learning_rate", cur_lr, params.batch_idx_train + ) + loss_info.write_summary( + tb_writer, "train/current_", params.batch_idx_train + ) + tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train) + if params.use_fp16: + tb_writer.add_scalar( + "train/grad_scale", + cur_grad_scale, + params.batch_idx_train, + ) + + loss_value = tot_loss["loss"] + params.train_loss = loss_value + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def compute_validation_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> MetricsTracker: + """Run the validation process.""" + + model.eval() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to summary the stats over iterations + tot_loss = MetricsTracker() + + for batch_idx, batch in enumerate(valid_dl): + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=False, + ) + assert loss.requires_grad is False + tot_loss = tot_loss + loss_info + + if world_size > 1: + tot_loss.reduce(loss.device) + + loss_value = tot_loss["loss"] + if loss_value < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = loss_value + + return tot_loss + + +def display_and_save_batch( + batch: dict, + params: AttributeDict, +) -> None: + """Display the batch statistics and save the batch into disk. + + Args: + batch: + A batch of data. See `lhotse.dataset.K2SpeechRecognitionDataset()` + for the content in it. + params: + Parameters for training. See :func:`get_params`. + sp: + The BPE model. + """ + from lhotse.utils import uuid4 + + filename = f"{params.exp_dir}/batch-{uuid4()}.pt" + logging.info(f"Saving batch to {filename}") + torch.save(batch, filename) + + features = batch["features"] + tokens = batch["tokens"] + + logging.info(f"features shape: {features.shape}") + num_tokens = sum(len(i) for i in tokens) + logging.info(f"num tokens: {num_tokens}") + + +def scan_pessimistic_batches_for_oom( + model: Union[nn.Module, DDP], + train_dl: torch.utils.data.DataLoader, + optimizer: torch.optim.Optimizer, + params: AttributeDict, +): + from lhotse.dataset import find_pessimistic_batches + + logging.info( + "Sanity check -- see if any of the batches in epoch 1 would cause OOM." + ) + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batches, crit_values = find_pessimistic_batches(train_dl.sampler) + for criterion, cuts in batches.items(): + batch = train_dl.dataset[cuts] + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + loss.backward() + optimizer.zero_grad() + except Exception as e: + if "CUDA out of memory" in str(e): + logging.error( + "Your GPU ran out of memory with the current " + "max_duration setting. We recommend decreasing " + "max_duration and trying again.\n" + f"Failing criterion: {criterion} " + f"(={crit_values[criterion]}) ..." + ) + display_and_save_batch(batch, params=params) + raise + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + + +def tokenize_text(c: Cut, tokenizer): + if hasattr(c.supervisions[0], "tokens"): + tokens = tokenizer.tokens_to_token_ids([c.supervisions[0].tokens]) + else: + tokens = tokenizer.texts_to_token_ids([c.supervisions[0].text]) + c.supervisions[0].tokens = tokens[0] + return c + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + params.valid_interval = params.save_every_n + # Set epoch to a large number to ignore it. + if params.num_iters > 0: + params.num_epochs = 1000000 + with open(params.model_config, "r") as f: + model_config = json.load(f) + params.update(model_config["model"]) + params.update(model_config["feature"]) + + fix_random_seed(params.seed) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + os.makedirs(f"{params.exp_dir}", exist_ok=True) + copyfile(src=params.model_config, dst=f"{params.exp_dir}/model.json") + copyfile(src=params.token_file, dst=f"{params.exp_dir}/tokens.txt") + setup_logger(f"{params.exp_dir}/log/log-train") + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + if torch.cuda.is_available(): + params.device = torch.device("cuda", rank) + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + if params.tokenizer == "emilia": + tokenizer = EmiliaTokenizer(token_file=params.token_file) + elif params.tokenizer == "libritts": + tokenizer = LibriTTSTokenizer(token_file=params.token_file) + elif params.tokenizer == "espeak": + tokenizer = EspeakTokenizer(token_file=params.token_file, lang=params.lang) + else: + assert params.tokenizer == "simple" + tokenizer = SimpleTokenizer(token_file=params.token_file) + + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + params.update(tokenizer_config) + + logging.info(params) + + logging.info("About to create model") + + model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + + if params.checkpoint is not None: + logging.info(f"Loading pre-trained model from {params.checkpoint}") + _ = load_checkpoint(filename=params.checkpoint, model=model, strict=True) + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of parameters : {num_param}") + + model_avg: Optional[nn.Module] = None + if rank == 0: + # model_avg is only used with rank 0 + model_avg = copy.deepcopy(model).to(torch.float64) + + assert params.start_epoch > 0, params.start_epoch + if params.start_epoch > 1: + checkpoints = resume_checkpoint(params=params, model=model, model_avg=model_avg) + + model = model.to(params.device) + if world_size > 1: + logging.info("Using DDP") + model = DDP(model, device_ids=[rank], find_unused_parameters=True) + + optimizer = ScaledAdam( + get_parameter_groups_with_lrs( + model, + lr=params.base_lr, + include_names=True, + ), + lr=params.base_lr, # should have no effect + clipping_scale=2.0, + ) + + assert params.lr_hours >= 0 + + if params.finetune: + scheduler = FixedLRScheduler(optimizer) + elif params.lr_hours > 0: + scheduler = Eden(optimizer, params.lr_batches, params.lr_hours) + else: + scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs) + + scaler = create_grad_scaler(enabled=params.use_fp16) + + if params.start_epoch > 1 and checkpoints is not None: + # load state_dict for optimizers + if "optimizer" in checkpoints: + logging.info("Loading optimizer state dict") + optimizer.load_state_dict(checkpoints["optimizer"]) + + # load state_dict for schedulers + if "scheduler" in checkpoints: + logging.info("Loading scheduler state dict") + scheduler.load_state_dict(checkpoints["scheduler"]) + + if "grad_scaler" in checkpoints: + logging.info("Loading grad scaler state dict") + scaler.load_state_dict(checkpoints["grad_scaler"]) + + if params.print_diagnostics: + opts = diagnostics.TensorDiagnosticOptions( + 512 + ) # allow 4 megabytes per sub-module + diagnostic = diagnostics.attach_diagnostics(model, opts) + + if params.inf_check: + register_inf_check_hooks(model) + + def remove_short_and_long_utt(c: Cut, min_len: float, max_len: float): + if c.duration < min_len or c.duration > max_len: + return False + return True + + _remove_short_and_long_utt = partial( + remove_short_and_long_utt, min_len=params.min_len, max_len=params.max_len + ) + + datamodule = TtsDataModule(args) + if params.dataset == "emilia": + train_cuts = CutSet.mux( + datamodule.train_emilia_EN_cuts(), + datamodule.train_emilia_ZH_cuts(), + weights=[46000, 49000], + ) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = CutSet.mux( + datamodule.dev_emilia_EN_cuts(), + datamodule.dev_emilia_ZH_cuts(), + weights=[0.5, 0.5], + ) + elif params.dataset == "libritts": + train_cuts = datamodule.train_libritts_cuts() + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_libritts_cuts() + else: + assert params.dataset == "custom" + train_cuts = datamodule.train_custom_cuts(params.train_manifest) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_custom_cuts(params.dev_manifest) + # To avoid OOM issues due to too long dev cuts + dev_cuts = dev_cuts.filter(_remove_short_and_long_utt) + + if params.tokenizer in ["emilia", "espeak", "dialog"]: + if not hasattr(train_cuts[0].supervisions[0], "tokens") or not hasattr( + dev_cuts[0].supervisions[0], "tokens" + ): + logging.warning( + f"Using {params.tokenizer} tokenizer but tokens are not prepared," + f"will tokenize on-the-fly, which can slow down training significantly." + ) + _tokenize_text = partial(tokenize_text, tokenizer=tokenizer) + train_cuts = train_cuts.map(_tokenize_text) + dev_cuts = dev_cuts.map(_tokenize_text) + + train_dl = datamodule.train_dataloaders(train_cuts) + + valid_dl = datamodule.dev_dataloaders(dev_cuts) + + if params.scan_oom: + scan_pessimistic_batches_for_oom( + model=model, + train_dl=train_dl, + optimizer=optimizer, + params=params, + ) + + logging.info("Training started") + + for epoch in range(params.start_epoch, params.num_epochs + 1): + logging.info(f"Start epoch {epoch}") + + if params.lr_hours == 0: + scheduler.step_epoch(epoch - 1) + fix_random_seed(params.seed + epoch - 1) + train_dl.sampler.set_epoch(epoch - 1) + + params.cur_epoch = epoch + + if tb_writer is not None: + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + train_one_epoch( + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + train_dl=train_dl, + valid_dl=valid_dl, + scaler=scaler, + tb_writer=tb_writer, + world_size=world_size, + rank=rank, + ) + + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + + if params.print_diagnostics: + diagnostic.print_diagnostics() + break + + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint( + filename=filename, + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + + if rank == 0: + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + logging.info("Done!") + + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + TtsDataModule.add_arguments(parser) + args = parser.parse_args() + args.exp_dir = Path(args.exp_dir) + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + main() diff --git a/zipvoice/bin/train_zipvoice_dialog.py b/zipvoice/bin/train_zipvoice_dialog.py new file mode 100644 index 0000000..c401220 --- /dev/null +++ b/zipvoice/bin/train_zipvoice_dialog.py @@ -0,0 +1,980 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script trains a ZipVoice-Dialog model. + +Usage: + +python3 -m zipvoice.bin.train_zipvoice_dialog \ + --world-size 8 \ + --use-fp16 1 \ + --base-lr 0.0001 \ + --max-duration 500 \ + --checkpoint download/zipvoice/model.pt \ + --model-config conf/zipvoice_base.json \ + --token-file "data/tokens_dialog.txt" \ + --dataset opendialog \ + --manifest-dir data/fbank \ + --exp-dir exp/zipvoice_dialog +""" + +import argparse +import copy +import json +import logging +import os +from functools import partial +from pathlib import Path +from shutil import copyfile +from typing import List, Optional, Tuple, Union + +import torch +import torch.multiprocessing as mp +import torch.nn as nn +from lhotse.cut import Cut, CutSet +from lhotse.utils import fix_random_seed +from torch import Tensor +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.utils.tensorboard import SummaryWriter + +import zipvoice.utils.diagnostics as diagnostics +from zipvoice.bin.train_zipvoice import ( + display_and_save_batch, + get_params, + tokenize_text, +) +from zipvoice.dataset.datamodule import TtsDataModule +from zipvoice.models.zipvoice_dialog import ZipVoiceDialog +from zipvoice.tokenizer.tokenizer import DialogTokenizer +from zipvoice.utils.checkpoint import ( + load_checkpoint, + load_checkpoint_extend_vocab_size, + remove_checkpoints, + resume_checkpoint, + save_checkpoint, + save_checkpoint_with_global_batch_idx, + update_averaged_model, +) +from zipvoice.utils.common import ( + AttributeDict, + GradScaler, + MetricsTracker, + cleanup_dist, + create_grad_scaler, + get_adjusted_batch_count, + get_parameter_groups_with_lrs, + prepare_input, + set_batch_count, + setup_dist, + setup_logger, + str2bool, + torch_autocast, +) +from zipvoice.utils.hooks import register_inf_check_hooks +from zipvoice.utils.lr_scheduler import FixedLRScheduler, LRScheduler +from zipvoice.utils.optim import ScaledAdam + +LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, LRScheduler] + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12356, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + parser.add_argument( + "--num-epochs", + type=int, + default=8, + help="Number of epochs to train.", + ) + + parser.add_argument( + "--num-iters", + type=int, + default=60000, + help="Number of iter to train, will ignore num_epochs if > 0.", + ) + + parser.add_argument( + "--start-epoch", + type=int, + default=1, + help="""Resume training from this epoch. It should be positive. + If larger than 1, it will load checkpoint from + exp-dir/epoch-{start_epoch-1}.pt + """, + ) + + parser.add_argument( + "--checkpoint", + type=str, + required=True, + help="""Checkpoints of pre-trained models, either a ZipVoice model or a + ZipVoice-Dialog model. + """, + ) + + parser.add_argument( + "--exp-dir", + type=str, + default="exp/zipvoice_dialog", + help="""The experiment dir. + It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + """, + ) + + parser.add_argument( + "--base-lr", type=float, default=0.0001, help="The base learning rate." + ) + + parser.add_argument( + "--ref-duration", + type=float, + default=50, + help="""Reference batch duration for purposes of adjusting batch counts for" + setting various schedules inside the model". + """, + ) + + parser.add_argument( + "--finetune", + type=str2bool, + default=False, + help="Whether to fine-tune from our pre-traied ZipVoice-Dialog model." + "False means to fine-tune from a pre-trained ZipVoice model.", + ) + + parser.add_argument( + "--seed", + type=int, + default=42, + help="The seed for random generators intended for reproducibility", + ) + + parser.add_argument( + "--print-diagnostics", + type=str2bool, + default=False, + help="Accumulate stats on activations, print them and exit.", + ) + + parser.add_argument( + "--scan-oom", + type=str2bool, + default=False, + help="Scan pessimistic batches to see whether they cause OOMs.", + ) + + parser.add_argument( + "--inf-check", + type=str2bool, + default=False, + help="Add hooks to check for infinite module outputs and gradients.", + ) + + parser.add_argument( + "--save-every-n", + type=int, + default=5000, + help="""Save checkpoint after processing this number of batches" + periodically. We save checkpoint to exp-dir/ whenever + params.batch_idx_train % save_every_n == 0. The checkpoint filename + has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt' + Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the + end of each epoch where `xxx` is the epoch number counting from 1. + """, + ) + + parser.add_argument( + "--keep-last-k", + type=int, + default=30, + help="""Only keep this number of checkpoints on disk. + For instance, if it is 3, there are only 3 checkpoints + in the exp-dir with filenames `checkpoint-xxx.pt`. + It does not affect checkpoints with name `epoch-xxx.pt`. + """, + ) + + parser.add_argument( + "--average-period", + type=int, + default=200, + help="""Update the averaged model, namely `model_avg`, after processing + this number of batches. `model_avg` is a separate version of model, + in which each floating-point parameter is the average of all the + parameters from the start of training. Each time we take the average, + we do: `model_avg = model * (average_period / batch_idx_train) + + model_avg * ((batch_idx_train - average_period) / batch_idx_train)`. + """, + ) + + parser.add_argument( + "--use-fp16", + type=str2bool, + default=True, + help="Whether to use half precision training.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--condition-drop-ratio", + type=float, + default=0.2, + help="The drop rate of text condition during training.", + ) + + parser.add_argument( + "--dataset", + type=str, + default="opendialog", + choices=["opendialog", "custom"], + help="The used training dataset", + ) + + parser.add_argument( + "--train-manifest", + type=str, + help="Path of the training manifest", + ) + + parser.add_argument( + "--dev-manifest", + type=str, + help="Path of the validation manifest", + ) + + parser.add_argument( + "--min-len", + type=float, + default=1.0, + help="The minimum audio length used for training", + ) + + parser.add_argument( + "--max-len", + type=float, + default=30.0, + help="The maximum audio length used for training", + ) + + parser.add_argument( + "--model-config", + type=str, + default="zipvoice_base.json", + help="The model configuration file.", + ) + + parser.add_argument( + "--token-file", + type=str, + default="data/tokens_dialog.txt", + help="The file that contains information that maps tokens to ids," + "which is a text file with '{token}\t{token_id}' per line.", + ) + + return parser + + +def compute_fbank_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + features: Tensor, + features_lens: Tensor, + tokens: List[List[int]], + is_training: bool, +) -> Tuple[Tensor, MetricsTracker]: + """ + Compute loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. + features: + The target acoustic feature. + features_lens: + The number of frames of each utterance. + tokens: + Input tokens that representing the transcripts. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + """ + + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batch_size, num_frames, _ = features.shape + + noise = torch.randn_like(features) # (B, T, F) + + # Sampling t from uniform distribution + if is_training: + t = torch.rand(batch_size, 1, 1, device=device) + else: + t = ( + (torch.arange(batch_size, device=device) / batch_size) + .unsqueeze(1) + .unsqueeze(2) + ) + with torch.set_grad_enabled(is_training): + + loss = model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=noise, + t=t, + condition_drop_ratio=params.condition_drop_ratio, + ) + + assert loss.requires_grad == is_training + info = MetricsTracker() + num_frames = features_lens.sum().item() + info["frames"] = num_frames + info["loss"] = loss.detach().cpu().item() * num_frames + + return loss, info + + +def train_one_epoch( + params: AttributeDict, + model: Union[nn.Module, DDP], + optimizer: Optimizer, + scheduler: LRSchedulerType, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + scaler: GradScaler, + model_avg: Optional[nn.Module] = None, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, + rank: int = 0, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + optimizer: + The optimizer. + scheduler: + The learning rate scheduler, we call step() every epoch. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + scaler: + The scaler used for mix precision training. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + rank: + The rank of the node in DDP training. If no DDP is used, it should + be set to 0. + """ + model.train() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to track the stats over iterations in one epoch + tot_loss = MetricsTracker() + + saved_bad_model = False + + def save_bad_model(suffix: str = ""): + save_checkpoint( + filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt", + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=0, + ) + + for batch_idx, batch in enumerate(train_dl): + + if batch_idx % 10 == 0: + set_batch_count(model, get_adjusted_batch_count(params) + 100000) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.valid_interval == 0 + and not params.print_diagnostics + ): + logging.info("Computing validation loss") + valid_info = compute_validation_loss( + params=params, + model=model, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, global_batch_idx: {params.batch_idx_train}," + f" validation: {valid_info}" + ) + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + if tb_writer is not None: + valid_info.write_summary( + tb_writer, "train/valid_", params.batch_idx_train + ) + + params.batch_idx_train += 1 + + batch_size = len(batch["text"]) + + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + + tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info + + scaler.scale(loss).backward() + + scheduler.step_batch(params.batch_idx_train) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + except Exception as e: + logging.info(f"Caught exception : {e}.") + save_bad_model() + raise + + if params.print_diagnostics and batch_idx == 5: + return + + if ( + rank == 0 + and params.batch_idx_train > 0 + and params.batch_idx_train % params.average_period == 0 + ): + update_averaged_model( + params=params, + model_cur=model, + model_avg=model_avg, + ) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.save_every_n == 0 + ): + save_checkpoint_with_global_batch_idx( + out_dir=params.exp_dir, + global_batch_idx=params.batch_idx_train, + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + remove_checkpoints( + out_dir=params.exp_dir, + topk=params.keep_last_k, + rank=rank, + ) + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + if params.batch_idx_train % 100 == 0 and params.use_fp16: + # If the grad scale was less than 1, try increasing it. The _growth_interval + # of the grad scaler is configurable, but we can't configure it to have + # different behavior depending on the current grad scale. + cur_grad_scale = scaler._scale.item() + + if cur_grad_scale < 1024.0 or ( + cur_grad_scale < 4096.0 and params.batch_idx_train % 400 == 0 + ): + scaler.update(cur_grad_scale * 2.0) + if cur_grad_scale < 0.01: + if not saved_bad_model: + save_bad_model(suffix="-first-warning") + saved_bad_model = True + logging.warning(f"Grad scale is small: {cur_grad_scale}") + if cur_grad_scale < 1.0e-05: + save_bad_model() + raise RuntimeError( + f"grad_scale is too small, exiting: {cur_grad_scale}" + ) + + if params.batch_idx_train % params.log_interval == 0: + cur_lr = max(scheduler.get_last_lr()) + cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0 + + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"global_batch_idx: {params.batch_idx_train}, " + f"batch size: {batch_size}, " + f"loss[{loss_info}], tot_loss[{tot_loss}], " + f"cur_lr: {cur_lr:.2e}, " + + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "") + ) + + if tb_writer is not None: + tb_writer.add_scalar( + "train/learning_rate", cur_lr, params.batch_idx_train + ) + loss_info.write_summary( + tb_writer, "train/current_", params.batch_idx_train + ) + tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train) + if params.use_fp16: + tb_writer.add_scalar( + "train/grad_scale", + cur_grad_scale, + params.batch_idx_train, + ) + + loss_value = tot_loss["loss"] + params.train_loss = loss_value + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def compute_validation_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> MetricsTracker: + """Run the validation process.""" + + model.eval() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to summary the stats over iterations + tot_loss = MetricsTracker() + + for batch_idx, batch in enumerate(valid_dl): + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=False, + ) + assert loss.requires_grad is False + tot_loss = tot_loss + loss_info + + if world_size > 1: + tot_loss.reduce(loss.device) + + loss_value = tot_loss["loss"] + if loss_value < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = loss_value + + return tot_loss + + +def scan_pessimistic_batches_for_oom( + model: Union[nn.Module, DDP], + train_dl: torch.utils.data.DataLoader, + optimizer: torch.optim.Optimizer, + params: AttributeDict, +): + from lhotse.dataset import find_pessimistic_batches + + logging.info( + "Sanity check -- see if any of the batches in epoch 1 would cause OOM." + ) + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batches, crit_values = find_pessimistic_batches(train_dl.sampler) + for criterion, cuts in batches.items(): + batch = train_dl.dataset[cuts] + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + loss.backward() + optimizer.zero_grad() + except Exception as e: + if "CUDA out of memory" in str(e): + logging.error( + "Your GPU ran out of memory with the current " + "max_duration setting. We recommend decreasing " + "max_duration and trying again.\n" + f"Failing criterion: {criterion} " + f"(={crit_values[criterion]}) ..." + ) + display_and_save_batch(batch, params=params) + raise + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + params.valid_interval = params.save_every_n + # Set epoch to a large number to ignore it. + if params.num_iters > 0: + params.num_epochs = 1000000 + with open(params.model_config, "r") as f: + model_config = json.load(f) + params.update(model_config["model"]) + params.update(model_config["feature"]) + + fix_random_seed(params.seed) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + os.makedirs(f"{params.exp_dir}", exist_ok=True) + copyfile(src=params.model_config, dst=f"{params.exp_dir}/model.json") + copyfile(src=params.token_file, dst=f"{params.exp_dir}/tokens.txt") + setup_logger(f"{params.exp_dir}/log/log-train") + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + if torch.cuda.is_available(): + params.device = torch.device("cuda", rank) + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + tokenizer = DialogTokenizer(token_file=params.token_file) + tokenizer_config = { + "vocab_size": tokenizer.vocab_size, + "pad_id": tokenizer.pad_id, + "spk_a_id": tokenizer.spk_a_id, + "spk_b_id": tokenizer.spk_b_id, + } + params.update(tokenizer_config) + + logging.info(params) + + logging.info("About to create model") + + model = ZipVoiceDialog( + **model_config["model"], + **tokenizer_config, + ) + + assert params.checkpoint is not None, ( + "require a pre-trained checkpoint, as training from random initialization " + "leads to uninteligible dialogue speech" + ) + logging.info(f"Loading pre-trained model from {params.checkpoint}") + + if params.finetune: + # load a pre-trained ZipVoice-Dialog model + _ = load_checkpoint(filename=params.checkpoint, model=model, strict=True) + else: + # load a pre-trained ZipVoice model, extend the vocab size for additional tokens + _ = load_checkpoint_extend_vocab_size( + filename=params.checkpoint, + extend_size=28, + model=model, + strict=True, + ) + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of parameters : {num_param}") + + model_avg: Optional[nn.Module] = None + if rank == 0: + # model_avg is only used with rank 0 + model_avg = copy.deepcopy(model).to(torch.float64) + + assert params.start_epoch > 0, params.start_epoch + if params.start_epoch > 1: + checkpoints = resume_checkpoint(params=params, model=model, model_avg=model_avg) + + model = model.to(params.device) + if world_size > 1: + logging.info("Using DDP") + model = DDP(model, device_ids=[rank], find_unused_parameters=True) + + optimizer = ScaledAdam( + get_parameter_groups_with_lrs( + model, + lr=params.base_lr, + include_names=True, + ), + lr=params.base_lr, # should have no effect + clipping_scale=2.0, + ) + + scheduler = FixedLRScheduler(optimizer) + + scaler = create_grad_scaler(enabled=params.use_fp16) + + if params.start_epoch > 1 and checkpoints is not None: + # load state_dict for optimizers + if "optimizer" in checkpoints: + logging.info("Loading optimizer state dict") + optimizer.load_state_dict(checkpoints["optimizer"]) + + # load state_dict for schedulers + if "scheduler" in checkpoints: + logging.info("Loading scheduler state dict") + scheduler.load_state_dict(checkpoints["scheduler"]) + + if "grad_scaler" in checkpoints: + logging.info("Loading grad scaler state dict") + scaler.load_state_dict(checkpoints["grad_scaler"]) + + if params.print_diagnostics: + opts = diagnostics.TensorDiagnosticOptions( + 512 + ) # allow 4 megabytes per sub-module + diagnostic = diagnostics.attach_diagnostics(model, opts) + + if params.inf_check: + register_inf_check_hooks(model) + + def remove_short_and_long_utt(c: Cut, min_len: float, max_len: float): + if c.duration < min_len or c.duration > max_len: + return False + return True + + _remove_short_and_long_utt = partial( + remove_short_and_long_utt, min_len=params.min_len, max_len=params.max_len + ) + + datamodule = TtsDataModule(args) + if params.dataset == "opendialog": + train_opendialog_en_cuts = datamodule.train_opendialog_en_cuts() + train_opendialog_zh_cuts = datamodule.train_opendialog_zh_cuts().repeat(2) + + train_cuts = CutSet.mux( + train_opendialog_en_cuts, + train_opendialog_zh_cuts, + weights=[ + len(train_opendialog_en_cuts), + len(train_opendialog_zh_cuts), + ], + ) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + + dev_cuts = CutSet.mux( + datamodule.dev_opendialog_en_cuts(), + datamodule.dev_opendialog_zh_cuts(), + ) + else: + assert params.dataset == "custom" + train_cuts = datamodule.train_custom_cuts(params.train_manifest) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_custom_cuts(params.dev_manifest) + # To avoid OOM issues due to too long dev cuts + dev_cuts = dev_cuts.filter(_remove_short_and_long_utt) + + if not hasattr(train_cuts[0].supervisions[0], "tokens") or not hasattr( + dev_cuts[0].supervisions[0], "tokens" + ): + logging.warning( + "Tokens are not prepared, will tokenize on-the-fly, " + "which can slow down training significantly." + ) + _tokenize_text = partial(tokenize_text, tokenizer=tokenizer) + train_cuts = train_cuts.map(_tokenize_text) + dev_cuts = dev_cuts.map(_tokenize_text) + + train_dl = datamodule.train_dataloaders(train_cuts) + + valid_dl = datamodule.dev_dataloaders(dev_cuts) + + if params.scan_oom: + scan_pessimistic_batches_for_oom( + model=model, + train_dl=train_dl, + optimizer=optimizer, + params=params, + ) + + logging.info("Training started") + + for epoch in range(params.start_epoch, params.num_epochs + 1): + logging.info(f"Start epoch {epoch}") + scheduler.step_epoch(epoch - 1) + fix_random_seed(params.seed + epoch - 1) + train_dl.sampler.set_epoch(epoch - 1) + + params.cur_epoch = epoch + + if tb_writer is not None: + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + train_one_epoch( + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + train_dl=train_dl, + valid_dl=valid_dl, + scaler=scaler, + tb_writer=tb_writer, + world_size=world_size, + rank=rank, + ) + + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + + if params.print_diagnostics: + diagnostic.print_diagnostics() + break + + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint( + filename=filename, + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + + if rank == 0: + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + logging.info("Done!") + + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + TtsDataModule.add_arguments(parser) + args = parser.parse_args() + args.exp_dir = Path(args.exp_dir) + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + main() diff --git a/zipvoice/bin/train_zipvoice_dialog_stereo.py b/zipvoice/bin/train_zipvoice_dialog_stereo.py new file mode 100644 index 0000000..12ce915 --- /dev/null +++ b/zipvoice/bin/train_zipvoice_dialog_stereo.py @@ -0,0 +1,963 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script trains a ZipVoice-Dialog model. + +Usage: + +python3 -m zipvoice.bin.train_zipvoice_dialog_stereo \ + --world-size 8 \ + --use-fp16 1 \ + --base-lr 0.002 \ + --max-duration 500 \ + --model-config conf/zipvoice_base.json \ + --token-file "data/tokens_dialog.txt" \ + --manifest-dir data/fbank \ + --exp-dir exp/zipvoice_dialog_stereo +""" + +import argparse +import copy +import json +import logging +import os +from functools import partial +from pathlib import Path +from shutil import copyfile +from typing import List, Optional, Tuple, Union + +import torch +import torch.multiprocessing as mp +import torch.nn as nn +from lhotse.cut import Cut +from lhotse.utils import fix_random_seed +from torch import Tensor +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.utils.tensorboard import SummaryWriter + +import zipvoice.utils.diagnostics as diagnostics +from zipvoice.bin.train_zipvoice import ( + display_and_save_batch, + get_params, + tokenize_text, +) +from zipvoice.dataset.datamodule import TtsDataModule +from zipvoice.models.zipvoice_dialog import ZipVoiceDialogStereo +from zipvoice.tokenizer.tokenizer import DialogTokenizer +from zipvoice.utils.checkpoint import ( + load_checkpoint, + load_checkpoint_copy_proj_three_channel_alter, + remove_checkpoints, + resume_checkpoint, + save_checkpoint, + save_checkpoint_with_global_batch_idx, + update_averaged_model, +) +from zipvoice.utils.common import ( + AttributeDict, + GradScaler, + MetricsTracker, + cleanup_dist, + create_grad_scaler, + get_adjusted_batch_count, + get_parameter_groups_with_lrs, + prepare_input, + set_batch_count, + setup_dist, + setup_logger, + str2bool, + torch_autocast, +) +from zipvoice.utils.hooks import register_inf_check_hooks +from zipvoice.utils.lr_scheduler import FixedLRScheduler, LRScheduler +from zipvoice.utils.optim import ScaledAdam + +LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, LRScheduler] + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12356, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + parser.add_argument( + "--num-epochs", + type=int, + default=8, + help="Number of epochs to train.", + ) + + parser.add_argument( + "--num-iters", + type=int, + default=25000, + help="Number of iter to train, will ignore num_epochs if > 0.", + ) + + parser.add_argument( + "--start-epoch", + type=int, + default=1, + help="""Resume training from this epoch. It should be positive. + If larger than 1, it will load checkpoint from + exp-dir/epoch-{start_epoch-1}.pt + """, + ) + + parser.add_argument( + "--checkpoint", + type=str, + required=True, + help="""Checkpoints of pre-trained models, either a ZipVoice model or a + ZipVoice-Dialog model. + """, + ) + + parser.add_argument( + "--exp-dir", + type=str, + default="exp/zipvoice_dialog", + help="""The experiment dir. + It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + """, + ) + + parser.add_argument( + "--base-lr", type=float, default=0.002, help="The base learning rate." + ) + + parser.add_argument( + "--ref-duration", + type=float, + default=50, + help="""Reference batch duration for purposes of adjusting batch counts for" + setting various schedules inside the model". + """, + ) + + parser.add_argument( + "--finetune", + type=str2bool, + default=False, + help="Whether to fine-tune from our pre-traied ZipVoice-Dialog model." + "False means to fine-tune from a pre-trained ZipVoice model.", + ) + + parser.add_argument( + "--seed", + type=int, + default=42, + help="The seed for random generators intended for reproducibility", + ) + + parser.add_argument( + "--print-diagnostics", + type=str2bool, + default=False, + help="Accumulate stats on activations, print them and exit.", + ) + + parser.add_argument( + "--scan-oom", + type=str2bool, + default=False, + help="Scan pessimistic batches to see whether they cause OOMs.", + ) + + parser.add_argument( + "--inf-check", + type=str2bool, + default=False, + help="Add hooks to check for infinite module outputs and gradients.", + ) + + parser.add_argument( + "--save-every-n", + type=int, + default=5000, + help="""Save checkpoint after processing this number of batches" + periodically. We save checkpoint to exp-dir/ whenever + params.batch_idx_train % save_every_n == 0. The checkpoint filename + has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt' + Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the + end of each epoch where `xxx` is the epoch number counting from 1. + """, + ) + + parser.add_argument( + "--keep-last-k", + type=int, + default=30, + help="""Only keep this number of checkpoints on disk. + For instance, if it is 3, there are only 3 checkpoints + in the exp-dir with filenames `checkpoint-xxx.pt`. + It does not affect checkpoints with name `epoch-xxx.pt`. + """, + ) + + parser.add_argument( + "--average-period", + type=int, + default=200, + help="""Update the averaged model, namely `model_avg`, after processing + this number of batches. `model_avg` is a separate version of model, + in which each floating-point parameter is the average of all the + parameters from the start of training. Each time we take the average, + we do: `model_avg = model * (average_period / batch_idx_train) + + model_avg * ((batch_idx_train - average_period) / batch_idx_train)`. + """, + ) + + parser.add_argument( + "--use-fp16", + type=str2bool, + default=True, + help="Whether to use half precision training.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--condition-drop-ratio", + type=float, + default=0.2, + help="The drop rate of text condition during training.", + ) + + parser.add_argument( + "--train-manifest", + type=str, + help="Path of the training manifest", + ) + + parser.add_argument( + "--dev-manifest", + type=str, + help="Path of the validation manifest", + ) + + parser.add_argument( + "--min-len", + type=float, + default=1.0, + help="The minimum audio length used for training", + ) + + parser.add_argument( + "--max-len", + type=float, + default=60.0, + help="The maximum audio length used for training", + ) + + parser.add_argument( + "--model-config", + type=str, + default="zipvoice_base.json", + help="The model configuration file.", + ) + + parser.add_argument( + "--token-file", + type=str, + default="data/tokens_dialog.txt", + help="The file that contains information that maps tokens to ids," + "which is a text file with '{token}\t{token_id}' per line.", + ) + + return parser + + +def compute_fbank_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + features: Tensor, + features_lens: Tensor, + tokens: List[List[int]], + is_training: bool, + use_two_channel: bool, +) -> Tuple[Tensor, MetricsTracker]: + """ + Compute loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. + features: + The target acoustic feature. + features_lens: + The number of frames of each utterance. + tokens: + Input tokens that representing the transcripts. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + use_two_channel: + True for using two channel features, False for using one channel features. + """ + + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batch_size, num_frames, _ = features.shape + + assert ( + features.size(2) == 3 * params.feat_dim + ), "we assume three channel features, the last channel is the mixed-channel feature" + if use_two_channel: + features = features[:, :, : params.feat_dim * 2] + else: + features = features[:, :, params.feat_dim * 2 :] + + noise = torch.randn_like(features) # (B, T, F) + + # Sampling t from uniform distribution + if is_training: + t = torch.rand(batch_size, 1, 1, device=device) + else: + t = ( + (torch.arange(batch_size, device=device) / batch_size) + .unsqueeze(1) + .unsqueeze(2) + ) + with torch.set_grad_enabled(is_training): + + loss = model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=noise, + t=t, + condition_drop_ratio=params.condition_drop_ratio, + se_weight=1 if use_two_channel else 0, + ) + + assert loss.requires_grad == is_training + info = MetricsTracker() + num_frames = features_lens.sum().item() + info["frames"] = num_frames + info["loss"] = loss.detach().cpu().item() * num_frames + + return loss, info + + +def train_one_epoch( + params: AttributeDict, + model: Union[nn.Module, DDP], + optimizer: Optimizer, + scheduler: LRSchedulerType, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + scaler: GradScaler, + model_avg: Optional[nn.Module] = None, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, + rank: int = 0, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + optimizer: + The optimizer. + scheduler: + The learning rate scheduler, we call step() every epoch. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + scaler: + The scaler used for mix precision training. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + rank: + The rank of the node in DDP training. If no DDP is used, it should + be set to 0. + """ + model.train() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to track the stats over iterations in one epoch + tot_loss = MetricsTracker() + + saved_bad_model = False + + def save_bad_model(suffix: str = ""): + save_checkpoint( + filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt", + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=0, + ) + + for batch_idx, batch in enumerate(train_dl): + + if batch_idx % 10 == 0: + set_batch_count(model, get_adjusted_batch_count(params) + 100000) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.valid_interval == 0 + and not params.print_diagnostics + ): + logging.info("Computing validation loss") + valid_info = compute_validation_loss( + params=params, + model=model, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, global_batch_idx: {params.batch_idx_train}," + f" validation: {valid_info}" + ) + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + if tb_writer is not None: + valid_info.write_summary( + tb_writer, "train/valid_", params.batch_idx_train + ) + + params.batch_idx_train += 1 + + batch_size = len(batch["text"]) + + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + use_two_channel=(batch_idx % 2 == 1), + ) + + tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info + + scaler.scale(loss).backward() + + scheduler.step_batch(params.batch_idx_train) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + except Exception as e: + logging.info(f"Caught exception : {e}.") + save_bad_model() + raise + + if params.print_diagnostics and batch_idx == 5: + return + + if ( + rank == 0 + and params.batch_idx_train > 0 + and params.batch_idx_train % params.average_period == 0 + ): + update_averaged_model( + params=params, + model_cur=model, + model_avg=model_avg, + ) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.save_every_n == 0 + ): + save_checkpoint_with_global_batch_idx( + out_dir=params.exp_dir, + global_batch_idx=params.batch_idx_train, + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + remove_checkpoints( + out_dir=params.exp_dir, + topk=params.keep_last_k, + rank=rank, + ) + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + if params.batch_idx_train % 100 == 0 and params.use_fp16: + # If the grad scale was less than 1, try increasing it. The _growth_interval + # of the grad scaler is configurable, but we can't configure it to have + # different behavior depending on the current grad scale. + cur_grad_scale = scaler._scale.item() + + if cur_grad_scale < 1024.0 or ( + cur_grad_scale < 4096.0 and params.batch_idx_train % 400 == 0 + ): + scaler.update(cur_grad_scale * 2.0) + if cur_grad_scale < 0.01: + if not saved_bad_model: + save_bad_model(suffix="-first-warning") + saved_bad_model = True + logging.warning(f"Grad scale is small: {cur_grad_scale}") + if cur_grad_scale < 1.0e-05: + save_bad_model() + raise RuntimeError( + f"grad_scale is too small, exiting: {cur_grad_scale}" + ) + + if params.batch_idx_train % params.log_interval == 0: + cur_lr = max(scheduler.get_last_lr()) + cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0 + + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"global_batch_idx: {params.batch_idx_train}, " + f"batch size: {batch_size}, " + f"loss[{loss_info}], tot_loss[{tot_loss}], " + f"cur_lr: {cur_lr:.2e}, " + + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "") + ) + + if tb_writer is not None: + tb_writer.add_scalar( + "train/learning_rate", cur_lr, params.batch_idx_train + ) + loss_info.write_summary( + tb_writer, "train/current_", params.batch_idx_train + ) + tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train) + if params.use_fp16: + tb_writer.add_scalar( + "train/grad_scale", + cur_grad_scale, + params.batch_idx_train, + ) + + loss_value = tot_loss["loss"] + params.train_loss = loss_value + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def compute_validation_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> MetricsTracker: + """Run the validation process.""" + + model.eval() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to summary the stats over iterations + tot_loss = MetricsTracker() + + for batch_idx, batch in enumerate(valid_dl): + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=False, + use_two_channel=True, + ) + assert loss.requires_grad is False + tot_loss = tot_loss + loss_info + + if world_size > 1: + tot_loss.reduce(loss.device) + + loss_value = tot_loss["loss"] + if loss_value < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = loss_value + + return tot_loss + + +def scan_pessimistic_batches_for_oom( + model: Union[nn.Module, DDP], + train_dl: torch.utils.data.DataLoader, + optimizer: torch.optim.Optimizer, + params: AttributeDict, +): + from lhotse.dataset import find_pessimistic_batches + + logging.info( + "Sanity check -- see if any of the batches in epoch 1 would cause OOM." + ) + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batches, crit_values = find_pessimistic_batches(train_dl.sampler) + for criterion, cuts in batches.items(): + batch = train_dl.dataset[cuts] + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + use_two_channel=True, + ) + loss.backward() + optimizer.zero_grad() + except Exception as e: + if "CUDA out of memory" in str(e): + logging.error( + "Your GPU ran out of memory with the current " + "max_duration setting. We recommend decreasing " + "max_duration and trying again.\n" + f"Failing criterion: {criterion} " + f"(={crit_values[criterion]}) ..." + ) + display_and_save_batch(batch, params=params) + raise + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + params.valid_interval = params.save_every_n + # Set epoch to a large number to ignore it. + if params.num_iters > 0: + params.num_epochs = 1000000 + with open(params.model_config, "r") as f: + model_config = json.load(f) + params.update(model_config["model"]) + params.update(model_config["feature"]) + + fix_random_seed(params.seed) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + os.makedirs(f"{params.exp_dir}", exist_ok=True) + copyfile(src=params.model_config, dst=f"{params.exp_dir}/model.json") + copyfile(src=params.token_file, dst=f"{params.exp_dir}/tokens.txt") + setup_logger(f"{params.exp_dir}/log/log-train") + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + if torch.cuda.is_available(): + params.device = torch.device("cuda", rank) + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + tokenizer = DialogTokenizer(token_file=params.token_file) + tokenizer_config = { + "vocab_size": tokenizer.vocab_size, + "pad_id": tokenizer.pad_id, + "spk_a_id": tokenizer.spk_a_id, + "spk_b_id": tokenizer.spk_b_id, + } + params.update(tokenizer_config) + + logging.info(params) + + logging.info("About to create model") + + model = ZipVoiceDialogStereo( + **model_config["model"], + **tokenizer_config, + ) + + assert params.checkpoint is not None + logging.info(f"Loading pre-trained model from {params.checkpoint}") + + if params.finetune: + # load a pre-trained ZipVoice-Dialog-Stereo model + _ = load_checkpoint(filename=params.checkpoint, model=model, strict=True) + else: + # load a pre-trained ZipVoice-Dialog model, duplicate the proj layers + load_checkpoint_copy_proj_three_channel_alter( + filename=params.checkpoint, + in_proj_key="fm_decoder.in_proj", + out_proj_key="fm_decoder.out_proj", + dim=params.feat_dim, + model=model, + ) + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of parameters : {num_param}") + + model_avg: Optional[nn.Module] = None + if rank == 0: + # model_avg is only used with rank 0 + model_avg = copy.deepcopy(model).to(torch.float64) + + assert params.start_epoch > 0, params.start_epoch + if params.start_epoch > 1: + checkpoints = resume_checkpoint(params=params, model=model, model_avg=model_avg) + + model = model.to(params.device) + if world_size > 1: + logging.info("Using DDP") + model = DDP(model, device_ids=[rank], find_unused_parameters=True) + + optimizer = ScaledAdam( + get_parameter_groups_with_lrs( + model, + lr=params.base_lr, + include_names=True, + ), + lr=params.base_lr, # should have no effect + clipping_scale=2.0, + ) + + scheduler = FixedLRScheduler(optimizer) + + scaler = create_grad_scaler(enabled=params.use_fp16) + + if params.start_epoch > 1 and checkpoints is not None: + # load state_dict for optimizers + if "optimizer" in checkpoints: + logging.info("Loading optimizer state dict") + optimizer.load_state_dict(checkpoints["optimizer"]) + + # load state_dict for schedulers + if "scheduler" in checkpoints: + logging.info("Loading scheduler state dict") + scheduler.load_state_dict(checkpoints["scheduler"]) + + if "grad_scaler" in checkpoints: + logging.info("Loading grad scaler state dict") + scaler.load_state_dict(checkpoints["grad_scaler"]) + + if params.print_diagnostics: + opts = diagnostics.TensorDiagnosticOptions( + 512 + ) # allow 4 megabytes per sub-module + diagnostic = diagnostics.attach_diagnostics(model, opts) + + if params.inf_check: + register_inf_check_hooks(model) + + def remove_short_and_long_utt(c: Cut, min_len: float, max_len: float): + if c.duration < min_len or c.duration > max_len: + return False + return True + + _remove_short_and_long_utt = partial( + remove_short_and_long_utt, min_len=params.min_len, max_len=params.max_len + ) + + datamodule = TtsDataModule(args) + train_cuts = datamodule.train_custom_cuts(params.train_manifest) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_custom_cuts(params.dev_manifest) + # To avoid OOM issues due to too long dev cuts + dev_cuts = dev_cuts.filter(_remove_short_and_long_utt) + + if not hasattr(train_cuts[0].supervisions[0], "tokens") or not hasattr( + dev_cuts[0].supervisions[0], "tokens" + ): + logging.warning( + "Tokens are not prepared, will tokenize on-the-fly, " + "which can slow down training significantly." + ) + _tokenize_text = partial(tokenize_text, tokenizer=tokenizer) + train_cuts = train_cuts.map(_tokenize_text) + dev_cuts = dev_cuts.map(_tokenize_text) + + train_dl = datamodule.train_dataloaders(train_cuts) + + valid_dl = datamodule.dev_dataloaders(dev_cuts) + + if params.scan_oom: + scan_pessimistic_batches_for_oom( + model=model, + train_dl=train_dl, + optimizer=optimizer, + params=params, + ) + + logging.info("Training started") + + for epoch in range(params.start_epoch, params.num_epochs + 1): + logging.info(f"Start epoch {epoch}") + scheduler.step_epoch(epoch - 1) + fix_random_seed(params.seed + epoch - 1) + train_dl.sampler.set_epoch(epoch - 1) + + params.cur_epoch = epoch + + if tb_writer is not None: + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + train_one_epoch( + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + train_dl=train_dl, + valid_dl=valid_dl, + scaler=scaler, + tb_writer=tb_writer, + world_size=world_size, + rank=rank, + ) + + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + + if params.print_diagnostics: + diagnostic.print_diagnostics() + break + + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint( + filename=filename, + params=params, + model=model, + model_avg=model_avg, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + + if rank == 0: + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + logging.info("Done!") + + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + TtsDataModule.add_arguments(parser) + args = parser.parse_args() + args.exp_dir = Path(args.exp_dir) + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + main() diff --git a/zipvoice/bin/train_zipvoice_distill.py b/zipvoice/bin/train_zipvoice_distill.py new file mode 100644 index 0000000..0a6fc9f --- /dev/null +++ b/zipvoice/bin/train_zipvoice_distill.py @@ -0,0 +1,1158 @@ +#!/usr/bin/env python3 +# Copyright 2024 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +""" +This script trains a ZipVoice-Distill model starting from a ZipVoice model. +It has two distillation stages. + +Usage: + +(1) The first distillation stage with a fixed ZipVoice model as the teacher. + +python3 -m zipvoice.bin.train_zipvoice_distill \ + --world-size 8 \ + --use-fp16 1 \ + --num-iters 60000 \ + --max-duration 500 \ + --base-lr 0.0005 \ + --tokenizer emilia \ + --token-file data/tokens_emilia.txt \ + --dataset emilia \ + --manifest-dir data/fbank \ + --teacher-model exp/zipvoice/epoch-11-avg-4.pt \ + --distill-stage first \ + --exp-dir exp/zipvoice_distill_1stage + +(2) The second distillation stage with a EMA model as the teacher. +python3 -m zipvoice.bin.train_zipvoice_distill \ + --world-size 8 \ + --use-fp16 1 \ + --num-iters 2000 \ + --save-every-n 1000 \ + --max-duration 500 \ + --base-lr 0.0001 \ + --model-config conf/zipvoice_base.json \ + --tokenizer emilia \ + --token-file data/tokens_emilia.txt \ + --dataset emilia \ + --manifest-dir data/fbank \ + --teacher-model exp/zipvoice_distill_1stage/iter-60000-avg-7.pt \ + --distill-stage second \ + --exp-dir exp/zipvoice_distill +""" + +import argparse +import copy +import json +import logging +import os +import random +from functools import partial +from pathlib import Path +from shutil import copyfile +from typing import List, Optional, Tuple, Union + +import torch +import torch.multiprocessing as mp +import torch.nn as nn +from lhotse.cut import Cut, CutSet +from lhotse.utils import fix_random_seed +from torch import Tensor +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer +from torch.utils.tensorboard import SummaryWriter + +import zipvoice.utils.diagnostics as diagnostics +from zipvoice.bin.train_zipvoice import ( + display_and_save_batch, + get_params, + tokenize_text, +) +from zipvoice.dataset.datamodule import TtsDataModule +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import ( + EmiliaTokenizer, + EspeakTokenizer, + LibriTTSTokenizer, + SimpleTokenizer, +) +from zipvoice.utils.checkpoint import ( + load_checkpoint, + remove_checkpoints, + resume_checkpoint, + save_checkpoint, + save_checkpoint_with_global_batch_idx, + update_averaged_model, +) +from zipvoice.utils.common import ( + AttributeDict, + GradScaler, + MetricsTracker, + cleanup_dist, + condition_time_mask, + create_grad_scaler, + get_adjusted_batch_count, + get_parameter_groups_with_lrs, + make_pad_mask, + prepare_input, + set_batch_count, + setup_dist, + setup_logger, + str2bool, + torch_autocast, +) +from zipvoice.utils.hooks import register_inf_check_hooks +from zipvoice.utils.lr_scheduler import FixedLRScheduler, LRScheduler +from zipvoice.utils.optim import ScaledAdam + +LRSchedulerType = Union[torch.optim.lr_scheduler._LRScheduler, LRScheduler] + + +def get_parser(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + + parser.add_argument( + "--world-size", + type=int, + default=1, + help="Number of GPUs for DDP training.", + ) + + parser.add_argument( + "--master-port", + type=int, + default=12356, + help="Master port to use for DDP training.", + ) + + parser.add_argument( + "--tensorboard", + type=str2bool, + default=True, + help="Should various information be logged in tensorboard.", + ) + + parser.add_argument( + "--num-epochs", + type=int, + default=1, + help="Number of epochs to train.", + ) + + parser.add_argument( + "--num-iters", + type=int, + default=0, + help="Number of iter to train, will ignore num_epochs if > 0.", + ) + + parser.add_argument( + "--start-epoch", + type=int, + default=1, + help="""Resume training from this epoch. It should be positive. + If larger than 1, it will load checkpoint from + exp-dir/epoch-{start_epoch-1}.pt + """, + ) + + parser.add_argument( + "--teacher-model", + type=str, + help="""Checkpoints of pre-trained teacher model""", + ) + + parser.add_argument( + "--exp-dir", + type=str, + default="exp/zipvoice_distill", + help="""The experiment dir. + It specifies the directory where all training related + files, e.g., checkpoints, log, etc, are saved + """, + ) + + parser.add_argument( + "--base-lr", type=float, default=0.001, help="The base learning rate." + ) + + parser.add_argument( + "--ref-duration", + type=float, + default=50, + help="Reference batch duration for purposes of adjusting batch counts for " + "setting various schedules inside the model", + ) + + parser.add_argument( + "--seed", + type=int, + default=42, + help="The seed for random generators intended for reproducibility", + ) + + parser.add_argument( + "--print-diagnostics", + type=str2bool, + default=False, + help="Accumulate stats on activations, print them and exit.", + ) + + parser.add_argument( + "--scan-oom", + type=str2bool, + default=False, + help="Scan pessimistic batches to see whether they cause OOMs.", + ) + + parser.add_argument( + "--inf-check", + type=str2bool, + default=False, + help="Add hooks to check for infinite module outputs and gradients.", + ) + + parser.add_argument( + "--save-every-n", + type=int, + default=1000, + help="""Save checkpoint after processing this number of batches" + periodically. We save checkpoint to exp-dir/ whenever + params.batch_idx_train % save_every_n == 0. The checkpoint filename + has the form: f'exp-dir/checkpoint-{params.batch_idx_train}.pt' + Note: It also saves checkpoint to `exp-dir/epoch-xxx.pt` at the + end of each epoch where `xxx` is the epoch number counting from 1. + """, + ) + + parser.add_argument( + "--keep-last-k", + type=int, + default=30, + help="""Only keep this number of checkpoints on disk. + For instance, if it is 3, there are only 3 checkpoints + in the exp-dir with filenames `checkpoint-xxx.pt`. + It does not affect checkpoints with name `epoch-xxx.pt`. + """, + ) + + parser.add_argument( + "--average-period", + type=int, + default=200, + help="""Update the averaged model, namely `model_avg`, after processing + this number of batches. `model_avg` is a separate version of model, + in which each floating-point parameter is the average of all the + parameters from the start of training. Each time we take the average, + we do: `model_avg = model * (average_period / batch_idx_train) + + model_avg * ((batch_idx_train - average_period) / batch_idx_train)`. + """, + ) + + parser.add_argument( + "--use-fp16", + type=str2bool, + default=True, + help="Whether to use half precision training.", + ) + + parser.add_argument( + "--feat-scale", + type=float, + default=0.1, + help="The scale factor of fbank feature", + ) + + parser.add_argument( + "--ema-decay", + type=float, + default=0.9999, + help="The EMA decay factor of target model in distillation.", + ) + parser.add_argument( + "--distill-stage", + type=str, + choices=["first", "second"], + help="The stage of distillation.", + ) + + parser.add_argument( + "--dataset", + type=str, + default="emilia", + choices=["emilia", "libritts", "custom"], + help="The used training dataset", + ) + + parser.add_argument( + "--train-manifest", + type=str, + help="Path of the training manifest", + ) + + parser.add_argument( + "--dev-manifest", + type=str, + help="Path of the validation manifest", + ) + + parser.add_argument( + "--min-len", + type=float, + default=1.0, + help="The minimum audio length used for training", + ) + + parser.add_argument( + "--max-len", + type=float, + default=30.0, + help="The maximum audio length used for training", + ) + + parser.add_argument( + "--model-config", + type=str, + default="conf/zipvoice_base.json", + help="The model configuration file.", + ) + + parser.add_argument( + "--tokenizer", + type=str, + default="emilia", + choices=["emilia", "libritts", "espeak", "simple"], + help="Tokenizer type.", + ) + + parser.add_argument( + "--lang", + type=str, + default="en-us", + help="Language identifier, used when tokenizer type is espeak. see" + "https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md", + ) + + parser.add_argument( + "--token-file", + type=str, + default="data/tokens_emilia.txt", + help="The file that contains information that maps tokens to ids," + "which is a text file with '{token}\t{token_id}' per line.", + ) + + return parser + + +def ema(new_model, ema_model, decay): + if isinstance(new_model, DDP): + new_model = new_model.module + if isinstance(ema_model, DDP): + ema_model = ema_model.module + new_model_dict = new_model.state_dict() + ema_model_dict = ema_model.state_dict() + for key in new_model_dict.keys(): + ema_model_dict[key].data.copy_( + ema_model_dict[key].data * decay + new_model_dict[key].data * (1 - decay) + ) + + +def compute_fbank_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + teacher_model: Union[nn.Module, DDP], + features: Tensor, + features_lens: Tensor, + tokens: List[List[int]], + is_training: bool, +) -> Tuple[Tensor, MetricsTracker]: + """ + Compute loss given the model and its inputs. + + Args: + params: + Parameters for training. See :func:`get_params`. + model: + The model for training. + teacher_model: + The teacher model for distillation. + features: + The target acoustic feature. + features_lens: + The number of frames of each utterance. + tokens: + Input tokens that representing the transcripts. + is_training: + True for training. False for validation. When it is True, this + function enables autograd during computation; when it is False, it + disables autograd. + """ + + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batch_size, num_frames, _ = features.shape + + noise = torch.randn_like(features) # (B, T, F) + + # Sampling t and guidance_scale from uniform distribution + + t_value = random.random() + t = torch.ones(batch_size, 1, 1, device=device) * t_value + if params.distill_stage == "first": + guidance_scale = torch.rand(batch_size, 1, 1, device=device) * 2 + else: + guidance_scale = torch.rand(batch_size, 1, 1, device=device) * 2 + 1 + xt = features * t + noise * (1 - t) + t_delta_fix = random.uniform(0.0, min(0.3, 1 - t_value)) + t_delta_ema = random.uniform(0.0, min(0.3, 1 - t_value - t_delta_fix)) + t_dest = t_value + t_delta_fix + t_delta_ema + + with torch.no_grad(): + speech_condition_mask = condition_time_mask( + features_lens=features_lens, + mask_percent=(0.7, 1.0), + max_len=features.size(1), + ) + + if params.distill_stage == "first": + teacher_x_t_mid, _ = teacher_model.sample_intermediate( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=xt, + speech_condition_mask=speech_condition_mask, + t_start=t_value, + t_end=t_value + t_delta_fix, + num_step=1, + guidance_scale=guidance_scale, + ) + + target_x1, _ = teacher_model.sample_intermediate( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=teacher_x_t_mid, + speech_condition_mask=speech_condition_mask, + t_start=t_value + t_delta_fix, + t_end=t_dest, + num_step=1, + guidance_scale=guidance_scale, + ) + else: + teacher_x_t_mid, _ = teacher_model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=xt, + speech_condition_mask=speech_condition_mask, + t_start=t_value, + t_end=t_value + t_delta_fix, + num_step=1, + guidance_scale=guidance_scale, + ) + + target_x1, _ = teacher_model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=teacher_x_t_mid, + speech_condition_mask=speech_condition_mask, + t_start=t_value + t_delta_fix, + t_end=t_dest, + num_step=1, + guidance_scale=guidance_scale, + ) + + with torch.set_grad_enabled(is_training): + + pred_x1, _ = model( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=xt, + speech_condition_mask=speech_condition_mask, + t_start=t_value, + t_end=t_dest, + num_step=1, + guidance_scale=guidance_scale, + ) + pred_v = (pred_x1 - xt) / (t_dest - t) + + padding_mask = make_pad_mask(features_lens, max_len=num_frames) # (B, T) + loss_mask = speech_condition_mask & (~padding_mask) + + target_v = (target_x1 - xt) / (t_dest - t) + loss = torch.mean((pred_v[loss_mask] - target_v[loss_mask]) ** 2) + + ut = features - noise # (B, T, F) + + ref_loss = torch.mean((pred_v[loss_mask] - ut[loss_mask]) ** 2) + + assert loss.requires_grad == is_training + info = MetricsTracker() + num_frames = features_lens.sum().item() + info["frames"] = num_frames + info["loss"] = loss.detach().cpu().item() * num_frames + info["ref_loss"] = ref_loss.detach().cpu().item() * num_frames + return loss, info + + +def train_one_epoch( + params: AttributeDict, + model: Union[nn.Module, DDP], + teacher_model: Union[nn.Module, DDP], + optimizer: Optimizer, + scheduler: LRSchedulerType, + train_dl: torch.utils.data.DataLoader, + valid_dl: torch.utils.data.DataLoader, + scaler: GradScaler, + model_avg: Optional[nn.Module] = None, + tb_writer: Optional[SummaryWriter] = None, + world_size: int = 1, + rank: int = 0, +) -> None: + """Train the model for one epoch. + + The training loss from the mean of all frames is saved in + `params.train_loss`. It runs the validation process every + `params.valid_interval` batches. + + Args: + params: + It is returned by :func:`get_params`. + model: + The model for training. + teacher_model: + The model for distillation. + Used to convert text to tokens. + optimizer: + The optimizer. + scheduler: + The learning rate scheduler, we call step() every epoch. + train_dl: + Dataloader for the training dataset. + valid_dl: + Dataloader for the validation dataset. + scaler: + The scaler used for mix precision training. + tb_writer: + Writer to write log messages to tensorboard. + world_size: + Number of nodes in DDP training. If it is 1, DDP is disabled. + rank: + The rank of the node in DDP training. If no DDP is used, it should + be set to 0. + """ + model.train() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to track the stats over iterations in one epoch + tot_loss = MetricsTracker() + + saved_bad_model = False + + def save_bad_model(suffix: str = ""): + save_checkpoint( + filename=params.exp_dir / f"bad-model{suffix}-{rank}.pt", + model=model, + model_avg=model_avg, + model_ema=teacher_model, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=0, + ) + + for batch_idx, batch in enumerate(train_dl): + + if batch_idx % 10 == 0: + set_batch_count(model, get_adjusted_batch_count(params) + 100000) + + if ( + params.batch_idx_train % params.valid_interval == 0 + and not params.print_diagnostics + ): + logging.info("Computing validation loss") + valid_info = compute_validation_loss( + params=params, + model=model, + teacher_model=teacher_model, + valid_dl=valid_dl, + world_size=world_size, + ) + model.train() + logging.info( + f"Epoch {params.cur_epoch}, global_batch_idx: {params.batch_idx_train}," + f" validation: {valid_info}" + ) + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + if tb_writer is not None: + valid_info.write_summary( + tb_writer, "train/valid_", params.batch_idx_train + ) + + params.batch_idx_train += 1 + + batch_size = len(batch["text"]) + + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + teacher_model=teacher_model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + + tot_loss = (tot_loss * (1 - 1 / params.reset_interval)) + loss_info + + scaler.scale(loss).backward() + + scheduler.step_batch(params.batch_idx_train) + scaler.step(optimizer) + scaler.update() + optimizer.zero_grad() + if params.distill_stage == "second": + ema(model, teacher_model, params.ema_decay) + except Exception as e: + logging.info(f"Caught exception : {e}.") + save_bad_model() + raise + + if params.print_diagnostics and batch_idx == 5: + return + + if ( + rank == 0 + and params.batch_idx_train > 0 + and params.batch_idx_train % params.average_period == 0 + ): + update_averaged_model( + params=params, + model_cur=model, + model_avg=model_avg, + ) + + if ( + params.batch_idx_train > 0 + and params.batch_idx_train % params.save_every_n == 0 + ): + save_checkpoint_with_global_batch_idx( + out_dir=params.exp_dir, + global_batch_idx=params.batch_idx_train, + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + remove_checkpoints( + out_dir=params.exp_dir, + topk=params.keep_last_k, + rank=rank, + ) + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + if params.batch_idx_train % 100 == 0 and params.use_fp16: + # If the grad scale was less than 1, try increasing it. The _growth_interval + # of the grad scaler is configurable, but we can't configure it to have + # different behavior depending on the current grad scale. + cur_grad_scale = scaler._scale.item() + + if cur_grad_scale < 1024.0 or ( + cur_grad_scale < 4096.0 and params.batch_idx_train % 400 == 0 + ): + scaler.update(cur_grad_scale * 2.0) + if cur_grad_scale < 0.01: + if not saved_bad_model: + save_bad_model(suffix="-first-warning") + saved_bad_model = True + logging.warning(f"Grad scale is small: {cur_grad_scale}") + if cur_grad_scale < 1.0e-05: + save_bad_model() + raise RuntimeError( + f"grad_scale is too small, exiting: {cur_grad_scale}" + ) + + if params.batch_idx_train % params.log_interval == 0: + cur_lr = max(scheduler.get_last_lr()) + cur_grad_scale = scaler._scale.item() if params.use_fp16 else 1.0 + + logging.info( + f"Epoch {params.cur_epoch}, batch {batch_idx}, " + f"global_batch_idx: {params.batch_idx_train}, " + f"batch size: {batch_size}, " + f"loss[{loss_info}], tot_loss[{tot_loss}], " + f"cur_lr: {cur_lr:.2e}, " + + (f"grad_scale: {scaler._scale.item()}" if params.use_fp16 else "") + ) + + if tb_writer is not None: + tb_writer.add_scalar( + "train/learning_rate", cur_lr, params.batch_idx_train + ) + loss_info.write_summary( + tb_writer, "train/current_", params.batch_idx_train + ) + tot_loss.write_summary(tb_writer, "train/tot_", params.batch_idx_train) + if params.use_fp16: + tb_writer.add_scalar( + "train/grad_scale", + cur_grad_scale, + params.batch_idx_train, + ) + + loss_value = tot_loss["loss"] + params.train_loss = loss_value + if params.train_loss < params.best_train_loss: + params.best_train_epoch = params.cur_epoch + params.best_train_loss = params.train_loss + + +def compute_validation_loss( + params: AttributeDict, + model: Union[nn.Module, DDP], + teacher_model: Optional[nn.Module], + valid_dl: torch.utils.data.DataLoader, + world_size: int = 1, +) -> MetricsTracker: + """Run the validation process.""" + + model.eval() + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + # used to summary the stats over iterations + tot_loss = MetricsTracker() + + for batch_idx, batch in enumerate(valid_dl): + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + teacher_model=teacher_model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=False, + ) + assert loss.requires_grad is False + tot_loss = tot_loss + loss_info + + if world_size > 1: + tot_loss.reduce(loss.device) + + loss_value = tot_loss["loss"] + if loss_value < params.best_valid_loss: + params.best_valid_epoch = params.cur_epoch + params.best_valid_loss = loss_value + + return tot_loss + + +def scan_pessimistic_batches_for_oom( + model: Union[nn.Module, DDP], + teacher_model: nn.Module, + train_dl: torch.utils.data.DataLoader, + optimizer: torch.optim.Optimizer, + params: AttributeDict, +): + from lhotse.dataset import find_pessimistic_batches + + logging.info( + "Sanity check -- see if any of the batches in epoch 1 would cause OOM." + ) + device = model.device if isinstance(model, DDP) else next(model.parameters()).device + + batches, crit_values = find_pessimistic_batches(train_dl.sampler) + for criterion, cuts in batches.items(): + batch = train_dl.dataset[cuts] + tokens, features, features_lens = prepare_input( + params=params, + batch=batch, + device=device, + return_tokens=True, + return_feature=True, + ) + try: + with torch_autocast(dtype=torch.float16, enabled=params.use_fp16): + + loss, loss_info = compute_fbank_loss( + params=params, + model=model, + teacher_model=teacher_model, + features=features, + features_lens=features_lens, + tokens=tokens, + is_training=True, + ) + loss.backward() + optimizer.zero_grad() + except Exception as e: + if "CUDA out of memory" in str(e): + logging.error( + "Your GPU ran out of memory with the current " + "max_duration setting. We recommend decreasing " + "max_duration and trying again.\n" + f"Failing criterion: {criterion} " + f"(={crit_values[criterion]}) ..." + ) + display_and_save_batch(batch, params=params) + raise + logging.info( + f"Maximum memory allocated so far is " + f"{torch.cuda.max_memory_allocated() // 1000000}MB" + ) + + +def run(rank, world_size, args): + """ + Args: + rank: + It is a value between 0 and `world_size-1`, which is + passed automatically by `mp.spawn()` in :func:`main`. + The node with rank 0 is responsible for saving checkpoint. + world_size: + Number of GPUs for DDP training. + args: + The return value of get_parser().parse_args() + """ + params = get_params() + params.update(vars(args)) + params.valid_interval = params.save_every_n + # Set epoch to a large number to ignore it. + if params.num_iters > 0: + params.num_epochs = 1000000 + with open(params.model_config, "r") as f: + model_config = json.load(f) + params.update(model_config["model"]) + params.update(model_config["feature"]) + + fix_random_seed(params.seed) + if world_size > 1: + setup_dist(rank, world_size, params.master_port) + + os.makedirs(f"{params.exp_dir}", exist_ok=True) + copyfile(src=params.model_config, dst=f"{params.exp_dir}/model.json") + copyfile(src=params.token_file, dst=f"{params.exp_dir}/tokens.txt") + setup_logger(f"{params.exp_dir}/log/log-train") + + if args.tensorboard and rank == 0: + tb_writer = SummaryWriter(log_dir=f"{params.exp_dir}/tensorboard") + else: + tb_writer = None + + if torch.cuda.is_available(): + params.device = torch.device("cuda", rank) + else: + params.device = torch.device("cpu") + logging.info(f"Device: {params.device}") + + if params.tokenizer == "emilia": + tokenizer = EmiliaTokenizer(token_file=params.token_file) + elif params.tokenizer == "libritts": + tokenizer = LibriTTSTokenizer(token_file=params.token_file) + elif params.tokenizer == "espeak": + tokenizer = EspeakTokenizer(token_file=params.token_file, lang=params.lang) + else: + assert params.tokenizer == "simple" + tokenizer = SimpleTokenizer(token_file=params.token_file) + + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + params.update(tokenizer_config) + + logging.info(params) + + logging.info("About to create model") + + assert params.teacher_model is not None + logging.info(f"Loading pre-trained model from {params.teacher_model}") + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + _ = load_checkpoint( + filename=params.teacher_model, + model=model, + strict=(params.distill_stage == "second"), + ) + + if params.distill_stage == "first": + teacher_model = ZipVoice( + **model_config["model"], + **tokenizer_config, + ) + _ = load_checkpoint( + filename=params.teacher_model, model=teacher_model, strict=True + ) + else: + teacher_model = copy.deepcopy(model) + + num_param = sum([p.numel() for p in model.parameters()]) + logging.info(f"Number of parameters : {num_param}") + + model_avg: Optional[nn.Module] = None + if rank == 0: + # model_avg is only used with rank 0 + model_avg = copy.deepcopy(model).to(torch.float64) + assert params.start_epoch > 0, params.start_epoch + if params.start_epoch > 1: + logging.info(f"Resuming from epoch {params.start_epoch}") + if params.distill_stage == "first": + checkpoints = resume_checkpoint( + params=params, model=model, model_avg=model_avg + ) + else: + checkpoints = resume_checkpoint( + params=params, + model=model, + model_avg=model_avg, + model_ema=teacher_model, + ) + + model = model.to(params.device) + teacher_model.to(params.device) + teacher_model.eval() + + if world_size > 1: + logging.info("Using DDP") + model = DDP(model, device_ids=[rank], find_unused_parameters=True) + + # only update the fm_decoder + num_trainable = 0 + for name, p in model.named_parameters(): + if "fm_decoder" in name: + p.requires_grad = True + num_trainable += p.numel() + else: + p.requires_grad = False + + logging.info( + "A total of {} trainable parameters ({:.3f}% of the whole model)".format( + num_trainable, num_trainable / num_param * 100 + ) + ) + + optimizer = ScaledAdam( + get_parameter_groups_with_lrs( + model, + lr=params.base_lr, + include_names=True, + ), + lr=params.base_lr, # should have no effect + clipping_scale=2.0, + ) + + scheduler = FixedLRScheduler(optimizer) + + scaler = create_grad_scaler(enabled=params.use_fp16) + + if params.start_epoch > 1 and checkpoints is not None: + # load state_dict for optimizers + if "optimizer" in checkpoints: + logging.info("Loading optimizer state dict") + optimizer.load_state_dict(checkpoints["optimizer"]) + + # load state_dict for schedulers + if "scheduler" in checkpoints: + logging.info("Loading scheduler state dict") + scheduler.load_state_dict(checkpoints["scheduler"]) + + if "grad_scaler" in checkpoints: + logging.info("Loading grad scaler state dict") + scaler.load_state_dict(checkpoints["grad_scaler"]) + + if params.print_diagnostics: + opts = diagnostics.TensorDiagnosticOptions( + 512 + ) # allow 4 megabytes per sub-module + diagnostic = diagnostics.attach_diagnostics(model, opts) + + if params.inf_check: + register_inf_check_hooks(model) + + def remove_short_and_long_utt(c: Cut, min_len: float, max_len: float): + if c.duration < min_len or c.duration > max_len: + return False + return True + + _remove_short_and_long_utt = partial( + remove_short_and_long_utt, min_len=params.min_len, max_len=params.max_len + ) + + datamodule = TtsDataModule(args) + if params.dataset == "emilia": + train_cuts = CutSet.mux( + datamodule.train_emilia_EN_cuts(), + datamodule.train_emilia_ZH_cuts(), + weights=[46000, 49000], + ) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = CutSet.mux( + datamodule.dev_emilia_EN_cuts(), + datamodule.dev_emilia_ZH_cuts(), + weights=[0.5, 0.5], + ) + elif params.dataset == "libritts": + train_cuts = datamodule.train_libritts_cuts() + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_libritts_cuts() + else: + assert params.dataset == "custom" + train_cuts = datamodule.train_custom_cuts(params.train_manifest) + train_cuts = train_cuts.filter(_remove_short_and_long_utt) + dev_cuts = datamodule.dev_custom_cuts(params.dev_manifest) + # To avoid OOM issues due to too long dev cuts + dev_cuts = dev_cuts.filter(_remove_short_and_long_utt) + + if params.tokenizer in ["emilia", "espeak", "dialog"]: + if not hasattr(train_cuts[0].supervisions[0], "tokens") or not hasattr( + dev_cuts[0].supervisions[0], "tokens" + ): + logging.warning( + f"Using {params.tokenizer} tokenizer but tokens are not prepared," + f"will tokenize on-the-fly, which can slow down training significantly." + ) + _tokenize_text = partial(tokenize_text, tokenizer=tokenizer) + train_cuts = train_cuts.map(_tokenize_text) + dev_cuts = dev_cuts.map(_tokenize_text) + + train_dl = datamodule.train_dataloaders(train_cuts) + + valid_dl = datamodule.dev_dataloaders(dev_cuts) + + if params.scan_oom: + scan_pessimistic_batches_for_oom( + model=model, + teacher_model=teacher_model, + train_dl=train_dl, + optimizer=optimizer, + params=params, + ) + logging.info("Training started") + + for epoch in range(params.start_epoch, params.num_epochs + 1): + logging.info(f"Start epoch {epoch}") + + scheduler.step_epoch(epoch - 1) + fix_random_seed(params.seed + epoch - 1) + train_dl.sampler.set_epoch(epoch - 1) + + params.cur_epoch = epoch + + if tb_writer is not None: + tb_writer.add_scalar("train/epoch", epoch, params.batch_idx_train) + + train_one_epoch( + params=params, + model=model, + model_avg=model_avg, + teacher_model=teacher_model, + optimizer=optimizer, + scheduler=scheduler, + train_dl=train_dl, + valid_dl=valid_dl, + scaler=scaler, + tb_writer=tb_writer, + world_size=world_size, + rank=rank, + ) + + if params.num_iters > 0 and params.batch_idx_train > params.num_iters: + break + + if params.print_diagnostics: + diagnostic.print_diagnostics() + break + + filename = params.exp_dir / f"epoch-{params.cur_epoch}.pt" + save_checkpoint( + filename=filename, + params=params, + model=model, + model_avg=model_avg, + model_ema=teacher_model, + optimizer=optimizer, + scheduler=scheduler, + sampler=train_dl.sampler, + scaler=scaler, + rank=rank, + ) + + if rank == 0: + if params.best_train_epoch == params.cur_epoch: + best_train_filename = params.exp_dir / "best-train-loss.pt" + copyfile(src=filename, dst=best_train_filename) + + if params.best_valid_epoch == params.cur_epoch: + best_valid_filename = params.exp_dir / "best-valid-loss.pt" + copyfile(src=filename, dst=best_valid_filename) + + logging.info("Done!") + + if world_size > 1: + torch.distributed.barrier() + cleanup_dist() + + +def main(): + parser = get_parser() + TtsDataModule.add_arguments(parser) + args = parser.parse_args() + args.exp_dir = Path(args.exp_dir) + + world_size = args.world_size + assert world_size >= 1 + if world_size > 1: + mp.spawn(run, args=(world_size, args), nprocs=world_size, join=True) + else: + run(rank=0, world_size=1, args=args) + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + main() diff --git a/zipvoice/luxvoice.py b/zipvoice/luxvoice.py new file mode 100644 index 0000000..72db595 --- /dev/null +++ b/zipvoice/luxvoice.py @@ -0,0 +1,52 @@ +from zipvoice.modeling_utils import process_audio, generate, load_models_gpu, load_models_cpu +from zipvoice.onnx_modeling import generate_cpu + +class LuxTTS: + """ + LuxTTS class for encoding prompt and generating speech on cpu/cuda. + """ + + def __init__(self, model_path='YatharthS/LuxTTS', device='cuda', threads=4): + if model_path == 'YatharthS/LuxTTS': + model_path = None + + if device == 'cpu': + model, feature_extractor, vocos, tokenizer, transcriber = load_models_cpu(model_path, threads) + print("Loading model on CPU") + else: + model, feature_extractor, vocos, tokenizer, transcriber = load_models_gpu(model_path) + print("Loading model on GPU") + + self.model = model + self.feature_extractor = feature_extractor + self.vocos = vocos + self.tokenizer = tokenizer + self.transcriber = transcriber + self.device = device + self.vocos.freq_range = 12000 + + + + def encode_prompt(self, prompt_audio, duration=5, rms=0.001): + """encodes audio prompt according to duration and rms(volume control)""" + prompt_tokens, prompt_features_lens, prompt_features, prompt_rms = process_audio(prompt_audio, self.transcriber, self.tokenizer, self.feature_extractor, self.device, target_rms=rms, duration=duration) + encode_dict = {"prompt_tokens": prompt_tokens, 'prompt_features_lens': prompt_features_lens, 'prompt_features': prompt_features, 'prompt_rms': prompt_rms} + + return encode_dict + + def generate_speech(self, text, encode_dict, num_steps=4, guidance_scale=3.0, t_shift=0.5, speed=1.0, return_smooth=False): + """encodes text and generates speech using flow matching model according to steps, guidance scale, and t_shift(like temp)""" + + prompt_tokens, prompt_features_lens, prompt_features, prompt_rms = encode_dict.values() + + if return_smooth == True: + self.vocos.return_48k = False + else: + self.vocos.return_48k = True + + if self.device == 'cpu': + final_wav = generate_cpu(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, self.model, self.vocos, self.tokenizer, num_step=num_steps, guidance_scale=guidance_scale, t_shift=t_shift, speed=speed) + else: + final_wav = generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, self.model, self.vocos, self.tokenizer, num_step=num_steps, guidance_scale=guidance_scale, t_shift=t_shift, speed=speed) + + return final_wav.cpu() diff --git a/zipvoice/modeling_utils.py b/zipvoice/modeling_utils.py new file mode 100644 index 0000000..893e3fe --- /dev/null +++ b/zipvoice/modeling_utils.py @@ -0,0 +1,157 @@ +import argparse +import datetime as dt +import json +import logging +import os +from pathlib import Path +from typing import Optional + +import numpy as np +import safetensors.torch +import torch +import librosa +import torchaudio +from transformers import pipeline +from huggingface_hub import snapshot_download +from lhotse.utils import fix_random_seed + +from zipvoice.models.zipvoice_distill import ZipVoiceDistill +from zipvoice.tokenizer.tokenizer import EmiliaTokenizer +from zipvoice.utils.checkpoint import load_checkpoint +from zipvoice.utils.common import AttributeDict, str2bool +from zipvoice.utils.feature import VocosFbank +from zipvoice.utils.infer import rms_norm + +from dataclasses import dataclass, field +from typing import Optional, List + +from linacodec.vocoder.vocos import Vocos +from zipvoice.onnx_modeling import OnnxModel +from torch.nn.utils import parametrize + + +@dataclass +class LuxTTSConfig: + # Model Setup + model_dir: Optional[str] = None + checkpoint_name: str = "model.pt" + vocoder_path: Optional[str] = None + trt_engine_path: Optional[str] = None + + # Tokenizer & Language + tokenizer: str = "emilia" # choices: ["emilia", "libritts", "espeak", "simple"] + lang: str = "en-us" + + +@torch.inference_mode +def process_audio(audio, transcriber, tokenizer, feature_extractor, device, target_rms=0.1, duration=4, feat_scale=0.1): + prompt_wav, sr = librosa.load(audio, sr=24000, duration=duration) + prompt_wav2, sr = librosa.load(audio, sr=16000, duration=duration) + prompt_text = transcriber(prompt_wav2)["text"] + print(prompt_text) + + prompt_wav = torch.from_numpy(prompt_wav).unsqueeze(0) + prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms) + + prompt_features = feature_extractor.extract( + prompt_wav, sampling_rate=24000 + ).to(device) + prompt_features = prompt_features.unsqueeze(0) * feat_scale + prompt_features_lens = torch.tensor([prompt_features.size(1)], device=device) + prompt_tokens = tokenizer.texts_to_token_ids([prompt_text]) + return prompt_tokens, prompt_features_lens, prompt_features, prompt_rms + +def generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, model, vocoder, tokenizer, num_step=4, guidance_scale=3.0, speed=1.0, t_shift=0.5, target_rms=0.1): + tokens = tokenizer.texts_to_token_ids([text]) + device = next(model.parameters()).device # Auto-detect device + + speed = speed * 1.3 + + with torch.inference_mode(): + (pred_features, _, _, _) = model.sample( + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + prompt_features_lens=prompt_features_lens, + speed=speed, + t_shift=t_shift, + duration='predict', + num_step=num_step, + guidance_scale=guidance_scale, + ) + + # Convert to waveform + pred_features = pred_features.permute(0, 2, 1) / 0.1 + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + + # Volume matching + if prompt_rms < target_rms: + wav = wav * (prompt_rms / target_rms) + + return wav + +def load_models_gpu(model_path=None): + params = LuxTTSConfig() + if model_path is None: + model_path = snapshot_download("YatharthS/LuxTTS") + + token_file = f"{model_path}/tokens.txt" + model_ckpt = f"{model_path}/model.pt" + model_config = f"{model_path}/config.json" + + transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base", device='cuda') + tokenizer = EmiliaTokenizer(token_file=token_file) + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + + with open(model_config, "r") as f: + model_config = json.load(f) + + model = ZipVoiceDistill( + **model_config["model"], + **tokenizer_config, + ) + load_checkpoint(filename=model_ckpt, model=model, strict=True) + params.device = torch.device("cuda", 0) + + model = model.to(params.device).eval() + feature_extractor = VocosFbank() + + vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').cuda() + parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight") + parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight") + vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin')) + + params.sampling_rate = model_config["feature"]["sampling_rate"] + return model, feature_extractor, vocos, tokenizer, transcriber + +def load_models_cpu(model_path = None, num_thread=2): + params = LuxTTSConfig() + params.seed = 42 + + model_path = snapshot_download('YatharthS/LuxTTS') + + token_file = f"{model_path}/tokens.txt" + text_encoder_path = f"{model_path}/text_encoder.onnx" + fm_decoder_path = f"{model_path}/fm_decoder.onnx" + model_config = f"{model_path}/config.json" + + transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device='cpu') + + tokenizer = EmiliaTokenizer(token_file=token_file) + tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id} + + with open(model_config, "r") as f: + model_config = json.load(f) + + model = OnnxModel(text_encoder_path, fm_decoder_path, num_thread=num_thread) + + vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').eval() + parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight") + parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight") + vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin', map_location=torch.device('cpu'))) + + feature_extractor = VocosFbank() + + params.sampling_rate = model_config["feature"]["sampling_rate"] + params.onnx_int8 = True + return model, feature_extractor, vocos, tokenizer, transcriber diff --git a/zipvoice/models/modules/scaling.py b/zipvoice/models/modules/scaling.py new file mode 100644 index 0000000..580402b --- /dev/null +++ b/zipvoice/models/modules/scaling.py @@ -0,0 +1,1590 @@ +# Copyright 2022-2025 Xiaomi Corp. (authors: Daniel Povey +# Wei Kang) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import logging +import math +import random +import sys +from typing import Optional, Tuple, Union + +try: + import k2 +except Exception as e: + logging.warning( + f"Failed import k2 with error {e}. Swoosh functions will fallback to PyTorch" + f" implementation, leading to slower speed and higher memory consumption." + ) +import torch +import torch.nn as nn +from torch import Tensor + + +def custom_amp_decorator(dec, cuda_amp_deprecated): + def decorator(func): + return dec(func) if not cuda_amp_deprecated else dec(device_type="cuda")(func) + + return decorator + + +if hasattr(torch.amp, "custom_fwd"): + deprecated = True + from torch.amp import custom_bwd, custom_fwd +else: + deprecated = False + from torch.cuda.amp import custom_bwd, custom_fwd + +custom_fwd = custom_amp_decorator(custom_fwd, deprecated) +custom_bwd = custom_amp_decorator(custom_bwd, deprecated) + + +def logaddexp_onnx(x: Tensor, y: Tensor) -> Tensor: + max_value = torch.max(x, y) + diff = torch.abs(x - y) + return max_value + torch.log1p(torch.exp(-diff)) + + +# RuntimeError: Exporting the operator logaddexp to ONNX opset version +# 14 is not supported. Please feel free to request support or submit +# a pull request on PyTorch GitHub. +# +# The following function is to solve the above error when exporting +# models to ONNX via torch.jit.trace() +def logaddexp(x: Tensor, y: Tensor) -> Tensor: + # Caution(fangjun): Put torch.jit.is_scripting() before + # torch.onnx.is_in_onnx_export(); + # otherwise, it will cause errors for torch.jit.script(). + # + # torch.logaddexp() works for both torch.jit.script() and + # torch.jit.trace() but it causes errors for ONNX export. + # + if torch.jit.is_scripting(): + # Note: We cannot use torch.jit.is_tracing() here as it also + # matches torch.onnx.export(). + return torch.logaddexp(x, y) + elif torch.onnx.is_in_onnx_export(): + return logaddexp_onnx(x, y) + else: + # for torch.jit.trace() + return torch.logaddexp(x, y) + + +class PiecewiseLinear(object): + """ + Piecewise linear function, from float to float, specified as nonempty list of (x,y) + pairs with the x values in order. x values <[initial x] or >[final x] are map to + [initial y], [final y] respectively. + """ + + def __init__(self, *args): + assert len(args) >= 1, len(args) + if len(args) == 1 and isinstance(args[0], PiecewiseLinear): + self.pairs = list(args[0].pairs) + else: + self.pairs = [(float(x), float(y)) for x, y in args] + for x, y in self.pairs: + assert isinstance(x, (float, int)), type(x) + assert isinstance(y, (float, int)), type(y) + + for i in range(len(self.pairs) - 1): + assert self.pairs[i + 1][0] > self.pairs[i][0], ( + i, + self.pairs[i], + self.pairs[i + 1], + ) + + def __str__(self): + # e.g. 'PiecewiseLinear((0., 10.), (100., 0.))' + return f"PiecewiseLinear({str(self.pairs)[1:-1]})" + + def __call__(self, x): + if x <= self.pairs[0][0]: + return self.pairs[0][1] + elif x >= self.pairs[-1][0]: + return self.pairs[-1][1] + else: + cur_x, cur_y = self.pairs[0] + for i in range(1, len(self.pairs)): + next_x, next_y = self.pairs[i] + if x >= cur_x and x <= next_x: + return cur_y + (next_y - cur_y) * (x - cur_x) / (next_x - cur_x) + cur_x, cur_y = next_x, next_y + assert False + + def __mul__(self, alpha): + return PiecewiseLinear(*[(x, y * alpha) for x, y in self.pairs]) + + def __add__(self, x): + if isinstance(x, (float, int)): + return PiecewiseLinear(*[(p[0], p[1] + x) for p in self.pairs]) + s, x = self.get_common_basis(x) + return PiecewiseLinear( + *[(sp[0], sp[1] + xp[1]) for sp, xp in zip(s.pairs, x.pairs)] + ) + + def max(self, x): + if isinstance(x, (float, int)): + x = PiecewiseLinear((0, x)) + s, x = self.get_common_basis(x, include_crossings=True) + return PiecewiseLinear( + *[(sp[0], max(sp[1], xp[1])) for sp, xp in zip(s.pairs, x.pairs)] + ) + + def min(self, x): + if isinstance(x, float) or isinstance(x, int): + x = PiecewiseLinear((0, x)) + s, x = self.get_common_basis(x, include_crossings=True) + return PiecewiseLinear( + *[(sp[0], min(sp[1], xp[1])) for sp, xp in zip(s.pairs, x.pairs)] + ) + + def __eq__(self, other): + return self.pairs == other.pairs + + def get_common_basis(self, p: "PiecewiseLinear", include_crossings: bool = False): + """ + Returns (self_mod, p_mod) which are equivalent piecewise linear + functions to self and p, but with the same x values. + + p: the other piecewise linear function + include_crossings: if true, include in the x values positions + where the functions indicate by this and p crosss. + """ + assert isinstance(p, PiecewiseLinear), type(p) + + # get sorted x-values without repetition. + x_vals = sorted(set([x for x, _ in self.pairs] + [x for x, _ in p.pairs])) + y_vals1 = [self(x) for x in x_vals] + y_vals2 = [p(x) for x in x_vals] + + if include_crossings: + extra_x_vals = [] + for i in range(len(x_vals) - 1): + if (y_vals1[i] > y_vals2[i]) != (y_vals1[i + 1] > y_vals2[i + 1]): + # if the two lines in this subsegment potentially cross each other.. + diff_cur = abs(y_vals1[i] - y_vals2[i]) + diff_next = abs(y_vals1[i + 1] - y_vals2[i + 1]) + # `pos`, between 0 and 1, gives the relative x position, + # with 0 being x_vals[i] and 1 being x_vals[i+1]. + pos = diff_cur / (diff_cur + diff_next) + extra_x_val = x_vals[i] + pos * (x_vals[i + 1] - x_vals[i]) + extra_x_vals.append(extra_x_val) + if len(extra_x_vals) > 0: + x_vals = sorted(set(x_vals + extra_x_vals)) + y_vals1 = [self(x) for x in x_vals] + y_vals2 = [p(x) for x in x_vals] + return ( + PiecewiseLinear(*zip(x_vals, y_vals1)), + PiecewiseLinear(*zip(x_vals, y_vals2)), + ) + + +class ScheduledFloat(torch.nn.Module): + """ + This object is a torch.nn.Module only because we want it to show up in + [top_level module].modules(); it does not have a working forward() function. + You are supposed to cast it to float, as in, float(parent_module.whatever), and use + it as something like a dropout prob. + + It is a floating point value whose value changes depending on the batch count of the + training loop. It is a piecewise linear function where you specify the (x,y) pairs + in sorted order on x; x corresponds to the batch index. For batch-index values + before the first x or after the last x, we just use the first or last y value. + + Example: + self.dropout = ScheduledFloat((0.0, 0.2), (4000.0, 0.0), default=0.0) + + `default` is used when self.batch_count is not set or not in training mode or in + torch.jit scripting mode. + """ + + def __init__(self, *args, default: float = 0.0): + super().__init__() + # self.batch_count and self.name will be written to in the training loop. + self.batch_count = None + self.name = None + self.default = default + self.schedule = PiecewiseLinear(*args) + + def extra_repr(self) -> str: + return ( + f"batch_count={self.batch_count}, schedule={str(self.schedule.pairs[1:-1])}" + ) + + def __float__(self): + batch_count = self.batch_count + if ( + batch_count is None + or not self.training + or torch.jit.is_scripting() + or torch.jit.is_tracing() + ): + return float(self.default) + else: + ans = self.schedule(self.batch_count) + if random.random() < 0.0002: + logging.debug( + f"ScheduledFloat: name={self.name}, " + f"batch_count={self.batch_count}, ans={ans}" + ) + return ans + + def __add__(self, x): + if isinstance(x, float) or isinstance(x, int): + return ScheduledFloat(self.schedule + x, default=self.default) + else: + return ScheduledFloat( + self.schedule + x.schedule, default=self.default + x.default + ) + + def max(self, x): + if isinstance(x, float) or isinstance(x, int): + return ScheduledFloat(self.schedule.max(x), default=self.default) + else: + return ScheduledFloat( + self.schedule.max(x.schedule), + default=max(self.default, x.default), + ) + + +FloatLike = Union[float, ScheduledFloat] + + +class CutoffEstimator: + """ + Estimates cutoffs of an arbitrary numerical quantity such that a specified + proportion of items will be above the cutoff on average. + + p is the proportion of items that should be above the cutoff. + """ + + def __init__(self, p: float): + self.p = p + # total count of items + self.count = 0 + # total count of items that were above the cutoff + self.count_above = 0 + # initial cutoff value + self.cutoff = 0 + + def __call__(self, x: float) -> bool: + """ + Returns true if x is above the cutoff. + """ + ans = x > self.cutoff + self.count += 1 + if ans: + self.count_above += 1 + cur_p = self.count_above / self.count + delta_p = cur_p - self.p + if (delta_p > 0) == ans: + q = abs(delta_p) + self.cutoff = x * q + self.cutoff * (1 - q) + return ans + + +class SoftmaxFunction(torch.autograd.Function): + """ + Tries to handle half-precision derivatives in a randomized way that should + be more accurate for training than the default behavior. + """ + + @staticmethod + def forward(ctx, x: Tensor, dim: int): + ans = x.softmax(dim=dim) + # if x dtype is float16, x.softmax() returns a float32 because + # (presumably) that op does not support float16, and autocast + # is enabled. + if torch.is_autocast_enabled(): + ans = ans.to(torch.float16) + ctx.save_for_backward(ans) + ctx.x_dtype = x.dtype + ctx.dim = dim + return ans + + @staticmethod + def backward(ctx, ans_grad: Tensor): + (ans,) = ctx.saved_tensors + with torch.amp.autocast("cuda", enabled=False): + ans_grad = ans_grad.to(torch.float32) + ans = ans.to(torch.float32) + x_grad = ans_grad * ans + x_grad = x_grad - ans * x_grad.sum(dim=ctx.dim, keepdim=True) + return x_grad, None + + +def softmax(x: Tensor, dim: int): + if not x.requires_grad or torch.jit.is_scripting() or torch.jit.is_tracing(): + return x.softmax(dim=dim) + + return SoftmaxFunction.apply(x, dim) + + +class BiasNormFunction(torch.autograd.Function): + # This computes: + # scales = (torch.mean((x - bias) ** 2, keepdim=True)) ** -0.5 * log_scale.exp() + # return x * scales + # (after unsqueezing the bias), but it does it in a memory-efficient way so that + # it can just store the returned value (chances are, this will also be needed for + # some other reason, related to the next operation, so we can save memory). + @staticmethod + def forward( + ctx, + x: Tensor, + bias: Tensor, + log_scale: Tensor, + channel_dim: int, + store_output_for_backprop: bool, + ) -> Tensor: + assert bias.ndim == 1 + if channel_dim < 0: + channel_dim = channel_dim + x.ndim + ctx.store_output_for_backprop = store_output_for_backprop + ctx.channel_dim = channel_dim + for _ in range(channel_dim + 1, x.ndim): + bias = bias.unsqueeze(-1) + scales = ( + torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5 + ) * log_scale.exp() + ans = x * scales + ctx.save_for_backward( + ans.detach() if store_output_for_backprop else x, + scales.detach(), + bias.detach(), + log_scale.detach(), + ) + return ans + + @staticmethod + def backward(ctx, ans_grad: Tensor) -> Tensor: + ans_or_x, scales, bias, log_scale = ctx.saved_tensors + if ctx.store_output_for_backprop: + x = ans_or_x / scales + else: + x = ans_or_x + x = x.detach() + x.requires_grad = True + bias.requires_grad = True + log_scale.requires_grad = True + with torch.enable_grad(): + # recompute scales from x, bias and log_scale. + scales = ( + torch.mean((x - bias) ** 2, dim=ctx.channel_dim, keepdim=True) ** -0.5 + ) * log_scale.exp() + ans = x * scales + ans.backward(gradient=ans_grad) + return x.grad, bias.grad.flatten(), log_scale.grad, None, None + + +class BiasNorm(torch.nn.Module): + """ + This is intended to be a simpler, and hopefully cheaper, replacement for + LayerNorm. The observation this is based on, is that Transformer-type + networks, especially with pre-norm, sometimes seem to set one of the + feature dimensions to a large constant value (e.g. 50), which "defeats" + the LayerNorm because the output magnitude is then not strongly dependent + on the other (useful) features. Presumably the weight and bias of the + LayerNorm are required to allow it to do this. + + Instead, we give the BiasNorm a trainable bias that it can use when + computing the scale for normalization. We also give it a (scalar) + trainable scale on the output. + + + Args: + num_channels: the number of channels, e.g. 512. + channel_dim: the axis/dimension corresponding to the channel, + interpreted as an offset from the input's ndim if negative. + This is NOT the num_channels; it should typically be one of + {-2, -1, 0, 1, 2, 3}. + log_scale: the initial log-scale that we multiply the output by; this + is learnable. + log_scale_min: FloatLike, minimum allowed value of log_scale + log_scale_max: FloatLike, maximum allowed value of log_scale + store_output_for_backprop: only possibly affects memory use; recommend + to set to True if you think the output of this module is more likely + than the input of this module to be required to be stored for the + backprop. + """ + + def __init__( + self, + num_channels: int, + channel_dim: int = -1, # CAUTION: see documentation. + log_scale: float = 1.0, + log_scale_min: float = -1.5, + log_scale_max: float = 1.5, + store_output_for_backprop: bool = False, + ) -> None: + super(BiasNorm, self).__init__() + self.num_channels = num_channels + self.channel_dim = channel_dim + self.log_scale = nn.Parameter(torch.tensor(log_scale)) + self.bias = nn.Parameter(torch.zeros(num_channels)) + + self.log_scale_min = log_scale_min + self.log_scale_max = log_scale_max + + self.store_output_for_backprop = store_output_for_backprop + + def forward(self, x: Tensor) -> Tensor: + assert x.shape[self.channel_dim] == self.num_channels + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + channel_dim = self.channel_dim + if channel_dim < 0: + channel_dim += x.ndim + bias = self.bias + for _ in range(channel_dim + 1, x.ndim): + bias = bias.unsqueeze(-1) + scales = ( + torch.mean((x - bias) ** 2, dim=channel_dim, keepdim=True) ** -0.5 + ) * self.log_scale.exp() + return x * scales + + log_scale = limit_param_value( + self.log_scale, + min=float(self.log_scale_min), + max=float(self.log_scale_max), + training=self.training, + ) + + return BiasNormFunction.apply( + x, + self.bias, + log_scale, + self.channel_dim, + self.store_output_for_backprop, + ) + + +def ScaledLinear(*args, initial_scale: float = 1.0, **kwargs) -> nn.Linear: + """ + Behaves like a constructor of a modified version of nn.Linear + that gives an easy way to set the default initial parameter scale. + + Args: + Accepts the standard args and kwargs that nn.Linear accepts + e.g. in_features, out_features, bias=False. + + initial_scale: you can override this if you want to increase + or decrease the initial magnitude of the module's output + (affects the initialization of weight_scale and bias_scale). + Another option, if you want to do something like this, is + to re-initialize the parameters. + """ + ans = nn.Linear(*args, **kwargs) + with torch.no_grad(): + ans.weight[:] *= initial_scale + if ans.bias is not None: + torch.nn.init.uniform_(ans.bias, -0.1 * initial_scale, 0.1 * initial_scale) + return ans + + +class BalancerFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + x: Tensor, + min_mean: float, + max_mean: float, + min_rms: float, + max_rms: float, + grad_scale: float, + channel_dim: int, + ) -> Tensor: + if channel_dim < 0: + channel_dim += x.ndim + ctx.channel_dim = channel_dim + ctx.save_for_backward(x) + ctx.config = ( + min_mean, + max_mean, + min_rms, + max_rms, + grad_scale, + channel_dim, + ) + return x + + @staticmethod + def backward(ctx, x_grad: Tensor) -> Tuple[Tensor, None, None, None, None, None]: + (x,) = ctx.saved_tensors + ( + min_mean, + max_mean, + min_rms, + max_rms, + grad_scale, + channel_dim, + ) = ctx.config + + try: + with torch.enable_grad(): + with torch.amp.autocast("cuda", enabled=False): + x = x.to(torch.float32) + x = x.detach() + x.requires_grad = True + mean_dims = [i for i in range(x.ndim) if i != channel_dim] + uncentered_var = (x**2).mean(dim=mean_dims, keepdim=True) + mean = x.mean(dim=mean_dims, keepdim=True) + stddev = (uncentered_var - (mean * mean)).clamp(min=1.0e-20).sqrt() + rms = uncentered_var.clamp(min=1.0e-20).sqrt() + + m = mean / stddev + # part of loss that relates to mean / stddev + m_loss = (m - m.clamp(min=min_mean, max=max_mean)).abs() + + # put a much larger scale on the RMS-max-limit loss, so that if both + # it and the m_loss are violated we fix the RMS loss first. + rms_clamped = rms.clamp(min=min_rms, max=max_rms) + r_loss = (rms_clamped / rms).log().abs() + + loss = m_loss + r_loss + + loss.backward(gradient=torch.ones_like(loss)) + loss_grad = x.grad + loss_grad_rms = ( + (loss_grad**2) + .mean(dim=mean_dims, keepdim=True) + .sqrt() + .clamp(min=1.0e-20) + ) + + loss_grad = loss_grad * (grad_scale / loss_grad_rms) + + x_grad_float = x_grad.to(torch.float32) + # scale each element of loss_grad by the absolute value of the + # corresponding element of x_grad, which we view as a noisy estimate + # of its magnitude for that (frame and dimension). later we can + # consider factored versions. + x_grad_mod = x_grad_float + (x_grad_float.abs() * loss_grad) + x_grad = x_grad_mod.to(x_grad.dtype) + except Exception as e: + logging.info( + f"Caught exception in Balancer backward: {e}, " + f"size={list(x_grad.shape)}, will continue." + ) + + return x_grad, None, None, None, None, None, None + + +class Balancer(torch.nn.Module): + """ + Modifies the backpropped derivatives of a function to try to encourage, for + each channel, that it is positive at least a proportion `threshold` of the + time. It does this by multiplying negative derivative values by up to + (1+max_factor), and positive derivative values by up to (1-max_factor), + interpolated from 1 at the threshold to those extremal values when none + of the inputs are positive. + + Args: + num_channels: the number of channels + channel_dim: the dimension/axis corresponding to the channel, e.g. + -1, 0, 1, 2; will be interpreted as an offset from x.ndim if negative. + min_positive: the minimum, per channel, of the proportion of the time + that (x > 0), below which we start to modify the derivatives. + max_positive: the maximum, per channel, of the proportion of the time + that (x > 0), above which we start to modify the derivatives. + scale_gain_factor: determines the 'gain' with which we increase the + change in gradient once the constraints on min_abs and max_abs + are violated. + min_abs: the minimum average-absolute-value difference from the mean + value per channel, which we allow, before we start to modify + the derivatives to prevent this. + max_abs: the maximum average-absolute-value difference from the mean + value per channel, which we allow, before we start to modify + the derivatives to prevent this. + prob: determines the minimum probability with which we modify the + gradients for the {min,max}_positive and {min,max}_abs constraints, + on each forward(). This is done randomly to prevent all layers + from doing it at the same time. + """ + + def __init__( + self, + num_channels: int, + channel_dim: int, + min_positive: FloatLike = 0.05, + max_positive: FloatLike = 0.95, + min_abs: FloatLike = 0.2, + max_abs: FloatLike = 100.0, + grad_scale: FloatLike = 0.04, + prob: Optional[FloatLike] = None, + ): + super().__init__() + + if prob is None: + prob = ScheduledFloat((0.0, 0.5), (8000.0, 0.125), default=0.4) + self.prob = prob + # 5% of the time we will return and do nothing because memory usage is + # too high. + self.mem_cutoff = CutoffEstimator(0.05) + + # actually self.num_channels is no longer needed except for an assertion. + self.num_channels = num_channels + self.channel_dim = channel_dim + self.min_positive = min_positive + self.max_positive = max_positive + self.min_abs = min_abs + self.max_abs = max_abs + self.grad_scale = grad_scale + + def forward(self, x: Tensor) -> Tensor: + if ( + torch.jit.is_scripting() + or not x.requires_grad + or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated())) + ): + return _no_op(x) + + prob = float(self.prob) + if random.random() < prob: + # The following inner-functions convert from the way we historically + # specified these limitations, as limits on the absolute value and the + # proportion of positive values, to limits on the RMS value and + # the (mean / stddev). + def _abs_to_rms(x): + # for normally distributed data, if the expected absolute value is x, + # the expected rms value will be sqrt(pi/2) * x. + return 1.25331413732 * x + + def _proportion_positive_to_mean(x): + def _atanh(x): + eps = 1.0e-10 + # eps is to prevent crashes if x is exactly 0 or 1. + # we'll just end up returning a fairly large value. + return (math.log(1 + x + eps) - math.log(1 - x + eps)) / 2.0 + + def _approx_inverse_erf(x): + # 1 / (sqrt(pi) * ln(2)), + # see https://math.stackexchange.com/questions/321569/ + # approximating-the-error-function-erf-by-analytical-functions + # this approximation is extremely crude and gets progressively worse + # for x very close to -1 or +1, but we mostly care about the + # "middle" region + # e.g. _approx_inverse_erf(0.05) = 0.0407316414078772, + # and math.erf(0.0407316414078772) = 0.045935330944660666, + # which is pretty close to 0.05. + return 0.8139535143 * _atanh(x) + + # first convert x from the range 0..1 to the range -1..1 which the error + # function returns + x = -1 + (2 * x) + return _approx_inverse_erf(x) + + min_mean = _proportion_positive_to_mean(float(self.min_positive)) + max_mean = _proportion_positive_to_mean(float(self.max_positive)) + min_rms = _abs_to_rms(float(self.min_abs)) + max_rms = _abs_to_rms(float(self.max_abs)) + grad_scale = float(self.grad_scale) + + assert x.shape[self.channel_dim] == self.num_channels + + return BalancerFunction.apply( + x, + min_mean, + max_mean, + min_rms, + max_rms, + grad_scale, + self.channel_dim, + ) + else: + return _no_op(x) + + +def penalize_abs_values_gt( + x: Tensor, limit: float, penalty: float, name: str = None +) -> Tensor: + """ + Returns x unmodified, but in backprop will put a penalty for the excess of + the absolute values of elements of x over the limit "limit". E.g. if + limit == 10.0, then if x has any values over 10 it will get a penalty. + + Caution: the value of this penalty will be affected by grad scaling used + in automatic mixed precision training. For this reasons we use this, + it shouldn't really matter, or may even be helpful; we just use this + to disallow really implausible values of scores to be given to softmax. + + The name is for randomly printed debug info. + """ + x_sign = x.sign() + over_limit = (x.abs() - limit) > 0 + # The following is a memory efficient way to penalize the absolute values of + # x that's over the limit. (The memory efficiency comes when you think + # about which items torch needs to cache for the autograd, and which ones it + # can throw away). The numerical value of aux_loss as computed here will + # actually be larger than it should be, by limit * over_limit.sum(), but it + # has the same derivative as the real aux_loss which is penalty * (x.abs() - + # limit).relu(). + aux_loss = penalty * ((x_sign * over_limit).to(torch.int8) * x) + # note: we don't do sum() here on aux)_loss, but it's as if we had done + # sum() due to how with_loss() works. + x = with_loss(x, aux_loss, name) + # you must use x for something, or this will be ineffective. + return x + + +def _diag(x: Tensor): # like .diag(), but works for tensors with 3 dims. + if x.ndim == 2: + return x.diag() + else: + (batch, dim, dim) = x.shape + x = x.reshape(batch, dim * dim) + x = x[:, :: dim + 1] + assert x.shape == (batch, dim) + return x + + +def _whitening_metric(x: Tensor, num_groups: int): + """ + Computes the "whitening metric", a value which will be 1.0 if all the eigenvalues of + of the centered feature covariance are the same within each group's covariance + matrix and also between groups. + Args: + x: a Tensor of shape (*, num_channels) + num_groups: the number of groups of channels, a number >=1 that divides + num_channels + Returns: + Returns a scalar Tensor that will be 1.0 if the data is "perfectly white" and + greater than 1.0 otherwise. + """ + assert x.dtype != torch.float16 + x = x.reshape(-1, x.shape[-1]) + (num_frames, num_channels) = x.shape + assert num_channels % num_groups == 0 + channels_per_group = num_channels // num_groups + x = x.reshape(num_frames, num_groups, channels_per_group).transpose(0, 1) + # x now has shape (num_groups, num_frames, channels_per_group) + # subtract the mean so we use the centered, not uncentered, covariance. + # My experience has been that when we "mess with the gradients" like this, + # it's better not do anything that tries to move the mean around, because + # that can easily cause instability. + x = x - x.mean(dim=1, keepdim=True) + # x_covar: (num_groups, channels_per_group, channels_per_group) + x_covar = torch.matmul(x.transpose(1, 2), x) + x_covar_mean_diag = _diag(x_covar).mean() + # the following expression is what we'd get if we took the matrix product + # of each covariance and measured the mean of its trace, i.e. + # the same as _diag(torch.matmul(x_covar, x_covar)).mean(). + x_covarsq_mean_diag = (x_covar**2).sum() / (num_groups * channels_per_group) + # this metric will be >= 1.0; the larger it is, the less 'white' the data was. + metric = x_covarsq_mean_diag / (x_covar_mean_diag**2 + 1.0e-20) + return metric + + +class WhiteningPenaltyFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, module: nn.Module) -> Tensor: + ctx.save_for_backward(x) + ctx.module = module + return x + + @staticmethod + def backward(ctx, x_grad: Tensor): + (x_orig,) = ctx.saved_tensors + w = ctx.module + + try: + with torch.enable_grad(): + with torch.amp.autocast("cuda", enabled=False): + x_detached = x_orig.to(torch.float32).detach() + x_detached.requires_grad = True + + metric = _whitening_metric(x_detached, w.num_groups) + + if random.random() < 0.005 or __name__ == "__main__": + logging.debug( + f"Whitening: name={w.name}, num_groups={w.num_groups}," + f"num_channels={x_orig.shape[-1]}, " + f"metric={metric.item():.2f}" + f" vs. limit={float(w.whitening_limit)}" + ) + + if metric < float(w.whitening_limit): + w.prob = w.min_prob + return x_grad, None + else: + w.prob = w.max_prob + metric.backward() + penalty_grad = x_detached.grad + scale = w.grad_scale * ( + x_grad.to(torch.float32).norm() + / (penalty_grad.norm() + 1.0e-20) + ) + penalty_grad = penalty_grad * scale + return x_grad + penalty_grad.to(x_grad.dtype), None + except Exception as e: + logging.info( + f"Caught exception in Whiten backward: {e}, " + f"size={list(x_grad.shape)}, will continue." + ) + return x_grad, None + + +class Whiten(nn.Module): + def __init__( + self, + num_groups: int, + whitening_limit: FloatLike, + prob: Union[float, Tuple[float, float]], + grad_scale: FloatLike, + ): + """ + Args: + num_groups: the number of groups to divide the channel dim into before + whitening. We will attempt to make the feature covariance + within each group, after mean subtraction, as "white" as possible, + while having the same trace across all groups. + whitening_limit: a value greater than 1.0, that dictates how much + freedom we have to violate the constraints. 1.0 would mean perfectly + white, with exactly the same trace across groups; larger values + give more freedom. E.g. 2.0. + prob: the probability with which we apply the gradient modification + (also affects the grad scale). May be supplied as a float, + or as a pair (min_prob, max_prob) + + grad_scale: determines the scale on the gradient term from this object, + relative to the rest of the gradient on the attention weights. + E.g. 0.02 (you may want to use smaller values than this if prob is large) + """ + super(Whiten, self).__init__() + assert num_groups >= 1 + assert float(whitening_limit) >= 1 + assert grad_scale >= 0 + self.num_groups = num_groups + self.whitening_limit = whitening_limit + self.grad_scale = grad_scale + + if isinstance(prob, float): + prob = (prob, prob) + (self.min_prob, self.max_prob) = prob + assert 0 < self.min_prob <= self.max_prob <= 1 + self.prob = self.max_prob + self.name = None # will be set in training loop + + def forward(self, x: Tensor) -> Tensor: + """ + In the forward pass, this function just returns the input unmodified. + In the backward pass, it will modify the gradients to ensure that the + distribution in each group has close to (lambda times I) as the covariance + after mean subtraction, with the same lambda across groups. + For whitening_limit > 1, there will be more freedom to violate this + constraint. + + Args: + x: the input of shape (*, num_channels) + + Returns: + x, unmodified. You should make sure + you use the returned value, or the graph will be freed + and nothing will happen in backprop. + """ + grad_scale = float(self.grad_scale) + if not x.requires_grad or random.random() > self.prob or grad_scale == 0: + return _no_op(x) + else: + return WhiteningPenaltyFunction.apply(x, self) + + +class WithLoss(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, y: Tensor, name: str): + ctx.y_shape = y.shape + if random.random() < 0.002 and name is not None: + loss_sum = y.sum().item() + logging.debug(f"WithLoss: name={name}, loss-sum={loss_sum:.3e}") + return x + + @staticmethod + def backward(ctx, ans_grad: Tensor): + return ( + ans_grad, + torch.ones(ctx.y_shape, dtype=ans_grad.dtype, device=ans_grad.device), + None, + ) + + +def with_loss(x, y, name): + # returns x but adds y.sum() to the loss function. + return WithLoss.apply(x, y, name) + + +class LimitParamValue(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, min: float, max: float): + ctx.save_for_backward(x) + assert max >= min + ctx.min = min + ctx.max = max + return x + + @staticmethod + def backward(ctx, x_grad: Tensor): + (x,) = ctx.saved_tensors + # where x < ctx.min, ensure all grads are negative (this will tend to make + # x more positive). + x_grad = x_grad * torch.where( + torch.logical_and(x_grad > 0, x < ctx.min), -1.0, 1.0 + ) + # where x > ctx.max, ensure all grads are positive (this will tend to make + # x more negative). + x_grad *= torch.where(torch.logical_and(x_grad < 0, x > ctx.max), -1.0, 1.0) + return x_grad, None, None + + +def limit_param_value( + x: Tensor, min: float, max: float, prob: float = 0.6, training: bool = True +): + # You apply this to (typically) an nn.Parameter during training to ensure that its + # (elements mostly) stays within a supplied range. This is done by modifying the + # gradients in backprop. + # It's not necessary to do this on every batch: do it only some of the time, + # to save a little time. + if training and random.random() < prob: + return LimitParamValue.apply(x, min, max) + else: + return x + + +def _no_op(x: Tensor) -> Tensor: + if torch.jit.is_scripting() or torch.jit.is_tracing(): + return x + else: + # a no-op function that will have a node in the autograd graph, + # to avoid certain bugs relating to backward hooks + return x.chunk(1, dim=-1)[0] + + +# Identity more friendly to backward hooks than nn.Identity(), for diagnostic reasons. +class Identity(torch.nn.Module): + def __init__(self): + super(Identity, self).__init__() + + def forward(self, x): + return _no_op(x) + + +# Dropout2 is just like normal dropout, except it supports schedules +# on the dropout rates. +class Dropout2(nn.Module): + def __init__(self, p: FloatLike): + super().__init__() + self.p = p + + def forward(self, x: Tensor) -> Tensor: + return torch.nn.functional.dropout(x, p=float(self.p), training=self.training) + + +class MulForDropout3(torch.autograd.Function): + # returns (x * y * alpha) where alpha is a float and y doesn't require + # grad and is zero-or-one. + @staticmethod + @custom_fwd + def forward(ctx, x, y, alpha): + assert not y.requires_grad + ans = x * y * alpha + ctx.save_for_backward(ans) + ctx.alpha = alpha + return ans + + @staticmethod + @custom_bwd + def backward(ctx, ans_grad): + (ans,) = ctx.saved_tensors + x_grad = ctx.alpha * ans_grad * (ans != 0) + return x_grad, None, None + + +# Dropout3 is just like normal dropout, except it supports schedules on the dropout +# rates, and it lets you choose one dimension to share the dropout mask over +class Dropout3(nn.Module): + def __init__(self, p: FloatLike, shared_dim: int): + super().__init__() + self.p = p + self.shared_dim = shared_dim + + def forward(self, x: Tensor) -> Tensor: + p = float(self.p) + if not self.training or p == 0: + return _no_op(x) + scale = 1.0 / (1 - p) + rand_shape = list(x.shape) + rand_shape[self.shared_dim] = 1 + mask = torch.rand(*rand_shape, device=x.device) > p + ans = MulForDropout3.apply(x, mask, scale) + return ans + + +class SwooshLFunction(torch.autograd.Function): + """ + swoosh_l(x) = log(1 + exp(x-4)) - 0.08*x - 0.035 + """ + + @staticmethod + def forward(ctx, x: Tensor) -> Tensor: + requires_grad = x.requires_grad + if x.dtype == torch.float16: + x = x.to(torch.float32) + + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + + coeff = -0.08 + + with torch.amp.autocast("cuda", enabled=False): + with torch.enable_grad(): + x = x.detach() + x.requires_grad = True + y = torch.logaddexp(zero, x - 4.0) + coeff * x - 0.035 + + if not requires_grad: + return y + + y.backward(gradient=torch.ones_like(y)) + + grad = x.grad + floor = coeff + ceil = 1.0 + coeff + 0.005 + + d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like( + grad + ) + if __name__ == "__main__": + # for self-testing only. + assert d_scaled.min() >= 0.0 + assert d_scaled.max() < 256.0 + + d_int = d_scaled.to(torch.uint8) + ctx.save_for_backward(d_int) + if x.dtype == torch.float16 or torch.is_autocast_enabled(): + y = y.to(torch.float16) + return y + + @staticmethod + def backward(ctx, y_grad: Tensor) -> Tensor: + (d,) = ctx.saved_tensors + # the same constants as used in forward pass. + coeff = -0.08 + floor = coeff + ceil = 1.0 + coeff + 0.005 + d = d * ((ceil - floor) / 255.0) + floor + return y_grad * d + + +class SwooshL(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + """Return Swoosh-L activation.""" + if torch.jit.is_scripting() or torch.jit.is_tracing(): + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + return logaddexp(zero, x - 4.0) - 0.08 * x - 0.035 + elif "k2" not in sys.modules: + return SwooshLFunction.apply(x) + else: + if not x.requires_grad: + return k2.swoosh_l_forward(x) + else: + return k2.swoosh_l(x) + + +class SwooshLOnnx(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + """Return Swoosh-L activation.""" + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + return logaddexp_onnx(zero, x - 4.0) - 0.08 * x - 0.035 + + +class SwooshRFunction(torch.autograd.Function): + """ + swoosh_r(x) = log(1 + exp(x-1)) - 0.08*x - 0.313261687 + + derivatives are between -0.08 and 0.92. + """ + + @staticmethod + def forward(ctx, x: Tensor) -> Tensor: + requires_grad = x.requires_grad + + if x.dtype == torch.float16: + x = x.to(torch.float32) + + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + + with torch.amp.autocast("cuda", enabled=False): + with torch.enable_grad(): + x = x.detach() + x.requires_grad = True + y = torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687 + + if not requires_grad: + return y + y.backward(gradient=torch.ones_like(y)) + + grad = x.grad + floor = -0.08 + ceil = 0.925 + + d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like( + grad + ) + if __name__ == "__main__": + # for self-testing only. + assert d_scaled.min() >= 0.0 + assert d_scaled.max() < 256.0 + + d_int = d_scaled.to(torch.uint8) + ctx.save_for_backward(d_int) + if x.dtype == torch.float16 or torch.is_autocast_enabled(): + y = y.to(torch.float16) + return y + + @staticmethod + def backward(ctx, y_grad: Tensor) -> Tensor: + (d,) = ctx.saved_tensors + # the same constants as used in forward pass. + floor = -0.08 + ceil = 0.925 + d = d * ((ceil - floor) / 255.0) + floor + return y_grad * d + + +class SwooshR(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + """Return Swoosh-R activation.""" + if torch.jit.is_scripting() or torch.jit.is_tracing(): + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + return logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687 + elif "k2" not in sys.modules: + return SwooshRFunction.apply(x) + else: + if not x.requires_grad: + return k2.swoosh_r_forward(x) + else: + return k2.swoosh_r(x) + + +class SwooshROnnx(torch.nn.Module): + def forward(self, x: Tensor) -> Tensor: + """Return Swoosh-R activation.""" + zero = torch.tensor(0.0, dtype=x.dtype, device=x.device) + return logaddexp_onnx(zero, x - 1.0) - 0.08 * x - 0.313261687 + + +# simple version of SwooshL that does not redefine the backprop, used in +# ActivationDropoutAndLinearFunction. +def SwooshLForward(x: Tensor): + with torch.amp.autocast("cuda", enabled=False): + x = x.to(torch.float32) + x_offset = x - 4.0 + log_sum = (1.0 + x_offset.exp()).log().to(x.dtype) + log_sum = torch.where(log_sum == float("inf"), x_offset, log_sum) + return log_sum - 0.08 * x - 0.035 + + +# simple version of SwooshR that does not redefine the backprop, used in +# ActivationDropoutAndLinearFunction. +def SwooshRForward(x: Tensor): + with torch.amp.autocast("cuda", enabled=False): + x = x.to(torch.float32) + x_offset = x - 1.0 + log_sum = (1.0 + x_offset.exp()).log().to(x.dtype) + log_sum = torch.where(log_sum == float("inf"), x_offset, log_sum) + return log_sum - 0.08 * x - 0.313261687 + + +class ActivationDropoutAndLinearFunction(torch.autograd.Function): + @staticmethod + @custom_fwd + def forward( + ctx, + x: Tensor, + weight: Tensor, + bias: Optional[Tensor], + activation: str, + dropout_p: float, + dropout_shared_dim: Optional[int], + ): + if dropout_p != 0.0: + dropout_shape = list(x.shape) + if dropout_shared_dim is not None: + dropout_shape[dropout_shared_dim] = 1 + # else it won't be very memory efficient. + dropout_mask = (1.0 / (1.0 - dropout_p)) * ( + torch.rand(*dropout_shape, device=x.device, dtype=x.dtype) > dropout_p + ) + else: + dropout_mask = None + + ctx.save_for_backward(x, weight, bias, dropout_mask) + + ctx.activation = activation + + forward_activation_dict = { + "SwooshL": k2.swoosh_l_forward, + "SwooshR": k2.swoosh_r_forward, + } + # it will raise a KeyError if this fails. This will be an error. We let it + # propagate to the user. + activation_func = forward_activation_dict[activation] + x = activation_func(x) + if dropout_mask is not None: + x = x * dropout_mask + x = torch.nn.functional.linear(x, weight, bias) + return x + + @staticmethod + @custom_bwd + def backward(ctx, ans_grad: Tensor): + saved = ctx.saved_tensors + (x, weight, bias, dropout_mask) = saved + + forward_and_deriv_activation_dict = { + "SwooshL": k2.swoosh_l_forward_and_deriv, + "SwooshR": k2.swoosh_r_forward_and_deriv, + } + # the following lines a KeyError if the activation is unrecognized. + # This will be an error. We let it propagate to the user. + func = forward_and_deriv_activation_dict[ctx.activation] + + y, func_deriv = func(x) + if dropout_mask is not None: + y = y * dropout_mask + # now compute derivative of y w.r.t. weight and bias.. + # y: (..., in_channels), ans_grad: (..., out_channels), + (out_channels, in_channels) = weight.shape + + in_channels = y.shape[-1] + g = ans_grad.reshape(-1, out_channels) + weight_deriv = torch.matmul(g.t(), y.reshape(-1, in_channels)) + y_deriv = torch.matmul(ans_grad, weight) + bias_deriv = None if bias is None else g.sum(dim=0) + x_deriv = y_deriv * func_deriv + if dropout_mask is not None: + # order versus func_deriv does not matter + x_deriv = x_deriv * dropout_mask + + return x_deriv, weight_deriv, bias_deriv, None, None, None + + +class ActivationDropoutAndLinear(torch.nn.Module): + """ + This merges an activation function followed by dropout and then a nn.Linear module; + it does so in a memory efficient way so that it only stores the input to the whole + module. If activation == SwooshL and dropout_shared_dim != None, this will be + equivalent to: + nn.Sequential(SwooshL(), + Dropout3(dropout_p, shared_dim=dropout_shared_dim), + ScaledLinear(in_channels, out_channels, bias=bias, + initial_scale=initial_scale)) + If dropout_shared_dim is None, the dropout would be equivalent to + Dropout2(dropout_p). Note: Dropout3 will be more memory efficient as the dropout + mask is smaller. + + Args: + in_channels: number of input channels, e.g. 256 + out_channels: number of output channels, e.g. 256 + bias: if true, have a bias + activation: the activation function, for now just support SwooshL. + dropout_p: the dropout probability or schedule (happens after nonlinearity). + dropout_shared_dim: the dimension, if any, across which the dropout mask is + shared (e.g. the time dimension). If None, this may be less memory + efficient if there are modules before this one that cache the input + for their backprop (e.g. Balancer or Whiten). + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + bias: bool = True, + activation: str = "SwooshL", + dropout_p: FloatLike = 0.0, + dropout_shared_dim: Optional[int] = -1, + initial_scale: float = 1.0, + ): + super().__init__() + # create a temporary module of nn.Linear that we'll steal the + # weights and bias from + l = ScaledLinear( + in_channels, out_channels, bias=bias, initial_scale=initial_scale + ) + + self.weight = l.weight + # register_parameter properly handles making it a parameter when l.bias + # is None. I think there is some reason for doing it this way rather + # than just setting it to None but I don't know what it is, maybe + # something to do with exporting the module.. + self.register_parameter("bias", l.bias) + + self.activation = activation + self.dropout_p = dropout_p + self.dropout_shared_dim = dropout_shared_dim + + def forward(self, x: Tensor): + if ( + torch.jit.is_scripting() + or torch.jit.is_tracing() + or "k2" not in sys.modules + ): + if self.activation == "SwooshL": + x = SwooshLForward(x) + elif self.activation == "SwooshR": + x = SwooshRForward(x) + else: + assert False, self.activation + return torch.nn.functional.linear(x, self.weight, self.bias) + + return ActivationDropoutAndLinearFunction.apply( + x, + self.weight, + self.bias, + self.activation, + float(self.dropout_p), + self.dropout_shared_dim, + ) + + +def _test_whiten(): + for proportion in [0.1, 0.5, 10.0]: + logging.info(f"_test_whiten(): proportion = {proportion}") + x = torch.randn(100, 128) + direction = torch.randn(128) + coeffs = torch.randn(100, 1) + x += proportion * direction * coeffs + + x.requires_grad = True + + m = Whiten( + 1, 5.0, prob=1.0, grad_scale=0.1 # num_groups # whitening_limit, + ) # grad_scale + + for _ in range(4): + y = m(x) + + y_grad = torch.randn_like(x) + y.backward(gradient=y_grad) + + if proportion < 0.2: + assert torch.allclose(x.grad, y_grad) + elif proportion > 1.0: + assert not torch.allclose(x.grad, y_grad) + + +def _test_balancer_sign(): + probs = torch.arange(0, 1, 0.01) + N = 1000 + x = 1.0 * ((2.0 * (torch.rand(probs.numel(), N) < probs.unsqueeze(-1))) - 1.0) + x = x.detach() + x.requires_grad = True + m = Balancer( + probs.numel(), + channel_dim=0, + min_positive=0.05, + max_positive=0.95, + min_abs=0.0, + prob=1.0, + ) + + y_grad = torch.sign(torch.randn(probs.numel(), N)) + + y = m(x) + y.backward(gradient=y_grad) + print("_test_balancer_sign: x = ", x) + print("_test_balancer_sign: y grad = ", y_grad) + print("_test_balancer_sign: x grad = ", x.grad) + + +def _test_balancer_magnitude(): + magnitudes = torch.arange(0, 1, 0.01) + N = 1000 + x = torch.sign(torch.randn(magnitudes.numel(), N)) * magnitudes.unsqueeze(-1) + x = x.detach() + x.requires_grad = True + m = Balancer( + magnitudes.numel(), + channel_dim=0, + min_positive=0.0, + max_positive=1.0, + min_abs=0.2, + max_abs=0.7, + prob=1.0, + ) + + y_grad = torch.sign(torch.randn(magnitudes.numel(), N)) + + y = m(x) + y.backward(gradient=y_grad) + print("_test_balancer_magnitude: x = ", x) + print("_test_balancer_magnitude: y grad = ", y_grad) + print("_test_balancer_magnitude: x grad = ", x.grad) + + +def _test_swooshl_deriv(): + x = torch.randn(10, 12, dtype=torch.double) * 3.0 + x.requires_grad = True + m = SwooshL() + + tol = 1.0 / 255.0 + torch.autograd.gradcheck(m, x, atol=tol, eps=0.01) + + # for self-test. + x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 + x.requires_grad = True + y = m(x) + return y + + +def _test_swooshr_deriv(): + x = torch.randn(10, 12, dtype=torch.double) * 3.0 + x.requires_grad = True + m = SwooshR() + + tol = 1.0 / 255.0 + torch.autograd.gradcheck(m, x, atol=tol, eps=0.01) + + # for self-test. + x = torch.randn(1000, 1000, dtype=torch.double) * 3.0 + x.requires_grad = True + y = m(x) + return y + + +def _test_softmax(): + a = torch.randn(2, 10, dtype=torch.float64) + b = a.clone() + a.requires_grad = True + b.requires_grad = True + a.softmax(dim=1)[:, 0].sum().backward() + print("a grad = ", a.grad) + softmax(b, dim=1)[:, 0].sum().backward() + print("b grad = ", b.grad) + assert torch.allclose(a.grad, b.grad) + + +def _test_piecewise_linear(): + p = PiecewiseLinear((0, 10.0)) + for x in [-100, 0, 100]: + assert p(x) == 10.0 + p = PiecewiseLinear((0, 10.0), (1, 0.0)) + for x, y in [(-100, 10.0), (0, 10.0), (0.5, 5.0), (1, 0.0), (2, 0.0)]: + print("x, y = ", x, y) + assert p(x) == y, (x, p(x), y) + + q = PiecewiseLinear((0.5, 15.0), (0.6, 1.0)) + x_vals = [-1.0, 0.0, 0.1, 0.2, 0.5, 0.6, 0.7, 0.9, 1.0, 2.0] + pq = p.max(q) + for x in x_vals: + y1 = max(p(x), q(x)) + y2 = pq(x) + assert abs(y1 - y2) < 0.001 + pq = p.min(q) + for x in x_vals: + y1 = min(p(x), q(x)) + y2 = pq(x) + assert abs(y1 - y2) < 0.001 + pq = p + q + for x in x_vals: + y1 = p(x) + q(x) + y2 = pq(x) + assert abs(y1 - y2) < 0.001 + + +def _test_activation_dropout_and_linear(): + in_channels = 20 + out_channels = 30 + + for bias in [True, False]: + # actually we don't test for dropout_p != 0.0 because forward functions will + # different answers. This is because we are using the k2 implementation of + # swoosh_l an swoosh_r inside SwooshL() and SwooshR(), and they call randn() + # internally, messing up the random state. + for dropout_p in [0.0]: + for activation in ["SwooshL", "SwooshR"]: + m1 = nn.Sequential( + SwooshL() if activation == "SwooshL" else SwooshR(), + Dropout3(p=dropout_p, shared_dim=-1), + ScaledLinear( + in_channels, out_channels, bias=bias, initial_scale=0.5 + ), + ) + m2 = ActivationDropoutAndLinear( + in_channels, + out_channels, + bias=bias, + initial_scale=0.5, + activation=activation, + dropout_p=dropout_p, + ) + with torch.no_grad(): + m2.weight[:] = m1[2].weight + if bias: + m2.bias[:] = m1[2].bias + # make sure forward gives same result. + x1 = torch.randn(10, in_channels) + x1.requires_grad = True + + # TEMP. + assert torch.allclose( + SwooshRFunction.apply(x1), SwooshRForward(x1), atol=1.0e-03 + ) + + x2 = x1.clone().detach() + x2.requires_grad = True + seed = 10 + torch.manual_seed(seed) + y1 = m1(x1) + y_grad = torch.randn_like(y1) + y1.backward(gradient=y_grad) + torch.manual_seed(seed) + y2 = m2(x2) + y2.backward(gradient=y_grad) + + print( + f"bias = {bias}, dropout_p = {dropout_p}, activation = {activation}" + ) + print("y1 = ", y1) + print("y2 = ", y2) + assert torch.allclose(y1, y2, atol=0.02) + assert torch.allclose(m1[2].weight.grad, m2.weight.grad, atol=1.0e-05) + if bias: + assert torch.allclose(m1[2].bias.grad, m2.bias.grad, atol=1.0e-05) + print("x1.grad = ", x1.grad) + print("x2.grad = ", x2.grad) + + def isclose(a, b): + # return true if cosine similarity is > 0.9. + return (a * b).sum() > 0.9 * ( + (a**2).sum() * (b**2).sum() + ).sqrt() + + # the SwooshL() implementation has a noisy gradient due to 1-byte + # storage of it. + assert isclose(x1.grad, x2.grad) + + +if __name__ == "__main__": + logging.getLogger().setLevel(logging.DEBUG) + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + _test_piecewise_linear() + _test_softmax() + _test_whiten() + _test_balancer_sign() + _test_balancer_magnitude() + _test_swooshr_deriv() + _test_swooshl_deriv() + _test_activation_dropout_and_linear() diff --git a/zipvoice/models/modules/solver.py b/zipvoice/models/modules/solver.py new file mode 100644 index 0000000..ca71ec0 --- /dev/null +++ b/zipvoice/models/modules/solver.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +# Copyright 2024 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Optional, Union + +import torch + + +class DiffusionModel(torch.nn.Module): + """A wrapper of diffusion models for inference. + Args: + model: The diffusion model. + func_name: The function name to call. + """ + + def __init__( + self, + model: torch.nn.Module, + func_name: str = "forward_fm_decoder", + ): + super().__init__() + self.model = model + self.func_name = func_name + self.model_func = getattr(self.model, func_name) + + def forward( + self, + t: torch.Tensor, + x: torch.Tensor, + text_condition: torch.Tensor, + speech_condition: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + guidance_scale: Union[float, torch.Tensor] = 0.0, + **kwargs + ) -> torch.Tensor: + """ + Forward function that Handles the classifier-free guidance. + Args: + t: The current timestep, a tensor of a tensor of a single float. + x: The initial value, with the shape (batch, seq_len, emb_dim). + text_condition: The text_condition of the diffision model, with + the shape (batch, seq_len, emb_dim). + speech_condition: The speech_condition of the diffision model, with the + shape (batch, seq_len, emb_dim). + padding_mask: The mask for padding; True means masked position, with the + shape (batch, seq_len). + guidance_scale: The scale of classifier-free guidance, a float or a tensor + of shape (batch, 1, 1). + Retrun: + The prediction with the shape (batch, seq_len, emb_dim). + """ + if not torch.is_tensor(guidance_scale): + guidance_scale = torch.tensor( + guidance_scale, dtype=t.dtype, device=t.device + ) + + if (guidance_scale == 0.0).all(): + return self.model_func( + t=t, + xt=x, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + **kwargs + ) + else: + assert t.dim() == 0 + + x = torch.cat([x] * 2, dim=0) + padding_mask = torch.cat([padding_mask] * 2, dim=0) + + text_condition = torch.cat( + [torch.zeros_like(text_condition), text_condition], dim=0 + ) + + if t > 0.5: + speech_condition = torch.cat( + [torch.zeros_like(speech_condition), speech_condition], dim=0 + ) + else: + guidance_scale = guidance_scale * 2 + speech_condition = torch.cat( + [speech_condition, speech_condition], dim=0 + ) + + data_uncond, data_cond = self.model_func( + t=t, + xt=x, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + **kwargs + ).chunk(2, dim=0) + + res = (1 + guidance_scale) * data_cond - guidance_scale * data_uncond + return res + + +class DistillDiffusionModel(DiffusionModel): + """A wrapper of distilled diffusion models for inference. + Args: + model: The distilled diffusion model. + func_name: The function name to call. + """ + + def __init__( + self, + model: torch.nn.Module, + func_name: str = "forward_fm_decoder", + ): + super().__init__(model=model, func_name=func_name) + + def forward( + self, + t: torch.Tensor, + x: torch.Tensor, + text_condition: torch.Tensor, + speech_condition: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + guidance_scale: Union[float, torch.Tensor] = 0.0, + **kwargs + ) -> torch.Tensor: + """ + Forward function that Handles the classifier-free guidance. + Args: + t: The current timestep, a tensor of a single float. + x: The initial value, with the shape (batch, seq_len, emb_dim). + text_condition: The text_condition of the diffision model, with + the shape (batch, seq_len, emb_dim). + speech_condition: The speech_condition of the diffision model, with the + shape (batch, seq_len, emb_dim). + padding_mask: The mask for padding; True means masked position, with the + shape (batch, seq_len). + guidance_scale: The scale of classifier-free guidance, a float or a tensor + of shape (batch, 1, 1). + Retrun: + The prediction with the shape (batch, seq_len, emb_dim). + """ + if not torch.is_tensor(guidance_scale): + guidance_scale = torch.tensor( + guidance_scale, dtype=t.dtype, device=t.device + ) + return self.model_func( + t=t, + xt=x, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + guidance_scale=guidance_scale, + **kwargs + ) + + +class EulerSolver: + def __init__( + self, + model: torch.nn.Module, + func_name: str = "forward_fm_decoder", + ): + """Construct a Euler Solver + Args: + model: The diffusion model. + func_name: The function name to call. + """ + self.model = DiffusionModel(model, func_name=func_name) + + def sample( + self, + x: torch.Tensor, + text_condition: torch.Tensor, + speech_condition: torch.Tensor, + padding_mask: torch.Tensor, + num_step: int = 10, + guidance_scale: Union[float, torch.Tensor] = 0.0, + t_start: float = 0.0, + t_end: float = 1.0, + t_shift: float = 1.0, + **kwargs + ) -> torch.Tensor: + device = x.device + assert isinstance(t_start, float) and isinstance(t_end, float) + + # Generate the schedule of timesteps + timesteps = get_time_steps( + t_start=t_start, + t_end=t_end, + num_step=num_step, + t_shift=t_shift, + device=device, + ) + + for step in range(num_step): + t_cur = timesteps[step] + t_next = timesteps[step + 1] + + # Predict velocity (v) + v = self.model( + t=t_cur, + x=x, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + guidance_scale=guidance_scale, + **kwargs + ) + + # 1. Predict the clean 'data' (x_1) and 'noise' (x_0) + # Flow matching formulation: x_t = (1 - t) * x_0 + t * x_1 + # Therefore: v = x_1 - x_0 + x_1_pred = x + (1.0 - t_cur) * v + x_0_pred = x - t_cur * v + + if step < num_step - 1: + # 2. Probability Flow ODE update (Anchor-based) + # This 'anchors' the next point along the predicted line, + # making it more robust than simple Euler integration. + x = (1.0 - t_next) * x_0_pred + t_next * x_1_pred + else: + # Final step: Snap directly to the predicted clean data + x = x_1_pred + + return x + + +class DistillEulerSolver(EulerSolver): + def __init__( + self, + model: torch.nn.Module, + func_name: str = "forward_fm_decoder", + ): + """Construct a Euler Solver for distilled diffusion models. + Args: + model: The diffusion model. + """ + self.model = DistillDiffusionModel(model, func_name=func_name) + + +def get_time_steps( + t_start: float = 0.0, + t_end: float = 1.0, + num_step: int = 10, + t_shift: float = 1.0, + device: torch.device = torch.device("cpu"), +) -> torch.Tensor: + """Compute the intermediate time steps for sampling. + + Args: + t_start: The starting time of the sampling (default is 0). + t_end: The starting time of the sampling (default is 1). + num_step: The number of sampling. + t_shift: shift the t toward smaller numbers so that the sampling + will emphasize low SNR region. Should be in the range of (0, 1]. + The shifting will be more significant when the number is smaller. + device: A torch device. + Returns: + The time step with the shape (num_step + 1,). + """ + + timesteps = torch.linspace(t_start, t_end, num_step + 1).to(device) + + timesteps = t_shift * timesteps / (1 + (t_shift - 1) * timesteps) + + return timesteps diff --git a/zipvoice/models/modules/zipformer.py b/zipvoice/models/modules/zipformer.py new file mode 100644 index 0000000..8c00bd1 --- /dev/null +++ b/zipvoice/models/modules/zipformer.py @@ -0,0 +1,1680 @@ +#!/usr/bin/env python3 +# Copyright 2022-2024 Xiaomi Corp. (authors: Daniel Povey, +# Zengwei Yao, +# Wei Kang +# Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import logging +import math +import random +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor, nn + +from zipvoice.models.modules.scaling import ( + ActivationDropoutAndLinear, + Balancer, + BiasNorm, + Dropout2, + FloatLike, + Identity, + ScaledLinear, + ScheduledFloat, + SwooshR, + Whiten, + limit_param_value, + penalize_abs_values_gt, + softmax, +) + + +def timestep_embedding(timesteps, dim, max_period=10000): + """Create sinusoidal timestep embeddings. + + :param timesteps: shape of (N) or (N, T) + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an Tensor of positional embeddings. shape of (N, dim) or (T, N, dim) + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) + / half + ) + + if timesteps.dim() == 2: + timesteps = timesteps.transpose(0, 1) # (N, T) -> (T, N) + + args = timesteps[..., None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[..., :1])], dim=-1) + return embedding + + +class TTSZipformer(nn.Module): + """ + Args: + + Note: all "int or Tuple[int]" arguments below will be treated as lists of the same + length as downsampling_factor if they are single ints or one-element tuples. + The length of downsampling_factor defines the number of stacks. + + downsampling_factor (Tuple[int]): downsampling factor for each encoder stack. + Note: this is in addition to the downsampling factor of 2 that is applied in + the frontend (self.encoder_embed). + encoder_dim (Tuple[int]): embedding dimension of each of the encoder stacks, + one per encoder stack. + num_encoder_layers (int or Tuple[int])): number of encoder layers for each stack + query_head_dim (int or Tuple[int]): dimension of query and key per attention + head: per stack, if a tuple.. + pos_head_dim (int or Tuple[int]): dimension of positional-encoding projection + per attention head + value_head_dim (int or Tuple[int]): dimension of value in each attention head + num_heads: (int or Tuple[int]): number of heads in the self-attention mechanism. + Must be at least 4. + feedforward_dim (int or Tuple[int]): hidden dimension in feedforward modules + cnn_module_kernel (int or Tuple[int])): Kernel size of convolution module + + pos_dim (int): the dimension of each positional-encoding vector prior to + projection, e.g. 128. + + dropout (float): dropout rate + warmup_batches (float): number of batches to warm up over; this controls + dropout of encoder layers. + use_time_embed: (bool): if True, take time embedding as an additional input. + time_embed_dim: (int): the dimension of the time embedding. + use_guidance_scale_embed (bool): if True, take guidance scale embedding as + an additional input. + guidance_scale_embed_dim: (int): the dimension of the guidance scale embedding. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + downsampling_factor: Union[int, Tuple[int]] = (2, 4), + num_encoder_layers: Union[int, Tuple[int]] = 4, + cnn_module_kernel: Union[int, Tuple[int]] = 31, + encoder_dim: int = 384, + query_head_dim: int = 24, + pos_head_dim: int = 4, + value_head_dim: int = 12, + num_heads: int = 8, + feedforward_dim: int = 1536, + pos_dim: int = 192, + dropout: FloatLike = None, # see code below for default + warmup_batches: float = 4000.0, + use_time_embed: bool = True, + time_embed_dim: int = 192, + use_guidance_scale_embed: bool = False, + guidance_scale_embed_dim: int = 192, + use_conv: bool = True, + ) -> None: + super(TTSZipformer, self).__init__() + + if dropout is None: + dropout = ScheduledFloat((0.0, 0.3), (20000.0, 0.1)) + if isinstance(downsampling_factor, int): + downsampling_factor = (downsampling_factor,) + + def _to_tuple(x): + """Converts a single int or a 1-tuple of an int to a tuple with the same + length as downsampling_factor""" + if isinstance(x, int): + x = (x,) + if len(x) == 1: + x = x * len(downsampling_factor) + else: + assert len(x) == len(downsampling_factor) and isinstance(x[0], int) + return x + + def _assert_downsampling_factor(factors): + """assert downsampling_factor follows u-net style""" + assert factors[0] == 1 and factors[-1] == 1 + + for i in range(1, len(factors) // 2 + 1): + assert factors[i] == factors[i - 1] * 2 + + for i in range(len(factors) // 2 + 1, len(factors)): + assert factors[i] * 2 == factors[i - 1] + + _assert_downsampling_factor(downsampling_factor) + self.downsampling_factor = downsampling_factor # tuple + num_encoder_layers = _to_tuple(num_encoder_layers) + self.cnn_module_kernel = cnn_module_kernel = _to_tuple(cnn_module_kernel) + self.encoder_dim = encoder_dim + self.num_encoder_layers = num_encoder_layers + self.query_head_dim = query_head_dim + self.value_head_dim = value_head_dim + self.num_heads = num_heads + + self.use_time_embed = use_time_embed + self.use_guidance_scale_embed = use_guidance_scale_embed + + self.time_embed_dim = time_embed_dim + if self.use_time_embed: + assert time_embed_dim != -1 + else: + time_embed_dim = -1 + self.guidance_scale_embed_dim = guidance_scale_embed_dim + + self.in_proj = nn.Linear(in_dim, encoder_dim) + self.out_proj = nn.Linear(encoder_dim, out_dim) + + # each one will be Zipformer2Encoder or DownsampledZipformer2Encoder + encoders = [] + + num_encoders = len(downsampling_factor) + for i in range(num_encoders): + encoder_layer = Zipformer2EncoderLayer( + embed_dim=encoder_dim, + pos_dim=pos_dim, + num_heads=num_heads, + query_head_dim=query_head_dim, + pos_head_dim=pos_head_dim, + value_head_dim=value_head_dim, + feedforward_dim=feedforward_dim, + use_conv=use_conv, + cnn_module_kernel=cnn_module_kernel[i], + dropout=dropout, + ) + + # For the segment of the warmup period, we let the Conv2dSubsampling + # layer learn something. Then we start to warm up the other encoders. + encoder = Zipformer2Encoder( + encoder_layer, + num_encoder_layers[i], + embed_dim=encoder_dim, + time_embed_dim=time_embed_dim, + pos_dim=pos_dim, + warmup_begin=warmup_batches * (i + 1) / (num_encoders + 1), + warmup_end=warmup_batches * (i + 2) / (num_encoders + 1), + final_layerdrop_rate=0.035 * (downsampling_factor[i] ** 0.5), + ) + + if downsampling_factor[i] != 1: + encoder = DownsampledZipformer2Encoder( + encoder, + dim=encoder_dim, + downsample=downsampling_factor[i], + ) + + encoders.append(encoder) + + self.encoders = nn.ModuleList(encoders) + if self.use_time_embed: + self.time_embed = nn.Sequential( + nn.Linear(time_embed_dim, time_embed_dim * 2), + SwooshR(), + nn.Linear(time_embed_dim * 2, time_embed_dim), + ) + else: + self.time_embed = None + + if self.use_guidance_scale_embed: + self.guidance_scale_embed = ScaledLinear( + guidance_scale_embed_dim, + time_embed_dim, + bias=False, + initial_scale=0.1, + ) + else: + self.guidance_scale_embed = None + + def forward( + self, + x: Tensor, + t: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + guidance_scale: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + x: + The input tensor. Its shape is (batch_size, seq_len, feature_dim). + t: + A t tensor of shape (batch_size,) or (batch_size, seq_len) + padding_mask: + The mask for padding, of shape (batch_size, seq_len); True means + masked position. May be None. + guidance_scale: + The guidance scale in classifier-free guidance of distillation model. + Returns: + Return the output embeddings. its shape is + (batch_size, output_seq_len, encoder_dim) + """ + x = x.permute(1, 0, 2) + x = self.in_proj(x) + + if t is not None: + assert t.dim() == 1 or t.dim() == 2, t.shape + time_emb = timestep_embedding(t, self.time_embed_dim) + if guidance_scale is not None: + assert ( + guidance_scale.dim() == 1 or guidance_scale.dim() == 2 + ), guidance_scale.shape + guidance_scale_emb = self.guidance_scale_embed( + timestep_embedding(guidance_scale, self.guidance_scale_embed_dim) + ) + time_emb = time_emb + guidance_scale_emb + time_emb = self.time_embed(time_emb) + else: + time_emb = None + + attn_mask = None + + for i, module in enumerate(self.encoders): + x = module( + x, + time_emb=time_emb, + src_key_padding_mask=padding_mask, + attn_mask=attn_mask, + ) + x = self.out_proj(x) + x = x.permute(1, 0, 2) + return x + + +def _whitening_schedule(x: float, ratio: float = 2.0) -> ScheduledFloat: + return ScheduledFloat((0.0, x), (20000.0, ratio * x), default=x) + + +class Zipformer2EncoderLayer(nn.Module): + """ + Args: + embed_dim: the number of expected features in the input (required). + nhead: the number of heads in the multiheadattention models (required). + feedforward_dim: the dimension of the feedforward network model (required). + dropout: the dropout value (default=0.1). + cnn_module_kernel (int): Kernel size of convolution module (default=31). + + Examples:: + >>> encoder_layer = Zipformer2EncoderLayer(embed_dim=512, nhead=8) + >>> src = torch.rand(10, 32, 512) + >>> pos_emb = torch.rand(32, 19, 512) + >>> out = encoder_layer(src, pos_emb) + """ + + def __init__( + self, + embed_dim: int, + pos_dim: int, + num_heads: int, + query_head_dim: int, + pos_head_dim: int, + value_head_dim: int, + feedforward_dim: int, + dropout: FloatLike = 0.1, + cnn_module_kernel: int = 31, + use_conv: bool = True, + attention_skip_rate: FloatLike = ScheduledFloat( + (0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0 + ), + conv_skip_rate: FloatLike = ScheduledFloat( + (0.0, 0.2), (4000.0, 0.05), (16000, 0.0), default=0 + ), + const_attention_rate: FloatLike = ScheduledFloat( + (0.0, 0.25), (4000.0, 0.025), default=0 + ), + ff2_skip_rate: FloatLike = ScheduledFloat( + (0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0) + ), + ff3_skip_rate: FloatLike = ScheduledFloat( + (0.0, 0.1), (4000.0, 0.01), (50000.0, 0.0) + ), + bypass_skip_rate: FloatLike = ScheduledFloat( + (0.0, 0.5), (4000.0, 0.02), default=0 + ), + ) -> None: + super(Zipformer2EncoderLayer, self).__init__() + self.embed_dim = embed_dim + + # self.bypass implements layer skipping as well as bypass. + self.bypass = BypassModule( + embed_dim, skip_rate=bypass_skip_rate, straight_through_rate=0 + ) + # bypass_mid is bypass used in the middle of the layer. + self.bypass_mid = BypassModule(embed_dim, straight_through_rate=0) + + # skip probability for dynamic modules (meaning: anything but feedforward). + self.attention_skip_rate = copy.deepcopy(attention_skip_rate) + # an additional skip probability that applies to ConvModule to stop it from + # contributing too much early on. + self.conv_skip_rate = copy.deepcopy(conv_skip_rate) + + # ff2_skip_rate is to prevent the ff2 module from having output that's too big + # compared to its residual. + self.ff2_skip_rate = copy.deepcopy(ff2_skip_rate) + self.ff3_skip_rate = copy.deepcopy(ff3_skip_rate) + + self.const_attention_rate = copy.deepcopy(const_attention_rate) + + self.self_attn_weights = RelPositionMultiheadAttentionWeights( + embed_dim, + pos_dim=pos_dim, + num_heads=num_heads, + query_head_dim=query_head_dim, + pos_head_dim=pos_head_dim, + dropout=0.0, + ) + + self.self_attn1 = SelfAttention(embed_dim, num_heads, value_head_dim) + + self.self_attn2 = SelfAttention(embed_dim, num_heads, value_head_dim) + + self.feed_forward1 = FeedforwardModule( + embed_dim, (feedforward_dim * 3) // 4, dropout + ) + + self.feed_forward2 = FeedforwardModule(embed_dim, feedforward_dim, dropout) + + self.feed_forward3 = FeedforwardModule( + embed_dim, (feedforward_dim * 5) // 4, dropout + ) + + self.nonlin_attention = NonlinAttention( + embed_dim, hidden_channels=3 * embed_dim // 4 + ) + + self.use_conv = use_conv + + if self.use_conv: + self.conv_module1 = ConvolutionModule(embed_dim, cnn_module_kernel) + + self.conv_module2 = ConvolutionModule(embed_dim, cnn_module_kernel) + + self.norm = BiasNorm(embed_dim) + + self.balancer1 = Balancer( + embed_dim, + channel_dim=-1, + min_positive=0.45, + max_positive=0.55, + min_abs=0.2, + max_abs=4.0, + ) + + # balancer for output of NonlinAttentionModule + self.balancer_na = Balancer( + embed_dim, + channel_dim=-1, + min_positive=0.3, + max_positive=0.7, + min_abs=ScheduledFloat((0.0, 0.004), (4000.0, 0.02)), + prob=0.05, # out of concern for memory usage + ) + + # balancer for output of feedforward2, prevent it from staying too + # small. give this a very small probability, even at the start of + # training, it's to fix a rare problem and it's OK to fix it slowly. + self.balancer_ff2 = Balancer( + embed_dim, + channel_dim=-1, + min_positive=0.3, + max_positive=0.7, + min_abs=ScheduledFloat((0.0, 0.0), (4000.0, 0.1), default=0.0), + max_abs=2.0, + prob=0.05, + ) + + self.balancer_ff3 = Balancer( + embed_dim, + channel_dim=-1, + min_positive=0.3, + max_positive=0.7, + min_abs=ScheduledFloat((0.0, 0.0), (4000.0, 0.2), default=0.0), + max_abs=4.0, + prob=0.05, + ) + + self.whiten = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(4.0, ratio=3.0), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + self.balancer2 = Balancer( + embed_dim, + channel_dim=-1, + min_positive=0.45, + max_positive=0.55, + min_abs=0.1, + max_abs=4.0, + ) + + def get_sequence_dropout_mask( + self, x: Tensor, dropout_rate: float + ) -> Optional[Tensor]: + if ( + dropout_rate == 0.0 + or not self.training + or torch.jit.is_scripting() + or torch.jit.is_tracing() + ): + return None + batch_size = x.shape[1] + mask = (torch.rand(batch_size, 1, device=x.device) > dropout_rate).to(x.dtype) + return mask + + def sequence_dropout(self, x: Tensor, dropout_rate: float) -> Tensor: + """ + Apply sequence-level dropout to x. + x shape: (seq_len, batch_size, embed_dim) + """ + dropout_mask = self.get_sequence_dropout_mask(x, dropout_rate) + if dropout_mask is None: + return x + else: + return x * dropout_mask + + def forward( + self, + src: Tensor, + pos_emb: Tensor, + time_emb: Optional[Tensor] = None, + attn_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """ + Pass the input through the encoder layer. + Args: + src: the sequence to the encoder (required): + shape (seq_len, batch_size, embedding_dim). + pos_emb: (1, 2*seq_len-1, pos_emb_dim) or + (batch_size, 2*seq_len-1, pos_emb_dim) + time_emb: the embedding representing the current timestep + shape (batch_size, embedding_dim) or (seq_len, batch_size, embedding_dim). + attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) + or (seq_len, seq_len), interpreted as (batch_size, tgt_seq_len, src_seq_len) + or (tgt_seq_len, src_seq_len). True means masked position. May be None. + src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); + True means masked position. May be None. + + Returns: + A tensor which has the same shape as src + """ + src_orig = src + + # dropout rate for non-feedforward submodules + if torch.jit.is_scripting() or torch.jit.is_tracing(): + attention_skip_rate = 0.0 + else: + attention_skip_rate = ( + float(self.attention_skip_rate) if self.training else 0.0 + ) + + # attn_weights: (num_heads, batch_size, seq_len, seq_len) + attn_weights = self.self_attn_weights( + src, + pos_emb=pos_emb, + attn_mask=attn_mask, + key_padding_mask=src_key_padding_mask, + ) + if time_emb is not None: + + src = src + time_emb + + src = src + self.feed_forward1(src) + + self_attn_dropout_mask = self.get_sequence_dropout_mask( + src, attention_skip_rate + ) + + selected_attn_weights = attn_weights[0:1] + if torch.jit.is_scripting() or torch.jit.is_tracing(): + pass + elif self.training and random.random() < float(self.const_attention_rate): + # Make attention weights constant. The intention is to + # encourage these modules to do something similar to an + # averaging-over-time operation. + # only need the mask, can just use the 1st one and expand later + selected_attn_weights = selected_attn_weights[0:1] + selected_attn_weights = (selected_attn_weights > 0.0).to( + selected_attn_weights.dtype + ) + selected_attn_weights = selected_attn_weights * ( + 1.0 / selected_attn_weights.sum(dim=-1, keepdim=True) + ) + + na = self.balancer_na(self.nonlin_attention(src, selected_attn_weights)) + + src = src + ( + na if self_attn_dropout_mask is None else na * self_attn_dropout_mask + ) + + self_attn = self.self_attn1(src, attn_weights) + + src = src + ( + self_attn + if self_attn_dropout_mask is None + else self_attn * self_attn_dropout_mask + ) + + if self.use_conv: + if torch.jit.is_scripting() or torch.jit.is_tracing(): + conv_skip_rate = 0.0 + else: + conv_skip_rate = float(self.conv_skip_rate) if self.training else 0.0 + + if time_emb is not None: + src = src + time_emb + + src = src + self.sequence_dropout( + self.conv_module1( + src, + src_key_padding_mask=src_key_padding_mask, + ), + conv_skip_rate, + ) + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + ff2_skip_rate = 0.0 + else: + ff2_skip_rate = float(self.ff2_skip_rate) if self.training else 0.0 + src = src + self.sequence_dropout( + self.balancer_ff2(self.feed_forward2(src)), ff2_skip_rate + ) + + # bypass in the middle of the layer. + src = self.bypass_mid(src_orig, src) + + self_attn = self.self_attn2(src, attn_weights) + + src = src + ( + self_attn + if self_attn_dropout_mask is None + else self_attn * self_attn_dropout_mask + ) + + if self.use_conv: + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + conv_skip_rate = 0.0 + else: + conv_skip_rate = float(self.conv_skip_rate) if self.training else 0.0 + + if time_emb is not None: + src = src + time_emb + + src = src + self.sequence_dropout( + self.conv_module2( + src, + src_key_padding_mask=src_key_padding_mask, + ), + conv_skip_rate, + ) + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + ff3_skip_rate = 0.0 + else: + ff3_skip_rate = float(self.ff3_skip_rate) if self.training else 0.0 + src = src + self.sequence_dropout( + self.balancer_ff3(self.feed_forward3(src)), ff3_skip_rate + ) + + src = self.balancer1(src) + src = self.norm(src) + + src = self.bypass(src_orig, src) + + src = self.balancer2(src) + src = self.whiten(src) + + return src + + +class Zipformer2Encoder(nn.Module): + r"""Zipformer2Encoder is a stack of N encoder layers + + Args: + encoder_layer: an instance of the Zipformer2EncoderLayer() class (required). + num_layers: the number of sub-encoder-layers in the encoder (required). + pos_dim: the dimension for the relative positional encoding + + Examples:: + >>> encoder_layer = Zipformer2EncoderLayer(embed_dim=512, nhead=8) + >>> zipformer_encoder = Zipformer2Encoder(encoder_layer, num_layers=6) + >>> src = torch.rand(10, 32, 512) + >>> out = zipformer_encoder(src) + """ + + def __init__( + self, + encoder_layer: nn.Module, + num_layers: int, + embed_dim: int, + time_embed_dim: int, + pos_dim: int, + warmup_begin: float, + warmup_end: float, + initial_layerdrop_rate: float = 0.5, + final_layerdrop_rate: float = 0.05, + ) -> None: + super().__init__() + self.encoder_pos = CompactRelPositionalEncoding( + pos_dim, dropout_rate=0.15, length_factor=1.0 + ) + if time_embed_dim != -1: + self.time_emb = nn.Sequential( + SwooshR(), + nn.Linear(time_embed_dim, embed_dim), + ) + else: + self.time_emb = None + + self.layers = nn.ModuleList( + [copy.deepcopy(encoder_layer) for i in range(num_layers)] + ) + self.num_layers = num_layers + + assert 0 <= warmup_begin <= warmup_end + + delta = (1.0 / num_layers) * (warmup_end - warmup_begin) + cur_begin = warmup_begin # interpreted as a training batch index + for i in range(num_layers): + cur_end = cur_begin + delta + self.layers[i].bypass.skip_rate = ScheduledFloat( + (cur_begin, initial_layerdrop_rate), + (cur_end, final_layerdrop_rate), + default=0.0, + ) + cur_begin = cur_end + + def forward( + self, + src: Tensor, + time_emb: Optional[Tensor] = None, + attn_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + r"""Pass the input through the encoder layers in turn. + + Args: + src: the sequence to the encoder (required): + shape (seq_len, batch_size, embedding_dim). + time_emb: the embedding representing the current timestep: + shape (batch_size, embedding_dim) + or (seq_len, batch_size, embedding_dim) . + attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) + or (seq_len, seq_len), interpreted as + (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len). + True means masked position. May be None. + src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); + True means masked position. May be None. + + Returns: a Tensor with the same shape as src. + """ + pos_emb = self.encoder_pos(src) + if self.time_emb is not None: + assert time_emb is not None + time_emb = self.time_emb(time_emb) + else: + assert time_emb is None + + output = src + + for i, mod in enumerate(self.layers): + output = mod( + output, + pos_emb, + time_emb=time_emb, + attn_mask=attn_mask, + src_key_padding_mask=src_key_padding_mask, + ) + + return output + + +class BypassModule(nn.Module): + """ + An nn.Module that implements a learnable bypass scale, and also randomized + per-sequence layer-skipping. The bypass is limited during early stages of training + to be close to "straight-through", i.e. to not do the bypass operation much + initially, in order to force all the modules to learn something. + """ + + def __init__( + self, + embed_dim: int, + skip_rate: FloatLike = 0.0, + straight_through_rate: FloatLike = 0.0, + scale_min: FloatLike = ScheduledFloat((0.0, 0.9), (20000.0, 0.2), default=0), + scale_max: FloatLike = 1.0, + ): + super().__init__() + self.bypass_scale = nn.Parameter(torch.full((embed_dim,), 0.5)) + self.skip_rate = copy.deepcopy(skip_rate) + self.straight_through_rate = copy.deepcopy(straight_through_rate) + self.scale_min = copy.deepcopy(scale_min) + self.scale_max = copy.deepcopy(scale_max) + + def _get_bypass_scale(self, batch_size: int): + # returns bypass-scale of shape (num_channels,), + # or (batch_size, num_channels,). This is actually the + # scale on the non-residual term, so 0 corresponds to bypassing + # this module. + if torch.jit.is_scripting() or torch.jit.is_tracing() or not self.training: + return self.bypass_scale + else: + ans = limit_param_value( + self.bypass_scale, + min=float(self.scale_min), + max=float(self.scale_max), + ) + skip_rate = float(self.skip_rate) + if skip_rate != 0.0: + mask = torch.rand((batch_size, 1), device=ans.device) > skip_rate + ans = ans * mask + # now ans is of shape (batch_size, num_channels), and is zero for + # sequences on which we have randomly chosen to do layer-skipping. + straight_through_rate = float(self.straight_through_rate) + if straight_through_rate != 0.0: + mask = ( + torch.rand((batch_size, 1), device=ans.device) + < straight_through_rate + ) + ans = torch.maximum(ans, mask.to(ans.dtype)) + return ans + + def forward(self, src_orig: Tensor, src: Tensor): + """ + Args: src_orig and src are both of shape (seq_len, batch_size, num_channels) + Returns: something with the same shape as src and src_orig + """ + bypass_scale = self._get_bypass_scale(src.shape[1]) + return src_orig + (src - src_orig) * bypass_scale + + +class DownsampledZipformer2Encoder(nn.Module): + r""" + DownsampledZipformer2Encoder is a zipformer encoder evaluated at a reduced frame + rate, after convolutional downsampling, and then upsampled again at the output, and + combined with the origin input, so that the output has the same shape as the input. + """ + + def __init__(self, encoder: nn.Module, dim: int, downsample: int): + super(DownsampledZipformer2Encoder, self).__init__() + self.downsample_factor = downsample + self.downsample = SimpleDownsample(downsample) + self.num_layers = encoder.num_layers + self.encoder = encoder + self.upsample = SimpleUpsample(downsample) + self.out_combiner = BypassModule(dim, straight_through_rate=0) + + def forward( + self, + src: Tensor, + time_emb: Optional[Tensor] = None, + attn_mask: Optional[Tensor] = None, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + r"""Downsample, go through encoder, upsample. + + Args: + src: the sequence to the encoder (required): + shape (seq_len, batch_size, embedding_dim). + time_emb: the embedding representing the current timestep: + shape (batch_size, embedding_dim) + or (seq_len, batch_size, embedding_dim) . + feature_mask: something that broadcasts with src, that we'll multiply `src` + by at every layer: if a Tensor, likely of shape + (seq_len, batch_size, embedding_dim) + attn_mask: the attention mask, of shape (batch_size, seq_len, seq_len) + or (seq_len, seq_len), interpreted as + (batch_size, tgt_seq_len, src_seq_len) or (tgt_seq_len, src_seq_len). + True means masked position. May be None. + src_key_padding_mask: the mask for padding, of shape (batch_size, seq_len); + True means masked position. May be None. + + Returns: a Tensor with the same shape as src. + """ + src_orig = src + src = self.downsample(src) + ds = self.downsample_factor + if time_emb is not None and time_emb.dim() == 3: + time_emb = time_emb[::ds] + if attn_mask is not None: + attn_mask = attn_mask[::ds, ::ds] + if src_key_padding_mask is not None: + src_key_padding_mask = src_key_padding_mask[..., ::ds] + + src = self.encoder( + src, + time_emb=time_emb, + attn_mask=attn_mask, + src_key_padding_mask=src_key_padding_mask, + ) + src = self.upsample(src) + # remove any extra frames that are not a multiple of downsample_factor + src = src[: src_orig.shape[0]] + + return self.out_combiner(src_orig, src) + + +class SimpleDownsample(torch.nn.Module): + """ + Does downsampling with attention, by weighted sum. + """ + + def __init__(self, downsample: int): + super(SimpleDownsample, self).__init__() + + self.bias = nn.Parameter(torch.zeros(downsample)) + + self.name = None # will be set from training code + + self.downsample = downsample + + def forward(self, src: Tensor) -> Tensor: + """ + x: (seq_len, batch_size, in_channels) + Returns a tensor of shape + ( (seq_len+downsample-1)//downsample, batch_size, channels) + """ + (seq_len, batch_size, in_channels) = src.shape + ds = self.downsample + d_seq_len = (seq_len + ds - 1) // ds + + # Pad to an exact multiple of self.downsample + # right-pad src, repeating the last element. + pad = d_seq_len * ds - seq_len + src_extra = src[src.shape[0] - 1 :].expand(pad, src.shape[1], src.shape[2]) + src = torch.cat((src, src_extra), dim=0) + assert src.shape[0] == d_seq_len * ds + + src = src.reshape(d_seq_len, ds, batch_size, in_channels) + + weights = self.bias.softmax(dim=0) + # weights: (downsample, 1, 1) + weights = weights.unsqueeze(-1).unsqueeze(-1) + + # ans1 is the first `in_channels` channels of the output + ans = (src * weights).sum(dim=1) + + return ans + + +class SimpleUpsample(torch.nn.Module): + """ + A very simple form of upsampling that just repeats the input. + """ + + def __init__(self, upsample: int): + super(SimpleUpsample, self).__init__() + self.upsample = upsample + + def forward(self, src: Tensor) -> Tensor: + """ + x: (seq_len, batch_size, num_channels) + Returns a tensor of shape + ( (seq_len*upsample), batch_size, num_channels) + """ + upsample = self.upsample + (seq_len, batch_size, num_channels) = src.shape + src = src.unsqueeze(1).expand(seq_len, upsample, batch_size, num_channels) + src = src.reshape(seq_len * upsample, batch_size, num_channels) + return src + + +class CompactRelPositionalEncoding(torch.nn.Module): + """ + Relative positional encoding module. This version is "compact" meaning it is able + to encode the important information about the relative position in a relatively + small number of dimensions. The goal is to make it so that small differences between + large relative offsets (e.g. 1000 vs. 1001) make very little difference to the + embedding. Such differences were potentially important when encoding absolute + position, but not important when encoding relative position because there is now no + need to compare two large offsets with each other. + + Our embedding works by projecting the interval [-infinity,infinity] to a finite + interval using the atan() function, before doing the Fourier transform of that fixed + interval. The atan() function would compress the "long tails" too small, making it + hard to distinguish between different magnitudes of large offsets, so we use a + logarithmic function to compress large offsets to a smaller range before applying + atan(). Scalings are chosen in such a way that the embedding can clearly distinguish + individual offsets as long as they are quite close to the origin, e.g. abs(offset) + <= about sqrt(embedding_dim) + + + Args: + embed_dim: Embedding dimension. + dropout_rate: Dropout rate. + max_len: Maximum input length: just a heuristic for initialization. + length_factor: a heuristic scale (should be >= 1.0) which, if larger, gives + less weight to small differences of offset near the origin. + """ + + def __init__( + self, + embed_dim: int, + dropout_rate: FloatLike, + max_len: int = 1000, + length_factor: float = 1.0, + ) -> None: + """Construct a CompactRelPositionalEncoding object.""" + super(CompactRelPositionalEncoding, self).__init__() + self.embed_dim = embed_dim + assert embed_dim % 2 == 0, embed_dim + self.dropout = Dropout2(dropout_rate) + self.pe = None + assert length_factor >= 1.0, length_factor + self.length_factor = length_factor + self.extend_pe(torch.tensor(0.0).expand(max_len)) + + def extend_pe(self, x: Tensor, left_context_len: int = 0) -> None: + """Reset the positional encodings.""" + T = x.size(0) + left_context_len + + if self.pe is not None: + # self.pe contains both positive and negative parts + # the length of self.pe is 2 * input_len - 1 + if self.pe.size(0) >= T * 2 - 1: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + + # if T == 4, x would contain [ -3, -2, 1, 0, 1, 2, 3 ] + x = torch.arange(-(T - 1), T, device=x.device).to(torch.float32).unsqueeze(1) + + freqs = 1 + torch.arange(self.embed_dim // 2, device=x.device) + + # `compression_length` this is arbitrary/heuristic, if it is larger we have more + # resolution for small time offsets but less resolution for large time offsets. + compression_length = self.embed_dim**0.5 + # x_compressed, like X, goes from -infinity to infinity as T goes from -infinity + # to infinity; but it does so more slowly than T for large absolute values of T. + # The formula is chosen so that d(x_compressed )/dx is 1 around x == 0, which is + # important. + x_compressed = ( + compression_length + * x.sign() + * ((x.abs() + compression_length).log() - math.log(compression_length)) + ) + + # if self.length_factor == 1.0, then length_scale is chosen so that the + # FFT can exactly separate points close to the origin (T == 0). So this + # part of the formulation is not really heuristic. + # But empirically, for ASR at least, length_factor > 1.0 seems to work better. + length_scale = self.length_factor * self.embed_dim / (2.0 * math.pi) + + # note for machine implementations: if atan is not available, we can use: + # x.sign() * ((1 / (x.abs() + 1)) - 1) * (-math.pi/2) + # check on wolframalpha.com: plot(sign(x) * (1 / ( abs(x) + 1) - 1 ) * -pi/2 , + # atan(x)) + x_atan = (x_compressed / length_scale).atan() # results between -pi and pi + + cosines = (x_atan * freqs).cos() + sines = (x_atan * freqs).sin() + + pe = torch.zeros(x.shape[0], self.embed_dim, device=x.device) + pe[:, 0::2] = cosines + pe[:, 1::2] = sines + pe[:, -1] = 1.0 # for bias. + + self.pe = pe.to(dtype=x.dtype) + + def forward(self, x: Tensor, left_context_len: int = 0) -> Tensor: + """Create positional encoding. + + Args: + x (Tensor): Input tensor (time, batch, `*`). + left_context_len: (int): Length of cached left context. + + Returns: + positional embedding, of shape (batch, left_context_len + 2*time-1, `*`). + """ + self.extend_pe(x, left_context_len) + x_size_left = x.size(0) + left_context_len + # length of positive side: x.size(0) + left_context_len + # length of negative side: x.size(0) + pos_emb = self.pe[ + self.pe.size(0) // 2 + - x_size_left + + 1 : self.pe.size(0) // 2 # noqa E203 + + x.size(0), + :, + ] + pos_emb = pos_emb.unsqueeze(0) + return self.dropout(pos_emb) + + +class RelPositionMultiheadAttentionWeights(nn.Module): + r"""Module that computes multi-head attention weights with relative position + encoding. Various other modules consume the resulting attention weights: + see, for example, the SimpleAttention module which allows you to compute + conventional attention. + + This is a quite heavily modified from: "Transformer-XL: Attentive Language + Models Beyond a Fixed-Length Context", + we have to write up the differences. + + + Args: + embed_dim: number of channels at the input to this module, e.g. 256 + pos_dim: dimension of the positional encoding vectors, e.g. 128. + num_heads: number of heads to compute weights for, e.g. 8 + query_head_dim: dimension of the query (and key), per head. e.g. 24. + pos_head_dim: dimension of the projected positional encoding per head, e.g. 4. + dropout: dropout probability for attn_output_weights. Default: 0.0. + pos_emb_skip_rate: probability for skipping the pos_emb part of the scores on + any given call to forward(), in training time. + """ + + def __init__( + self, + embed_dim: int, + pos_dim: int, + num_heads: int, + query_head_dim: int, + pos_head_dim: int, + dropout: float = 0.0, + pos_emb_skip_rate: FloatLike = ScheduledFloat((0.0, 0.5), (4000.0, 0.0)), + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.num_heads = num_heads + self.query_head_dim = query_head_dim + self.pos_head_dim = pos_head_dim + self.dropout = dropout + self.pos_emb_skip_rate = copy.deepcopy(pos_emb_skip_rate) + self.name = None # will be overwritten in training code; for diagnostics. + + key_head_dim = query_head_dim + in_proj_dim = (query_head_dim + key_head_dim + pos_head_dim) * num_heads + + # the initial_scale is supposed to take over the "scaling" factor of + # head_dim ** -0.5 that has been used in previous forms of attention, + # dividing it between the query and key. Note: this module is intended + # to be used with the ScaledAdam optimizer; with most other optimizers, + # it would be necessary to apply the scaling factor in the forward function. + self.in_proj = ScaledLinear( + embed_dim, + in_proj_dim, + bias=True, + initial_scale=query_head_dim**-0.25, + ) + + self.whiten_keys = Whiten( + num_groups=num_heads, + whitening_limit=_whitening_schedule(3.0), + prob=(0.025, 0.25), + grad_scale=0.025, + ) + + # add a balancer for the keys that runs with very small probability, and + # tries to enforce that all dimensions have mean around zero. The + # weights produced by this module are invariant to adding a constant to + # the keys, so the derivative of the bias is mathematically zero; but + # due to how Adam/ScaledAdam work, it can learn a fairly large nonzero + # bias because the small numerical roundoff tends to have a non-random + # sign. This module is intended to prevent that. Use a very small + # probability; that should be sufficient to fix the problem. + self.balance_keys = Balancer( + key_head_dim * num_heads, + channel_dim=-1, + min_positive=0.4, + max_positive=0.6, + min_abs=0.0, + max_abs=100.0, + prob=0.025, + ) + + # linear transformation for positional encoding. + self.linear_pos = ScaledLinear( + pos_dim, num_heads * pos_head_dim, bias=False, initial_scale=0.05 + ) + + # the following are for diagnostics only, see --print-diagnostics option + self.copy_pos_query = Identity() + self.copy_query = Identity() + + def forward( + self, + x: Tensor, + pos_emb: Tensor, + key_padding_mask: Optional[Tensor] = None, + attn_mask: Optional[Tensor] = None, + ) -> Tensor: + r""" + Args: + x: input of shape (seq_len, batch_size, embed_dim) + pos_emb: Positional embedding tensor, of shape (1, 2*seq_len - 1, pos_dim) + key_padding_mask: a bool tensor of shape (batch_size, seq_len). + Positions that are True in this mask will be ignored as sources in the + attention weighting. + attn_mask: mask of shape (seq_len, seq_len) or + (batch_size, seq_len, seq_len), interpreted as + ([batch_size,] tgt_seq_len, src_seq_len) + saying which positions are allowed to attend to which other positions. + Returns: + a tensor of attention weights, of + shape (hum_heads, batch_size, seq_len, seq_len) + interpreted as (hum_heads, batch_size, tgt_seq_len, src_seq_len). + """ + x = self.in_proj(x) + query_head_dim = self.query_head_dim + pos_head_dim = self.pos_head_dim + num_heads = self.num_heads + + seq_len, batch_size, _ = x.shape + + query_dim = query_head_dim * num_heads + + # self-attention + q = x[..., 0:query_dim] + k = x[..., query_dim : 2 * query_dim] + # p is the position-encoding query + p = x[..., 2 * query_dim :] + assert p.shape[-1] == num_heads * pos_head_dim, ( + p.shape[-1], + num_heads, + pos_head_dim, + ) + + q = self.copy_query(q) # for diagnostics only, does nothing. + k = self.whiten_keys(self.balance_keys(k)) # does nothing in the forward pass. + p = self.copy_pos_query(p) # for diagnostics only, does nothing. + + q = q.reshape(seq_len, batch_size, num_heads, query_head_dim) + p = p.reshape(seq_len, batch_size, num_heads, pos_head_dim) + k = k.reshape(seq_len, batch_size, num_heads, query_head_dim) + + # time1 refers to target, time2 refers to source. + q = q.permute(2, 1, 0, 3) # (head, batch, time1, query_head_dim) + p = p.permute(2, 1, 0, 3) # (head, batch, time1, pos_head_dim) + k = k.permute(2, 1, 3, 0) # (head, batch, d_k, time2) + + attn_scores = torch.matmul(q, k) + + use_pos_scores = False + if torch.jit.is_scripting() or torch.jit.is_tracing(): + # We can't put random.random() in the same line + use_pos_scores = True + elif not self.training or random.random() >= float(self.pos_emb_skip_rate): + use_pos_scores = True + + if use_pos_scores: + pos_emb = self.linear_pos(pos_emb) + seq_len2 = 2 * seq_len - 1 + pos_emb = pos_emb.reshape(-1, seq_len2, num_heads, pos_head_dim).permute( + 2, 0, 3, 1 + ) + # pos shape now: (head, {1 or batch_size}, pos_dim, seq_len2) + + # (head, batch, time1, pos_dim) x (head, 1, pos_dim, seq_len2) -> (head, + # batch, time1, seq_len2) [where seq_len2 represents relative position.] + pos_scores = torch.matmul(p, pos_emb) + # the following .as_strided() expression converts the last axis of + # pos_scores from relative to absolute position. I don't know whether I + # might have got the time-offsets backwards or not, but let this code define + # which way round it is supposed to be. + if torch.jit.is_tracing(): + (num_heads, batch_size, time1, n) = pos_scores.shape + rows = torch.arange(start=time1 - 1, end=-1, step=-1) + cols = torch.arange(seq_len) + rows = rows.repeat(batch_size * num_heads).unsqueeze(-1) + indexes = rows + cols + pos_scores = pos_scores.reshape(-1, n) + pos_scores = torch.gather(pos_scores, dim=1, index=indexes) + pos_scores = pos_scores.reshape(num_heads, batch_size, time1, seq_len) + else: + pos_scores = pos_scores.as_strided( + (num_heads, batch_size, seq_len, seq_len), + ( + pos_scores.stride(0), + pos_scores.stride(1), + pos_scores.stride(2) - pos_scores.stride(3), + pos_scores.stride(3), + ), + storage_offset=pos_scores.stride(3) * (seq_len - 1), + ) + + attn_scores = attn_scores + pos_scores + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + pass + elif self.training and random.random() < 0.1: + # This is a harder way of limiting the attention scores to not be + # too large. It incurs a penalty if any of them has an absolute + # value greater than 50.0. this should be outside the normal range + # of the attention scores. We use this mechanism instead of, say, + # something added to the loss function involving the entropy, + # because once the entropy gets very small gradients through the + # softmax can become very small, and we'd get zero derivatives. The + # choices of 1.0e-04 as the scale on the penalty makes this + # mechanism vulnerable to the absolute scale of the loss function, + # but we view this as a failsafe to avoid "implausible" parameter + # values rather than a regularization method that should be active + # under normal circumstances. + attn_scores = penalize_abs_values_gt( + attn_scores, limit=25.0, penalty=1.0e-04, name=self.name + ) + + assert attn_scores.shape == (num_heads, batch_size, seq_len, seq_len) + + if attn_mask is not None: + assert attn_mask.dtype == torch.bool + # use -1000 to avoid nan's where attn_mask and key_padding_mask make + # all scores zero. It's important that this be large enough that exp(-1000) + # is exactly zero, for reasons related to const_attention_rate, it + # compares the final weights with zero. + attn_scores = attn_scores.masked_fill(attn_mask, -1000) + + if key_padding_mask is not None: + assert key_padding_mask.shape == ( + batch_size, + seq_len, + ), key_padding_mask.shape + attn_scores = attn_scores.masked_fill( + key_padding_mask.unsqueeze(1), + -1000, + ) + + # We use our own version of softmax, defined in scaling.py, which should + # save a little of the memory used in backprop by, if we are in + # automatic mixed precision mode (amp / autocast), by only storing the + # half-precision output for backprop purposes. + attn_weights = softmax(attn_scores, dim=-1) + + if torch.jit.is_scripting() or torch.jit.is_tracing(): + pass + elif random.random() < 0.001 and not self.training: + self._print_attn_entropy(attn_weights) + + attn_weights = nn.functional.dropout( + attn_weights, p=self.dropout, training=self.training + ) + + return attn_weights + + def _print_attn_entropy(self, attn_weights: Tensor): + # attn_weights: (num_heads, batch_size, seq_len, seq_len) + (num_heads, batch_size, seq_len, seq_len) = attn_weights.shape + + with torch.no_grad(): + with torch.amp.autocast("cuda", enabled=False): + attn_weights = attn_weights.to(torch.float32) + attn_weights_entropy = ( + -((attn_weights + 1.0e-20).log() * attn_weights) + .sum(dim=-1) + .mean(dim=(1, 2)) + ) + logging.debug( + f"name={self.name}, attn_weights_entropy = {attn_weights_entropy}" + ) + + +class SelfAttention(nn.Module): + """ + The simplest possible attention module. This one works with already-computed + attention weights, e.g. as computed by RelPositionMultiheadAttentionWeights. + + Args: + embed_dim: the input and output embedding dimension + num_heads: the number of attention heads + value_head_dim: the value dimension per head + """ + + def __init__( + self, + embed_dim: int, + num_heads: int, + value_head_dim: int, + ) -> None: + super().__init__() + self.in_proj = nn.Linear(embed_dim, num_heads * value_head_dim, bias=True) + + self.out_proj = ScaledLinear( + num_heads * value_head_dim, + embed_dim, + bias=True, + initial_scale=0.05, + ) + + self.whiten = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(7.5, ratio=3.0), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + def forward( + self, + x: Tensor, + attn_weights: Tensor, + ) -> Tensor: + """ + Args: + x: input tensor, of shape (seq_len, batch_size, embed_dim) + attn_weights: a tensor of shape (num_heads, batch_size, seq_len, seq_len), + with seq_len being interpreted as (tgt_seq_len, src_seq_len). Expect + attn_weights.sum(dim=-1) == 1. + Returns: + a tensor with the same shape as x. + """ + (seq_len, batch_size, embed_dim) = x.shape + num_heads = attn_weights.shape[0] + assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len) + + x = self.in_proj(x) # (seq_len, batch_size, num_heads * value_head_dim) + x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3) + # now x: (num_heads, batch_size, seq_len, value_head_dim) + value_head_dim = x.shape[-1] + + # todo: see whether there is benefit in overriding matmul + x = torch.matmul(attn_weights, x) + # v: (num_heads, batch_size, seq_len, value_head_dim) + + x = ( + x.permute(2, 1, 0, 3) + .contiguous() + .view(seq_len, batch_size, num_heads * value_head_dim) + ) + + # returned value is of shape (seq_len, batch_size, embed_dim), like the input. + x = self.out_proj(x) + x = self.whiten(x) + + return x + + +class FeedforwardModule(nn.Module): + """Feedforward module in TTSZipformer model.""" + + def __init__(self, embed_dim: int, feedforward_dim: int, dropout: FloatLike): + super(FeedforwardModule, self).__init__() + self.in_proj = nn.Linear(embed_dim, feedforward_dim) + + self.hidden_balancer = Balancer( + feedforward_dim, + channel_dim=-1, + min_positive=0.3, + max_positive=1.0, + min_abs=0.75, + max_abs=5.0, + ) + + # shared_dim=0 means we share the dropout mask along the time axis + self.out_proj = ActivationDropoutAndLinear( + feedforward_dim, + embed_dim, + activation="SwooshL", + dropout_p=dropout, + dropout_shared_dim=0, + bias=True, + initial_scale=0.1, + ) + + self.out_whiten = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(7.5), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + def forward(self, x: Tensor): + x = self.in_proj(x) + x = self.hidden_balancer(x) + # out_proj contains SwooshL activation, then dropout, then linear. + x = self.out_proj(x) + x = self.out_whiten(x) + return x + + +class NonlinAttention(nn.Module): + """This is like the ConvolutionModule, but refactored so that we use multiplication + by attention weights (borrowed from the attention module) in place of actual + convolution. We also took out the second nonlinearity, the one after the + attention mechanism. + + Args: + channels (int): The number of channels of conv layers. + """ + + def __init__( + self, + channels: int, + hidden_channels: int, + ) -> None: + super().__init__() + + self.hidden_channels = hidden_channels + + self.in_proj = nn.Linear(channels, hidden_channels * 3, bias=True) + + # balancer that goes before the sigmoid. Have quite a large min_abs value, at + # 2.0, because we noticed that well-trained instances of this module have + # abs-value before the sigmoid starting from about 3, and poorly-trained + # instances of the module have smaller abs values before the sigmoid. + self.balancer = Balancer( + hidden_channels, + channel_dim=-1, + min_positive=ScheduledFloat((0.0, 0.25), (20000.0, 0.05)), + max_positive=ScheduledFloat((0.0, 0.75), (20000.0, 0.95)), + min_abs=0.5, + max_abs=5.0, + ) + self.tanh = nn.Tanh() + + self.identity1 = Identity() # for diagnostics. + self.identity2 = Identity() # for diagnostics. + self.identity3 = Identity() # for diagnostics. + + self.out_proj = ScaledLinear( + hidden_channels, channels, bias=True, initial_scale=0.05 + ) + + self.whiten1 = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(5.0), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + self.whiten2 = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(5.0, ratio=3.0), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + def forward( + self, + x: Tensor, + attn_weights: Tensor, + ) -> Tensor: + """. + Args: + x: a Tensor of shape (seq_len, batch_size, num_channels) + attn_weights: a Tensor of shape (num_heads, batch_size, seq_len, seq_len) + Returns: + a Tensor with the same shape as x + """ + x = self.in_proj(x) + + (seq_len, batch_size, _) = x.shape + hidden_channels = self.hidden_channels + + s, x, y = x.chunk(3, dim=2) + + # s will go through tanh. + + s = self.balancer(s) + s = self.tanh(s) + + s = s.unsqueeze(-1).reshape(seq_len, batch_size, hidden_channels) + x = self.whiten1(x) + x = x * s + x = self.identity1(x) # diagnostics only, it's the identity. + + (seq_len, batch_size, embed_dim) = x.shape + num_heads = attn_weights.shape[0] + assert attn_weights.shape == (num_heads, batch_size, seq_len, seq_len) + + x = x.reshape(seq_len, batch_size, num_heads, -1).permute(2, 1, 0, 3) + # now x: (num_heads, batch_size, seq_len, head_dim) + x = torch.matmul(attn_weights, x) + # now x: (num_heads, batch_size, seq_len, head_dim) + x = x.permute(2, 1, 0, 3).reshape(seq_len, batch_size, -1) + + y = self.identity2(y) + x = x * y + x = self.identity3(x) + + x = self.out_proj(x) + x = self.whiten2(x) + return x + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Zipformer2 model. + + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernerl size of conv layers. + bias (bool): Whether to use bias in conv layers (default=True). + + """ + + def __init__( + self, + channels: int, + kernel_size: int, + ) -> None: + """Construct a ConvolutionModule object.""" + super(ConvolutionModule, self).__init__() + # kernerl_size should be a odd number for 'SAME' padding + assert (kernel_size - 1) % 2 == 0 + + bottleneck_dim = channels + + self.in_proj = nn.Linear( + channels, + 2 * bottleneck_dim, + ) + # the gradients on in_proj are a little noisy, likely to do with the + # sigmoid in glu. + + # after in_proj we put x through a gated linear unit (nn.functional.glu). For + # most layers the normal rms value of channels of x seems to be in the range 1 + # to 4, but sometimes, for some reason, for layer 0 the rms ends up being very + # large, between 50 and 100 for different channels. This will cause very peaky + # and sparse derivatives for the sigmoid gating function, which will tend to + # make the loss function not learn effectively. (for most layers the average + # absolute values are in the range 0.5..9.0, and the average p(x>0), i.e. + # positive proportion, at the output of pointwise_conv1.output is around 0.35 to + # 0.45 for different layers, which likely breaks down as 0.5 for the "linear" + # half and 0.2 to 0.3 for the part that goes into the sigmoid. The idea is that + # if we constrain the rms values to a reasonable range via a constraint of + # max_abs=10.0, it will be in a better position to start learning something, + # i.e. to latch onto the correct range. + self.balancer1 = Balancer( + bottleneck_dim, + channel_dim=-1, + min_positive=ScheduledFloat((0.0, 0.05), (8000.0, 0.025)), + max_positive=1.0, + min_abs=1.5, + max_abs=ScheduledFloat((0.0, 5.0), (8000.0, 10.0), default=1.0), + ) + + self.activation1 = Identity() # for diagnostics + + self.sigmoid = nn.Sigmoid() + + self.activation2 = Identity() # for diagnostics + + assert kernel_size % 2 == 1 + + self.depthwise_conv = nn.Conv1d( + in_channels=bottleneck_dim, + out_channels=bottleneck_dim, + groups=bottleneck_dim, + kernel_size=kernel_size, + padding=kernel_size // 2, + ) + + self.balancer2 = Balancer( + bottleneck_dim, + channel_dim=1, + min_positive=ScheduledFloat((0.0, 0.1), (8000.0, 0.05)), + max_positive=1.0, + min_abs=ScheduledFloat((0.0, 0.2), (20000.0, 0.5)), + max_abs=10.0, + ) + + self.whiten = Whiten( + num_groups=1, + whitening_limit=_whitening_schedule(7.5), + prob=(0.025, 0.25), + grad_scale=0.01, + ) + + self.out_proj = ActivationDropoutAndLinear( + bottleneck_dim, + channels, + activation="SwooshR", + dropout_p=0.0, + initial_scale=0.05, + ) + + def forward( + self, + x: Tensor, + src_key_padding_mask: Optional[Tensor] = None, + ) -> Tensor: + """Compute convolution module. + + Args: + x: Input tensor (#time, batch, channels). + src_key_padding_mask: the mask for the src keys per batch (optional): + (batch, #time), contains True in masked positions. + + Returns: + Tensor: Output tensor (#time, batch, channels). + + """ + + x = self.in_proj(x) # (time, batch, 2*channels) + + x, s = x.chunk(2, dim=2) + s = self.balancer1(s) + s = self.sigmoid(s) + x = self.activation1(x) # identity. + x = x * s + x = self.activation2(x) # identity + + # (time, batch, channels) + + # exchange the temporal dimension and the feature dimension + x = x.permute(1, 2, 0) # (#batch, channels, time). + + if src_key_padding_mask is not None: + x = x.masked_fill(src_key_padding_mask.unsqueeze(1).expand_as(x), 0.0) + + x = self.depthwise_conv(x) + + x = self.balancer2(x) + x = x.permute(2, 0, 1) # (time, batch, channels) + + x = self.whiten(x) # (time, batch, channels) + x = self.out_proj(x) # (time, batch, channels) + + return x diff --git a/zipvoice/models/modules/zipformer_two_stream.py b/zipvoice/models/modules/zipformer_two_stream.py new file mode 100644 index 0000000..03dd050 --- /dev/null +++ b/zipvoice/models/modules/zipformer_two_stream.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Optional, Tuple, Union + +import torch +from torch import Tensor, nn + +from zipvoice.models.modules.scaling import FloatLike, ScheduledFloat, SwooshR +from zipvoice.models.modules.zipformer import ( + DownsampledZipformer2Encoder, + TTSZipformer, + Zipformer2Encoder, + Zipformer2EncoderLayer, +) + + +def timestep_embedding(timesteps, dim, max_period=10000): + """Create sinusoidal timestep embeddings. + + :param timesteps: shape of (N) or (N, T) + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an Tensor of positional embeddings. shape of (N, dim) or (T, N, dim) + """ + half = dim // 2 + freqs = torch.exp( + -math.log(max_period) + * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) + / half + ) + + if timesteps.dim() == 2: + timesteps = timesteps.transpose(0, 1) # (N, T) -> (T, N) + + args = timesteps[..., None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[..., :1])], dim=-1) + return embedding + + +class TTSZipformerTwoStream(TTSZipformer): + """ + Args: + + Note: all "int or Tuple[int]" arguments below will be treated as lists of the same + length as downsampling_factor if they are single ints or one-element tuples. + The length of downsampling_factor defines the number of stacks. + + downsampling_factor (Tuple[int]): downsampling factor for each encoder stack. + Note: this is in addition to the downsampling factor of 2 that is applied in + the frontend (self.encoder_embed). + encoder_dim (Tuple[int]): embedding dimension of each of the encoder stacks, + one per encoder stack. + num_encoder_layers (int or Tuple[int])): number of encoder layers for each stack + query_head_dim (int or Tuple[int]): dimension of query and key per attention + head: per stack, if a tuple.. + pos_head_dim (int or Tuple[int]): dimension of positional-encoding projection + per attention head + value_head_dim (int or Tuple[int]): dimension of value in each attention head + num_heads: (int or Tuple[int]): number of heads in the self-attention mechanism. + Must be at least 4. + feedforward_dim (int or Tuple[int]): hidden dimension in feedforward modules + cnn_module_kernel (int or Tuple[int])): Kernel size of convolution module + + pos_dim (int): the dimension of each positional-encoding vector prior to + projection, e.g. 128. + + dropout (float): dropout rate + warmup_batches (float): number of batches to warm up over; this controls + dropout of encoder layers. + use_time_embed: (bool): if True, do not take time embedding as additional input. + time_embed_dim: (int): the dimension of the time embedding. + """ + + def __init__( + self, + in_dim: Tuple[int], + out_dim: Tuple[int], + downsampling_factor: Tuple[int] = (2, 4), + num_encoder_layers: Union[int, Tuple[int]] = 4, + cnn_module_kernel: Union[int, Tuple[int]] = 31, + encoder_dim: int = 384, + query_head_dim: int = 24, + pos_head_dim: int = 4, + value_head_dim: int = 12, + num_heads: int = 8, + feedforward_dim: int = 1536, + pos_dim: int = 192, + dropout: FloatLike = None, # see code below for default + warmup_batches: float = 4000.0, + use_time_embed: bool = True, + time_embed_dim: int = 192, + use_conv: bool = True, + ) -> None: + nn.Module.__init__(self) + + if dropout is None: + dropout = ScheduledFloat((0.0, 0.3), (20000.0, 0.1)) + if isinstance(downsampling_factor, int): + downsampling_factor = (downsampling_factor,) + + def _to_tuple(x): + """Converts a single int or a 1-tuple of an int to a tuple with the same + length as downsampling_factor""" + if isinstance(x, int): + x = (x,) + if len(x) == 1: + x = x * len(downsampling_factor) + else: + assert len(x) == len(downsampling_factor) and isinstance(x[0], int) + return x + + def _assert_downsampling_factor(factors): + """assert downsampling_factor follows u-net style""" + assert factors[0] == 1 and factors[-1] == 1 + + for i in range(1, len(factors) // 2 + 1): + assert factors[i] == factors[i - 1] * 2 + + for i in range(len(factors) // 2 + 1, len(factors)): + assert factors[i] * 2 == factors[i - 1] + + _assert_downsampling_factor(downsampling_factor) + self.downsampling_factor = downsampling_factor # tuple + num_encoder_layers = _to_tuple(num_encoder_layers) + self.cnn_module_kernel = cnn_module_kernel = _to_tuple(cnn_module_kernel) + self.encoder_dim = encoder_dim + self.num_encoder_layers = num_encoder_layers + self.query_head_dim = query_head_dim + self.value_head_dim = value_head_dim + self.num_heads = num_heads + + self.use_time_embed = use_time_embed + + self.time_embed_dim = time_embed_dim + if self.use_time_embed: + assert time_embed_dim != -1 + else: + time_embed_dim = -1 + + assert len(in_dim) == len(out_dim) == 2 + + self.in_dim = in_dim + self.in_proj = nn.ModuleList( + [nn.Linear(in_dim[0], encoder_dim), nn.Linear(in_dim[1], encoder_dim)] + ) + self.out_dim = out_dim + self.out_proj = nn.ModuleList( + [nn.Linear(encoder_dim, out_dim[0]), nn.Linear(encoder_dim, out_dim[1])] + ) + + # each one will be Zipformer2Encoder or DownsampledZipformer2Encoder + encoders = [] + + num_encoders = len(downsampling_factor) + for i in range(num_encoders): + encoder_layer = Zipformer2EncoderLayer( + embed_dim=encoder_dim, + pos_dim=pos_dim, + num_heads=num_heads, + query_head_dim=query_head_dim, + pos_head_dim=pos_head_dim, + value_head_dim=value_head_dim, + feedforward_dim=feedforward_dim, + use_conv=use_conv, + cnn_module_kernel=cnn_module_kernel[i], + dropout=dropout, + ) + + # For the segment of the warmup period, we let the Conv2dSubsampling + # layer learn something. Then we start to warm up the other encoders. + encoder = Zipformer2Encoder( + encoder_layer, + num_encoder_layers[i], + embed_dim=encoder_dim, + time_embed_dim=time_embed_dim, + pos_dim=pos_dim, + warmup_begin=warmup_batches * (i + 1) / (num_encoders + 1), + warmup_end=warmup_batches * (i + 2) / (num_encoders + 1), + final_layerdrop_rate=0.035 * (downsampling_factor[i] ** 0.5), + ) + + if downsampling_factor[i] != 1: + encoder = DownsampledZipformer2Encoder( + encoder, + dim=encoder_dim, + downsample=downsampling_factor[i], + ) + + encoders.append(encoder) + + self.encoders = nn.ModuleList(encoders) + if self.use_time_embed: + self.time_embed = nn.Sequential( + nn.Linear(time_embed_dim, time_embed_dim * 2), + SwooshR(), + nn.Linear(time_embed_dim * 2, time_embed_dim), + ) + else: + self.time_embed = None + + def forward( + self, + x: Tensor, + t: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """ + Args: + x: + The input tensor. Its shape is (batch_size, seq_len, feature_dim). + t: + A t tensor of shape (batch_size,) or (batch_size, seq_len) + padding_mask: + The mask for padding, of shape (batch_size, seq_len); True means + masked position. May be None. + Returns: + Return the output embeddings. its shape is + (batch_size, output_seq_len, encoder_dim) + """ + assert x.size(2) in self.in_dim, f"{x.size(2)} in {self.in_dim}" + if x.size(2) == self.in_dim[0]: + index = 0 + else: + index = 1 + x = x.permute(1, 0, 2) + x = self.in_proj[index](x) + + if t is not None: + assert t.dim() == 1 or t.dim() == 2, t.shape + time_emb = timestep_embedding(t, self.time_embed_dim) + time_emb = self.time_embed(time_emb) + else: + time_emb = None + + attn_mask = None + + for i, module in enumerate(self.encoders): + x = module( + x, + time_emb=time_emb, + src_key_padding_mask=padding_mask, + attn_mask=attn_mask, + ) + x = self.out_proj[index](x) + x = x.permute(1, 0, 2) + return x diff --git a/zipvoice/models/zipvoice.py b/zipvoice/models/zipvoice.py new file mode 100644 index 0000000..24cbd24 --- /dev/null +++ b/zipvoice/models/zipvoice.py @@ -0,0 +1,534 @@ +# Copyright 2024 Xiaomi Corp. (authors: Wei Kang +# Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Optional + +import torch +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP + +from zipvoice.models.modules.solver import EulerSolver +from zipvoice.models.modules.zipformer import TTSZipformer +from zipvoice.utils.common import ( + condition_time_mask, + get_tokens_index, + make_pad_mask, + pad_labels, + prepare_avg_tokens_durations, +) + + +class ZipVoice(nn.Module): + """The ZipVoice model.""" + + def __init__( + self, + fm_decoder_downsampling_factor: List[int] = [1, 2, 4, 2, 1], + fm_decoder_num_layers: List[int] = [2, 2, 4, 4, 4], + fm_decoder_cnn_module_kernel: List[int] = [31, 15, 7, 15, 31], + fm_decoder_feedforward_dim: int = 1536, + fm_decoder_num_heads: int = 4, + fm_decoder_dim: int = 512, + text_encoder_num_layers: int = 4, + text_encoder_feedforward_dim: int = 512, + text_encoder_cnn_module_kernel: int = 9, + text_encoder_num_heads: int = 4, + text_encoder_dim: int = 192, + time_embed_dim: int = 192, + text_embed_dim: int = 192, + query_head_dim: int = 32, + value_head_dim: int = 12, + pos_head_dim: int = 4, + pos_dim: int = 48, + feat_dim: int = 100, + vocab_size: int = 26, + pad_id: int = 0, + ): + """ + Initialize the model with specified configuration parameters. + + Args: + fm_decoder_downsampling_factor: List of downsampling factors for each layer + in the flow-matching decoder. + fm_decoder_num_layers: List of the number of layers for each block in the + flow-matching decoder. + fm_decoder_cnn_module_kernel: List of kernel sizes for CNN modules in the + flow-matching decoder. + fm_decoder_feedforward_dim: Dimension of the feedforward network in the + flow-matching decoder. + fm_decoder_num_heads: Number of attention heads in the flow-matching + decoder. + fm_decoder_dim: Hidden dimension of the flow-matching decoder. + text_encoder_num_layers: Number of layers in the text encoder. + text_encoder_feedforward_dim: Dimension of the feedforward network in the + text encoder. + text_encoder_cnn_module_kernel: Kernel size for the CNN module in the + text encoder. + text_encoder_num_heads: Number of attention heads in the text encoder. + text_encoder_dim: Hidden dimension of the text encoder. + time_embed_dim: Dimension of the time embedding. + text_embed_dim: Dimension of the text embedding. + query_head_dim: Dimension of the query attention head. + value_head_dim: Dimension of the value attention head. + pos_head_dim: Dimension of the position attention head. + pos_dim: Dimension of the positional encoding. + feat_dim: Dimension of the acoustic features. + vocab_size: Size of the vocabulary. + pad_id: ID used for padding tokens. + """ + super().__init__() + + self.fm_decoder = TTSZipformer( + in_dim=feat_dim * 3, + out_dim=feat_dim, + downsampling_factor=fm_decoder_downsampling_factor, + num_encoder_layers=fm_decoder_num_layers, + cnn_module_kernel=fm_decoder_cnn_module_kernel, + encoder_dim=fm_decoder_dim, + feedforward_dim=fm_decoder_feedforward_dim, + num_heads=fm_decoder_num_heads, + query_head_dim=query_head_dim, + pos_head_dim=pos_head_dim, + value_head_dim=value_head_dim, + pos_dim=pos_dim, + use_time_embed=True, + time_embed_dim=time_embed_dim, + ) + + self.text_encoder = TTSZipformer( + in_dim=text_embed_dim, + out_dim=feat_dim, + downsampling_factor=1, + num_encoder_layers=text_encoder_num_layers, + cnn_module_kernel=text_encoder_cnn_module_kernel, + encoder_dim=text_encoder_dim, + feedforward_dim=text_encoder_feedforward_dim, + num_heads=text_encoder_num_heads, + query_head_dim=query_head_dim, + pos_head_dim=pos_head_dim, + value_head_dim=value_head_dim, + pos_dim=pos_dim, + use_time_embed=False, + ) + + self.feat_dim = feat_dim + self.text_embed_dim = text_embed_dim + self.pad_id = pad_id + + self.embed = nn.Embedding(vocab_size, text_embed_dim) + self.solver = EulerSolver(self, func_name="forward_fm_decoder") + + def forward_fm_decoder( + self, + t: torch.Tensor, + xt: torch.Tensor, + text_condition: torch.Tensor, + speech_condition: torch.Tensor, + padding_mask: Optional[torch.Tensor] = None, + guidance_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Compute velocity. + Args: + t: A tensor of shape (N, 1, 1) or a tensor of a float, + in the range of (0, 1). + xt: the input of the current timestep, including condition + embeddings and noisy acoustic features. + text_condition: the text condition embeddings, with the + shape (batch, seq_len, emb_dim). + speech_condition: the speech condition embeddings, with the + shape (batch, seq_len, emb_dim). + padding_mask: The mask for padding, True means masked + position, with the shape (N, T). + guidance_scale: The guidance scale in classifier-free guidance, + which is a tensor of shape (N, 1, 1) or a tensor of a float. + + Returns: + predicted velocity, with the shape (batch, seq_len, emb_dim). + """ + + xt = torch.cat([xt, text_condition, speech_condition], dim=2) + + assert t.dim() in (0, 3) + # Handle t with the shape (N, 1, 1): + # squeeze the last dimension if it's size is 1. + while t.dim() > 1 and t.size(-1) == 1: + t = t.squeeze(-1) + # Handle t with a single value: expand to the size of batch size. + if t.dim() == 0: + t = t.repeat(xt.shape[0]) + + if guidance_scale is not None: + while guidance_scale.dim() > 1 and guidance_scale.size(-1) == 1: + guidance_scale = guidance_scale.squeeze(-1) + if guidance_scale.dim() == 0: + guidance_scale = guidance_scale.repeat(xt.shape[0]) + + vt = self.fm_decoder( + x=xt, t=t, padding_mask=padding_mask, guidance_scale=guidance_scale + ) + else: + vt = self.fm_decoder(x=xt, t=t, padding_mask=padding_mask) + return vt + + def forward_text_embed( + self, + tokens: List[List[int]], + ): + """ + Get the text embeddings. + Args: + tokens: a list of list of token ids. + Returns: + embed: the text embeddings, shape (batch, seq_len, emb_dim). + tokens_lens: the length of each token sequence, shape (batch,). + """ + device = ( + self.device if isinstance(self, DDP) else next(self.parameters()).device + ) + tokens_padded = pad_labels(tokens, pad_id=self.pad_id, device=device) # (B, S) + embed = self.embed(tokens_padded) # (B, S, C) + tokens_lens = torch.tensor( + [len(token) for token in tokens], dtype=torch.int64, device=device + ) + tokens_padding_mask = make_pad_mask(tokens_lens, embed.shape[1]) # (B, S) + + embed = self.text_encoder( + x=embed, t=None, padding_mask=tokens_padding_mask + ) # (B, S, C) + return embed, tokens_lens + + def forward_text_condition( + self, + embed: torch.Tensor, + tokens_lens: torch.Tensor, + features_lens: torch.Tensor, + ): + """ + Get the text condition with the same length of the acoustic feature. + Args: + embed: the text embeddings, shape (batch, token_seq_len, emb_dim). + tokens_lens: the length of each token sequence, shape (batch,). + features_lens: the length of each acoustic feature sequence, + shape (batch,). + Returns: + text_condition: the text condition, shape + (batch, feature_seq_len, emb_dim). + padding_mask: the padding mask of text condition, shape + (batch, feature_seq_len). + """ + + num_frames = int(features_lens.max()) + + padding_mask = make_pad_mask(features_lens, max_len=num_frames) # (B, T) + + tokens_durations = prepare_avg_tokens_durations(features_lens, tokens_lens) + + tokens_index = get_tokens_index(tokens_durations, num_frames).to( + embed.device + ) # (B, T) + + text_condition = torch.gather( + embed, + dim=1, + index=tokens_index.unsqueeze(-1).expand( + embed.size(0), num_frames, embed.size(-1) + ), + ) # (B, T, F) + return text_condition, padding_mask + + def forward_text_train( + self, + tokens: List[List[int]], + features_lens: torch.Tensor, + ): + """ + Process text for training, given text tokens and real feature lengths. + """ + embed, tokens_lens = self.forward_text_embed(tokens) + text_condition, padding_mask = self.forward_text_condition( + embed, tokens_lens, features_lens + ) + return ( + text_condition, + padding_mask, + ) + + def forward_text_inference_gt_duration( + self, + tokens: List[List[int]], + features_lens: torch.Tensor, + prompt_tokens: List[List[int]], + prompt_features_lens: torch.Tensor, + ): + """ + Process text for inference, given text tokens, real feature lengths and prompts. + """ + tokens = [ + prompt_token + token for prompt_token, token in zip(prompt_tokens, tokens) + ] + features_lens = prompt_features_lens + features_lens + embed, tokens_lens = self.forward_text_embed(tokens) + text_condition, padding_mask = self.forward_text_condition( + embed, tokens_lens, features_lens + ) + return text_condition, padding_mask + + def forward_text_inference_ratio_duration( + self, + tokens: List[List[int]], + prompt_tokens: List[List[int]], + prompt_features_lens: torch.Tensor, + speed: float, + ): + """ + Process text for inference, given text tokens and prompts, + feature lengths are predicted with the ratio of token numbers. + """ + device = ( + self.device if isinstance(self, DDP) else next(self.parameters()).device + ) + + cat_tokens = [ + prompt_token + token for prompt_token, token in zip(prompt_tokens, tokens) + ] + + prompt_tokens_lens = torch.tensor( + [len(token) for token in prompt_tokens], + dtype=torch.int64, + device=device, + ) + + tokens_lens = torch.tensor( + [len(token) for token in tokens], + dtype=torch.int64, + device=device, + ) + + cat_embed, cat_tokens_lens = self.forward_text_embed(cat_tokens) + + features_lens = prompt_features_lens + torch.ceil( + (prompt_features_lens / prompt_tokens_lens * tokens_lens / speed) + ).to(dtype=torch.int64) + + text_condition, padding_mask = self.forward_text_condition( + cat_embed, cat_tokens_lens, features_lens + ) + return text_condition, padding_mask + + def forward( + self, + tokens: List[List[int]], + features: torch.Tensor, + features_lens: torch.Tensor, + noise: torch.Tensor, + t: torch.Tensor, + condition_drop_ratio: float = 0.0, + ) -> torch.Tensor: + """Forward pass of the model for training. + Args: + tokens: a list of list of token ids. + features: the acoustic features, with the shape (batch, seq_len, feat_dim). + features_lens: the length of each acoustic feature sequence, shape (batch,). + noise: the intitial noise, with the shape (batch, seq_len, feat_dim). + t: the time step, with the shape (batch, 1, 1). + condition_drop_ratio: the ratio of dropped text condition. + Returns: + fm_loss: the flow-matching loss. + """ + + (text_condition, padding_mask,) = self.forward_text_train( + tokens=tokens, + features_lens=features_lens, + ) + + speech_condition_mask = condition_time_mask( + features_lens=features_lens, + mask_percent=(0.7, 1.0), + max_len=features.size(1), + ) + speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features) + + if condition_drop_ratio > 0.0: + drop_mask = ( + torch.rand(text_condition.size(0), 1, 1).to(text_condition.device) + > condition_drop_ratio + ) + text_condition = text_condition * drop_mask + + xt = features * t + noise * (1 - t) + ut = features - noise # (B, T, F) + + vt = self.forward_fm_decoder( + t=t, + xt=xt, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + ) + + loss_mask = speech_condition_mask & (~padding_mask) + fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2) + + return fm_loss + + def sample( + self, + tokens: List[List[int]], + prompt_tokens: List[List[int]], + prompt_features: torch.Tensor, + prompt_features_lens: torch.Tensor, + features_lens: Optional[torch.Tensor] = None, + speed: float = 1.0, + t_shift: float = 1.0, + duration: str = "predict", + num_step: int = 5, + guidance_scale: float = 0.5, + ) -> torch.Tensor: + """ + Generate acoustic features, given text tokens, prompts feature + and prompt transcription's text tokens. + Args: + tokens: a list of list of text tokens. + prompt_tokens: a list of list of prompt tokens. + prompt_features: the prompt feature with the shape + (batch_size, seq_len, feat_dim). + prompt_features_lens: the length of each prompt feature, + with the shape (batch_size,). + features_lens: the length of the predicted eature, with the + shape (batch_size,). It is used only when duration is "real". + duration: "real" or "predict". If "real", the predicted + feature length is given by features_lens. + num_step: the number of steps to use in the ODE solver. + guidance_scale: the guidance scale for classifier-free guidance. + """ + + assert duration in ["real", "predict"] + + if duration == "predict": + ( + text_condition, + padding_mask, + ) = self.forward_text_inference_ratio_duration( + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features_lens=prompt_features_lens, + speed=speed, + ) + else: + assert features_lens is not None + text_condition, padding_mask = self.forward_text_inference_gt_duration( + tokens=tokens, + features_lens=features_lens, + prompt_tokens=prompt_tokens, + prompt_features_lens=prompt_features_lens, + ) + batch_size, num_frames, _ = text_condition.shape + + speech_condition = torch.nn.functional.pad( + prompt_features, (0, 0, 0, num_frames - prompt_features.size(1)) + ) # (B, T, F) + + # False means speech condition positions. + speech_condition_mask = make_pad_mask(prompt_features_lens, num_frames) + speech_condition = torch.where( + speech_condition_mask.unsqueeze(-1), + torch.zeros_like(speech_condition), + speech_condition, + ) + + x0 = torch.randn( + batch_size, + num_frames, + prompt_features.size(-1), + device=text_condition.device, + ) + + x1 = self.solver.sample( + x=x0, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + num_step=num_step, + guidance_scale=guidance_scale, + t_shift=t_shift, + ) + x1_wo_prompt_lens = (~padding_mask).sum(-1) - prompt_features_lens + x1_prompt = torch.zeros( + x1.size(0), prompt_features_lens.max(), x1.size(2), device=x1.device + ) + x1_wo_prompt = torch.zeros( + x1.size(0), x1_wo_prompt_lens.max(), x1.size(2), device=x1.device + ) + for i in range(x1.size(0)): + x1_wo_prompt[i, : x1_wo_prompt_lens[i], :] = x1[ + i, + prompt_features_lens[i] : prompt_features_lens[i] + + x1_wo_prompt_lens[i], + ] + x1_prompt[i, : prompt_features_lens[i], :] = x1[ + i, : prompt_features_lens[i] + ] + + return x1_wo_prompt, x1_wo_prompt_lens, x1_prompt, prompt_features_lens + + def sample_intermediate( + self, + tokens: List[List[int]], + features: torch.Tensor, + features_lens: torch.Tensor, + noise: torch.Tensor, + speech_condition_mask: torch.Tensor, + t_start: float, + t_end: float, + num_step: int = 1, + guidance_scale: torch.Tensor = None, + ) -> torch.Tensor: + """ + Generate acoustic features in intermediate timesteps. + Args: + tokens: List of list of token ids. + features: The acoustic features, with the shape (batch, seq_len, feat_dim). + features_lens: The length of each acoustic feature sequence, + with the shape (batch,). + noise: The initial noise, with the shape (batch, seq_len, feat_dim). + speech_condition_mask: The mask for speech condition, True means + non-condition positions, with the shape (batch, seq_len). + t_start: The start timestep. + t_end: The end timestep. + num_step: The number of steps for sampling. + guidance_scale: The scale for classifier-free guidance inference, + with the shape (batch, 1, 1). + """ + (text_condition, padding_mask,) = self.forward_text_train( + tokens=tokens, + features_lens=features_lens, + ) + + speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features) + + x_t_end = self.solver.sample( + x=noise, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + num_step=num_step, + guidance_scale=guidance_scale, + t_start=t_start, + t_end=t_end, + ) + x_t_end_lens = (~padding_mask).sum(-1) + return x_t_end, x_t_end_lens diff --git a/zipvoice/models/zipvoice_dialog.py b/zipvoice/models/zipvoice_dialog.py new file mode 100644 index 0000000..1137d5c --- /dev/null +++ b/zipvoice/models/zipvoice_dialog.py @@ -0,0 +1,358 @@ +# Copyright 2025 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List + +import torch +import torch.nn as nn +from torch.nn.parallel import DistributedDataParallel as DDP + +from zipvoice.models.modules.zipformer_two_stream import TTSZipformerTwoStream +from zipvoice.models.zipvoice import ZipVoice +from zipvoice.utils.common import condition_time_mask_suffix, make_pad_mask, pad_labels + + +class ZipVoiceDialog(ZipVoice): + """The ZipVoice-Dialog model.""" + + def __init__( + self, + fm_decoder_downsampling_factor: List[int] = [1, 2, 4, 2, 1], + fm_decoder_num_layers: List[int] = [2, 2, 4, 4, 4], + fm_decoder_cnn_module_kernel: List[int] = [31, 15, 7, 15, 31], + fm_decoder_feedforward_dim: int = 1536, + fm_decoder_num_heads: int = 4, + fm_decoder_dim: int = 512, + text_encoder_num_layers: int = 4, + text_encoder_feedforward_dim: int = 512, + text_encoder_cnn_module_kernel: int = 9, + text_encoder_num_heads: int = 4, + text_encoder_dim: int = 192, + time_embed_dim: int = 192, + text_embed_dim: int = 192, + query_head_dim: int = 32, + value_head_dim: int = 12, + pos_head_dim: int = 4, + pos_dim: int = 48, + feat_dim: int = 100, + vocab_size: int = 26, + pad_id: int = 0, + spk_a_id: int = 360, + spk_b_id: int = 361, + ): + """ + Initialize the model with specified configuration parameters. + + Args: + fm_decoder_downsampling_factor: List of downsampling factors for each layer + in the flow-matching decoder. + fm_decoder_num_layers: List of the number of layers for each block in the + flow-matching decoder. + fm_decoder_cnn_module_kernel: List of kernel sizes for CNN modules in the + flow-matching decoder. + fm_decoder_feedforward_dim: Dimension of the feedforward network in the + flow-matching decoder. + fm_decoder_num_heads: Number of attention heads in the flow-matching + decoder. + fm_decoder_dim: Hidden dimension of the flow-matching decoder. + text_encoder_num_layers: Number of layers in the text encoder. + text_encoder_feedforward_dim: Dimension of the feedforward network in the + text encoder. + text_encoder_cnn_module_kernel: Kernel size for the CNN module in the + text encoder. + text_encoder_num_heads: Number of attention heads in the text encoder. + text_encoder_dim: Hidden dimension of the text encoder. + time_embed_dim: Dimension of the time embedding. + text_embed_dim: Dimension of the text embedding. + query_head_dim: Dimension of the query attention head. + value_head_dim: Dimension of the value attention head. + pos_head_dim: Dimension of the position attention head. + pos_dim: Dimension of the positional encoding. + feat_dim: Dimension of the acoustic features. + vocab_size: Size of the vocabulary. + pad_id: ID used for padding tokens. + spk_a_id: ID of speaker A / [S1]. + spk_b_id: ID of speaker B / [S2]. + """ + super().__init__( + fm_decoder_downsampling_factor=fm_decoder_downsampling_factor, + fm_decoder_num_layers=fm_decoder_num_layers, + fm_decoder_cnn_module_kernel=fm_decoder_cnn_module_kernel, + fm_decoder_feedforward_dim=fm_decoder_feedforward_dim, + fm_decoder_num_heads=fm_decoder_num_heads, + fm_decoder_dim=fm_decoder_dim, + text_encoder_num_layers=text_encoder_num_layers, + text_encoder_feedforward_dim=text_encoder_feedforward_dim, + text_encoder_cnn_module_kernel=text_encoder_cnn_module_kernel, + text_encoder_num_heads=text_encoder_num_heads, + text_encoder_dim=text_encoder_dim, + time_embed_dim=time_embed_dim, + text_embed_dim=text_embed_dim, + query_head_dim=query_head_dim, + value_head_dim=value_head_dim, + pos_head_dim=pos_head_dim, + pos_dim=pos_dim, + feat_dim=feat_dim, + vocab_size=vocab_size, + pad_id=pad_id, + ) + + self.spk_a_id = spk_a_id + self.spk_b_id = spk_b_id + self.spk_embed = nn.Embedding(2, feat_dim) + torch.nn.init.normal_(self.spk_embed.weight, mean=0, std=0.1) + + def extract_spk_indices(self, tensor): + turn_mask = ((tensor == self.spk_a_id) | (tensor == self.spk_b_id)).long() + turn_counts = turn_mask.cumsum(dim=1) + spk_mask = turn_counts % 2 + spk_mask = torch.where(tensor == self.pad_id, -1, spk_mask) + spk_a_indices = torch.where(spk_mask == 0) + spk_b_indices = torch.where(spk_mask == 1) + return spk_a_indices, spk_b_indices + + def forward_text_embed( + self, + tokens: List[List[int]], + ): + """ + Get the text embeddings. + Args: + tokens: a list of list of token ids. + Returns: + embed: the text embeddings, shape (batch, seq_len, emb_dim). + tokens_lens: the length of each token sequence, shape (batch,). + """ + device = ( + self.device if isinstance(self, DDP) else next(self.parameters()).device + ) + tokens_padded = pad_labels(tokens, pad_id=self.pad_id, device=device) # (B, S) + embed = self.embed(tokens_padded) # (B, S, C) + spk_a_indices, spk_b_indices = self.extract_spk_indices(tokens_padded) + tokens_lens = torch.tensor( + [len(token) for token in tokens], dtype=torch.int64, device=device + ) + tokens_padding_mask = make_pad_mask(tokens_lens, embed.shape[1]) # (B, S) + + embed = self.text_encoder( + x=embed, t=None, padding_mask=tokens_padding_mask + ) # (B, S, C) + embed[spk_a_indices] += self.spk_embed(torch.tensor(0, device=device)).to( + embed.dtype + ) + embed[spk_b_indices] += self.spk_embed(torch.tensor(1, device=device)).to( + embed.dtype + ) + return embed, tokens_lens + + def forward( + self, + tokens: List[List[int]], + features: torch.Tensor, + features_lens: torch.Tensor, + noise: torch.Tensor, + t: torch.Tensor, + condition_drop_ratio: float = 0.0, + ) -> torch.Tensor: + """Forward pass of the model for training. + Args: + tokens: a list of list of token ids. + features: the acoustic features, with the shape (batch, seq_len, feat_dim). + features_lens: the length of each acoustic feature sequence, shape (batch,). + noise: the intitial noise, with the shape (batch, seq_len, feat_dim). + t: the time step, with the shape (batch, 1, 1). + condition_drop_ratio: the ratio of dropped text condition. + Returns: + fm_loss: the flow-matching loss. + """ + + (text_condition, padding_mask,) = self.forward_text_train( + tokens=tokens, + features_lens=features_lens, + ) + + speech_condition_mask = condition_time_mask_suffix( + features_lens=features_lens, + mask_percent=(0.5, 1.0), + max_len=features.size(1), + ) + speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features) + + if condition_drop_ratio > 0.0: + drop_mask = ( + torch.rand(text_condition.size(0), 1, 1).to(text_condition.device) + > condition_drop_ratio + ) + text_condition = text_condition * drop_mask + + xt = features * t + noise * (1 - t) + ut = features - noise # (B, T, F) + + vt = self.forward_fm_decoder( + t=t, + xt=xt, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + ) + + loss_mask = speech_condition_mask & (~padding_mask) + fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2) + + return fm_loss + + +class ZipVoiceDialogStereo(ZipVoiceDialog): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + required_params = { + "feat_dim", + "fm_decoder_downsampling_factor", + "fm_decoder_num_layers", + "fm_decoder_cnn_module_kernel", + "fm_decoder_dim", + "fm_decoder_feedforward_dim", + "fm_decoder_num_heads", + "query_head_dim", + "pos_head_dim", + "value_head_dim", + "pos_dim", + "time_embed_dim", + } + + missing = [p for p in required_params if p not in kwargs] + if missing: + raise ValueError(f"Missing required parameters: {', '.join(missing)}") + + self.fm_decoder = TTSZipformerTwoStream( + in_dim=(kwargs["feat_dim"] * 5, kwargs["feat_dim"] * 3), + out_dim=(kwargs["feat_dim"] * 2, kwargs["feat_dim"]), + downsampling_factor=kwargs["fm_decoder_downsampling_factor"], + num_encoder_layers=kwargs["fm_decoder_num_layers"], + cnn_module_kernel=kwargs["fm_decoder_cnn_module_kernel"], + encoder_dim=kwargs["fm_decoder_dim"], + feedforward_dim=kwargs["fm_decoder_feedforward_dim"], + num_heads=kwargs["fm_decoder_num_heads"], + query_head_dim=kwargs["query_head_dim"], + pos_head_dim=kwargs["pos_head_dim"], + value_head_dim=kwargs["value_head_dim"], + pos_dim=kwargs["pos_dim"], + use_time_embed=True, + time_embed_dim=kwargs["time_embed_dim"], + ) + + def forward( + self, + tokens: List[List[int]], + features: torch.Tensor, + features_lens: torch.Tensor, + noise: torch.Tensor, + t: torch.Tensor, + condition_drop_ratio: float = 0.0, + se_weight: float = 1.0, + ) -> torch.Tensor: + """Forward pass of the model for training. + Args: + tokens: a list of list of token ids. + features: the acoustic features, with the shape (batch, seq_len, feat_dim). + features_lens: the length of each acoustic feature sequence, shape (batch,). + noise: the intitial noise, with the shape (batch, seq_len, feat_dim). + t: the time step, with the shape (batch, 1, 1). + condition_drop_ratio: the ratio of dropped text condition. + se_weight: the weight of the speaker exclusive loss. + Returns: + fm_loss: the flow-matching loss. + """ + + (text_condition, padding_mask,) = self.forward_text_train( + tokens=tokens, + features_lens=features_lens, + ) + + speech_condition_mask = condition_time_mask_suffix( + features_lens=features_lens, + mask_percent=(0.5, 1.0), + max_len=features.size(1), + ) + speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features) + + if condition_drop_ratio > 0.0: + drop_mask = ( + torch.rand(text_condition.size(0), 1, 1).to(text_condition.device) + > condition_drop_ratio + ) + text_condition = text_condition * drop_mask + + xt = features * t + noise * (1 - t) + ut = features - noise # (B, T, F) + + vt = self.forward_fm_decoder( + t=t, + xt=xt, + text_condition=text_condition, + speech_condition=speech_condition, + padding_mask=padding_mask, + ) + + loss_mask = speech_condition_mask & (~padding_mask) + fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2) + + if se_weight > 0: + target = xt + vt * (1 - t) + fbank_1 = target[:, :, : self.feat_dim] + fbank_2 = target[:, :, self.feat_dim :] + energy_loss = torch.mean( + self.energy_based_loss(fbank_1, fbank_2, features)[loss_mask] + ) + loss = fm_loss + energy_loss * se_weight + else: + loss = fm_loss + + return loss + + def energy_based_loss(self, fbank1, fbank2, gt_fbank): + energy1 = self.energy(fbank1) + energy2 = self.energy(fbank2) + + energy_thresholds = self.adaptive_threshold_from_gt( + torch.cat( + [ + gt_fbank[:, :, : self.feat_dim], + gt_fbank[:, :, self.feat_dim :], + ], + dim=1, + ) + ) + + both_speaking = ( + (energy1 > energy_thresholds) & (energy2 > energy_thresholds) + ).float() + + penalty = ( + both_speaking + * (energy1 - energy_thresholds) + * (energy2 - energy_thresholds) + ) + return penalty + + def energy(self, fbank): + return torch.mean(fbank, dim=-1) + + def adaptive_threshold_from_gt(self, gt_fbank, percentile=50): + frame_energies = self.energy(gt_fbank) + thresholds = torch.quantile(frame_energies, q=percentile / 100, dim=1) + return thresholds.unsqueeze(1) diff --git a/zipvoice/models/zipvoice_distill.py b/zipvoice/models/zipvoice_distill.py new file mode 100644 index 0000000..29481bd --- /dev/null +++ b/zipvoice/models/zipvoice_distill.py @@ -0,0 +1,94 @@ +# Copyright 2024 Xiaomi Corp. (authors: Wei Kang +# Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List + +import torch + +from zipvoice.models.modules.solver import DistillEulerSolver +from zipvoice.models.modules.zipformer import TTSZipformer +from zipvoice.models.zipvoice import ZipVoice + + +class ZipVoiceDistill(ZipVoice): + """ZipVoice-Distill model.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + required_params = { + "feat_dim", + "fm_decoder_downsampling_factor", + "fm_decoder_num_layers", + "fm_decoder_cnn_module_kernel", + "fm_decoder_dim", + "fm_decoder_feedforward_dim", + "fm_decoder_num_heads", + "query_head_dim", + "pos_head_dim", + "value_head_dim", + "pos_dim", + "time_embed_dim", + } + + missing = [p for p in required_params if p not in kwargs] + if missing: + raise ValueError(f"Missing required parameters: {', '.join(missing)}") + + self.fm_decoder = TTSZipformer( + in_dim=kwargs["feat_dim"] * 3, + out_dim=kwargs["feat_dim"], + downsampling_factor=kwargs["fm_decoder_downsampling_factor"], + num_encoder_layers=kwargs["fm_decoder_num_layers"], + cnn_module_kernel=kwargs["fm_decoder_cnn_module_kernel"], + encoder_dim=kwargs["fm_decoder_dim"], + feedforward_dim=kwargs["fm_decoder_feedforward_dim"], + num_heads=kwargs["fm_decoder_num_heads"], + query_head_dim=kwargs["query_head_dim"], + pos_head_dim=kwargs["pos_head_dim"], + value_head_dim=kwargs["value_head_dim"], + pos_dim=kwargs["pos_dim"], + use_time_embed=True, + time_embed_dim=kwargs["time_embed_dim"], + use_guidance_scale_embed=True, + ) + self.solver = DistillEulerSolver(self, func_name="forward_fm_decoder") + + def forward( + self, + tokens: List[List[int]], + features: torch.Tensor, + features_lens: torch.Tensor, + noise: torch.Tensor, + speech_condition_mask: torch.Tensor, + t_start: float, + t_end: float, + num_step: int = 1, + guidance_scale: torch.Tensor = None, + ) -> torch.Tensor: + + return self.sample_intermediate( + tokens=tokens, + features=features, + features_lens=features_lens, + noise=noise, + speech_condition_mask=speech_condition_mask, + t_start=t_start, + t_end=t_end, + num_step=num_step, + guidance_scale=guidance_scale, + ) diff --git a/zipvoice/onnx_modeling.py b/zipvoice/onnx_modeling.py new file mode 100644 index 0000000..31f949b --- /dev/null +++ b/zipvoice/onnx_modeling.py @@ -0,0 +1,206 @@ +import argparse +import datetime as dt +import json +import logging +import os +from pathlib import Path +from typing import List, Tuple + +import numpy as np +import onnxruntime as ort +import torch +import torchaudio +from huggingface_hub import hf_hub_download +from lhotse.utils import fix_random_seed +from torch import Tensor, nn + +from zipvoice.bin.infer_zipvoice import get_vocoder +from zipvoice.models.modules.solver import get_time_steps +from zipvoice.tokenizer.tokenizer import ( + EmiliaTokenizer, + EspeakTokenizer, + LibriTTSTokenizer, + SimpleTokenizer, +) +from zipvoice.utils.common import AttributeDict, str2bool +from zipvoice.utils.feature import VocosFbank +from zipvoice.utils.infer import ( + add_punctuation, + chunk_tokens_punctuation, + cross_fade_concat, + load_prompt_wav, + remove_silence, + rms_norm, +) + +class OnnxModel: + def __init__( + self, + text_encoder_path: str, + fm_decoder_path: str, + num_thread: int = 1, + ): + session_opts = ort.SessionOptions() + session_opts.inter_op_num_threads = num_thread + session_opts.intra_op_num_threads = num_thread + + self.session_opts = session_opts + + self.init_text_encoder(text_encoder_path) + self.init_fm_decoder(fm_decoder_path) + + def init_text_encoder(self, model_path: str): + self.text_encoder = ort.InferenceSession( + model_path, + sess_options=self.session_opts, + providers=["CPUExecutionProvider"], + ) + + def init_fm_decoder(self, model_path: str): + self.fm_decoder = ort.InferenceSession( + model_path, + sess_options=self.session_opts, + providers=["CPUExecutionProvider"], + ) + meta = self.fm_decoder.get_modelmeta().custom_metadata_map + self.feat_dim = int(meta["feat_dim"]) + + def run_text_encoder( + self, + tokens: Tensor, + prompt_tokens: Tensor, + prompt_features_len: Tensor, + speed: Tensor, + ) -> Tuple[Tensor, Tensor]: + out = self.text_encoder.run( + [ + self.text_encoder.get_outputs()[0].name, + ], + { + self.text_encoder.get_inputs()[0].name: tokens.numpy(), + self.text_encoder.get_inputs()[1].name: prompt_tokens.numpy(), + self.text_encoder.get_inputs()[2].name: prompt_features_len.numpy(), + self.text_encoder.get_inputs()[3].name: speed.numpy(), + }, + ) + return torch.from_numpy(out[0]) + + def run_fm_decoder( + self, + t: Tensor, + x: Tensor, + text_condition: Tensor, + speech_condition: torch.Tensor, + guidance_scale: Tensor, + ) -> Tensor: + out = self.fm_decoder.run( + [ + self.fm_decoder.get_outputs()[0].name, + ], + { + self.fm_decoder.get_inputs()[0].name: t.numpy(), + self.fm_decoder.get_inputs()[1].name: x.numpy(), + self.fm_decoder.get_inputs()[2].name: text_condition.numpy(), + self.fm_decoder.get_inputs()[3].name: speech_condition.numpy(), + self.fm_decoder.get_inputs()[4].name: guidance_scale.numpy(), + }, + ) + return torch.from_numpy(out[0]) + +def sample( + model: OnnxModel, + tokens: List[List[int]], + prompt_tokens: List[List[int]], + prompt_features: Tensor, + speed: float = 1.3, + t_shift: float = 0.5, + guidance_scale: float = 1.0, + num_step: int = 16, + +) -> torch.Tensor: + # --- Preparation --- + assert len(tokens) == len(prompt_tokens) == 1 + tokens = torch.tensor(tokens, dtype=torch.int64) + prompt_tokens = torch.tensor(prompt_tokens, dtype=torch.int64) + prompt_features_len = torch.tensor(prompt_features.size(1), dtype=torch.int64) + speed = torch.tensor(speed, dtype=torch.float32) + + # Run text encoder + + text_condition = model.run_text_encoder( + tokens, prompt_tokens, prompt_features_len, speed + ) + batch_size, num_frames, _ = text_condition.shape + feat_dim = model.feat_dim + + # Get the time schedule + timesteps = get_time_steps( + t_start=0.0, + t_end=1.0, + num_step=num_step, + t_shift=t_shift, + ) + + # Initialize x with noise (x_0) + x = torch.randn(batch_size, num_frames, feat_dim) + speech_condition = torch.nn.functional.pad( + prompt_features, (0, 0, 0, num_frames - prompt_features.shape[1]) + ) + guidance_scale = torch.tensor(guidance_scale, dtype=torch.float32) + + # --- Sampling Loop --- + for step in range(num_step): + t_cur = timesteps[step] + t_next = timesteps[step + 1] + + # Predict velocity v + v = model.run_fm_decoder( + t=t_cur, + x=x, + text_condition=text_condition, + speech_condition=speech_condition, + guidance_scale=guidance_scale, + ) + + # Flow matching formula: x_t = (1 - t) * x_0 + t * x_1 + # Therefore: v = x_1 - x_0 + # This implies: + x_1_pred = x + (1.0 - t_cur) * v + x_0_pred = x - t_cur * v + + if step < num_step - 1: + # Anchor-based ODE update for the next step + x = (1.0 - t_next) * x_0_pred + t_next * x_1_pred + else: + # Final step: Snap directly to the predicted clean data (x_1) + x = x_1_pred + + # Remove the prompt portion from the generated sequence + x = x[:, prompt_features_len.item() :, :] + return x + +def generate_cpu(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, model, vocoder, tokenizer, num_step=4, guidance_scale=3.0, speed=1.0, t_shift=0.9, target_rms=0.1): + + tokens = tokenizer.texts_to_token_ids([text]) + speed = speed * 1.3 ## default is too slow + + pred_features = sample( + model=model, + tokens=tokens, + prompt_tokens=prompt_tokens, + prompt_features=prompt_features, + speed=speed, + t_shift=t_shift, + guidance_scale=guidance_scale, + num_step=num_step, + ) + + # Convert to waveform + pred_features = pred_features.permute(0, 2, 1) / 0.1 + wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1) + + # Volume matching + if prompt_rms < target_rms: + wav = wav * (prompt_rms / target_rms) + + return wav diff --git a/zipvoice/tokenizer/normalizer.py b/zipvoice/tokenizer/normalizer.py new file mode 100644 index 0000000..d295589 --- /dev/null +++ b/zipvoice/tokenizer/normalizer.py @@ -0,0 +1,170 @@ +import re +from abc import ABC, abstractmethod + +import cn2an +import inflect + + +class TextNormalizer(ABC): + """Abstract base class for text normalization, defining common interface.""" + + @abstractmethod + def normalize(self, text: str) -> str: + """Normalize text.""" + raise NotImplementedError + + +class EnglishTextNormalizer(TextNormalizer): + """ + A class to handle preprocessing of English text including normalization. Following: + https://github.com/espnet/espnet_tts_frontend/blob/master/tacotron_cleaner/cleaners.py + """ + + def __init__(self): + # List of (regular expression, replacement) pairs for abbreviations: + self._abbreviations = [ + (re.compile("\\b%s\\b" % x[0], re.IGNORECASE), x[1]) + for x in [ + ("mrs", "misess"), + ("mr", "mister"), + ("dr", "doctor"), + ("st", "saint"), + ("co", "company"), + ("jr", "junior"), + ("maj", "major"), + ("gen", "general"), + ("drs", "doctors"), + ("rev", "reverend"), + ("lt", "lieutenant"), + ("hon", "honorable"), + ("sgt", "sergeant"), + ("capt", "captain"), + ("esq", "esquire"), + ("ltd", "limited"), + ("col", "colonel"), + ("ft", "fort"), + ("etc", "et cetera"), + ("btw", "by the way"), + ] + ] + + self._inflect = inflect.engine() + self._comma_number_re = re.compile(r"([0-9][0-9\,]+[0-9])") + self._decimal_number_re = re.compile(r"([0-9]+\.[0-9]+)") + self._percent_number_re = re.compile(r"([0-9\.\,]*[0-9]+%)") + self._pounds_re = re.compile(r"£([0-9\,]*[0-9]+)") + self._dollars_re = re.compile(r"\$([0-9\.\,]*[0-9]+)") + self._fraction_re = re.compile(r"([0-9]+)/([0-9]+)") + self._ordinal_re = re.compile(r"[0-9]+(st|nd|rd|th)") + self._number_re = re.compile(r"[0-9]+") + self._whitespace_re = re.compile(r"\s+") + + def normalize(self, text: str) -> str: + """Custom pipeline for English text, + including number and abbreviation expansion.""" + text = self.expand_abbreviations(text) + text = self.normalize_numbers(text) + + return text + + def fraction_to_words(self, numerator, denominator): + if numerator == 1 and denominator == 2: + return " one half " + if numerator == 1 and denominator == 4: + return " one quarter " + if denominator == 2: + return " " + self._inflect.number_to_words(numerator) + " halves " + if denominator == 4: + return " " + self._inflect.number_to_words(numerator) + " quarters " + return ( + " " + + self._inflect.number_to_words(numerator) + + " " + + self._inflect.ordinal(self._inflect.number_to_words(denominator)) + + " " + ) + + def _remove_commas(self, m): + return m.group(1).replace(",", "") + + def _expand_dollars(self, m): + match = m.group(1) + parts = match.split(".") + if len(parts) > 2: + return " " + match + " dollars " # Unexpected format + dollars = int(parts[0]) if parts[0] else 0 + cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0 + if dollars and cents: + dollar_unit = "dollar" if dollars == 1 else "dollars" + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s, %s %s " % (dollars, dollar_unit, cents, cent_unit) + elif dollars: + dollar_unit = "dollar" if dollars == 1 else "dollars" + return " %s %s " % (dollars, dollar_unit) + elif cents: + cent_unit = "cent" if cents == 1 else "cents" + return " %s %s " % (cents, cent_unit) + else: + return " zero dollars " + + def _expand_fraction(self, m): + numerator = int(m.group(1)) + denominator = int(m.group(2)) + return self.fraction_to_words(numerator, denominator) + + def _expand_decimal_point(self, m): + return m.group(1).replace(".", " point ") + + def _expand_percent(self, m): + return m.group(1).replace("%", " percent ") + + def _expand_ordinal(self, m): + return " " + self._inflect.number_to_words(m.group(0)) + " " + + def _expand_number(self, m): + num = int(m.group(0)) + if num > 1000 and num < 3000: + if num == 2000: + return " two thousand " + elif num > 2000 and num < 2010: + return " two thousand " + self._inflect.number_to_words(num % 100) + " " + elif num % 100 == 0: + return " " + self._inflect.number_to_words(num // 100) + " hundred " + else: + return ( + " " + + self._inflect.number_to_words( + num, andword="", zero="oh", group=2 + ).replace(", ", " ") + + " " + ) + else: + return " " + self._inflect.number_to_words(num, andword="") + " " + + def normalize_numbers(self, text): + text = re.sub(self._comma_number_re, self._remove_commas, text) + text = re.sub(self._pounds_re, r"\1 pounds", text) + text = re.sub(self._dollars_re, self._expand_dollars, text) + text = re.sub(self._fraction_re, self._expand_fraction, text) + text = re.sub(self._decimal_number_re, self._expand_decimal_point, text) + text = re.sub(self._percent_number_re, self._expand_percent, text) + text = re.sub(self._ordinal_re, self._expand_ordinal, text) + text = re.sub(self._number_re, self._expand_number, text) + return text + + def expand_abbreviations(self, text): + for regex, replacement in self._abbreviations: + text = re.sub(regex, replacement, text) + return text + + +class ChineseTextNormalizer(TextNormalizer): + """ + A class to handle preprocessing of Chinese text including normalization. + """ + + def normalize(self, text: str) -> str: + """Normalize text.""" + # Convert numbers to Chinese + text = cn2an.transform(text, "an2cn") + return text diff --git a/zipvoice/tokenizer/tokenizer.py b/zipvoice/tokenizer/tokenizer.py new file mode 100644 index 0000000..b508176 --- /dev/null +++ b/zipvoice/tokenizer/tokenizer.py @@ -0,0 +1,648 @@ +# Copyright 2023-2024 Xiaomi Corp. (authors: Zengwei Yao +# Han Zhu, +# Wei Kang) +# +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import re +from abc import ABC, abstractmethod +from functools import reduce +from typing import Dict, List, Optional + +import jieba +from lhotse import CutSet +from pypinyin import Style, lazy_pinyin +from pypinyin.contrib.tone_convert import to_finals_tone3, to_initials + +from zipvoice.tokenizer.normalizer import ChineseTextNormalizer, EnglishTextNormalizer + +try: + from piper_phonemize import phonemize_espeak +except Exception as ex: + raise RuntimeError( + f"{ex}\nPlease run\n" + "pip install piper_phonemize -f \ + https://k2-fsa.github.io/icefall/piper_phonemize.html" + ) + +jieba.default_logger.setLevel(logging.INFO) + + +class Tokenizer(ABC): + """Abstract base class for tokenizers, defining common interface.""" + + @abstractmethod + def texts_to_token_ids(self, texts: List[str]) -> List[List[int]]: + """Convert list of texts to list of token id sequences.""" + raise NotImplementedError + + @abstractmethod + def texts_to_tokens(self, texts: List[str]) -> List[List[str]]: + """Convert list of texts to list of token sequences.""" + raise NotImplementedError + + @abstractmethod + def tokens_to_token_ids(self, tokens: List[List[str]]) -> List[List[int]]: + """Convert list of token sequences to list of token id sequences.""" + raise NotImplementedError + + +class SimpleTokenizer(Tokenizer): + """The simplpest tokenizer, treat every character as a token, + without text normalization. + """ + + def __init__(self, token_file: Optional[str] = None): + """ + Args: + tokens: the file that contains information that maps tokens to ids, + which is a text file with '{token}\t{token_id}' per line. + """ + # Parse token file + self.has_tokens = False + if token_file is None: + logging.debug( + "Initialize Tokenizer without tokens file, \ + will fail when map to ids." + ) + return + self.token2id: Dict[str, int] = {} + with open(token_file, "r", encoding="utf-8") as f: + for line in f.readlines(): + info = line.rstrip().split("\t") + token, id = info[0], int(info[1]) + assert token not in self.token2id, token + self.token2id[token] = id + self.pad_id = self.token2id["_"] # padding + self.vocab_size = len(self.token2id) + self.has_tokens = True + + def texts_to_token_ids( + self, + texts: List[str], + ) -> List[List[int]]: + return self.tokens_to_token_ids(self.texts_to_tokens(texts)) + + def texts_to_tokens( + self, + texts: List[str], + ) -> List[List[str]]: + tokens_list = [list(texts[i]) for i in range(len(texts))] + return tokens_list + + def tokens_to_token_ids( + self, + tokens_list: List[List[str]], + ) -> List[List[int]]: + assert self.has_tokens, "Please initialize Tokenizer with a tokens file." + + token_ids_list = [] + + for tokens in tokens_list: + token_ids = [] + for t in tokens: + if t not in self.token2id: + logging.debug(f"Skip OOV {t}") + continue + token_ids.append(self.token2id[t]) + + token_ids_list.append(token_ids) + + return token_ids_list + + +class EspeakTokenizer(Tokenizer): + """A simple tokenizer with Espeak g2p function.""" + + def __init__(self, token_file: Optional[str] = None, lang: str = "en-us"): + """ + Args: + tokens: the file that contains information that maps tokens to ids, + which is a text file with '{token}\t{token_id}' per line. + lang: the language identifier, see + https://github.com/rhasspy/espeak-ng/blob/master/docs/languages.md + """ + # Parse token file + self.has_tokens = False + self.lang = lang + if token_file is None: + logging.debug( + "Initialize Tokenizer without tokens file, \ + will fail when map to ids." + ) + return + self.token2id: Dict[str, int] = {} + with open(token_file, "r", encoding="utf-8") as f: + for line in f.readlines(): + info = line.rstrip().split("\t") + token, id = info[0], int(info[1]) + assert token not in self.token2id, token + self.token2id[token] = id + self.pad_id = self.token2id["_"] # padding + self.vocab_size = len(self.token2id) + self.has_tokens = True + + def g2p(self, text: str) -> List[str]: + try: + tokens = phonemize_espeak(text, self.lang) + tokens = reduce(lambda x, y: x + y, tokens) + return tokens + except Exception as ex: + logging.warning(f"Tokenization of {self.lang} texts failed: {ex}") + return [] + + def texts_to_token_ids( + self, + texts: List[str], + ) -> List[List[int]]: + return self.tokens_to_token_ids(self.texts_to_tokens(texts)) + + def texts_to_tokens( + self, + texts: List[str], + ) -> List[List[str]]: + tokens_list = [self.g2p(texts[i]) for i in range(len(texts))] + return tokens_list + + def tokens_to_token_ids( + self, + tokens_list: List[List[str]], + ) -> List[List[int]]: + assert self.has_tokens, "Please initialize Tokenizer with a tokens file." + + token_ids_list = [] + + for tokens in tokens_list: + token_ids = [] + for t in tokens: + if t not in self.token2id: + logging.debug(f"Skip OOV {t}") + continue + token_ids.append(self.token2id[t]) + + token_ids_list.append(token_ids) + + return token_ids_list + + +class EmiliaTokenizer(Tokenizer): + def __init__(self, token_file: Optional[str] = None, token_type="phone"): + """ + Args: + tokens: the file that contains information that maps tokens to ids, + which is a text file with '{token}\t{token_id}' per line. + """ + assert ( + token_type == "phone" + ), f"Only support phone tokenizer for Emilia, but get {token_type}." + + self.english_normalizer = EnglishTextNormalizer() + self.chinese_normalizer = ChineseTextNormalizer() + + self.has_tokens = False + if token_file is None: + logging.debug( + "Initialize Tokenizer without tokens file, \ + will fail when map to ids." + ) + return + self.token2id: Dict[str, int] = {} + with open(token_file, "r", encoding="utf-8") as f: + for line in f.readlines(): + info = line.rstrip().split("\t") + token, id = info[0], int(info[1]) + assert token not in self.token2id, token + self.token2id[token] = id + self.pad_id = self.token2id["_"] # padding + + self.vocab_size = len(self.token2id) + self.has_tokens = True + + def texts_to_token_ids( + self, + texts: List[str], + ) -> List[List[int]]: + return self.tokens_to_token_ids(self.texts_to_tokens(texts)) + + def preprocess_text( + self, + text: str, + ) -> str: + return self.map_punctuations(text) + + def texts_to_tokens( + self, + texts: List[str], + ) -> List[List[str]]: + for i in range(len(texts)): + # Text normalization + texts[i] = self.preprocess_text(texts[i]) + + phoneme_list = [] + for text in texts: + # now only en and ch + segments = self.get_segment(text) + all_phoneme = [] + for index in range(len(segments)): + seg = segments[index] + if seg[1] == "zh": + phoneme = self.tokenize_ZH(seg[0]) + elif seg[1] == "en": + phoneme = self.tokenize_EN(seg[0]) + elif seg[1] == "pinyin": + phoneme = self.tokenize_pinyin(seg[0]) + elif seg[1] == "tag": + phoneme = [seg[0]] + else: + logging.warning( + f"No English or Chinese characters found, \ + skipping segment of unknown language: {seg}" + ) + continue + all_phoneme += phoneme + phoneme_list.append(all_phoneme) + return phoneme_list + + def tokens_to_token_ids( + self, + tokens_list: List[List[str]], + ) -> List[List[int]]: + assert self.has_tokens, "Please initialize Tokenizer with a tokens file." + token_ids_list = [] + + for tokens in tokens_list: + token_ids = [] + for t in tokens: + if t not in self.token2id: + logging.debug(f"Skip OOV {t}") + continue + token_ids.append(self.token2id[t]) + + token_ids_list.append(token_ids) + + return token_ids_list + + def tokenize_ZH(self, text: str) -> List[str]: + try: + text = self.chinese_normalizer.normalize(text) + segs = list(jieba.cut(text)) + full = lazy_pinyin( + segs, + style=Style.TONE3, + tone_sandhi=True, + neutral_tone_with_five=True, + ) + phones = [] + for x in full: + # valid pinyin (in tone3 style) is alphabet + 1 number in [1-5]. + if not (x[0:-1].isalpha() and x[-1] in ("1", "2", "3", "4", "5")): + phones.append(x) + continue + else: + phones.extend(self.seperate_pinyin(x)) + return phones + except Exception as ex: + logging.warning(f"Tokenization of Chinese texts failed: {ex}") + return [] + + def tokenize_EN(self, text: str) -> List[str]: + try: + text = self.english_normalizer.normalize(text) + tokens = phonemize_espeak(text, "en-us") + tokens = reduce(lambda x, y: x + y, tokens) + return tokens + except Exception as ex: + logging.warning(f"Tokenization of English texts failed: {ex}") + return [] + + def tokenize_pinyin(self, text: str) -> List[str]: + try: + assert text.startswith("<") and text.endswith(">") + text = text.lstrip("<").rstrip(">") + # valid pinyin (in tone3 style) is alphabet + 1 number in [1-5]. + if not (text[0:-1].isalpha() and text[-1] in ("1", "2", "3", "4", "5")): + logging.warning( + f"Strings enclosed with <> should be pinyin, \ + but got: {text}. Skipped it. " + ) + return [] + else: + return self.seperate_pinyin(text) + except Exception as ex: + logging.warning(f"Tokenize pinyin failed: {ex}") + return [] + + def seperate_pinyin(self, text: str) -> List[str]: + """ + Separate pinyin into initial and final + """ + pinyins = [] + initial = to_initials(text, strict=False) + # don't want to share tokens with espeak tokens, + # so use tone3 style + final = to_finals_tone3( + text, + strict=False, + neutral_tone_with_five=True, + ) + if initial != "": + # don't want to share tokens with espeak tokens, + # so add a '0' after each initial + pinyins.append(initial + "0") + if final != "": + pinyins.append(final) + return pinyins + + def map_punctuations(self, text): + text = text.replace("，", ",") + text = text.replace("。", ".") + text = text.replace("！", "!") + text = text.replace("？", "?") + text = text.replace("；", ";") + text = text.replace("：", ":") + text = text.replace("、", ",") + text = text.replace("‘", "'") + text = text.replace("“", '"') + text = text.replace("”", '"') + text = text.replace("’", "'") + text = text.replace("⋯", "…") + text = text.replace("···", "…") + text = text.replace("・・・", "…") + text = text.replace("...", "…") + return text + + def get_segment(self, text: str) -> List[str]: + """ + Split a text into segments based on language types + (Chinese, English, Pinyin, tags, etc.) + + Args: + text (str): Input text to be segmented + + Returns: + List[str]: Segmented text parts with their language types + + Example: + Input: 我们是小米人,是吗? Yes I think so!霍...啦啦啦 + Output: [('我们是小米人,是吗? ', 'zh'), + ('Yes I think so!', 'en'), ('霍...啦啦啦', 'zh')] + """ + # Stores the final segmented parts and their language types + segments = [] + # Stores the language type of each character in the input text + types = [] + temp_seg = "" + temp_lang = "" + + # Each part is a character, or a special string enclosed in <> and [] + # <> denotes pinyin string, [] denotes other special strings. + _part_pattern = re.compile(r"[<[].*?[>\]]|.") + text = _part_pattern.findall(text) + + for i, part in enumerate(text): + if self.is_chinese(part) or self.is_pinyin(part): + types.append("zh") + elif self.is_alphabet(part): + types.append("en") + else: + types.append("other") + + assert len(types) == len(text) + + for i in range(len(types)): + # find the first char of the seg + if i == 0: + temp_seg += text[i] + temp_lang = types[i] + else: + if temp_lang == "other": + temp_seg += text[i] + temp_lang = types[i] + else: + if types[i] in [temp_lang, "other"]: + temp_seg += text[i] + else: + segments.append((temp_seg, temp_lang)) + temp_seg = text[i] + temp_lang = types[i] + + segments.append((temp_seg, temp_lang)) + + # Handle "pinyin" and "tag" types + segments = self.split_segments(segments) + return segments + + def split_segments(self, segments): + """ + split segments into smaller parts if special strings enclosed by [] or <> + are found, where <> denotes pinyin strings, [] denotes other special strings. + + Args: + segments (list): A list of tuples where each tuple contains: + - temp_seg (str): The text segment to be split. + - temp_lang (str): The language code associated with the segment. + + Returns: + list: A list of smaller segments. + """ + result = [] + for temp_seg, temp_lang in segments: + parts = re.split(r"([<[].*?[>\]])", temp_seg) + for part in parts: + if not part: + continue + if self.is_pinyin(part): + result.append((part, "pinyin")) + elif self.is_tag(part): + result.append((part, "tag")) + else: + result.append((part, temp_lang)) + return result + + def is_chinese(self, char: str) -> bool: + if char >= "\u4e00" and char <= "\u9fa5": + return True + else: + return False + + def is_alphabet(self, char: str) -> bool: + if (char >= "\u0041" and char <= "\u005a") or ( + char >= "\u0061" and char <= "\u007a" + ): + return True + else: + return False + + def is_pinyin(self, part: str) -> bool: + if part.startswith("<") and part.endswith(">"): + return True + else: + return False + + def is_tag(self, part: str) -> bool: + if part.startswith("[") and part.endswith("]"): + return True + else: + return False + + +class DialogTokenizer(EmiliaTokenizer): + def __init__(self, token_file: Optional[str] = None, token_type="phone"): + super().__init__(token_file=token_file, token_type=token_type) + if token_file: + self.spk_a_id = self.token2id["[S1]"] + self.spk_b_id = self.token2id["[S2]"] + + def preprocess_text( + self, + text: str, + ) -> str: + text = re.sub(r"\s*(\[S[12]\])\s*", r"\1", text) + text = self.map_punctuations(text) + return text + + +class LibriTTSTokenizer(Tokenizer): + def __init__(self, token_file: Optional[str] = None, token_type="char"): + """ + Args: + type: the type of tokenizer, e.g., bpe, char, phone. + tokens: the file that contains information that maps tokens to ids, + which is a text file with '{token}\t{token_id}' per line if type is + char or phone, otherwise it is a bpe_model file. + """ + self.type = token_type + assert token_type in ["bpe", "char", "phone"] + try: + import tacotron_cleaner.cleaners + except Exception as ex: + raise RuntimeError(f"{ex}\nPlease run\n" "pip install espnet_tts_frontend") + + self.normalize = tacotron_cleaner.cleaners.custom_english_cleaners + + self.has_tokens = False + if token_file is None: + logging.debug( + "Initialize Tokenizer without tokens file, \ + will fail when map to ids." + ) + return + if token_type == "bpe": + import sentencepiece as spm + + self.sp = spm.SentencePieceProcessor() + self.sp.load(token_file) + self.pad_id = self.sp.piece_to_id("") + self.vocab_size = self.sp.get_piece_size() + else: + self.token2id: Dict[str, int] = {} + with open(token_file, "r", encoding="utf-8") as f: + for line in f.readlines(): + info = line.rstrip().split("\t") + token, id = info[0], int(info[1]) + assert token not in self.token2id, token + self.token2id[token] = id + self.pad_id = self.token2id["_"] # padding + self.vocab_size = len(self.token2id) + self.has_tokens = True + + def texts_to_token_ids( + self, + texts: List[str], + ) -> List[List[int]]: + if self.type == "bpe": + for i in range(len(texts)): + texts[i] = self.normalize(texts[i]) + return self.sp.encode(texts) + else: + return self.tokens_to_token_ids(self.texts_to_tokens(texts)) + + def texts_to_tokens( + self, + texts: List[str], + ) -> List[List[str]]: + for i in range(len(texts)): + texts[i] = self.normalize(texts[i]) + + if self.type == "char": + tokens_list = [list(texts[i]) for i in range(len(texts))] + elif self.type == "phone": + tokens_list = [ + phonemize_espeak(texts[i].lower(), "en-us") for i in range(len(texts)) + ] + elif self.type == "bpe": + tokens_list = self.sp.encode(texts, out_type=str) + + return tokens_list + + def tokens_to_token_ids( + self, + tokens_list: List[List[str]], + ) -> List[List[int]]: + assert self.has_tokens, "Please initialize Tokenizer with a tokens file." + + assert self.type != "bpe", "BPE tokenizer does not support this function." + + token_ids_list = [] + + for tokens in tokens_list: + token_ids = [] + for t in tokens: + if t not in self.token2id: + logging.debug(f"Skip OOV {t}") + continue + token_ids.append(self.token2id[t]) + + token_ids_list.append(token_ids) + + return token_ids_list + + +def add_tokens(cut_set: CutSet, tokenizer: str, lang: str): + if tokenizer == "emilia": + tokenizer = EmiliaTokenizer() + elif tokenizer == "espeak": + tokenizer = EspeakTokenizer(lang=lang) + elif tokenizer == "dialog": + tokenizer = DialogTokenizer() + elif tokenizer == "libritts": + tokenizer = LibriTTSTokenizer() + elif tokenizer == "simple": + tokenizer = SimpleTokenizer() + else: + raise ValueError(f"Unsupported tokenizer: {tokenizer}.") + + def _prepare_cut(cut): + # Each cut only contains one supervision + assert len(cut.supervisions) == 1, (len(cut.supervisions), cut) + text = cut.supervisions[0].text + tokens = tokenizer.texts_to_tokens([text])[0] + cut.supervisions[0].tokens = tokens + return cut + + cut_set = cut_set.map(_prepare_cut) + return cut_set + + +if __name__ == "__main__": + text = ( + "我们是5年小米人,是吗? Yes I think so! " + "mr king, 5 years, from 2019 to 2024." + "霍...啦啦啦超过90%的人...?!9204" + ) + tokenizer = EmiliaTokenizer() + tokens = tokenizer.texts_to_tokens([text]) + print(f"tokens: {'|'.join(tokens[0])}") diff --git a/zipvoice/utils/checkpoint.py b/zipvoice/utils/checkpoint.py new file mode 100644 index 0000000..3ac150c --- /dev/null +++ b/zipvoice/utils/checkpoint.py @@ -0,0 +1,570 @@ +# Copyright 2021-2025 Xiaomi Corporation (authors: Fangjun Kuang, +# Zengwei Yao) +# +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import glob +import logging +import os +import re +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +import torch +import torch.nn as nn +from lhotse.dataset.sampling.base import CutSampler +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.optim import Optimizer + +from zipvoice.utils.common import AttributeDict, GradScaler + +# use duck typing for LRScheduler since we have different possibilities, see +# our class LRScheduler. +LRSchedulerType = object + + +def save_checkpoint( + filename: Path, + model: Union[nn.Module, DDP], + model_avg: Optional[nn.Module] = None, + model_ema: Optional[nn.Module] = None, + params: Optional[Dict[str, Any]] = None, + optimizer: Optional[Optimizer] = None, + scheduler: Optional[LRSchedulerType] = None, + scaler: Optional[GradScaler] = None, + sampler: Optional[CutSampler] = None, + rank: int = 0, +) -> None: + """Save training information to a file. + + Args: + filename: + The checkpoint filename. + model: + The model to be saved. We only save its `state_dict()`. + model_avg: + The stored model averaged from the start of training. + model_ema: + The EMA version of model. + params: + User defined parameters, e.g., epoch, loss. + optimizer: + The optimizer to be saved. We only save its `state_dict()`. + scheduler: + The scheduler to be saved. We only save its `state_dict()`. + scalar: + The GradScaler to be saved. We only save its `state_dict()`. + sampler: + The sampler used in the labeled training dataset. We only + save its `state_dict()`. + rank: + Used in DDP. We save checkpoint only for the node whose + rank is 0. + Returns: + Return None. + """ + if rank != 0: + return + + logging.info(f"Saving checkpoint to {filename}") + + if isinstance(model, DDP): + model = model.module + + checkpoint = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict() if optimizer is not None else None, + "scheduler": scheduler.state_dict() if scheduler is not None else None, + "grad_scaler": scaler.state_dict() if scaler is not None else None, + "sampler": sampler.state_dict() if sampler is not None else None, + } + + if model_avg is not None: + checkpoint["model_avg"] = model_avg.to(torch.float32).state_dict() + if model_ema is not None: + checkpoint["model_ema"] = model_ema.to(torch.float32).state_dict() + + if params: + for k, v in params.items(): + assert k not in checkpoint + checkpoint[k] = v + + torch.save(checkpoint, filename) + + +def load_checkpoint( + filename: Path, + model: Optional[nn.Module] = None, + model_avg: Optional[nn.Module] = None, + model_ema: Optional[nn.Module] = None, + strict: bool = False, +) -> Dict[str, Any]: + logging.info(f"Loading checkpoint from {filename}") + checkpoint = torch.load(filename, map_location="cpu", weights_only=False) + + if model is not None: + + if next(iter(checkpoint["model"])).startswith("module."): + logging.debug("Loading checkpoint saved by DDP") + dst_state_dict = model.state_dict() + src_state_dict = checkpoint["model"] + for key in dst_state_dict.keys(): + src_key = "{}.{}".format("module", key) + dst_state_dict[key] = src_state_dict.pop(src_key) + assert len(src_state_dict) == 0 + model.load_state_dict(dst_state_dict, strict=strict) + else: + logging.debug("Loading checkpoint") + model.load_state_dict(checkpoint["model"], strict=strict) + + checkpoint.pop("model") + + if model_avg is not None and "model_avg" in checkpoint: + logging.info("Loading averaged model") + model_avg.load_state_dict(checkpoint["model_avg"], strict=strict) + checkpoint.pop("model_avg") + + if model_ema is not None and "model_ema" in checkpoint: + logging.info("Loading ema model") + model_ema.load_state_dict(checkpoint["model_ema"], strict=strict) + checkpoint.pop("model_ema") + + return checkpoint + + +def load_checkpoint_extend_vocab_size( + filename: Path, extend_size: int, model: nn.Module, strict: bool = True +) -> Dict[str, Any]: + logging.info(f"Loading checkpoint from {filename}") + checkpoint = torch.load(filename, map_location="cpu", weights_only=False) + + if model is not None: + if next(iter(checkpoint["model"])).startswith("module."): + logging.info("Loading checkpoint saved by DDP") + dst_state_dict = model.state_dict() + src_state_dict = checkpoint["model"] + for key in dst_state_dict.keys(): + src_key = "{}.{}".format("module", key) + dst_state_dict[key] = src_state_dict.pop(src_key) + assert len(src_state_dict) == 0 + else: + logging.info("Loading checkpoint") + dst_state_dict = checkpoint["model"] + dst_state_dict["spk_embed.weight"] = model.state_dict()["spk_embed.weight"] + embed_weight = model.state_dict()["embed.weight"] + embed_weight[:-extend_size, :] = dst_state_dict["embed.weight"] + dst_state_dict["embed.weight"] = embed_weight + + model.load_state_dict(dst_state_dict, strict=strict) + + +def load_checkpoint_copy_proj_three_channel_alter( + filename: Path, + in_proj_key: str, + out_proj_key: str, + dim: int, + model: nn.Module, +) -> Dict[str, Any]: + logging.info(f"Loading checkpoint from {filename}") + checkpoint = torch.load(filename, map_location="cpu", weights_only=False) + + if model is not None: + if next(iter(checkpoint["model"])).startswith("module."): + logging.info("Loading checkpoint saved by DDP") + + dst_state_dict = dict() + src_state_dict = checkpoint["model"] + for key in src_state_dict.keys(): + dst_state_dict[key.lstrip("module.")] = src_state_dict.pop(key) + assert len(src_state_dict) == 0 + else: + logging.info("Loading checkpoint") + dst_state_dict = checkpoint["model"] + keys = list(dst_state_dict.keys()) + for key in keys: + if in_proj_key in key: + if "weight" in key: + weight = dst_state_dict.pop(key) + dst_state_dict[key.replace("weight", "0.weight")] = torch.cat( + [ + weight[:, :dim] / 2, + weight[:, :dim] / 2, + weight[:, dim : dim * 2], + weight[:, dim * 2 :] / 2, + weight[:, dim * 2 :] / 2, + ], + dim=-1, + ) + dst_state_dict[key.replace("weight", "1.weight")] = weight + if "bias" in key: + bias = dst_state_dict.pop(key) + dst_state_dict[key.replace("bias", "0.bias")] = bias + dst_state_dict[key.replace("bias", "1.bias")] = bias + if out_proj_key in key: + if "weight" in key: + weight = dst_state_dict.pop(key) + dst_state_dict[key.replace("weight", "0.weight")] = torch.cat( + [weight, weight], dim=0 + ) + dst_state_dict[key.replace("weight", "1.weight")] = weight + elif "bias" in key: + bias = dst_state_dict.pop(key) + dst_state_dict[key.replace("bias", "0.bias")] = torch.cat( + [bias, bias], dim=0 + ) + dst_state_dict[key.replace("bias", "1.bias")] = bias + + model.load_state_dict(dst_state_dict, strict=True) + + +def find_checkpoints(out_dir: Path, iteration: int = 0) -> List[str]: + """Find all available checkpoints in a directory. + + The checkpoint filenames have the form: `checkpoint-xxx.pt` + where xxx is a numerical value. + + Assume you have the following checkpoints in the folder `foo`: + + - checkpoint-1.pt + - checkpoint-20.pt + - checkpoint-300.pt + - checkpoint-4000.pt + + Case 1 (Return all checkpoints):: + + find_checkpoints(out_dir='foo') + + Case 2 (Return checkpoints newer than checkpoint-20.pt, i.e., + checkpoint-4000.pt, checkpoint-300.pt, and checkpoint-20.pt) + + find_checkpoints(out_dir='foo', iteration=20) + + Case 3 (Return checkpoints older than checkpoint-20.pt, i.e., + checkpoint-20.pt, checkpoint-1.pt):: + + find_checkpoints(out_dir='foo', iteration=-20) + + Args: + out_dir: + The directory where to search for checkpoints. + iteration: + If it is 0, return all available checkpoints. + If it is positive, return the checkpoints whose iteration number is + greater than or equal to `iteration`. + If it is negative, return the checkpoints whose iteration number is + less than or equal to `-iteration`. + Returns: + Return a list of checkpoint filenames, sorted in descending + order by the numerical value in the filename. + """ + checkpoints = list(glob.glob(f"{out_dir}/checkpoint-[0-9]*.pt")) + pattern = re.compile(r"checkpoint-([0-9]+).pt") + iter_checkpoints = [] + for c in checkpoints: + result = pattern.search(c) + if not result: + logging.warn(f"Invalid checkpoint filename {c}") + continue + + iter_checkpoints.append((int(result.group(1)), c)) + + # iter_checkpoints is a list of tuples. Each tuple contains + # two elements: (iteration_number, checkpoint-iteration_number.pt) + + iter_checkpoints = sorted(iter_checkpoints, reverse=True, key=lambda x: x[0]) + if iteration >= 0: + ans = [ic[1] for ic in iter_checkpoints if ic[0] >= iteration] + else: + ans = [ic[1] for ic in iter_checkpoints if ic[0] <= -iteration] + + return ans + + +def average_checkpoints_with_averaged_model( + filename_start: str, + filename_end: str, + device: torch.device = torch.device("cpu"), +) -> Dict[str, torch.Tensor]: + """Average model parameters over the range with given + start model (excluded) and end model. + + Let start = batch_idx_train of model-start; + end = batch_idx_train of model-end; + interval = end - start. + Then the average model over range from start (excluded) to end is + (1) avg = (model_end * end - model_start * start) / interval. + It can be written as + (2) avg = model_end * weight_end + model_start * weight_start, + where weight_end = end / interval, + weight_start = -start / interval = 1 - weight_end. + Since the terms `weight_end` and `weight_start` would be large + if the model has been trained for lots of batches, which would cause + overflow when multiplying the model parameters. + To avoid this, we rewrite (2) as: + (3) avg = (model_end + model_start * (weight_start / weight_end)) + * weight_end + + The model index could be epoch number or iteration number. + + Args: + filename_start: + Checkpoint filename of the start model. We assume it + is saved by :func:`save_checkpoint`. + filename_end: + Checkpoint filename of the end model. We assume it + is saved by :func:`save_checkpoint`. + device: + Move checkpoints to this device before averaging. + """ + state_dict_start = torch.load( + filename_start, map_location=device, weights_only=False + ) + state_dict_end = torch.load(filename_end, map_location=device, weights_only=False) + + average_period = state_dict_start["average_period"] + + batch_idx_train_start = state_dict_start["batch_idx_train"] + batch_idx_train_start = (batch_idx_train_start // average_period) * average_period + batch_idx_train_end = state_dict_end["batch_idx_train"] + batch_idx_train_end = (batch_idx_train_end // average_period) * average_period + interval = batch_idx_train_end - batch_idx_train_start + assert interval > 0, interval + weight_end = batch_idx_train_end / interval + weight_start = 1 - weight_end + + model_end = state_dict_end["model_avg"] + model_start = state_dict_start["model_avg"] + avg = model_end + + # scale the weight to avoid overflow + average_state_dict( + state_dict_1=avg, + state_dict_2=model_start, + weight_1=1.0, + weight_2=weight_start / weight_end, + scaling_factor=weight_end, + ) + + return avg + + +def remove_checkpoints( + out_dir: Path, + topk: int, + rank: int = 0, +): + """Remove checkpoints from the given directory. + + We assume that checkpoint filename has the form `checkpoint-xxx.pt` + where xxx is a number, representing the number of processed batches + when saving that checkpoint. We sort checkpoints by filename and keep + only the `topk` checkpoints with the highest `xxx`. + + Args: + out_dir: + The directory containing checkpoints to be removed. + topk: + Number of checkpoints to keep. + rank: + If using DDP for training, it is the rank of the current node. + Use 0 if no DDP is used for training. + """ + assert topk >= 1, topk + if rank != 0: + return + checkpoints = find_checkpoints(out_dir) + + if len(checkpoints) == 0: + logging.warn(f"No checkpoints found in {out_dir}") + return + + if len(checkpoints) <= topk: + return + + to_remove = checkpoints[topk:] + for c in to_remove: + os.remove(c) + + +def resume_checkpoint( + params: AttributeDict, + model: nn.Module, + model_avg: nn.Module, + model_ema: Optional[nn.Module] = None, +) -> Optional[Dict[str, Any]]: + """Load checkpoint from file. + + If params.start_epoch is larger than 1, it will load the checkpoint from + `params.start_epoch - 1`. + + Apart from loading state dict for `model` and `optimizer` it also updates + `best_train_epoch`, `best_train_loss`, `best_valid_epoch`, + and `best_valid_loss` in `params`. + + Args: + params: + The return value of :func:`get_params`. + model: + The training model. + Returns: + Return a dict containing previously saved training info. + """ + filename = params.exp_dir / f"epoch-{params.start_epoch - 1}.pt" + + assert filename.is_file(), f"{filename} does not exist!" + + saved_params = load_checkpoint( + filename, + model=model, + model_avg=model_avg, + model_ema=model_ema, + strict=True, + ) + + if params.start_epoch > 1: + keys = [ + "best_train_epoch", + "best_valid_epoch", + "batch_idx_train", + "best_train_loss", + "best_valid_loss", + ] + for k in keys: + params[k] = saved_params[k] + + return saved_params + + +def average_state_dict( + state_dict_1: Dict[str, torch.Tensor], + state_dict_2: Dict[str, torch.Tensor], + weight_1: float, + weight_2: float, + scaling_factor: float = 1.0, +) -> Dict[str, torch.Tensor]: + """Average two state_dict with given weights: + state_dict_1 = (state_dict_1 * weight_1 + state_dict_2 * weight_2) + * scaling_factor + It is an in-place operation on state_dict_1 itself. + """ + # Identify shared parameters. Two parameters are said to be shared + # if they have the same data_ptr + uniqued: Dict[int, str] = dict() + for k, v in state_dict_1.items(): + v_data_ptr = v.data_ptr() + if v_data_ptr in uniqued: + continue + uniqued[v_data_ptr] = k + + uniqued_names = list(uniqued.values()) + for k in uniqued_names: + v = state_dict_1[k] + if torch.is_floating_point(v): + v *= weight_1 + v += state_dict_2[k].to(device=state_dict_1[k].device) * weight_2 + v *= scaling_factor + + +def update_averaged_model( + params: Dict[str, torch.Tensor], + model_cur: Union[nn.Module, DDP], + model_avg: nn.Module, +) -> None: + """Update the averaged model: + model_avg = model_cur * (average_period / batch_idx_train) + + model_avg * ((batch_idx_train - average_period) / batch_idx_train) + + Args: + params: + User defined parameters, e.g., epoch, loss. + model_cur: + The current model. + model_avg: + The averaged model to be updated. + """ + weight_cur = params.average_period / params.batch_idx_train + weight_avg = 1 - weight_cur + + if isinstance(model_cur, DDP): + model_cur = model_cur.module + + cur = model_cur.state_dict() + avg = model_avg.state_dict() + + average_state_dict( + state_dict_1=avg, + state_dict_2=cur, + weight_1=weight_avg, + weight_2=weight_cur, + ) + + +def save_checkpoint_with_global_batch_idx( + out_dir: Path, + global_batch_idx: int, + model: Union[nn.Module, DDP], + model_avg: Optional[nn.Module] = None, + params: Optional[Dict[str, Any]] = None, + optimizer: Optional[Optimizer] = None, + scheduler: Optional[LRSchedulerType] = None, + scaler: Optional[GradScaler] = None, + sampler: Optional[CutSampler] = None, + rank: int = 0, +): + """Save training info after processing given number of batches. + + Args: + out_dir: + The directory to save the checkpoint. + global_batch_idx: + The number of batches processed so far from the very start of the + training. The saved checkpoint will have the following filename: + + f'out_dir / checkpoint-{global_batch_idx}.pt' + model: + The neural network model whose `state_dict` will be saved in the + checkpoint. + model_avg: + The stored model averaged from the start of training. + params: + A dict of training configurations to be saved. + optimizer: + The optimizer used in the training. Its `state_dict` will be saved. + scheduler: + The learning rate scheduler used in the training. Its `state_dict` will + be saved. + scaler: + The scaler used for mix precision training. Its `state_dict` will + be saved. + sampler: + The sampler used in the training dataset. + rank: + The rank ID used in DDP training of the current node. Set it to 0 + if DDP is not used. + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + filename = out_dir / f"checkpoint-{global_batch_idx}.pt" + save_checkpoint( + filename=filename, + model=model, + model_avg=model_avg, + params=params, + optimizer=optimizer, + scheduler=scheduler, + scaler=scaler, + sampler=sampler, + rank=rank, + ) diff --git a/zipvoice/utils/common.py b/zipvoice/utils/common.py new file mode 100644 index 0000000..75a33d9 --- /dev/null +++ b/zipvoice/utils/common.py @@ -0,0 +1,670 @@ +import argparse +import collections +import json +import logging +import os +import socket +import subprocess +import sys +import warnings +from collections import defaultdict +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Tuple, Union + +import torch +from packaging import version +from torch import distributed as dist +from torch import nn +from torch.nn.parallel import DistributedDataParallel as DDP +from torch.utils.tensorboard import SummaryWriter + + +if hasattr(torch.amp, "GradScaler"): + from torch.amp import GradScaler +else: + from torch.cuda.amp import GradScaler + +Pathlike = Union[str, Path] + + +class AttributeDict(dict): + def __getattr__(self, key): + if key in self: + return self[key] + raise AttributeError(f"No such attribute '{key}'") + + def __setattr__(self, key, value): + self[key] = value + + def __delattr__(self, key): + if key in self: + del self[key] + return + raise AttributeError(f"No such attribute '{key}'") + + def __str__(self, indent: int = 2): + tmp = {} + for k, v in self.items(): + # PosixPath is ont JSON serializable + if isinstance(v, (Path, torch.device, torch.dtype)): + v = str(v) + tmp[k] = v + return json.dumps(tmp, indent=indent, sort_keys=True) + + +class MetricsTracker(collections.defaultdict): + def __init__(self): + # Passing the type 'int' to the base-class constructor + # makes undefined items default to int() which is zero. + # This class will play a role as metrics tracker. + # It can record many metrics, including but not limited to loss. + super(MetricsTracker, self).__init__(int) + + def __add__(self, other: "MetricsTracker") -> "MetricsTracker": + ans = MetricsTracker() + for k, v in self.items(): + ans[k] = v + for k, v in other.items(): + if v - v == 0: + ans[k] = ans[k] + v + return ans + + def __mul__(self, alpha: float) -> "MetricsTracker": + ans = MetricsTracker() + for k, v in self.items(): + ans[k] = v * alpha + return ans + + def __str__(self) -> str: + ans_frames = "" + ans_utterances = "" + for k, v in self.norm_items(): + norm_value = "%.4g" % v + if "utt_" not in k: + ans_frames += str(k) + "=" + str(norm_value) + ", " + else: + ans_utterances += str(k) + "=" + str(norm_value) + if k == "utt_duration": + ans_utterances += " frames, " + elif k == "utt_pad_proportion": + ans_utterances += ", " + else: + raise ValueError(f"Unexpected key: {k}") + frames = "%.2f" % self["frames"] + ans_frames += "over " + str(frames) + " frames. " + if ans_utterances != "": + utterances = "%.2f" % self["utterances"] + ans_utterances += "over " + str(utterances) + " utterances." + + return ans_frames + ans_utterances + + def norm_items(self) -> List[Tuple[str, float]]: + """ + Returns a list of pairs, like: + [('ctc_loss', 0.1), ('att_loss', 0.07)] + """ + num_frames = self["frames"] if "frames" in self else 1 + num_utterances = self["utterances"] if "utterances" in self else 1 + ans = [] + for k, v in self.items(): + if k == "frames" or k == "utterances": + continue + norm_value = ( + float(v) / num_frames if "utt_" not in k else float(v) / num_utterances + ) + ans.append((k, norm_value)) + return ans + + def reduce(self, device): + """ + Reduce using torch.distributed, which I believe ensures that + all processes get the total. + """ + keys = sorted(self.keys()) + s = torch.tensor([float(self[k]) for k in keys], device=device) + dist.all_reduce(s, op=dist.ReduceOp.SUM) + for k, v in zip(keys, s.cpu().tolist()): + self[k] = v + + def write_summary( + self, + tb_writer: SummaryWriter, + prefix: str, + batch_idx: int, + ) -> None: + """Add logging information to a TensorBoard writer. + + Args: + tb_writer: a TensorBoard writer + prefix: a prefix for the name of the loss, e.g. "train/valid_", + or "train/current_" + batch_idx: The current batch index, used as the x-axis of the plot. + """ + for k, v in self.norm_items(): + tb_writer.add_scalar(prefix + k, v, batch_idx) + + +@contextmanager +def torch_autocast(device_type="cuda", **kwargs): + """ + To fix the following warnings: + FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. + Please use `torch.amp.autocast('cuda', args...)` instead. + with torch.cuda.amp.autocast(enabled=False): + """ + if version.parse(torch.__version__) >= version.parse("2.3.0"): + # Use new unified API + with torch.amp.autocast(device_type=device_type, **kwargs): + yield + else: + # Suppress deprecation warning and use old CUDA-specific autocast + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=FutureWarning) + with torch.cuda.amp.autocast(**kwargs): + yield + + +def create_grad_scaler(device="cuda", **kwargs): + """ + Creates a GradScaler compatible with both torch < 2.3.0 and >= 2.3.0. + Accepts all kwargs like: enabled, init_scale, growth_factor, etc. + + FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. + Please use `torch.amp.GradScaler('cuda', args...)` instead. + """ + if version.parse(torch.__version__) >= version.parse("2.3.0"): + from torch.amp import GradScaler + + return GradScaler(device=device, **kwargs) + else: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=FutureWarning) + return torch.cuda.amp.GradScaler(**kwargs) + + +def setup_dist( + rank=None, + world_size=None, + master_port=None, + use_ddp_launch=False, + master_addr=None, +): + """ + rank and world_size are used only if use_ddp_launch is False. + """ + if "MASTER_ADDR" not in os.environ: + os.environ["MASTER_ADDR"] = ( + "localhost" if master_addr is None else str(master_addr) + ) + + if "MASTER_PORT" not in os.environ: + os.environ["MASTER_PORT"] = "12354" if master_port is None else str(master_port) + + if use_ddp_launch is False: + dist.init_process_group("nccl", rank=rank, world_size=world_size) + torch.cuda.set_device(rank) + else: + dist.init_process_group("nccl") + + +def cleanup_dist(): + dist.destroy_process_group() + + +def prepare_input( + params: AttributeDict, + batch: dict, + device: torch.device, + return_tokens: bool = True, + return_feature: bool = True, + return_audio: bool = False, +): + """ + Parse the features and targets of the current batch. + Args: + params: + It is returned by :func:`get_params`. + batch: + It is the return value from iterating + `lhotse.dataset.K2SpeechRecognitionDataset`. See its documentation + for the format of the `batch`. + device: + The device of Tensor. + """ + return_list = [] + + if return_tokens: + return_list += [batch["tokens"]] + + if return_feature: + features = batch["features"].to(device) + features_lens = batch["features_lens"].to(device) + return_list += [features * params.feat_scale, features_lens] + + if return_audio: + return_list += [batch["audio"], batch["audio_lens"]] + + return return_list + + +def prepare_avg_tokens_durations(features_lens, tokens_lens): + tokens_durations = [] + for i in range(len(features_lens)): + utt_duration = features_lens[i] + avg_token_duration = utt_duration // tokens_lens[i] + tokens_durations.append([avg_token_duration] * tokens_lens[i]) + return tokens_durations + + +def pad_labels(y: List[List[int]], pad_id: int, device: torch.device): + """ + Pad the transcripts to the same length with zeros. + + Args: + y: the transcripts, which is a list of a list + + Returns: + Return a Tensor of padded transcripts. + """ + y = [token_ids + [pad_id] for token_ids in y] + length = max([len(token_ids) for token_ids in y]) + y = [token_ids + [pad_id] * (length - len(token_ids)) for token_ids in y] + return torch.tensor(y, dtype=torch.int64, device=device) + + +def get_tokens_index(durations: List[List[int]], num_frames: int) -> torch.Tensor: + """ + Gets position in the transcript for each frame, i.e. the position + in the symbol-sequence to look up. + + Args: + durations: + Duration of each token in transcripts. + num_frames: + The maximum frame length of the current batch. + + Returns: + Return a Tensor of shape (batch_size, num_frames) + """ + durations = [x + [num_frames - sum(x)] for x in durations] + batch_size = len(durations) + ans = torch.zeros(batch_size, num_frames, dtype=torch.int64) + for b in range(batch_size): + this_dur = durations[b] + cur_frame = 0 + for i, d in enumerate(this_dur): + ans[b, cur_frame : cur_frame + d] = i + cur_frame += d + assert cur_frame == num_frames, (cur_frame, num_frames) + return ans + + +def to_int_tuple(s: Union[str, int]): + if isinstance(s, int): + return (s,) + return tuple(map(int, s.split(","))) + + +def get_adjusted_batch_count(params: AttributeDict) -> float: + # returns the number of batches we would have used so far if we had used the + # reference duration. This is for purposes of set_batch_count(). + return ( + params.batch_idx_train + * (params.max_duration * params.world_size) + / params.ref_duration + ) + + +def set_batch_count(model: Union[nn.Module, DDP], batch_count: float) -> None: + if isinstance(model, DDP): + # get underlying nn.Module + model = model.module + for name, module in model.named_modules(): + if hasattr(module, "batch_count"): + module.batch_count = batch_count + if hasattr(module, "name"): + module.name = name + + +def condition_time_mask( + features_lens: torch.Tensor, + mask_percent: Tuple[float, float], + max_len: int = 0, +) -> torch.Tensor: + """ + Apply Time masking. + Args: + features_lens: + input tensor of shape ``(B)`` + mask_size: + the width size for masking. + max_len: + the maximum length of the mask. + Returns: + Return a 2-D bool tensor (B, T), where masked positions + are filled with `True` and non-masked positions are + filled with `False`. + """ + mask_size = ( + torch.zeros_like(features_lens, dtype=torch.float32).uniform_(*mask_percent) + * features_lens + ).to(torch.int64) + mask_starts = ( + torch.rand_like(mask_size, dtype=torch.float32) * (features_lens - mask_size) + ).to(torch.int64) + mask_ends = mask_starts + mask_size + max_len = max(max_len, features_lens.max()) + seq_range = torch.arange(0, max_len, device=features_lens.device) + mask = (seq_range[None, :] >= mask_starts[:, None]) & ( + seq_range[None, :] < mask_ends[:, None] + ) + return mask + + +def condition_time_mask_suffix( + features_lens: torch.Tensor, + mask_percent: Tuple[float, float], + max_len: int = 0, +) -> torch.Tensor: + """ + Apply Time masking, mask from the end time index. + Args: + features_lens: + input tensor of shape ``(B)`` + mask_size: + the width size for masking. + max_len: + the maximum length of the mask. + Returns: + Return a 2-D bool tensor (B, T), where masked positions + are filled with `True` and non-masked positions are + filled with `False`. + """ + mask_size = ( + torch.zeros_like(features_lens, dtype=torch.float32).uniform_(*mask_percent) + * features_lens + ).to(torch.int64) + mask_starts = ( + torch.ones_like(mask_size, dtype=torch.float32) * (features_lens - mask_size) + ).to(torch.int64) + mask_ends = mask_starts + mask_size + max_len = max(max_len, features_lens.max()) + seq_range = torch.arange(0, max_len, device=features_lens.device) + mask = (seq_range[None, :] >= mask_starts[:, None]) & ( + seq_range[None, :] < mask_ends[:, None] + ) + return mask + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """ + Args: + lengths: + A 1-D tensor containing sentence lengths. + max_len: + The length of masks. + Returns: + Return a 2-D bool tensor, where masked positions + are filled with `True` and non-masked positions are + filled with `False`. + + >>> lengths = torch.tensor([1, 3, 2, 5]) + >>> make_pad_mask(lengths) + tensor([[False, True, True, True, True], + [False, False, False, True, True], + [False, False, True, True, True], + [False, False, False, False, False]]) + """ + assert lengths.ndim == 1, lengths.ndim + max_len = max(max_len, lengths.max()) + n = lengths.size(0) + seq_range = torch.arange(0, max_len, device=lengths.device) + expaned_lengths = seq_range.unsqueeze(0).expand(n, max_len) + + return expaned_lengths >= lengths.unsqueeze(-1) + + +def str2bool(v): + """Used in argparse.ArgumentParser.add_argument to indicate + that a type is a bool type and user can enter + + - yes, true, t, y, 1, to represent True + - no, false, f, n, 0, to represent False + + See https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse # noqa + """ + if isinstance(v, bool): + return v + if v.lower() in ("yes", "true", "t", "y", "1"): + return True + elif v.lower() in ("no", "false", "f", "n", "0"): + return False + else: + raise argparse.ArgumentTypeError("Boolean value expected.") + + +def setup_logger( + log_filename: Pathlike, + log_level: str = "info", + use_console: bool = True, +) -> None: + """Setup log level. + + Args: + log_filename: + The filename to save the log. + log_level: + The log level to use, e.g., "debug", "info", "warning", "error", + "critical" + use_console: + True to also print logs to console. + """ + now = datetime.now() + date_time = now.strftime("%Y-%m-%d-%H-%M-%S") + if dist.is_available() and dist.is_initialized(): + world_size = dist.get_world_size() + rank = dist.get_rank() + formatter = f"%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] ({rank}/{world_size}) %(message)s" # noqa + log_filename = f"{log_filename}-{date_time}-{rank}" + else: + formatter = "%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s" + log_filename = f"{log_filename}-{date_time}" + + os.makedirs(os.path.dirname(log_filename), exist_ok=True) + + level = logging.ERROR + if log_level == "debug": + level = logging.DEBUG + elif log_level == "info": + level = logging.INFO + elif log_level == "warning": + level = logging.WARNING + elif log_level == "critical": + level = logging.CRITICAL + + logging.basicConfig( + filename=log_filename, + format=formatter, + level=level, + filemode="w", + force=True, + ) + if use_console: + console = logging.StreamHandler() + console.setLevel(level) + console.setFormatter(logging.Formatter(formatter)) + logging.getLogger("").addHandler(console) + + +def get_git_sha1(): + try: + git_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + check=True, + stdout=subprocess.PIPE, + ) + .stdout.decode() + .rstrip("\n") + .strip() + ) + dirty_commit = ( + len( + subprocess.run( + ["git", "diff", "--shortstat"], + check=True, + stdout=subprocess.PIPE, + ) + .stdout.decode() + .rstrip("\n") + .strip() + ) + > 0 + ) + git_commit = git_commit + "-dirty" if dirty_commit else git_commit + "-clean" + except: # noqa + return None + + return git_commit + + +def get_git_date(): + try: + git_date = ( + subprocess.run( + ["git", "log", "-1", "--format=%ad", "--date=local"], + check=True, + stdout=subprocess.PIPE, + ) + .stdout.decode() + .rstrip("\n") + .strip() + ) + except: # noqa + return None + + return git_date + + +def get_git_branch_name(): + try: + git_date = ( + subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + check=True, + stdout=subprocess.PIPE, + ) + .stdout.decode() + .rstrip("\n") + .strip() + ) + except: # noqa + return None + + return git_date + + +def get_env_info() -> Dict[str, Any]: + """Get the environment information.""" + return { + "torch-version": str(torch.__version__), + "torch-cuda-available": torch.cuda.is_available(), + "torch-cuda-version": torch.version.cuda, + "python-version": sys.version[:4], + "zipvoice-git-branch": get_git_branch_name(), + "zipvoice-git-sha1": get_git_sha1(), + "zipvoice-git-date": get_git_date(), + "zipvoice-path": str(Path(__file__).resolve().parent.parent), + "hostname": socket.gethostname(), + "IP address": socket.gethostbyname(socket.gethostname()), + } + + +def get_parameter_groups_with_lrs( + model: nn.Module, + lr: float, + include_names: bool = False, + freeze_modules: List[str] = [], + unfreeze_modules: List[str] = [], +) -> List[dict]: + """ + This is for use with the ScaledAdam optimizers (more recent versions that accept + lists of named-parameters; we can, if needed, create a version without the names). + + It provides a way to specify learning-rate scales inside the module, so that if + any nn.Module in the hierarchy has a floating-point parameter 'lr_scale', it will + scale the LR of any parameters inside that module or its submodules. Note: you + can set module parameters outside the __init__ function, e.g.: + >>> a = nn.Linear(10, 10) + >>> a.lr_scale = 0.5 + + Returns: a list of dicts, of the following form: + if include_names == False: + [ { 'params': [ tensor1, tensor2, ... ], 'lr': 0.01 }, + { 'params': [ tensor3, tensor4, ... ], 'lr': 0.005 }, + ... ] + if include_names == true: + [ { 'named_params': [ (name1, tensor1, (name2, tensor2), ... ], 'lr': 0.01 }, + { 'named_params': [ (name3, tensor3), (name4, tensor4), ... ], 'lr': 0.005 }, + ... ] + + """ + # Use freeze_modules or unfreeze_modules to freeze or unfreeze modules + assert not (len(freeze_modules) and len(unfreeze_modules)) + + # flat_lr_scale just contains the lr_scale explicitly specified + # for each prefix of the name, e.g. 'encoder.layers.3', these need + # to be multiplied for all prefix of the name of any given parameter. + flat_lr_scale = defaultdict(lambda: 1.0) + names = [] + for name, m in model.named_modules(): + names.append(name) + if hasattr(m, "lr_scale"): + flat_lr_scale[name] = m.lr_scale + + # lr_to_parames is a dict from learning rate (floating point) to: if + # include_names == true, a list of (name, parameter) for that learning rate; + # otherwise a list of parameters for that learning rate. + lr_to_params = defaultdict(list) + + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + logging.info(f"Remove {name} from parameter") + continue + split_name = name.split(".") + # caution: as a special case, if the name is '', split_name will be [ '' ]. + prefix = split_name[0] + if len(freeze_modules) > 0: + if prefix == "module": # DDP + module_name = split_name[1] + if module_name in freeze_modules: + logging.info(f"Remove {name} from parameters") + continue + else: + if prefix in freeze_modules: + logging.info(f"Remove {name} from parameters") + continue + elif len(unfreeze_modules) > 0: + if prefix == "module": # DDP + module_name = split_name[1] + if module_name not in unfreeze_modules: + logging.info(f"Remove {name} from parameters") + continue + else: + if prefix not in unfreeze_modules: + logging.info(f"Remove {name} from parameters") + continue + cur_lr = lr * flat_lr_scale[prefix] + if prefix != "": + cur_lr *= flat_lr_scale[""] + for part in split_name[1:]: + prefix = ".".join([prefix, part]) + cur_lr *= flat_lr_scale[prefix] + lr_to_params[cur_lr].append((name, parameter) if include_names else parameter) + + if include_names: + return [{"named_params": pairs, "lr": lr} for lr, pairs in lr_to_params.items()] + else: + return [{"params": params, "lr": lr} for lr, params in lr_to_params.items()] diff --git a/zipvoice/utils/diagnostics.py b/zipvoice/utils/diagnostics.py new file mode 100644 index 0000000..bdefaa3 --- /dev/null +++ b/zipvoice/utils/diagnostics.py @@ -0,0 +1,723 @@ +# Copyright 2022-2024 Xiaomi Corp. (authors: Daniel Povey +# Zengwei Yao +# Mingshuang Luo, +# Zengrui Jin,) +# +# See ../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import random +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +from torch import Tensor, nn + + +class TensorDiagnosticOptions(object): + """Options object for tensor diagnostics: + + Args: + max_eig_dim: + The maximum dimension for which we print out eigenvalues + (limited for speed reasons). + """ + + def __init__(self, max_eig_dim: int = 512): + self.max_eig_dim = max_eig_dim + + def dim_is_summarized(self, size: int): + return size > 10 and size != 31 + + +def get_tensor_stats( + x: Tensor, + dim: int, + stats_type: str, +) -> Tuple[Tensor, int]: + """ + Returns the specified transformation of the Tensor (either x or x.abs() + or (x > 0), summed over all but the index `dim`. + + Args: + x: + Tensor, tensor to be analyzed + dim: + Dimension with 0 <= dim < x.ndim + stats_type: + The stats_type includes several types: + "abs" -> take abs() before summing + "positive" -> take (x > 0) before summing + "rms" -> square before summing, we'll take sqrt later + "value" -> just sum x itself + "max", "min" -> take the maximum or minimum [over all other dims but dim] + instead of summing + "rms-sort" -> this is a bit different than the others, it's based on computing + the rms over the specified dim and returning percentiles of the result + (11 of them). + Returns: + stats: a Tensor of shape (x.shape[dim],). + count: an integer saying how many items were counted in each element + of stats. + """ + + if stats_type == "rms-sort": + rms = (x**2).mean(dim=dim).sqrt() + rms = rms.flatten() + rms = rms.sort()[0] + rms = rms[(torch.arange(11) * rms.numel() // 10).clamp(max=rms.numel() - 1)] + count = 1.0 + return rms, count + + count = x.numel() // x.shape[dim] + + if stats_type == "eigs": + x = x.transpose(dim, -1) + x = x.reshape(-1, x.shape[-1]) + # shape of returned tensor: (s, s), + # where s is size of dimension `dim` of original x. + return torch.matmul(x.transpose(0, 1), x), count + elif stats_type == "abs": + x = x.abs() + elif stats_type == "rms": + x = x**2 + elif stats_type == "positive": + x = (x > 0).to(dtype=torch.float) + else: + assert stats_type in ["value", "max", "min"] + + sum_dims = [d for d in range(x.ndim) if d != dim] + if len(sum_dims) > 0: + if stats_type == "max": + for dim in reversed(sum_dims): + x = torch.max(x, dim=dim)[0] + elif stats_type == "min": + for dim in reversed(sum_dims): + x = torch.min(x, dim=dim)[0] + else: + x = torch.sum(x, dim=sum_dims) + x = x.flatten().clone() + return x, count + + +@dataclass +class TensorAndCount: + tensor: Tensor + count: int + + +class TensorDiagnostic(object): + """This class is not directly used by the user, it is responsible for + collecting diagnostics for a module or parameter tensor of a torch.nn.Module. + + Args: + opts: + Options object. + name: + The name associated with this diagnostics object, will probably be + {module_name}.X where X is "output" or "grad", or {parameter_name}. + Y where Y is param_value or param_grad. + """ + + def __init__(self, opts: TensorDiagnosticOptions, name: str): + self.opts = opts + self.name = name + self.class_name = None # will assign in accumulate() + + self.stats = None # we'll later assign a list to self.stats. + # It's a list of dicts, indexed by dim (i.e. by the + # axis of the tensor). The dicts, in turn, are + # indexed by `stats-type` which are strings in + # ["abs", "max", "min", "positive", "value", "rms"]. + + # scalar_stats contains some analysis of the activations and gradients, + self.scalar_stats = None + + # the keys into self.stats[dim] are strings, whose values can be + # "abs", "max", "min" ,"value", "positive", "rms", "value". + # The values e.g. self.stats[dim]["rms"] are lists of dataclass TensorAndCount, + # containing a tensor and its associated count (which is the sum of the other + # dims that we aggregated over, e.g. the number of frames and/or batch elements + # and/or channels. + # ... we actually accumulate the Tensors / counts any time we have the same-dim + # tensor, only adding a new element to the list if there was a different dim. + # if the string in the key is "eigs", if we detect a length mismatch we put None + # as the value. + + def accumulate(self, x, class_name: Optional[str] = None): + """ + Accumulate tensors. + """ + if class_name is not None: + self.class_name = class_name + if isinstance(x, Tuple): + x = x[0] + if not isinstance(x, Tensor): + return + if x.numel() == 0: # for empty tensor + return + x = x.detach().clone() + if x.ndim == 0: + x = x.unsqueeze(0) + ndim = x.ndim + if self.stats is None: + self.stats = [dict() for _ in range(ndim)] + + for dim in range(ndim): + this_dim_stats = self.stats[dim] + if ndim > 1: + # rms-sort is different from the others, it's based on summing over just + # this dim, then sorting and returning the percentiles. + stats_types = [ + "abs", + "max", + "min", + "positive", + "value", + "rms", + "rms-sort", + ] + if x.shape[dim] <= self.opts.max_eig_dim: + stats_types.append("eigs") + else: + stats_types = ["value", "abs", "max", "min"] + + for stats_type in stats_types: + stats, count = get_tensor_stats(x, dim, stats_type) + if stats_type not in this_dim_stats: + this_dim_stats[stats_type] = [] # list of TensorAndCount + + done = False + if this_dim_stats[stats_type] is None: + # we can reach here if we detected for stats_type "eigs" that + # where was more than one different size for this dim. Then we + # disable accumulating this stats type, as it uses too much memory. + continue + for s in this_dim_stats[stats_type]: + if s.tensor.shape == stats.shape: + if stats_type == "max": + s.tensor = torch.maximum(s.tensor, stats) + + elif stats_type == "min": + s.tensor = torch.minimum(s.tensor, stats) + else: + assert stats_type != "max" + s.tensor += stats + s.count += count + done = True + break + if not done: + if this_dim_stats[stats_type] != [] and stats_type == "eigs": + # >1 size encountered on this dim, e.g. it's a batch or time + # dimension, don't accumulat "eigs" stats type, it uses too much + # memory + this_dim_stats[stats_type] = None + else: + this_dim_stats[stats_type].append(TensorAndCount(stats, count)) + + def print_diagnostics(self): + """Print diagnostics for each dimension of the tensor.""" + if self.stats is None: + print(f"Warning: the stats of {self.name} is None.") + return + for dim, this_dim_stats in enumerate(self.stats): + if "rms" in this_dim_stats and "value" in this_dim_stats: + # produce "stddev" stats, which is centered RMS. + rms_stats_list = this_dim_stats["rms"] + value_stats_list = this_dim_stats["value"] + if len(rms_stats_list) == len(value_stats_list): + stddev_stats_list = [] + for r, v in zip(rms_stats_list, value_stats_list): + stddev_stats_list.append( + # r.count and v.count should be the same, but we don't check + # this. + TensorAndCount( + r.tensor - v.tensor * v.tensor / (v.count + 1.0e-20), + r.count, + ) + ) + this_dim_stats["stddev"] = stddev_stats_list + + for stats_type, stats_list in this_dim_stats.items(): + # stats_type could be "rms", "value", "abs", "eigs", "positive", "min" + # or "max". "stats_list" could be a list of TensorAndCount (one list per + # distinct tensor shape of the stats), or None + if stats_list is None: + assert stats_type == "eigs" + continue + + def get_count(count): + return 1 if stats_type in ["max", "min"] else count + + if len(stats_list) == 1: + stats = stats_list[0].tensor / get_count(stats_list[0].count) + else: + # a dimension that has variable size in different nnet + # forwards, e.g. a time dimension in an ASR model. + stats = torch.cat( + [x.tensor / get_count(x.count) for x in stats_list], dim=0 + ) + + if stats_type == "eigs": + try: + if hasattr(torch, "linalg") and hasattr(torch.linalg, "eigh"): + eigs, _ = torch.linalg.eigh(stats) + else: + eigs, _ = torch.symeig(stats) + stats = eigs.abs().sqrt() + except: # noqa + print("Error getting eigenvalues, trying another method.") + if hasattr(torch, "linalg") and hasattr(torch.linalg, "eig"): + eigs, _ = torch.linalg.eig(stats) + eigs = eigs.abs() + else: + eigs, _ = torch.eig(stats) + eigs = eigs.norm(dim=1) + stats = eigs.sqrt() + # sqrt so it reflects data magnitude, like stddev- not variance + + if stats_type in ["rms", "stddev"]: + # we stored the square; after aggregation we need to take sqrt. + stats = stats.sqrt() + + # if `summarize` we print percentiles of the stats; else, + # we print out individual elements. + summarize = (len(stats_list) > 1) or self.opts.dim_is_summarized( + stats.numel() + ) + if summarize: # usually `summarize` will be true + # print out percentiles. + stats = stats.sort()[0] + num_percentiles = 10 + size = stats.numel() + percentiles = [] + for i in range(num_percentiles + 1): + index = (i * (size - 1)) // num_percentiles + percentiles.append(stats[index].item()) + percentiles = ["%.2g" % x for x in percentiles] + percentiles = " ".join(percentiles) + ans = f"percentiles: [{percentiles}]" + else: + ans = stats.tolist() + ans = ["%.2g" % x for x in ans] + ans = "[" + " ".join(ans) + "]" + if stats_type in ["value", "rms", "stddev", "eigs"]: + # This norm is useful because it is strictly less than the largest + # sqrt(eigenvalue) of the variance, which we print out, and shows, + # speaking in an approximate way, how much of that largest + # eigenvalue can be attributed to the mean of the distribution. + norm = (stats**2).sum().sqrt().item() + ans += f", norm={norm:.2g}" + mean = stats.mean().item() + rms = (stats**2).mean().sqrt().item() + ans += f", mean={mean:.3g}, rms={rms:.3g}" + + # OK, "ans" contains the actual stats, e.g. + # ans = "percentiles: \ + # [0.43 0.46 0.48 0.49 0.49 0.5 0.51 0.52 0.53 0.54 0.59], \ + # mean=0.5, rms=0.5" + + sizes = [x.tensor.shape[0] for x in stats_list] + size_str = ( + f"{sizes[0]}" if len(sizes) == 1 else f"{min(sizes)}..{max(sizes)}" + ) + maybe_class_name = ( + f" type={self.class_name}," if self.class_name is not None else "" + ) + print( + f"module={self.name},{maybe_class_name} dim={dim}, size={size_str}, " + f"{stats_type} {ans}" + ) + + +class ScalarDiagnostic(object): + """This class is not directly used by the user, it is responsible for + collecting diagnostics for a single module (subclass of torch.nn.Module) that + represents some kind of nonlinearity, e.g. ReLU, sigmoid, etc. + """ + + def __init__(self, opts: TensorDiagnosticOptions, name: str): + self.opts = opts + self.name = name + self.class_name = None # will assign in accumulate() + self.is_forward_pass = True + + self.tick_scale = None + + self.saved_inputs = [] + self.is_ok = True + + self.counts = None + self.sum_grad = None + self.sum_gradsq = None + self.sum_abs_grad = None + + def accumulate_input(self, x: Tensor, class_name: Optional[str] = None): + """ + Called in forward pass. + """ + if not self.is_forward_pass: + # in case we did a forward pass without a backward pass, for some reason. + self.saved_inputs = [] + self.is_forward_pass = True + + if class_name is not None: + self.class_name = class_name + if not self.is_ok: + return + + limit = 10 + if len(self.saved_inputs) > limit: + print( + f"ERROR: forward pass called for this module over {limit} times " + f"with no backward pass. Will not accumulate scalar stats." + ) + self.is_ok = False + return + self.saved_inputs.append(x) + + def accumulate_output_grad(self, grad: Tensor): + if not self.is_ok: + return + if self.is_forward_pass: + self.is_forward_pass = False + + last_shape = ( + "n/a" if len(self.saved_inputs) == 0 else self.saved_inputs[-1].shape + ) + if len(self.saved_inputs) == 0 or grad.shape != last_shape: + print( + f"ERROR: shape mismatch or no forward activation present when backward " + f"pass called: grad shape ={tuple(grad.shape)}" + f", num-saved-inputs={len(self.saved_inputs)}" + f", shape-of-last-saved-input={last_shape}" + ) + self.is_ok = False + return + + x = self.saved_inputs.pop() + self.process_input_and_grad(x, grad) + + def process_input_and_grad(self, x: Tensor, grad: Tensor): + assert x.shape == grad.shape + x = x.flatten() + grad = grad.flatten() + + num_ticks_per_side = 256 + + if self.tick_scale is None: + x_abs_sorted = x.abs().sort()[0] + # take the 98th percentile as the largest value we count separately. + index = int(x.numel() * 0.98) + self.tick_scale = float(x_abs_sorted[index] / num_ticks_per_side) + + # integerize from tick * (-num ticks_per_side .. num_ticks_per_side - 1] + self.counts = torch.zeros( + 2 * num_ticks_per_side, dtype=torch.long, device=x.device + ) + self.sum_grad = torch.zeros( + 2 * num_ticks_per_side, dtype=torch.double, device=x.device + ) + # sum_gradsq is for getting error bars. + self.sum_gradsq = torch.zeros( + 2 * num_ticks_per_side, dtype=torch.double, device=x.device + ) + self.sum_abs_grad = torch.zeros( + 2 * num_ticks_per_side, dtype=torch.double, device=x.device + ) + + # this will round down. + x = (x / self.tick_scale).to(torch.long) + x = x.clamp_(min=-num_ticks_per_side, max=num_ticks_per_side - 1) + x = x + num_ticks_per_side + + self.counts.index_add_(dim=0, index=x, source=torch.ones_like(x)) + self.sum_grad.index_add_(dim=0, index=x, source=grad.to(torch.double)) + self.sum_gradsq.index_add_( + dim=0, index=x, source=(grad * grad).to(torch.double) + ) + self.sum_abs_grad.index_add_(dim=0, index=x, source=grad.abs().to(torch.double)) + + def print_diagnostics(self): + """Print diagnostics.""" + if self.is_ok is False or self.counts is None: + print(f"Warning: no stats accumulated for {self.name}, is_ok={self.is_ok}") + return + + counts = self.counts.to("cpu") + sum_grad = self.sum_grad.to(device="cpu", dtype=torch.float32) + sum_gradsq = self.sum_gradsq.to(device="cpu", dtype=torch.float32) + sum_abs_grad = self.sum_abs_grad.to(device="cpu", dtype=torch.float32) + + counts_cumsum = counts.cumsum(dim=0) + counts_tot = counts_cumsum[-1] + + # subdivide the distribution up into `num_bins` intervals for analysis, for + # greater statistical significance. each bin corresponds to multiple of the + # original 'tick' intervals. + num_bins = 20 + + # integer division + counts_per_bin = (counts_tot // num_bins) + 1 + bin_indexes = counts_cumsum // counts_per_bin + bin_indexes = bin_indexes.clamp(min=0, max=num_bins).to(torch.long) + + bin_counts = torch.zeros(num_bins, dtype=torch.long) + bin_counts.index_add_(dim=0, index=bin_indexes, source=counts) + bin_grad = torch.zeros(num_bins) + bin_grad.index_add_(dim=0, index=bin_indexes, source=sum_grad) + bin_gradsq = torch.zeros(num_bins) + bin_gradsq.index_add_(dim=0, index=bin_indexes, source=sum_gradsq) + bin_abs_grad = torch.zeros(num_bins) + bin_abs_grad.index_add_(dim=0, index=bin_indexes, source=sum_abs_grad) + + bin_boundary_counts = ( + torch.arange(num_bins + 1, dtype=torch.long) * counts_per_bin + ) + bin_tick_indexes = torch.searchsorted(counts_cumsum, bin_boundary_counts) + # boundaries are the "x" values between the bins, e.g. corresponding to the + # locations of percentiles of the distribution. + num_ticks_per_side = counts.numel() // 2 + bin_boundaries = (bin_tick_indexes - num_ticks_per_side) * self.tick_scale + + bin_grad = bin_grad / (bin_counts + 1) + bin_conf_interval = bin_gradsq.sqrt() / ( + bin_counts + 1 + ) # consider this a standard deviation. + # bin_grad / bin_abs_grad will give us a sense for how important in a practical + # sense, the gradients are. + bin_abs_grad = bin_abs_grad / (bin_counts + 1) + + bin_rel_grad = bin_grad / (bin_abs_grad + 1.0e-20) + bin_conf = bin_grad / (bin_conf_interval + 1.0e-20) + + def tensor_to_str(x: Tensor): + x = ["%.2g" % f for f in x] + x = "[" + " ".join(x) + "]" + return x + + maybe_class_name = ( + f" type={self.class_name}," if self.class_name is not None else "" + ) + + print( + f"module={self.name},{maybe_class_name} " + f"bin-boundaries={tensor_to_str(bin_boundaries)}, " + f"rel_grad={tensor_to_str(bin_rel_grad)}, " + f"grad_conf={tensor_to_str(bin_conf)}" + ) + + +class ModelDiagnostic(object): + """This class stores diagnostics for all tensors in the torch.nn.Module. + + Args: + opts: + Options object. + """ + + def __init__(self, opts: Optional[TensorDiagnosticOptions] = None): + # In this dictionary, the keys are tensors names and the values + # are corresponding TensorDiagnostic objects. + if opts is None: + self.opts = TensorDiagnosticOptions() + else: + self.opts = opts + self.diagnostics = dict() + + def __getitem__(self, name: str): + T = ScalarDiagnostic if name[-7:] == ".scalar" else TensorDiagnostic + if name not in self.diagnostics: + self.diagnostics[name] = T(self.opts, name) + return self.diagnostics[name] + + def print_diagnostics(self): + """Print diagnostics for each tensor.""" + for k in sorted(self.diagnostics.keys()): + self.diagnostics[k].print_diagnostics() + + +def get_class_name(module: nn.Module): + ans = type(module).__name__ + # we put the below in try blocks in case anyone is using a different version of + # these modules that might have different member names. + if ans == "Balancer" or ans == "ActivationBalancer": + try: + ans += f"[{float(module.min_positive)},{float(module.max_positive)}," + f"{float(module.min_abs)},{float(module.max_abs)}]" + except: + pass + elif ans == "AbsValuePenalizer": + try: + ans += f"[{module.limit}]" + except: + pass + return ans + + +def attach_diagnostics( + model: nn.Module, opts: Optional[TensorDiagnosticOptions] = None +) -> ModelDiagnostic: + """Attach a ModelDiagnostic object to the model by + 1) registering forward hook and backward hook on each module, to accumulate + its output tensors and gradient tensors, respectively; + 2) registering backward hook on each module parameter, to accumulate its + values and gradients. + + Args: + model: + the model to be analyzed. + opts: + Options object. + + Returns: + The ModelDiagnostic object attached to the model. + """ + + ans = ModelDiagnostic(opts) + for name, module in model.named_modules(): + if name == "": + name = "" + + # Setting model_diagnostic=ans and n=name below, instead of trying to + # capture the variables, ensures that we use the current values. + # (this matters for `name`, since the variable gets overwritten). + # These closures don't really capture by value, only by + # "the final value the variable got in the function" :-( + def forward_hook(_module, _input, _output, _model_diagnostic=ans, _name=name): + if isinstance(_output, tuple) and len(_output) == 1: + _output = _output[0] + + if isinstance(_output, Tensor) and _output.dtype in ( + torch.float32, + torch.float16, + torch.float64, + ): + _model_diagnostic[f"{_name}.output"].accumulate( + _output, class_name=get_class_name(_module) + ) + elif isinstance(_output, tuple): + for i, o in enumerate(_output): + if isinstance(o, Tensor) and o.dtype in ( + torch.float32, + torch.float16, + torch.float64, + ): + _model_diagnostic[f"{_name}.output[{i}]"].accumulate( + o, class_name=get_class_name(_module) + ) + + def backward_hook(_module, _input, _output, _model_diagnostic=ans, _name=name): + if isinstance(_output, tuple) and len(_output) == 1: + _output = _output[0] + if isinstance(_output, Tensor) and _output.dtype in ( + torch.float32, + torch.float16, + torch.float64, + ): + _model_diagnostic[f"{_name}.grad"].accumulate( + _output, class_name=get_class_name(_module) + ) + elif isinstance(_output, tuple): + for i, o in enumerate(_output): + if isinstance(o, Tensor) and o.dtype in ( + torch.float32, + torch.float16, + torch.float64, + ): + _model_diagnostic[f"{_name}.grad[{i}]"].accumulate( + o, class_name=get_class_name(_module) + ) + + module.register_forward_hook(forward_hook) + module.register_backward_hook(backward_hook) + + if type(module).__name__ in [ + "Sigmoid", + "Tanh", + "ReLU", + "TanSwish", + "Swish", + "DoubleSwish", + "Swoosh", + ]: + # For these specific module types, accumulate some additional diagnostics + # that can help us improve the activation function. These require a lot of + # memory, to save the forward activations, so limit this to some select + # classes. Note: this will not work correctly for all model types. + def scalar_forward_hook( + _module, _input, _output, _model_diagnostic=ans, _name=name + ): + if isinstance(_input, tuple): + (_input,) = _input + assert isinstance(_input, Tensor) + _model_diagnostic[f"{_name}.scalar"].accumulate_input( + _input, class_name=get_class_name(_module) + ) + + def scalar_backward_hook( + _module, _input, _output, _model_diagnostic=ans, _name=name + ): + if isinstance(_output, tuple): + (_output,) = _output + assert isinstance(_output, Tensor) + _model_diagnostic[f"{_name}.scalar"].accumulate_output_grad(_output) + + module.register_forward_hook(scalar_forward_hook) + module.register_backward_hook(scalar_backward_hook) + + for name, parameter in model.named_parameters(): + + def param_backward_hook( + grad, _parameter=parameter, _model_diagnostic=ans, _name=name + ): + _model_diagnostic[f"{_name}.param_value"].accumulate(_parameter) + _model_diagnostic[f"{_name}.param_grad"].accumulate(grad) + + try: + parameter.register_hook(param_backward_hook) + except: + logging.warning( + f"Warning: could not register backward hook for parameter {name}, " + f"it might not be differentiable." + ) + + return ans + + +def _test_tensor_diagnostic(): + opts = TensorDiagnosticOptions(512) + + diagnostic = TensorDiagnostic(opts, "foo") + + for _ in range(10): + diagnostic.accumulate(torch.randn(50, 100) * 10.0) + + diagnostic.print_diagnostics() + + model = nn.Sequential(nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 80)) + + diagnostic = attach_diagnostics(model, opts) + for _ in range(10): + T = random.randint(200, 300) + x = torch.randn(T, 100) + y = model(x) + y.sum().backward() + + diagnostic.print_diagnostics() + + +if __name__ == "__main__": + _test_tensor_diagnostic() diff --git a/zipvoice/utils/feature.py b/zipvoice/utils/feature.py new file mode 100644 index 0000000..f6cbc34 --- /dev/null +++ b/zipvoice/utils/feature.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python3 +# Copyright 2024 Xiaomi Corp. (authors: Han Zhu) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from typing import Union + +import numpy as np +import torch +import torchaudio +from lhotse.features.base import FeatureExtractor, register_extractor +from lhotse.utils import Seconds, compute_num_frames + + +@dataclass +class VocosFbankConfig: + sampling_rate: int = 24000 + n_mels: int = 100 + n_fft: int = 1024 + hop_length: int = 256 + + +@register_extractor +class VocosFbank(FeatureExtractor): + + name = "VocosFbank" + config_type = VocosFbankConfig + + def __init__(self, num_channels: int = 1): + config = VocosFbankConfig + super().__init__(config=config) + assert num_channels in (1, 2) + self.num_channels = num_channels + self.fbank = torchaudio.transforms.MelSpectrogram( + sample_rate=self.config.sampling_rate, + n_fft=self.config.n_fft, + hop_length=self.config.hop_length, + n_mels=self.config.n_mels, + center=True, + power=1, + ) + + def _feature_fn(self, sample): + mel = self.fbank(sample) + logmel = mel.clamp(min=1e-7).log() + + return logmel + + @property + def device(self) -> Union[str, torch.device]: + return self.config.device + + def feature_dim(self, sampling_rate: int) -> int: + return self.config.n_mels + + def extract( + self, + samples: Union[np.ndarray, torch.Tensor], + sampling_rate: int, + ) -> Union[np.ndarray, torch.Tensor]: + # Check for sampling rate compatibility. + expected_sr = self.config.sampling_rate + assert sampling_rate == expected_sr, ( + f"Mismatched sampling rate: extractor expects {expected_sr}, " + f"got {sampling_rate}" + ) + is_numpy = False + if not isinstance(samples, torch.Tensor): + samples = torch.from_numpy(samples) + is_numpy = True + + if len(samples.shape) == 1: + samples = samples.unsqueeze(0) + else: + assert samples.ndim == 2, samples.shape + + if self.num_channels == 1: + if samples.shape[0] == 2: + samples = samples.mean(dim=0, keepdims=True) + else: + assert samples.shape[0] == 2, samples.shape + + mel = self._feature_fn(samples) + # (1, n_mels, time) or (2, n_mels, time) + mel = mel.reshape(-1, mel.shape[-1]).t() + # (time, n_mels) or (time, 2 * n_mels) + + num_frames = compute_num_frames( + samples.shape[1] / sampling_rate, self.frame_shift, sampling_rate + ) + + if mel.shape[0] > num_frames: + mel = mel[:num_frames] + elif mel.shape[0] < num_frames: + mel = mel.unsqueeze(0) + mel = torch.nn.functional.pad( + mel, (0, 0, 0, num_frames - mel.shape[1]), mode="replicate" + ).squeeze(0) + + if is_numpy: + return mel.cpu().numpy() + else: + return mel + + @property + def frame_shift(self) -> Seconds: + return self.config.hop_length / self.config.sampling_rate diff --git a/zipvoice/utils/hooks.py b/zipvoice/utils/hooks.py new file mode 100644 index 0000000..3a16581 --- /dev/null +++ b/zipvoice/utils/hooks.py @@ -0,0 +1,111 @@ +# Copyright 2021-2024 Xiaomi Corporation (authors: Zengwei Yao, +# Daniel Povey, +# Zengrui Jin,) +# +# See ../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +import random + +import torch +from torch import Tensor, nn + + +def register_inf_check_hooks(model: nn.Module) -> None: + """Registering forward hook on each module, to check + whether its output tensors is not finite. + + Args: + model: + the model to be analyzed. + """ + + for name, module in model.named_modules(): + if name == "": + name = "" + + # default param _name is a way to capture the current value of the variable + # "name". + def forward_hook(_module, _input, _output, _name=name): + if isinstance(_output, Tensor): + try: + if not torch.isfinite(_output.to(torch.float32).sum()): + logging.warning(f"The sum of {_name}.output is not finite") + except RuntimeError: # e.g. CUDA out of memory + pass + elif isinstance(_output, tuple): + for i, o in enumerate(_output): + if isinstance(o, tuple): + o = o[0] + if not isinstance(o, Tensor): + continue + try: + if not torch.isfinite(o.to(torch.float32).sum()): + logging.warning( + f"The sum of {_name}.output[{i}] is not finite" + ) + except RuntimeError: # e.g. CUDA out of memory + pass + + # default param _name is a way to capture the current value of the variable + # "name". + def backward_hook(_module, _input, _output, _name=name): + if isinstance(_output, Tensor): + try: + if not torch.isfinite(_output.to(torch.float32).sum()): + logging.warning(f"The sum of {_name}.grad is not finite") + except RuntimeError: # e.g. CUDA out of memory + pass + + elif isinstance(_output, tuple): + for i, o in enumerate(_output): + if isinstance(o, tuple): + o = o[0] + if not isinstance(o, Tensor): + continue + if not torch.isfinite(o.to(torch.float32).sum()): + logging.warning(f"The sum of {_name}.grad[{i}] is not finite") + + module.register_forward_hook(forward_hook) + module.register_backward_hook(backward_hook) + + for name, parameter in model.named_parameters(): + + def param_backward_hook(grad, _name=name): + if not torch.isfinite(grad.to(torch.float32).sum()): + logging.warning(f"The sum of {_name}.param_grad is not finite") + + try: + parameter.register_hook(param_backward_hook) + except Exception as e: + logging.warning( + f"Warning: could not register backward hook for parameter {name}" + f" with error {e}, it might not be differentiable." + ) + + +def _test_inf_check_hooks(): + model = nn.Sequential(nn.Linear(100, 50), nn.Linear(50, 80)) + + register_inf_check_hooks(model) + for _ in range(10): + T = random.randint(200, 300) + x = torch.randn(T, 100) + float("inf") * (T % 2) + y = model(x) + y.sum().backward() + + +if __name__ == "__main__": + _test_inf_check_hooks() diff --git a/zipvoice/utils/infer.py b/zipvoice/utils/infer.py new file mode 100644 index 0000000..1dcc9ca --- /dev/null +++ b/zipvoice/utils/infer.py @@ -0,0 +1,414 @@ +from typing import List + +import numpy as np +import torch +import torchaudio +from pydub import AudioSegment +from pydub.silence import detect_leading_silence, split_on_silence + +punctuation = {";", ":", ",", ".", "!", "?", "；", "：", "，", "。", "！", "？"} + + +def chunk_tokens_punctuation(tokens_list: List[str], max_tokens: int = 100): + """ + Splits the input tokens list into chunks according to punctuations, + each with a maximum number of tokens. + + Args: + token_list (list of str): The list of tokens to be split. + max_tokens (int): The maximum number of tokens per chunk. + + Returns: + List[str]: A list of text chunks. + """ + + # 1. Split the tokens according to punctuations. + sentences = [] + current_sentence = [] + for token in tokens_list: + # If the first token of current sentence is punctuation or blank, + # append it to the end of the previous sentence. + if ( + len(current_sentence) == 0 + and len(sentences) != 0 + and (token in punctuation or token == " ") + ): + sentences[-1].append(token) + # Otherwise, append the current token to the current sentence. + else: + current_sentence.append(token) + # Split the sentence in positions of punctuations. + if token in punctuation: + sentences.append(current_sentence) + current_sentence = [] + # Assume the last few tokens are also a sentence + if len(current_sentence) != 0: + sentences.append(current_sentence) + + # 2. Merge short sentences. + chunks = [] + current_chunk = [] + for sentence in sentences: + if len(current_chunk) + len(sentence) <= max_tokens: + current_chunk.extend(sentence) + else: + if len(current_chunk) > 0: + chunks.append(current_chunk) + current_chunk = sentence + + if len(current_chunk) > 0: + chunks.append(current_chunk) + + return chunks + + +def chunk_tokens_dialog(tokens_list: List[str], max_tokens: int = 100): + """ + Splits the input tokens list into chunks according to speaker-turn + symbol [S1], each with a maximum number of tokens. + + Args: + token_list (list of str): The list of tokens to be split. + max_tokens (int): The maximum number of tokens per chunk. + + Returns: + List[str]: A list of text chunks. + """ + + # 1. Split the tokens according to speaker-turn symbol [S1]. + dialogs = [] + current_dialog = [] + for token in tokens_list: + if token == "[S1]": + if len(current_dialog) != 0: + dialogs.append(current_dialog) + current_dialog = [] + current_dialog.append(token) + # Assume the last few tokens are also a dialog + if len(current_dialog) != 0: + dialogs.append(current_dialog) + + # 2. Merge short dialogs. + chunks = [] + current_chunk = [] + for dialog in dialogs: + if len(current_chunk) + len(dialog) <= max_tokens: + current_chunk.extend(dialog) + else: + if len(current_chunk) > 0: + chunks.append(current_chunk) + current_chunk = dialog + + if len(current_chunk) > 0: + chunks.append(current_chunk) + + return chunks + + +def batchify_tokens( + tokens_list: List[List[int]], + max_duration: float, + prompt_duration: float, + token_duration: float, +): + """ + Sort and group the input list of token sequences into batches, where each batch's + total duration does not exceed the maximum. + + Args: + tokens_list (List[List[int]]): A list of token sequences, where each inner + list represents a sequence of tokens. + max_duration (float): The maximum allowed total duration for each batch. + prompt_duration (float): The duration cost per prompt in the batch. + token_duration (float): The duration cost per token. + + Returns: + batches: List[List[List[int]]]: A list of batches, where each batch is a list of + token sequences that fit within the max duration. + index: List[int]: The original index of each sentence, used to recover the + sequential order in the future. + """ + # Create index for each sentence + indexed_tokens = list(enumerate(tokens_list)) + + # Sort according to sentence length (for less padding) + indexed_sorted_tokens = sorted(indexed_tokens, key=lambda x: len(x[1])) + index = [indexed_sorted_tokens[i][0] for i in range(len(indexed_sorted_tokens))] + sorted_tokens = [ + indexed_sorted_tokens[i][1] for i in range(len(indexed_sorted_tokens)) + ] + + batches = [] + batch = [] + batch_size = 0 # Total number of tokens in current batch + + for tokens in sorted_tokens: + # Calculate if adding current token sequence would exceed max duration + # Formula considers: existing tokens' duration + existing + # prompts' duration + new tokens' duration + if ( + batch_size * token_duration + + len(batch) * prompt_duration + + len(tokens) * token_duration + <= max_duration + ): + # Add to current batch if within duration limit + batch.append(tokens) + batch_size += len(tokens) + else: + # If exceeding limit, finalize current batch (if not empty) + if len(batch) > 0: + batches.append(batch) + # Start new batch with current token sequence + batch = [tokens] + batch_size = len(tokens) + + # Add the last batch if it's not empty + if len(batch) > 0: + batches.append(batch) + + return batches, index + + +def cross_fade_concat( + chunks: List[torch.Tensor], fade_duration: float = 0.1, sample_rate: int = 24000 +) -> torch.Tensor: + """ + Concatenates audio chunks with cross-fading between consecutive chunks. + + Args: + chunks: List of audio tensors, each with shape (C, T) where + C = number of channel, T = time dimension (samples) + fade_duration: Duration of cross-fade in seconds + sample_rate: Audio sample rate in Hz + + Returns: + Concatenated audio tensor with shape (N, T_total) + """ + # Handle edge cases: empty input or single chunk + if len(chunks) <= 1: + return chunks[0] if chunks else torch.tensor([]) + + # Calculate total fade samples from duration and sample rate + fade_samples = int(fade_duration * sample_rate) + + # Use simple concatenation if fade duration is non-positive + if fade_samples <= 0: + return torch.cat(chunks, dim=-1) + + # Initialize final tensor with the first chunk + final = chunks[0] + + # Iterate through remaining chunks to apply cross-fading + for next_chunk in chunks[1:]: + # Calculate safe fade length (cannot exceed either chunk's duration) + k = min(fade_samples, final.shape[-1], next_chunk.shape[-1]) + + # Fall back to simple concatenation if safe fade length is invalid + if k <= 0: + final = torch.cat([final, next_chunk], dim=-1) + continue + + # Create fade curve (1 -> 0) with shape (1, k) for broadcasting + fade = torch.linspace(1, 0, k, device=final.device)[None] + + # Concatenate three parts: + # 1. Non-overlapping part of previous audio + # 2. Cross-faded overlapping region + # 3. Non-overlapping part of next audio + final = torch.cat( + [ + final[..., :-k], # All samples except last k from previous + final[..., -k:] * fade + + next_chunk[..., :k] * (1 - fade), # Cross-fade region + next_chunk[..., k:], # All samples except first k from next + ], + dim=-1, + ) + + return final + + +def add_punctuation(text: str): + """Add punctuation if there is not in the end of text""" + text = text.strip() + if text[-1] not in punctuation: + text += "." + return text + + +def load_prompt_wav(prompt_wav: str, sampling_rate: int): + """ + Load the waveform with torchaudio and resampling if needed. + + Parameters: + prompt_wav: path of the prompt wav. + sampling_rate: target sampling rate. + + Returns: + Loaded prompt waveform with target sampling rate, + PyTorch tensor of shape (C, T) + """ + prompt_wav, prompt_sampling_rate = torchaudio.load(prompt_wav) + + if prompt_sampling_rate != sampling_rate: + resampler = torchaudio.transforms.Resample( + orig_freq=prompt_sampling_rate, new_freq=sampling_rate + ) + prompt_wav = resampler(prompt_wav) + return prompt_wav + + +def rms_norm(prompt_wav: torch.Tensor, target_rms: float): + """ + Normalize the rms of prompt_wav is it is smaller than target rms. + + Parameters: + prompt_wav: PyTorch tensor with shape (C, T). + target_rms: target rms value + + Returns: + prompt_wav: normalized prompt wav with shape (C, T). + promt_rms: rms of original prompt wav. Will be used to + re-normalize the generated wav. + """ + prompt_rms = torch.sqrt(torch.mean(torch.square(prompt_wav))) + if prompt_rms < target_rms: + prompt_wav = prompt_wav * target_rms / prompt_rms + return prompt_wav, prompt_rms + + +def remove_silence( + audio: torch.Tensor, + sampling_rate: int, + only_edge: bool = False, + trail_sil: float = 0, +): + """ + Remove silences longer than 1 second, and edge silences longer than 0.1 seconds + + Parameters: + audio: PyTorch tensor with shape (C, T). + sampling_rate: sampling rate of the audio. + only_edge: If true, only remove edge silences. + trail_sil: the duration of added trailing silence in ms. + + Returns: + PyTorch tensor with shape (C, T), where C is number of channels + and T is number of audio samples + """ + # Load audio file + wave = tensor_to_audiosegment(audio, sampling_rate) + + if not only_edge: + # Split audio using silences longer than 1 second + non_silent_segs = split_on_silence( + wave, + min_silence_len=1000, # Silences longer than 1 second (1000ms) + silence_thresh=-50, + keep_silence=1000, # Keep 1.0 second of silence around segments + seek_step=10, + ) + + # Concatenate all non-silent segments + wave = AudioSegment.silent(duration=0) + for seg in non_silent_segs: + wave += seg + + # Remove silence longer than 0.1 seconds in the begining and ending of wave + wave = remove_silence_edges(wave, 100, -50) + + # Add trailing silence to avoid leaking prompt to generated speech. + wave = wave + AudioSegment.silent(duration=trail_sil) + + # Convert to PyTorch tensor + return audiosegment_to_tensor(wave) + + +def remove_silence_edges( + audio: AudioSegment, keep_silence: int = 100, silence_threshold: float = -50 +): + """ + Remove edge silences longer than `keep_silence` ms. + + Parameters: + audio: an AudioSegment object. + keep_silence: kept silence in the edge. + only_edge: If true, only remove edge silences. + silence_threshold: the threshold of silence. + + Returns: + An AudioSegment object + """ + # Remove leading silence + start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold) + start_idx = max(0, start_idx - keep_silence) + audio = audio[start_idx:] + + # Remove trailing silence + audio = audio.reverse() + start_idx = detect_leading_silence(audio, silence_threshold=silence_threshold) + start_idx = max(0, start_idx - keep_silence) + audio = audio[start_idx:] + audio = audio.reverse() + + return audio + + +def audiosegment_to_tensor(aseg): + """ + Convert a pydub.AudioSegment to PyTorch audio tensor + """ + audio_data = np.array(aseg.get_array_of_samples()) + + # Convert to float32 and normalize to [-1, 1] range + audio_data = audio_data.astype(np.float32) / 32768.0 + + # Handle channels + if aseg.channels == 1: + # Mono channel: add channel dimension (T) -> (1, T) + tensor_data = torch.from_numpy(audio_data).unsqueeze(0) + else: + # Multi-channel: reshape to (C, T) + tensor_data = torch.from_numpy(audio_data.reshape(-1, aseg.channels).T) + + return tensor_data + + +def tensor_to_audiosegment(tensor, sample_rate): + """ + Convert a PyTorch audio tensor to pydub.AudioSegment + + Parameters: + tensor: Tensor with shape (C, T), where C is the number of channels + and T is the time steps + sample_rate: Audio sample rate + """ + # Convert tensor to numpy array + audio_np = tensor.cpu().numpy() + + # Add channel dimension if single channel + if audio_np.ndim == 1: + audio_np = audio_np[np.newaxis, :] + + # Convert to int16 type (common format for pydub) + # Assumes tensor values are in [-1, 1] range as floating point + audio_np = (audio_np * 32768.0).clip(-32768, 32767).astype(np.int16) + + # Convert to byte stream + # For multi-channel audio, pydub requires interleaved format + # (e.g., left-right-left-right) + if audio_np.shape[0] > 1: + # Convert to interleaved format + audio_np = audio_np.transpose(1, 0).flatten() + audio_bytes = audio_np.tobytes() + + # Create AudioSegment + audio_segment = AudioSegment( + data=audio_bytes, + sample_width=2, + frame_rate=sample_rate, + channels=tensor.shape[0], + ) + + return audio_segment diff --git a/zipvoice/utils/lr_scheduler.py b/zipvoice/utils/lr_scheduler.py new file mode 100644 index 0000000..c6b6b7b --- /dev/null +++ b/zipvoice/utils/lr_scheduler.py @@ -0,0 +1,245 @@ +# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey) +# +# See ../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import logging +from typing import List, Optional, Union + +import torch +from torch.optim import Optimizer + + +class LRScheduler(object): + """ + Base-class for learning rate schedulers where the learning-rate depends on both the + batch and the epoch. + """ + + def __init__(self, optimizer: Optimizer, verbose: bool = False): + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError("{} is not an Optimizer".format(type(optimizer).__name__)) + self.optimizer = optimizer + self.verbose = verbose + + for group in optimizer.param_groups: + group.setdefault("base_lr", group["lr"]) + + self.base_lrs = [group["base_lr"] for group in optimizer.param_groups] + + self.epoch = 0 + self.batch = 0 + + def state_dict(self): + """Returns the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in self.__dict__ which + is not the optimizer. + """ + return { + # the user might try to override the base_lr, so don't include this in the + # state. previously they were included. + # "base_lrs": self.base_lrs, + "epoch": self.epoch, + "batch": self.batch, + } + + def load_state_dict(self, state_dict): + """Loads the schedulers state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + # the things with base_lrs are a work-around for a previous problem + # where base_lrs were written with the state dict. + base_lrs = self.base_lrs + self.__dict__.update(state_dict) + self.base_lrs = base_lrs + + def get_last_lr(self) -> List[float]: + """Return last computed learning rate by current scheduler. + Will be a list of float.""" + return self._last_lr + + def get_lr(self): + # Compute list of learning rates from self.epoch and self.batch and + # self.base_lrs; this must be overloaded by the user. + # e.g. return [some_formula(self.batch, self.epoch, base_lr) + # for base_lr in self.base_lrs ] + raise NotImplementedError + + def step_batch(self, batch: Optional[int] = None) -> None: + # Step the batch index, or just set it. If `batch` is specified, it + # must be the batch index from the start of training, i.e. summed over + # all epochs. + # You can call this in any order; if you don't provide 'batch', it should + # of course be called once per batch. + if batch is not None: + self.batch = batch + else: + self.batch = self.batch + 1 + self._set_lrs() + + def step_epoch(self, epoch: Optional[int] = None): + # Step the epoch index, or just set it. If you provide the 'epoch' arg, you + # should call this at the start of the epoch; if you don't provide the 'epoch' + # arg, you should call it at the end of the epoch. + if epoch is not None: + self.epoch = epoch + else: + self.epoch = self.epoch + 1 + self._set_lrs() + + def _set_lrs(self): + values = self.get_lr() + assert len(values) == len(self.optimizer.param_groups) + + for i, data in enumerate(zip(self.optimizer.param_groups, values)): + param_group, lr = data + param_group["lr"] = lr + self.print_lr(self.verbose, i, lr) + self._last_lr = [group["lr"] for group in self.optimizer.param_groups] + + def print_lr(self, is_verbose, group, lr): + """Display the current learning rate.""" + if is_verbose: + logging.warning( + f"Epoch={self.epoch}, batch={self.batch}: adjusting learning rate" + f" of group {group} to {lr:.4e}." + ) + + +class Eden(LRScheduler): + """ + Eden scheduler. + The basic formula (before warmup) is: + lr = base_lr * (((batch**2 + lr_batches**2) / lr_batches**2) ** -0.25 * + (((epoch**2 + lr_epochs**2) / lr_epochs**2) ** -0.25)) * warmup + where `warmup` increases from linearly 0.5 to 1 over `warmup_batches` batches + and then stays constant at 1. + + If you don't have the concept of epochs, or one epoch takes a very long time, + you can replace the notion of 'epoch' with some measure of the amount of data + processed, e.g. hours of data or frames of data, with 'lr_epochs' being set to + some measure representing "quite a lot of data": say, one fifth or one third + of an entire training run, but it doesn't matter much. You could also use + Eden2 which has only the notion of batches. + + We suggest base_lr = 0.04 (passed to optimizer) if used with ScaledAdam + + Args: + optimizer: the optimizer to change the learning rates on + lr_batches: the number of batches after which we start significantly + decreasing the learning rate, suggest 5000. + lr_epochs: the number of epochs after which we start significantly + decreasing the learning rate, suggest 6 if you plan to do e.g. + 20 to 40 epochs, but may need smaller number if dataset is huge + and you will do few epochs. + """ + + def __init__( + self, + optimizer: Optimizer, + lr_batches: Union[int, float], + lr_epochs: Union[int, float], + warmup_batches: Union[int, float] = 500.0, + warmup_start: float = 0.5, + verbose: bool = False, + ): + super(Eden, self).__init__(optimizer, verbose) + self.lr_batches = lr_batches + self.lr_epochs = lr_epochs + self.warmup_batches = warmup_batches + + assert 0.0 <= warmup_start <= 1.0, warmup_start + self.warmup_start = warmup_start + + def get_lr(self): + factor = ( + (self.batch**2 + self.lr_batches**2) / self.lr_batches**2 + ) ** -0.25 * ( + ((self.epoch**2 + self.lr_epochs**2) / self.lr_epochs**2) ** -0.25 + ) + warmup_factor = ( + 1.0 + if self.batch >= self.warmup_batches + else self.warmup_start + + (1.0 - self.warmup_start) * (self.batch / self.warmup_batches) + # else 0.5 + 0.5 * (self.batch / self.warmup_batches) + ) + + return [x * factor * warmup_factor for x in self.base_lrs] + + +class FixedLRScheduler(LRScheduler): + """ + Fixed learning rate scheduler. + + Args: + optimizer: the optimizer to change the learning rates on + """ + + def __init__( + self, + optimizer: Optimizer, + verbose: bool = False, + ): + super(FixedLRScheduler, self).__init__(optimizer, verbose) + + def get_lr(self): + + return [x for x in self.base_lrs] + + +def _test_eden(): + m = torch.nn.Linear(100, 100) + from zipvoice.utils.optim import ScaledAdam + + optim = ScaledAdam(m.parameters(), lr=0.03) + + scheduler = Eden(optim, lr_batches=100, lr_epochs=2, verbose=True) + + for epoch in range(10): + scheduler.step_epoch(epoch) # sets epoch to `epoch` + + for step in range(20): + x = torch.randn(200, 100).detach() + x.requires_grad = True + y = m(x) + dy = torch.randn(200, 100).detach() + f = (y * dy).sum() + f.backward() + + optim.step() + scheduler.step_batch() + optim.zero_grad() + + logging.info(f"last lr = {scheduler.get_last_lr()}") + logging.info(f"state dict = {scheduler.state_dict()}") + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + logging.getLogger().setLevel(logging.INFO) + import subprocess + + s = subprocess.check_output( + "git status -uno .; git log -1; git diff HEAD .", shell=True + ) + logging.info(s) + + _test_eden() diff --git a/zipvoice/utils/optim.py b/zipvoice/utils/optim.py new file mode 100644 index 0000000..f90466a --- /dev/null +++ b/zipvoice/utils/optim.py @@ -0,0 +1,868 @@ +# Copyright 2022 Xiaomi Corp. (authors: Daniel Povey) +# +# See ../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import contextlib +import logging +from collections import defaultdict +from typing import Dict, List, Tuple + +import torch +from lhotse.utils import fix_random_seed +from torch import Tensor +from torch.optim import Optimizer + + +class BatchedOptimizer(Optimizer): + """ + This class adds to class Optimizer the capability to optimize parameters in batches: + it will stack the parameters and their grads for you so the optimizer can work + on tensors with an extra leading dimension. This is intended for speed with GPUs, + as it reduces the number of kernels launched in the optimizer. + + Args: + params: + """ + + def __init__(self, params, defaults): + super(BatchedOptimizer, self).__init__(params, defaults) + + @contextlib.contextmanager + def batched_params(self, param_group, group_params_names): + """ + This function returns (technically, yields) a list of + of tuples (p, state), where + p is a `fake` parameter that is stacked (over axis 0) from real parameters + that share the same shape, and its gradient is also stacked; + `state` is the state corresponding to this batch of parameters + (it will be physically located in the "state" for one of the real + parameters, the last one that has any particular shape and dtype). + + This function is decorated as a context manager so that it can + write parameters back to their "real" locations. + + The idea is, instead of doing: + + for p in group["params"]: + state = self.state[p] + ... + + you can do: + + with self.batched_params(group["params"]) as batches: + for p, state, p_names in batches: + ... + + + Args: + group: a parameter group, which is a list of parameters; should be + one of self.param_groups. + group_params_names: name for each parameter in group, + which is List[str]. + """ + batches = defaultdict( + list + ) # `batches` maps from tuple (dtype_as_str,*shape) to list of nn.Parameter + batches_names = defaultdict( + list + ) # `batches` maps from tuple (dtype_as_str,*shape) to list of str + + assert len(param_group) == len(group_params_names) + for p, named_p in zip(param_group, group_params_names): + key = (str(p.dtype), *p.shape) + batches[key].append(p) + batches_names[key].append(named_p) + + batches_names_keys = list(batches_names.keys()) + sorted_idx = sorted( + range(len(batches_names)), key=lambda i: batches_names_keys[i] + ) + batches_names = [batches_names[batches_names_keys[idx]] for idx in sorted_idx] + batches = [batches[batches_names_keys[idx]] for idx in sorted_idx] + + stacked_params_dict = dict() + + # turn batches into a list, in deterministic order. + # tuples will contain tuples of (stacked_param, state, stacked_params_names), + # one for each batch in `batches`. + tuples = [] + + for batch, batch_names in zip(batches, batches_names): + p = batch[0] + # we arbitrarily store the state in the + # state corresponding to the 1st parameter in the + # group. class Optimizer will take care of saving/loading state. + state = self.state[p] + p_stacked = torch.stack(batch) + grad = torch.stack( + [torch.zeros_like(p) if p.grad is None else p.grad for p in batch] + ) + p_stacked.grad = grad + stacked_params_dict[key] = p_stacked + tuples.append((p_stacked, state, batch_names)) + + yield tuples # <-- calling code will do the actual optimization here! + + for (stacked_params, _state, _names), batch in zip(tuples, batches): + for i, p in enumerate(batch): # batch is list of Parameter + p.copy_(stacked_params[i]) + + +def basic_step(group, p, state, grad): + # computes basic Adam update using beta2 (dividing by gradient stddev) only. no + # momentum yet. + lr = group["lr"] + if p.numel() == p.shape[0]: + lr = lr * group["scalar_lr_scale"] + beta2 = group["betas"][1] + eps = group["eps"] + # p shape: (batch_size,) or (batch_size, 1, [1,..]) + try: + exp_avg_sq = state[ + "exp_avg_sq" + ] # shape: (batch_size,) or (batch_size, 1, [1,..]) + except KeyError: + exp_avg_sq = torch.zeros(*p.shape, device=p.device, dtype=torch.float) + state["exp_avg_sq"] = exp_avg_sq + + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + # bias_correction2 is like in Adam. + # slower update at the start will help stability anyway. + bias_correction2 = 1 - beta2 ** (state["step"] + 1) + if bias_correction2 < 0.99: + # note: not in-place. + exp_avg_sq = exp_avg_sq * (1.0 / bias_correction2) + denom = exp_avg_sq.sqrt().add_(eps) + + return -lr * grad / denom + + +def scaling_step(group, p, state, grad): + delta = basic_step(group, p, state, grad) + if p.numel() == p.shape[0]: + return delta + # there is no scaling for scalar parameters. + # (p.shape[0] is the batch of parameters.) + + step = state["step"] + size_update_period = group["size_update_period"] + + try: + param_rms = state["param_rms"] + scale_grads = state["scale_grads"] + scale_exp_avg_sq = state["scale_exp_avg_sq"] + except KeyError: + # we know p.ndim > 1 because we'd have returned above if not, so don't worry + # about the speial case of dim=[] that pytorch treats inconsistently. + param_rms = (p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt() + param_rms = param_rms.to(torch.float) + scale_exp_avg_sq = torch.zeros_like(param_rms) + scale_grads = torch.zeros( + size_update_period, + *param_rms.shape, + dtype=torch.float, + device=p.device, + ) + state["param_rms"] = param_rms + state["scale_grads"] = scale_grads + state["scale_exp_avg_sq"] = scale_exp_avg_sq + + # on every step, update the gradient w.r.t. the scale of the parameter, we + # store these as a batch and periodically update the size (for speed only, to + # avoid too many operations). + scale_grads[step % size_update_period] = (p * grad).sum( + dim=list(range(1, p.ndim)), keepdim=True + ) + + # periodically recompute the value of param_rms. + if step % size_update_period == size_update_period - 1: + param_rms.copy_((p**2).mean(dim=list(range(1, p.ndim)), keepdim=True).sqrt()) + + param_min_rms = group["param_min_rms"] + + # scale the step size by param_rms. This is the most important "scaling" part of + # ScaledAdam + delta *= param_rms.clamp(min=param_min_rms) + + if step % size_update_period == size_update_period - 1 and step > 0: + # This block updates the size of parameter by adding a step ("delta") value in + # the direction of either shrinking or growing it. + beta2 = group["betas"][1] + size_lr = group["lr"] * group["scalar_lr_scale"] + param_max_rms = group["param_max_rms"] + eps = group["eps"] + # correct beta2 for the size update period: we will have + # faster decay at this level. + beta2_corr = beta2**size_update_period + scale_exp_avg_sq.mul_(beta2_corr).add_( + (scale_grads**2).mean(dim=0), # mean over dim `size_update_period` + alpha=1 - beta2_corr, + ) # shape is (batch_size, 1, 1, ...) + + # The 1st time we reach here is when size_step == 1. + size_step = (step + 1) // size_update_period + bias_correction2 = 1 - beta2_corr**size_step + + denom = scale_exp_avg_sq.sqrt() + eps + + scale_step = ( + -size_lr * (bias_correction2**0.5) * scale_grads.sum(dim=0) / denom + ) + + is_too_small = param_rms < param_min_rms + + # when the param gets too small, just don't shrink it any further. + scale_step.masked_fill_(is_too_small, 0.0) + + # The following may help prevent instability: don't allow the scale step to be + # too large in either direction. + scale_step.clamp_(min=-0.1, max=0.1) + + # and ensure the parameter rms after update never exceeds param_max_rms. + # We have to look at the trained model for parameters at or around the + # param_max_rms, because sometimes they can indicate a problem with the + # topology or settings. + scale_step = torch.minimum(scale_step, (param_max_rms - param_rms) / param_rms) + + delta.add_(p * scale_step) + + return delta + + +def momentum_step(group, p, state, grad): + delta = scaling_step(group, p, state, grad) + beta1 = group["betas"][0] + try: + stored_delta = state["delta"] + except KeyError: + stored_delta = torch.zeros(*p.shape, device=p.device, dtype=torch.float) + state["delta"] = stored_delta + stored_delta.mul_(beta1) + stored_delta.add_(delta, alpha=(1 - beta1)) + # we don't bother doing the "bias correction" part of Adam for beta1 because this is + # just an edge effect that affects the first 10 or so batches; and the effect of not + # doing it is just to do a slower update for the first few batches, which will help + # stability. + return stored_delta + + +class ScaledAdam(BatchedOptimizer): + """ + Implements 'Scaled Adam', a variant of Adam where we scale each parameter's update + proportional to the norm of that parameter; and also learn the scale of the + parameter, in log space, subject to upper and lower limits (as if we had factored + each parameter as param = underlying_param * log_scale.exp()) + + + Args: + params: The parameters or param_groups to optimize (like other Optimizer + subclasses) Unlike common optimizers, which accept + model.parameters() or groups of parameters(), this optimizer + could accept model.named_parameters() or groups of + named_parameters(). See comments of function + _get_names_of_parameters for its 4 possible cases. + lr: The learning rate. We will typically use a learning rate schedule + that starts at 0.03 and decreases over time, i.e. much higher + than other common optimizers. + clipping_scale: (e.g. 2.0) + A scale for gradient-clipping: if specified, the normalized gradients + over the whole model will be clipped to have 2-norm equal to + `clipping_scale` times the median 2-norm over the most recent period + of `clipping_update_period` minibatches. By "normalized gradients", + we mean after multiplying by the rms parameter value for this tensor + [for non-scalars]; this is appropriate because our update is scaled + by this quantity. + betas: beta1,beta2 are momentum constants for regular momentum, and moving + sum-sq grad. Must satisfy 0 < beta <= beta2 < 1. + scalar_lr_scale: A scaling factor on the learning rate, that we use to update the + scale of each parameter tensor and scalar parameters of the mode.. + If each parameter were decomposed as p * p_scale.exp(), + where (p**2).mean().sqrt() == 1.0, scalar_lr_scale would be a the + scaling factor on the learning rate of p_scale. + eps: A general-purpose epsilon to prevent division by zero + param_min_rms: Minimum root-mean-square value of parameter tensor, for purposes of + learning the scale on the parameters (we'll constrain the rms of + each non-scalar parameter tensor to be >= this value) + param_max_rms: Maximum root-mean-square value of parameter tensor, for purposes of + learning the scale on the parameters (we'll constrain the rms of + each non-scalar parameter tensor to be <= this value) + scalar_max: Maximum absolute value for scalar parameters (applicable if your + model has any parameters with numel() == 1). + size_update_period: The periodicity, in steps, with which we update the size (scale) + of the parameter tensor. This is provided to save a little time + in the update. + clipping_update_period: if clipping_scale is specified, this is the period + """ + + def __init__( + self, + params, + lr=3e-02, + clipping_scale=None, + betas=(0.9, 0.98), + scalar_lr_scale=0.1, + eps=1.0e-08, + param_min_rms=1.0e-05, + param_max_rms=3.0, + scalar_max=10.0, + size_update_period=4, + clipping_update_period=100, + ): + + defaults = dict( + lr=lr, + clipping_scale=clipping_scale, + betas=betas, + scalar_lr_scale=scalar_lr_scale, + eps=eps, + param_min_rms=param_min_rms, + param_max_rms=param_max_rms, + scalar_max=scalar_max, + size_update_period=size_update_period, + clipping_update_period=clipping_update_period, + ) + + # If params only contains parameters or group of parameters, + # i.e when parameter names are not given, + # this flag will be set to False in funciton _get_names_of_parameters. + self.show_dominant_parameters = True + param_groups, parameters_names = self._get_names_of_parameters(params) + super(ScaledAdam, self).__init__(param_groups, defaults) + assert len(self.param_groups) == len(parameters_names) + self.parameters_names = parameters_names + + def _get_names_of_parameters( + self, params_or_named_params + ) -> Tuple[List[Dict], List[List[str]]]: + """ + Args: + params_or_named_params: according to the way ScaledAdam is initialized + in train.py, this argument could be one of following 4 cases, + case 1, a generator of parameter, e.g.: + optimizer = ScaledAdam(model.parameters(), lr=params.base_lr, + clipping_scale=3.0) + + case 2, a list of parameter groups with different config, e.g.: + model_param_groups = [ + {'params': model.encoder.parameters(), 'lr': 0.05}, + {'params': model.decoder.parameters(), 'lr': 0.01}, + {'params': model.joiner.parameters(), 'lr': 0.03}, + ] + optimizer = ScaledAdam(model_param_groups, lr=params.base_lr, + clipping_scale=3.0) + + case 3, a generator of named_parameter, e.g.: + optimizer = ScaledAdam(model.named_parameters(), lr=params.base_lr, + clipping_scale=3.0) + + case 4, a list of named_parameter groups with different config, e.g.: + model_named_param_groups = [ + {'named_params': model.encoder.named_parameters(), 'lr': 0.05}, + {'named_params': model.decoder.named_parameters(), 'lr': 0.01}, + {'named_params': model.joiner.named_parameters(), 'lr': 0.03}, + ] + optimizer = ScaledAdam(model_named_param_groups, lr=params.base_lr, + clipping_scale=3.0) + + For case 1 and case 2, input params is used to initialize the underlying + torch.optimizer. + For case 3 and case 4, firstly, names and params are extracted from input + named_params, then, these extracted params are used to initialize the + underlying torch.optimizer, and these extracted names are mainly used by + function `_show_gradient_dominating_parameter` + + Returns: + Returns a tuple containing 2 elements: + - `param_groups` with type List[Dict], each Dict element is a parameter + group. An example of `param_groups` could be: + [ + {'params': `one iterable of Parameter`, 'lr': 0.05}, + {'params': `another iterable of Parameter`, 'lr': 0.08}, + {'params': `a third iterable of Parameter`, 'lr': 0.1}, + ] + - `param_gruops_names` with type List[List[str]], + each `List[str]` is for a group['params'] in param_groups, + and each `str` is the name of a parameter. + A dummy name "foo" is related to each parameter, + if input are params without names, i.e. case 1 or case 2. + """ + # variable naming convention in this function: + # p is short for param. + # np is short for named_param. + # p_or_np is short for param_or_named_param. + # cur is short for current. + # group is a dict, + # e.g. {'params': iterable of parameter, 'lr': 0.05, other fields}. + # groups is a List[group] + + iterable_or_groups = list(params_or_named_params) + if len(iterable_or_groups) == 0: + raise ValueError("optimizer got an empty parameter list") + + # The first value of returned tuple. A list of dicts containing at + # least 'params' as a key. + param_groups = [] + + # The second value of returned tuple, + # a List[List[str]], each sub-List is for a group. + param_groups_names = [] + + if not isinstance(iterable_or_groups[0], dict): + # case 1 or case 3, + # the input is an iterable of parameter or named parameter. + param_iterable_cur_group = [] + param_names_cur_group = [] + for p_or_np in iterable_or_groups: + if isinstance(p_or_np, tuple): + # case 3 + name, param = p_or_np + else: + # case 1 + assert isinstance(p_or_np, torch.Tensor) + param = p_or_np + # Assign a dummy name as a placeholder + name = "foo" + self.show_dominant_parameters = False + param_iterable_cur_group.append(param) + param_names_cur_group.append(name) + param_groups.append({"params": param_iterable_cur_group}) + param_groups_names.append(param_names_cur_group) + else: + # case 2 or case 4 + # the input is groups of parameter or named parameter. + for cur_group in iterable_or_groups: + if "named_params" in cur_group: + name_list = [x[0] for x in cur_group["named_params"]] + p_list = [x[1] for x in cur_group["named_params"]] + del cur_group["named_params"] + cur_group["params"] = p_list + else: + assert "params" in cur_group + name_list = ["foo" for _ in cur_group["params"]] + param_groups.append(cur_group) + param_groups_names.append(name_list) + + return param_groups, param_groups_names + + def __setstate__(self, state): + super(ScaledAdam, self).__setstate__(state) + + @torch.no_grad() + def step(self, closure=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group, group_params_names in zip(self.param_groups, self.parameters_names): + + with self.batched_params(group["params"], group_params_names) as batches: + + # batches is list of pairs (stacked_param, state). stacked_param is + # like a regular parameter, and will have a .grad, but the 1st dim + # corresponds to a stacking dim, it is not a real dim. + + if ( + len(batches[0][1]) == 0 + ): # if len(first state) == 0: not yet initialized + clipping_scale = 1 + else: + clipping_scale = self._get_clipping_scale(group, batches) + + for p, state, _ in batches: + # Perform optimization step. + # grad is not going to be None, we handled that when creating the + # batches. + grad = p.grad + if grad.is_sparse: + raise RuntimeError( + "ScaledAdam optimizer does not support sparse gradients" + ) + + try: + cur_step = state["step"] + except KeyError: + state["step"] = 0 + cur_step = 0 + + grad = ( + p.grad if clipping_scale == 1.0 else p.grad.mul_(clipping_scale) + ) + p += momentum_step(group, p.detach(), state, grad) + + if p.numel() == p.shape[0]: # scalar parameter + scalar_max = group["scalar_max"] + p.clamp_(min=-scalar_max, max=scalar_max) + + state["step"] = cur_step + 1 + + return loss + + def _get_clipping_scale( + self, group: dict, tuples: List[Tuple[Tensor, dict, List[str]]] + ) -> float: + """ + Returns a scalar factor <= 1.0 that dictates gradient clipping, i.e. we will + scale the gradients by this amount before applying the rest of the update. + + Args: + group: the parameter group, an item in self.param_groups + tuples: a list of tuples of (param, state, param_names) + where param is a batched set of parameters, + with a .grad (1st dim is batch dim) + and state is the state-dict where optimization parameters are kept. + param_names is a List[str] while each str is name for a parameter + in batched set of parameters "param". + """ + assert len(tuples) >= 1 + clipping_scale = group["clipping_scale"] + (first_p, first_state, _) = tuples[0] + step = first_state["step"] + if clipping_scale is None or step == 0: + # no clipping. return early on step == 0 because the other + # parameters' state won't have been initialized yet. + return 1.0 + clipping_update_period = group["clipping_update_period"] + scalar_lr_scale = group["scalar_lr_scale"] + + tot_sumsq = torch.tensor(0.0, device=first_p.device) + for p, state, param_names in tuples: + grad = p.grad + if grad.is_sparse: + raise RuntimeError( + "ScaledAdam optimizer does not support sparse gradients" + ) + if p.numel() == p.shape[0]: # a batch of scalars + tot_sumsq += (grad**2).sum() * ( + scalar_lr_scale**2 + ) # sum() to change shape [1] to [] + else: + tot_sumsq += ((grad * state["param_rms"]) ** 2).sum() + + tot_norm = tot_sumsq.sqrt() + if "model_norms" not in first_state: + first_state["model_norms"] = torch.zeros( + clipping_update_period, device=p.device + ) + first_state["model_norms"][step % clipping_update_period] = tot_norm + + irregular_estimate_steps = [ + i for i in [10, 20, 40] if i < clipping_update_period + ] + if step % clipping_update_period == 0 or step in irregular_estimate_steps: + # Print some stats. + # We don't reach here if step == 0 because we would have returned + # above. + sorted_norms = first_state["model_norms"].sort()[0].to("cpu") + if step in irregular_estimate_steps: + sorted_norms = sorted_norms[-step:] + num_norms = sorted_norms.numel() + quartiles = [] + for n in range(0, 5): + index = min(num_norms - 1, (num_norms // 4) * n) + quartiles.append(sorted_norms[index].item()) + + median = quartiles[2] + if median - median != 0: + raise RuntimeError("Too many grads were not finite") + threshold = clipping_scale * median + if step in irregular_estimate_steps: + # use larger thresholds on first few steps of estimating threshold, + # as norm may be changing rapidly. + threshold = threshold * 2.0 + first_state["model_norm_threshold"] = threshold + percent_clipped = ( + first_state["num_clipped"] * 100.0 / num_norms + if "num_clipped" in first_state + else 0.0 + ) + first_state["num_clipped"] = 0 + quartiles = " ".join(["%.3e" % x for x in quartiles]) + logging.warning( + f"Clipping_scale={clipping_scale}, grad-norm quartiles {quartiles}, " + f"threshold={threshold:.3e}, percent-clipped={percent_clipped:.1f}" + ) + + try: + model_norm_threshold = first_state["model_norm_threshold"] + except KeyError: + return 1.0 # threshold has not yet been set. + + ans = min(1.0, (model_norm_threshold / (tot_norm + 1.0e-20)).item()) + if ans != ans: # e.g. ans is nan + ans = 0.0 + if ans < 1.0: + first_state["num_clipped"] += 1 + if ans < 0.5: + logging.debug( + f"Scaling gradients by {ans}, " + f"model_norm_threshold={model_norm_threshold}" + ) + if self.show_dominant_parameters: + assert p.shape[0] == len(param_names) + self._show_gradient_dominating_parameter( + tuples, tot_sumsq, group["scalar_lr_scale"] + ) + self._show_param_with_unusual_grad(tuples) + + if ans == 0.0: + for p, state, param_names in tuples: + p.grad.zero_() # get rid of infinity() + + return ans + + def _show_param_with_unusual_grad( + self, + tuples: List[Tuple[Tensor, dict, List[str]]], + ): + """ + Print information about parameter which has the largest ratio of + grad-on-this-batch divided by normal grad size. + tuples: a list of tuples of (param, state, param_names) + where param is a batched set of parameters, + with a .grad (1st dim is batch dim) + and state is the state-dict where optimization parameters are kept. + param_names is a List[str] while each str is name for a parameter + in batched set of parameters "param". + """ + # ratios_names is a list of 3-tuples: (grad_ratio, param_name, tensor) + ratios_names = [] + for p, state, batch_param_names in tuples: + dims = list(range(1, p.ndim)) + + def mean(x): + # workaround for bad interface of torch's "mean" for when dims is the + # empty list. + if len(dims) > 0: + return x.mean(dim=dims) + else: + return x + + grad_ratio = ( + (mean(p.grad**2) / state["exp_avg_sq"].mean(dim=dims)) + .sqrt() + .to("cpu") + ) + + ratios_names += zip( + grad_ratio.tolist(), batch_param_names, p.grad.unbind(dim=0) + ) + + ratios_names = sorted(ratios_names, reverse=True) + ratios_names = ratios_names[:10] + ratios_names = [ + (ratio, name, largest_index(tensor)) + for (ratio, name, tensor) in ratios_names + ] + + logging.debug( + f"Parameters with most larger-than-usual grads, with ratios, " + f"are: {ratios_names}" + ) + + def _show_gradient_dominating_parameter( + self, + tuples: List[Tuple[Tensor, dict, List[str]]], + tot_sumsq: Tensor, + scalar_lr_scale: float, + ): + """ + Show information of parameter which dominates tot_sumsq. + + Args: + tuples: a list of tuples of (param, state, param_names) + where param is a batched set of parameters, + with a .grad (1st dim is batch dim) + and state is the state-dict where optimization parameters are kept. + param_names is a List[str] while each str is name for a parameter + in batched set of parameters "param". + tot_sumsq: sumsq of all parameters. Though it's could be calculated + from tuples, we still pass it to save some time. + """ + all_sumsq_orig = {} + for p, state, batch_param_names in tuples: + # p is a stacked batch parameters. + batch_grad = p.grad + if p.numel() == p.shape[0]: # a batch of scalars + # Dummy values used by following `zip` statement. + batch_rms_orig = torch.full( + p.shape, scalar_lr_scale, device=batch_grad.device + ) + else: + batch_rms_orig = state["param_rms"] + batch_sumsq_orig = (batch_grad * batch_rms_orig) ** 2 + if batch_grad.ndim > 1: + # need to guard it with if-statement because sum() sums over + # all dims if dim == (). + batch_sumsq_orig = batch_sumsq_orig.sum( + dim=list(range(1, batch_grad.ndim)) + ) + for name, sumsq_orig, rms, grad in zip( + batch_param_names, batch_sumsq_orig, batch_rms_orig, batch_grad + ): + + proportion_orig = sumsq_orig / tot_sumsq + all_sumsq_orig[name] = (proportion_orig, sumsq_orig, rms, grad) + + sorted_by_proportion = { + k: v + for k, v in sorted( + all_sumsq_orig.items(), + key=lambda item: item[1][0], + reverse=True, + ) + } + dominant_param_name = next(iter(sorted_by_proportion)) + ( + dominant_proportion, + dominant_sumsq, + dominant_rms, + dominant_grad, + ) = sorted_by_proportion[dominant_param_name] + logging.debug( + f"Parameter dominating tot_sumsq {dominant_param_name}" + f" with proportion {dominant_proportion:.2f}," + f" where dominant_sumsq=(grad_sumsq*orig_rms_sq)" + f"={dominant_sumsq:.3e}," + f" grad_sumsq={(dominant_grad**2).sum():.3e}," + f" orig_rms_sq={(dominant_rms**2).item():.3e}" + ) + + +def largest_index(x: Tensor): + x = x.contiguous() + argmax = x.abs().argmax().item() + return [(argmax // x.stride(i)) % x.size(i) for i in range(x.ndim)] + + +def _test_scaled_adam(hidden_dim: int): + import timeit + + from zipvoice.models.modules.scaling import ScaledLinear + from zipvoice.utils.lr_scheduler import Eden + + E = 100 + B = 4 + T = 2 + logging.info("in test_eve_cain") + # device = torch.device('cuda') + device = torch.device("cpu") + dtype = torch.float32 + + fix_random_seed(42) + # these input_magnitudes and output_magnitudes are to test that + # Abel is working as we expect and is able to adjust scales of + # different dims differently. + input_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp() + output_magnitudes = (1.0 * torch.randn(E, dtype=dtype, device=device)).exp() + + fix_random_seed(42) + Linear = ScaledLinear + + m = torch.nn.Sequential( + Linear(E, hidden_dim), + torch.nn.PReLU(), + Linear(hidden_dim, hidden_dim), + torch.nn.PReLU(), + Linear(hidden_dim, E), + ).to(device) + + train_pairs = [ + ( + 100.0 * torch.randn(B, T, E, device=device, dtype=dtype) * input_magnitudes, + torch.randn(B, T, E, device=device, dtype=dtype) * output_magnitudes, + ) + for _ in range(20) + ] + optim = ScaledAdam(m.named_parameters(), lr=0.03, clipping_scale=2.0) + scheduler = Eden(optim, lr_batches=200, lr_epochs=5, verbose=False) + + start = timeit.default_timer() + avg_loss = 0.0 + for epoch in range(180): + scheduler.step_epoch() + # if epoch == 100 and iter in [2,3]: + # optim.reset_speedup() # check it doesn't crash. + + # if epoch == 130: + # opts = diagnostics.TensorDiagnosticOptions( + # 512 + # ) # allow 4 megabytes per sub-module + # diagnostic = diagnostics.attach_diagnostics(m, opts) + + for n, (x, y) in enumerate(train_pairs): + y_out = m(x) + loss = ((y_out - y) ** 2).mean() * 100.0 + if epoch == 0 and n == 0: + avg_loss = loss.item() + else: + avg_loss = 0.98 * avg_loss + 0.02 * loss.item() + if n == 0 and epoch % 5 == 0: + # norm1 = '%.2e' % (m[0].weight**2).mean().sqrt().item() + # norm1b = '%.2e' % (m[0].bias**2).mean().sqrt().item() + # norm2 = '%.2e' % (m[2].weight**2).mean().sqrt().item() + # norm2b = '%.2e' % (m[2].bias**2).mean().sqrt().item() + # scale1 = '%.2e' % (m[0].weight_scale.exp().item()) + # scale1b = '%.2e' % (m[0].bias_scale.exp().item()) + # scale2 = '%.2e' % (m[2].weight_scale.exp().item()) + # scale2b = '%.2e' % (m[2].bias_scale.exp().item()) + lr = scheduler.get_last_lr()[0] + logging.info( + f"Iter {iter}, epoch {epoch}, batch {n}, " + f"avg_loss {avg_loss:.4g}, lr={lr:.4e}" + ) # , norms={norm1,norm1b,norm2,norm2b}") + # scales={scale1,scale1b,scale2,scale2b} + loss.log().backward() + optim.step() + optim.zero_grad() + scheduler.step_batch() + + # diagnostic.print_diagnostics() + + stop = timeit.default_timer() + logging.info(f"Iter={iter}, Time taken: {stop - start}") + + logging.info(f"last lr = {scheduler.get_last_lr()}") + # logging.info("state dict = ", scheduler.state_dict()) + # logging.info("optim state_dict = ", optim.state_dict()) + logging.info(f"input_magnitudes = {input_magnitudes}") + logging.info(f"output_magnitudes = {output_magnitudes}") + + +if __name__ == "__main__": + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + logging.getLogger().setLevel(logging.INFO) + import subprocess + + s = subprocess.check_output( + "git status -uno .; git log -1; git diff HEAD .", shell=True + ) + logging.info(s) + import sys + + if len(sys.argv) > 1: + hidden_dim = int(sys.argv[1]) + else: + hidden_dim = 200 + + _test_scaled_adam(hidden_dim) diff --git a/zipvoice/utils/scaling_converter.py b/zipvoice/utils/scaling_converter.py new file mode 100644 index 0000000..0f9ff21 --- /dev/null +++ b/zipvoice/utils/scaling_converter.py @@ -0,0 +1,105 @@ +# Copyright 2022-2023 Xiaomi Corp. (authors: Fangjun Kuang, +# Zengwei Yao) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This file replaces various modules in a model. +Specifically, ActivationBalancer is replaced with an identity operator; +Whiten is also replaced with an identity operator; +BasicNorm is replaced by a module with `exp` removed. +""" + +import copy +from typing import List + +import torch +import torch.nn as nn + +from zipvoice.models.modules.scaling import ( + Balancer, + Dropout3, + SwooshL, + SwooshLOnnx, + SwooshR, + SwooshROnnx, + Whiten, +) +from zipvoice.models.modules.zipformer import CompactRelPositionalEncoding + + +# Copied from https://pytorch.org/docs/1.9.0/_modules/torch/nn/modules/module.html#Module.get_submodule # noqa +# get_submodule was added to nn.Module at v1.9.0 +def get_submodule(model, target): + if target == "": + return model + atoms: List[str] = target.split(".") + mod: torch.nn.Module = model + for item in atoms: + if not hasattr(mod, item): + raise AttributeError( + mod._get_name() + " has no " "attribute `" + item + "`" + ) + mod = getattr(mod, item) + if not isinstance(mod, torch.nn.Module): + raise AttributeError("`" + item + "` is not " "an nn.Module") + return mod + + +def convert_scaled_to_non_scaled( + model: nn.Module, + inplace: bool = False, + is_pnnx: bool = False, + is_onnx: bool = False, +): + """ + Args: + model: + The model to be converted. + inplace: + If True, the input model is modified inplace. + If False, the input model is copied and we modify the copied version. + is_pnnx: + True if we are going to export the model for PNNX. + is_onnx: + True if we are going to export the model for ONNX. + Return: + Return a model without scaled layers. + """ + if not inplace: + model = copy.deepcopy(model) + + d = {} + for name, m in model.named_modules(): + if isinstance(m, (Balancer, Dropout3, Whiten)): + d[name] = nn.Identity() + elif is_onnx and isinstance(m, SwooshR): + d[name] = SwooshROnnx() + elif is_onnx and isinstance(m, SwooshL): + d[name] = SwooshLOnnx() + elif is_onnx and isinstance(m, CompactRelPositionalEncoding): + # We want to recreate the positional encoding vector when + # the input changes, so we have to use torch.jit.script() + # to replace torch.jit.trace() + d[name] = torch.jit.script(m) + + for k, v in d.items(): + if "." in k: + parent, child = k.rsplit(".", maxsplit=1) + setattr(get_submodule(model, parent), child, v) + else: + setattr(model, k, v) + + return model diff --git a/zipvoice/utils/tensorrt.py b/zipvoice/utils/tensorrt.py new file mode 100644 index 0000000..649bd66 --- /dev/null +++ b/zipvoice/utils/tensorrt.py @@ -0,0 +1,143 @@ +# Copyright 2025 Nvidia Corp. (authors: Yuekai Zhang) +# +# See ../../../../LICENSE for clarification regarding multiple authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This script provides utility functions for working with TensorRT in ZipVoice. +""" + +import logging +import os +import queue +from typing import Any, Tuple, Optional + +import torch +import torch.nn as nn + + +class TrtContextWrapper: + """A wrapper class for managing TensorRT execution contexts.""" + + def __init__( + self, trt_engine: Any, trt_concurrent: int = 1, device: str = "cuda:0" + ): + """ + Initializes the TrtContextWrapper. + + Args: + trt_engine (Any): The TensorRT engine. + trt_concurrent (int, optional): The number of concurrent contexts. Defaults to 1. + device (str, optional): The device to use. Defaults to 'cuda:0'. + """ + self.trt_context_pool = queue.Queue(maxsize=trt_concurrent) + self.trt_engine = trt_engine + self.device = device + for _ in range(trt_concurrent): + trt_context = trt_engine.create_execution_context() + trt_stream = torch.cuda.stream(torch.cuda.Stream(torch.device(device))) + assert trt_context is not None, 'failed to create trt context, maybe not enough CUDA memory, try reduce current trt concurrent {}'.format(trt_concurrent) + self.trt_context_pool.put([trt_context, trt_stream]) + assert self.trt_context_pool.empty() is False, 'no avaialbe estimator context' + self.feat_dim = 100 + + def acquire_estimator(self) -> Tuple[list, Any]: + """Acquires a TensorRT context from the pool.""" + return self.trt_context_pool.get(), self.trt_engine + + def release_estimator(self, context: Any, stream: Any): + """ + Releases a TensorRT context back to the pool. + + Args: + context (Any): The TensorRT context. + stream (Any): The CUDA stream. + """ + self.trt_context_pool.put([context, stream]) + + def __call__( + self, + x: torch.Tensor, + t: torch.Tensor, + padding_mask: torch.Tensor, + guidance_scale: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """ + Executes the TensorRT engine. + + Args: + x (torch.Tensor): The input tensor. + t (torch.Tensor): The time tensor. + padding_mask (torch.Tensor): The padding mask tensor. + guidance_scale (torch.Tensor): The guidance scale tensor. + + Returns: + torch.Tensor: The output tensor. + """ + x = x.to(torch.float16) + t = t.to(torch.float16) + padding_mask = padding_mask.to(torch.float16) + if guidance_scale is not None: + guidance_scale = guidance_scale.to(torch.float16) + [estimator, stream], trt_engine = self.acquire_estimator() + # NOTE need to synchronize when switching stream + torch.cuda.current_stream().synchronize() + batch_size = x.size(0) + seq_len = x.size(1) + + # Create output tensor with shape (N, T, 100) + output = torch.empty(batch_size, seq_len, self.feat_dim, dtype=x.dtype, device=x.device) + + with stream: + estimator.set_input_shape('x', (batch_size, x.size(1), x.size(2))) + estimator.set_input_shape('t', (batch_size,)) + estimator.set_input_shape('padding_mask', (batch_size, padding_mask.size(1))) + if guidance_scale is not None: + estimator.set_input_shape('guidance_scale', (batch_size,)) + + # Set input tensor addresses + input_data_ptrs = [x.contiguous().data_ptr(), t.contiguous().data_ptr(), padding_mask.contiguous().data_ptr()] + if guidance_scale is not None: + input_data_ptrs.append(guidance_scale.contiguous().data_ptr()) + for i, j in enumerate(input_data_ptrs): + estimator.set_tensor_address(trt_engine.get_tensor_name(i), j) + + # Set output tensor address + # The output tensor name should be the last tensor name in the engine + num_tensors = trt_engine.num_io_tensors + output_tensor_name = trt_engine.get_tensor_name(num_tensors - 1) # Last tensor is output + estimator.set_tensor_address(output_tensor_name, output.contiguous().data_ptr()) + + # run trt engine + assert estimator.execute_async_v3(torch.cuda.current_stream().cuda_stream) is True + torch.cuda.current_stream().synchronize() + self.release_estimator(estimator, stream) + return output.to(torch.float32) + +def load_trt(model: nn.Module, trt_model: str, trt_concurrent: int = 1): + """ + Loads a TensorRT engine and replaces the model's fm_decoder with a TrtContextWrapper. + + Args: + model (nn.Module): The model to modify. + trt_model (str): The path to the TensorRT engine file. + trt_concurrent (int, optional): The number of concurrent contexts. Defaults to 1. + """ + assert os.path.exists(trt_model), f"Please export trt model first." + import tensorrt as trt + with open(trt_model, 'rb') as f: + estimator_engine = trt.Runtime(trt.Logger(trt.Logger.INFO)).deserialize_cuda_engine(f.read()) + assert estimator_engine is not None, 'failed to load trt {}'.format(trt_model) + del model.fm_decoder + model.fm_decoder = TrtContextWrapper(estimator_engine, trt_concurrent=trt_concurrent, device='cuda')