TTS/tts/layers/overflow/neural_hmm.py

from typing import List

import torch
import torch.distributions as tdist
import torch.nn.functional as F
from torch import nn
from torch.utils.checkpoint import checkpoint

from TTS.tts.layers.overflow.common_layers import Outputnet, OverflowUtils
from TTS.tts.layers.tacotron.common_layers import Prenet
from TTS.tts.utils.helpers import sequence_mask


class NeuralHMM(nn.Module):
    """Autoregressive left to right HMM model primarily used in "Neural HMMs are all you need (for high-quality attention-free TTS)"

    Paper::
        https://arxiv.org/abs/2108.13320

    Paper abstract::
        Neural sequence-to-sequence TTS has achieved significantly better output quality than statistical speech synthesis using
        HMMs. However, neural TTS is generally not probabilistic and uses non-monotonic attention. Attention failures increase
        training time and can make synthesis babble incoherently. This paper describes how the old and new paradigms can be
        combined to obtain the advantages of both worlds, by replacing attention in neural TTS with an autoregressive left-right
        no-skip hidden Markov model defined by a neural network. Based on this proposal, we modify Tacotron 2 to obtain an
        HMM-based neural TTS model with monotonic alignment, trained to maximise the full sequence likelihood without
        approximation. We also describe how to combine ideas from classical and contemporary TTS for best results. The resulting
        example system is smaller and simpler than Tacotron 2, and learns to speak with fewer iterations and less data, whilst
        achieving comparable naturalness prior to the post-net. Our approach also allows easy control over speaking rate.

    Args:
        frame_channels (int): Output dimension to generate.
        ar_order (int): Autoregressive order of the model. In ablations of Neural HMM it was found that more autoregression while giving more variation hurts naturalness of the synthesised audio.
        deterministic_transition (bool): deterministic duration generation based on duration quantiles as defiend in "S. Ronanki, O. Watts, S. King, and G. E. Henter, “Medianbased generation of synthetic speech durations using a nonparametric approach,” in Proc. SLT, 2016.". Defaults to True.
        encoder_dim (int): Channels of encoder input and character embedding tensors. Defaults to 512.
        prenet_type (str): `original` or `bn`. `original` sets the default Prenet and `bn` uses Batch Normalization version of the Prenet.
        prenet_dim (int): Dimension of the Prenet.
        prenet_n_layers (int): Number of layers in the Prenet.
        prenet_dropout (float): Dropout probability of the Prenet.
        prenet_dropout_at_inference (bool): If True, dropout is applied at inference time.
        memory_rnn_dim (int): Size of the memory RNN to process output of prenet.
        outputnet_size (List[int]): Size of the output network inside the neural HMM.
        flat_start_params (dict): Parameters for the flat start initialization of the neural HMM.
        std_floor (float): Floor value for the standard deviation of the neural HMM. Prevents model cheating by putting point mass and getting infinite likelihood at any datapoint.
        use_grad_checkpointing (bool, optional): Use gradient checkpointing to save memory. Defaults to True.
    """

    def __init__(
        self,
        frame_channels: int,
        ar_order: int,
        deterministic_transition: bool,
        encoder_dim: int,
        prenet_type: str,
        prenet_dim: int,
        prenet_n_layers: int,
        prenet_dropout: float,
        prenet_dropout_at_inference: bool,
        memory_rnn_dim: int,
        outputnet_size: List[int],
        flat_start_params: dict,
        std_floor: float,
        use_grad_checkpointing: bool = True,
    ):
        super().__init__()

        self.frame_channels = frame_channels
        self.ar_order = ar_order
        self.deterministic_transition = deterministic_transition
        self.prenet_dim = prenet_dim
        self.memory_rnn_dim = memory_rnn_dim
        self.use_grad_checkpointing = use_grad_checkpointing

        self.transition_model = TransitionModel()
        self.emission_model = EmissionModel()

        assert ar_order > 0, f"AR order must be greater than 0 provided {ar_order}"

        self.ar_order = ar_order
        self.prenet = Prenet(
            in_features=frame_channels * ar_order,
            prenet_type=prenet_type,
            prenet_dropout=prenet_dropout,
            dropout_at_inference=prenet_dropout_at_inference,
            out_features=[self.prenet_dim for _ in range(prenet_n_layers)],
            bias=False,
        )
        self.memory_rnn = nn.LSTMCell(input_size=prenet_dim, hidden_size=memory_rnn_dim)
        self.output_net = Outputnet(
            encoder_dim, memory_rnn_dim, frame_channels, outputnet_size, flat_start_params, std_floor
        )
        self.register_buffer("go_tokens", torch.zeros(ar_order, 1))

    def forward(self, inputs, inputs_len, mels, mel_lens):
        r"""HMM forward algorithm for training uses logarithmic version of Rabiner (1989) forward algorithm.

        Args:
            inputs (torch.FloatTensor): Encoder outputs
            inputs_len (torch.LongTensor): Encoder output lengths
            mels (torch.FloatTensor): Mel inputs
            mel_lens (torch.LongTensor): Length of mel inputs

        Shapes:
            - inputs: (B, T, D_out_enc)
            - inputs_len: (B)
            - mels: (B, D_mel, T_mel)
            - mel_lens: (B)

        Returns:
            log_prob (torch.FloatTensor): Log probability of the sequence
        """
        # Get dimensions of inputs
        batch_size, N, _ = inputs.shape
        T_max = torch.max(mel_lens)
        mels = mels.permute(0, 2, 1)

        # Intialize forward algorithm
        log_state_priors = self._initialize_log_state_priors(inputs)
        log_c, log_alpha_scaled, transition_matrix, means = self._initialize_forward_algorithm_variables(mels, N)

        # Initialize autoregression elements
        ar_inputs = self._add_go_token(mels)
        h_memory, c_memory = self._init_lstm_states(batch_size, self.memory_rnn_dim, mels)

        for t in range(T_max):
            # Process Autoregression
            h_memory, c_memory = self._process_ar_timestep(t, ar_inputs, h_memory, c_memory)
            # Get mean, std and transition vector from decoder for this timestep
            # Note: Gradient checkpointing currently doesn't works with multiple gpus inside a loop
            if self.use_grad_checkpointing and self.training:
                mean, std, transition_vector = checkpoint(self.output_net, h_memory, inputs)
            else:
                mean, std, transition_vector = self.output_net(h_memory, inputs)

            if t == 0:
                log_alpha_temp = log_state_priors + self.emission_model(mels[:, 0], mean, std, inputs_len)
            else:
                log_alpha_temp = self.emission_model(mels[:, t], mean, std, inputs_len) + self.transition_model(
                    log_alpha_scaled[:, t - 1, :], transition_vector, inputs_len
                )
            log_c[:, t] = torch.logsumexp(log_alpha_temp, dim=1)
            log_alpha_scaled[:, t, :] = log_alpha_temp - log_c[:, t].unsqueeze(1)
            transition_matrix[:, t] = transition_vector  # needed for absorption state calculation

            # Save for plotting
            means.append(mean.detach())

        log_c, log_alpha_scaled = self._mask_lengths(mel_lens, log_c, log_alpha_scaled)

        sum_final_log_c = self.get_absorption_state_scaling_factor(
            mel_lens, log_alpha_scaled, inputs_len, transition_matrix
        )

        log_probs = torch.sum(log_c, dim=1) + sum_final_log_c

        return log_probs, log_alpha_scaled, transition_matrix, means

    @staticmethod
    def _mask_lengths(mel_lens, log_c, log_alpha_scaled):
        """
        Mask the lengths of the forward variables so that the variable lenghts
        do not contribute in the loss calculation
        Args:
            mel_inputs (torch.FloatTensor): (batch, T, frame_channels)
            mel_inputs_lengths (torch.IntTensor): (batch)
            log_c (torch.FloatTensor): (batch, T)
        Returns:
            log_c (torch.FloatTensor) : scaled probabilities (batch, T)
            log_alpha_scaled (torch.FloatTensor): forward probabilities (batch, T, N)
        """
        mask_log_c = sequence_mask(mel_lens)
        log_c = log_c * mask_log_c
        mask_log_alpha_scaled = mask_log_c.unsqueeze(2)
        log_alpha_scaled = log_alpha_scaled * mask_log_alpha_scaled
        return log_c, log_alpha_scaled

    def _process_ar_timestep(
        self,
        t,
        ar_inputs,
        h_memory,
        c_memory,
    ):
        """
        Process autoregression in timestep
        1. At a specific t timestep
        2. Perform data dropout if applied (we did not use it)
        3. Run the autoregressive frame through the prenet (has dropout)
        4. Run the prenet output through the post prenet rnn

        Args:
            t (int): mel-spec timestep
            ar_inputs (torch.FloatTensor): go-token appended mel-spectrograms
                - shape: (b, D_out, T_out)
            h_post_prenet (torch.FloatTensor): previous timestep rnn hidden state
                - shape: (b, memory_rnn_dim)
            c_post_prenet (torch.FloatTensor): previous timestep rnn cell state
                - shape: (b, memory_rnn_dim)

        Returns:
            h_post_prenet (torch.FloatTensor): rnn hidden state of the current timestep
            c_post_prenet (torch.FloatTensor): rnn cell state of the current timestep
        """
        prenet_input = ar_inputs[:, t : t + self.ar_order].flatten(1)
        memory_inputs = self.prenet(prenet_input)
        h_memory, c_memory = self.memory_rnn(memory_inputs, (h_memory, c_memory))
        return h_memory, c_memory

    def _add_go_token(self, mel_inputs):
        """Append the go token to create the autoregressive input
        Args:
            mel_inputs (torch.FloatTensor): (batch_size, T, n_mel_channel)
        Returns:
            ar_inputs (torch.FloatTensor): (batch_size, T, n_mel_channel)
        """
        batch_size, T, _ = mel_inputs.shape
        go_tokens = self.go_tokens.unsqueeze(0).expand(batch_size, self.ar_order, self.frame_channels)
        ar_inputs = torch.cat((go_tokens, mel_inputs), dim=1)[:, :T]
        return ar_inputs

    @staticmethod
    def _initialize_forward_algorithm_variables(mel_inputs, N):
        r"""Initialize placeholders for forward algorithm variables, to use a stable
                version we will use log_alpha_scaled and the scaling constant

        Args:
            mel_inputs (torch.FloatTensor): (b, T_max, frame_channels)
            N (int): number of states
        Returns:
            log_c (torch.FloatTensor): Scaling constant (b, T_max)
        """
        b, T_max, _ = mel_inputs.shape
        log_alpha_scaled = mel_inputs.new_zeros((b, T_max, N))
        log_c = mel_inputs.new_zeros(b, T_max)
        transition_matrix = mel_inputs.new_zeros((b, T_max, N))

        # Saving for plotting later, will not have gradient tapes
        means = []
        return log_c, log_alpha_scaled, transition_matrix, means

    @staticmethod
    def _init_lstm_states(batch_size, hidden_state_dim, device_tensor):
        r"""
        Initialize Hidden and Cell states for LSTM Cell

        Args:
            batch_size (Int): batch size
            hidden_state_dim (Int): dimensions of the h and c
            device_tensor (torch.FloatTensor): useful for the device and type

        Returns:
            (torch.FloatTensor): shape (batch_size, hidden_state_dim)
                can be hidden state for LSTM
            (torch.FloatTensor): shape (batch_size, hidden_state_dim)
                can be the cell state for LSTM
        """
        return (
            device_tensor.new_zeros(batch_size, hidden_state_dim),
            device_tensor.new_zeros(batch_size, hidden_state_dim),
        )

    def get_absorption_state_scaling_factor(self, mels_len, log_alpha_scaled, inputs_len, transition_vector):
        """Returns the final scaling factor of absorption state

        Args:
            mels_len (torch.IntTensor): Input size of mels to
                    get the last timestep of log_alpha_scaled
            log_alpha_scaled (torch.FloatTEnsor): State probabilities
            text_lengths (torch.IntTensor): length of the states to
                    mask the values of states lengths
                (
                    Useful when the batch has very different lengths,
                    when the length of an observation is less than
                    the number of max states, then the log alpha after
                    the state value is filled with -infs. So we mask
                    those values so that it only consider the states
                    which are needed for that length
                )
            transition_vector (torch.FloatTensor): transtiion vector for each state per timestep

        Shapes:
            - mels_len: (batch_size)
            - log_alpha_scaled: (batch_size, N, T)
            - text_lengths: (batch_size)
            - transition_vector: (batch_size, N, T)

        Returns:
            sum_final_log_c (torch.FloatTensor): (batch_size)

        """
        N = torch.max(inputs_len)
        max_inputs_len = log_alpha_scaled.shape[2]
        state_lengths_mask = sequence_mask(inputs_len, max_len=max_inputs_len)

        last_log_alpha_scaled_index = (
            (mels_len - 1).unsqueeze(-1).expand(-1, N).unsqueeze(1)
        )  # Batch X Hidden State Size
        last_log_alpha_scaled = torch.gather(log_alpha_scaled, 1, last_log_alpha_scaled_index).squeeze(1)
        last_log_alpha_scaled = last_log_alpha_scaled.masked_fill(~state_lengths_mask, -float("inf"))

        last_transition_vector = torch.gather(transition_vector, 1, last_log_alpha_scaled_index).squeeze(1)
        last_transition_probability = torch.sigmoid(last_transition_vector)
        log_probability_of_transitioning = OverflowUtils.log_clamped(last_transition_probability)

        last_transition_probability_index = self.get_mask_for_last_item(inputs_len, inputs_len.device)
        log_probability_of_transitioning = log_probability_of_transitioning.masked_fill(
            ~last_transition_probability_index, -float("inf")
        )
        final_log_c = last_log_alpha_scaled + log_probability_of_transitioning

        # If the length of the mel is less than the number of states it will select the -inf values leading to nan gradients
        # Ideally, we should clean the dataset otherwise this is a little hack uncomment the line below
        final_log_c = final_log_c.clamp(min=torch.finfo(final_log_c.dtype).min)

        sum_final_log_c = torch.logsumexp(final_log_c, dim=1)
        return sum_final_log_c

    @staticmethod
    def get_mask_for_last_item(lengths, device, out_tensor=None):
        """Returns n-1 mask for the last item in the sequence.

        Args:
            lengths (torch.IntTensor): lengths in a batch
            device (str, optional): Defaults to "cpu".
            out_tensor (torch.Tensor, optional): uses the memory of a specific tensor.
                Defaults to None.

        Returns:
            - Shape: :math:`(b, max_len)`
        """
        max_len = torch.max(lengths).item()
        ids = (
            torch.arange(0, max_len, device=device) if out_tensor is None else torch.arange(0, max_len, out=out_tensor)
        )
        mask = ids == lengths.unsqueeze(1) - 1
        return mask

    @torch.inference_mode()
    def inference(
        self,
        inputs: torch.FloatTensor,
        input_lens: torch.LongTensor,
        sampling_temp: float,
        max_sampling_time: int,
        duration_threshold: float,
    ):
        """Inference from autoregressive neural HMM

        Args:
            inputs (torch.FloatTensor): input states
                - shape: :math:`(b, T, d)`
            input_lens (torch.LongTensor): input state lengths
                - shape: :math:`(b)`
            sampling_temp (float): sampling temperature
            max_sampling_temp (int): max sampling temperature
            duration_threshold (float): duration threshold to switch to next state
                - Use this to change the spearking rate of the synthesised audio
        """

        b = inputs.shape[0]
        outputs = {
            "hmm_outputs": [],
            "hmm_outputs_len": [],
            "alignments": [],
            "input_parameters": [],
            "output_parameters": [],
        }
        for i in range(b):
            neural_hmm_outputs, states_travelled, input_parameters, output_parameters = self.sample(
                inputs[i : i + 1], input_lens[i], sampling_temp, max_sampling_time, duration_threshold
            )

            outputs["hmm_outputs"].append(neural_hmm_outputs)
            outputs["hmm_outputs_len"].append(neural_hmm_outputs.shape[0])
            outputs["alignments"].append(states_travelled)
            outputs["input_parameters"].append(input_parameters)
            outputs["output_parameters"].append(output_parameters)

        outputs["hmm_outputs"] = nn.utils.rnn.pad_sequence(outputs["hmm_outputs"], batch_first=True)
        outputs["hmm_outputs_len"] = torch.tensor(
            outputs["hmm_outputs_len"], dtype=input_lens.dtype, device=input_lens.device
        )
        return outputs

    @torch.inference_mode()
    def sample(self, inputs, input_lens, sampling_temp, max_sampling_time, duration_threshold):
        """Samples an output from the parameter models

        Args:
            inputs (torch.FloatTensor): input states
                - shape: :math:`(1, T, d)`
            input_lens (torch.LongTensor): input state lengths
                - shape: :math:`(1)`
            sampling_temp (float): sampling temperature
            max_sampling_time (int): max sampling time
            duration_threshold (float): duration threshold to switch to next state

        Returns:
            outputs (torch.FloatTensor): Output Observations
                - Shape: :math:`(T, output_dim)`
            states_travelled (list[int]): Hidden states travelled
                - Shape: :math:`(T)`
            input_parameters (list[torch.FloatTensor]): Input parameters
            output_parameters (list[torch.FloatTensor]): Output parameters
        """
        states_travelled, outputs, t = [], [], 0

        # Sample initial state
        current_state = 0
        states_travelled.append(current_state)

        # Prepare autoregression
        prenet_input = self.go_tokens.unsqueeze(0).expand(1, self.ar_order, self.frame_channels)
        h_memory, c_memory = self._init_lstm_states(1, self.memory_rnn_dim, prenet_input)

        input_parameter_values = []
        output_parameter_values = []
        quantile = 1
        while True:
            memory_input = self.prenet(prenet_input.flatten(1).unsqueeze(0))
            # will be 1 while sampling
            h_memory, c_memory = self.memory_rnn(memory_input.squeeze(0), (h_memory, c_memory))

            z_t = inputs[:, current_state].unsqueeze(0)  # Add fake time dimension
            mean, std, transition_vector = self.output_net(h_memory, z_t)

            transition_probability = torch.sigmoid(transition_vector.flatten())
            staying_probability = torch.sigmoid(-transition_vector.flatten())

            # Save for plotting
            input_parameter_values.append([prenet_input, current_state])
            output_parameter_values.append([mean, std, transition_probability])

            x_t = self.emission_model.sample(mean, std, sampling_temp=sampling_temp)

            # Prepare autoregressive input for next iteration
            prenet_input = torch.cat((prenet_input, x_t), dim=1)[:, 1:]

            outputs.append(x_t.flatten())

            transition_matrix = torch.cat((staying_probability, transition_probability))
            quantile *= staying_probability
            if not self.deterministic_transition:
                switch = transition_matrix.multinomial(1)[0].item()
            else:
                switch = quantile < duration_threshold

            if switch:
                current_state += 1
                quantile = 1

            states_travelled.append(current_state)

            if (current_state == input_lens) or (max_sampling_time and t == max_sampling_time - 1):
                break

            t += 1

        return (
            torch.stack(outputs, dim=0),
            F.one_hot(input_lens.new_tensor(states_travelled)),
            input_parameter_values,
            output_parameter_values,
        )

    @staticmethod
    def _initialize_log_state_priors(text_embeddings):
        """Creates the log pi in forward algorithm.

        Args:
            text_embeddings (torch.FloatTensor): used to create the log pi
                    on current device

        Shapes:
            - text_embeddings: (B, T, D_out_enc)
        """
        N = text_embeddings.shape[1]
        log_state_priors = text_embeddings.new_full([N], -float("inf"))
        log_state_priors[0] = 0.0
        return log_state_priors


