<gh_stars>1-10
# Adapted from
# https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechSynthesis/FastPitch

# *****************************************************************************
#  Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# *****************************************************************************

from typing import Tuple
import torch
from torch import nn
from torch.nn.utils.rnn import pad_sequence

from common.layers import ConvReLUNorm
from common.utils import mask_from_lens
from fastpitch.transformer import FFTransformer


def regulate_len(durations, enc_out, pace: float = 1.0,
                 mel_max_len: int = None) -> Tuple[torch.Tensor, torch.Tensor]:
    """Regulate the length of durations.

    Args:
        durations ([type]): Durations, ground truth or predicted.
        enc_out ([type]): [description]
        pace (float, optional): Target pace. Defaults to 1.0.
        mel_max_len (int, optional): Max length of mel. Defaults to None.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: enc_rep, dec_lens
    """

    reps = torch.round(durations.float() / pace).long()
    dec_lens = reps.sum(dim=1)

    enc_rep = pad_sequence([torch.repeat_interleave(o, r, dim=0)
                            for o, r in zip(enc_out, reps)],
                           batch_first=True)
    if mel_max_len:
        enc_rep = enc_rep[:, :mel_max_len]
        dec_lens = torch.clamp_max(dec_lens, mel_max_len)
    return enc_rep, dec_lens


class TemporalPredictor(nn.Module):
    """Predicts a single float per each temporal location.
    Used as duration predictor and pitch predictor.
    """

    def __init__(self, input_size, filter_size, kernel_size, dropout,
                 n_layers=2):
        super(TemporalPredictor, self).__init__()

        self.layers = nn.Sequential(*[
            ConvReLUNorm(input_size if i == 0 else filter_size, filter_size,
                         kernel_size=kernel_size, dropout=dropout)
            for i in range(n_layers)]
        )
        self.fc = nn.Linear(filter_size, 1, bias=True)

    def forward(self, enc_out, enc_out_mask):
        out = enc_out * enc_out_mask
        out = self.layers(out.transpose(1, 2)).transpose(1, 2)
        out = self.fc(out) * enc_out_mask
        return out.squeeze(-1)


