Add inference notebooks

This commit is contained in:
ZDisket
2023-03-17 18:31:39 -03:00
committed by Brandon Thomas
parent e928e6b44a
commit 4aa96d4697
4 changed files with 1501 additions and 24 deletions
+8
View File
@@ -80,6 +80,14 @@ Use
tensorboard --logdir output/log
```
# Inference
Use
```
python3 infer_arpa.py --mode single --dataset SingleSpk22 --checkpoint tacotron2.pth --out-path MyAud --text "there are no faces"
```
Or check `inference.ipynb` for more lean, hands-on notebook version
to serve TensorBoard on your localhost.
<!-- The loss curves, synthesized mel-spectrograms, and audios are shown.
@@ -11,11 +11,24 @@ from torch.utils.data import DataLoader
# from pypinyin import pinyin, Style
from utils.model import get_model, get_vocoder
from utils.tools import get_configs_of, to_device, infer_one_sample #, read_lexicon
from utils.tools import get_configs_of, to_device, infer_one_sample_s #, read_lexicon
from dataset import TextDataset
from text import text_to_sequence, sequence_to_text
from text.cleaners import to_arpa
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
from model import Tacotron2
def get_tac2(configs, device):
(preprocess_config, model_config, train_config) = configs
model = Tacotron2(preprocess_config, model_config, train_config).to(device)
model.eval()
model.requires_grad_ = False
return model
def preprocess_english(text, preprocess_config):
@@ -36,14 +49,14 @@ def synthesize(model, args, configs, mel_stats, vocoder, batchs):
with torch.no_grad():
# Forward
output = model.inference(*batch[2:4], *batch[5:7])
infer_one_sample(
infer_one_sample_s(
batch,
output,
vocoder,
mel_stats,
model_config,
preprocess_config,
train_config["path"]["result_path"],
args.out_path,
args,
)
@@ -51,7 +64,6 @@ def synthesize(model, args, configs, mel_stats, vocoder, batchs):
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--restore_step", type=int, required=True)
parser.add_argument(
"--mode",
type=str,
@@ -83,6 +95,18 @@ if __name__ == "__main__":
required=True,
help="name of dataset",
)
parser.add_argument(
"--checkpoint",
type=str,
required=True,
help="Path to model file.",
)
parser.add_argument(
"--out-path",
type=str,
required=True,
help="Path to store files",
)
args = parser.parse_args()
# Check source texts
@@ -94,16 +118,17 @@ if __name__ == "__main__":
# Read Config
preprocess_config, model_config, train_config = get_configs_of(args.dataset)
configs = (preprocess_config, model_config, train_config)
with open(
os.path.join(preprocess_config["path"]["preprocessed_path"], "stats.json")
) as f:
stats = json.load(f)
mel_stats = stats["mel"]
stats = None
mel_stats = None
os.makedirs(
os.path.join(train_config["path"]["result_path"], str(args.restore_step)), exist_ok=True)
args.out_path, exist_ok=True)
# Get model
model = get_model(args, configs, device, train=False)
model = get_tac2(configs, device)
ckpt = torch.load(args.checkpoint)
model.load_state_dict(ckpt["model"])
# Load vocoder
vocoder = get_vocoder(model_config, device)
@@ -121,21 +146,26 @@ if __name__ == "__main__":
ids = raw_texts = [args.text[:100]]
# Speaker Info
load_spker_embed = model_config["multi_speaker"] \
and preprocess_config["preprocessing"]["speaker_embedder"] != 'none'
with open(os.path.join(preprocess_config["path"]["preprocessed_path"], "speakers.json")) as f:
speaker_map = json.load(f)
speakers = np.array([speaker_map[args.speaker_id]]) if model_config["multi_speaker"] else np.array([0]) # single speaker is allocated 0
spker_embed = np.load(os.path.join(
preprocess_config["path"]["preprocessed_path"],
"spker_embed",
"{}-spker_embed.npy".format(args.speaker_id),
)) if load_spker_embed else None
is_multi_speaker = model_config["multi_speaker"]
if preprocess_config["preprocessing"]["text"]["language"] == "en":
texts = np.array([preprocess_english(args.text, preprocess_config)])
if is_multi_speaker:
load_spker_embed = model_config["multi_speaker"] \
and preprocess_config["preprocessing"]["speaker_embedder"] != 'none'
with open(os.path.join(preprocess_config["path"]["preprocessed_path"], "speakers.json")) as f:
speaker_map = json.load(f)
speakers = np.array([speaker_map[args.speaker_id]]) if model_config["multi_speaker"] else np.array([0]) # single speaker is allocated 0
spker_embed = np.load(os.path.join(
preprocess_config["path"]["preprocessed_path"],
"spker_embed",
"{}-spker_embed.npy".format(args.speaker_id),
)) if load_spker_embed else None
else:
raise NotImplementedError
speakers = np.array([0])
spker_embed = None
texts = np.array([preprocess_english(to_arpa(args.text), preprocess_config)])
text_lens = np.array([len(texts[0])])
batchs = [(ids, raw_texts, speakers, texts, text_lens, max(text_lens), spker_embed)]
File diff suppressed because one or more lines are too long
+42
View File
@@ -224,7 +224,49 @@ def infer_one_sample(targets, predictions, vocoder, mel_stats, model_config, pre
if multi_speaker and args.mode == "single" else "{}.wav".format(basename)),
sampling_rate, wav_predictions[0])
def infer_one_sample_s(targets, predictions, vocoder, mel_stats, model_config, preprocess_config, path, args):
normalize = preprocess_config["preprocessing"]["mel"]["normalize"]
n_frames_per_step = model_config["decoder"]["n_frames_per_step"]
multi_speaker = model_config["multi_speaker"]
mel_predictions = predictions[1]
if normalize:
mel_predictions = Audio.tools.mel_denormalize(predictions[1], *mel_stats)
basename = targets[0][0]
src_len = targets[4][0].item()
mel_len = mel_predictions[0].shape[0]
reduced_mel_len = mel_len // n_frames_per_step
mel_prediction = mel_predictions[0].detach().transpose(0, 1)
attention = predictions[3][0, :reduced_mel_len, :src_len].detach().transpose(0, 1) # [seq_len, mel_len]
fig = plot_mel(
[
mel_prediction.cpu().numpy(),
attention.cpu().numpy(),
],
["Synthetized Spectrogram", "Decoder Alignment"],
attention=True,
)
plt.savefig(os.path.join(
path, "{}_{}.png".format(basename, args.speaker_id)\
if multi_speaker and args.mode == "single" else "{}.png".format(basename)))
plt.close()
from .model import vocoder_infer
lengths = mel_len * preprocess_config["preprocessing"]["stft"]["hop_length"]
wav_predictions = vocoder_infer(
mel_predictions.transpose(1, 2), vocoder, model_config, preprocess_config, lengths=[lengths]
)
sampling_rate = preprocess_config["preprocessing"]["audio"]["sampling_rate"]
wavfile.write(os.path.join(
path, "{}_{}.wav".format(basename, args.speaker_id)\
if multi_speaker and args.mode == "single" else "{}.wav".format(basename)),
sampling_rate, wav_predictions[0])
def plot_mel(data, titles, attention=False):
fig, axes = plt.subplots(len(data), 1, squeeze=False)
if titles is None: