fix noise

WaveGenAI · Jul 2, 2024 · f8ce00f · f8ce00f
1 parent 1ab490b
commit f8ce00f
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/audioenhancer/dataset/loader.py b/audioenhancer/dataset/loader.py
@@ -140,13 +140,14 @@ def __getitem__(self, index: int) -> tuple:
         base_waveform = base_waveform.transpose(0, 1).cuda()
 
         if random.random() < self._prob:
-            strength = torch.rand(compressed_waveform.shape[:2]) * 1.5
-            noise = torch.randn_like(compressed_waveform) * 0.2
-            compressed_waveform = torchaudio.functional.add_noise(
-                compressed_waveform,
-                noise=noise.cuda(),
-                snr=strength.cuda(),
-            )
+            strength = torch.rand(compressed_waveform.shape[:2]) * 0.01
+            strength_expanded = strength.unsqueeze(2).expand(-1, -1, compressed_waveform.shape[2]).cuda()
+            noise = torch.randn_like(compressed_waveform).cuda()
+            compressed_waveform = compressed_waveform + noise * strength_expanded
+
+        if self._mono:
+            compressed_waveform = compressed_waveform.mean(dim=1)
+            base_waveform = base_waveform.mean(dim=1)
 
         encoded_compressed_waveform, _, _, _, _ = self.autoencoder.encode(
             compressed_waveform