class FastPitch(nn.Module):
    """FastPitch module definition."""

    def __init__(self, n_mel_channels: int, max_seq_len: int, n_symbols: int,
                 symbols_embedding_dim: int, in_fft_n_layers: int, in_fft_n_heads: int,
                 in_fft_d_head: int,
                 in_fft_conv1d_kernel_size: int, in_fft_conv1d_filter_size: int,
                 in_fft_output_size: int,
                 p_in_fft_dropout: float, p_in_fft_dropatt, p_in_fft_dropemb,
                 out_fft_n_layers: int, out_fft_n_heads: int, out_fft_d_head: int,
                 out_fft_conv1d_kernel_size: int, out_fft_conv1d_filter_size: int,
                 out_fft_output_size: int, p_out_fft_dropout: float,
                 p_out_fft_dropatt: float, p_out_fft_dropemb: float,
                 dur_predictor_kernel_size: int, dur_predictor_filter_size: int,
                 p_dur_predictor_dropout: float, dur_predictor_n_layers: int,
                 pitch_predictor_kernel_size: int, pitch_predictor_filter_size: int,
                 p_pitch_predictor_dropout: float, pitch_predictor_n_layers: int):
        """FastPitch's constructor

        Args:
            n_mel_channels (int): Number of bins in mel-spectrograms
            max_seq_len (int): Unused.
            n_symbols (int): Unused.
            symbols_embedding_dim (int): Input embedding dimension

            in_fft_n_layers (int): Number of FFT blocks
            in_fft_n_heads (int): Number of attention heads
            in_fft_d_head (int): Dim of attention heads
            in_fft_conv1d_kernel_size (int): Conv-1D kernel size
            in_fft_conv1d_filter_size (int): Conv-1D filter size
            in_fft_output_size (int): Output dim

            p_in_fft_dropout (float): Dropout probability
            p_in_fft_dropatt (float): Multi-head attention dropout
            p_in_fft_dropemb (float): Dropout added to word+positional embeddings')


            out_fft_n_layers (int): Number of FFT blocks
            out_fft_n_heads (int): Number of attention heads
            out_fft_d_head (int): Dim of attention head
            out_fft_conv1d_kernel_size (int): Conv-1D kernel size
            out_fft_conv1d_filter_size (int): Conv-1D filter size
            out_fft_output_size (int): Output dim

            p_out_fft_dropout (float): Dropout probability for out_fft
            p_out_fft_dropatt (float): Multi-head attention dropout
            p_out_fft_dropemb (float): Dropout added to word+positional embeddings

            dur_predictor_kernel_size (int): Duration predictor conv-1D kernel size
            dur_predictor_filter_size (int): Duration predictor conv-1D filter size
            p_dur_predictor_dropout (float): Dropout probability for duration predictor
            dur_predictor_n_layers (int): Number of conv-1D layers

            pitch_predictor_kernel_size (int): Pitch predictor conv-1D kernel size
            pitch_predictor_filter_size (int): Pitch predictor conv-1D filter size
            p_pitch_predictor_dropout (float): Dropout probability for pitch predictor
            pitch_predictor_n_layers (int): Number of conv-1D layers
        """

        super(FastPitch, self).__init__()
        del max_seq_len  # unused
        del n_symbols  # unused

        self.encoder = FFTransformer(
            n_layer=in_fft_n_layers,
            n_head=in_fft_n_heads,
            d_model=symbols_embedding_dim,
            d_head=in_fft_d_head,
            d_inner=in_fft_conv1d_filter_size,
            kernel_size=in_fft_conv1d_kernel_size,
            dropout=p_in_fft_dropout,
            dropatt=p_in_fft_dropatt,
            dropemb=p_in_fft_dropemb,
            d_embed=symbols_embedding_dim,
            embed_input=True
        )

        self.duration_predictor = TemporalPredictor(
            in_fft_output_size,
            filter_size=dur_predictor_filter_size,
            kernel_size=dur_predictor_kernel_size,
            dropout=p_dur_predictor_dropout, n_layers=dur_predictor_n_layers
        )

        self.decoder = FFTransformer(
            n_layer=out_fft_n_layers, n_head=out_fft_n_heads,
            d_model=symbols_embedding_dim,
            d_head=out_fft_d_head,
            d_inner=out_fft_conv1d_filter_size,
            kernel_size=out_fft_conv1d_kernel_size,
            dropout=p_out_fft_dropout,
            dropatt=p_out_fft_dropatt,
            dropemb=p_out_fft_dropemb,
            d_embed=symbols_embedding_dim,
            embed_input=False
        )

        self.pitch_predictor = TemporalPredictor(
            in_fft_output_size,
            filter_size=pitch_predictor_filter_size,
            kernel_size=pitch_predictor_kernel_size,
            dropout=p_pitch_predictor_dropout, n_layers=pitch_predictor_n_layers
        )

        self.pitch_emb = nn.Conv1d(1, symbols_embedding_dim, kernel_size=3,
                                   padding=1)

        # Store values precomputed for training data within the model
        self.register_buffer('pitch_mean', torch.zeros(1))
        self.register_buffer('pitch_std', torch.zeros(1))

        self.proj = nn.Linear(out_fft_output_size, n_mel_channels, bias=True)

    def forward(self, inputs: Tuple, use_gt_durations: bool = True, 
                use_gt_pitch: bool = True,
                pace: float = 1.0, max_duration: int = 75):
        """A forward step.

        Args:
            inputs (tuple): inputs, _, mel_tgt, _, dur_tgt, _, pitch_tgt
            use_gt_durations (bool, optional): Use ground truth durations. Defaults to True.
            use_gt_pitch (bool, optional): Use ground truth pitch. Defaults to True.
            pace (float, optional): Pace. Defaults to 1.0.
            max_duration (int, optional): Max duration. Defaults to 75.

        Returns:
            tuple: mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred
        """

        inputs, _, mel_tgt, _, dur_tgt, _, pitch_tgt = inputs
        mel_max_len = mel_tgt.size(2)

        # Input FFT
        enc_out, enc_mask = self.encoder(inputs)

        # Embedded for predictors
        pred_enc_out, pred_enc_mask = enc_out, enc_mask

        # Predict durations
        log_dur_pred = self.duration_predictor(pred_enc_out, pred_enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, 0, max_duration)

        # Predict pitch
        pitch_pred = self.pitch_predictor(enc_out, enc_mask)

        if use_gt_pitch and pitch_tgt is not None:
            pitch_emb = self.pitch_emb(pitch_tgt.unsqueeze(1))
        else:
            pitch_emb = self.pitch_emb(pitch_pred.unsqueeze(1))
        enc_out = enc_out + pitch_emb.transpose(1, 2)

        len_regulated, dec_lens = regulate_len(
            dur_tgt if use_gt_durations else dur_pred,
            enc_out, pace, mel_max_len)

        # Output FFT
        dec_out, dec_mask = self.decoder(len_regulated, dec_lens)
        mel_out = self.proj(dec_out)
        return mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred

    def infer(self, inputs, input_lens, pace: float=1.0, dur_tgt=None, pitch_tgt=None,
              pitch_transform=None, max_duration: int = 75):
        """Inference.

        Args:
            inputs (tuple): [description]
            input_lens ([type]): [description]
            pace (float, optional): Target pace. Defaults to 1.0.
            dur_tgt ([type], optional): Target duration. Defaults to None.
            pitch_tgt ([type], optional): Target pitch. Defaults to None.
            pitch_transform ([type], optional): [description]. Defaults to None.
            max_duration (int, optional): [description]. Defaults to 75.

        Returns:
            tuple: mel_out, dec_lens, dur_pred, pitch_pred
        """

        del input_lens  # unused

        # Input FFT
        enc_out, enc_mask = self.encoder(inputs)

        # Embedded for predictors
        pred_enc_out, pred_enc_mask = enc_out, enc_mask

        # Predict durations
        log_dur_pred = self.duration_predictor(pred_enc_out, pred_enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, 0, max_duration)

        # Pitch over chars
        pitch_pred = self.pitch_predictor(enc_out, enc_mask)

        if pitch_transform is not None:
            if self.pitch_std[0] == 0.0:
                # XXX LJSpeech-1.1 defaults
                mean, std = 218.14, 67.24
            else:
                mean, std = self.pitch_mean[0], self.pitch_std[0]
            pitch_pred = pitch_transform(pitch_pred, mean, std)

        if pitch_tgt is None:
            pitch_emb = self.pitch_emb(pitch_pred.unsqueeze(1)).transpose(1, 2)
        else:
            pitch_emb = self.pitch_emb(pitch_tgt.unsqueeze(1)).transpose(1, 2)

        enc_out = enc_out + pitch_emb

        len_regulated, dec_lens = regulate_len(
            dur_pred if dur_tgt is None else dur_tgt,
            enc_out, pace, mel_max_len=None)

        dec_out, dec_mask = self.decoder(len_regulated, dec_lens)
        mel_out = self.proj(dec_out)
        # mel_lens = dec_mask.squeeze(2).sum(axis=1).long()
        mel_out = mel_out.permute(0, 2, 1)  # For inference.py
        return mel_out, dec_lens, dur_pred, pitch_pred

