add conditional training

README.md

@@ -23,7 +23,9 @@ Go to https://soundcloud.com/teticio2/sets/audio-diffusion-loops for more exampl
 ---
 #### Updates
 
-**5/12/2022** 🤗 Exciting news! `AudioDiffusionPipeline` has been migrated to the Hugging Face `diffusers` package so that it is even easier for others to use and contribute.
+**25/12/2022**. Now it is possible to train models conditional on an encoding (of text or audio, for example). See the section on Conditional Audio Generation below.
+
+**5/12/2022**. 🤗 Exciting news! `AudioDiffusionPipeline` has been migrated to the Hugging Face `diffusers` package so that it is even easier for others to use and contribute.
 
 **2/12/2022**. Added Mel to pipeline and updated the pretrained models to save Mel config (they are now no longer compatible with previous versions of this repo). It is relatively straightforward to migrate previously trained models to the new format (see https://huggingface.co/teticio/audio-diffusion-256).
 
@@ -58,7 +60,8 @@ You can play around with some pre-trained models on [Google Colab](https://colab
 | [teticio/audio-diffusion-instrumental-hiphop-256](https://huggingface.co/teticio/audio-diffusion-instrumental-hiphop-256) | [teticio/audio-diffusion-instrumental-hiphop-256](https://huggingface.co/datasets/teticio/audio-diffusion-instrumental-hiphop-256) | Instrumental Hip Hop music |
 | [teticio/audio-diffusion-ddim-256](https://huggingface.co/teticio/audio-diffusion-ddim-256) | [teticio/audio-diffusion-256](https://huggingface.co/datasets/teticio/audio-diffusion-256) | De-noising Diffusion Implicit Model |
 | [teticio/latent-audio-diffusion-256](https://huggingface.co/teticio/latent-audio-diffusion-256) | [teticio/audio-diffusion-256](https://huggingface.co/datasets/teticio/audio-diffusion-256) | Latent Audio Diffusion model |
-| [teticio/latent-audio-diffusion-ddim-256](https://huggingface.co/teticio/latent-audio-diffusion-ddim-256) | [teticio/audio-diffusion-256](https://huggingface.co/datasets/teticio/audio-diffusion-256) | Latent Audio Diffusion De-noising Diffusion Implicit Model |
+| [teticio/latent-audio-diffusion-ddim-256](https://huggingface.co/teticio/latent-audio-diffusion-ddim-256) | [teticio/audio-diffusion-256](https://huggingface.co/datasets/teticio/audio-diffusion-256) | Latent Audio Diffusion Implicit Model |
+| [teticio/conditional-latent-audio-diffusion-512](https://huggingface.co/teticio/latent-audio-diffusion-512) | [teticio/audio-diffusion-512](https://huggingface.co/datasets/teticio/audio-diffusion-512) | Conditional Latent Audio Diffusion Model |
 
 ---
 
@@ -106,7 +109,7 @@ Note that the default `sample_rate` is 22050 and audios will be resampled if the
 
 ```bash
 accelerate launch --config_file config/accelerate_local.yaml \
-scripts/train_unconditional.py \
+scripts/train_unet.py \
 --dataset_name data/audio-diffusion-64 \
 --hop_length 1024 \
 --output_dir models/ddpm-ema-audio-64 \
@@ -122,7 +125,7 @@ scripts/train_unconditional.py \
 
 ```bash
 accelerate launch --config_file config/accelerate_local.yaml \
-scripts/train_unconditional.py \
+scripts/train_unet.py \
 --dataset_name teticio/audio-diffusion-256 \
 --output_dir models/audio-diffusion-256 \
 --num_epochs 100 \
@@ -141,7 +144,7 @@ scripts/train_unconditional.py \
 
 ```bash
 accelerate launch --config_file config/accelerate_sagemaker.yaml \
-scripts/train_unconditional.py \
+scripts/train_unet.py \
 --dataset_name teticio/audio-diffusion-256 \
 --output_dir models/ddpm-ema-audio-256 \
 --train_batch_size 16 \
@@ -200,3 +203,33 @@ accelerate launch ...
 ...
 --vae models/autoencoder-kl
 ```
+
+## Conditional Audio Generation
+
+We can generate audio conditional on a text prompt - or indeed anything which can be encoded into a bunch of numbers - much like DALL-E2 and Midjourney. It is generally harder to find good quality datasets of audios together with descriptions, although the people behind the dataset used to train Midjourney are making some very interesting progress [here](https://github.com/LAION-AI/audio-dataset). I have chosen to encode the audio directly instead based on "how it sounds", using a [model which I trained on hundreds of thousands of Spotify playlists](https://github.com/teticio/Deej-AI). To encode an audio into a 100 dimensional vector
+
+```python
+from diffusers import Mel
+from audiodiffusion.audio_encoder import AudioEncoder
+
+audio_encoder = AudioEncoder.from_pretrained("teticio/audio-encoder")
+audio_encoder.encode(['/home/teticio/Music/liked/Agua Re - Holy Dance - Large Sound Mix.mp3'])
+```
+
+One you have prepared a dataset, you can encode the audio files with this script
+
+```bash
+python scripts/encode_audio \
+--dataset_name teticio/audio-diffusion-256 \
+--out_file data/encodings.p
+```
+
+Then you can train a model with
+
+```bash
+accelerate launch ...
+...
+--encodings data/encodings.p
+```
+
+When generating audios, you will need to pass an `encodings` Tensor. See the [`conditional_generation.ipynb`](https://colab.research.google.com/github/teticio/audio-diffusion/blob/master/notebooks/conditional_generation.ipynb) notebook for an example that uses encodings of Spotify track previews to influence the generation.

audiodiffusion/__init__.py

@@ -1,21 +1,24 @@
 from typing import Iterable, Tuple
 
-import torch
 import numpy as np
+import torch
+from librosa.beat import beat_track
 from PIL import Image
 from tqdm.auto import tqdm
-from librosa.beat import beat_track
-from diffusers import AudioDiffusionPipeline
 
-VERSION = "1.3.2"
+# from diffusers import AudioDiffusionPipeline
+from .pipeline_audio_diffusion import AudioDiffusionPipeline
 
+VERSION = "1.4.0"
 
-class AudioDiffusion:
 
-    def __init__(self,
-                 model_id: str = "teticio/audio-diffusion-256",
-                 cuda: bool = torch.cuda.is_available(),
-                 progress_bar: Iterable = tqdm):
+class AudioDiffusion:
+    def __init__(
+        self,
+        model_id: str = "teticio/audio-diffusion-256",
+        cuda: bool = torch.cuda.is_available(),
+        progress_bar: Iterable = tqdm,
+    ):
         """Class for generating audio using De-noising Diffusion Probabilistic Models.
 
         Args:
@@ -35,7 +38,8 @@ class AudioDiffusion:
         generator: torch.Generator = None,
         step_generator: torch.Generator = None,
         eta: float = 0,
-        noise: torch.Tensor = None
+        noise: torch.Tensor = None,
+        encoding: torch.Tensor = None,
     ) -> Tuple[Image.Image, Tuple[int, np.ndarray]]:
         """Generate random mel spectrogram and convert to audio.
 
@@ -45,19 +49,22 @@ class AudioDiffusion:
             step_generator (torch.Generator): random number generator used to de-noise or None
             eta (float): parameter between 0 and 1 used with DDIM scheduler
             noise (torch.Tensor): noisy image or None
+            encoding (`torch.Tensor`): for UNet2DConditionModel shape (batch_size, seq_length, cross_attention_dim)
 
         Returns:
             PIL Image: mel spectrogram
             (float, np.ndarray): sample rate and raw audio
         """
-        images, (sample_rate,
-                 audios) = self.pipe(batch_size=1,
-                                     steps=steps,
-                                     generator=generator,
-                                     step_generator=step_generator,
-                                     eta=eta,
-                                     noise=noise,
-                                     return_dict=False)
+        images, (sample_rate, audios) = self.pipe(
+            batch_size=1,
+            steps=steps,
+            generator=generator,
+            step_generator=step_generator,
+            eta=eta,
+            noise=noise,
+            encoding=encoding,
+            return_dict=False,
+        )
         return images[0], (sample_rate, audios[0])
 
     def generate_spectrogram_and_audio_from_audio(
@@ -72,7 +79,8 @@ class AudioDiffusion:
         mask_end_secs: float = 0,
         step_generator: torch.Generator = None,
         eta: float = 0,
-        noise: torch.Tensor = None
+        encoding: torch.Tensor = None,
+        noise: torch.Tensor = None,
     ) -> Tuple[Image.Image, Tuple[int, np.ndarray]]:
         """Generate random mel spectrogram from audio input and convert to audio.
 
@@ -87,6 +95,7 @@ class AudioDiffusion:
             mask_end_secs (float): number of seconds of audio to mask (not generate) at end
             step_generator (torch.Generator): random number generator used to de-noise or None
             eta (float): parameter between 0 and 1 used with DDIM scheduler
+            encoding (`torch.Tensor`): for UNet2DConditionModel shape (batch_size, seq_length, cross_attention_dim)
             noise (torch.Tensor): noisy image or None
 
         Returns:
@@ -94,26 +103,26 @@ class AudioDiffusion:
             (float, np.ndarray): sample rate and raw audio
         """
 
-        images, (sample_rate,
-                 audios) = self.pipe(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,
-                                     step_generator=step_generator,
-                                     eta=eta,
-                                     noise=noise,
-                                     return_dict=False)
+        images, (sample_rate, audios) = self.pipe(
+            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,
+            step_generator=step_generator,
+            eta=eta,
+            noise=noise,
+            encoding=encoding,
+            return_dict=False,
+        )
         return images[0], (sample_rate, audios[0])
 
     @staticmethod
-    def loop_it(audio: np.ndarray,
-                sample_rate: int,
-                loops: int = 12) -> np.ndarray:
+    def loop_it(audio: np.ndarray, sample_rate: int, loops: int = 12) -> np.ndarray:
         """Loop audio
 
         Args:
@@ -124,403 +133,8 @@ class AudioDiffusion:
         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)
+        _, beats = beat_track(y=audio, sr=sample_rate, units="samples")
+        beats_in_bar = (len(beats) - 1) // 4 * 4
+        if beats_in_bar > 0:
+            return np.tile(audio[beats[0] : beats[beats_in_bar]], loops)
         return None
-
-
-'''
-# This code will be migrated to diffusers shortly
-
-#-----------------------------------------------------------------------------#
-
-import os
-import warnings
-from typing import Any, Dict, Optional, Union
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.schedulers.scheduling_utils import SchedulerMixin
-
-
-warnings.filterwarnings("ignore")
-
-import numpy as np  # noqa: E402
-
-import librosa  # noqa: E402
-from PIL import Image  # noqa: E402
-
-
-class Mel(ConfigMixin, SchedulerMixin):
-    """
-    Parameters:
-        x_res (`int`): x resolution of spectrogram (time)
-        y_res (`int`): y resolution of spectrogram (frequency bins)
-        sample_rate (`int`): sample rate of audio
-        n_fft (`int`): number of Fast Fourier Transforms
-        hop_length (`int`): hop length (a higher number is recommended for lower than 256 y_res)
-        top_db (`int`): loudest in decibels
-        n_iter (`int`): number of iterations for Griffin Linn mel inversion
-    """
-
-    config_name = "mel_config.json"
-
-    @register_to_config
-    def __init__(
-        self,
-        x_res: int = 256,
-        y_res: int = 256,
-        sample_rate: int = 22050,
-        n_fft: int = 2048,
-        hop_length: int = 512,
-        top_db: int = 80,
-        n_iter: int = 32,
-    ):
-        self.hop_length = hop_length
-        self.sr = sample_rate
-        self.n_fft = n_fft
-        self.top_db = top_db
-        self.n_iter = n_iter
-        self.set_resolution(x_res, y_res)
-        self.audio = None
-
-    def set_resolution(self, x_res: int, y_res: int):
-        """Set resolution.
-
-        Args:
-            x_res (`int`): x resolution of spectrogram (time)
-            y_res (`int`): y resolution of spectrogram (frequency bins)
-        """
-        self.x_res = x_res
-        self.y_res = y_res
-        self.n_mels = self.y_res
-        self.slice_size = self.x_res * self.hop_length - 1
-
-    def load_audio(self, audio_file: str = None, raw_audio: np.ndarray = None):
-        """Load audio.
-
-        Args:
-            audio_file (`str`): must be a file on disk due to Librosa limitation or
-            raw_audio (`np.ndarray`): audio as numpy array
-        """
-        if audio_file is not None:
-            self.audio, _ = librosa.load(audio_file, mono=True, sr=self.sr)
-        else:
-            self.audio = raw_audio
-
-        # Pad with silence if necessary.
-        if len(self.audio) < self.x_res * self.hop_length:
-            self.audio = np.concatenate([self.audio, np.zeros((self.x_res * self.hop_length - len(self.audio),))])
-
-    def get_number_of_slices(self) -> int:
-        """Get number of slices in audio.
-
-        Returns:
-            `int`: number of spectograms audio can be sliced into
-        """
-        return len(self.audio) // self.slice_size
-
-    def get_audio_slice(self, slice: int = 0) -> np.ndarray:
-        """Get slice of audio.
-
-        Args:
-            slice (`int`): slice number of audio (out of get_number_of_slices())
-
-        Returns:
-            `np.ndarray`: audio as numpy array
-        """
-        return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)]
-
-    def get_sample_rate(self) -> int:
-        """Get sample rate:
-
-        Returns:
-            `int`: sample rate of audio
-        """
-        return self.sr
-
-    def audio_slice_to_image(self, slice: int) -> Image.Image:
-        """Convert slice of audio to spectrogram.
-
-        Args:
-            slice (`int`): slice number of audio to convert (out of get_number_of_slices())
-
-        Returns:
-            `PIL Image`: grayscale image of x_res x y_res
-        """
-        S = librosa.feature.melspectrogram(
-            y=self.get_audio_slice(slice), sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_mels=self.n_mels
-        )
-        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.fromarray(bytedata)
-        return image
-
-    def image_to_audio(self, image: Image.Image) -> np.ndarray:
-        """Converts spectrogram to audio.
-
-        Args:
-            image (`PIL Image`): x_res x y_res grayscale image
-
-        Returns:
-            audio (`np.ndarray`): raw audio
-        """
-        bytedata = np.frombuffer(image.tobytes(), dtype="uint8").reshape((image.height, image.width))
-        log_S = bytedata.astype("float") * self.top_db / 255 - self.top_db
-        S = librosa.db_to_power(log_S)
-        audio = librosa.feature.inverse.mel_to_audio(
-            S, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_iter=self.n_iter
-        )
-        return audio
-
-#-----------------------------------------------------------------------------#
-
-from math import acos, sin
-from typing import List, Tuple, Union
-
-import numpy as np
-import torch
-
-from PIL import Image
-
-from diffusers import AutoencoderKL, UNet2DConditionModel, DiffusionPipeline, DDIMScheduler, DDPMScheduler
-from diffusers.pipeline_utils import AudioPipelineOutput, BaseOutput, ImagePipelineOutput
-
-
-class AudioDiffusionPipeline(DiffusionPipeline):
-    """
-    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
-    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
-
-    Parameters:
-        vqae ([`AutoencoderKL`]): Variational AutoEncoder for Latent Audio Diffusion or None
-        unet ([`UNet2DConditionModel`]): UNET model
-        mel ([`Mel`]): transform audio <-> spectrogram
-        scheduler ([`DDIMScheduler` or `DDPMScheduler`]): de-noising scheduler
-    """
-
-    _optional_components = ["vqvae"]
-
-    def __init__(
-        self,
-        vqvae: AutoencoderKL,
-        unet: UNet2DConditionModel,
-        mel: Mel,
-        scheduler: Union[DDIMScheduler, DDPMScheduler],
-    ):
-        super().__init__()
-        self.register_modules(unet=unet, scheduler=scheduler, mel=mel, vqvae=vqvae)
-
-    def get_input_dims(self) -> Tuple:
-        """Returns dimension of input image
-
-        Returns:
-            `Tuple`: (height, width)
-        """
-        input_module = self.vqvae if self.vqvae is not None else self.unet
-        # For backwards compatibility
-        sample_size = (
-            (input_module.sample_size, input_module.sample_size)
-            if type(input_module.sample_size) == int
-            else input_module.sample_size
-        )
-        return sample_size
-
-    def get_default_steps(self) -> int:
-        """Returns default number of steps recommended for inference
-
-        Returns:
-            `int`: number of steps
-        """
-        return 50 if isinstance(self.scheduler, DDIMScheduler) else 1000
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        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,
-        step_generator: torch.Generator = None,
-        eta: float = 0,
-        noise: torch.Tensor = None,
-        return_dict=True,
-    ) -> Union[
-        Union[AudioPipelineOutput, ImagePipelineOutput], Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]]
-    ]:
-        """Generate random mel spectrogram from audio input and convert to audio.
-
-        Args:
-            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 (defaults to 50 for DDIM, 1000 for DDPM)
-            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
-            step_generator (`torch.Generator`): random number generator used to de-noise or None
-            eta (`float`): parameter between 0 and 1 used with DDIM scheduler
-            noise (`torch.Tensor`): noise tensor of shape (batch_size, 1, height, width) or None
-            return_dict (`bool`): if True return AudioPipelineOutput, ImagePipelineOutput else Tuple
-
-        Returns:
-            `List[PIL Image]`: mel spectrograms (`float`, `List[np.ndarray]`): sample rate and raw audios
-        """
-
-        steps = steps or self.get_default_steps()
-        self.scheduler.set_timesteps(steps)
-        step_generator = step_generator or generator
-        # For backwards compatibility
-        if type(self.unet.sample_size) == int:
-            self.unet.sample_size = (self.unet.sample_size, self.unet.sample_size)
-        input_dims = self.get_input_dims()
-        self.mel.set_resolution(x_res=input_dims[1], y_res=input_dims[0])
-        if noise is None:
-            noise = torch.randn(
-                (batch_size, self.unet.in_channels, self.unet.sample_size[0], self.unet.sample_size[1]),
-                generator=generator,
-                device=self.device,
-            )
-        images = noise
-        mask = None
-
-        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)
-            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 = torch.tensor(input_image[np.newaxis, :, :], dtype=torch.float).to(self.device)
-
-            if self.vqvae is not None:
-                input_images = self.vqvae.encode(torch.unsqueeze(input_images, 0)).latent_dist.sample(
-                    generator=generator
-                )[0]
-                input_images = 0.18215 * input_images
-
-            if start_step > 0:
-                images[0, 0] = self.scheduler.add_noise(input_images, noise, self.scheduler.timesteps[start_step - 1])
-
-            pixels_per_second = (
-                self.unet.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
-            )
-            mask_start = int(mask_start_secs * pixels_per_second)
-            mask_end = int(mask_end_secs * pixels_per_second)
-            mask = self.scheduler.add_noise(input_images, noise, torch.tensor(self.scheduler.timesteps[start_step:]))
-
-        for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:])):
-            model_output = self.unet(images, t)["sample"]
-
-            if isinstance(self.scheduler, DDIMScheduler):
-                images = self.scheduler.step(
-                    model_output=model_output, timestep=t, sample=images, eta=eta, generator=step_generator
-                )["prev_sample"]
-            else:
-                images = self.scheduler.step(
-                    model_output=model_output, timestep=t, sample=images, generator=step_generator
-                )["prev_sample"]
-
-            if mask is not None:
-                if mask_start > 0:
-                    images[:, :, :, :mask_start] = mask[:, step, :, :mask_start]
-                if mask_end > 0:
-                    images[:, :, :, -mask_end:] = mask[:, step, :, -mask_end:]
-
-        if self.vqvae is not None:
-            # 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)
-        )
-
-        audios = list(map(lambda _: self.mel.image_to_audio(_), images))
-        if not return_dict:
-            return images, (self.mel.get_sample_rate(), audios)
-
-        return BaseOutput(**AudioPipelineOutput(np.array(audios)[:, np.newaxis, :]), **ImagePipelineOutput(images))
-
-    @torch.no_grad()
-    def encode(self, images: List[Image.Image], steps: int = 50) -> np.ndarray:
-        """Reverse step process: recover noisy image from generated image.
-
-        Args:
-            images (`List[PIL Image]`): list of images to encode
-            steps (`int`): number of encoding steps to perform (defaults to 50)
-
-        Returns:
-            `np.ndarray`: noise tensor of shape (batch_size, 1, height, width)
-        """
-
-        # Only works with DDIM as this method is deterministic
-        assert isinstance(self.scheduler, DDIMScheduler)
-        self.scheduler.set_timesteps(steps)
-        sample = np.array(
-            [np.frombuffer(image.tobytes(), dtype="uint8").reshape((1, image.height, image.width)) for image in images]
-        )
-        sample = (sample / 255) * 2 - 1
-        sample = torch.Tensor(sample).to(self.device)
-
-        for t in self.progress_bar(torch.flip(self.scheduler.timesteps, (0,))):
-            prev_timestep = t - self.scheduler.num_train_timesteps // self.scheduler.num_inference_steps
-            alpha_prod_t = self.scheduler.alphas_cumprod[t]
-            alpha_prod_t_prev = (
-                self.scheduler.alphas_cumprod[prev_timestep]
-                if prev_timestep >= 0
-                else self.scheduler.final_alpha_cumprod
-            )
-            beta_prod_t = 1 - alpha_prod_t
-            model_output = self.unet(sample, t)["sample"]
-            pred_sample_direction = (1 - alpha_prod_t_prev) ** (0.5) * model_output
-            sample = (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
-            sample = sample * alpha_prod_t ** (0.5) + beta_prod_t ** (0.5) * model_output
-
-        return sample
-
-    @staticmethod
-    def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor:
-        """Spherical Linear intERPolation
-
-        Args:
-            x0 (`torch.Tensor`): first tensor to interpolate between
-            x1 (`torch.Tensor`): seconds tensor to interpolate between
-            alpha (`float`): interpolation between 0 and 1
-
-        Returns:
-            `torch.Tensor`: interpolated tensor
-        """
-
-        theta = acos(torch.dot(torch.flatten(x0), torch.flatten(x1)) / torch.norm(x0) / torch.norm(x1))
-        return sin((1 - alpha) * theta) * x0 / sin(theta) + sin(alpha * theta) * x1 / sin(theta)
-
-
-import sys
-import diffusers
-
-class audio_diffusion():
-    __name__ = 'audio_diffusion'
-    pass
-
-
-sys.modules['audio_diffusion'] = audio_diffusion
-setattr(audio_diffusion, Mel.__name__, Mel)
-diffusers.AudioDiffusionPipeline = AudioDiffusionPipeline
-setattr(diffusers, AudioDiffusionPipeline.__name__, AudioDiffusionPipeline)
-diffusers.pipeline_utils.LOADABLE_CLASSES['audio_diffusion'] = {}
-diffusers.pipeline_utils.LOADABLE_CLASSES['audio_diffusion']['Mel'] = ["save_pretrained", "from_pretrained"]
-'''

audiodiffusion/audio_encoder.py

@@ -0,0 +1,107 @@
+import numpy as np
+import torch
+from diffusers import ConfigMixin, Mel, ModelMixin
+from torch import nn
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size):
+        super(SeparableConv2d, self).__init__()
+        self.depthwise = nn.Conv2d(
+            in_channels,
+            in_channels,
+            kernel_size=kernel_size,
+            groups=in_channels,
+            bias=False,
+            padding=1,
+        )
+        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+
+    def forward(self, x):
+        out = self.depthwise(x)
+        out = self.pointwise(out)
+        return out
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, dropout_rate):
+        super(ConvBlock, self).__init__()
+        self.sep_conv = SeparableConv2d(in_channels, out_channels, (3, 3))
+        self.leaky_relu = nn.LeakyReLU(0.2)
+        self.batch_norm = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.01)
+        self.max_pool = nn.MaxPool2d((2, 2))
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        x = self.sep_conv(x)
+        x = self.leaky_relu(x)
+        x = self.batch_norm(x)
+        x = self.max_pool(x)
+        x = self.dropout(x)
+        return x
+
+
+class DenseBlock(nn.Module):
+    def __init__(self, in_features, out_features, dropout_rate):
+        super(DenseBlock, self).__init__()
+        self.flatten = nn.Flatten()
+        self.dense = nn.Linear(in_features, out_features)
+        self.leaky_relu = nn.LeakyReLU(0.2)
+        self.batch_norm = nn.BatchNorm1d(out_features, eps=0.001, momentum=0.01)
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(self, x):
+        x = self.flatten(x.permute(0, 2, 3, 1))
+        x = self.dense(x)
+        x = self.leaky_relu(x)
+        x = self.batch_norm(x)
+        x = self.dropout(x)
+        return x
+
+
+class AudioEncoder(ModelMixin, ConfigMixin):
+    def __init__(self):
+        super().__init__()
+        self.mel = Mel(
+            x_res=216,
+            y_res=96,
+            sample_rate=22050,
+            n_fft=2048,
+            hop_length=512,
+            top_db=80,
+        )
+        self.conv_blocks = nn.ModuleList([ConvBlock(1, 32, 0.2), ConvBlock(32, 64, 0.3), ConvBlock(64, 128, 0.4)])
+        self.dense_block = DenseBlock(41472, 1024, 0.5)
+        self.embedding = nn.Linear(1024, 100)
+
+    def forward(self, x):
+        for conv_block in self.conv_blocks:
+            x = conv_block(x)
+        x = self.dense_block(x)
+        x = self.embedding(x)
+        return x
+
+    @torch.no_grad()
+    def encode(self, audio_files):
+        self.eval()
+        y = []
+        for audio_file in audio_files:
+            self.mel.load_audio(audio_file)
+            x = [
+                np.expand_dims(
+                    np.frombuffer(self.mel.audio_slice_to_image(slice).tobytes(), dtype="uint8").reshape(
+                        (self.mel.y_res, self.mel.x_res)
+                    )
+                    / 255,
+                    axis=0,
+                )
+                for slice in range(self.mel.get_number_of_slices())
+            ]
+            y += [torch.mean(self(torch.Tensor(x)), dim=0)]
+        return torch.stack(y)
+
+
+# from diffusers import Mel
+# from audiodiffusion.audio_encoder import AudioEncoder
+# audio_encoder = AudioEncoder.from_pretrained("teticio/audio-encoder")
+# audio_encoder.encode(['/home/teticio/Music/liked/Agua Re - Holy Dance - Large Sound Mix.mp3'])

audiodiffusion/mel.py

@@ -0,0 +1,164 @@
+# This code has been migrated to diffusers but can be run locally with
+# pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-256", custom_pipeline="audio-diffusion/audiodiffusion/pipeline_audio_diffusion.py")
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import warnings
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+warnings.filterwarnings("ignore")
+
+import librosa  # noqa: E402
+import numpy as np  # noqa: E402
+from PIL import Image  # noqa: E402
+
+
+class Mel(ConfigMixin, SchedulerMixin):
+    """
+    Parameters:
+        x_res (`int`): x resolution of spectrogram (time)
+        y_res (`int`): y resolution of spectrogram (frequency bins)
+        sample_rate (`int`): sample rate of audio
+        n_fft (`int`): number of Fast Fourier Transforms
+        hop_length (`int`): hop length (a higher number is recommended for lower than 256 y_res)
+        top_db (`int`): loudest in decibels
+        n_iter (`int`): number of iterations for Griffin Linn mel inversion
+    """
+
+    config_name = "mel_config.json"
+
+    @register_to_config
+    def __init__(
+        self,
+        x_res: int = 256,
+        y_res: int = 256,
+        sample_rate: int = 22050,
+        n_fft: int = 2048,
+        hop_length: int = 512,
+        top_db: int = 80,
+        n_iter: int = 32,
+    ):
+        self.hop_length = hop_length
+        self.sr = sample_rate
+        self.n_fft = n_fft
+        self.top_db = top_db
+        self.n_iter = n_iter
+        self.set_resolution(x_res, y_res)
+        self.audio = None
+
+    def set_resolution(self, x_res: int, y_res: int):
+        """Set resolution.
+
+        Args:
+            x_res (`int`): x resolution of spectrogram (time)
+            y_res (`int`): y resolution of spectrogram (frequency bins)
+        """
+        self.x_res = x_res
+        self.y_res = y_res
+        self.n_mels = self.y_res
+        self.slice_size = self.x_res * self.hop_length - 1
+
+    def load_audio(self, audio_file: str = None, raw_audio: np.ndarray = None):
+        """Load audio.
+
+        Args:
+            audio_file (`str`): must be a file on disk due to Librosa limitation or
+            raw_audio (`np.ndarray`): audio as numpy array
+        """
+        if audio_file is not None:
+            self.audio, _ = librosa.load(audio_file, mono=True, sr=self.sr)
+        else:
+            self.audio = raw_audio
+
+        # Pad with silence if necessary.
+        if len(self.audio) < self.x_res * self.hop_length:
+            self.audio = np.concatenate(
+                [
+                    self.audio,
+                    np.zeros((self.x_res * self.hop_length - len(self.audio),)),
+                ]
+            )
+
+    def get_number_of_slices(self) -> int:
+        """Get number of slices in audio.
+
+        Returns:
+            `int`: number of spectograms audio can be sliced into
+        """
+        return len(self.audio) // self.slice_size
+
+    def get_audio_slice(self, slice: int = 0) -> np.ndarray:
+        """Get slice of audio.
+
+        Args:
+            slice (`int`): slice number of audio (out of get_number_of_slices())
+
+        Returns:
+            `np.ndarray`: audio as numpy array
+        """
+        return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)]
+
+    def get_sample_rate(self) -> int:
+        """Get sample rate:
+
+        Returns:
+            `int`: sample rate of audio
+        """
+        return self.sr
+
+    def audio_slice_to_image(self, slice: int) -> Image.Image:
+        """Convert slice of audio to spectrogram.
+
+        Args:
+            slice (`int`): slice number of audio to convert (out of get_number_of_slices())
+
+        Returns:
+            `PIL Image`: grayscale image of x_res x y_res
+        """
+        S = librosa.feature.melspectrogram(
+            y=self.get_audio_slice(slice),
+            sr=self.sr,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            n_mels=self.n_mels,
+        )
+        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.fromarray(bytedata)
+        return image
+
+    def image_to_audio(self, image: Image.Image) -> np.ndarray:
+        """Converts spectrogram to audio.
+
+        Args:
+            image (`PIL Image`): x_res x y_res grayscale image
+
+        Returns:
+            audio (`np.ndarray`): raw audio
+        """
+        bytedata = np.frombuffer(image.tobytes(), dtype="uint8").reshape((image.height, image.width))
+        log_S = bytedata.astype("float") * self.top_db / 255 - self.top_db
+        S = librosa.db_to_power(log_S)
+        audio = librosa.feature.inverse.mel_to_audio(
+            S,
+            sr=self.sr,
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            n_iter=self.n_iter,
+        )
+        return audio

audiodiffusion/pipeline_audio_diffusion.py

@@ -0,0 +1,267 @@
+# This code has been migrated to diffusers but can be run locally with
+# pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-256", custom_pipeline="audio-diffusion/audiodiffusion/pipeline_audio_diffusion.py")
+
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from math import acos, sin
+from typing import List, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, Mel, UNet2DConditionModel
+from diffusers.pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
+from PIL import Image
+
+from .mel import Mel
+
+
+class AudioDiffusionPipeline(DiffusionPipeline):
+    """
+    This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+    library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+    Parameters:
+        vqae ([`AutoencoderKL`]): Variational AutoEncoder for Latent Audio Diffusion or None
+        unet ([`UNet2DConditionModel`]): UNET model
+        mel ([`Mel`]): transform audio <-> spectrogram
+        scheduler ([`DDIMScheduler` or `DDPMScheduler`]): de-noising scheduler
+    """
+
+    _optional_components = ["vqvae"]
+
+    def __init__(
+        self,
+        vqvae: AutoencoderKL,
+        unet: UNet2DConditionModel,
+        mel: Mel,
+        scheduler: Union[DDIMScheduler, DDPMScheduler],
+    ):
+        super().__init__()
+        self.register_modules(unet=unet, scheduler=scheduler, mel=mel, vqvae=vqvae)
+
+    def get_input_dims(self) -> Tuple:
+        """Returns dimension of input image
+
+        Returns:
+            `Tuple`: (height, width)
+        """
+        input_module = self.vqvae if self.vqvae is not None else self.unet
+        # For backwards compatibility
+        sample_size = (
+            (input_module.sample_size, input_module.sample_size)
+            if type(input_module.sample_size) == int
+            else input_module.sample_size
+        )
+        return sample_size
+
+    def get_default_steps(self) -> int:
+        """Returns default number of steps recommended for inference
+
+        Returns:
+            `int`: number of steps
+        """
+        return 50 if isinstance(self.scheduler, DDIMScheduler) else 1000
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        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,
+        step_generator: torch.Generator = None,
+        eta: float = 0,
+        noise: torch.Tensor = None,
+        encoding: torch.Tensor = None,
+        return_dict=True,
+    ) -> Union[
+        Union[AudioPipelineOutput, ImagePipelineOutput],
+        Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
+    ]:
+        """Generate random mel spectrogram from audio input and convert to audio.
+
+        Args:
+            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 (defaults to 50 for DDIM, 1000 for DDPM)
+            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
+            step_generator (`torch.Generator`): random number generator used to de-noise or None
+            eta (`float`): parameter between 0 and 1 used with DDIM scheduler
+            noise (`torch.Tensor`): noise tensor of shape (batch_size, 1, height, width) or None
+            encoding (`torch.Tensor`): for UNet2DConditionModel shape (batch_size, seq_length, cross_attention_dim)
+            return_dict (`bool`): if True return AudioPipelineOutput, ImagePipelineOutput else Tuple
+
+        Returns:
+            `List[PIL Image]`: mel spectrograms (`float`, `List[np.ndarray]`): sample rate and raw audios
+        """
+
+        steps = steps or self.get_default_steps()
+        self.scheduler.set_timesteps(steps)
+        step_generator = step_generator or generator
+        # For backwards compatibility
+        if type(self.unet.sample_size) == int:
+            self.unet.sample_size = (self.unet.sample_size, self.unet.sample_size)
+        input_dims = self.get_input_dims()
+        self.mel.set_resolution(x_res=input_dims[1], y_res=input_dims[0])
+        if noise is None:
+            noise = torch.randn(
+                (
+                    batch_size,
+                    self.unet.in_channels,
+                    self.unet.sample_size[0],
+                    self.unet.sample_size[1],
+                ),
+                generator=generator,
+                device=self.device,
+            )
+        images = noise
+        mask = None
+
+        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)
+            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 = torch.tensor(input_image[np.newaxis, :, :], dtype=torch.float).to(self.device)
+
+            if self.vqvae is not None:
+                input_images = self.vqvae.encode(torch.unsqueeze(input_images, 0)).latent_dist.sample(
+                    generator=generator
+                )[0]
+                input_images = 0.18215 * input_images
+
+            if start_step > 0:
+                images[0, 0] = self.scheduler.add_noise(input_images, noise, self.scheduler.timesteps[start_step - 1])
+
+            pixels_per_second = (
+                self.unet.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
+            )
+            mask_start = int(mask_start_secs * pixels_per_second)
+            mask_end = int(mask_end_secs * pixels_per_second)
+            mask = self.scheduler.add_noise(input_images, noise, torch.tensor(self.scheduler.timesteps[start_step:]))
+
+        for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:])):
+            if isinstance(self.unet, UNet2DConditionModel):
+                model_output = self.unet(images, t, encoding)["sample"]
+            else:
+                model_output = self.unet(images, t)["sample"]
+
+            if isinstance(self.scheduler, DDIMScheduler):
+                images = self.scheduler.step(
+                    model_output=model_output,
+                    timestep=t,
+                    sample=images,
+                    eta=eta,
+                    generator=step_generator,
+                )["prev_sample"]
+            else:
+                images = self.scheduler.step(
+                    model_output=model_output,
+                    timestep=t,
+                    sample=images,
+                    generator=step_generator,
+                )["prev_sample"]
+
+            if mask is not None:
+                if mask_start > 0:
+                    images[:, :, :, :mask_start] = mask[:, step, :, :mask_start]
+                if mask_end > 0:
+                    images[:, :, :, -mask_end:] = mask[:, step, :, -mask_end:]
+
+        if self.vqvae is not None:
+            # 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)
+        )
+
+        audios = list(map(lambda _: self.mel.image_to_audio(_), images))
+        if not return_dict:
+            return images, (self.mel.get_sample_rate(), audios)
+
+        return BaseOutput(**AudioPipelineOutput(np.array(audios)[:, np.newaxis, :]), **ImagePipelineOutput(images))
+
+    @torch.no_grad()
+    def encode(self, images: List[Image.Image], steps: int = 50) -> np.ndarray:
+        """Reverse step process: recover noisy image from generated image.
+
+        Args:
+            images (`List[PIL Image]`): list of images to encode
+            steps (`int`): number of encoding steps to perform (defaults to 50)
+
+        Returns:
+            `np.ndarray`: noise tensor of shape (batch_size, 1, height, width)
+        """
+
+        # Only works with DDIM as this method is deterministic
+        assert isinstance(self.scheduler, DDIMScheduler)
+        self.scheduler.set_timesteps(steps)
+        sample = np.array(
+            [np.frombuffer(image.tobytes(), dtype="uint8").reshape((1, image.height, image.width)) for image in images]
+        )
+        sample = (sample / 255) * 2 - 1
+        sample = torch.Tensor(sample).to(self.device)
+
+        for t in self.progress_bar(torch.flip(self.scheduler.timesteps, (0,))):
+            prev_timestep = t - self.scheduler.num_train_timesteps // self.scheduler.num_inference_steps
+            alpha_prod_t = self.scheduler.alphas_cumprod[t]
+            alpha_prod_t_prev = (
+                self.scheduler.alphas_cumprod[prev_timestep]
+                if prev_timestep >= 0
+                else self.scheduler.final_alpha_cumprod
+            )
+            beta_prod_t = 1 - alpha_prod_t
+            model_output = self.unet(sample, t)["sample"]
+            pred_sample_direction = (1 - alpha_prod_t_prev) ** (0.5) * model_output
+            sample = (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
+            sample = sample * alpha_prod_t ** (0.5) + beta_prod_t ** (0.5) * model_output
+
+        return sample
+
+    @staticmethod
+    def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor:
+        """Spherical Linear intERPolation
+
+        Args:
+            x0 (`torch.Tensor`): first tensor to interpolate between
+            x1 (`torch.Tensor`): seconds tensor to interpolate between
+            alpha (`float`): interpolation between 0 and 1
+
+        Returns:
+            `torch.Tensor`: interpolated tensor
+        """
+
+        theta = acos(torch.dot(torch.flatten(x0), torch.flatten(x1)) / torch.norm(x0) / torch.norm(x1))
+        return sin((1 - alpha) * theta) * x0 / sin(theta) + sin(alpha * theta) * x1 / sin(theta)

audiodiffusion/utils.py

@@ -23,8 +23,7 @@ def renew_vae_resnet_paths(old_list, n_shave_prefix_segments=0):
         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)
