update

toolbox/torchaudio/models/nx_mpnet/discriminator.py

@@ -5,7 +5,10 @@ from typing import Optional, Union
 
 import torch
 import torch.nn as nn
-import torchaudio
+import numpy as np
+import torch.nn.functional as F
+from pesq import pesq
+from joblib import Parallel, delayed
 
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.nx_mpnet.configuration_nx_mpnet import NXMPNetConfig
@@ -16,23 +19,10 @@ class MetricDiscriminator(nn.Module):
     def __init__(self, config: NXMPNetConfig):
         super(MetricDiscriminator, self).__init__()
         dim = config.discriminator_dim
-        self.in_channel = config.discriminator_in_channel
-
-        self.n_fft = config.n_fft
-        self.win_length = config.win_length
-        self.hop_length = config.hop_length
-
-        self.transform = torchaudio.transforms.Spectrogram(
-            n_fft=self.n_fft,
-            win_length=self.win_length,
-            hop_length=self.hop_length,
-            power=1.0,
-            window_fn=torch.hann_window,
-            # window_fn=torch.hamming_window if window_fn == "hamming" else torch.hann_window,
-        )
+        in_channel = config.discriminator_in_channel
 
         self.layers = nn.Sequential(
-            nn.utils.spectral_norm(nn.Conv2d(self.in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
+            nn.utils.spectral_norm(nn.Conv2d(in_channel, dim, (4,4), (2,2), (1,1), bias=False)),
             nn.InstanceNorm2d(dim, affine=True),
             nn.PReLU(dim),
             nn.utils.spectral_norm(nn.Conv2d(dim, dim*2, (4,4), (2,2), (1,1), bias=False)),
@@ -54,9 +44,6 @@ class MetricDiscriminator(nn.Module):
         )
 
     def forward(self, x, y):
-        x = self.transform.forward(x)
-        y = self.transform.forward(y)
-
         xy = torch.stack((x, y), dim=1)
         return self.layers(xy)
 
@@ -111,22 +98,5 @@ class MetricDiscriminatorPretrainedModel(MetricDiscriminator):
         return save_directory
 
 
-def main():
-    config = NXMPNetConfig()
-    discriminator = MetricDiscriminator(config=config)
-
-    # shape: [batch_size, num_samples]
-    # x = torch.ones([4, int(4.5 * 16000)])
-    # y = torch.ones([4, int(4.5 * 16000)])
-    x = torch.ones([4, 16000])
-    y = torch.ones([4, 16000])
-
-    output = discriminator.forward(x, y)
-    print(output.shape)
-    print(output)
-
-    return
-
-
-if __name__ == "__main__":
-    main()
+if __name__ == '__main__':
+    pass