# Adapted from
# https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechSynthesis/FastPitch

# *****************************************************************************
#  Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#      * Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#      * Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# *****************************************************************************

from typing import Tuple
import torch
from torch import nn
from torch.nn.utils.rnn import pad_sequence

from common.layers import ConvReLUNorm
from common.utils import mask_from_lens
from fastpitch.transformer import FFTransformer


def regulate_len(durations, enc_out, pace: float = 1.0,
                 mel_max_len: int = None) -> Tuple[torch.Tensor, torch.Tensor]:
    """Regulate the length of durations.

    Args:
        durations ([type]): Durations, ground truth or predicted.
        enc_out ([type]): [description]
        pace (float, optional): Target pace. Defaults to 1.0.
        mel_max_len (int, optional): Max length of mel. Defaults to None.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: enc_rep, dec_lens
    """

    reps = torch.round(durations.float() / pace).long()
    dec_lens = reps.sum(dim=1)

    enc_rep = pad_sequence([torch.repeat_interleave(o, r, dim=0)
                            for o, r in zip(enc_out, reps)],
                           batch_first=True)
    if mel_max_len:
        enc_rep = enc_rep[:, :mel_max_len]
        dec_lens = torch.clamp_max(dec_lens, mel_max_len)
    return enc_rep, dec_lens


