correct imports after refactoring, add AlignTTS (old SSMAS) and some formatting

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)

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)

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

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.feed_forward.decoder import Decoder
-from TTS.tts.layers.speedy_speech.duration_predictor import DurationPredictor
+from TTS.tts.layers.feed_forward.duration_predictor import DurationPredictor
-from TTS.tts.layers.speedy_speech.encoder import Encoder, PositionalEncoding
+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
 

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.gst_layers import GST
-from TTS.tts.layers.tacotron import Decoder, Encoder, PostCBHG
+from TTS.tts.layers.tacotron.tacotron import Decoder, Encoder, PostCBHG
 from TTS.tts.models.tacotron_abstract import TacotronAbstract
 
 

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.tacotron.gst_layers import GST
-from TTS.tts.layers.tacotron2 import Decoder, Encoder, Postnet
+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".

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
 
 

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
 

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.feed_forward.encoder import Encoder
-from TTS.tts.layers.speedy_speech.decoder import Decoder
+from TTS.tts.layers.feed_forward.decoder import Decoder
-from TTS.tts.layers.speedy_speech.duration_predictor import DurationPredictor
+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