class TransitionModel(nn.Module):
    """Transition Model of the HMM, it represents the probability of transitioning
    form current state to all other states"""

    def forward(self, log_alpha_scaled, transition_vector, inputs_len):  # pylint: disable=no-self-use
        r"""
        product of the past state with transitional probabilities in log space

        Args:
            log_alpha_scaled (torch.Tensor): Multiply previous timestep's alphas by
                        transition matrix (in log domain)
                - shape: (batch size, N)
            transition_vector (torch.tensor): transition vector for each state
                - shape: (N)
            inputs_len (int tensor): Lengths of states in a batch
                - shape: (batch)

        Returns:
            out (torch.FloatTensor): log probability of transitioning to each state
        """
        transition_p = torch.sigmoid(transition_vector)
        staying_p = torch.sigmoid(-transition_vector)

        log_staying_probability = OverflowUtils.log_clamped(staying_p)
        log_transition_probability = OverflowUtils.log_clamped(transition_p)

        staying = log_alpha_scaled + log_staying_probability
        leaving = log_alpha_scaled + log_transition_probability
        leaving = leaving.roll(1, dims=1)
        leaving[:, 0] = -float("inf")
        inputs_len_mask = sequence_mask(inputs_len)
        out = OverflowUtils.logsumexp(torch.stack((staying, leaving), dim=2), dim=2)
        out = out.masked_fill(~inputs_len_mask, -float("inf"))  # There are no states to contribute to the loss
        return out


