Enables non square mel spectrograms

audiodiffusion/__init__.py

@@ -181,8 +181,8 @@ class AudioDiffusionPipeline(DiffusionPipeline):
             self.scheduler.set_timesteps(steps)
         mask = None
         images = noise = torch.randn(
-            (batch_size, self.unet.in_channels, self.unet.sample_size,
-             self.unet.sample_size),
+            (batch_size, self.unet.in_channels, self.unet.sample_size[0],
+             self.unet.sample_size[1]),
             generator=generator)
 
         if audio_file is not None or raw_audio is not None:
@@ -206,7 +206,7 @@ class AudioDiffusionPipeline(DiffusionPipeline):
                     noise, torch.tensor(steps - start_step))
 
             pixels_per_second = (mel.get_sample_rate() *
-                                 self.unet.sample_size / mel.hop_length /
+                                 mel.y_res / mel.hop_length /
                                  mel.x_res)
             mask_start = int(mask_start_secs * pixels_per_second)
             mask_end = int(mask_end_secs * pixels_per_second)

scripts/train_unconditional.py

@@ -26,6 +26,9 @@ import numpy as np
 from tqdm.auto import tqdm
 from librosa.util import normalize
 
+import sys
+sys.path.append('.')
+sys.path.append('..')
 from audiodiffusion.mel import Mel
 from audiodiffusion import LatentAudioDiffusionPipeline, AudioDiffusionPipeline
 
@@ -42,6 +45,18 @@ def main(args):
         logging_dir=logging_dir,
     )
 
+    # Handle the resolutions.
+    try:
+        args.resolution = (int(args.resolution), int(args.resolution))
+    except:
+        try :
+            args.resolution = tuple(int(x) for x in args.resolution.split(","))
+            if len(args.resolution) != 2:
+                raise ValueError("Resolution must be a tuple of two integers or a single integer.")
+        except:
+            raise ValueError("Resolution must be a tuple of two integers or a single integer.")
+    assert isinstance(args.resolution, tuple)
+
     if args.vae is not None:
         vqvae = AutoencoderKL.from_pretrained(args.vae)
 
@@ -156,9 +171,9 @@ def main(args):
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run)
 
-    mel = Mel(x_res=args.resolution,
-              y_res=args.resolution,
-              hop_length=args.hop_length)
+    mel = Mel(x_res=args.resolution[0],
+            y_res=args.resolution[1],
+            hop_length=args.hop_length)
 
     global_step = 0
     for epoch in range(args.num_epochs):
@@ -311,7 +326,7 @@ if __name__ == "__main__":
     parser.add_argument("--output_dir", type=str, default="ddpm-model-64")
     parser.add_argument("--overwrite_output_dir", type=bool, default=False)
     parser.add_argument("--cache_dir", type=str, default=None)
-    parser.add_argument("--resolution", type=int, default=256)
+    parser.add_argument("--resolution", type=str, default="256")
     parser.add_argument("--train_batch_size", type=int, default=16)
     parser.add_argument("--eval_batch_size", type=int, default=16)
     parser.add_argument("--num_epochs", type=int, default=100)