various fixes

README.md

@@ -18,7 +18,7 @@ license: gpl-3.0
 **UPDATES**:
 
 15/10/2022
-Added latent audio diffusion (see below).
+Added latent audio diffusion (see below). Also added the possibility to train a model to use DDIM ([Denoising Diffusion Implicit Models](https://arxiv.org/pdf/2010.02502.pdf)) by setting `--scheduler ddim`. These have the benefit that samples can be generated with much fewer steps (~50) than used in training.
 
 4/10/2022
 It is now possible to mask parts of the input audio during generation which means you can stitch several samples together (think "out-painting").
@@ -138,5 +138,6 @@ python scripts/train_vae.py \
 #### Train latent diffusion model.
 ```bash
 accelerate launch ...
-  --vae models/autoencoder-kl 
+  --vae models/autoencoder-kl
+  --latent_resoultion 32
 ```

audiodiffusion/__init__.py

@@ -1,15 +1,16 @@
-from typing import Iterable, Tuple
+from typing import Iterable, Tuple, Union, List
 
 import torch
 import numpy as np
 from PIL import Image
 from tqdm.auto import tqdm
 from librosa.beat import beat_track
-from diffusers import DiffusionPipeline
+from diffusers import (DiffusionPipeline, DDPMPipeline, UNet2DConditionModel,
+                       DDIMScheduler, DDPMScheduler, AutoencoderKL)
 
 from .mel import Mel
 
-VERSION = "1.1.5"
+VERSION = "1.2.0"
 
 
 class AudioDiffusion:
@@ -42,7 +43,11 @@ class AudioDiffusion:
                        hop_length=hop_length,
                        top_db=top_db)
         self.model_id = model_id
-        self.pipe = DiffusionPipeline.from_pretrained(self.model_id)
+        try:  # a bit hacky
+            self.pipe = LatentAudioDiffusionPipeline.from_pretrained(self.model_id)
+        except:
+            self.pipe = AudioDiffusionPipeline.from_pretrained(self.model_id)
+
         if cuda:
             self.pipe.to("cuda")
         self.progress_bar = progress_bar or (lambda _: _)
@@ -55,20 +60,18 @@ class AudioDiffusion:
         """Generate random mel spectrogram and convert to audio.
 
         Args:
+            steps (int): number of de-noising steps to perform (defaults to num_train_timesteps)
             generator (torch.Generator): random number generator or None
 
         Returns:
             PIL Image: mel spectrogram
             (float, np.ndarray): sample rate and raw audio
         """
-        images = self.pipe(output_type="numpy",
-                           generator=generator,
-                           num_inference_steps=self.pipe.scheduler.
-                           num_train_timesteps)["sample"]
-        images = (images * 255).round().astype("uint8").transpose(0, 3, 1, 2)
-        image = Image.fromarray(images[0][0])
-        audio = self.mel.image_to_audio(image)
-        return image, (self.mel.get_sample_rate(), audio)
+        images, (sample_rate, audios) = self.pipe(mel=self.mel,
+                                                  batch_size=1,
+                                                  steps=steps,
+                                                  generator=generator)
+        return images[0], (sample_rate, audios[0])
 
     @torch.no_grad()
     def generate_spectrogram_and_audio_from_audio(
@@ -99,51 +102,124 @@ class AudioDiffusion:
             (float, np.ndarray): sample rate and raw audio
         """
 
-        # It would be better to derive a class from DiffusionPipeline
-        # but currently the return type ImagePipelineOutput cannot be imported
+        images, (sample_rate,
+                 audios) = self.pipe(mel=self.mel,
+                                     batch_size=1,
+                                     audio_file=audio_file,
+                                     raw_audio=raw_audio,
+                                     slice=slice,
+                                     start_step=start_step,
+                                     steps=steps,
+                                     generator=generator,
+                                     mask_start_secs=mask_start_secs,
+                                     mask_end_secs=mask_end_secs)
+        return images[0], (sample_rate, audios[0])
+
+    @staticmethod
+    def loop_it(audio: np.ndarray,
+                sample_rate: int,
+                loops: int = 12) -> np.ndarray:
+        """Loop audio
+
+        Args:
+            audio (np.ndarray): audio as numpy array
+            sample_rate (int): sample rate of audio
+            loops (int): number of times to loop
+
+        Returns:
+            (float, np.ndarray): sample rate and raw audio or None
+        """
+        _, beats = beat_track(y=audio, sr=sample_rate, units='samples')
+        for beats_in_bar in [16, 12, 8, 4]:
+            if len(beats) > beats_in_bar:
+                return np.tile(audio[beats[0]:beats[beats_in_bar]], loops)
+        return None
+
+
+class AudioDiffusionPipeline(DiffusionPipeline):
+
+    def __init__(self, unet: UNet2DConditionModel,
+                 scheduler: Union[DDIMScheduler, DDPMScheduler]):
+        super().__init__()
+        self.register_modules(unet=unet, scheduler=scheduler)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        mel: Mel,
+        batch_size: int = 1,
+        audio_file: str = None,
+        raw_audio: np.ndarray = None,
+        slice: int = 0,
+        start_step: int = 0,
+        steps: int = None,
+        generator: torch.Generator = None,
+        mask_start_secs: float = 0,
+        mask_end_secs: float = 0
+    ) -> Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]]:
+        """Generate random mel spectrogram from audio input and convert to audio.
+
+        Args:
+            mel (Mel): instance of Mel class to perform image <-> audio
+            batch_size (int): number of samples to generate
+            audio_file (str): must be a file on disk due to Librosa limitation or
+            raw_audio (np.ndarray): audio as numpy array
+            slice (int): slice number of audio to convert
+            start_step (int): step to start from
+            steps (int): number of de-noising steps to perform (defaults to num_train_timesteps)
+            generator (torch.Generator): random number generator or None
+            mask_start_secs (float): number of seconds of audio to mask (not generate) at start
+            mask_end_secs (float): number of seconds of audio to mask (not generate) at end
+
+        Returns:
+            List[PIL Image]: mel spectrograms
+            (float, List[np.ndarray]): sample rate and raw audios
+        """
+
         if steps is None:
-            steps = self.pipe.scheduler.num_train_timesteps
+            steps = self.scheduler.num_train_timesteps
         # Unfortunately, the schedule is set up in the constructor
-        scheduler = self.pipe.scheduler.__class__(num_train_timesteps=steps)
+        scheduler = self.scheduler.__class__(num_train_timesteps=steps)
         scheduler.set_timesteps(steps)
         mask = None
         images = noise = torch.randn(
-            (1, self.pipe.unet.in_channels, self.pipe.unet.sample_size,
-             self.pipe.unet.sample_size),
+            (batch_size, self.unet.in_channels, self.unet.sample_size,
+             self.unet.sample_size),
             generator=generator)
 
         if audio_file is not None or raw_audio is not None:
-            self.mel.load_audio(audio_file, raw_audio)
-            input_image = self.mel.audio_slice_to_image(slice)
+            mel.load_audio(audio_file, raw_audio)
+            input_image = mel.audio_slice_to_image(slice)
             input_image = np.frombuffer(input_image.tobytes(),
                                         dtype="uint8").reshape(
                                             (input_image.height,
                                              input_image.width))
             input_image = ((input_image / 255) * 2 - 1)
+            input_images = np.tile(input_image, (batch_size, 1, 1, 1))
 
-            if hasattr(self.pipe, 'vqvae'):
-                input_image = self.pipe.vqvae.encode(
-                    input_image).latent_dist.sample(generator=generator)
-                input_image = 0.18215 * input_image
+            if hasattr(self, 'vqvae'):
+                input_images = self.vqvae.encode(
+                    input_images).latent_dist.sample(generator=generator)
+                input_images = 0.18215 * input_images
 
             if start_step > 0:
                 images[0, 0] = scheduler.add_noise(
-                    torch.tensor(input_image[np.newaxis, np.newaxis, :]),
+                    torch.tensor(input_images[:, np.newaxis, np.newaxis, :]),
                     noise, torch.tensor(steps - start_step))
 
-            pixels_per_second = (self.mel.get_sample_rate() *
-                                 self.pipe.unet.sample_size /
-                                 self.mel.hop_length / self.mel.x_res)
+            pixels_per_second = (mel.get_sample_rate() *
+                                 self.unet.sample_size / 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 = scheduler.add_noise(
-                torch.tensor(input_image[np.newaxis, np.newaxis, :]), noise,
+                torch.tensor(input_images[:, np.newaxis, :]), noise,
                 torch.tensor(scheduler.timesteps[start_step:]))
 
-        images = images.to(self.pipe.device)
+        images = images.to(self.device)
         for step, t in enumerate(
                 self.progress_bar(scheduler.timesteps[start_step:])):
-            model_output = self.pipe.unet(images, t)['sample']
+            model_output = self.unet(images, t)['sample']
             images = scheduler.step(model_output,
                                     t,
                                     images,
@@ -151,45 +227,36 @@ class AudioDiffusion:
 
             if mask is not None:
                 if mask_start > 0:
-                    images[0, 0, :, :mask_start] = mask[step,
-                                                        0, :, :mask_start]
+                    images[:, :, :, :mask_start] = mask[
+                        step, :, :, :, :mask_start]
                 if mask_end > 0:
-                    images[0, 0, :, -mask_end:] = mask[step, 0, :, -mask_end:]
+                    images[:, :, :, -mask_end:] = mask[step, :, :, :,
+                                                       -mask_end:]
 
-        if hasattr(self.pipe, 'vqvae'):
+        if hasattr(self, 'vqvae'):
             # 0.18215 was scaling factor used in training to ensure unit variance
-            # This is also currently hardcoded in diffusers pipeline
             images = 1 / 0.18215 * images
-            images = self.pipe.vqvae.decode(images)['sample']
+            images = self.vqvae.decode(images)['sample']
 
         images = (images / 2 + 0.5).clamp(0, 1)
         images = images.cpu().permute(0, 2, 3, 1).numpy()
-        images = (images * 255).round().astype("uint8").transpose(0, 3, 1, 2)
-        image = Image.fromarray(images[0][0])
+        images = (images * 255).round().astype("uint8")
+        images = list(
+            map(lambda _: Image.fromarray(_[:, :, 0]), images) if images.
+            shape[3] == 1 else map(
+                lambda _: Image.fromarray(_, mode='RGB').convert('L'), images))
 
-        if hasattr(self.pipe,
-                   'vqvae') and self.pipe.vqvae.config['out_channels'] == 3:
-            image = image.convert('L')
+        audios = list(map(lambda _: mel.image_to_audio(_), images))
+        return images, (mel.get_sample_rate(), audios)
 
-        audio = self.mel.image_to_audio(image)
-        return image, (self.mel.get_sample_rate(), audio)
 
-    @staticmethod
-    def loop_it(audio: np.ndarray,
-                sample_rate: int,
-                loops: int = 12) -> np.ndarray:
-        """Loop audio
+class LatentAudioDiffusionPipeline(AudioDiffusionPipeline):
 
-        Args:
-            audio (np.ndarray): audio as numpy array
-            sample_rate (int): sample rate of audio
-            loops (int): number of times to loop
+    def __init__(self, unet: UNet2DConditionModel,
+                 scheduler: Union[DDIMScheduler,
+                                  DDPMScheduler], vqvae: AutoencoderKL):
+        super().__init__(unet=unet, scheduler=scheduler)
+        self.register_modules(vqvae=vqvae)
 
-        Returns:
-            (float, np.ndarray): sample rate and raw audio or None
-        """
-        _, beats = beat_track(y=audio, sr=sample_rate, units='samples')
-        for beats_in_bar in [16, 12, 8, 4]:
-            if len(beats) > beats_in_bar:
-                return np.tile(audio[beats[0]:beats[beats_in_bar]], loops)
-        return None
+    def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)

config/ldm_autoencoder_kl.yaml

@@ -14,12 +14,12 @@ model:
 
     ddconfig:
       double_z: True
-      z_channels: 3
+      z_channels: 4
       resolution: 256
       in_channels: 3
       out_ch: 3
       ch: 128
-      ch_mult: [ 1,2,4 ]  # num_down = len(ch_mult)-1
+      ch_mult: [ 1,2,4,4 ]  # num_down = len(ch_mult)-1
       num_res_blocks: 2
       attn_resolutions: [ ]
       dropout: 0.0

scripts/train_unconditional.py

@@ -5,12 +5,11 @@ import os
 
 import torch
 import torch.nn.functional as F
-from PIL import Image
 
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from datasets import load_from_disk, load_dataset
-from diffusers import (DDPMPipeline, DDPMScheduler, UNet2DModel, LDMPipeline,
+from diffusers import (DiffusionPipeline, DDPMScheduler, UNet2DModel,
                        DDIMScheduler, AutoencoderKL)
 from diffusers.hub_utils import init_git_repo, push_to_hub
 from diffusers.optimization import get_scheduler
@@ -23,10 +22,12 @@ from torchvision.transforms import (
     Resize,
     ToTensor,
 )
+import numpy as np
 from tqdm.auto import tqdm
 from librosa.util import normalize
 
 from audiodiffusion.mel import Mel
+from audiodiffusion import LatentAudioDiffusionPipeline, AudioDiffusionPipeline
 
 logger = get_logger(__name__)
 
@@ -45,7 +46,7 @@ def main(args):
         vqvae = AutoencoderKL.from_pretrained(args.vae)
 
     if args.from_pretrained is not None:
-        model = DDPMPipeline.from_pretrained(args.from_pretrained).unet
+        model = DiffusionPipeline.from_pretrained(args.from_pretrained).unet
     else:
         model = UNet2DModel(
             sample_size=args.resolution
@@ -237,12 +238,14 @@ def main(args):
         if accelerator.is_main_process:
             if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
                 if args.vae is not None:
-                    pipeline = LDMPipeline(unet=accelerator.unwrap_model(
-                        ema_model.averaged_model if args.use_ema else model),
-                                           vqvae=vqvae,
-                                           scheduler=noise_scheduler)
+                    pipeline = LatentAudioDiffusionPipeline(
+                        unet=accelerator.unwrap_model(
+                            ema_model.averaged_model if args.use_ema else model
+                        ),
+                        vqvae=vqvae,
+                        scheduler=noise_scheduler)
                 else:
-                    pipeline = DDPMPipeline(
+                    pipeline = AudioDiffusionPipeline(
                         unet=accelerator.unwrap_model(
                             ema_model.averaged_model if args.use_ema else model
                         ),
@@ -267,33 +270,27 @@ def main(args):
             if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
                 generator = torch.manual_seed(42)
                 # run pipeline in inference (sample random noise and denoise)
-                with torch.no_grad():
-                    images = pipeline(
-                        generator=generator,
-                        batch_size=args.eval_batch_size,
-                        output_type="numpy",
-                        num_inference_steps=args.num_train_steps,
-                    )["sample"]
+                images, (sample_rate, audios) = pipeline(
+                    mel=mel,
+                    generator=generator,
+                    batch_size=args.eval_batch_size,
+                    steps=args.num_train_steps,
+                )
 
                 # denormalize the images and save to tensorboard
-                images_processed = ((images *
-                                     255).round().astype("uint8").transpose(
-                                         0, 3, 1, 2))
+                images = np.array([
+                    np.frombuffer(image.tobytes(), dtype="uint8").reshape(
+                        (len(image.getbands()), image.height, image.width))
+                    for image in images
+                ])
                 accelerator.trackers[0].writer.add_images(
-                    "test_samples", images_processed, epoch)
-                for _, image in enumerate(images_processed):
-                    image = Image.fromarray(image[0])
-
-                    if args.vae is not None and vqvae.config[
-                            'out_channels'] == 3:
-                        image = image.convert('L')
-
-                    audio = mel.image_to_audio(image)
+                    "test_samples", images, epoch)
+                for _, audio in enumerate(audios):
                     accelerator.trackers[0].writer.add_audio(
                         f"test_audio_{_}",
                         normalize(audio),
                         epoch,
-                        sample_rate=mel.get_sample_rate(),
+                        sample_rate=sample_rate,
                     )
         accelerator.wait_for_everyone()
 
@@ -353,7 +350,7 @@ 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=64)
+    parser.add_argument("--latent_resolution", type=int, default=None)
     parser.add_argument("--scheduler",
                         type=str,
                         default="ddpm",

scripts/train_vae.py

@@ -1,10 +1,8 @@
 # pip install -e git+https://github.com/CompVis/stable-diffusion.git@master
 # pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
-# convert_original_stable_diffusion_to_diffusers.py
 
 # TODO
 # grayscale
-# update generate from audio to include vae step
 
 import os
 import argparse