Merge pull request #9 from teticio/latent-audio-diffusion

.gitignore

@@ -1,8 +1,11 @@
 .vscode
 __pycache__
 .ipynb_checkpoints
-data*
-ddpm-ema-audio-*
+data
+models
 flagged
 build
 audiodiffusion.egg-info
+lightning_logs
+taming
+checkpoints

README.md

@@ -15,7 +15,10 @@ license: gpl-3.0
 
 ---
 
-**UPDATES**: 
+**UPDATES**:
+
+15/10/2022
+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").
@@ -45,35 +48,39 @@ You can play around with some pretrained models on [Google Colab](https://colab.
 ---
 
 ## Generate Mel spectrogram dataset from directory of audio files
+#### Install
+```bash
+pip install .
+```
+
 #### Training can be run with Mel spectrograms of resolution 64x64 on a single commercial grade GPU (e.g. RTX 2080 Ti). The `hop_length` should be set to 1024 for better results.
 
 ```bash
-python audio_to_images.py \
+python scripts/audio_to_images.py \
   --resolution 64 \
   --hop_length 1024 \
   --input_dir path-to-audio-files \
-  --output_dir data-test
+  --output_dir path-to-output-data
 ```
 
 #### Generate dataset of 256x256 Mel spectrograms and push to hub (you will need to be authenticated with `huggingface-cli login`).
-
 ```bash
-python audio_to_images.py \
+python scripts/audio_to_images.py \
   --resolution 256 \
   --input_dir path-to-audio-files \
-  --output_dir data-256 \
+  --output_dir data/audio-diffusion-256 \
   --push_to_hub teticio/audio-diffusion-256
 ```
+  
 ## Train model
 #### Run training on local machine.
-
 ```bash
-accelerate launch --config_file accelerate_local.yaml \
-  train_unconditional.py \
-  --dataset_name data-64 \
+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 ddpm-ema-audio-64 \
+  --output_dir models/ddpm-ema-audio-64 \
   --train_batch_size 16 \
   --num_epochs 100 \
   --gradient_accumulation_steps 1 \
@@ -83,13 +90,12 @@ accelerate launch --config_file accelerate_local.yaml \
 ```
 
 #### Run training on local machine with `batch_size` of 2 and `gradient_accumulation_steps` 8 to compensate, so that 256x256 resolution model fits on commercial grade GPU and push to hub.
-
 ```bash
-accelerate launch --config_file accelerate_local.yaml \
-  train_unconditional.py \
+accelerate launch --config_file config/accelerate_local.yaml \
+  scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
   --resolution 256 \
-  --output_dir ddpm-ema-audio-256 \
+  --output_dir models/audio-diffusion-256 \
   --num_epochs 100 \
   --train_batch_size 2 \
   --eval_batch_size 2 \
@@ -103,13 +109,12 @@ accelerate launch --config_file accelerate_local.yaml \
 ```
 
 #### Run training on SageMaker.
-
 ```bash
-accelerate launch --config_file accelerate_sagemaker.yaml \
-  strain_unconditional.py \
+accelerate launch --config_file config/accelerate_sagemaker.yaml \
+  scripts/train_unconditional.py \
   --dataset_name teticio/audio-diffusion-256 \
   --resolution 256 \
-  --output_dir ddpm-ema-audio-256 \
+  --output_dir models/ddpm-ema-audio-256 \
   --train_batch_size 16 \
   --num_epochs 100 \
   --gradient_accumulation_steps 1 \
@@ -117,3 +122,22 @@ accelerate launch --config_file accelerate_sagemaker.yaml \
   --lr_warmup_steps 500 \
   --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, 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.
+
+#### Train an autoencoder.
+```bash
+python scripts/train_vae.py \
+  --dataset_name teticio/audio-diffusion-256 \
+  --batch_size 2 \
+  --gradient_accumulation_steps 12
+```
+
+#### Train latent diffusion model.
+```bash
+accelerate launch ...
+  --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 DDPMPipeline, DDPMScheduler
+from diffusers import (DiffusionPipeline, DDPMPipeline, UNet2DConditionModel,
+                       DDIMScheduler, DDPMScheduler, AutoencoderKL)
 
 from .mel import Mel
 
-VERSION = "1.1.5"
+VERSION = "1.2.0"
 
 
 class AudioDiffusion:
@@ -42,29 +43,35 @@ class AudioDiffusion:
                        hop_length=hop_length,
                        top_db=top_db)
         self.model_id = model_id
