Fix loss issue in chain_inference

remfx/models.py

@@ -12,6 +12,7 @@ from remfx.utils import FADLoss, spectrogram
 from remfx.tcn import TCN
 from remfx.utils import causal_crop
 from remfx.callbacks import log_wandb_audio_batch
+from einops import rearrange
 from remfx import effects
 import asteroid
 
@@ -47,6 +48,23 @@ class RemFXChainInference(pl.LightningModule):
             for effect_label in rem_fx_labels
         ]
         output = []
+        input_samples = rearrange(x, "b c t -> c (b t)").unsqueeze(0)
+        target_samples = rearrange(y, "b c t -> c (b t)").unsqueeze(0)
+
+        log_wandb_audio_batch(
+            logger=self.logger,
+            id="input_effected_audio",
+            samples=input_samples.cpu(),
+            sampling_rate=self.sample_rate,
+            caption=effects,
+        )
+        log_wandb_audio_batch(
+            logger=self.logger,
+            id="target_audio",
+            samples=target_samples.cpu(),
+            sampling_rate=self.sample_rate,
+            caption="Target Data",
+        )
         with torch.no_grad():
             for i, (elem, effects_list) in enumerate(zip(x, effects)):
                 elem = elem.unsqueeze(0)  # Add batch dim
@@ -56,33 +74,41 @@ class RemFXChainInference(pl.LightningModule):
                     effect for effect in effects_order if effect in effect_list_names
                 ]
 
-                log_wandb_audio_batch(
-                    logger=self.logger,
-                    id=f"{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"{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,
-                )
+                #     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)
-
+        output_samples = rearrange(output, "b c t -> c (b t)").unsqueeze(0)
+
+        log_wandb_audio_batch(
+            logger=self.logger,
+            id="output_audio",
+            samples=output_samples.cpu(),
+            sampling_rate=self.sample_rate,
+            caption="Output Data",
+        )
         loss = self.mrstftloss(output, y) + self.l1loss(output, y) * 100
         return loss, output
 
@@ -112,6 +138,16 @@ class RemFXChainInference(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
 
     def sample(self, batch):
         return self.forward(batch, 0)[1]