Fix custom inferencing issues

README.md

@@ -47,6 +47,23 @@ see `cfg/exp/default.yaml` for an example.
 - `reverb`
 - `delay`
 
+## Run inference on directory
+Assumes directory is structured as
+- root
+    - clean
+        - file1.wav
+        - file2.wav
+        - file3.wav
+    - effected
+        - file1.wav
+        - file2.wav
+        - file3.wav
+
+Change root path in `shell_vars.sh` and `source shell_vars.sh`
+
+`python scripts/chain_inference.py +exp=chain_inference_custom`
+
+
 ## Misc.
 By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
 

cfg/exp/chain_inference.yaml

@@ -28,10 +28,17 @@ datamodule:
   num_workers: 8
   train_dataset: None
   val_dataset: None
+
 ckpts:
   RandomPedalboardDistortion: "ckpts/distortion.ckpt"
   RandomPedalboardCompressor: "ckpts/compressor.ckpt"
   RandomPedalboardReverb: "ckpts/reverb.ckpt"
   RandomPedalboardChorus: "ckpts/chorus.ckpt"
   RandomPedalboardDelay: "ckpts/delay.ckpt"
+inference_effects_ordering:
+  - "RandomPedalboardDistortion"
+  - "RandomPedalboardCompressor"
+  - "RandomPedalboardReverb"
+  - "RandomPedalboardChorus"
+  - "RandomPedalboardDelay"
 num_bins: 1025

cfg/exp/chain_inference_custom.yaml

@@ -24,13 +24,13 @@ effects_to_remove:
   - chorus
   - delay
 datamodule:
-  batch_size: 16
+  batch_size: 1
   num_workers: 8
   train_dataset: None
   val_dataset: None
   test_dataset:
     _target_: remfx.datasets.InferenceDataset
-    root: "./data/fx-examples"
+    root: ${oc.env:DATASET_ROOT}
     sample_rate: ${sample_rate}
 ckpts:
   RandomPedalboardDistortion: "ckpts/distortion.ckpt"
@@ -38,4 +38,10 @@ ckpts:
   RandomPedalboardReverb: "ckpts/reverb.ckpt"
   RandomPedalboardChorus: "ckpts/chorus.ckpt"
   RandomPedalboardDelay: "ckpts/delay.ckpt"
+inference_effects_ordering:
+  - "RandomPedalboardDistortion"
+  - "RandomPedalboardCompressor"
+  - "RandomPedalboardReverb"
+  - "RandomPedalboardChorus"
+  - "RandomPedalboardDelay"
 num_bins: 1025

remfx/datasets.py

@@ -361,14 +361,14 @@ class EffectDataset(Dataset):
 
 
 class InferenceDataset(Dataset):
-    def __init__(self, root: str, sample_rate: int):
+    def __init__(self, root: str, sample_rate: int, **kwargs):
         self.root = Path(root)
         self.sample_rate = sample_rate
-        self.clean_paths = list(self.root.glob("clean/*.wav"))
-        self.effected_paths = list(self.root.glob("effected/*.wav"))
+        self.clean_paths = sorted(list(self.root.glob("clean/*.wav")))
+        self.effected_paths = sorted(list(self.root.glob("effected/*.wav")))
 
     def __len__(self) -> int:
-        return len(self.audio_paths)
+        return len(self.clean_paths)
 
     def __getitem__(self, idx: int) -> torch.Tensor:
         clean_path = self.clean_paths[idx]
@@ -379,21 +379,20 @@ class InferenceDataset(Dataset):
         effected = torchaudio.functional.resample(effected_audio, sr, self.sample_rate)
 
         # Sum to mono
-        clean = torch.sum(clean, dim=0)
-        effected = torch.sum(effected, dim=0)
+        clean = torch.sum(clean, dim=0, keepdim=True)
+        effected = torch.sum(effected, dim=0, keepdim=True)
 
         # Pad or trim effected to clean
