Add demucs model

README.md

@@ -19,3 +19,20 @@ To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-li
 
 Ex. `python train.py exp=umx trainer.accelerator='gpu' trainer.devices=-1`
 
+### Effects
+Default effect is RAT (distortion). Effect choices:
+- BluesDriver
+- Clean
+- Flanger
+- Phaser
+- RAT
+- Sweep Echo
+- TubeScreamer
+- Chorus
+- Digital Delay
+- Hall Reverb
+- Plate Reverb
+- Spring Reverb
+- TapeEcho
+
+Change effect by adding `+datamodule.dataset.effect_types=["{Effect}"]` to the command-line

exp/demucs.yaml

@@ -1 +1,20 @@
 # @package _global_
+model:
+  _target_: remfx.models.RemFXModel
+  lr: 1e-4
+  lr_beta1: 0.95
+  lr_beta2: 0.999
+  lr_eps: 1e-6
+  lr_weight_decay: 1e-3
+  sample_rate: ${sample_rate}
+  network:
+    _target_: remfx.models.DemucsModel
+    sources: ["other"]
+    audio_channels: 1
+    nfft: 4096
+    sample_rate: ${sample_rate}
+
+
+datamodule:
+  dataset:
+    effect_types: ["RAT"]

remfx/models.py

@@ -9,6 +9,7 @@ from auraloss.freq import MultiResolutionSTFTLoss, STFTLoss
 from torch.nn import L1Loss
 
 from umx.openunmix.model import OpenUnmix, Separator
+from torchaudio.models import HDemucs
 
 
 class RemFXModel(pl.LightningModule):
@@ -148,6 +149,25 @@ class OpenUnmixModel(torch.nn.Module):
         return self.separator(x).squeeze(1)
 
 
+class DemucsModel(torch.nn.Module):
+    def __init__(self, sample_rate, **kwargs) -> None:
+        super().__init__()
+        self.model = HDemucs(**kwargs)
+        self.num_bins = kwargs["nfft"] // 2 + 1
+        self.loss_fn = MultiResolutionSTFTLoss(
+            n_bins=self.num_bins, sample_rate=sample_rate
+        )
+
+    def forward(self, batch):
+        x, target, label = batch
+        output = self.model(x).squeeze(1)
+        loss = self.loss_fn(output, target)
+        return loss, output
+
+    def sample(self, x: Tensor) -> Tensor:
+        return self.model(x).squeeze(1)
+
+
 class DiffusionGenerationModel(nn.Module):
     def __init__(self, n_channels: int = 1):
         super().__init__()

scripts/train.py

@@ -2,6 +2,7 @@ import pytorch_lightning as pl
 import hydra
 from omegaconf import DictConfig
 import remfx.utils as utils
+from pytorch_lightning.utilities.model_summary import ModelSummary
 
 log = utils.get_logger(__name__)
 
@@ -39,6 +40,8 @@ def main(cfg: DictConfig):
         callbacks=callbacks,
         logger=logger,
     )
+    summary = ModelSummary(model)
+    print(summary)
     trainer.fit(model=model, datamodule=datamodule)