-        self.ddpm = DDPMPipeline.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.ddpm.to("cuda")
+            self.pipe.to("cuda")
         self.progress_bar = progress_bar or (lambda _: _)
 
     def generate_spectrogram_and_audio(
         self,
+        steps: int = None,
         generator: torch.Generator = None
     ) -> Tuple[Image.Image, Tuple[int, np.ndarray]]:
         """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.ddpm(output_type="numpy", generator=generator)["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(
@@ -95,44 +102,124 @@ class AudioDiffusion:
             (float, np.ndarray): sample rate and raw audio
         """
 
-        # It would be better to derive a class from DDPMDiffusionPipeline
-        # 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.ddpm.scheduler.num_train_timesteps
-        scheduler = DDPMScheduler(num_train_timesteps=steps)
+            steps = self.scheduler.num_train_timesteps
+        # Unfortunately, the schedule is set up in the constructor
+        scheduler = self.scheduler.__class__(num_train_timesteps=steps)
         scheduler.set_timesteps(steps)
         mask = None
         images = noise = torch.randn(
-            (1, self.ddpm.unet.in_channels, self.ddpm.unet.sample_size,
-             self.ddpm.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, '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))
 
-            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 = (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.ddpm.device)
+        images = images.to(self.device)
         for step, t in enumerate(
                 self.progress_bar(scheduler.timesteps[start_step:])):
-            model_output = self.ddpm.unet(images, t)['sample']
+            model_output = self.unet(images, t)['sample']
             images = scheduler.step(model_output,
                                     t,
                                     images,
@@ -140,35 +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, 'vqvae'):
+            # 0.18215 was scaling factor used in training to ensure unit variance
+            images = 1 / 0.18215 * images
+            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")
+        images = list(
+            map(lambda _: Image.fromarray(_[:, :, 0]), images) if images.
+            shape[3] == 1 else map(
+                lambda _: Image.fromarray(_, mode='RGB').convert('L'), images))
 
-        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)
+        audios = list(map(lambda _: mel.image_to_audio(_), images))
+        return images, (mel.get_sample_rate(), audios)
 
-    @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
+class LatentAudioDiffusionPipeline(AudioDiffusionPipeline):
 
-        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 __init__(self, unet: UNet2DConditionModel,
+                 scheduler: Union[DDIMScheduler,
+                                  DDPMScheduler], vqvae: AutoencoderKL):
+        super().__init__(unet=unet, scheduler=scheduler)
+        self.register_modules(vqvae=vqvae)
+
+    def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)

audiodiffusion/utils.py

