Remove umx

remfx/models.py

@@ -6,7 +6,6 @@ from torch import Tensor, nn
 from torchaudio.models import HDemucs
 from auraloss.time import SISDRLoss
 from auraloss.freq import MultiResolutionSTFTLoss
-from umx.openunmix.model import OpenUnmix, Separator
 
 from remfx.utils import spectrogram
 from remfx.tcn import TCN
@@ -256,54 +255,6 @@ class RemFX(pl.LightningModule):
         return loss
 
 
-class OpenUnmixModel(nn.Module):
-    def __init__(
-        self,
-        n_fft: int = 2048,
-        hop_length: int = 512,
-        n_channels: int = 1,
-        alpha: float = 0.3,
-        sample_rate: int = 22050,
-    ):
-        super().__init__()
-        self.n_channels = n_channels
-        self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.alpha = alpha
-        window = torch.hann_window(n_fft)
-        self.register_buffer("window", window)
-
-        self.num_bins = self.n_fft // 2 + 1
-        self.sample_rate = sample_rate
-        self.model = OpenUnmix(
-            nb_channels=self.n_channels,
-            nb_bins=self.num_bins,
-        )
-        self.separator = Separator(
-            target_models={"other": self.model},
-            nb_channels=self.n_channels,
-            sample_rate=self.sample_rate,
-            n_fft=self.n_fft,
-            n_hop=self.hop_length,
-        )
-        self.mrstftloss = MultiResolutionSTFTLoss(
-            n_bins=self.num_bins, sample_rate=self.sample_rate
-        )
-        self.l1loss = nn.L1Loss()
-
-    def forward(self, batch):
-        x, target = batch
-        X = spectrogram(x, self.window, self.n_fft, self.hop_length, self.alpha)
-        Y = self.model(X)
-        sep_out = self.separator(x).squeeze(1)
-        loss = self.mrstftloss(sep_out, target) + self.l1loss(sep_out, target) * 100
-
-        return loss, sep_out
-
-    def sample(self, x: Tensor) -> Tensor:
-        return self.separator(x).squeeze(1)
-
-
 class DemucsModel(nn.Module):
     def __init__(self, sample_rate, **kwargs) -> None:
         super().__init__()