Fix input metrics

cfg/model/demucs.yaml

@@ -13,5 +13,5 @@ model:
     audio_channels: 1
     nfft: 4096
     sample_rate: ${sample_rate}
-    channels: 64
+    channels: 48
 

remfx/callbacks.py

@@ -71,48 +71,6 @@ class AudioCallback(Callback):
         self.on_validation_batch_start(*args)
 
 
-class MetricCallback(Callback):
-    def on_validation_batch_start(
-        self, trainer, pl_module, batch, batch_idx, dataloader_idx
-    ):
-        x, target, _, _ = batch
-        # Log Input Metrics
-        for metric in pl_module.metrics:
-            # SISDR returns negative values, so negate them
-            if metric == "SISDR":
-                negate = -1
-            else:
-                negate = 1
-            # Only Log FAD on test set
-            if metric == "FAD":
-                continue
-            pl_module.log(
-                f"Input_{metric}",
-                negate * pl_module.metrics[metric](x, target),
-                on_step=False,
-                on_epoch=True,
-                logger=True,
-                prog_bar=True,
-                sync_dist=True,
-            )
-
-    def on_test_batch_start(self, trainer, pl_module, batch, batch_idx, dataloader_idx):
-        self.on_validation_batch_start(
-            trainer, pl_module, batch, batch_idx, dataloader_idx
-        )
-        # Log FAD
-        x, target, _, _ = batch
-        pl_module.log(
-            "Input_FAD",
-            pl_module.metrics["FAD"](x, target),
-            on_step=False,
-            on_epoch=True,
-            logger=True,
-            prog_bar=True,
-            sync_dist=True,
-        )
-
-
 def log_wandb_audio_batch(
     logger: pl.loggers.WandbLogger,
     id: str,

remfx/models.py

@@ -14,7 +14,6 @@ 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
 
@@ -52,31 +51,36 @@ 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
-                # ]
+                # Get the correct effect by search for names in effects_order
                 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"{batch_idx}_{i}_Before",
-                #     samples=elem.cpu(),
-                #     sampling_rate=self.sample_rate,
-                #     caption=effects,
-                # )
+                log_wandb_audio_batch(
+                    logger=self.logger,
+                    id=f"{i}_Before",
+                    samples=elem.cpu(),
+                    sampling_rate=self.sample_rate,
+                    caption=effects,
+                )
                 for effect in effects:
                     # Sample the model
                     elem = self.model[effect].model.sample(elem)
-                    # log_wandb_audio_batch(
-                    #     logger=self.logger,
-                    #     id=f"{batch_idx}_{i}_{effect}",
-                    #     samples=elem.cpu(),
-                    #     sampling_rate=self.sample_rate,
-                    #     caption=effects,
-                    # )
+                    log_wandb_audio_batch(
+                        logger=self.logger,
+                        id=f"{i}_{effect}",
+                        samples=elem.cpu(),
+                        sampling_rate=self.sample_rate,
+                        caption=effects,
+                    )
+                log_wandb_audio_batch(
+                    logger=self.logger,
+                    id=f"{i}_After",
+                    samples=elem.cpu(),
+                    sampling_rate=self.sample_rate,
+                    caption=effects,
+                )
                 output.append(elem.squeeze(0))
         output = torch.stack(output)
 
@@ -111,7 +115,7 @@ class RemFXChainInference(pl.LightningModule):
                 )
 
     def sample(self, batch):
-        return self.forward(batch)[1]
+        return self.forward(batch, 0)[1]
 
 
 class RemFX(pl.LightningModule):
@@ -207,6 +211,15 @@ class RemFX(pl.LightningModule):
                     prog_bar=True,
                     sync_dist=True,
                 )
+                self.log(
+                    f"Input_{metric}",
+                    negate * self.metrics[metric](x, y),
+                    on_step=False,
+                    on_epoch=True,
+                    logger=True,
+                    prog_bar=True,
+                    sync_dist=True,
+                )
         return loss
 
 

scripts/chain_inference.py

@@ -20,9 +20,10 @@ def main(cfg: DictConfig):
     for effect in cfg.ckpts:
         ckpt_path = cfg.ckpts[effect]
         model = hydra.utils.instantiate(cfg.model, _convert_="partial")
-        state_dict = torch.load(ckpt_path)["state_dict"]
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        state_dict = torch.load(ckpt_path, map_location=device)["state_dict"]
         model.load_state_dict(state_dict)
-        model.to("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device)
         models[effect] = model
 
     callbacks = []