Merge branch 'move-resampling-to-network'

src/torch_wae/cli/export_to_onnx.py

@@ -4,9 +4,8 @@
 
 import torch
 import typer
-from torchaudio import functional as F
 
-from torch_wae.network import WAENet
+from torch_wae.network import WAENet, WithResample
 
 app = typer.Typer()
 
@@ -51,18 +50,5 @@ def main(
     )
 
 
-class WithResample(torch.nn.Module):
-    def __init__(self, f: WAENet, sample_rate: int) -> None:
-        super().__init__()
-
-        self.f = f
-        self.sample_rate = sample_rate
-
-    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
-        h = F.resample(waveform, self.sample_rate, self.f.SAMPLE_RATE)
-        z = self.f(h)
-        return z
-
-
 if __name__ == "__main__":
     app()

src/torch_wae/network.py

@@ -4,6 +4,7 @@
 from convmelspec.stft import ConvertibleSpectrogram as Spectrogram
 from torch import nn
 from torch.nn import functional as F
+from torchaudio import functional as FA
 
 
 # Wowrd Audio Encoder - A network for audio similar to MobileNet V2 for images.
@@ -137,3 +138,16 @@ def __init__(self, dim: int = 1, eps: float = 1e-12):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return F.normalize(x, p=2, dim=self.dim, eps=self.eps)
+
+
+class WithResample(torch.nn.Module):
+    def __init__(self, f: WAENet, sample_rate: int) -> None:
+        super().__init__()
+
+        self.f = f
+        self.sample_rate = sample_rate
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        h = FA.resample(waveform, self.sample_rate, self.f.SAMPLE_RATE)
+        z = self.f(h)
+        return z