remove transcript caches

Author: Flux9665

InferenceInterfaces/ToucanTTSInterface.py

@@ -166,7 +166,7 @@ def forward(self,
                                                            pause_duration_scaling_factor=pause_duration_scaling_factor,
                                                            prosody_creativity=prosody_creativity)
 
-            wave, _, _ = self.vocoder(mel.unsqueeze(0))
+            wave = self.vocoder(mel.unsqueeze(0))
             wave = wave.squeeze().cpu()
         wave = wave.numpy()
         sr = 24000

Modules/Vocoder/HiFiGAN_Discriminators.py

@@ -9,8 +9,6 @@
 import torch
 import torch.nn.functional as F
 
-from Modules.Vocoder.Avocodo_Discriminators import MultiCoMBDiscriminator
-from Modules.Vocoder.Avocodo_Discriminators import MultiSubBandDiscriminator
 from Modules.Vocoder.SAN_modules import SANConv1d
 from Modules.Vocoder.SAN_modules import SANConv2d
 
@@ -456,10 +454,13 @@ def forward(self, x):
 
 
 class AvocodoHiFiGANJointDiscriminator(torch.nn.Module):
+    """
+    Contradicting the legacy name, the Avocodo parts were removed again for stability
+    """
 
     def __init__(self,
                  # Multi-scale discriminator related
-                 scales=3,
+                 scales=4,
                  scale_downsample_pooling="AvgPool1d",
                  scale_downsample_pooling_params={"kernel_size": 4,
                                                   "stride"     : 2,
@@ -471,7 +472,7 @@ def __init__(self,
                                              "max_downsample_channels"    : 1024,
                                              "max_groups"                 : 16,
                                              "bias"                       : True,
-                                             "downsample_scales"          : [4, 4, 4, 4, 1],
+                                             "downsample_scales"          : [4, 4, 4, 1],
                                              "nonlinear_activation"       : "LeakyReLU",
                                              "nonlinear_activation_params": {"negative_slope": 0.1}, },
                  follow_official_norm=True,
@@ -481,41 +482,14 @@ def __init__(self,
                                               "out_channels"               : 1,
                                               "kernel_sizes"               : [5, 3],
                                               "channels"                   : 32,
-                                              "downsample_scales"          : [3, 3, 3, 3, 1],
+                                              "downsample_scales"          : [3, 3, 3, 1],
                                               "max_downsample_channels"    : 1024,
                                               "bias"                       : True,
                                               "nonlinear_activation"       : "LeakyReLU",
                                               "nonlinear_activation_params": {"negative_slope": 0.1},
                                               "use_weight_norm"            : True,
                                               "use_spectral_norm"          : False, },
-                 # CoMB discriminator related
-                 kernels=((7, 11, 11, 11, 11, 5),
-                          (11, 21, 21, 21, 21, 5),
-                          (15, 41, 41, 41, 41, 5)),
-                 channels=(16, 64, 256, 1024, 1024, 1024),
-                 groups=(1, 4, 16, 64, 256, 1),
-                 strides=(1, 1, 4, 4, 4, 1),
-                 # Sub-Band discriminator related
-                 tkernels=(7, 5, 3),
-                 fkernel=5,
-                 tchannels=(64, 128, 256, 256, 256),
-                 fchannels=(32, 64, 128, 128, 128),
-                 tstrides=((1, 1, 3, 3, 1),
-                           (1, 1, 3, 3, 1),
-                           (1, 1, 3, 3, 1)),
-                 fstride=(1, 1, 3, 3, 1),
-                 tdilations=(((5, 7, 11), (5, 7, 11), (5, 7, 11), (5, 7, 11), (5, 7, 11), (5, 7, 11)),
-                             ((3, 5, 7), (3, 5, 7), (3, 5, 7), (3, 5, 7), (3, 5, 7)),
-                             ((1, 2, 3), (1, 2, 3), (1, 2, 3), (1, 2, 3), (1, 2, 3))),
-                 fdilations=((1, 2, 3),
-                             (1, 2, 3),
-                             (1, 2, 3),
-                             (2, 3, 5),
-                             (2, 3, 5)),
-                 tsubband=(6, 11, 16),
-                 n=16,
-                 m=64,
-                 freq_init_ch=192):
+                 ):
         super().__init__()
         self.msd = HiFiGANMultiScaleDiscriminator(scales=scales,
                                                   downsample_pooling=scale_downsample_pooling,
@@ -524,10 +498,8 @@ def __init__(self,
                                                   follow_official_norm=follow_official_norm, )
         self.mpd = HiFiGANMultiPeriodDiscriminator(periods=periods,
                                                    discriminator_params=period_discriminator_params, )
-        self.mcmbd = MultiCoMBDiscriminator(kernels, channels, groups, strides)
-        self.msbd = MultiSubBandDiscriminator(tkernels, fkernel, tchannels, fchannels, tstrides, fstride, tdilations, fdilations, tsubband, n, m, freq_init_ch)
 
-    def forward(self, wave, intermediate_wave_upsampled_twice=None, intermediate_wave_upsampled_once=None, discriminator_train_flag=False):
+    def forward(self, wave, discriminator_train_flag=False):
         """
         Calculate forward propagation.
 
@@ -542,9 +514,9 @@ def forward(self, wave, intermediate_wave_upsampled_twice=None, intermediate_wav
         """
         msd_outs, msd_feats = self.msd(wave, discriminator_train_flag)
         mpd_outs, mpd_feats = self.mpd(wave, discriminator_train_flag)
-        mcmbd_outs, mcmbd_feats = self.mcmbd(wave_final=wave,
-                                             intermediate_wave_upsampled_twice=intermediate_wave_upsampled_twice,
-                                             intermediate_wave_upsampled_once=intermediate_wave_upsampled_once,
-                                             discriminator_train_flag=discriminator_train_flag)
-        msbd_outs, msbd_feats = self.msbd(wave, discriminator_train_flag)
-        return msd_outs + mpd_outs + mcmbd_outs + msbd_outs, msd_feats + mpd_feats + mcmbd_feats + msbd_feats
+        return msd_outs + mpd_outs, msd_feats + mpd_feats
+
+
+if __name__ == '__main__':
+    d = AvocodoHiFiGANJointDiscriminator()
+    print(d(torch.randn([2, 1, 12288 * 2])))

Modules/Vocoder/HiFiGAN_E2E_Dataset.py

@@ -1,10 +1,9 @@
-import os
-import random
 from multiprocessing import Manager
 from multiprocessing import Process
 
 import librosa
 import numpy
+import numpy as np
 import soundfile as sf
 import torch
 from torch.utils.data import Dataset
@@ -19,8 +18,8 @@ def __init__(self,
                  list_of_original_paths,
                  list_of_synthetic_paths,
                  desired_samplingrate=24000,
-                 samples_per_segment=12288,  # = (8192 * 3) 2 , as I used 8192 for 16kHz previously
-                 loading_processes=max(os.cpu_count() - 2, 1)):
+                 samples_per_segment=12288 * 2,  # = (8192 * 3) 2 , as I used 8192 for 16kHz previously
+                 loading_processes=1):
         self.samples_per_segment = samples_per_segment
         self.desired_samplingrate = desired_samplingrate
         self.melspec_ap = AudioPreprocessor(input_sr=self.desired_samplingrate,
@@ -53,19 +52,18 @@ def cache_builder_process(self, path_split):
             try:
                 path1, path2 = path
 
-                wave1, sr = sf.read(path1)
-                if len(wave1.shape) == 2:
-                    wave1 = librosa.to_mono(numpy.transpose(wave1))
+                wave, sr = sf.read(path1)
+                if len(wave.shape) == 2:
+                    wave = librosa.to_mono(numpy.transpose(wave))
                 if sr != self.desired_samplingrate:
-                    wave1 = librosa.resample(y=wave1, orig_sr=sr, target_sr=self.desired_samplingrate)
+                    wave = librosa.resample(y=wave, orig_sr=sr, target_sr=self.desired_samplingrate)
 
-                wave2, sr = sf.read(path2)
-                if len(wave2.shape) == 2:
-                    wave2 = librosa.to_mono(numpy.transpose(wave2))
-                if sr != self.desired_samplingrate:
-                    wave2 = librosa.resample(y=wave2, orig_sr=sr, target_sr=self.desired_samplingrate)
+                if len(wave) > self.samples_per_segment + 2000:
+                    spec = torch.load(path2, map_location="cpu")
+                    self.waves.append((wave, spec))
+                else:
+                    print("excluding short sample")
 
-                self.waves.append((wave1, wave2))
             except RuntimeError:
                 print(f"Problem with the following path: {path}")
 
@@ -77,32 +75,68 @@ def __getitem__(self, index):
 
         return a pair of high-res audio and corresponding low-res spectrogram as if it was predicted by the TTS
         """
-        try:
-            wave1 = self.waves[index][0]
-            wave2 = self.waves[index][1]
-            while len(wave1) < self.samples_per_segment + 50:  # + 50 is just to be extra sure
-                # catch files that are too short to apply meaningful signal processing and make them longer
-                wave1 = numpy.concatenate([wave1, numpy.zeros(shape=1000), wave1])
-                wave2 = numpy.concatenate([wave2, numpy.zeros(shape=1000), wave2])
-                # add some true silence in the mix, so the vocoder is exposed to that as well during training
-            wave1 = torch.Tensor(wave1)
-            wave2 = torch.Tensor(wave2)
-
-            max_audio_start = len(wave1) - self.samples_per_segment
-            audio_start = random.randint(0, max_audio_start)
-            segment1 = wave1[audio_start: audio_start + self.samples_per_segment]
-            segment2 = wave2[audio_start: audio_start + self.samples_per_segment]
-
-            resampled_segment = self.melspec_ap.resample(segment2).float()  # 16kHz spectrogram as input, 24kHz wave as output, see Blizzard 2021 DelightfulTTS
-            melspec = self.melspec_ap.audio_to_mel_spec_tensor(resampled_segment,
-                                                               explicit_sampling_rate=16000,
-                                                               normalize=False).transpose(0, 1)[:-1].transpose(0, 1)
-            return segment1.detach(), melspec.detach()
-        except RuntimeError:
-            print("encountered a runtime error, using fallback strategy")
-            if index == 0:
-                index = len(self.waves) - 1
-            return self.__getitem__(index - 1)
+        wave = self.waves[index][0]
+        wave = torch.Tensor(wave)
+
+        spec = self.waves[index][1]
+
+        spec_win, wave_win = get_matching_windows(waveform=wave, spectrogram=spec)
+        return wave_win.detach(), spec_win.detach()
 
     def __len__(self):
         return len(self.waves)
+
+
+def get_matching_windows(spectrogram, waveform, window_size_wave=24576, hop_length_spec=256, sample_rate_wave=24000, sample_rate_spec=16000):
+    """
+    Cut random matching windows from a spectrogram and waveform with perfectly aligned time axes.
+
+    Parameters:
+    - spectrogram: 2D numpy array (frames x freq_bins) of the spectrogram.
+    - waveform: 1D numpy array of the ground truth waveform.
+    - window_size_wave: Size of the window in waveform samples (default: 24576).
+    - hop_length_spec: Hop length used for spectrogram extraction (default: 200 samples for 16 kHz).
+    - sample_rate_wave: Sample rate of the waveform (default: 24000 Hz).
+    - sample_rate_spec: Sample rate used to create the spectrogram (default: 16000 Hz).
+
+    Returns:
+    - spec_window: A window cut from the spectrogram.
+    - wave_window: A window cut from the waveform.
+    """
+    spectrogram = spectrogram.transpose(0, 1)
+
+    # Calculate the number of samples per spectrogram frame in waveform's time
+    spec_frame_duration = hop_length_spec / sample_rate_spec
+    wave_sample_duration = 1 / sample_rate_wave
+    spec_to_wave_conversion_factor = wave_sample_duration / spec_frame_duration
+
+    num_frames = int(window_size_wave * spec_to_wave_conversion_factor)
+
+    # Ensure we can extract a full window from the spectrogram
+    max_start_frame = spectrogram.shape[0] - num_frames
+    if max_start_frame <= 0:
+        print(f"desired num frames: {num_frames}")
+        print(f"spec_to_wave_conversion_factor: {spec_to_wave_conversion_factor}")
+        print(f"spec_len: {spectrogram.shape[0]}")
+        raise ValueError("Spectrogram is too short to extract the desired window size.")
+
+    # Randomly choose a start frame from the spectrogram
+    start_frame = np.random.randint(0, max_start_frame)
+
+    # Calculate the start sample for the waveform based on the chosen start frame
+    start_sample = int(start_frame // spec_to_wave_conversion_factor)
+    end_sample = start_sample + window_size_wave
+
+    # Ensure the waveform can be fully sliced
+    if end_sample > len(waveform):
+        print(f"start_sample: {start_sample}")
+        print(f"end_sample: {end_sample}")
+        print(f"start_frame: {start_frame}")
+        print(f"spec_to_wave_conversion_factor: {spec_to_wave_conversion_factor}")
+        raise ValueError("Waveform is too short to extract the desired window size.")
+
+    # Extract matching windows
+    spec_window = spectrogram[start_frame:start_frame + num_frames, :].transpose(0, 1)
+    wave_window = waveform[start_sample:end_sample]
+
+    return spec_window, wave_window

Modules/Vocoder/HiFiGAN_Generator.py

@@ -15,10 +15,10 @@ class HiFiGAN(torch.nn.Module):
     def __init__(self,
                  in_channels=128,
                  out_channels=1,
-                 channels=512,
+                 channels=768,
                  kernel_size=7,
-                 upsample_scales=(8, 6, 4, 2),  # CAREFUL: Avocodo assumes that there are always 4 upsample scales, because it takes intermediate results.
-                 upsample_kernel_sizes=(16, 12, 8, 4),
+                 upsample_scales=(8, 6, 2, 2, 2),  # CAREFUL: Avocodo assumes that there are always 4 upsample scales, because it takes intermediate results.
+                 upsample_kernel_sizes=(16, 12, 4, 4, 4),
                  resblock_kernel_sizes=(3, 7, 11),
                  resblock_dilations=((1, 3, 5), (1, 3, 5), (1, 3, 5)),
                  use_additional_convs=True,
@@ -87,9 +87,6 @@ def __init__(self,
                             1,
                             padding=(kernel_size - 1) // 2, ), torch.nn.Tanh(), )
 
-        self.out_proj_x1 = torch.nn.Conv1d(channels // 4, 1, 7, 1, padding=3)
-        self.out_proj_x2 = torch.nn.Conv1d(channels // 8, 1, 7, 1, padding=3)
-
         # apply weight norm
         self.apply_weight_norm()
 
@@ -118,13 +115,9 @@ def forward(self, c):
             for j in range(self.num_blocks):
                 cs += self.blocks[i * self.num_blocks + j](c)
             c = cs / self.num_blocks
-            if i == 1:
-                x1 = self.out_proj_x1(c)
-            elif i == 2:
-                x2 = self.out_proj_x2(c)
         c = self.output_conv(c)
 
-        return c, x2, x1
+        return c
 
     def reset_parameters(self):
         """

Modules/Vocoder/HiFiGAN_train_loop.py

@@ -19,6 +19,10 @@
 from run_weight_averaging import load_net_bigvgan
 
 
+def collate_fn(batch):
+    return torch.stack([x[0] for x in batch]), torch.stack([x[1] for x in batch])
+
+
 def train_loop(generator,
                discriminator,
                train_dataset,
@@ -29,7 +33,7 @@ def train_loop(generator,
                batch_size=32,
                epochs=100,
                resume=False,
-               generator_steps_per_discriminator_step=5,
+               generator_steps_per_discriminator_step=2,
                generator_warmup=30000,
                use_wandb=False,
                finetune=False
@@ -52,11 +56,12 @@ def train_loop(generator,
     train_loader = DataLoader(dataset=train_dataset,
                               batch_size=batch_size,
                               shuffle=True,
-                              num_workers=8,
+                              num_workers=16,
                               pin_memory=True,
                               drop_last=True,
                               prefetch_factor=2,
-                              persistent_workers=True)
+                              persistent_workers=True,
+                              collate_fn=collate_fn)
 
     if resume:
         path_to_checkpoint = get_most_recent_checkpoint(checkpoint_dir=model_save_dir)
@@ -93,18 +98,16 @@ def train_loop(generator,
 
             gold_wave = datapoint[0].to(device).unsqueeze(1)
             melspec = datapoint[1].to(device)
-            pred_wave, intermediate_wave_upsampled_twice, intermediate_wave_upsampled_once = g(melspec)
+            pred_wave = g(melspec)
             if torch.any(torch.isnan(pred_wave)):
                 print("A NaN in the wave! Skipping...")
                 continue
 
             mel_loss = mel_l1(pred_wave.squeeze(1), gold_wave)
-            generator_total_loss = mel_loss * 85.0
+            generator_total_loss = mel_loss * 45.0
 
             if step_counter > generator_warmup + 100:  # a bit of warmup helps, but it's not that important
-                d_outs, d_fmaps = d(wave=pred_wave,
-                                    intermediate_wave_upsampled_twice=intermediate_wave_upsampled_twice,
-                                    intermediate_wave_upsampled_once=intermediate_wave_upsampled_once)
+                d_outs, d_fmaps = d(wave=pred_wave)
                 adversarial_loss = generator_adv_loss(d_outs)
                 adversarial_losses.append(adversarial_loss.item())
                 generator_total_loss = generator_total_loss + adversarial_loss * 2  # based on own experience
@@ -136,8 +139,6 @@ def train_loop(generator,
 
             if step_counter > generator_warmup and step_counter % generator_steps_per_discriminator_step == 0:
                 d_outs, d_fmaps = d(wave=pred_wave.detach(),
-                                    intermediate_wave_upsampled_twice=intermediate_wave_upsampled_twice.detach(),
-                                    intermediate_wave_upsampled_once=intermediate_wave_upsampled_once.detach(),
                                     discriminator_train_flag=True)
                 d_gold_outs, d_gold_fmaps = d(gold_wave,
                                               discriminator_train_flag=True)  # have to recompute unfortunately due to autograd behaviour
@@ -168,7 +169,7 @@ def train_loop(generator,
             g.train()
             delete_old_checkpoints(model_save_dir, keep=5)
 
-            checkpoint_paths = get_n_recent_checkpoints_paths(checkpoint_dir=model_save_dir, n=2)
+            checkpoint_paths = get_n_recent_checkpoints_paths(checkpoint_dir=model_save_dir, n=1)
             averaged_model, _ = average_checkpoints(checkpoint_paths, load_func=load_net_bigvgan)
             torch.save(averaged_model.state_dict(), os.path.join(model_save_dir, "best.pt"))