Refactor dataset to apply multiple effects at a time

README.md

@@ -40,6 +40,6 @@ Experiment dictates data, ckpt dictates model
 `python scripts/test.py +exp=umx_distortion.yaml +ckpt_path=test_ckpts/umx_dist.ckpt`
 
 ## Misc.
-By default, files are rendered to `input_dir / processed / train/val/test`.
+By default, files are rendered to `input_dir / processed / {string_of_effects} / {train|val|test}`.
 To skip rendering files (use previously rendered), add `render_files=False` to the command-line (added to test by default).
 To change the rendered location, add `render_root={path/to/dir}` to the command-line (use this for train and test)

cfg/{effects → applied_effects}/all.yaml

@@ -1,5 +1,5 @@
 # @package _global_
-effects:
+applied_effects:
   Chorus:
     _target_: remfx.effects.RandomPedalboardChorus
     sample_rate: ${sample_rate}

cfg/config.yaml

@@ -1,15 +1,17 @@
 defaults:
   - _self_
   - model: null
-  - effects: null
+  - applied_effects: null
+  - effect_to_remove: null
 
+max_effects_per_file: 3
 seed: 12345
 train: True
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
 render_files: True
-render_root: "./data/processed"
+render_root: "./data"
 
 callbacks:
   model_checkpoint:
@@ -33,7 +35,9 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "train"
-    effect_types: ${effects}
+    applied_effects: ${applied_effects}
+    effect_to_remove: ${effect_to_remove}
+    max_effects_per_file: ${max_effects_per_file}
     render_files: ${render_files}
     render_root: ${render_root}
   val_dataset:
@@ -42,7 +46,9 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "val"
-    effect_types: ${effects}
+    applied_effects: ${applied_effects}
+    effect_to_remove: ${effect_to_remove}
+    max_effects_per_file: ${max_effects_per_file}
     render_files: ${render_files}
     render_root: ${render_root}
   test_dataset:
@@ -51,7 +57,9 @@ datamodule:
     root: ${oc.env:DATASET_ROOT}
     chunk_size: ${chunk_size}
     mode: "test"
-    effect_types: ${effects}
+    applied_effects: ${applied_effects}
+    effect_to_remove: ${effect_to_remove}
+    max_effects_per_file: ${max_effects_per_file}
     render_files: ${render_files}
     render_root: ${render_root}
 

cfg/effect_to_remove/all.yaml

@@ -0,0 +1,31 @@
+# @package _global_
+effect_to_remove:
+  Chorus:
+    _target_: remfx.effects.RandomPedalboardChorus
+    sample_rate: ${sample_rate}
+    min_depth: 0.2
+    min_mix: 0.3
+  Distortion:
+    _target_: remfx.effects.RandomPedalboardDistortion
+    sample_rate: ${sample_rate}
+    min_drive_db: 10
+    max_drive_db: 50
+  Compressor:
+    _target_: remfx.effects.RandomPedalboardCompressor
+    sample_rate: ${sample_rate}
+    min_threshold_db: -42.0
+    max_threshold_db: -20.0
+    min_ratio: 1.5
+    max_ratio: 6.0
+  Reverb:
+    _target_: remfx.effects.RandomPedalboardReverb
+    sample_rate: ${sample_rate}
+    min_room_size: 0.3
+    max_room_size: 1.0
+    min_damping: 0.2
+    max_damping: 1.0
+    min_wet_dry: 0.2
+    max_wet_dry: 0.8
+    min_width: 0.2
+    max_width: 1.0
+

cfg/{effects → effect_to_remove}/chorus.yaml

@@ -1,5 +1,5 @@
 # @package _global_
-effects:
+effect_to_remove:
   Chorus:
     _target_: remfx.effects.RandomPedalboardChorus
     sample_rate: ${sample_rate}

cfg/{effects → effect_to_remove}/compressor.yaml

@@ -1,5 +1,5 @@
 # @package _global_
-effects:
+effect_to_remove:
   Compressor:
     _target_: remfx.effects.RandomPedalboardCompressor
     sample_rate: ${sample_rate}

cfg/{effects → effect_to_remove}/distortion.yaml

@@ -1,5 +1,5 @@
 # @package _global_
-effects:
+effect_to_remove:
   Distortion:
     _target_: remfx.effects.RandomPedalboardDistortion
     sample_rate: ${sample_rate}

cfg/{effects → effect_to_remove}/reverb.yaml

@@ -1,5 +1,5 @@
 # @package _global_
-effects:
+effect_to_remove:
   Reverb:
     _target_: remfx.effects.RandomPedalboardReverb
     sample_rate: ${sample_rate}

cfg/exp/demucs_all.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: demucs
-  - override /effects: all
+  - override /applied_effects: all
+  - override /effect_to_remove: all

cfg/exp/demucs_chorus.yaml

@@ -1,4 +1,6 @@
 # @package _global_
 defaults:
   - override /model: demucs
