move resolution specification to dataset generation

README.md

@@ -57,7 +57,7 @@ pip install .
 
 ```bash
 python scripts/audio_to_images.py \
-  --resolution 64 \
+  --resolution 64,64 \
   --hop_length 1024 \
   --input_dir path-to-audio-files \
   --output_dir path-to-output-data
@@ -78,7 +78,6 @@ python scripts/audio_to_images.py \
 accelerate launch --config_file config/accelerate_local.yaml \
   scripts/train_unconditional.py \
   --dataset_name data/audio-diffusion-64 \
-  --resolution 64 \
   --hop_length 1024 \
   --output_dir models/ddpm-ema-audio-64 \
   --train_batch_size 16 \
@@ -94,7 +93,6 @@ accelerate launch --config_file config/accelerate_local.yaml \
 accelerate launch --config_file config/accelerate_local.yaml \
   scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
-  --resolution 256 \
   --output_dir models/audio-diffusion-256 \
   --num_epochs 100 \
   --train_batch_size 2 \
@@ -113,7 +111,6 @@ accelerate launch --config_file config/accelerate_local.yaml \
 accelerate launch --config_file config/accelerate_sagemaker.yaml \
   scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
-  --resolution 256 \
   --output_dir models/ddpm-ema-audio-256 \
   --train_batch_size 16 \
   --num_epochs 100 \
@@ -147,5 +144,4 @@ python scripts/train_vae.py \
 accelerate launch ...
   ...
   --vae models/autoencoder-kl
-  --latent_resoultion 32
 ```

audiodiffusion/__init__.py

@@ -180,10 +180,12 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         if steps is not None:
             self.scheduler.set_timesteps(steps)
         mask = None
-        images = noise = torch.randn(
-            (batch_size, self.unet.in_channels, self.unet.sample_size[0],
-             self.unet.sample_size[1]),
-            generator=generator)
+        # For backwards compatibility
+        sample_size = (self.unet.sample_size, self.unet.sample_size) if type(
+            self.unet.sample_size) == int else self.unet.sample_size
+        images = noise = torch.randn((batch_size, self.unet.in_channels) +
+                                     sample_size,
+                                     generator=generator)
 
         if audio_file is not None or raw_audio is not None:
             mel.load_audio(audio_file, raw_audio)
@@ -205,9 +207,8 @@ class AudioDiffusionPipeline(DiffusionPipeline):
                     torch.tensor(input_images[:, np.newaxis, np.newaxis, :]),
                     noise, torch.tensor(steps - start_step))
 
-            pixels_per_second = (mel.get_sample_rate() *
-                                 mel.y_res / mel.hop_length /
-                                 mel.x_res)
+            pixels_per_second = (mel.get_sample_rate() * sample_size[1] /
+                                 mel.hop_length / mel.x_res)
             mask_start = int(mask_start_secs * pixels_per_second)
             mask_end = int(mask_end_secs * pixels_per_second)
             mask = self.scheduler.add_noise(

audiodiffusion/mel.py

@@ -106,7 +106,7 @@ class Mel:
         log_S = librosa.power_to_db(S, ref=np.max, top_db=self.top_db)
         bytedata = (((log_S + self.top_db) * 255 / self.top_db).clip(0, 255) +
                     0.5).astype(np.uint8)
-        image = Image.frombytes("L", log_S.shape, bytedata.tobytes())
+        image = Image.fromarray(bytedata)
         return image
 
     def image_to_audio(self, image: Image.Image) -> np.ndarray:

scripts/audio_to_images.py

@@ -16,9 +16,9 @@ logger = logging.getLogger('audio_to_images')
 
 
 def main(args):
-    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)   
     os.makedirs(args.output_dir, exist_ok=True)
     audio_files = [
         os.path.join(root, file) for root, _, files in os.walk(args.input_dir)
@@ -35,8 +35,8 @@ def main(args):
                 continue
             for slice in range(mel.get_number_of_slices()):
                 image = mel.audio_slice_to_image(slice)
-                assert (image.width == args.resolution
-                        and image.height == args.resolution)
+                assert (image.width == args.resolution[0] and image.height
+                        == args.resolution[1]), "Wrong resolution"
                 # skip completely silent slices
                 if all(np.frombuffer(image.tobytes(), dtype=np.uint8) == 255):
                     logger.warn('File %s slice %d is completely silent',
@@ -52,6 +52,8 @@ def main(args):
                     "audio_file": audio_file,
                     "slice": slice,
                 }])
+    except Exception as e:
+        print(e)
     finally:
         if len(examples) == 0:
             logger.warn('No valid audio files were found.')
@@ -76,12 +78,30 @@ if __name__ == "__main__":
         "Create dataset of Mel spectrograms from directory of audio files.")
     parser.add_argument("--input_dir", type=str)
     parser.add_argument("--output_dir", type=str, default="data")