class EmissionModel(nn.Module):
    """Emission Model of the HMM, it represents the probability of
    emitting an observation based on the current state"""

    def __init__(self) -> None:
        super().__init__()
        self.distribution_function: tdist.Distribution = tdist.normal.Normal

    def sample(self, means, stds, sampling_temp):
        return self.distribution_function(means, stds * sampling_temp).sample() if sampling_temp > 0 else means

    def forward(self, x_t, means, stds, state_lengths):
        r"""Calculates the log probability of the the given data (x_t)
            being observed from states with given means and stds
        Args:
            x_t (float tensor) : observation at current time step
                - shape: (batch, feature_dim)
            means (float tensor): means of the distributions of hidden states
                - shape: (batch, hidden_state, feature_dim)
            stds (float tensor): standard deviations of the distributions of the hidden states
                - shape: (batch, hidden_state, feature_dim)
            state_lengths (int tensor): Lengths of states in a batch
                - shape: (batch)

        Returns:
            out (float tensor): observation log likelihoods,
                                    expressing the probability of an observation
                being generated from a state i
                shape: (batch, hidden_state)
        """
        emission_dists = self.distribution_function(means, stds)
        out = emission_dists.log_prob(x_t.unsqueeze(1))
        state_lengths_mask = sequence_mask(state_lengths).unsqueeze(2)
        out = torch.sum(out * state_lengths_mask, dim=2)
        return out