mirror of
https://github.com/k2-fsa/icefall.git
synced 2025-08-09 10:02:22 +00:00
281 lines
8.1 KiB
Python
Executable File
281 lines
8.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
#
|
|
# Copyright 2023 Xiaomi Corporation (Author: Zengwei Yao,
|
|
# Zengrui Jin,)
|
|
#
|
|
# See ../../../../LICENSE for clarification regarding multiple authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""
|
|
This script performs model inference on test set.
|
|
|
|
Usage:
|
|
./vits/infer.py \
|
|
--epoch 1000 \
|
|
--exp-dir ./vits/exp \
|
|
--max-duration 500
|
|
"""
|
|
|
|
|
|
import argparse
|
|
import logging
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import k2
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchaudio
|
|
from lhotse.features.io import KaldiReader
|
|
from tokenizer import Tokenizer
|
|
from train import get_model, get_params
|
|
from tts_datamodule import LibrittsTtsDataModule
|
|
|
|
from icefall.checkpoint import load_checkpoint
|
|
from icefall.utils import AttributeDict, setup_logger
|
|
|
|
|
|
def get_parser():
|
|
parser = argparse.ArgumentParser(
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--epoch",
|
|
type=int,
|
|
default=1000,
|
|
help="""It specifies the checkpoint to use for decoding.
|
|
Note: Epoch counts from 1.
|
|
""",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--exp-dir",
|
|
type=str,
|
|
default="vits/exp",
|
|
help="The experiment dir",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--tokens",
|
|
type=str,
|
|
default="data/tokens.txt",
|
|
help="""Path to vocabulary.""",
|
|
)
|
|
|
|
return parser
|
|
|
|
|
|
def infer_dataset(
|
|
dl: torch.utils.data.DataLoader,
|
|
subset: str,
|
|
params: AttributeDict,
|
|
model: nn.Module,
|
|
tokenizer: Tokenizer,
|
|
speaker_map: KaldiReader,
|
|
) -> None:
|
|
"""Decode dataset.
|
|
The ground-truth and generated audio pairs will be saved to `params.save_wav_dir`.
|
|
|
|
Args:
|
|
dl:
|
|
PyTorch's dataloader containing the dataset to decode.
|
|
params:
|
|
It is returned by :func:`get_params`.
|
|
model:
|
|
The neural model.
|
|
tokenizer:
|
|
Used to convert text to phonemes.
|
|
"""
|
|
|
|
# Background worker save audios to disk.
|
|
def _save_worker(
|
|
subset: str,
|
|
batch_size: int,
|
|
cut_ids: List[str],
|
|
audio: torch.Tensor,
|
|
audio_pred: torch.Tensor,
|
|
audio_lens: List[int],
|
|
audio_lens_pred: List[int],
|
|
):
|
|
for i in range(batch_size):
|
|
torchaudio.save(
|
|
str(params.save_wav_dir / subset / f"{cut_ids[i]}_gt.wav"),
|
|
audio[i : i + 1, : audio_lens[i]],
|
|
sample_rate=params.sampling_rate,
|
|
)
|
|
torchaudio.save(
|
|
str(params.save_wav_dir / subset / f"{cut_ids[i]}_pred.wav"),
|
|
audio_pred[i : i + 1, : audio_lens_pred[i]],
|
|
sample_rate=params.sampling_rate,
|
|
)
|
|
|
|
device = next(model.parameters()).device
|
|
num_cuts = 0
|
|
log_interval = 5
|
|
|
|
try:
|
|
num_batches = len(dl)
|
|
except TypeError:
|
|
num_batches = "?"
|
|
|
|
futures = []
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
for batch_idx, batch in enumerate(dl):
|
|
batch_size = len(batch["tokens"])
|
|
|
|
tokens = batch["tokens"]
|
|
tokens = tokenizer.tokens_to_token_ids(
|
|
tokens, intersperse_blank=True, add_sos=True, add_eos=True
|
|
)
|
|
tokens = k2.RaggedTensor(tokens)
|
|
row_splits = tokens.shape.row_splits(1)
|
|
tokens_lens = row_splits[1:] - row_splits[:-1]
|
|
tokens = tokens.to(device)
|
|
tokens_lens = tokens_lens.to(device)
|
|
# tensor of shape (B, T)
|
|
tokens = tokens.pad(mode="constant", padding_value=tokenizer.pad_id)
|
|
|
|
audio = batch["audio"]
|
|
audio_lens = batch["audio_lens"].tolist()
|
|
cut_ids = [cut.id for cut in batch["cut"]]
|
|
sids = ["_".join(cut_id.split("_")[:2]) for cut_id in cut_ids]
|
|
spembs = (
|
|
torch.Tensor(np.array([speaker_map.read(sid) for sid in sids]))
|
|
.squeeze(1)
|
|
.to(device)
|
|
)
|
|
|
|
audio_pred, _, durations = model.inference_batch(
|
|
text=tokens,
|
|
text_lengths=tokens_lens,
|
|
spembs=spembs,
|
|
)
|
|
audio_pred = audio_pred.detach().cpu()
|
|
# convert to samples
|
|
audio_lens_pred = (
|
|
(durations.sum(1) * params.frame_shift).to(dtype=torch.int64).tolist()
|
|
)
|
|
|
|
futures.append(
|
|
executor.submit(
|
|
_save_worker,
|
|
subset,
|
|
batch_size,
|
|
cut_ids,
|
|
audio,
|
|
audio_pred,
|
|
audio_lens,
|
|
audio_lens_pred,
|
|
)
|
|
)
|
|
|
|
num_cuts += batch_size
|
|
|
|
if batch_idx % log_interval == 0:
|
|
batch_str = f"{batch_idx}/{num_batches}"
|
|
|
|
logging.info(
|
|
f"batch {batch_str}, cuts processed until now is {num_cuts}"
|
|
)
|
|
# return results
|
|
for f in futures:
|
|
f.result()
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
parser = get_parser()
|
|
LibrittsTtsDataModule.add_arguments(parser)
|
|
args = parser.parse_args()
|
|
args.exp_dir = Path(args.exp_dir)
|
|
|
|
params = get_params()
|
|
params.update(vars(args))
|
|
|
|
params.suffix = f"epoch-{params.epoch}"
|
|
|
|
params.res_dir = params.exp_dir / "infer" / params.suffix
|
|
params.save_wav_dir = params.res_dir / "wav"
|
|
params.save_wav_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
setup_logger(f"{params.res_dir}/log-infer-{params.suffix}")
|
|
logging.info("Infer started")
|
|
|
|
device = torch.device("cpu")
|
|
if torch.cuda.is_available():
|
|
device = torch.device("cuda", 0)
|
|
|
|
tokenizer = Tokenizer(params.tokens)
|
|
params.blank_id = tokenizer.pad_id
|
|
params.vocab_size = tokenizer.vocab_size
|
|
|
|
# we need cut ids to display recognition results.
|
|
args.return_cuts = True
|
|
libritts = LibrittsTtsDataModule(args)
|
|
|
|
logging.info(f"Device: {device}")
|
|
logging.info(params)
|
|
|
|
logging.info("About to create model")
|
|
model = get_model(params)
|
|
|
|
load_checkpoint(f"{params.exp_dir}/epoch-{params.epoch}.pt", model)
|
|
|
|
model.to(device)
|
|
model.eval()
|
|
|
|
num_param_g = sum([p.numel() for p in model.generator.parameters()])
|
|
logging.info(f"Number of parameters in generator: {num_param_g}")
|
|
num_param_d = sum([p.numel() for p in model.discriminator.parameters()])
|
|
logging.info(f"Number of parameters in discriminator: {num_param_d}")
|
|
logging.info(f"Total number of parameters: {num_param_g + num_param_d}")
|
|
|
|
test_clean_cuts = libritts.test_clean_cuts()
|
|
test_clean_speaker_map = libritts.test_clean_xvector()
|
|
test_clean_dl = libritts.test_dataloaders(test_clean_cuts)
|
|
|
|
dev_clean_cuts = libritts.dev_clean_cuts()
|
|
dev_clean_speaker_map = libritts.dev_clean_xvector()
|
|
dev_clean_dl = libritts.dev_dataloaders(dev_clean_cuts)
|
|
|
|
infer_sets = {
|
|
"test-clean": (test_clean_dl, test_clean_speaker_map),
|
|
"dev-clean": (dev_clean_dl, dev_clean_speaker_map),
|
|
}
|
|
|
|
for subset, data in infer_sets.items():
|
|
save_wav_dir = params.res_dir / "wav" / subset
|
|
save_wav_dir.mkdir(parents=True, exist_ok=True)
|
|
dl, speaker_map = data
|
|
|
|
logging.info(f"Processing {subset} set, saving to {save_wav_dir}")
|
|
|
|
infer_dataset(
|
|
dl=dl,
|
|
subset=subset,
|
|
params=params,
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
speaker_map=speaker_map,
|
|
)
|
|
|
|
logging.info(f"Wav files are saved to {params.save_wav_dir}")
|
|
logging.info("Done!")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|