Make style (#1405)

TTS/encoder/configs/base_encoder_config.py

@@ -33,10 +33,7 @@ class BaseEncoderConfig(BaseTrainingConfig):
     grad_clip: float = 3.0
     lr: float = 0.0001
     optimizer: str = "radam"
-    optimizer_params: Dict = field(default_factory=lambda: {
-        "betas": [0.9, 0.999],
-        "weight_decay": 0
-    })
+    optimizer_params: Dict = field(default_factory=lambda: {"betas": [0.9, 0.999], "weight_decay": 0})
     lr_decay: bool = False
     warmup_steps: int = 4000
 

TTS/encoder/dataset.py

@@ -5,6 +5,7 @@ from torch.utils.data import Dataset
 
 from TTS.encoder.utils.generic_utils import AugmentWAV
 
+
 class EncoderDataset(Dataset):
     def __init__(
         self,
@@ -57,7 +58,6 @@ class EncoderDataset(Dataset):
             print(f" | > Num Classes: {len(self.classes)}")
             print(f" | > Classes: {self.classes}")
 
-
     def load_wav(self, filename):
         audio = self.ap.load_wav(filename, sr=self.ap.sample_rate)
         return audio
@@ -75,9 +75,7 @@ class EncoderDataset(Dataset):
                 ]
 
         # skip classes with number of samples >= self.num_utter_per_class
-        class_to_utters = {
-            k: v for (k, v) in class_to_utters.items() if len(v) >= self.num_utter_per_class
-        }
+        class_to_utters = {k: v for (k, v) in class_to_utters.items() if len(v) >= self.num_utter_per_class}
 
         classes = list(class_to_utters.keys())
         classes.sort()
@@ -105,11 +103,11 @@ class EncoderDataset(Dataset):
 
     def get_class_list(self):
         return self.classes
+
     def set_classes(self, classes):
         self.classes = classes
         self.classname_to_classid = {key: i for i, key in enumerate(self.classes)}
 
-
     def get_map_classid_to_classname(self):
         return dict((c_id, c_n) for c_n, c_id in self.classname_to_classid.items())
 

TTS/encoder/losses.py

@@ -195,6 +195,7 @@ class SoftmaxLoss(nn.Module):
         class_id = torch.argmax(activations)
         return class_id
 
+
 class SoftmaxAngleProtoLoss(nn.Module):
     """
     Implementation of the Softmax AnglePrototypical loss as defined in https://arxiv.org/abs/2009.14153

TTS/encoder/models/base_encoder.py

@@ -1,12 +1,13 @@
+import numpy as np
 import torch
 import torchaudio
-import numpy as np
+from coqpit import Coqpit
 from torch import nn
 
-from TTS.utils.io import load_fsspec
 from TTS.encoder.losses import AngleProtoLoss, GE2ELoss, SoftmaxAngleProtoLoss
 from TTS.utils.generic_utils import set_init_dict
-from coqpit import Coqpit
+from TTS.utils.io import load_fsspec
+
 
 class PreEmphasis(nn.Module):
     def __init__(self, coefficient=0.97):
@@ -20,6 +21,7 @@ class PreEmphasis(nn.Module):
         x = torch.nn.functional.pad(x.unsqueeze(1), (1, 0), "reflect")
         return torch.nn.functional.conv1d(x, self.filter).squeeze(1)
 
+
 class BaseEncoder(nn.Module):
     """Base `encoder` class. Every new `encoder` model must inherit this.
 
@@ -32,31 +34,31 @@ class BaseEncoder(nn.Module):
 
     def get_torch_mel_spectrogram_class(self, audio_config):
         return torch.nn.Sequential(
-                PreEmphasis(audio_config["preemphasis"]),
-                # TorchSTFT(
-                #     n_fft=audio_config["fft_size"],
-                #     hop_length=audio_config["hop_length"],
-                #     win_length=audio_config["win_length"],
-                #     sample_rate=audio_config["sample_rate"],
-                #     window="hamming_window",
-                #     mel_fmin=0.0,
-                #     mel_fmax=None,
-                #     use_htk=True,
-                #     do_amp_to_db=False,
-                #     n_mels=audio_config["num_mels"],
-                #     power=2.0,
-                #     use_mel=True,
-                #     mel_norm=None,
-                # )
-                torchaudio.transforms.MelSpectrogram(
-                    sample_rate=audio_config["sample_rate"],
-                    n_fft=audio_config["fft_size"],
-                    win_length=audio_config["win_length"],
-                    hop_length=audio_config["hop_length"],
-                    window_fn=torch.hamming_window,
-                    n_mels=audio_config["num_mels"],
-                )
-            )
+            PreEmphasis(audio_config["preemphasis"]),
+            # TorchSTFT(
+            #     n_fft=audio_config["fft_size"],
+            #     hop_length=audio_config["hop_length"],
+            #     win_length=audio_config["win_length"],
+            #     sample_rate=audio_config["sample_rate"],
+            #     window="hamming_window",
+            #     mel_fmin=0.0,
+            #     mel_fmax=None,
+            #     use_htk=True,
+            #     do_amp_to_db=False,
+            #     n_mels=audio_config["num_mels"],
+            #     power=2.0,
+            #     use_mel=True,
+            #     mel_norm=None,
+            # )
+            torchaudio.transforms.MelSpectrogram(
+                sample_rate=audio_config["sample_rate"],
+                n_fft=audio_config["fft_size"],
+                win_length=audio_config["win_length"],
+                hop_length=audio_config["hop_length"],
+                window_fn=torch.hamming_window,
+                n_mels=audio_config["num_mels"],
+            ),
+        )
 
     @torch.no_grad()
     def inference(self, x, l2_norm=True):
