diff --git a/TTS/bin/train_align_tts.py b/TTS/bin/train_align_tts.py new file mode 100644 index 00000000..3e88c673 --- /dev/null +++ b/TTS/bin/train_align_tts.py @@ -0,0 +1,541 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import argparse +import glob +import os +import sys +import time +import traceback +import numpy as np +from random import randrange + +import torch +from TTS.utils.arguments import parse_arguments, process_args +# DISTRIBUTED +from torch.nn.parallel import DistributedDataParallel as DDP_th +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler +from TTS.tts.datasets.preprocess import load_meta_data +from TTS.tts.datasets.TTSDataset import MyDataset +from TTS.tts.layers.losses import AlignTTSLoss +from TTS.tts.utils.generic_utils import setup_model +from TTS.tts.utils.io import save_best_model, save_checkpoint +from TTS.tts.utils.measures import alignment_diagonal_score +from TTS.tts.utils.speakers import parse_speakers +from TTS.tts.utils.synthesis import synthesis +from TTS.tts.utils.text.symbols import make_symbols, phonemes, symbols +from TTS.tts.utils.visual import plot_alignment, plot_spectrogram +from TTS.utils.audio import AudioProcessor +from TTS.utils.distribute import init_distributed, reduce_tensor +from TTS.utils.generic_utils import (KeepAverage, count_parameters, + remove_experiment_folder, set_init_dict) +from TTS.utils.radam import RAdam +from TTS.utils.training import NoamLR, setup_torch_training_env + + +if __name__ == '__main__': + use_cuda, num_gpus = setup_torch_training_env(True, False) + # torch.autograd.set_detect_anomaly(True) + + def setup_loader(ap, r, is_val=False, verbose=False): + if is_val and not c.run_eval: + loader = None + else: + dataset = MyDataset( + r, + c.text_cleaner, + compute_linear_spec=False, + meta_data=meta_data_eval if is_val else meta_data_train, + ap=ap, + tp=c.characters if 'characters' in c.keys() else None, + add_blank=c['add_blank'] if 'add_blank' in c.keys() else False, + batch_group_size=0 if is_val else c.batch_group_size * + c.batch_size, + min_seq_len=c.min_seq_len, + max_seq_len=c.max_seq_len, + phoneme_cache_path=c.phoneme_cache_path, + use_phonemes=c.use_phonemes, + phoneme_language=c.phoneme_language, + enable_eos_bos=c.enable_eos_bos_chars, + use_noise_augment=not is_val, + verbose=verbose, + speaker_mapping=speaker_mapping if c.use_speaker_embedding and c.use_external_speaker_embedding_file else None) + + if c.use_phonemes and c.compute_input_seq_cache: + # precompute phonemes to have a better estimate of sequence lengths. + dataset.compute_input_seq(c.num_loader_workers) + dataset.sort_items() + + sampler = DistributedSampler(dataset) if num_gpus > 1 else None + loader = DataLoader( + dataset, + batch_size=c.eval_batch_size if is_val else c.batch_size, + shuffle=False, + collate_fn=dataset.collate_fn, + drop_last=False, + sampler=sampler, + num_workers=c.num_val_loader_workers + if is_val else c.num_loader_workers, + pin_memory=False) + return loader + + + def format_data(data): + # setup input data + text_input = data[0] + text_lengths = data[1] + speaker_names = data[2] + mel_input = data[4].permute(0, 2, 1) # B x D x T + mel_lengths = data[5] + item_idx = data[7] + attn_mask = data[9] + avg_text_length = torch.mean(text_lengths.float()) + avg_spec_length = torch.mean(mel_lengths.float()) + + if c.use_speaker_embedding: + if c.use_external_speaker_embedding_file: + # return precomputed embedding vector + speaker_c = data[8] + else: + # return speaker_id to be used by an embedding layer + speaker_c = [ + speaker_mapping[speaker_name] for speaker_name in speaker_names + ] + speaker_c = torch.LongTensor(speaker_c) + else: + speaker_c = None + # dispatch data to GPU + if use_cuda: + text_input = text_input.cuda(non_blocking=True) + text_lengths = text_lengths.cuda(non_blocking=True) + mel_input = mel_input.cuda(non_blocking=True) + mel_lengths = mel_lengths.cuda(non_blocking=True) + if speaker_c is not None: + speaker_c = speaker_c.cuda(non_blocking=True) + return text_input, text_lengths, mel_input, mel_lengths, speaker_c,\ + avg_text_length, avg_spec_length, item_idx + + + def train(data_loader, model, criterion, optimizer, scheduler, + ap, global_step, epoch): + + model.train() + epoch_time = 0 + keep_avg = KeepAverage() + if use_cuda: + batch_n_iter = int( + len(data_loader.dataset) / (c.batch_size * num_gpus)) + else: + batch_n_iter = int(len(data_loader.dataset) / c.batch_size) + end_time = time.time() + c_logger.print_train_start() + scaler = torch.cuda.amp.GradScaler() if c.mixed_precision else None + for num_iter, data in enumerate(data_loader): + start_time = time.time() + + # format data + text_input, text_lengths, mel_targets, mel_lengths, speaker_c,\ + avg_text_length, avg_spec_length, _ = format_data(data) + + loader_time = time.time() - end_time + + global_step += 1 + optimizer.zero_grad() + + # forward pass model + with torch.cuda.amp.autocast(enabled=c.mixed_precision): + decoder_output, dur_output, dur_mas_output, alignments, mu, log_sigma, logp_max_path = model.forward( + text_input, text_lengths, mel_targets, mel_lengths, g=speaker_c) + + # compute loss + loss_dict = criterion(mu, log_sigma, logp_max_path, decoder_output, mel_targets, mel_lengths, dur_output, dur_mas_output, text_lengths, global_step) + + # backward pass with loss scaling + if c.mixed_precision: + scaler.scale(loss_dict['loss']).backward() + scaler.unscale_(optimizer) + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), + c.grad_clip) + scaler.step(optimizer) + scaler.update() + else: + loss_dict['loss'].backward() + grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), + c.grad_clip) + optimizer.step() + + # setup lr + if c.noam_schedule: + scheduler.step() + + # current_lr + current_lr = optimizer.param_groups[0]['lr'] + + # compute alignment error (the lower the better ) + align_error = 1 - alignment_diagonal_score(alignments, binary=True) + loss_dict['align_error'] = align_error + + step_time = time.time() - start_time + epoch_time += step_time + + # aggregate losses from processes + if num_gpus > 1: + loss_dict['loss_l1'] = reduce_tensor(loss_dict['loss_l1'].data, num_gpus) + loss_dict['loss_ssim'] = reduce_tensor(loss_dict['loss_ssim'].data, num_gpus) + loss_dict['loss_dur'] = reduce_tensor(loss_dict['loss_dur'].data, num_gpus) + loss_dict['loss'] = reduce_tensor(loss_dict['loss'] .data, num_gpus) + + # detach loss values + loss_dict_new = dict() + for key, value in loss_dict.items(): + if isinstance(value, (int, float)): + loss_dict_new[key] = value + else: + loss_dict_new[key] = value.item() + loss_dict = loss_dict_new + + # update avg stats + update_train_values = dict() + for key, value in loss_dict.items(): + update_train_values['avg_' + key] = value + update_train_values['avg_loader_time'] = loader_time + update_train_values['avg_step_time'] = step_time + keep_avg.update_values(update_train_values) + + # print training progress + if global_step % c.print_step == 0: + log_dict = { + + "avg_spec_length": [avg_spec_length, 1], # value, precision + "avg_text_length": [avg_text_length, 1], + "step_time": [step_time, 4], + "loader_time": [loader_time, 2], + "current_lr": current_lr, + } + c_logger.print_train_step(batch_n_iter, num_iter, global_step, + log_dict, loss_dict, keep_avg.avg_values) + + if args.rank == 0: + # Plot Training Iter Stats + # reduce TB load + if global_step % c.tb_plot_step == 0: + iter_stats = { + "lr": current_lr, + "grad_norm": grad_norm, + "step_time": step_time + } + iter_stats.update(loss_dict) + tb_logger.tb_train_iter_stats(global_step, iter_stats) + + if global_step % c.save_step == 0: + if c.checkpoint: + # save model + save_checkpoint(model, optimizer, global_step, epoch, 1, OUT_PATH, model_characters, + model_loss=loss_dict['loss']) + + # wait all kernels to be completed + torch.cuda.synchronize() + + # Diagnostic visualizations + idx = np.random.randint(mel_targets.shape[0]) + pred_spec = decoder_output[idx].detach().data.cpu().numpy().T + gt_spec = mel_targets[idx].data.cpu().numpy().T + align_img = alignments[idx].data.cpu() + + figures = { + "prediction": plot_spectrogram(pred_spec, ap), + "ground_truth": plot_spectrogram(gt_spec, ap), + "alignment": plot_alignment(align_img), + } + + tb_logger.tb_train_figures(global_step, figures) + + # Sample audio + train_audio = ap.inv_melspectrogram(pred_spec.T) + tb_logger.tb_train_audios(global_step, + {'TrainAudio': train_audio}, + c.audio["sample_rate"]) + end_time = time.time() + + # print epoch stats + c_logger.print_train_epoch_end(global_step, epoch, epoch_time, keep_avg) + + # Plot Epoch Stats + if args.rank == 0: + epoch_stats = {"epoch_time": epoch_time} + epoch_stats.update(keep_avg.avg_values) + tb_logger.tb_train_epoch_stats(global_step, epoch_stats) + if c.tb_model_param_stats: + tb_logger.tb_model_weights(model, global_step) + return keep_avg.avg_values, global_step + + + @torch.no_grad() + def evaluate(data_loader, model, criterion, ap, global_step, epoch): + model.eval() + epoch_time = 0 + keep_avg = KeepAverage() + c_logger.print_eval_start() + if data_loader is not None: + for num_iter, data in enumerate(data_loader): + start_time = time.time() + + # format data + text_input, text_lengths, mel_targets, mel_lengths, speaker_c,\ + avg_text_length, avg_spec_length, _ = format_data(data) + + # forward pass model + with torch.cuda.amp.autocast(enabled=c.mixed_precision): + decoder_output, dur_output, dur_mas_output, alignments, mu, log_sigma, logp_max_path = model.forward( + text_input, text_lengths, mel_targets, mel_lengths, g=speaker_c) + + # compute loss + loss_dict = criterion(mu, log_sigma, logp_max_path, decoder_output, mel_targets, mel_lengths, dur_output, dur_mas_output, text_lengths, global_step) + + # step time + step_time = time.time() - start_time + epoch_time += step_time + + # compute alignment score + align_error = 1 - alignment_diagonal_score(alignments, binary=True) + loss_dict['align_error'] = align_error + + # aggregate losses from processes + if num_gpus > 1: + loss_dict['loss_l1'] = reduce_tensor(loss_dict['loss_l1'].data, num_gpus) + loss_dict['loss_ssim'] = reduce_tensor(loss_dict['loss_ssim'].data, num_gpus) + loss_dict['loss_dur'] = reduce_tensor(loss_dict['loss_dur'].data, num_gpus) + loss_dict['loss'] = reduce_tensor(loss_dict['loss'] .data, num_gpus) + + # detach loss values + loss_dict_new = dict() + for key, value in loss_dict.items(): + if isinstance(value, (int, float)): + loss_dict_new[key] = value + else: + loss_dict_new[key] = value.item() + loss_dict = loss_dict_new + + # update avg stats + update_train_values = dict() + for key, value in loss_dict.items(): + update_train_values['avg_' + key] = value + keep_avg.update_values(update_train_values) + + if c.print_eval: + c_logger.print_eval_step(num_iter, loss_dict, keep_avg.avg_values) + + if args.rank == 0: + # Diagnostic visualizations + idx = np.random.randint(mel_targets.shape[0]) + pred_spec = decoder_output[idx].detach().data.cpu().numpy().T + gt_spec = mel_targets[idx].data.cpu().numpy().T + align_img = alignments[idx].data.cpu() + + eval_figures = { + "prediction": plot_spectrogram(pred_spec, ap, output_fig=False), + "ground_truth": plot_spectrogram(gt_spec, ap, output_fig=False), + "alignment": plot_alignment(align_img, output_fig=False) + } + + # Sample audio + eval_audio = ap.inv_melspectrogram(pred_spec.T) + tb_logger.tb_eval_audios(global_step, {"ValAudio": eval_audio}, + c.audio["sample_rate"]) + + # Plot Validation Stats + tb_logger.tb_eval_stats(global_step, keep_avg.avg_values) + tb_logger.tb_eval_figures(global_step, eval_figures) + + if args.rank == 0 and epoch >= c.test_delay_epochs: + if c.test_sentences_file is None: + test_sentences = [ + "It took me quite a long time to develop a voice, and now that I have it I'm not going to be silent.", + "Be a voice, not an echo.", + "I'm sorry Dave. I'm afraid I can't do that.", + "This cake is great. It's so delicious and moist.", + "Prior to November 22, 1963." + ] + else: + with open(c.test_sentences_file, "r") as f: + test_sentences = [s.strip() for s in f.readlines()] + + # test sentences + test_audios = {} + test_figures = {} + print(" | > Synthesizing test sentences") + if c.use_speaker_embedding: + if c.use_external_speaker_embedding_file: + speaker_embedding = speaker_mapping[list(speaker_mapping.keys())[randrange(len(speaker_mapping)-1)]]['embedding'] + speaker_id = None + else: + speaker_id = 0 + speaker_embedding = None + else: + speaker_id = None + speaker_embedding = None + + style_wav = c.get("style_wav_for_test") + for idx, test_sentence in enumerate(test_sentences): + try: + wav, alignment, _, postnet_output, _, _ = synthesis( + model, + test_sentence, + c, + use_cuda, + ap, + speaker_id=speaker_id, + speaker_embedding=speaker_embedding, + style_wav=style_wav, + truncated=False, + enable_eos_bos_chars=c.enable_eos_bos_chars, #pylint: disable=unused-argument + use_griffin_lim=True, + do_trim_silence=False) + + file_path = os.path.join(AUDIO_PATH, str(global_step)) + os.makedirs(file_path, exist_ok=True) + file_path = os.path.join(file_path, + "TestSentence_{}.wav".format(idx)) + ap.save_wav(wav, file_path) + test_audios['{}-audio'.format(idx)] = wav + test_figures['{}-prediction'.format(idx)] = plot_spectrogram( + postnet_output, ap) + test_figures['{}-alignment'.format(idx)] = plot_alignment( + alignment) + except: #pylint: disable=bare-except + print(" !! Error creating Test Sentence -", idx) + traceback.print_exc() + tb_logger.tb_test_audios(global_step, test_audios, + c.audio['sample_rate']) + tb_logger.tb_test_figures(global_step, test_figures) + return keep_avg.avg_values + + + # FIXME: move args definition/parsing inside of main? + def main(args): # pylint: disable=redefined-outer-name + # pylint: disable=global-variable-undefined + global meta_data_train, meta_data_eval, symbols, phonemes, model_characters, speaker_mapping + # Audio processor + ap = AudioProcessor(**c.audio) + if 'characters' in c.keys(): + symbols, phonemes = make_symbols(**c.characters) + + # DISTRUBUTED + if num_gpus > 1: + init_distributed(args.rank, num_gpus, args.group_id, + c.distributed["backend"], c.distributed["url"]) + + # set model characters + model_characters = phonemes if c.use_phonemes else symbols + num_chars = len(model_characters) + + # load data instances + meta_data_train, meta_data_eval = load_meta_data(c.datasets, eval_split=True) + + # set the portion of the data used for training if set in config.json + if 'train_portion' in c.keys(): + meta_data_train = meta_data_train[:int(len(meta_data_train) * c.train_portion)] + if 'eval_portion' in c.keys(): + meta_data_eval = meta_data_eval[:int(len(meta_data_eval) * c.eval_portion)] + + # parse speakers + num_speakers, speaker_embedding_dim, speaker_mapping = parse_speakers(c, args, meta_data_train, OUT_PATH) + + # setup model + model = setup_model(num_chars, num_speakers, c, speaker_embedding_dim=speaker_embedding_dim) + optimizer = RAdam(model.parameters(), lr=c.lr, weight_decay=0, betas=(0.9, 0.98), eps=1e-9) + criterion = AlignTTSLoss(c) + + if args.restore_path: + print(f" > Restoring from {os.path.basename(args.restore_path)} ...") + checkpoint = torch.load(args.restore_path, map_location='cpu') + try: + # TODO: fix optimizer init, model.cuda() needs to be called before + # optimizer restore + optimizer.load_state_dict(checkpoint['optimizer']) + if c.reinit_layers: + raise RuntimeError + model.load_state_dict(checkpoint['model']) + except: #pylint: disable=bare-except + print(" > Partial model initialization.") + model_dict = model.state_dict() + model_dict = set_init_dict(model_dict, checkpoint['model'], c) + model.load_state_dict(model_dict) + del model_dict + + for group in optimizer.param_groups: + group['initial_lr'] = c.lr + print(" > Model restored from step %d" % checkpoint['step'], + flush=True) + args.restore_step = checkpoint['step'] + else: + args.restore_step = 0 + + if use_cuda: + model.cuda() + criterion.cuda() + + # DISTRUBUTED + if num_gpus > 1: + model = DDP_th(model, device_ids=[args.rank]) + + if c.noam_schedule: + scheduler = NoamLR(optimizer, + warmup_steps=c.warmup_steps, + last_epoch=args.restore_step - 1) + else: + scheduler = None + + num_params = count_parameters(model) + print("\n > Model has {} parameters".format(num_params), flush=True) + + if args.restore_step == 0 or not args.best_path: + best_loss = float('inf') + print(" > Starting with inf best loss.") + else: + print(" > Restoring best loss from " + f"{os.path.basename(args.best_path)} ...") + best_loss = torch.load(args.best_path, + map_location='cpu')['model_loss'] + print(f" > Starting with loaded last best loss {best_loss}.") + keep_all_best = c.get('keep_all_best', False) + keep_after = c.get('keep_after', 10000) # void if keep_all_best False + + # define dataloaders + train_loader = setup_loader(ap, 1, is_val=False, verbose=True) + eval_loader = setup_loader(ap, 1, is_val=True, verbose=True) + + global_step = args.restore_step + for epoch in range(0, c.epochs): + c_logger.print_epoch_start(epoch, c.epochs) + train_avg_loss_dict, global_step = train(train_loader, model, criterion, optimizer, + scheduler, ap, global_step, + epoch) + eval_avg_loss_dict = evaluate(eval_loader, model, criterion, ap, + global_step, epoch) + c_logger.print_epoch_end(epoch, eval_avg_loss_dict) + target_loss = train_avg_loss_dict['avg_loss'] + if c.run_eval: + target_loss = eval_avg_loss_dict['avg_loss'] + best_loss = save_best_model(target_loss, best_loss, model, optimizer, + global_step, epoch, c.r, OUT_PATH, model_characters, + keep_all_best=keep_all_best, keep_after=keep_after) + + + args = parse_arguments(sys.argv) + c, OUT_PATH, AUDIO_PATH, c_logger, tb_logger = process_args( + args, model_type='tts') + + try: + main(args) + except KeyboardInterrupt: + remove_experiment_folder(OUT_PATH) + try: + sys.exit(0) + except SystemExit: + os._exit(0) # pylint: disable=protected-access + except Exception: # pylint: disable=broad-except + remove_experiment_folder(OUT_PATH) + traceback.print_exc() + sys.exit(1) \ No newline at end of file diff --git a/TTS/tts/layers/glow_tts/monotonic_align/__init__.py b/TTS/tts/layers/glow_tts/monotonic_align/__init__.py index a2912a98..78fa0fbf 100644 --- a/TTS/tts/layers/glow_tts/monotonic_align/__init__.py +++ b/TTS/tts/layers/glow_tts/monotonic_align/__init__.py @@ -23,7 +23,6 @@ def generate_path(duration, mask): mask: [b, t_x, t_y] """ device = duration.device - b, t_x, t_y = mask.shape cum_duration = torch.cumsum(duration, 1) path = torch.zeros(b, t_x, t_y, dtype=mask.dtype).to(device=device) diff --git a/TTS/tts/models/align_tts.py b/TTS/tts/models/align_tts.py new file mode 100644 index 00000000..c2ba8bf2 --- /dev/null +++ b/TTS/tts/models/align_tts.py @@ -0,0 +1,257 @@ +import torch +import math +from torch import nn +from TTS.tts.layers.feed_forward.decoder import Decoder +from TTS.tts.layers.feed_forward.duration_predictor import DurationPredictor +from TTS.tts.layers.feed_forward.encoder import Encoder, PositionalEncoding +from TTS.tts.utils.generic_utils import sequence_mask +from TTS.tts.layers.glow_tts.monotonic_align import maximum_path, generate_path +from TTS.tts.layers.align_tts.mdn import MDNBlock + + + + +class AlignTTS(nn.Module): + """Speedy Speech model with Monotonic Alignment Search + https://arxiv.org/abs/2008.03802 + https://arxiv.org/pdf/2005.11129.pdf + + Encoder -> DurationPredictor -> Decoder + + This model is able to achieve a reasonable performance with only + ~3M model parameters and convolutional layers. + + This model requires precomputed phoneme durations to train a duration predictor. At inference + it only uses the duration predictor to compute durations and expand encoder outputs respectively. + + Args: + num_chars (int): number of unique input to characters + out_channels (int): number of output tensor channels. It is equal to the expected spectrogram size. + hidden_channels (int): number of channels in all the model layers. + positional_encoding (bool, optional): enable/disable Positional encoding on encoder outputs. Defaults to True. + length_scale (int, optional): coefficient to set the speech speed. <1 slower, >1 faster. Defaults to 1. + encoder_type (str, optional): set the encoder type. Defaults to 'residual_conv_bn'. + encoder_params (dict, optional): set encoder parameters depending on 'encoder_type'. Defaults to { "kernel_size": 4, "dilations": 4 * [1, 2, 4] + [1], "num_conv_blocks": 2, "num_res_blocks": 13 }. + decoder_type (str, optional): decoder type. Defaults to 'residual_conv_bn'. + decoder_params (dict, optional): set decoder parameters depending on 'decoder_type'. Defaults to { "kernel_size": 4, "dilations": 4 * [1, 2, 4, 8] + [1], "num_conv_blocks": 2, "num_res_blocks": 17 }. + num_speakers (int, optional): number of speakers for multi-speaker training. Defaults to 0. + external_c (bool, optional): enable external speaker embeddings. Defaults to False. + c_in_channels (int, optional): number of channels in speaker embedding vectors. Defaults to 0. + """ + # pylint: disable=dangerous-default-value + + def __init__( + self, + num_chars, + out_channels, + hidden_channels, + positional_encoding=True, + length_scale=1, + encoder_type='residual_conv_bn', + encoder_params={ + "kernel_size": 4, + "dilations": 4 * [1, 2, 4] + [1], + "num_conv_blocks": 2, + "num_res_blocks": 13 + }, + decoder_type='residual_conv_bn', + decoder_params={ + "kernel_size": 4, + "dilations": 4 * [1, 2, 4, 8] + [1], + "num_conv_blocks": 2, + "num_res_blocks": 17 + }, + num_speakers=0, + external_c=False, + c_in_channels=0): + + super().__init__() + self.length_scale = float(length_scale) if isinstance( + length_scale, int) else length_scale + self.emb = nn.Embedding(num_chars, hidden_channels) + self.encoder = Encoder(hidden_channels, hidden_channels, encoder_type, + encoder_params, c_in_channels) + if positional_encoding: + self.pos_encoder = PositionalEncoding(hidden_channels) + self.decoder = Decoder(out_channels, hidden_channels, decoder_type, + decoder_params) + self.duration_predictor = DurationPredictor(hidden_channels + + c_in_channels) + + self.mod_layer = nn.Conv1d(hidden_channels, hidden_channels, 1) + # self.wn_spec_encoder = WNSpecEncoder(out_channels, hidden_channels, c_in_channels=c_in_channels) + self.mdn_block = MDNBlock(hidden_channels, 2*out_channels) + + if num_speakers > 1 and not external_c: + # speaker embedding layer + self.emb_g = nn.Embedding(num_speakers, c_in_channels) + nn.init.uniform_(self.emb_g.weight, -0.1, 0.1) + + if c_in_channels > 0 and c_in_channels != hidden_channels: + self.proj_g = nn.Conv1d(c_in_channels, hidden_channels, 1) + + def compute_mas_path(self, mu, log_sigma, y, x_mask, y_mask): + # find the max alignment path + attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2) + with torch.no_grad(): + scale = torch.exp(-2 * log_sigma) + # [B, T_en, 1] + logp1 = torch.sum(-0.5 * math.log(2 * math.pi) - log_sigma, + [1]).unsqueeze(-1) + # [B, T_en, D] x [B, D, T_dec] = [B, T_en, T_dec] + logp2 = torch.matmul(scale.transpose(1, 2), -0.5 * (y**2)) + # [B, T_en, D] x [B, D, T_dec] = [B, T_en, T_dec] + logp3 = torch.matmul((mu * scale).transpose(1, 2), y) + # [B, T_en, 1] + logp4 = torch.sum(-0.5 * (mu**2) * scale, + [1]).unsqueeze(-1) + # [B, T_en, T_dec] + logp = logp1 + logp2 + logp3 + logp4 + # import pdb; pdb.set_trace() + # [B, T_en, T_dec] + attn = maximum_path(logp, + attn_mask.squeeze(1)).unsqueeze(1).detach() + # logp_max_path = logp.new_ones(logp.shape) * -1e4 + # logp_max_path += logp * attn.squeeze(1) + logp_max_path = None + dr_mas = torch.sum(attn, -1) + return dr_mas.squeeze(1), logp_max_path + + @staticmethod + def expand_encoder_outputs(en, dr, x_mask, y_mask): + """Generate attention alignment map from durations and + expand encoder outputs + + Example: + encoder output: [a,b,c,d] + durations: [1, 3, 2, 1] + + expanded: [a, b, b, b, c, c, d] + attention map: [[0, 0, 0, 0, 0, 0, 1], + [0, 0, 0, 0, 1, 1, 0], + [0, 1, 1, 1, 0, 0, 0], + [1, 0, 0, 0, 0, 0, 0]] + """ + attn_mask = torch.unsqueeze(x_mask, -1) * torch.unsqueeze(y_mask, 2) + attn = generate_path(dr, attn_mask.squeeze(1)).to(en.dtype) + o_en_ex = torch.matmul( + attn.squeeze(1).transpose(1, 2), en.transpose(1, + 2)).transpose(1, 2) + return o_en_ex, attn + + def format_durations(self, o_dr_log, x_mask): + o_dr = (torch.exp(o_dr_log) - 1) * x_mask * self.length_scale + o_dr[o_dr < 1] = 1.0 + o_dr = torch.round(o_dr) + return o_dr + + @staticmethod + def _concat_speaker_embedding(o_en, g): + g_exp = g.expand(-1, -1, o_en.size(-1)) # [B, C, T_en] + o_en = torch.cat([o_en, g_exp], 1) + return o_en + + def _sum_speaker_embedding(self, x, g): + # project g to decoder dim. + if hasattr(self, 'proj_g'): + g = self.proj_g(g) + return x + g + + def _forward_encoder(self, x, x_lengths, g=None): + if hasattr(self, 'emb_g'): + g = nn.functional.normalize(self.emb_g(g)) # [B, C, 1] + + if g is not None: + g = g.unsqueeze(-1) + + # [B, T, C] + x_emb = self.emb(x) + # [B, C, T] + x_emb = torch.transpose(x_emb, 1, -1) + + # compute sequence masks + x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.shape[1]), + 1).to(x.dtype) + + # encoder pass + o_en = self.encoder(x_emb, x_mask) + + # speaker conditioning for duration predictor + if g is not None: + o_en_dp = self._concat_speaker_embedding(o_en, g) + else: + o_en_dp = o_en + return o_en, o_en_dp, x_mask, g + + def _forward_decoder(self, o_en, o_en_dp, dr, x_mask, y_lengths, g): + y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), + 1).to(o_en_dp.dtype) + # expand o_en with durations + o_en_ex, attn = self.expand_encoder_outputs(o_en, dr, x_mask, y_mask) + # positional encoding + if hasattr(self, 'pos_encoder'): + o_en_ex = self.pos_encoder(o_en_ex, y_mask) + # speaker embedding + if g is not None: + o_en_ex = self._sum_speaker_embedding(o_en_ex, g) + # decoder pass + o_de = self.decoder(o_en_ex, y_mask, g=g) + + return o_de, attn.transpose(1, 2) + + # def _forward_mas(self, o_en, y, y_lengths, x_mask): + # # MAS potentials and alignment + # o_en_mean = self.mod_layer(o_en) + # y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), + # 1).to(o_en.dtype) + # z = self.wn_spec_encoder(y) + # dr_mas, y_mean, y_scale = self.compute_mas_path(o_en_mean, z, x_mask, y_mask) + # return dr_mas, z, y_mean, y_scale + + def _forward_mdn(self, o_en, y, y_lengths, x_mask): + # MAS potentials and alignment + mu, log_sigma = self.mdn_block(o_en) + y_mask = torch.unsqueeze(sequence_mask(y_lengths, None), 1).to(o_en.dtype) + dr_mas, logp_max_path = self.compute_mas_path(mu, log_sigma, y, x_mask, y_mask) + return dr_mas, mu, log_sigma, logp_max_path + + def forward(self, x, x_lengths, y, y_lengths, g=None): # pylint: disable=unused-argument + """ + Shapes: + x: [B, T_max] + x_lengths: [B] + y_lengths: [B] + dr: [B, T_max] + g: [B, C] + """ + o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g) + dr_mas, mu, log_sigma, logp_max_path = self._forward_mdn(o_en, y, y_lengths, x_mask) + o_dr_log = self.duration_predictor(o_en_dp.detach(), x_mask) + # TODO: compute attn once + o_de, attn = self._forward_decoder(o_en, o_en_dp, dr_mas, x_mask, y_lengths, g=g) + dr_mas_log = torch.log(1 + dr_mas).squeeze(1) + return o_de, o_dr_log.squeeze(1), dr_mas_log, attn, mu, log_sigma, logp_max_path + + def inference(self, x, x_lengths, g=None): # pylint: disable=unused-argument + """ + Shapes: + x: [B, T_max] + x_lengths: [B] + g: [B, C] + """ + # pad input to prevent dropping the last word + x = torch.nn.functional.pad(x, pad=(0, 5), mode='constant', value=0) + o_en, o_en_dp, x_mask, g = self._forward_encoder(x, x_lengths, g) + # duration predictor pass + o_dr_log = self.duration_predictor(o_en_dp.detach(), x_mask) + o_dr = self.format_durations(o_dr_log, x_mask).squeeze(1) + y_lengths = o_dr.sum(1) + o_de, attn = self._forward_decoder(o_en, o_en_dp, o_dr, x_mask, y_lengths, g=g) + return o_de, attn + + def load_checkpoint(self, config, checkpoint_path, eval=False): # pylint: disable=unused-argument, redefined-builtin + state = torch.load(checkpoint_path, map_location=torch.device('cpu')) + self.load_state_dict(state['model']) + if eval: + self.eval() + assert not self.training \ No newline at end of file diff --git a/TTS/tts/models/speedy_speech.py b/TTS/tts/models/speedy_speech.py index 101d77a0..afb0245a 100644 --- a/TTS/tts/models/speedy_speech.py +++ b/TTS/tts/models/speedy_speech.py @@ -1,8 +1,8 @@ import torch from torch import nn -from TTS.tts.layers.speedy_speech.decoder import Decoder -from TTS.tts.layers.speedy_speech.duration_predictor import DurationPredictor -from TTS.tts.layers.speedy_speech.encoder import Encoder, PositionalEncoding +from TTS.tts.layers.feed_forward.decoder import Decoder +from TTS.tts.layers.feed_forward.duration_predictor import DurationPredictor +from TTS.tts.layers.feed_forward.encoder import Encoder, PositionalEncoding from TTS.tts.utils.generic_utils import sequence_mask from TTS.tts.layers.glow_tts.monotonic_align import generate_path diff --git a/TTS/tts/models/tacotron.py b/TTS/tts/models/tacotron.py index 0b68a96c..541c4159 100644 --- a/TTS/tts/models/tacotron.py +++ b/TTS/tts/models/tacotron.py @@ -2,8 +2,8 @@ import torch from torch import nn -from TTS.tts.layers.gst_layers import GST -from TTS.tts.layers.tacotron import Decoder, Encoder, PostCBHG +from TTS.tts.layers.tacotron.gst_layers import GST +from TTS.tts.layers.tacotron.tacotron import Decoder, Encoder, PostCBHG from TTS.tts.models.tacotron_abstract import TacotronAbstract diff --git a/TTS/tts/models/tacotron2.py b/TTS/tts/models/tacotron2.py index e56e4ca0..0e751c32 100644 --- a/TTS/tts/models/tacotron2.py +++ b/TTS/tts/models/tacotron2.py @@ -1,8 +1,8 @@ import torch from torch import nn -from TTS.tts.layers.gst_layers import GST -from TTS.tts.layers.tacotron2 import Decoder, Encoder, Postnet +from TTS.tts.layers.tacotron.gst_layers import GST +from TTS.tts.layers.tacotron.tacotron2 import Decoder, Encoder, Postnet from TTS.tts.models.tacotron_abstract import TacotronAbstract # TODO: match function arguments with tacotron @@ -17,7 +17,7 @@ class Tacotron2(TacotronAbstract): r (int): initial model reduction rate. postnet_output_dim (int, optional): postnet output channels. Defaults to 80. decoder_output_dim (int, optional): decoder output channels. Defaults to 80. - attn_type (str, optional): attention type. Check ```TTS.tts.layers.common_layers.init_attn```. Defaults to 'original'. + attn_type (str, optional): attention type. Check ```TTS.tts.layers.tacotron.common_layers.init_attn```. Defaults to 'original'. attn_win (bool, optional): enable/disable attention windowing. It especially useful at inference to keep attention alignment diagonal. Defaults to False. attn_norm (str, optional): Attention normalization method. "sigmoid" or "softmax". Defaults to "softmax". diff --git a/TTS/tts/utils/synthesis.py b/TTS/tts/utils/synthesis.py index a0524e8f..f825d61c 100644 --- a/TTS/tts/utils/synthesis.py +++ b/TTS/tts/utils/synthesis.py @@ -77,7 +77,7 @@ def run_model_torch(model, inputs, CONFIG, truncated, speaker_id=None, style_mel # these only belong to tacotron models. decoder_output = None stop_tokens = None - elif 'speedy_speech' in CONFIG.model.lower(): + elif CONFIG.model.lower() in ['speedy_speech', 'align_tts']: inputs_lengths = torch.tensor(inputs.shape[1:2]).to(inputs.device) # pylint: disable=not-callable if hasattr(model, 'module'): # distributed model @@ -88,6 +88,8 @@ def run_model_torch(model, inputs, CONFIG, truncated, speaker_id=None, style_mel # these only belong to tacotron models. decoder_output = None stop_tokens = None + else: + raise ValueError('[!] Unknown model name.') return decoder_output, postnet_output, alignments, stop_tokens diff --git a/tests/test_layers.py b/tests/test_layers.py index 1a07b750..582ca8be 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -1,7 +1,7 @@ import unittest import torch as T -from TTS.tts.layers.tacotron import Prenet, CBHG, Decoder, Encoder +from TTS.tts.layers.tacotron.tacotron import Prenet, CBHG, Decoder, Encoder from TTS.tts.layers.losses import L1LossMasked, SSIMLoss from TTS.tts.utils.generic_utils import sequence_mask diff --git a/tests/test_speedy_speech_layers.py b/tests/test_speedy_speech_layers.py index 53351fff..b93d4766 100644 --- a/tests/test_speedy_speech_layers.py +++ b/tests/test_speedy_speech_layers.py @@ -1,8 +1,8 @@ import torch -from TTS.tts.layers.speedy_speech.encoder import Encoder -from TTS.tts.layers.speedy_speech.decoder import Decoder -from TTS.tts.layers.speedy_speech.duration_predictor import DurationPredictor +from TTS.tts.layers.feed_forward.encoder import Encoder +from TTS.tts.layers.feed_forward.decoder import Decoder +from TTS.tts.layers.feed_forward.duration_predictor import DurationPredictor from TTS.tts.utils.generic_utils import sequence_mask from TTS.tts.models.speedy_speech import SpeedySpeech