Adding overflow setting up flat start

TTS/tts/configs/overflow_config.py

@@ -0,0 +1,110 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from TTS.tts.configs.shared_configs import BaseTTSConfig
+
+
+@dataclass
+class OverFlowConfig(BaseTTSConfig):
+    """
+    Define parameters for OverFlow model.
+
+    Args:
+        BaseTTSConfig (_type_): _description_
+    """
+    model: str = "overflow"
+
+    # data parameters
+    normalize_mel: bool = True
+    normalized_mel_parameter_path: str = None
+
+    # Encoder parameters
+    num_chars: int = None
+    state_per_phone: int = 2
+    encoder_in_out_features: int = 512
+    encoder_n_convolutions: int = 3
+
+    # HMM parameters
+    out_channels: int = 80
+    ar_order: int = 1
+    sampling_temp: float = 0.667
+    deterministic_transition: bool = True
+    duration_threshold: float = 0.55
+    use_grad_checkpointing: bool = True
+
+    ## Prenet parameters
+    prenet_type: str = "original"
+    prenet_dim: int = 256
+    prenet_n_layers: int = 2
+    prenet_dropout: float = 0.5
+    prenet_dropout_at_inference: bool = False
+    memory_rnn_dim: int = 1024
+
+    ## Outputnet parameters
+    outputnet_size: List[int] = field(default_factory=lambda: [256, 256])
+    flat_start_params: dict = field(
+        default_factory=lambda: {
+            "mean": 0.0,
+            "std": 1.0,
+            "transition_p": 0.14
+        }
+    )
+    std_floor: float = 0.01
+
+    # Decoder parameters
+    hidden_channels_dec: int = 150
+    kernel_size_dec: int = 5
+    dilation_rate: int = 1
+    num_flow_blocks_dec: int = 12
+    num_block_layers: int = 4
+    dropout_p_dec: float = 0.05
+    num_splits: int = 4
+    num_squeeze: int = 2
+    sigmoid_scale: bool = False
+    c_in_channels: int = 0
+
+    # optimizer parameters
+    optimizer: str = "RAdam"
+    optimizer_params: dict = field(default_factory=lambda: {"betas": [0.9, 0.998], "weight_decay": 1e-6})
+    lr_scheduler: str = "NoamLR"
+    lr_scheduler_params: dict = field(default_factory=lambda: {"warmup_steps": 4000})
+    grad_clip: float = 40000.0
+    lr: float = 1e-3
+
+    # overrides
+    min_seq_len: int = 3
+    max_seq_len: int = 500
+
+    # testing
+    test_sentences: List[str] = field(
+        default_factory=lambda: [
+            "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.",
+        ]
+    ) 
+
+
+    # Extra needed config
+    # Do not change overflow does not use them
+    r: int = 1 
+    use_d_vector_file: bool = False
+
+    def check_values(self):
+        """Validate the hyperparameters.
+
+        Raises:
+            AssertionError: when the parameters network is not defined
+            AssertionError: transition probability is not between 0 and 1
+        """
+        assert (
+            self.parameternetwork >= 1
+        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
+        assert (
+            0 < self.flat_start_params["transition_p"] < 1
+        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"
+
+        if self.normalize_mel:
+            assert self.normalized_mel_parameter_path is not None, "Normalized mel parameter path must be provided when normalize_mel is True."

TTS/tts/datasets/dataset.py

@@ -65,6 +65,7 @@ def __init__(
         use_noise_augment: bool = False,
         start_by_longest: bool = False,
         verbose: bool = False,
+        compute_mel_statistics: bool = False,
     ):
         """Generic 📂 data loader for `tts` models. It is configurable for different outputs and needs.
 
@@ -140,6 +141,7 @@ def __init__(
         self.language_id_mapping = language_id_mapping
         self.use_noise_augment = use_noise_augment
         self.start_by_longest = start_by_longest
+        self.compute_mel_statistics = compute_mel_statistics
 
         self.verbose = verbose
         self.rescue_item_idx = 1

TTS/tts/layers/losses.py

@@ -872,3 +872,25 @@ def forward(
 
         return_dict["loss"] = loss
         return return_dict
+
+
+class NLLLoss(nn.Module):
+    """Negative log likelihood loss."""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, log_prob: torch.Tensor) -> dict:
+        """Compute the loss.
+
+        Args:
+            logits (Tensor): [B, T, D]
+
+        Returns:
+            Tensor: [1]
+
+        """
+        return_dict = {}
+        return_dict["loss"] = - log_prob.mean()
+        return return_dict 
+

TTS/tts/layers/neural_hmm/common_layers.py

@@ -25,17 +25,18 @@ class Encoder(nn.Module):
     def __init__(
         self, num_chars,
         state_per_phone,
-        in_out_channels=512
+        in_out_channels=512,
+        n_convolutions=3
     ):
 
         super().__init__()
 
         self.state_per_phone = state_per_phone
         self.in_out_channels = in_out_channels
 
-        self.emb = nn.Embedding(num_chars, hidden_channels)
+        self.emb = nn.Embedding(num_chars, in_out_channels)
         self.convolutions = nn.ModuleList()
-        for _ in range(3):
+        for _ in range(n_convolutions):
             self.convolutions.append(ConvBNBlock(in_out_channels, in_out_channels, 5, "relu"))
         self.lstm = nn.LSTM(
             in_out_channels,
@@ -68,22 +69,20 @@ class ParameterModel(nn.Module):
     Note: Do not put dropout layers here, the model will not converge.
 
     Args:
-            parameternetwork (List[int]): the architecture of the parameter model
+            outputnet_size (List[int]): the architecture of the parameter model
             input_size (int): size of input for the first layer
             output_size (int): size of output i.e size of the feature dim
             frame_channels (int): feature dim to set the flat start bias
-            init_transition_probability (float): flat start transition probability
-            init_mean (float): flat start mean
-            init_std (float): flat start std
+            flat_start_params (dict): flat start parameters to set the bias
     """
 
     def __init__(
         self,
         outputnet_size: List[int],
         input_size: int,
         output_size: int,
-        flat_start_params: dict,
         frame_channels: int,
+        flat_start_params: dict,
     ):
         super().__init__()
         self.flat_start_params = flat_start_params
@@ -134,8 +133,6 @@ def __init__(
         input_size = memory_rnn_dim + encoder_dim
         output_size = 2 * frame_channels + 1
 
-        self._validate_parameters()
-
         self.parametermodel = ParameterModel(
             outputnet_size=outputnet_size,
             input_size=input_size,
@@ -144,20 +141,6 @@ def __init__(
             frame_channels=frame_channels,
         )
 
-    def _validate_parameters(self):
-        """Validate the hyperparameters.
-
-        Raises:
-            AssertionError: when the parameters network is not defined
-            AssertionError: transition probability is not between 0 and 1
-        """
-        assert (
-            self.parameternetwork >= 1
-        ), f"Parameter Network must have atleast one layer check the config file for parameter network. Provided: {self.parameternetwork}"
-        assert (
-            0 < self.flat_start_params["transition_p"] < 1
-        ), f"Transition probability must be between 0 and 1. Provided: {self.flat_start_params['transition_p']}"
-
     def forward(self, ar_mels, inputs):
         r"""Inputs observation and returns the means, stds and transition probability for the current state
 
@@ -205,3 +188,46 @@ def _floor_std(self, std):
                 "[*] Standard deviation was floored! The model is preventing overfitting, nothing serious to worry about"
             )
         return std
+
+
+class OverFlowUtils:
+    @staticmethod
+    def get_data_parameters_for_flat_start(data_loader: torch.utils.data.DataLoader, out_channels: int, states_per_phone: int):
+        """Generates data parameters for flat starting the HMM.
+
+        Args:
+            data_loader (torch.utils.data.Dataloader): _description_
+            out_channels (int): mel spectrogram channels 
+            states_per_phone (_type_): HMM states per phone
+        """
+
+        # State related information for transition_p
+        total_state_len = 0
+        total_mel_len = 0
+
+        # Useful for data mean an std
+        total_mel_sum = 0
+        total_mel_sq_sum = 0
+
+        for batch in tqdm(data_loader, leave=False):
+            text_lengths = batch['token_id_lengths']
+            mels = batch['mel']
+            mel_lengths = batch['mel_lengths']
+
+            total_state_len += torch.sum(text_lengths)
+            total_mel_len += torch.sum(mel_lengths)
+            total_mel_sum += torch.sum(mels)
+            total_mel_sq_sum += torch.sum(torch.pow(mels, 2))
+
+        data_mean = total_mel_sum / (total_mel_len * out_channels)
+        data_std = torch.sqrt((total_mel_sq_sum / (total_mel_len * out_channels)) - torch.pow(data_mean, 2))
+        average_num_states = total_state_len / len(data_loader.dataset)
+        average_mel_len = total_mel_len / len(data_loader.dataset)
+        average_duration_each_state = average_mel_len / average_num_states
+        init_transition_prob = 1 / average_duration_each_state
+
+        return data_mean, data_std, init_transition_prob 
+
+
+
+