fixes

README.md

@@ -71,7 +71,7 @@ python scripts/audio_to_images.py \
   --output_dir data/audio-diffusion-256 \
   --push_to_hub teticio/audio-diffusion-256
 ```
-
+  
 ## Train model
 #### Run training on local machine.
 ```bash
@@ -123,7 +123,7 @@ accelerate launch --config_file config/accelerate_sagemaker.yaml \
   --mixed_precision no
 ```
 ## Latent Audio Diffusion
-Rather than denoising images directly, it is interesting to work in the "latent space" after first encoding images using an autoencoder. This has a number of advantages. Firstly, the information in the images is compressed into a latent space of a much lower dimension, so it is much faster to train denoising diffusion models and run inference with them. Secondly, as the latent space is really a array (tensor) of guassian variables with a particular mean, the decoder is invariant to guassian noise. And thirdly, similar images tend to be clustered together and interpolating between two images in latent space can produce meaningful combinations.
+Rather than denoising images directly, it is interesting to work in the "latent space" after first encoding images using an autoencoder. This has a number of advantages. Firstly, the information in the images is compressed into a latent space of a much lower dimension, so it is much faster to train denoising diffusion models and run inference with them. Secondly, similar images tend to be clustered together and interpolating between two images in latent space can produce meaningful combinations.
 
 At the time of writing, the Hugging Face `diffusers` library is geared towards inference and lacking in training functionality, rather like its cousin `transformers` in the early days of development. In order to train a VAE (Variational Autoencoder), I use the [stable-diffusion](https://github.com/CompVis/stable-diffusion) repo from CompVis and convert the checkpoints to `diffusers` format. Note that it uses a perceptual loss function for images; it would be nice to try a perceptual *audio* loss function.
 

audiodiffusion/__init__.py

@@ -100,10 +100,10 @@ class AudioDiffusion:
         """
 
         # It would be better to derive a class from DiffusionPipeline
-        # but currently the return type ImagePipelineOutput cannot be imported.
+        # but currently the return type ImagePipelineOutput cannot be imported
         if steps is None:
             steps = self.pipe.scheduler.num_train_timesteps
-        # Unfortunately, the schedule is set up in the constructor.
+        # Unfortunately, the schedule is set up in the constructor
         scheduler = self.pipe.scheduler.__class__(num_train_timesteps=steps)
         scheduler.set_timesteps(steps)
         mask = None
@@ -121,15 +121,21 @@ class AudioDiffusion:
                                              input_image.width))
             input_image = ((input_image / 255) * 2 - 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 start_step > 0:
                 images[0, 0] = scheduler.add_noise(
                     torch.tensor(input_image[np.newaxis, np.newaxis, :]),
                     noise, torch.tensor(steps - start_step))
 
-            mask_start = int(mask_start_secs * self.mel.get_sample_rate() /
-                             self.mel.hop_length)
-            mask_end = int(mask_end_secs * self.mel.get_sample_rate() /
-                           self.mel.hop_length)
+            pixels_per_second = (self.mel.get_sample_rate() *
+                                 self.pipe.unet.sample_size /
+                                 self.mel.hop_length / self.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(scheduler.timesteps[start_step:]))
@@ -150,11 +156,21 @@ class AudioDiffusion:
                 if mask_end > 0:
                     images[0, 0, :, -mask_end:] = mask[step, 0, :, -mask_end:]
 
+        if hasattr(self.pipe, '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 = (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])
+
+        if hasattr(self.pipe,
+                   'vqvae') and self.pipe.vqvae.config['out_channels'] == 3:
+            image = image.convert('L')
+
         audio = self.mel.image_to_audio(image)
         return image, (self.mel.get_sample_rate(), audio)
 

scripts/train_unconditional.py

@@ -50,8 +50,10 @@ def main(args):
         model = UNet2DModel(
             sample_size=args.resolution
             if args.vae is None else args.latent_resolution,
-            in_channels=1 if args.vae is None else 3,
-            out_channels=1 if args.vae is None else 3,
+            in_channels=1
+            if args.vae is None else vqvae.config['latent_channels'],
+            out_channels=1
+            if args.vae is None else vqvae.config['latent_channels'],
             layers_per_block=2,
             block_out_channels=(128, 128, 256, 256, 512, 512),
             down_block_types=(
@@ -115,9 +117,9 @@ def main(args):
         )
 
     def transforms(examples):
-        if args.vae is not None:
+        if args.vae is not None and vqvae.config['in_channels'] == 3:
             images = [
-                augmentations(image.convert("RGB"))
+                augmentations(image.convert('RGB'))
                 for image in examples["image"]
             ]
         else:
@@ -182,6 +184,8 @@ def main(args):
                 with torch.no_grad():
                     clean_images = vqvae.encode(
                         clean_images).latent_dist.sample()
+                # Scale latent images to ensure approximately unit variance
+                clean_images = clean_images * 0.18215
 
             # Sample noise that we'll add to the images
             noise = torch.randn(clean_images.shape).to(clean_images.device)
@@ -231,9 +235,7 @@ def main(args):
 
         # Generate sample images for visual inspection
         if accelerator.is_main_process:
-            if (
-                    epoch + 1
-            ) % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
+            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),
@@ -262,14 +264,16 @@ def main(args):
                 else:
                     pipeline.save_pretrained(output_dir)
 
-                generator = torch.manual_seed(0)
+            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)
-                images = pipeline(
-                    generator=generator,
-                    batch_size=args.eval_batch_size,
-                    output_type="numpy",
-                    num_inference_steps=args.num_train_steps,
-                )["sample"]
+                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"]
 
                 # denormalize the images and save to tensorboard
                 images_processed = ((images *
@@ -278,7 +282,13 @@ def main(args):
                 accelerator.trackers[0].writer.add_images(
                     "test_samples", images_processed, epoch)
                 for _, image in enumerate(images_processed):
-                    audio = mel.image_to_audio(Image.fromarray(image[0]))
+                    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)
                     accelerator.trackers[0].writer.add_audio(
                         f"test_audio_{_}",
                         normalize(audio),

scripts/train_vae.py

@@ -4,6 +4,7 @@
 
 # TODO
 # grayscale
+# update generate from audio to include vae step
 
 import os
 import argparse