class TemporalPredictor(nn.Module):
    """Predicts a single float per each temporal location.
    Used as duration predictor and pitch predictor.
    """

    def __init__(self, input_size, filter_size, kernel_size, dropout,
                 n_layers=2):
        super(TemporalPredictor, self).__init__()

        self.layers = nn.Sequential(*[
            ConvReLUNorm(input_size if i == 0 else filter_size, filter_size,
                         kernel_size=kernel_size, dropout=dropout)
            for i in range(n_layers)]
        )
        self.fc = nn.Linear(filter_size, 1, bias=True)

    def forward(self, enc_out, enc_out_mask):
        out = enc_out * enc_out_mask
        out = self.layers(out.transpose(1, 2)).transpose(1, 2)
        out = self.fc(out) * enc_out_mask
        return out.squeeze(-1)


class FastPitch(nn.Module):
    """FastPitch module definition."""

    def __init__(self, n_mel_channels: int, max_seq_len: int, n_symbols: int,
                 symbols_embedding_dim: int, in_fft_n_layers: int, in_fft_n_heads: int,
                 in_fft_d_head: int,
                 in_fft_conv1d_kernel_size: int, in_fft_conv1d_filter_size: int,
                 in_fft_output_size: int,
                 p_in_fft_dropout: float, p_in_fft_dropatt, p_in_fft_dropemb,
                 out_fft_n_layers: int, out_fft_n_heads: int, out_fft_d_head: int,
                 out_fft_conv1d_kernel_size: int, out_fft_conv1d_filter_size: int,
                 out_fft_output_size: int, p_out_fft_dropout: float,
                 p_out_fft_dropatt: float, p_out_fft_dropemb: float,
                 dur_predictor_kernel_size: int, dur_predictor_filter_size: int,
                 p_dur_predictor_dropout: float, dur_predictor_n_layers: int,
                 pitch_predictor_kernel_size: int, pitch_predictor_filter_size: int,
                 p_pitch_predictor_dropout: float, pitch_predictor_n_layers: int):
        """FastPitch's constructor

        Args:
            n_mel_channels (int): Number of bins in mel-spectrograms
            max_seq_len (int): Unused.
            n_symbols (int): Unused.
            symbols_embedding_dim (int): Input embedding dimension

            in_fft_n_layers (int): Number of FFT blocks
            in_fft_n_heads (int): Number of attention heads
            in_fft_d_head (int): Dim of attention heads
            in_fft_conv1d_kernel_size (int): Conv-1D kernel size
            in_fft_conv1d_filter_size (int): Conv-1D filter size
            in_fft_output_size (int): Output dim

            p_in_fft_dropout (float): Dropout probability
            p_in_fft_dropatt (float): Multi-head attention dropout
            p_in_fft_dropemb (float): Dropout added to word+positional embeddings')


            out_fft_n_layers (int): Number of FFT blocks
            out_fft_n_heads (int): Number of attention heads
            out_fft_d_head (int): Dim of attention head
            out_fft_conv1d_kernel_size (int): Conv-1D kernel size
            out_fft_conv1d_filter_size (int): Conv-1D filter size
            out_fft_output_size (int): Output dim

            p_out_fft_dropout (float): Dropout probability for out_fft
            p_out_fft_dropatt (float): Multi-head attention dropout
            p_out_fft_dropemb (float): Dropout added to word+positional embeddings

            dur_predictor_kernel_size (int): Duration predictor conv-1D kernel size
            dur_predictor_filter_size (int): Duration predictor conv-1D filter size
            p_dur_predictor_dropout (float): Dropout probability for duration predictor
            dur_predictor_n_layers (int): Number of conv-1D layers

            pitch_predictor_kernel_size (int): Pitch predictor conv-1D kernel size
            pitch_predictor_filter_size (int): Pitch predictor conv-1D filter size
            p_pitch_predictor_dropout (float): Dropout probability for pitch predictor
            pitch_predictor_n_layers (int): Number of conv-1D layers
        """

        super(FastPitch, self).__init__()
        del max_seq_len  # unused
        del n_symbols  # unused

        self.encoder = FFTransformer(
            n_layer=in_fft_n_layers,
            n_head=in_fft_n_heads,
            d_model=symbols_embedding_dim,
            d_head=in_fft_d_head,
            d_inner=in_fft_conv1d_filter_size,
            kernel_size=in_fft_conv1d_kernel_size,
            dropout=p_in_fft_dropout,
            dropatt=p_in_fft_dropatt,
            dropemb=p_in_fft_dropemb,
            d_embed=symbols_embedding_dim,
            embed_input=True
        )

        self.duration_predictor = TemporalPredictor(
            in_fft_output_size,
            filter_size=dur_predictor_filter_size,
            kernel_size=dur_predictor_kernel_size,
            dropout=p_dur_predictor_dropout, n_layers=dur_predictor_n_layers
        )

        self.decoder = FFTransformer(
            n_layer=out_fft_n_layers, n_head=out_fft_n_heads,
            d_model=symbols_embedding_dim,
            d_head=out_fft_d_head,
            d_inner=out_fft_conv1d_filter_size,
            kernel_size=out_fft_conv1d_kernel_size,
            dropout=p_out_fft_dropout,
            dropatt=p_out_fft_dropatt,
            dropemb=p_out_fft_dropemb,
            d_embed=symbols_embedding_dim,
            embed_input=False
        )

        self.pitch_predictor = TemporalPredictor(
            in_fft_output_size,
            filter_size=pitch_predictor_filter_size,
            kernel_size=pitch_predictor_kernel_size,
            dropout=p_pitch_predictor_dropout, n_layers=pitch_predictor_n_layers
        )

        self.pitch_emb = nn.Conv1d(1, symbols_embedding_dim, kernel_size=3,
                                   padding=1)

        # Store values precomputed for training data within the model
        self.register_buffer('pitch_mean', torch.zeros(1))
        self.register_buffer('pitch_std', torch.zeros(1))

        self.proj = nn.Linear(out_fft_output_size, n_mel_channels, bias=True)

    def forward(self, inputs: Tuple, use_gt_durations: bool = True, 
                use_gt_pitch: bool = True,
                pace: float = 1.0, max_duration: int = 75):
        """A forward step.

        Args:
            inputs (tuple): inputs, _, mel_tgt, _, dur_tgt, _, pitch_tgt
            use_gt_durations (bool, optional): Use ground truth durations. Defaults to True.
            use_gt_pitch (bool, optional): Use ground truth pitch. Defaults to True.
            pace (float, optional): Pace. Defaults to 1.0.
            max_duration (int, optional): Max duration. Defaults to 75.

        Returns:
            tuple: mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred
        """

        inputs, _, mel_tgt, _, dur_tgt, _, pitch_tgt = inputs
        mel_max_len = mel_tgt.size(2)

        # Input FFT
        enc_out, enc_mask = self.encoder(inputs)

        # Embedded for predictors
        pred_enc_out, pred_enc_mask = enc_out, enc_mask

        # Predict durations
        log_dur_pred = self.duration_predictor(pred_enc_out, pred_enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, 0, max_duration)

        # Predict pitch
        pitch_pred = self.pitch_predictor(enc_out, enc_mask)

        if use_gt_pitch and pitch_tgt is not None:
            pitch_emb = self.pitch_emb(pitch_tgt.unsqueeze(1))
        else:
            pitch_emb = self.pitch_emb(pitch_pred.unsqueeze(1))
        enc_out = enc_out + pitch_emb.transpose(1, 2)

        len_regulated, dec_lens = regulate_len(
            dur_tgt if use_gt_durations else dur_pred,
            enc_out, pace, mel_max_len)

        # Output FFT
        dec_out, dec_mask = self.decoder(len_regulated, dec_lens)
        mel_out = self.proj(dec_out)
        return mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred

    def infer(self, inputs, input_lens, pace: float=1.0, dur_tgt=None, pitch_tgt=None,
              pitch_transform=None, max_duration: int = 75):
        """Inference.

        Args:
            inputs (tuple): [description]
            input_lens ([type]): [description]
            pace (float, optional): Target pace. Defaults to 1.0.
            dur_tgt ([type], optional): Target duration. Defaults to None.
            pitch_tgt ([type], optional): Target pitch. Defaults to None.
            pitch_transform ([type], optional): [description]. Defaults to None.
            max_duration (int, optional): [description]. Defaults to 75.

        Returns:
            tuple: mel_out, dec_lens, dur_pred, pitch_pred
        """

        del input_lens  # unused

        # Input FFT
        enc_out, enc_mask = self.encoder(inputs)

        # Embedded for predictors
        pred_enc_out, pred_enc_mask = enc_out, enc_mask

        # Predict durations
        log_dur_pred = self.duration_predictor(pred_enc_out, pred_enc_mask)
        dur_pred = torch.clamp(torch.exp(log_dur_pred) - 1, 0, max_duration)

        # Pitch over chars
        pitch_pred = self.pitch_predictor(enc_out, enc_mask)

        if pitch_transform is not None:
            if self.pitch_std[0] == 0.0:
                # XXX LJSpeech-1.1 defaults
                mean, std = 218.14, 67.24
            else:
                mean, std = self.pitch_mean[0], self.pitch_std[0]
            pitch_pred = pitch_transform(pitch_pred, mean, std)

        if pitch_tgt is None:
            pitch_emb = self.pitch_emb(pitch_pred.unsqueeze(1)).transpose(1, 2)
        else:
            pitch_emb = self.pitch_emb(pitch_tgt.unsqueeze(1)).transpose(1, 2)

        enc_out = enc_out + pitch_emb

        len_regulated, dec_lens = regulate_len(
            dur_pred if dur_tgt is None else dur_tgt,
            enc_out, pace, mel_max_len=None)

        dec_out, dec_mask = self.decoder(len_regulated, dec_lens)
        mel_out = self.proj(dec_out)
        # mel_lens = dec_mask.squeeze(2).sum(axis=1).long()
        mel_out = mel_out.permute(0, 2, 1)  # For inference.py
        return mel_out, dec_lens, dur_pred, pitch_pred
