Fix circular import issue

cfg/exp/chain_inference.yaml

@@ -6,7 +6,7 @@ seed: 12345
 sample_rate: 48000
 chunk_size: 262144 # 5.5s
 logs_dir: "./logs"
-render_files: True
+render_files: False
 render_root: "/scratch/EffectSet"
 accelerator: "gpu"
 log_audio: True
@@ -26,10 +26,12 @@ effects_to_remove:
 datamodule:
   batch_size: 16
   num_workers: 8
+  train_dataset: None
+  val_dataset: None
 ckpts:
-  RandomPedalboardChorus: "ckpts/chorus.ckpt"
-  RandomPedalboardDelay: "ckpts/delay.ckpt"
   RandomPedalboardDistortion: "ckpts/distortion.ckpt"
   RandomPedalboardCompressor: "ckpts/compressor.ckpt"
   RandomPedalboardReverb: "ckpts/reverb.ckpt"
+  RandomPedalboardChorus: "ckpts/chorus.ckpt"
+  RandomPedalboardDelay: "ckpts/delay.ckpt"
 num_bins: 1025

remfx/callbacks.py

@@ -4,7 +4,6 @@ from einops import rearrange
 import torch
 import wandb
 from torch import Tensor
-from remfx.models import RemFXChainInference
 
 
 class AudioCallback(Callback):
@@ -47,6 +46,9 @@ class AudioCallback(Callback):
         # Only run on first batch
         if batch_idx == 0 and self.log_audio:
             with torch.no_grad():
+                # Avoids circular import
+                from remfx.models import RemFXChainInference
+
                 if type(pl_module) == RemFXChainInference:
                     y = pl_module.sample(batch)
                 else:

remfx/models.py

@@ -14,12 +14,13 @@ from remfx.utils import causal_crop
 from remfx.callbacks import log_wandb_audio_batch
 from remfx import effects
 import asteroid
+import random
 
 ALL_EFFECTS = effects.Pedalboard_Effects
 
 
 class RemFXChainInference(pl.LightningModule):
-    def __init__(self, models, sample_rate, num_bins):
+    def __init__(self, models, sample_rate, num_bins, effect_order):
         super().__init__()
         self.model = models
         self.mrstftloss = MultiResolutionSTFTLoss(
@@ -33,36 +34,45 @@ class RemFXChainInference(pl.LightningModule):
                 "FAD": FADLoss(sample_rate=sample_rate),
             }
         )
+        self.sample_rate = sample_rate
+        self.effect_order = effect_order
 
-    def forward(self, batch):
+    def forward(self, batch, order=None):
         x, y, _, rem_fx_labels = batch
         # Use chain of effects defined in config
+        if order:
+            effects_order = order
+        else:
+            effects_order = self.effect_order
         effects = [
             [ALL_EFFECTS[i] for i, effect in enumerate(effect_label) if effect == 1.0]
             for effect_label in rem_fx_labels
         ]
         output = []
         with torch.no_grad():
-            for i, (elem, effect_chain) in enumerate(zip(x, effects)):
+            for i, (elem, effects_list) in enumerate(zip(x, effects)):
                 elem = elem.unsqueeze(0)  # Add batch dim
-                log_wandb_audio_batch(
-                    logger=self.logger,
-                    id=f"{i}_Before",
-                    samples=elem.cpu(),
-                    sampling_rate=self.sample_rate,
-                    caption=effect_chain,
-                )
+                effect_chain = [
+                    effects_order.index(effect.__name__) for effect in effects_list
+                ]
+                # log_wandb_audio_batch(
+                #     logger=self.logger,
+                #     id=f"{i}_Before",
+                #     samples=elem.cpu(),
+                #     sampling_rate=self.sample_rate,
+                #     caption=effect_chain,
+                # )
+                effect_chain
                 for effect in effect_chain:
-                    # Get correct model based on effect name. This is a bit hacky
-                    # Then sample the model
-                    elem = self.model[effect.__name__].model.sample(elem)
-                    log_wandb_audio_batch(
-                        logger=self.logger,
-                        id=f"{i}_{effect}",
-                        samples=elem.cpu(),
-                        sampling_rate=self.sample_rate,
-                        caption=effect_chain,
-                    )
+                    # Sample the model
+                    elem = self.model[effect].model.sample(elem)
+                    # log_wandb_audio_batch(
+                    #     logger=self.logger,
+                    #     id=f"{i}_{effect}",
+                    #     samples=elem.cpu(),
+                    #     sampling_rate=self.sample_rate,
+                    #     caption=effect_chain,
+                    # )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
 
@@ -71,8 +81,9 @@ class RemFXChainInference(pl.LightningModule):
 
     def test_step(self, batch, batch_idx):
         x, y, _, _ = batch  # x, y = (B, C, T), (B, C, T)
-
-        loss, output = self.forward(batch)
+        # Random order
+        order = random.shuffle(self.effect_order)
+        loss, output = self.forward(batch, order=order)
         # Crop target to match output
         if output.shape[-1] < y.shape[-1]:
             y = causal_crop(y, output.shape[-1])
@@ -86,7 +97,7 @@ class RemFXChainInference(pl.LightningModule):
                 else:
                     negate = 1
                 self.log(
-                    f"test_{metric}",
+                    f"test_{metric}_" + "".join(order),
                     negate * self.metrics[metric](output, y),
                     on_step=False,
                     on_epoch=True,

scripts/chain_inference.py

@@ -22,7 +22,7 @@ def main(cfg: DictConfig):
         model = hydra.utils.instantiate(cfg.model, _convert_="partial")
         state_dict = torch.load(ckpt_path)["state_dict"]
         model.load_state_dict(state_dict)
-        model.to(cfg.device)
+        model.to("cuda" if torch.cuda.is_available() else "cpu")
         models[effect] = model
 
     callbacks = []
@@ -48,7 +48,10 @@ def main(cfg: DictConfig):
     )
 
     inference_model = RemFXChainInference(
-        models, sample_rate=cfg.sample_rate, num_bins=cfg.num_bins
+        models,
+        sample_rate=cfg.sample_rate,
+        num_bins=cfg.num_bins,
+        order=["Distortion", "Compressor", "Reverb", "Chorus", "Delay"],
     )
     trainer.test(model=inference_model, datamodule=datamodule)