@@ -0,0 +1,363 @@
+# adpated from https://github.com/huggingface/diffusers/blob/main/scripts/convert_original_stable_diffusion_to_diffusers.py
+
+import torch
+from diffusers import AutoencoderKL
+
+
+def shave_segments(path, n_shave_prefix_segments=1):
+    """
+    Removes segments. Positive values shave the first segments, negative shave the last segments.
+    """
+    if n_shave_prefix_segments >= 0:
+        return ".".join(path.split(".")[n_shave_prefix_segments:])
+    else:
+        return ".".join(path.split(".")[:n_shave_prefix_segments])
+
+
+def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside resnets to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("nin_shortcut", "conv_shortcut")
+        new_item = shave_segments(
+            new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
+        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
+
+        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
+        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
+
+        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
+    """
+    Updates paths inside attentions to the new naming scheme (local renaming)
+    """
+    mapping = []
+    for old_item in old_list:
+        new_item = old_item
+
+        new_item = new_item.replace("norm.weight", "group_norm.weight")
+        new_item = new_item.replace("norm.bias", "group_norm.bias")
+
+        new_item = new_item.replace("q.weight", "query.weight")
+        new_item = new_item.replace("q.bias", "query.bias")
+
+        new_item = new_item.replace("k.weight", "key.weight")
+        new_item = new_item.replace("k.bias", "key.bias")
+
+        new_item = new_item.replace("v.weight", "value.weight")
+        new_item = new_item.replace("v.bias", "value.bias")
+
+        new_item = new_item.replace("proj_out.weight", "proj_attn.weight")
+        new_item = new_item.replace("proj_out.bias", "proj_attn.bias")
+
+        new_item = shave_segments(
+            new_item, n_shave_prefix_segments=n_shave_prefix_segments)
+
+        mapping.append({"old": old_item, "new": new_item})
+
+    return mapping
+
+
+def assign_to_checkpoint(paths,
+                         checkpoint,
+                         old_checkpoint,
+                         attention_paths_to_split=None,
+                         additional_replacements=None,
+                         config=None):
+    """
+    This does the final conversion step: take locally converted weights and apply a global renaming
+    to them. It splits attention layers, and takes into account additional replacements
+    that may arise.
+
+    Assigns the weights to the new checkpoint.
+    """
+    assert isinstance(
+        paths, list
+    ), "Paths should be a list of dicts containing 'old' and 'new' keys."
+
+    # Splits the attention layers into three variables.
+    if attention_paths_to_split is not None:
+        for path, path_map in attention_paths_to_split.items():
+            old_tensor = old_checkpoint[path]
+            channels = old_tensor.shape[0] // 3
+
+            target_shape = (-1,
+                            channels) if len(old_tensor.shape) == 3 else (-1)
+
+            num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3
+
+            old_tensor = old_tensor.reshape((num_heads, 3 * channels //
+                                             num_heads) + old_tensor.shape[1:])
+            query, key, value = old_tensor.split(channels // num_heads, dim=1)
+
+            checkpoint[path_map["query"]] = query.reshape(target_shape)
+            checkpoint[path_map["key"]] = key.reshape(target_shape)
+            checkpoint[path_map["value"]] = value.reshape(target_shape)
+
+    for path in paths:
+        new_path = path["new"]
+
+        # These have already been assigned
+        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
+            continue
+
+        # Global renaming happens here
+        new_path = new_path.replace("middle_block.0", "mid_block.resnets.0")
+        new_path = new_path.replace("middle_block.1", "mid_block.attentions.0")
+        new_path = new_path.replace("middle_block.2", "mid_block.resnets.1")
+
+        if additional_replacements is not None:
+            for replacement in additional_replacements:
+                new_path = new_path.replace(replacement["old"],
+                                            replacement["new"])
+
+        # proj_attn.weight has to be converted from conv 1D to linear
+        if "proj_attn.weight" in new_path:
+            checkpoint[new_path] = old_checkpoint[path["old"]][:, :, 0]
+        else:
+            checkpoint[new_path] = old_checkpoint[path["old"]]
+
+
+def conv_attn_to_linear(checkpoint):
+    keys = list(checkpoint.keys())
+    attn_keys = ["query.weight", "key.weight", "value.weight"]
+    for key in keys:
+        if ".".join(key.split(".")[-2:]) in attn_keys:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0, 0]
+        elif "proj_attn.weight" in key:
+            if checkpoint[key].ndim > 2:
+                checkpoint[key] = checkpoint[key][:, :, 0]
+
+
+def create_vae_diffusers_config(original_config):
+    """
+    Creates a config for the diffusers based on the config of the LDM model.
+    """
+    vae_params = original_config.model.params.ddconfig
+    _ = original_config.model.params.embed_dim
+
+    block_out_channels = [vae_params.ch * mult for mult in vae_params.ch_mult]
+    down_block_types = ["DownEncoderBlock2D"] * len(block_out_channels)
+    up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
+
+    config = dict(
+        sample_size=vae_params.resolution,
+        in_channels=vae_params.in_channels,
+        out_channels=vae_params.out_ch,
+        down_block_types=tuple(down_block_types),
+        up_block_types=tuple(up_block_types),
+        block_out_channels=tuple(block_out_channels),
+        latent_channels=vae_params.z_channels,
+        layers_per_block=vae_params.num_res_blocks,
+    )
+    return config
+
+
+def convert_ldm_vae_checkpoint(checkpoint, config):
+    # extract state dict for VAE
+    vae_state_dict = checkpoint
+
+    new_checkpoint = {}
+
+    new_checkpoint["encoder.conv_in.weight"] = vae_state_dict[
+        "encoder.conv_in.weight"]
+    new_checkpoint["encoder.conv_in.bias"] = vae_state_dict[
+        "encoder.conv_in.bias"]
+    new_checkpoint["encoder.conv_out.weight"] = vae_state_dict[
+        "encoder.conv_out.weight"]
+    new_checkpoint["encoder.conv_out.bias"] = vae_state_dict[
+        "encoder.conv_out.bias"]
+    new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict[
+        "encoder.norm_out.weight"]
+    new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict[
+        "encoder.norm_out.bias"]
+
+    new_checkpoint["decoder.conv_in.weight"] = vae_state_dict[
+        "decoder.conv_in.weight"]
+    new_checkpoint["decoder.conv_in.bias"] = vae_state_dict[
+        "decoder.conv_in.bias"]
+    new_checkpoint["decoder.conv_out.weight"] = vae_state_dict[
+        "decoder.conv_out.weight"]
+    new_checkpoint["decoder.conv_out.bias"] = vae_state_dict[
+        "decoder.conv_out.bias"]
+    new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict[
+        "decoder.norm_out.weight"]
+    new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict[
+        "decoder.norm_out.bias"]
+
+    new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+    new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+    new_checkpoint["post_quant_conv.weight"] = vae_state_dict[
+        "post_quant_conv.weight"]
+    new_checkpoint["post_quant_conv.bias"] = vae_state_dict[
+        "post_quant_conv.bias"]
+
+    # Retrieves the keys for the encoder down blocks only
+    num_down_blocks = len({
+        ".".join(layer.split(".")[:3])
+        for layer in vae_state_dict if "encoder.down" in layer
+    })
+    down_blocks = {
+        layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key]
+        for layer_id in range(num_down_blocks)
+    }
+
+    # Retrieves the keys for the decoder up blocks only
+    num_up_blocks = len({
+        ".".join(layer.split(".")[:3])
+        for layer in vae_state_dict if "decoder.up" in layer
+    })
+    up_blocks = {
+        layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key]
+        for layer_id in range(num_up_blocks)
+    }
+
+    for i in range(num_down_blocks):
+        resnets = [
+            key for key in down_blocks[i]
+            if f"down.{i}" in key and f"down.{i}.downsample" not in key
+        ]
+
+        if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+            new_checkpoint[
+                f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(
+                    f"encoder.down.{i}.downsample.conv.weight")
+            new_checkpoint[
+                f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(
+                    f"encoder.down.{i}.downsample.conv.bias")
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {
+            "old": f"down.{i}.block",
+            "new": f"down_blocks.{i}.resnets"
+        }
+        assign_to_checkpoint(paths,
+                             new_checkpoint,
+                             vae_state_dict,
+                             additional_replacements=[meta_path],
+                             config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [
+            key for key in mid_resnets if f"encoder.mid.block_{i}" in key
+        ]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {
+            "old": f"mid.block_{i}",
+            "new": f"mid_block.resnets.{i - 1}"
+        }
+        assign_to_checkpoint(paths,
+                             new_checkpoint,
+                             vae_state_dict,
+                             additional_replacements=[meta_path],
+                             config=config)
+
+    mid_attentions = [
+        key for key in vae_state_dict if "encoder.mid.attn" in key
+    ]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths,
+                         new_checkpoint,
+                         vae_state_dict,
+                         additional_replacements=[meta_path],
+                         config=config)
+    conv_attn_to_linear(new_checkpoint)
+
+    for i in range(num_up_blocks):
+        block_id = num_up_blocks - 1 - i
+        resnets = [
+            key for key in up_blocks[block_id]
+            if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+        ]
+
+        if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+            new_checkpoint[
+                f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+                    f"decoder.up.{block_id}.upsample.conv.weight"]
+            new_checkpoint[
+                f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+                    f"decoder.up.{block_id}.upsample.conv.bias"]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {
+            "old": f"up.{block_id}.block",
+            "new": f"up_blocks.{i}.resnets"
+        }
+        assign_to_checkpoint(paths,
+                             new_checkpoint,
+                             vae_state_dict,
+                             additional_replacements=[meta_path],
+                             config=config)
+
+    mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+    num_mid_res_blocks = 2
+    for i in range(1, num_mid_res_blocks + 1):
+        resnets = [
+            key for key in mid_resnets if f"decoder.mid.block_{i}" in key
+        ]
+
+        paths = renew_vae_resnet_paths(resnets)
+        meta_path = {
+            "old": f"mid.block_{i}",
+            "new": f"mid_block.resnets.{i - 1}"
+        }
+        assign_to_checkpoint(paths,
+                             new_checkpoint,
+                             vae_state_dict,
+                             additional_replacements=[meta_path],
+                             config=config)
+
+    mid_attentions = [
+        key for key in vae_state_dict if "decoder.mid.attn" in key
+    ]
+    paths = renew_vae_attention_paths(mid_attentions)
+    meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+    assign_to_checkpoint(paths,
+                         new_checkpoint,
+                         vae_state_dict,
+                         additional_replacements=[meta_path],
+                         config=config)
+    conv_attn_to_linear(new_checkpoint)
+    return new_checkpoint
+
+def convert_ldm_to_hf_vae(ldm_checkpoint, ldm_config, hf_checkpoint):
+    checkpoint = torch.load(ldm_checkpoint)["state_dict"]
+
+    # Convert the VAE model.
+    vae_config = create_vae_diffusers_config(ldm_config)
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(
+        checkpoint, vae_config)
+
+    vae = AutoencoderKL(**vae_config)
+    vae.load_state_dict(converted_vae_checkpoint)
+    vae.save_pretrained(hf_checkpoint)

config/ldm_autoencoder_kl.yaml

@@ -0,0 +1,31 @@
+
+model:
+  base_learning_rate: 4.5e-6
+  target: ldm.models.autoencoder.AutoencoderKL
+  params:
+    monitor: "val/rec_loss"
+    embed_dim: 3
+    lossconfig:
+      target: ldm.modules.losses.LPIPSWithDiscriminator
+      params:
+        disc_start: 50001
+        kl_weight: 0.000001
+        disc_weight: 0.5
+
+    ddconfig:
+      double_z: True
+      z_channels: 4
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [ 1,2,4,4 ]  # num_down = len(ch_mult)-1
+      num_res_blocks: 2
+      attn_resolutions: [ ]
+      dropout: 0.0
+
+lightning:
+  trainer:
+    benchmark: True
+    accelerator: gpu
+    devices: 1

notebooks/test_vae.ipynb

@@ -0,0 +1,169 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bcbbe26c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import sys\n",
+    "sys.path.insert(0, os.path.dirname(os.path.abspath(\"\")))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b451ab22",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import random\n",
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "from datasets import load_dataset\n",
+    "from IPython.display import Audio\n",
+    "from diffusers import AutoencoderKL\n",
+    "from audiodiffusion.mel import Mel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "324cef44",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mel = Mel()\n",
+    "vae = AutoencoderKL.from_pretrained('../models/autoencoder-kl')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da55ce79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "vae.config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5fea99ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds = load_dataset('teticio/audio-diffusion-256')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "426c6edd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image = random.choice(ds['train'])['image']\n",
+    "display(image)\n",
+    "Audio(data=mel.image_to_audio(image), rate=mel.get_sample_rate())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d123f8a0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# encode\n",
+    "input_image = np.frombuffer(image.convert('RGB').tobytes(), dtype=\"uint8\").reshape(\n",
+    "    (image.height, image.width, 3))\n",
+    "input_image = ((input_image / 255) * 2 - 1).transpose(2, 0, 1)\n",
+    "posterior = vae.encode(torch.tensor([input_image], dtype=torch.float32)).latent_dist\n",
+    "latents = posterior.sample()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "482c458f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# reconstruct\n",
+    "output_image = vae.decode(latents)['sample']\n",
+    "output_image = torch.clamp(output_image, -1., 1.)\n",
+    "output_image = (output_image + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w\n",
+    "output_image = (output_image.detach().cpu().numpy() *\n",
+    "                255).round().astype(\"uint8\").transpose(0, 2, 3, 1)[0]\n",
+    "output_image = Image.fromarray(output_image).convert('L')\n",
+    "display(output_image)\n",
+    "Audio(data=mel.image_to_audio(output_image), rate=mel.get_sample_rate())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f10db020",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# sample\n",
+    "output_image = vae.decode(torch.randn_like(posterior.sample()))['sample']\n",
+    "output_image = torch.clamp(output_image, -1., 1.)\n",
+    "output_image = (output_image + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w\n",
+    "output_image = (output_image.detach().cpu().numpy() *\n",
+    "                255).round().astype(\"uint8\").transpose(0, 2, 3, 1)[0]\n",
+    "output_image = Image.fromarray(output_image).convert('L')\n",
+    "display(output_image)\n",
+    "Audio(data=mel.image_to_audio(output_image), rate=mel.get_sample_rate())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46019770",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "huggingface",
+   "language": "python",
+   "name": "huggingface"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.6"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": false
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

train_unconditional.py → scripts/train_unconditional.py

@@ -5,12 +5,12 @@ 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
+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
 from diffusers.training_utils import EMAModel
@@ -22,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__)
 
@@ -34,18 +36,25 @@ def main(args):
     output_dir = os.environ.get("SM_MODEL_DIR", None) or args.output_dir
     logging_dir = os.path.join(output_dir, args.logging_dir)
     accelerator = Accelerator(
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
         log_with="tensorboard",
         logging_dir=logging_dir,
     )
 
+    if args.vae is not None:
+        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,
-            in_channels=1,
-            out_channels=1,
+            sample_size=args.resolution
+            if args.vae is None else args.latent_resolution,
+            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=(
@@ -65,8 +74,14 @@ def main(args):
                 "UpBlock2D",
             ),
         )
