Initial commit

zipvoice/__init__.py

@@ -0,0 +1,7 @@
+import warnings
+
+warnings.filterwarnings(
+    "ignore",
+    category=UserWarning,
+    message="pkg_resources is deprecated as an API.*",
+)

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!")

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()

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()

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()

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()

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,
+    )

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!")

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()

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()

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()

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()

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

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

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

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

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)

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,
+        )

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

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

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("<pad>")
+            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%的人<le5>...?!9204"
+    )
+    tokenizer = EmiliaTokenizer()
+    tokens = tokenizer.texts_to_tokens([text])
+    print(f"tokens: {'|'.join(tokens[0])}")

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,
+    )

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()]

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 = "<top-level>"
+
+        # 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()

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

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 = "<top-level>"
+
+        # 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()

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

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()

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:
+        <code>
+          for p in group["params"]:
+             state = self.state[p]
+             ...
+        </code>
+        you can do:
+        <code>
+          with self.batched_params(group["params"]) as batches:
+             for p, state, p_names in batches:
+                 ...
+        </code>
+
+        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)

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

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')