-  - override /effects: chorus
+  - override /applied_effects: all
+  - override /effect_to_remove: chorus
+

cfg/exp/demucs_compressor.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: demucs
-  - override /effects: compressor
+  - override /applied_effects: all
+  - override /effect_to_remove: distortion

cfg/exp/demucs_distortion.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: demucs
-  - override /effects: distortion
+  - override /applied_effects: all
+  - override /effect_to_remove: distortion

cfg/exp/demucs_reverb.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: demucs
-  - override /effects: reverb
+  - override /applied_effects: all
+  - override /effect_to_remove: reverb

cfg/exp/umx_all.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: umx
-  - override /effects: all
+  - override /applied_effects: all
+  - override /effect_to_remove: all

cfg/exp/umx_chorus.yaml

@@ -1,4 +1,6 @@
 # @package _global_
 defaults:
   - override /model: umx
-  - override /effects: chorus
+  - override /applied_effects: all
+  - override /effect_to_remove: chorus
+

cfg/exp/umx_compressor.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: umx
-  - override /effects: compressor
+  - override /applied_effects: all
+  - override /effect_to_remove: distortion

cfg/exp/umx_distortion.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: umx
-  - override /effects: distortion
+  - override /applied_effects: all
+  - override /effect_to_remove: distortion

cfg/exp/umx_reverb.yaml

@@ -1,4 +1,5 @@
 # @package _global_
 defaults:
   - override /model: umx
-  - override /effects: reverb
+  - override /applied_effects: all
+  - override /effect_to_remove: reverb

remfx/datasets.py

@@ -1,16 +1,20 @@
 import torch
 from torch.utils.data import Dataset, DataLoader
-import torchaudio
 import torch.nn.functional as F
+import torchaudio
 from pathlib import Path
 import pytorch_lightning as pl
-from typing import Any, List
+import sys
+from typing import Any, Dict
 from remfx import effects
 from tqdm import tqdm
 from remfx.utils import create_sequential_chunks
+import shutil
 
 # https://zenodo.org/record/1193957 -> VocalSet
 
