WIP: demucs_fix

README.md

@@ -10,10 +10,13 @@
 `./scripts/download_egfx.sh`
 
 ## Train model
-1. Change Wandb variables in `shell_vars.sh`
+1. Change Wandb variables in `shell_vars.sh` and `source shell_vars.sh`
 2. `python scripts/train.py exp=audio_diffusion`
 or
 2. `python scripts/train.py exp=umx`
+or
+2. `python scripts/train.py exp=demucs`
+
 
 To add gpu, add `trainer.accelerator='gpu' trainer.devices=-1` to the command-line
 

exp/demucs.yaml

@@ -9,7 +9,7 @@ model:
   sample_rate: ${sample_rate}
   network:
     _target_: remfx.models.DemucsModel
-    sources: ["other"]
+    sources: ["mixture"]
     audio_channels: 1
     nfft: 4096
     sample_rate: ${sample_rate}

remfx/models.py

@@ -86,7 +86,7 @@ class RemFXModel(pl.LightningModule):
 
         return loss
 
-    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+    def on_train_batch_start(self, batch, batch_idx):
         if self.log_first:
             x, target, label = batch
             for metric in self.metrics: