Initial commit
This commit is contained in:
Yatharth Sharma
272
zipvoice/bin/compute_fbank.py
Executable file
272
zipvoice/bin/compute_fbank.py
Executable file
@@ -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!")
|
||||
229
zipvoice/bin/generate_averaged_model.py
Normal file
229
zipvoice/bin/generate_averaged_model.py
Normal file
@@ -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()
|
||||
899
zipvoice/bin/infer_zipvoice.py
Normal file
899
zipvoice/bin/infer_zipvoice.py
Normal file
@@ -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()
|
||||
1286
zipvoice/bin/infer_zipvoice_dialog.py
Normal file
1286
zipvoice/bin/infer_zipvoice_dialog.py
Normal file
File diff suppressed because it is too large
Load Diff
924
zipvoice/bin/infer_zipvoice_onnx.py
Normal file
924
zipvoice/bin/infer_zipvoice_onnx.py
Normal file
@@ -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()
|
||||
429
zipvoice/bin/onnx_export.py
Normal file
429
zipvoice/bin/onnx_export.py
Normal file
@@ -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()
|
||||
274
zipvoice/bin/prepare_dataset.py
Normal file
274
zipvoice/bin/prepare_dataset.py
Normal file
@@ -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,
|
||||
)
|
||||
103
zipvoice/bin/prepare_tokens.py
Normal file
103
zipvoice/bin/prepare_tokens.py
Normal file
@@ -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!")
|
||||
382
zipvoice/bin/tensorrt_export.py
Normal file
382
zipvoice/bin/tensorrt_export.py
Normal file
@@ -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()
|
||||
1130
zipvoice/bin/train_zipvoice.py
Normal file
1130
zipvoice/bin/train_zipvoice.py
Normal file
File diff suppressed because it is too large
Load Diff
980
zipvoice/bin/train_zipvoice_dialog.py
Normal file
980
zipvoice/bin/train_zipvoice_dialog.py
Normal file
@@ -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()
|
||||
963
zipvoice/bin/train_zipvoice_dialog_stereo.py
Normal file
963
zipvoice/bin/train_zipvoice_dialog_stereo.py
Normal file
@@ -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()
|
||||
1158
zipvoice/bin/train_zipvoice_distill.py
Normal file
1158
zipvoice/bin/train_zipvoice_distill.py
Normal file
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user