Merge pull request #49 from idiap/vc-refactors

VC-related refactors and fixes

Author: eginhard

TTS/bin/train_encoder.py

@@ -161,9 +161,6 @@ def train(model, optimizer, scheduler, criterion, data_loader, eval_data_loader,
             loader_time = time.time() - end_time
             global_step += 1
 
-            # setup lr
-            if c.lr_decay:
-                scheduler.step()
             optimizer.zero_grad()
 
             # dispatch data to GPU
@@ -182,6 +179,10 @@ def train(model, optimizer, scheduler, criterion, data_loader, eval_data_loader,
             grad_norm, _ = check_update(model, c.grad_clip)
             optimizer.step()
 
+            # setup lr
+            if c.lr_decay:
+                scheduler.step()
+
             step_time = time.time() - start_time
             epoch_time += step_time
 

TTS/encoder/dataset.py

@@ -55,7 +55,7 @@ def __init__(
         logger.info(" | Number of instances: %d", len(self.items))
         logger.info(" | Sequence length: %d", self.seq_len)
         logger.info(" | Number of classes: %d", len(self.classes))
-        logger.info(" | Classes: %d", self.classes)
+        logger.info(" | Classes: %s", self.classes)
 
     def load_wav(self, filename):
         audio = self.ap.load_wav(filename, sr=self.ap.sample_rate)

TTS/model.py

@@ -1,5 +1,6 @@
+import os
 from abc import abstractmethod
-from typing import Dict
+from typing import Any, Union
 
 import torch
 from coqpit import Coqpit
@@ -16,15 +17,15 @@ class BaseTrainerModel(TrainerModel):
 
     @staticmethod
     @abstractmethod
-    def init_from_config(config: Coqpit):
+    def init_from_config(config: Coqpit) -> "BaseTrainerModel":
         """Init the model and all its attributes from the given config.
 
         Override this depending on your model.
         """
         ...
 
     @abstractmethod
-    def inference(self, input: torch.Tensor, aux_input={}) -> Dict:
+    def inference(self, input: torch.Tensor, aux_input: dict[str, Any] = {}) -> dict[str, Any]:
         """Forward pass for inference.
 
         It must return a dictionary with the main model output and all the auxiliary outputs. The key ```model_outputs```
@@ -45,13 +46,18 @@ def inference(self, input: torch.Tensor, aux_input={}) -> Dict:
 
     @abstractmethod
     def load_checkpoint(
-        self, config: Coqpit, checkpoint_path: str, eval: bool = False, strict: bool = True, cache=False
+        self,
+        config: Coqpit,
+        checkpoint_path: Union[str, os.PathLike[Any]],
+        eval: bool = False,
+        strict: bool = True,
+        cache: bool = False,
     ) -> None:
-        """Load a model checkpoint gile and get ready for training or inference.
+        """Load a model checkpoint file and get ready for training or inference.
 
         Args:
             config (Coqpit): Model configuration.
-            checkpoint_path (str): Path to the model checkpoint file.
+            checkpoint_path (str | os.PathLike): Path to the model checkpoint file.
             eval (bool, optional): If true, init model for inference else for training. Defaults to False.
             strict (bool, optional): Match all checkpoint keys to model's keys. Defaults to True.
             cache (bool, optional): If True, cache the file locally for subsequent calls. It is cached under `get_user_data_dir()/tts_cache`. Defaults to False.

TTS/tts/layers/glow_tts/transformer.py

@@ -5,6 +5,7 @@
 from torch.nn import functional as F
 
 from TTS.tts.layers.generic.normalization import LayerNorm, LayerNorm2
+from TTS.tts.utils.helpers import convert_pad_shape
 
 
 class RelativePositionMultiHeadAttention(nn.Module):
@@ -300,7 +301,7 @@ def _causal_padding(self, x):
         pad_l = self.kernel_size - 1
         pad_r = 0
         padding = [[0, 0], [0, 0], [pad_l, pad_r]]
-        x = F.pad(x, self._pad_shape(padding))
+        x = F.pad(x, convert_pad_shape(padding))
         return x
 
     def _same_padding(self, x):
@@ -309,15 +310,9 @@ def _same_padding(self, x):
         pad_l = (self.kernel_size - 1) // 2
         pad_r = self.kernel_size // 2
         padding = [[0, 0], [0, 0], [pad_l, pad_r]]
-        x = F.pad(x, self._pad_shape(padding))
+        x = F.pad(x, convert_pad_shape(padding))
         return x
 
-    @staticmethod
-    def _pad_shape(padding):
-        l = padding[::-1]
-        pad_shape = [item for sublist in l for item in sublist]
-        return pad_shape
-
 
 class RelativePositionTransformer(nn.Module):
     """Transformer with Relative Potional Encoding.

TTS/tts/layers/losses.py

@@ -255,7 +255,7 @@ def forward(self, att_ws, ilens, olens):
 
     @staticmethod
     def _make_ga_mask(ilen, olen, sigma):
-        grid_x, grid_y = torch.meshgrid(torch.arange(olen).to(olen), torch.arange(ilen).to(ilen))
+        grid_x, grid_y = torch.meshgrid(torch.arange(olen).to(olen), torch.arange(ilen).to(ilen), indexing="ij")
         grid_x, grid_y = grid_x.float(), grid_y.float()
         return 1.0 - torch.exp(-((grid_y / ilen - grid_x / olen) ** 2) / (2 * (sigma**2)))
 

TTS/tts/layers/overflow/neural_hmm.py

@@ -128,7 +128,8 @@ def forward(self, inputs, inputs_len, mels, mel_lens):
             # 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)
+                # TODO: use_reentrant=False is recommended
+                mean, std, transition_vector = checkpoint(self.output_net, h_memory, inputs, use_reentrant=True)
             else:
                 mean, std, transition_vector = self.output_net(h_memory, inputs)
 

TTS/tts/layers/vits/networks.py

@@ -10,22 +10,6 @@
 LRELU_SLOPE = 0.1
 
 
-def convert_pad_shape(pad_shape):
-    l = pad_shape[::-1]
-    pad_shape = [item for sublist in l for item in sublist]
-    return pad_shape
-
-
-def init_weights(m, mean=0.0, std=0.01):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        m.weight.data.normal_(mean, std)
-
-
-def get_padding(kernel_size, dilation=1):
-    return int((kernel_size * dilation - dilation) / 2)
-
-
 class TextEncoder(nn.Module):
     def __init__(
         self,

TTS/tts/layers/xtts/hifigan_decoder.py

@@ -9,16 +9,13 @@
 from torch.nn.utils.parametrize import remove_parametrizations
 
 from TTS.utils.io import load_fsspec
+from TTS.vocoder.models.hifigan_generator import get_padding
 
 logger = logging.getLogger(__name__)
 
 LRELU_SLOPE = 0.1
 
 
-def get_padding(k, d):
-    return int((k * d - d) / 2)
-
-
 class ResBlock1(torch.nn.Module):
     """Residual Block Type 1. It has 3 convolutional layers in each convolutional block.
 

TTS/tts/models/base_tts.py

@@ -144,7 +144,7 @@ def get_aux_input_from_test_sentences(self, sentence_info):
                 if speaker_name is None:
                     d_vector = self.speaker_manager.get_random_embedding()
                 else:
-                    d_vector = self.speaker_manager.get_d_vector_by_name(speaker_name)
+                    d_vector = self.speaker_manager.get_mean_embedding(speaker_name)
             elif config.use_speaker_embedding:
                 if speaker_name is None:
                     speaker_id = self.speaker_manager.get_random_id()

TTS/tts/models/delightful_tts.py

@@ -88,12 +88,6 @@ def pad(input_ele: List[torch.Tensor], max_len: int) -> torch.Tensor:
     return out_padded
 
 
-def init_weights(m: nn.Module, mean: float = 0.0, std: float = 0.01):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        m.weight.data.normal_(mean, std)
-
-
 def stride_lens(lens: torch.Tensor, stride: int = 2) -> torch.Tensor:
     return torch.ceil(lens / stride).int()
 

TTS/tts/utils/helpers.py

@@ -145,10 +145,9 @@ def average_over_durations(values, durs):
     return avg
 
 
-def convert_pad_shape(pad_shape):
+def convert_pad_shape(pad_shape: list[list]) -> list:
     l = pad_shape[::-1]
-    pad_shape = [item for sublist in l for item in sublist]
-    return pad_shape
+    return [item for sublist in l for item in sublist]
 
 
 def generate_path(duration, mask):

TTS/vc/models/base_vc.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import random
-from typing import Dict, List, Tuple, Union
+from typing import Any, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -10,6 +10,7 @@
 from torch.utils.data import DataLoader
 from torch.utils.data.sampler import WeightedRandomSampler
 from trainer.torch import DistributedSampler, DistributedSamplerWrapper
+from trainer.trainer import Trainer
 
 from TTS.model import BaseTrainerModel
 from TTS.tts.datasets.dataset import TTSDataset
@@ -18,6 +19,7 @@
 from TTS.tts.utils.speakers import SpeakerManager, get_speaker_balancer_weights
 from TTS.tts.utils.synthesis import synthesis
 from TTS.tts.utils.visual import plot_alignment, plot_spectrogram
+from TTS.utils.audio.processor import AudioProcessor
 
 # pylint: skip-file
 
@@ -35,18 +37,18 @@ class BaseVC(BaseTrainerModel):
     def __init__(
         self,
         config: Coqpit,
-        ap: "AudioProcessor",
-        speaker_manager: SpeakerManager = None,
-        language_manager: LanguageManager = None,
-    ):
+        ap: AudioProcessor,
+        speaker_manager: Optional[SpeakerManager] = None,
+        language_manager: Optional[LanguageManager] = None,
+    ) -> None:
         super().__init__()
         self.config = config
         self.ap = ap
         self.speaker_manager = speaker_manager
         self.language_manager = language_manager
         self._set_model_args(config)
 
-    def _set_model_args(self, config: Coqpit):
+    def _set_model_args(self, config: Coqpit) -> None:
         """Setup model args based on the config type (`ModelConfig` or `ModelArgs`).
 
         `ModelArgs` has all the fields reuqired to initialize the model architecture.
@@ -67,7 +69,7 @@ def _set_model_args(self, config: Coqpit):
         else:
             raise ValueError("config must be either a *Config or *Args")
 
-    def init_multispeaker(self, config: Coqpit, data: List = None):
+    def init_multispeaker(self, config: Coqpit, data: Optional[list[Any]] = None) -> None:
         """Initialize a speaker embedding layer if needen and define expected embedding channel size for defining
         `in_channels` size of the connected layers.
 
@@ -100,11 +102,11 @@ def init_multispeaker(self, config: Coqpit, data: List = None):
             self.speaker_embedding = nn.Embedding(self.num_speakers, self.embedded_speaker_dim)
             self.speaker_embedding.weight.data.normal_(0, 0.3)
 
-    def get_aux_input(self, **kwargs) -> Dict:
+    def get_aux_input(self, **kwargs: Any) -> dict[str, Any]:
         """Prepare and return `aux_input` used by `forward()`"""
         return {"speaker_id": None, "style_wav": None, "d_vector": None, "language_id": None}
 
-    def get_aux_input_from_test_sentences(self, sentence_info):
+    def get_aux_input_from_test_sentences(self, sentence_info: Union[str, list[str]]) -> dict[str, Any]:
         if hasattr(self.config, "model_args"):
             config = self.config.model_args
         else:
@@ -132,7 +134,7 @@ def get_aux_input_from_test_sentences(self, sentence_info):
                 if speaker_name is None:
                     d_vector = self.speaker_manager.get_random_embedding()
                 else:
-                    d_vector = self.speaker_manager.get_d_vector_by_name(speaker_name)
+                    d_vector = self.speaker_manager.get_mean_embedding(speaker_name)
             elif config.use_speaker_embedding:
                 if speaker_name is None:
                     speaker_id = self.speaker_manager.get_random_id()
@@ -151,16 +153,16 @@ def get_aux_input_from_test_sentences(self, sentence_info):
             "language_id": language_id,
         }
 
-    def format_batch(self, batch: Dict) -> Dict:
+    def format_batch(self, batch: dict[str, Any]) -> dict[str, Any]:
         """Generic batch formatting for `VCDataset`.
 
         You must override this if you use a custom dataset.
 
         Args:
-            batch (Dict): [description]
+            batch (dict): [description]
 
         Returns:
-            Dict: [description]
+            dict: [description]
         """
         # setup input batch
         text_input = batch["token_id"]
@@ -230,7 +232,7 @@ def format_batch(self, batch: Dict) -> Dict:
             "audio_unique_names": batch["audio_unique_names"],
         }
 
-    def get_sampler(self, config: Coqpit, dataset: TTSDataset, num_gpus=1):
+    def get_sampler(self, config: Coqpit, dataset: TTSDataset, num_gpus: int = 1):
         weights = None
         data_items = dataset.samples
 
@@ -271,12 +273,12 @@ def get_sampler(self, config: Coqpit, dataset: TTSDataset, num_gpus=1):
     def get_data_loader(
         self,
         config: Coqpit,
-        assets: Dict,
+        assets: dict,
         is_eval: bool,
-        samples: Union[List[Dict], List[List]],
+        samples: Union[list[dict], list[list]],
         verbose: bool,
         num_gpus: int,
-        rank: int = None,
+        rank: Optional[int] = None,
     ) -> "DataLoader":
         if is_eval and not config.run_eval:
             loader = None
@@ -352,9 +354,9 @@ def get_data_loader(
 
     def _get_test_aux_input(
         self,
-    ) -> Dict:
+    ) -> dict[str, Any]:
         d_vector = None
-        if self.config.use_d_vector_file:
+        if self.speaker_manager is not None and self.config.use_d_vector_file:
             d_vector = [self.speaker_manager.embeddings[name]["embedding"] for name in self.speaker_manager.embeddings]
             d_vector = (random.sample(sorted(d_vector), 1),)
 
@@ -369,7 +371,7 @@ def _get_test_aux_input(
         }
         return aux_inputs
 
-    def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
+    def test_run(self, assets: dict) -> tuple[dict, dict]:
         """Generic test run for `vc` models used by `Trainer`.
 
         You can override this for a different behaviour.
@@ -378,7 +380,7 @@ def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
             assets (dict): A dict of training assets. For `vc` models, it must include `{'audio_processor': ap}`.
 
         Returns:
-            Tuple[Dict, Dict]: Test figures and audios to be projected to Tensorboard.
+            tuple[dict, dict]: Test figures and audios to be projected to Tensorboard.
         """
         logger.info("Synthesizing test sentences.")
         test_audios = {}
@@ -409,7 +411,7 @@ def test_run(self, assets: Dict) -> Tuple[Dict, Dict]:
             )
         return test_figures, test_audios
 
-    def on_init_start(self, trainer):
+    def on_init_start(self, trainer: Trainer) -> None:
         """Save the speaker.pth and language_ids.json at the beginning of the training. Also update both paths."""
         if self.speaker_manager is not None:
             output_path = os.path.join(trainer.output_path, "speakers.pth")

TTS/vc/models/freevc.py

@@ -14,14 +14,16 @@
 
 import TTS.vc.modules.freevc.commons as commons
 import TTS.vc.modules.freevc.modules as modules
+from TTS.tts.utils.helpers import sequence_mask
 from TTS.tts.utils.speakers import SpeakerManager
 from TTS.utils.io import load_fsspec
 from TTS.vc.configs.freevc_config import FreeVCConfig
 from TTS.vc.models.base_vc import BaseVC
-from TTS.vc.modules.freevc.commons import get_padding, init_weights
+from TTS.vc.modules.freevc.commons import init_weights
 from TTS.vc.modules.freevc.mel_processing import mel_spectrogram_torch
 from TTS.vc.modules.freevc.speaker_encoder.speaker_encoder import SpeakerEncoder as SpeakerEncoderEx
 from TTS.vc.modules.freevc.wavlm import get_wavlm
+from TTS.vocoder.models.hifigan_generator import get_padding
 
 logger = logging.getLogger(__name__)
 
@@ -80,7 +82,7 @@ def __init__(
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
     def forward(self, x, x_lengths, g=None):
-        x_mask = torch.unsqueeze(commons.sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
+        x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype)
         x = self.pre(x) * x_mask
         x = self.enc(x, x_mask, g=g)
         stats = self.proj(x) * x_mask