# Adapted from # https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechSynthesis/FastPitch # ***************************************************************************** #  Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved. # #  Redistribution and use in source and binary forms, with or without #  modification, are permitted provided that the following conditions are met: #      * Redistributions of source code must retain the above copyright #        notice, this list of conditions and the following disclaimer. #      * Redistributions in binary form must reproduce the above copyright #        notice, this list of conditions and the following disclaimer in the #        documentation and/or other materials provided with the distribution. #      * Neither the name of the NVIDIA CORPORATION nor the #        names of its contributors may be used to endorse or promote products #        derived from this software without specific prior written permission. # #  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND #  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED #  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE #  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY #  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES #  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; #  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND #  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT #  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS #  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # # ***************************************************************************** Regulate the length of durations.
 
     Args:
         durations ([type]): Durations, ground truth or predicted.
         enc_out ([type]): [description]
         pace (float, optional): Target pace. Defaults to 1.0.
         mel_max_len (int, optional): Max length of mel. Defaults to None.
 
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: enc_rep, dec_lens Predicts a single float per each temporal location.
     Used as duration predictor and pitch predictor. FastPitch module definition. FastPitch's constructor
 
         Args:
             n_mel_channels (int): Number of bins in mel-spectrograms
             max_seq_len (int): Unused.
             n_symbols (int): Unused.
             symbols_embedding_dim (int): Input embedding dimension
 
             in_fft_n_layers (int): Number of FFT blocks
             in_fft_n_heads (int): Number of attention heads
             in_fft_d_head (int): Dim of attention heads
             in_fft_conv1d_kernel_size (int): Conv-1D kernel size
             in_fft_conv1d_filter_size (int): Conv-1D filter size
             in_fft_output_size (int): Output dim
 
             p_in_fft_dropout (float): Dropout probability
             p_in_fft_dropatt (float): Multi-head attention dropout
             p_in_fft_dropemb (float): Dropout added to word+positional embeddings')
 
 
             out_fft_n_layers (int): Number of FFT blocks
             out_fft_n_heads (int): Number of attention heads
             out_fft_d_head (int): Dim of attention head
             out_fft_conv1d_kernel_size (int): Conv-1D kernel size
             out_fft_conv1d_filter_size (int): Conv-1D filter size
             out_fft_output_size (int): Output dim
 
             p_out_fft_dropout (float): Dropout probability for out_fft
             p_out_fft_dropatt (float): Multi-head attention dropout
             p_out_fft_dropemb (float): Dropout added to word+positional embeddings
 
             dur_predictor_kernel_size (int): Duration predictor conv-1D kernel size
             dur_predictor_filter_size (int): Duration predictor conv-1D filter size
             p_dur_predictor_dropout (float): Dropout probability for duration predictor
             dur_predictor_n_layers (int): Number of conv-1D layers
 
             pitch_predictor_kernel_size (int): Pitch predictor conv-1D kernel size
             pitch_predictor_filter_size (int): Pitch predictor conv-1D filter size
             p_pitch_predictor_dropout (float): Dropout probability for pitch predictor
             pitch_predictor_n_layers (int): Number of conv-1D layers # unused # unused # Store values precomputed for training data within the model A forward step.
 
         Args:
             inputs (tuple): inputs, _, mel_tgt, _, dur_tgt, _, pitch_tgt
             use_gt_durations (bool, optional): Use ground truth durations. Defaults to True.
             use_gt_pitch (bool, optional): Use ground truth pitch. Defaults to True.
             pace (float, optional): Pace. Defaults to 1.0.
             max_duration (int, optional): Max duration. Defaults to 75.
 
         Returns:
             tuple: mel_out, dec_mask, dur_pred, log_dur_pred, pitch_pred # Input FFT # Embedded for predictors # Predict durations # Predict pitch # Output FFT Inference.
 
         Args:
             inputs (tuple): [description]
             input_lens ([type]): [description]
             pace (float, optional): Target pace. Defaults to 1.0.
             dur_tgt ([type], optional): Target duration. Defaults to None.
             pitch_tgt ([type], optional): Target pitch. Defaults to None.
             pitch_transform ([type], optional): [description]. Defaults to None.
             max_duration (int, optional): [description]. Defaults to 75.
 
         Returns:
             tuple: mel_out, dec_lens, dur_pred, pitch_pred # unused # Input FFT # Embedded for predictors # Predict durations # Pitch over chars # XXX LJSpeech-1.1 defaults # mel_lens = dec_mask.squeeze(2).sum(axis=1).long() # For inference.py