+ALL_EFFECTS = effects.Pedalboard_Effects
+
 
 class VocalSet(Dataset):
     def __init__(
@@ -18,7 +22,9 @@ class VocalSet(Dataset):
         root: str,
         sample_rate: int,
         chunk_size: int = 3,
-        effect_types: List[torch.nn.Module] = None,
+        applied_effects: Dict[str, torch.nn.Module] = None,
+        effect_to_remove: Dict[str, torch.nn.Module] = None,
+        max_effects_per_file: int = 1,
         render_files: bool = True,
         render_root: str = None,
         mode: str = "train",
@@ -31,22 +37,36 @@ class VocalSet(Dataset):
         self.chunk_size = chunk_size
         self.sample_rate = sample_rate
         self.mode = mode
-
+        self.max_effects_per_file = max_effects_per_file
+        self.effect_to_remove = effect_to_remove
         mode_path = self.root / self.mode
         self.files = sorted(list(mode_path.glob("./**/*.wav")))
         self.normalize = effects.LoudnessNormalize(sample_rate, target_lufs_db=-20)
-        self.effect_types = effect_types
-        effect_str = "_".join([e for e in self.effect_types])
-        self.processed_root = self.render_root / "processed" / effect_str / self.mode
-        if self.processed_root.exists():
+        self.applied_effects = applied_effects
+        self.effect_to_remove_name = list(effect_to_remove.keys())[0]
+
+        effect_str = "_".join([e for e in self.applied_effects])
+        effect_str += f"_{self.effect_to_remove_name}"
+        self.proc_root = self.render_root / "processed" / effect_str / self.mode
+
+        if self.proc_root.exists() and len(list(self.proc_root.iterdir())) > 0:
             print("Found processed files.")
-            render_files = False
+            if render_files:
+                re_render = input(
+                    "WARNING: By default, will re-render files.\n"
+                    "Set render_files=False to skip re-rendering.\n"
+                    "Are you sure you want to re-render? (y/n): "
+                )
+                if re_render != "y":
+                    sys.exit()
+                shutil.rmtree(self.proc_root)
+
         self.num_chunks = 0
         print("Total files:", len(self.files))
         print("Processing files...")
         if render_files:
             # Split audio file into chunks, resample, then apply random effects
-            self.processed_root.mkdir(parents=True, exist_ok=True)
+            self.proc_root.mkdir(parents=True, exist_ok=True)
             for audio_file in tqdm(self.files, total=len(self.files)):
                 chunks, orig_sr = create_sequential_chunks(audio_file, self.chunk_size)
                 for chunk in chunks:
@@ -56,27 +76,16 @@ class VocalSet(Dataset):
                     if resampled_chunk.shape[-1] < chunk_size:
                         # Skip if chunk is too small
                         continue
-                    # Apply effect
-                    effect_idx = torch.rand(1).item() * len(self.effect_types.keys())
-                    effect_name = list(self.effect_types.keys())[int(effect_idx)]
-                    effect = self.effect_types[effect_name]
-                    effected_input = effect(resampled_chunk)
-                    # Normalize
-                    normalized_input = self.normalize(effected_input)
-                    normalized_target = self.normalize(resampled_chunk)
-
-                    output_dir = self.processed_root / str(self.num_chunks)
+
+                    x, y, effect = self.process_effects(resampled_chunk)
+                    output_dir = self.proc_root / str(self.num_chunks)
                     output_dir.mkdir(exist_ok=True)
-                    torchaudio.save(
-                        output_dir / "input.wav", normalized_input, self.sample_rate
-                    )
-                    torchaudio.save(
-                        output_dir / "target.wav", normalized_target, self.sample_rate
-                    )
-                    torch.save(effect_name, output_dir / "effect_name.pt")
+                    torchaudio.save(output_dir / "input.wav", x, self.sample_rate)
+                    torchaudio.save(output_dir / "target.wav", y, self.sample_rate)
+                    torch.save(effect, output_dir / "effect.pt")
                     self.num_chunks += 1
         else:
-            self.num_chunks = len(list(self.processed_root.iterdir()))
+            self.num_chunks = len(list(self.proc_root.iterdir()))
 
         print(
             f"Found {len(self.files)} {self.mode} files .\n"
@@ -87,13 +96,47 @@ class VocalSet(Dataset):
         return self.num_chunks
 
     def __getitem__(self, idx):
-        input_file = self.processed_root / str(idx) / "input.wav"
-        target_file = self.processed_root / str(idx) / "target.wav"
-        effect_name = torch.load(self.processed_root / str(idx) / "effect_name.pt")
+        input_file = self.proc_root / str(idx) / "input.wav"
+        target_file = self.proc_root / str(idx) / "target.wav"
+        effect_name = torch.load(self.proc_root / str(idx) / "effect.pt")
         input, sr = torchaudio.load(input_file)
         target, sr = torchaudio.load(target_file)
         return (input, target, effect_name)
 
+    def process_effects(self, dry: torch.Tensor):
+        # Apply random number of effects up to num_effects - 1 (excluding effect_to_remove)
+        if self.max_effects_per_file > 1:
+            num_effects = torch.randint(self.max_effects_per_file - 1, (1,)).item()
+            # Remove effect to remove from applied effects if present
+            self.applied_effects.pop(self.effect_to_remove_name, None)
+
+            # Choose random effects to apply
+            effect_indices = torch.randperm(len(self.applied_effects.keys()))[
+                :num_effects
+            ]
+            effects_to_apply = [
+                list(self.applied_effects.keys())[i] for i in effect_indices
+            ]
+            labels = []
+            for effect_name in effects_to_apply:
+                effect = self.applied_effects[effect_name]
+                dry = effect(dry)
+                labels.append(ALL_EFFECTS.index(type(effect)))
+
+        # Apply effect_to_remove
+        effect = self.effect_to_remove[self.effect_to_remove_name]
+        wet = effect(torch.clone(dry))
+        labels.append(ALL_EFFECTS.index(type(effect)))
+
+        # Convert labels to one-hot
+        one_hot = F.one_hot(torch.tensor(labels), num_classes=len(ALL_EFFECTS))
+        effects_present = torch.sum(one_hot, dim=0).float()
+
+        # Normalize
+        normalized_dry = self.normalize(dry)
+        normalized_wet = self.normalize(wet)
+        return normalized_dry, normalized_wet, effects_present
+
 
 class VocalSetDatamodule(pl.LightningDataModule):
     def __init__(

remfx/effects.py

@@ -675,7 +675,7 @@ class RandomAudioEffectsChannel(torch.nn.Module):
                     p=compressor_prob,
                 ),
                 RandomApply(
-                    [RandomPebalboardReverb(sample_rate)],
+                    [RandomPedalboardReverb(sample_rate)],
                     p=reverb_prob,
                 ),
                 RandomApply(
@@ -696,3 +696,14 @@ class RandomAudioEffectsChannel(torch.nn.Module):
 
     def forward(self, x: torch.Tensor):
         return self.transforms(x)
+
+
+Pedalboard_Effects = [
+    RandomPedalboardReverb,
+    RandomPedalboardChorus,
+    RandomPedalboardDelay,
+    RandomPedalboardDistortion,
+    RandomPedalboardCompressor,
+    RandomPedalboardPhaser,
+    RandomPedalboardLimiter,
+]

remfx/models.py

@@ -64,7 +64,6 @@ class RemFXModel(pl.LightningModule):
         optimizer_idx,
         optimizer_closure,
         on_tpu,
-        using_native_amp,
         using_lbfgs,
     ):
         # update params

setup.py

@@ -30,7 +30,7 @@ setup(
     packages=find_packages(),
     install_requires=[
         "torch>=1.11.0",
-        "torchaudio",
+        "torchaudio>=0.13.0",
         "functorch",
         "scipy",
         "numpy",