-    noise_scheduler = DDPMScheduler(num_train_timesteps=1000,
-                                    tensor_format="pt")
+
+    if args.scheduler == "ddpm":
+        noise_scheduler = DDPMScheduler(
+            num_train_timesteps=args.num_train_steps, tensor_format="pt")
+    else:
+        noise_scheduler = DDIMScheduler(
+            num_train_timesteps=args.num_train_steps, tensor_format="pt")
+
     optimizer = torch.optim.AdamW(
         model.parameters(),
         lr=args.learning_rate,
@@ -103,7 +118,13 @@ def main(args):
         )
 
     def transforms(examples):
-        images = [augmentations(image) for image in examples["image"]]
+        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)
@@ -158,6 +179,15 @@ def main(args):
         model.train()
         for step, batch in enumerate(train_dataloader):
             clean_images = batch["input"]
+
+            if args.vae is not None:
+                vqvae.to(clean_images.device)
+                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)
             bsz = clean_images.shape[0]
@@ -180,7 +210,8 @@ def main(args):
                 loss = F.mse_loss(noise_pred, noise)
                 accelerator.backward(loss)
 
-                accelerator.clip_grad_norm_(model.parameters(), 1.0)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                 optimizer.step()
                 lr_scheduler.step()
                 if args.use_ema:
@@ -188,6 +219,8 @@ def main(args):
                 optimizer.zero_grad()
 
             progress_bar.update(1)
+            global_step += 1
+
             logs = {
                 "loss": loss.detach().item(),
                 "lr": lr_scheduler.get_last_lr()[0],
@@ -197,7 +230,6 @@ def main(args):
                 logs["ema_decay"] = ema_model.decay
             progress_bar.set_postfix(**logs)
             accelerator.log(logs, step=global_step)
-            global_step += 1
         progress_bar.close()
 
         accelerator.wait_for_everyone()
@@ -205,11 +237,20 @@ 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:
-                pipeline = DDPMPipeline(
-                    unet=accelerator.unwrap_model(
-                        ema_model.averaged_model if args.use_ema else model),
-                    scheduler=noise_scheduler,
-                )
+                if args.vae is not None:
+                    pipeline = LatentAudioDiffusionPipeline(
+                        unet=accelerator.unwrap_model(
+                            ema_model.averaged_model if args.use_ema else model
+                        ),
+                        vqvae=vqvae,
+                        scheduler=noise_scheduler)
+                else:
+                    pipeline = AudioDiffusionPipeline(
+                        unet=accelerator.unwrap_model(
+                            ema_model.averaged_model if args.use_ema else model
+                        ),
+                        scheduler=noise_scheduler,
+                    )
 
                 # save the model
                 if args.push_to_hub:
@@ -226,27 +267,30 @@ 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(
+                images, (sample_rate, audios) = pipeline(
+                    mel=mel,
                     generator=generator,
                     batch_size=args.eval_batch_size,
-                    output_type="numpy",
-                )["sample"]
+                    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):
-                    audio = mel.image_to_audio(Image.fromarray(image[0]))
+                    "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()
 
@@ -268,7 +312,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=64)
+    parser.add_argument("--resolution", type=int, 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)
@@ -305,6 +349,16 @@ if __name__ == "__main__":
     parser.add_argument("--hop_length", type=int, default=512)
     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",