-        if len(clean) > len(effected):
-            effected = torch.nn.functional.pad(
-                effected, (0, len(clean) - len(effected))
-            )
-        elif len(effected) > len(clean):
-            effected = effected[: len(clean)]
+        if effected.shape[1] > clean.shape[1]:
+            effected = effected[:, : clean.shape[1]]
+        elif effected.shape[1] < clean.shape[1]:
+            pad_size = clean.shape[1] - effected.shape[1]
+            effected = torch.nn.functional.pad(effected, (0, pad_size))
 
         dry_labels_tensor = torch.zeros(len(ALL_EFFECTS))
         wet_labels_tensor = torch.ones(len(ALL_EFFECTS))
 
-        return clean, effected, dry_labels_tensor, wet_labels_tensor
+        return effected, clean, dry_labels_tensor, wet_labels_tensor
 
 
 class EffectDatamodule(pl.LightningDataModule):

remfx/models.py

@@ -37,7 +37,7 @@ class RemFXChainInference(pl.LightningModule):
         self.sample_rate = sample_rate
         self.effect_order = effect_order
 
-    def forward(self, batch, order=None):
+    def forward(self, batch, batch_idx, order=None):
         x, y, _, rem_fx_labels = batch
         # Use chain of effects defined in config
         if order:
@@ -52,25 +52,30 @@ class RemFXChainInference(pl.LightningModule):
         with torch.no_grad():
             for i, (elem, effects_list) in enumerate(zip(x, effects)):
                 elem = elem.unsqueeze(0)  # Add batch dim
-                effect_chain_idx = [
-                    effects_order.index(effect.__name__) for effect in effects_list
+                # effect_chain_idx = [
+                #     effects_order.index(effect.__name__) for effect in effects_list
+                # ]
+                effect_list_names = [effect.__name__ for effect in effects_list]
+                effects = [
+                    effect for effect in effects_order if effect in effect_list_names
                 ]
+
                 # log_wandb_audio_batch(
                 #     logger=self.logger,
-                #     id=f"{i}_Before",
+                #     id=f"{batch_idx}_{i}_Before",
                 #     samples=elem.cpu(),
                 #     sampling_rate=self.sample_rate,
-                #     caption=effect_chain,
+                #     caption=effects,
                 # )
-                for idx in effect_chain_idx:
+                for effect in effects:
                     # Sample the model
-                    elem = self.model[effects_order[idx]].model.sample(elem)
+                    elem = self.model[effect].model.sample(elem)
                     # log_wandb_audio_batch(
                     #     logger=self.logger,
-                    #     id=f"{i}_{effect}",
+                    #     id=f"{batch_idx}_{i}_{effect}",
                     #     samples=elem.cpu(),
                     #     sampling_rate=self.sample_rate,
-                    #     caption=effect_chain,
+                    #     caption=effects,
                     # )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
@@ -81,8 +86,8 @@ class RemFXChainInference(pl.LightningModule):
     def test_step(self, batch, batch_idx):
         x, y, _, _ = batch  # x, y = (B, C, T), (B, C, T)
         # Random order
-        random.shuffle(self.effect_order)
-        loss, output = self.forward(batch, order=self.effect_order)
+        # random.shuffle(self.effect_order)
+        loss, output = self.forward(batch, batch_idx, order=self.effect_order)
         # Crop target to match output
         if output.shape[-1] < y.shape[-1]:
             y = causal_crop(y, output.shape[-1])
@@ -96,8 +101,7 @@ class RemFXChainInference(pl.LightningModule):
                 else:
                     negate = 1
                 self.log(
-                    f"test_{metric}_"
-                    + "".join(self.effect_order).replace("RandomPedalboard", ""),
+                    f"test_{metric}",  # + "".join(self.effect_order).replace("RandomPedalboard", ""),
                     negate * self.metrics[metric](output, y),
                     on_step=False,
                     on_epoch=True,

scripts/chain_inference.py

@@ -47,17 +47,12 @@ def main(cfg: DictConfig):
         logger=logger,
     )
 
+    log.info("Instantiating Inference Model")
     inference_model = RemFXChainInference(
         models,
         sample_rate=cfg.sample_rate,
         num_bins=cfg.num_bins,
-        effect_order=[
-            "RandomPedalboardDistortion",
-            "RandomPedalboardCompressor",
-            "RandomPedalboardReverb",
-            "RandomPedalboardChorus",
-            "RandomPedalboardDelay",
-        ],
+        effect_order=cfg.inference_effects_ordering,
     )
     trainer.test(model=inference_model, datamodule=datamodule)