-    parser.add_argument("--resolution", type=int, default=256)
+    parser.add_argument("--resolution",
+                        type=str,
+                        default="256",
+                        help="Either square resolution or width,height.")
     parser.add_argument("--hop_length", type=int, default=512)
     parser.add_argument("--push_to_hub", type=str, default=None)
     args = parser.parse_args()
+
     if args.input_dir is None:
         raise ValueError(
-            "You must specify an input directory for the audio files."
-        )    
+            "You must specify an input directory for the audio files.")
+
+    # Handle the resolutions.
+    try:
+        args.resolution = (int(args.resolution), int(args.resolution))
+    except ValueError:
+        try:
+            args.resolution = tuple(int(x) for x in args.resolution.split(","))
+            if len(args.resolution) != 2:
+                raise ValueError
+        except ValueError:
+            raise ValueError(
+                "Resolution must be a tuple of two integers or a single integer."
+            )
+    assert isinstance(args.resolution, tuple)
+
     main(args)

scripts/train_unconditional.py

@@ -26,9 +26,6 @@ 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
 
@@ -45,31 +42,68 @@ 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.dataset_name is not None:
+        if os.path.exists(args.dataset_name):
+            dataset = load_from_disk(args.dataset_name,
+                                     args.dataset_config_name)["train"]
+        else:
+            dataset = load_dataset(
+                args.dataset_name,
+                args.dataset_config_name,
+                cache_dir=args.cache_dir,
+                use_auth_token=True if args.use_auth_token else None,
+                split="train",
+            )
+    else:
+        dataset = load_dataset(
+            "imagefolder",
+            data_dir=args.train_data_dir,
+            cache_dir=args.cache_dir,
+            split="train",
+        )
+    # Determine image resolution
+    resolution = dataset[0]['image'].height, dataset[0]['image'].width
 
+    augmentations = Compose([
+        ToTensor(),
+        Normalize([0.5], [0.5]),
+    ])
+
+    def transforms(examples):
+        if args.vae is not None and vqvae.config['in_channels'] == 3:
+            images = [
+                augmentations(image.convert('RGB'))
+                for image in examples["image"]
+            ]
+        else:
+            images = [augmentations(image) for image in examples["image"]]
+        return {"input": images}
+
+    dataset.set_transform(transforms)
+    train_dataloader = torch.utils.data.DataLoader(
+        dataset, batch_size=args.train_batch_size, shuffle=True)
+
+    vqvae = None
     if args.vae is not None:
         vqvae = AutoencoderKL.from_pretrained(args.vae)
+        # Determine latent resolution
+        with torch.no_grad():
+            latent_resolution = vqvae.encode(
+                torch.zeros((1, 1) +
+                            resolution)).latent_dist.sample().shape[2:]
 
     if args.from_pretrained is not None:
-        model = DiffusionPipeline.from_pretrained(args.from_pretrained).unet
+        pipeline = DiffusionPipeline.from_pretrained(args.from_pretrained)
+        model = pipeline.unet
+        if hasattr(pipeline, 'vqvae'):
+            vqvae = AutoencoderKL.from_pretrained(args.vae)
     else:
         model = UNet2DModel(
-            sample_size=args.resolution
-            if args.vae is None else args.latent_resolution,
+            sample_size=resolution if vqvae is None else latent_resolution,
             in_channels=1
-            if args.vae is None else vqvae.config['latent_channels'],
+            if vqvae is None else vqvae.config['latent_channels'],
             out_channels=1
-            if args.vae is None else vqvae.config['latent_channels'],
+            if vqvae is None else vqvae.config['latent_channels'],
             layers_per_block=2,
             block_out_channels=(128, 128, 256, 256, 512, 512),
             down_block_types=(
@@ -105,47 +139,6 @@ def main(args):
         eps=args.adam_epsilon,
     )
 
-    augmentations = Compose([
-        Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
-        CenterCrop(args.resolution),
-        ToTensor(),
-        Normalize([0.5], [0.5]),
-    ])
-
-    if args.dataset_name is not None:
-        if os.path.exists(args.dataset_name):
-            dataset = load_from_disk(args.dataset_name,
-                                     args.dataset_config_name)["train"]
-        else:
-            dataset = load_dataset(
-                args.dataset_name,
-                args.dataset_config_name,
-                cache_dir=args.cache_dir,
-                use_auth_token=True if args.use_auth_token else None,
-                split="train",
-            )
-    else:
-        dataset = load_dataset(
-            "imagefolder",
-            data_dir=args.train_data_dir,
-            cache_dir=args.cache_dir,
-            split="train",
-        )
-
-    def transforms(examples):
-        if args.vae is not None and vqvae.config['in_channels'] == 3:
-            images = [
-                augmentations(image.convert('RGB'))
-                for image in examples["image"]
-            ]
-        else:
-            images = [augmentations(image) for image in examples["image"]]
-        return {"input": images}
-
-    dataset.set_transform(transforms)
-    train_dataloader = torch.utils.data.DataLoader(
-        dataset, batch_size=args.train_batch_size, shuffle=True)
-
     lr_scheduler = get_scheduler(
         args.lr_scheduler,
         optimizer=optimizer,
@@ -171,9 +164,9 @@ def main(args):
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run)
 
-    mel = Mel(x_res=args.resolution[0],
-            y_res=args.resolution[1],
-            hop_length=args.hop_length)
+    mel = Mel(x_res=resolution[1],
+              y_res=resolution[0],
+              hop_length=args.hop_length)
 
     global_step = 0
     for epoch in range(args.num_epochs):
@@ -195,7 +188,7 @@ def main(args):
         for step, batch in enumerate(train_dataloader):
             clean_images = batch["input"]
 
-            if args.vae is not None:
+            if vqvae is not None:
                 vqvae.to(clean_images.device)
                 with torch.no_grad():
                     clean_images = vqvae.encode(
@@ -252,7 +245,7 @@ def main(args):
         # Generate sample images for visual inspection
         if accelerator.is_main_process:
             if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
-                if args.vae is not None:
+                if vqvae is not None:
                     pipeline = LatentAudioDiffusionPipeline(
                         unet=accelerator.unwrap_model(
                             ema_model.averaged_model if args.use_ema else model
@@ -326,7 +319,6 @@ 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=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)
@@ -364,7 +356,6 @@ if __name__ == "__main__":
     parser.add_argument("--from_pretrained", type=str, default=None)
     parser.add_argument("--start_epoch", type=int, default=0)
     parser.add_argument("--num_train_steps", type=int, default=1000)
-    parser.add_argument("--latent_resolution", type=int, default=None)
     parser.add_argument("--scheduler",
                         type=str,
                         default="ddpm",

scripts/train_vae.py

@@ -58,13 +58,10 @@ class AudioDiffusionDataModule(pl.LightningDataModule):
 
 class ImageLogger(Callback):
 
-    def __init__(self, every=1000, channels=3, resolution=256, hop_length=512):
+    def __init__(self, every=1000, hop_length=512):
         super().__init__()
-        self.mel = Mel(x_res=resolution,
-                       y_res=resolution,
-                       hop_length=hop_length)
         self.every = every
-        self.channels = channels
+        self.hop_length = hop_length
 
     @rank_zero_only
     def log_images_and_audios(self, pl_module, batch):
@@ -73,6 +70,12 @@ class ImageLogger(Callback):
             images = pl_module.log_images(batch, split='train')
         pl_module.train()
 
+        image_shape = next(iter(images.values())).shape
+        channels = image_shape[1]
+        mel = Mel(x_res=image_shape[2],
+                  y_res=image_shape[3],
+                  hop_length=self.hop_length)
+
         for k in images:
             images[k] = images[k].detach().cpu()
             images[k] = torch.clamp(images[k], -1., 1.)
@@ -86,14 +89,14 @@ class ImageLogger(Callback):
             images[k] = (images[k].numpy() *
                          255).round().astype("uint8").transpose(0, 2, 3, 1)
             for _, image in enumerate(images[k]):
-                audio = self.mel.image_to_audio(
-                    Image.fromarray(image, mode='RGB').convert('L') if self.
-                    channels == 3 else Image.fromarray(image[0]))
+                audio = mel.image_to_audio(
+                    Image.fromarray(image, mode='RGB').convert('L')
+                    if channels == 3 else Image.fromarray(image[0]))
                 pl_module.logger.experiment.add_audio(
                     tag + f"/{_}",
                     normalize(audio),
                     global_step=pl_module.global_step,
-                    sample_rate=self.mel.get_sample_rate())
+                    sample_rate=mel.get_sample_rate())
 
     def on_train_batch_end(self, trainer, pl_module, outputs, batch,
                            batch_idx):
@@ -139,7 +142,6 @@ if __name__ == "__main__":
                         "--gradient_accumulation_steps",
                         type=int,
                         default=1)
-    parser.add_argument("--resolution", type=int, default=256)
     parser.add_argument("--hop_length", type=int, default=512)
     parser.add_argument("--save_images_batches", type=int, default=1000)
     args = parser.parse_args()
@@ -160,8 +162,6 @@ if __name__ == "__main__":
         resume_from_checkpoint=args.resume_from_checkpoint,
         callbacks=[
             ImageLogger(every=args.save_images_batches,
-                        channels=config.model.params.ddconfig.out_ch,
-                        resolution=args.resolution,
                         hop_length=args.hop_length),
             HFModelCheckpoint(ldm_config=config,
                               hf_checkpoint=args.hf_checkpoint_dir,