+                        help="ddpm or ddim")
+    parser.add_argument("--vae",
+                        type=str,
+                        default=None,
+                        help="pretrained VAE model for latent diffusion")
 
     args = parser.parse_args()
     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))

scripts/train_vae.py

@@ -0,0 +1,166 @@
+# 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
+
+# TODO
+# grayscale
+
+import os
+import argparse
+
+import torch
+import torchvision
+import numpy as np
+from PIL import Image
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+from librosa.util import normalize
+from ldm.util import instantiate_from_config
+from pytorch_lightning.trainer import Trainer
+from torch.utils.data import DataLoader, Dataset
+from datasets import load_from_disk, load_dataset
+from pytorch_lightning.callbacks import Callback, ModelCheckpoint
+from pytorch_lightning.utilities.distributed import rank_zero_only
+
+from audiodiffusion.mel import Mel
+from audiodiffusion.utils import convert_ldm_to_hf_vae
+
+
+class AudioDiffusion(Dataset):
+
+    def __init__(self, model_id):
+        super().__init__()
+        if os.path.exists(model_id):
+            self.hf_dataset = load_from_disk(model_id)['train']
+        else:
+            self.hf_dataset = load_dataset(model_id)['train']
+
+    def __len__(self):
+        return len(self.hf_dataset)
+
+    def __getitem__(self, idx):
+        image = self.hf_dataset[idx]['image'].convert('RGB')
+        image = np.frombuffer(image.tobytes(), dtype="uint8").reshape(
+            (image.height, image.width, 3))
+        image = ((image / 255) * 2 - 1)
+        return {'image': image}
+
+
+class AudioDiffusionDataModule(pl.LightningDataModule):
+
+    def __init__(self, model_id, batch_size):
+        super().__init__()
+        self.batch_size = batch_size
+        self.dataset = AudioDiffusion(model_id)
+        self.num_workers = 1
+
+    def train_dataloader(self):
+        return DataLoader(self.dataset,
+                          batch_size=self.batch_size,
+                          num_workers=self.num_workers)
+
+
+class ImageLogger(Callback):
+
+    def __init__(self, every=1000, resolution=256, hop_length=512):
+        super().__init__()
+        self.mel = Mel(x_res=resolution,
+                       y_res=resolution,
+                       hop_length=hop_length)
+        self.every = every
+
+    @rank_zero_only
+    def log_images_and_audios(self, pl_module, batch):
+        pl_module.eval()
+        with torch.no_grad():
+            images = pl_module.log_images(batch, split='train')
+        pl_module.train()
+
+        for k in images:
+            images[k] = images[k].detach().cpu()
+            images[k] = torch.clamp(images[k], -1., 1.)
+            images[k] = (images[k] + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
+            grid = torchvision.utils.make_grid(images[k])
+
+            tag = f"train/{k}"
+            pl_module.logger.experiment.add_image(
+                tag, grid, global_step=pl_module.global_step)
+
+            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'))
+                pl_module.logger.experiment.add_audio(
+                    tag + f"/{_}",
+                    normalize(audio),
+                    global_step=pl_module.global_step,
+                    sample_rate=self.mel.get_sample_rate())
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch,
+                           batch_idx):
+        if (batch_idx + 1) % self.every != 0:
+            return
+        self.log_images_and_audios(pl_module, batch)
+
+
+class HFModelCheckpoint(ModelCheckpoint):
+
+    def __init__(self, ldm_config, hf_checkpoint, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.ldm_config = ldm_config
+        self.hf_checkpoint = hf_checkpoint
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        super().on_train_epoch_end(trainer, pl_module)
+        ldm_checkpoint = self.format_checkpoint_name(
+            {'epoch': trainer.current_epoch})
+        convert_ldm_to_hf_vae(ldm_checkpoint, self.ldm_config,
+                              self.hf_checkpoint)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train VAE using ldm.")
+    parser.add_argument("-d", "--dataset_name", type=str, default=None)
+    parser.add_argument("-b", "--batch_size", type=int, default=1)
+    parser.add_argument("-c",
+                        "--ldm_config_file",
+                        type=str,
+                        default="config/ldm_autoencoder_kl.yaml")
+    parser.add_argument("--ldm_checkpoint_dir",
+                        type=str,
+                        default="models/ldm-autoencoder-kl")
+    parser.add_argument("--hf_checkpoint_dir",
+                        type=str,
+                        default="models/autoencoder-kl")
+    parser.add_argument("-r",
+                        "--resume_from_checkpoint",
+                        type=str,
+                        default=None)
+    parser.add_argument("-g",
+                        "--gradient_accumulation_steps",
+                        type=int,
+                        default=1)
+    args = parser.parse_args()
+
+    config = OmegaConf.load(args.ldm_config_file)
+    lightning_config = config.pop("lightning", OmegaConf.create())
+    trainer_config = lightning_config.get("trainer", OmegaConf.create())
+    trainer_config.accumulate_grad_batches = args.gradient_accumulation_steps
+    trainer_opt = argparse.Namespace(**trainer_config)
+    trainer = Trainer.from_argparse_args(
+        trainer_opt,
+        resume_from_checkpoint=args.resume_from_checkpoint,
+        callbacks=[
+            ImageLogger(),
+            HFModelCheckpoint(ldm_config=config,
+                              hf_checkpoint=args.hf_checkpoint_dir,
+                              dirpath=args.ldm_checkpoint_dir,
+                              filename='{epoch:06}',
+                              verbose=True,
+                              save_last=True)
+        ])
+    model = instantiate_from_config(config.model)
+    model.learning_rate = config.model.base_learning_rate
+    data = AudioDiffusionDataModule(args.dataset_name,
+                                    batch_size=args.batch_size)
+    trainer.fit(model, data)