@@ -104,7 +106,9 @@ class BaseEncoder(nn.Module):
             raise Exception("The %s  not is a loss supported" % c.loss)
         return criterion
 
-    def load_checkpoint(self, config: Coqpit, checkpoint_path: str, eval: bool = False, use_cuda: bool = False, criterion=None):
+    def load_checkpoint(
+        self, config: Coqpit, checkpoint_path: str, eval: bool = False, use_cuda: bool = False, criterion=None
+    ):
         state = load_fsspec(checkpoint_path, map_location=torch.device("cpu"))
         try:
             self.load_state_dict(state["model"])
@@ -127,7 +131,12 @@ class BaseEncoder(nn.Module):
                 print(" > Criterion load ignored because of:", error)
 
         # instance and load the criterion for the encoder classifier in inference time
-        if eval and criterion is None and "criterion" in state and getattr(config, 'map_classid_to_classname', None) is not None:
+        if (
+            eval
+            and criterion is None
+            and "criterion" in state
+            and getattr(config, "map_classid_to_classname", None) is not None
+        ):
             criterion = self.get_criterion(config, len(config.map_classid_to_classname))
             criterion.load_state_dict(state["criterion"])
 

TTS/encoder/models/resnet.py

@@ -4,6 +4,7 @@ from torch import nn
 # from TTS.utils.audio import TorchSTFT
 from TTS.encoder.models.base_encoder import BaseEncoder
 
+
 class SELayer(nn.Module):
     def __init__(self, channel, reduction=8):
         super(SELayer, self).__init__()

TTS/encoder/utils/samplers.py

@@ -1,4 +1,5 @@
 import random
+
 from torch.utils.data.sampler import Sampler, SubsetRandomSampler
 
 
@@ -34,10 +35,21 @@ class PerfectBatchSampler(Sampler):
         drop_last (bool): if True, drops last incomplete batch.
     """
 
-    def __init__(self, dataset_items, classes, batch_size, num_classes_in_batch, num_gpus=1, shuffle=True, drop_last=False, label_key="class_name"):
+    def __init__(
+        self,
+        dataset_items,
+        classes,
+        batch_size,
+        num_classes_in_batch,
+        num_gpus=1,
+        shuffle=True,
+        drop_last=False,
+        label_key="class_name",
+    ):
         super().__init__(dataset_items)
-        assert batch_size % (num_classes_in_batch * num_gpus) == 0, (
-            'Batch size must be divisible by number of classes times the number of data parallel devices (if enabled).')
+        assert (
+            batch_size % (num_classes_in_batch * num_gpus) == 0
+        ), "Batch size must be divisible by number of classes times the number of data parallel devices (if enabled)."
 
         label_indices = {}
         for idx, item in enumerate(dataset_items):
@@ -93,7 +105,7 @@ class PerfectBatchSampler(Sampler):
                 if groups % self._dp_devices == 0:
                     yield batch
                 else:
-                    batch = batch[:(groups // self._dp_devices) * self._dp_devices * self._num_classes_in_batch]
+                    batch = batch[: (groups // self._dp_devices) * self._dp_devices * self._num_classes_in_batch]
                     if len(batch) > 0:
                         yield batch