Shadhil's picture
voice-clone with single audio sample input
9b2107c
raw
history blame
24.7 kB
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