Add flag to not log audio

cfg/config.yaml

@@ -11,6 +11,7 @@ logs_dir: "./logs"
 render_files: True
 render_root: "./data"
 accelerator: null
+log_audio: True
 
 # Effects
 max_kept_effects: -1
@@ -47,6 +48,7 @@ callbacks:
   audio_logging:
     _target_: remfx.callbacks.AudioCallback
     sample_rate: ${sample_rate}
+    log_audio: ${log_audio}
   metric_logging:
     _target_: remfx.callbacks.MetricCallback
 

cfg/exp/default.yaml

@@ -9,6 +9,8 @@ logs_dir: "./logs"
 render_files: True
 render_root: "./data"
 accelerator: null
+log_audio: True
+# Effects
 max_kept_effects: -1
 max_removed_effects: -1
 shuffle_kept_effects: True

remfx/callbacks.py

@@ -7,10 +7,13 @@ from torch import Tensor
 
 
 class AudioCallback(Callback):
-    def __init__(self, sample_rate, *args, **kwargs):
+    def __init__(self, sample_rate, log_audio, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.log_audio = log_audio
         self.log_train_audio = True
         self.sample_rate = sample_rate
+        if not self.log_audio:
+            self.log_train_audio = False
 
     def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
         # Log initial audio
@@ -41,7 +44,7 @@ class AudioCallback(Callback):
     ):
         x, target, _, _ = batch
         # Only run on first batch
-        if batch_idx == 0:
+        if batch_idx == 0 and self.log_audio:
             with torch.no_grad():
                 y = pl_module.model.sample(x)
             # Concat samples together for easier viewing in dashboard