+        new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
 
         mapping.append({"old": old_item, "new": new_item})
 
@@ -54,20 +53,21 @@ def renew_vae_attention_paths(old_list, n_shave_prefix_segments=0):
         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)
+        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):
+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
@@ -75,9 +75,7 @@ def assign_to_checkpoint(paths,
 
     Assigns the weights to the new checkpoint.
     """
-    assert isinstance(
-        paths, list
-    ), "Paths should be a list of dicts containing 'old' and 'new' keys."
+    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:
@@ -85,13 +83,11 @@ def assign_to_checkpoint(paths,
             old_tensor = old_checkpoint[path]
             channels = old_tensor.shape[0] // 3
 
-            target_shape = (-1,
-                            channels) if len(old_tensor.shape) == 3 else (-1)
+            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:])
+            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)
@@ -112,8 +108,7 @@ def assign_to_checkpoint(paths,
 
         if additional_replacements is not None:
             for replacement in additional_replacements:
-                new_path = new_path.replace(replacement["old"],
-                                            replacement["new"])
+                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:
@@ -146,7 +141,7 @@ def create_vae_diffusers_config(original_config):
     up_block_types = ["UpDecoderBlock2D"] * len(block_out_channels)
 
     config = dict(
-        sample_size=vae_params.resolution,
+        sample_size=tuple(vae_params.resolution),
         in_channels=vae_params.in_channels,
         out_channels=vae_params.out_ch,
         down_block_types=tuple(down_block_types),
@@ -164,178 +159,144 @@ def convert_ldm_vae_checkpoint(checkpoint, config):
 
     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["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"]
+    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
-    })
+    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)
+        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
-    })
+    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)
+        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
-        ]
+        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")
+            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)
+        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
-        ]
+        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
-    ]
+        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)
+    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
+            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"]
+            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)
+        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
-        ]
+        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
-    ]
+        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)
+    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):
+
+def convert_ldm_to_hf_vae(ldm_checkpoint, ldm_config, hf_checkpoint, sample_size):
     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)
+    converted_vae_checkpoint = convert_ldm_vae_checkpoint(checkpoint, vae_config)
 
     vae = AutoencoderKL(**vae_config)
     vae.load_state_dict(converted_vae_checkpoint)

config/ldm_autoencoder_kl.yaml

@@ -18,7 +18,7 @@ model:
     ddconfig:
       double_z: True
       z_channels: 1  # must = embed_dim due to HF limitation
-      resolution: 256
+      resolution: 256  # overriden by input image size
       in_channels: 1
       out_ch: 1
       ch: 128

notebooks/audio_encoder.ipynb

@@ -0,0 +1,69 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "592fff30",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from diffusers import Mel\n",
+    "from audiodiffusion.audio_encoder import AudioEncoder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d99ef523",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio_encoder = AudioEncoder.from_pretrained(\"teticio/audio-encoder\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4eb3bbd7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio_encoder.encode(['/home/teticio/Music/liked/Agua Re - Holy Dance - Large Sound Mix.mp3'])"
+   ]
+  }
+ ],
+ "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
+}

notebooks/conditional_generation.ipynb

@@ -0,0 +1,221 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "2a44739f",
+   "metadata": {},
+   "source": [
+    "<a href=\"https://colab.research.google.com/github/teticio/audio-diffusion/blob/master/notebooks/condtional_generation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1935544",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "try:\n",
+    "    # are we running on Google Colab?\n",
+    "    import google.colab\n",
+    "    !git clone -q https://github.com/teticio/audio-diffusion.git\n",
+    "    %cd audio-diffusion\n",
+    "    %pip install -q -r requirements.txt\n",
+    "except:\n",
+    "    pass"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b0e656c9",
+   "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": "d448b299",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import urllib\n",
+    "import requests\n",
+    "from IPython.display import Audio\n",
+    "from audiodiffusion import AudioDiffusion\n",
+    "from audiodiffusion.audio_encoder import AudioEncoder"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f1548971",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "generator = torch.Generator(device=device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "056f179c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio_diffusion = AudioDiffusion(model_id=\"teticio/conditional-latent-audio-diffusion-512\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4a08500",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "audio_encoder = AudioEncoder.from_pretrained(\"teticio/audio-encoder\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "387550ac",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Uncomment for faster (but slightly lower quality) generation\n",
+    "#from diffusers import DDIMScheduler\n",
+    "#audio_diffusion.pipe.scheduler = DDIMScheduler()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9936a72f",
+   "metadata": {},
+   "source": [
+    "## Download and encode preview track from Spotify"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "57a9b134",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Get temporary API credentials\n",
+    "credentials = requests.get(\n",
+    "    \"https://open.spotify.com/get_access_token?reason=transport&productType=embed\"\n",
+    ").json()\n",
+    "headers = {\n",
+    "    \"Accept\": \"application/json\",\n",
+    "    \"Content-Type\": \"application/json\",\n",
+    "    \"Authorization\": \"Bearer \" + credentials[\"accessToken\"]\n",
+    "}\n",
+    "\n",
+    "# Search for tracks\n",
+    "search_string = input(\"Search: \")\n",
+    "response = requests.get(\n",
+    "    f\"https://api.spotify.com/v1/search?q={urllib.parse.quote(search_string)}&type=track\",\n",
+    "    headers=headers).json()\n",
+    "\n",
+    "# List results\n",
+    "for _, track in enumerate(response[\"tracks\"][\"items\"]):\n",
+    "    print(f\"{_ + 1}. {track['artists'][0]['name']} - {track['name']}\")\n",
+    "selection = input(\"Select a track: \")\n",
+    "\n",
+    "# Download and encode selection\n",
+    "r = requests.get(response[\"tracks\"][\"items\"][int(selection) -\n",
+    "                                             1][\"preview_url\"],\n",
+    "                 stream=True)\n",
+    "with open(\"temp.mp3\", \"wb\") as f:\n",
+    "    for chunk in r:\n",
+    "        f.write(chunk)\n",
+    "encoding = torch.unsqueeze(audio_encoder.encode([\"temp.mp3\"]),\n",
+    "                           axis=1).to(device)\n",
+    "os.remove(\"temp.mp3\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8af863f5",
+   "metadata": {},
+   "source": [
+    "## Conditional Generation\n",
+    "Bear in mind that the generative model can only generate music similar to that on which it was trained. The audio encoding will influence the generation within those limitations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8f119ddd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for _ in range(10):\n",
+    "    seed = generator.seed()\n",
+    "    print(f'Seed = {seed}')\n",
+    "    generator.manual_seed(seed)\n",
+    "    image, (sample_rate,\n",
+    "            audio) = audio_diffusion.generate_spectrogram_and_audio(\n",
+    "                generator=generator, encoding=encoding)\n",
+    "    display(image)\n",
+    "    display(Audio(audio, rate=sample_rate))\n",
+    "    loop = AudioDiffusion.loop_it(audio, sample_rate)\n",
+    "    if loop is not None:\n",
+    "        display(Audio(loop, rate=sample_rate))\n",
+    "    else:\n",
+    "        print(\"Unable to determine loop points\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d0bd18c0",
+   "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
+}

pyproject.toml

@@ -0,0 +1,3 @@
+[tool.black]
+line-length = 119
+target-version = ['py36']

scripts/audio_to_images.py

@@ -1,29 +1,33 @@
-import os
-import re
+import argparse
 import io
 import logging
-import argparse
+import os
+import re
 
 import numpy as np
 import pandas as pd
-from tqdm.auto import tqdm
-from diffusers.pipelines.audio_diffusion import Mel
 from datasets import Dataset, DatasetDict, Features, Image, Value
+from diffusers.pipelines.audio_diffusion import Mel
+from tqdm.auto import tqdm
 
 logging.basicConfig(level=logging.WARN)
-logger = logging.getLogger('audio_to_images')
+logger = logging.getLogger("audio_to_images")
 
 
 def main(args):
-    mel = Mel(x_res=args.resolution[0],
-              y_res=args.resolution[1],
-              hop_length=args.hop_length,
-              sample_rate=args.sample_rate,
-              n_fft=args.n_fft)
+    mel = Mel(
+        x_res=args.resolution[0],
+        y_res=args.resolution[1],
+        hop_length=args.hop_length,
+        sample_rate=args.sample_rate,
+        n_fft=args.n_fft,
+    )
     os.makedirs(args.output_dir, exist_ok=True)
     audio_files = [
-        os.path.join(root, file) for root, _, files in os.walk(args.input_dir)
-        for file in files if re.search("\.(mp3|wav|m4a)$", file, re.IGNORECASE)
+        os.path.join(root, file)
+        for root, _, files in os.walk(args.input_dir)
+        for file in files
+        if re.search("\.(mp3|wav|m4a)$", file, re.IGNORECASE)
     ]
     examples = []
     try:
@@ -36,36 +40,38 @@ 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[0] and image.height
-                        == args.resolution[1]), "Wrong 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',
-                                audio_file, slice)
+                    logger.warn("File %s slice %d is completely silent", audio_file, slice)
                     continue
                 with io.BytesIO() as output:
                     image.save(output, format="PNG")
                     bytes = output.getvalue()
-                examples.extend([{
-                    "image": {
-                        "bytes": bytes
-                    },
-                    "audio_file": audio_file,
-                    "slice": slice,
-                }])
+                examples.extend(
+                    [
+                        {
+                            "image": {"bytes": bytes},
+                            "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.')
+            logger.warn("No valid audio files were found.")
             return
         ds = Dataset.from_pandas(
             pd.DataFrame(examples),
-            features=Features({
-                "image": Image(),
-                "audio_file": Value(dtype="string"),
-                "slice": Value(dtype="int16"),
-            }),
+            features=Features(
+                {
+                    "image": Image(),
+                    "audio_file": Value(dtype="string"),
+                    "slice": Value(dtype="int16"),
+                }
+            ),
         )
         dsd = DatasetDict({"train": ds})
         dsd.save_to_disk(os.path.join(args.output_dir))
@@ -74,15 +80,15 @@ def main(args):
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description=
-        "Create dataset of Mel spectrograms from directory of audio files.")
+    parser = argparse.ArgumentParser(description="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=str,
-                        default="256",
-                        help="Either square resolution or width,height.")
+    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)
     parser.add_argument("--sample_rate", type=int, default=22050)
@@ -90,8 +96,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     if args.input_dir is None:
-        raise ValueError(
-            "You must specify an input directory for the audio files.")
+        raise ValueError("You must specify an input directory for the audio files.")
 
     # Handle the resolutions.
     try:
@@ -102,9 +107,7 @@ if __name__ == "__main__":
             if len(args.resolution) != 2:
                 raise ValueError
         except ValueError:
-            raise ValueError(
-                "Resolution must be a tuple of two integers or a single integer."
-            )
+            raise ValueError("Resolution must be a tuple of two integers or a single integer.")
     assert isinstance(args.resolution, tuple)
 
     main(args)

scripts/encode_audio.py

@@ -0,0 +1,38 @@
+import argparse
+import os
+import pickle
+
+from datasets import load_dataset, load_from_disk
+from tqdm.auto import tqdm
+
+from audiodiffusion.audio_encoder import AudioEncoder
+
+
+def main(args):
+    audio_encoder = AudioEncoder.from_pretrained("teticio/audio-encoder")
+
+    if args.dataset_name is not None:
+        if os.path.exists(args.dataset_name):
+            dataset = load_from_disk(args.dataset_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",
+            )
+
+    encodings = {}
+    for audio_file in tqdm(dataset.to_pandas()["audio_file"].unique()):
+        encodings[audio_file] = audio_encoder.encode([audio_file])
+    pickle.dump(encodings, open(args.output_file, "wb"))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Create pickled audio encodings for dataset of audio files.")
+    parser.add_argument("--dataset_name", type=str, default=None)
+    parser.add_argument("--output_file", type=str, default="data/encodings.p")
+    parser.add_argument("--use_auth_token", type=bool, default=False)
+    args = parser.parse_args()
+    main(args)

scripts/{train_unconditional.py → train_unet.py}

@@ -2,34 +2,29 @@
 
 import argparse
 import os
+import pickle
+import random
 from pathlib import Path
 from typing import Optional
 
+import numpy as np
+import torch
+import torch.nn.functional as F
 from accelerate import Accelerator
 from accelerate.logging import get_logger
-from datasets import load_from_disk, load_dataset
-from diffusers import (
-    AudioDiffusionPipeline,
-    DDPMScheduler,
-    UNet2DModel,
-    DDIMScheduler,
-    AutoencoderKL,
-)
-from diffusers.pipelines.audio_diffusion import Mel
+from datasets import load_dataset, load_from_disk
+from diffusers import (AutoencoderKL, DDIMScheduler, DDPMScheduler,
+                       UNet2DConditionModel, UNet2DModel)
 from diffusers.optimization import get_scheduler
+from diffusers.pipelines.audio_diffusion import Mel
 from diffusers.training_utils import EMAModel
 from huggingface_hub import HfFolder, Repository, whoami
 from librosa.util import normalize
-import numpy as np
-import torch
-import torch.nn.functional as F
-from torchvision.transforms import (
-    Compose,
-    Normalize,
-    ToTensor,
-)
+from torchvision.transforms import Compose, Normalize, ToTensor
 from tqdm.auto import tqdm
 
+from audiodiffusion.pipeline_audio_diffusion import AudioDiffusionPipeline
+
 logger = get_logger(__name__)
 
 
@@ -90,12 +85,18 @@ def main(args):
             ]
         else:
             images = [augmentations(image) for image in examples["image"]]
+        if args.encodings is not None:
+            encoding = [encodings[file] for file in examples["audio_file"]]
+            return {"input": images, "encoding": encoding}
         return {"input": images}
 
     dataset.set_transform(transforms)
     train_dataloader = torch.utils.data.DataLoader(
         dataset, batch_size=args.train_batch_size, shuffle=True)
 
+    if args.encodings is not None:
+        encodings = pickle.load(open(args.encodings, "rb"))
+
     vqvae = None
     if args.vae is not None:
         try:
@@ -104,9 +105,9 @@ def main(args):
             vqvae = AudioDiffusionPipeline.from_pretrained(args.vae).vqvae
         # Determine latent resolution
         with torch.no_grad():
-            latent_resolution = (vqvae.encode(
+            latent_resolution = vqvae.encode(
                 torch.zeros((1, 1) +
-                            resolution)).latent_dist.sample().shape[2:])
+                            resolution)).latent_dist.sample().shape[2:]
 
     if args.from_pretrained is not None:
         pipeline = AudioDiffusionPipeline.from_pretrained(args.from_pretrained)
@@ -114,32 +115,58 @@ def main(args):
         model = pipeline.unet
         if hasattr(pipeline, "vqvae"):
             vqvae = pipeline.vqvae
+
     else:
-        model = UNet2DModel(
-            sample_size=resolution if vqvae is None else latent_resolution,
-            in_channels=1
-            if vqvae is None else vqvae.config["latent_channels"],
-            out_channels=1
-            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=(
-                "DownBlock2D",
-                "DownBlock2D",
-                "DownBlock2D",
-                "DownBlock2D",
-                "AttnDownBlock2D",
-                "DownBlock2D",
-            ),
-            up_block_types=(
-                "UpBlock2D",
-                "AttnUpBlock2D",
-                "UpBlock2D",
-                "UpBlock2D",
-                "UpBlock2D",
-                "UpBlock2D",
-            ),
-        )
+        if args.encodings is None:
+            model = UNet2DModel(
+                sample_size=resolution if vqvae is None else latent_resolution,
+                in_channels=1
+                if vqvae is None else vqvae.config["latent_channels"],
+                out_channels=1
+                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=(
+                    "DownBlock2D",
+                    "DownBlock2D",
+                    "DownBlock2D",
+                    "DownBlock2D",
+                    "AttnDownBlock2D",
+                    "DownBlock2D",
+                ),
+                up_block_types=(
+                    "UpBlock2D",
+                    "AttnUpBlock2D",
+                    "UpBlock2D",
+                    "UpBlock2D",
+                    "UpBlock2D",
+                    "UpBlock2D",
+                ),
+            )
+
+        else:
+            model = UNet2DConditionModel(
+                sample_size=resolution if vqvae is None else latent_resolution,
+                in_channels=1
+                if vqvae is None else vqvae.config["latent_channels"],
+                out_channels=1
+                if vqvae is None else vqvae.config["latent_channels"],
+                layers_per_block=2,
+                block_out_channels=(128, 256, 512, 512),
+                down_block_types=(
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "CrossAttnDownBlock2D",
+                    "DownBlock2D",
+                ),
+                up_block_types=(
+                    "UpBlock2D",
+                    "CrossAttnUpBlock2D",
+                    "CrossAttnUpBlock2D",
+                    "CrossAttnUpBlock2D",
+                ),
+                cross_attention_dim=list(encodings.values())[0].shape[-1],
+            )
 
     if args.scheduler == "ddpm":
         noise_scheduler = DDPMScheduler(
@@ -240,7 +267,11 @@ def main(args):
 
             with accelerator.accumulate(model):
                 # Predict the noise residual
-                noise_pred = model(noisy_images, timesteps)["sample"]
+                if args.encodings is not None:
+                    noise_pred = model(noisy_images, timesteps,
+                                       batch["encoding"])["sample"]
+                else:
+                    noise_pred = model(noisy_images, timesteps)["sample"]
                 loss = F.mse_loss(noise_pred, noise)
                 accelerator.backward(loss)
 
@@ -270,9 +301,9 @@ def main(args):
 
         # Generate sample images for visual inspection
         if accelerator.is_main_process:
-            if (epoch + 1) % args.save_model_epochs == 0 or (
-                    epoch + 1
-            ) % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
+            if ((epoch + 1) % args.save_model_epochs == 0
+                    or (epoch + 1) % args.save_images_epochs == 0
+                    or epoch == args.num_epochs - 1):
                 pipeline = AudioDiffusionPipeline(
                     vqvae=vqvae,
                     unet=accelerator.unwrap_model(
@@ -288,18 +319,32 @@ def main(args):
 
                 # save the model
                 if args.push_to_hub:
-                    repo.push_to_hub(commit_message=f"Epoch {epoch}",
-                                     blocking=False,
-                                     auto_lfs_prune=True)
+                    repo.push_to_hub(
+                        commit_message=f"Epoch {epoch}",
+                        blocking=False,
+                        auto_lfs_prune=True,
+                    )
 
             if (epoch + 1) % args.save_images_epochs == 0:
                 generator = torch.Generator(
                     device=clean_images.device).manual_seed(42)
+
+                if args.encodings is not None:
+                    random.seed(42)
+                    encoding = torch.stack(
+                        random.sample(list(encodings.values()),
+                                      args.eval_batch_size)).to(
+                                          clean_images.device)
+                else:
+                    encoding = None
+
                 # run pipeline in inference (sample random noise and denoise)
-                images, (sample_rate,
-                         audios) = pipeline(generator=generator,
-                                            batch_size=args.eval_batch_size,
-                                            return_dict=False)
+                images, (sample_rate, audios) = pipeline(
+                    generator=generator,
+                    batch_size=args.eval_batch_size,
+                    return_dict=False,
+                    encoding=encoding,
+                )
 
                 # denormalize the images and save to tensorboard
                 images = np.array([
@@ -385,6 +430,12 @@ if __name__ == "__main__":
         default=None,
         help="pretrained VAE model for latent diffusion",
     )
+    parser.add_argument(
+        "--encodings",
+        type=str,
+        default=None,
+        help="picked dictionary mapping audio_file to encoding",
+    )
 
     args = parser.parse_args()
     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))

scripts/train_vae.py

@@ -1,50 +1,48 @@
 # based on https://github.com/CompVis/stable-diffusion/blob/main/main.py
 
-import os
 import argparse
+import os
 
-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
+import torch
+import torchvision
+from datasets import load_dataset, load_from_disk
+from diffusers.pipelines.audio_diffusion import Mel
 from ldm.util import instantiate_from_config
+from librosa.util import normalize
+from omegaconf import OmegaConf
+from PIL import Image
+from pytorch_lightning.callbacks import Callback, ModelCheckpoint
 from pytorch_lightning.trainer import Trainer
+from pytorch_lightning.utilities.distributed import rank_zero_only
 from torch.utils.data import DataLoader, Dataset
-from datasets import load_from_disk, load_dataset
-from diffusers.pipelines.audio_diffusion import Mel
+
 from audiodiffusion.utils import convert_ldm_to_hf_vae
-from pytorch_lightning.callbacks import Callback, ModelCheckpoint
-from pytorch_lightning.utilities.distributed import rank_zero_only
 
 
 class AudioDiffusion(Dataset):
-
     def __init__(self, model_id, channels=3):
         super().__init__()
         self.channels = channels
         if os.path.exists(model_id):
-            self.hf_dataset = load_from_disk(model_id)['train']
+            self.hf_dataset = load_from_disk(model_id)["train"]
         else:
-            self.hf_dataset = load_dataset(model_id)['train']
+            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']
+        image = self.hf_dataset[idx]["image"]
         if self.channels == 3:
-            image = image.convert('RGB')
-        image = np.frombuffer(image.tobytes(), dtype="uint8").reshape(
-            (image.height, image.width, self.channels))
-        image = ((image / 255) * 2 - 1)
-        return {'image': image}
+            image = image.convert("RGB")
+        image = np.frombuffer(image.tobytes(), dtype="uint8").reshape((image.height, image.width, self.channels))
+        image = (image / 255) * 2 - 1
+        return {"image": image}
 
 
 class AudioDiffusionDataModule(pl.LightningDataModule):
-
     def __init__(self, model_id, batch_size, channels):
         super().__init__()
         self.batch_size = batch_size
@@ -52,18 +50,11 @@ class AudioDiffusionDataModule(pl.LightningDataModule):
         self.num_workers = 1
 
     def train_dataloader(self):
-        return DataLoader(self.dataset,
-                          batch_size=self.batch_size,
-                          num_workers=self.num_workers)
+        return DataLoader(self.dataset, batch_size=self.batch_size, num_workers=self.num_workers)
 
 
 class ImageLogger(Callback):
-
-    def __init__(self,
-                 every=1000,
-                 hop_length=512,
-                 sample_rate=22050,
-                 n_fft=2048):
+    def __init__(self, every=1000, hop_length=512, sample_rate=22050, n_fft=2048):
         super().__init__()
         self.every = every
         self.hop_length = hop_length
@@ -74,83 +65,75 @@ class ImageLogger(Callback):
     def log_images_and_audios(self, pl_module, batch):
         pl_module.eval()
         with torch.no_grad():
-            images = pl_module.log_images(batch, split='train')
+            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,
-                  sample_rate=self.sample_rate,
-                  n_fft=self.n_fft)
+        mel = Mel(
+            x_res=image_shape[2],
+            y_res=image_shape[3],
+            hop_length=self.hop_length,
+            sample_rate=self.sample_rate,
+            n_fft=self.n_fft,
+        )
 
         for k in images:
             images[k] = images[k].detach().cpu()
-            images[k] = torch.clamp(images[k], -1., 1.)
+            images[k] = torch.clamp(images[k], -1.0, 1.0)
             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)
+            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)
+            images[k] = (images[k].numpy() * 255).round().astype("uint8").transpose(0, 2, 3, 1)
             for _, image in enumerate(images[k]):
                 audio = mel.image_to_audio(
-                    Image.fromarray(image, mode='RGB').convert('L')
-                    if channels == 3 else Image.fromarray(image[:, :, 0]))
+                    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=mel.get_sample_rate())
+                    sample_rate=mel.get_sample_rate(),
+                )
 
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch,
-                           batch_idx):
+    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
+        self.sample_size = None
+
+    def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
+        if self.sample_size is None:
+            self.sample_size = list(batch["image"].shape[1:3])
 
     def on_train_epoch_end(self, trainer, pl_module):
-        ldm_checkpoint = self._get_metric_interpolated_filepath_name(
-            {'epoch': trainer.current_epoch}, trainer)
+        ldm_checkpoint = self._get_metric_interpolated_filepath_name({"epoch": trainer.current_epoch}, trainer)
         super().on_train_epoch_end(trainer, pl_module)
-        convert_ldm_to_hf_vae(ldm_checkpoint, self.ldm_config,
-                              self.hf_checkpoint)
+        self.ldm_config.model.params.ddconfig.resolution = self.sample_size
+        convert_ldm_to_hf_vae(ldm_checkpoint, self.ldm_config, self.hf_checkpoint, self.sample_size)
 
 
 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)
+    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)
     parser.add_argument("--hop_length", type=int, default=512)
     parser.add_argument("--sample_rate", type=int, default=22050)
     parser.add_argument("--n_fft", type=int, default=2048)
@@ -164,7 +147,8 @@ if __name__ == "__main__":
     data = AudioDiffusionDataModule(
         model_id=args.dataset_name,
         batch_size=args.batch_size,
-        channels=config.model.params.ddconfig.in_channels)
+        channels=config.model.params.ddconfig.in_channels,
+    )
     lightning_config = config.pop("lightning", OmegaConf.create())
     trainer_config = lightning_config.get("trainer", OmegaConf.create())
     trainer_config.accumulate_grad_batches = args.gradient_accumulation_steps
@@ -174,15 +158,20 @@ if __name__ == "__main__":
         max_epochs=args.max_epochs,
         resume_from_checkpoint=args.resume_from_checkpoint,
         callbacks=[
-            ImageLogger(every=args.save_images_batches,
-                        hop_length=args.hop_length,
-                        sample_rate=args.sample_rate,
-                        n_fft=args.n_fft),
-            HFModelCheckpoint(ldm_config=config,
-                              hf_checkpoint=args.hf_checkpoint_dir,
-                              dirpath=args.ldm_checkpoint_dir,
-                              filename='{epoch:06}',
-                              verbose=True,
-                              save_last=True)
-        ])
+            ImageLogger(
+                every=args.save_images_batches,
+                hop_length=args.hop_length,
+                sample_rate=args.sample_rate,
+                n_fft=args.n_fft,
+            ),
+            HFModelCheckpoint(
+                ldm_config=config,
+                hf_checkpoint=args.hf_checkpoint_dir,
+                dirpath=args.ldm_checkpoint_dir,
+                filename="{epoch:06}",
+                verbose=True,
+                save_last=True,
+            ),
+        ],
+    )
     trainer.fit(model, data)