get ready for migration to diffusers

README.md

@@ -23,6 +23,8 @@ Go to https://soundcloud.com/teticio2/sets/audio-diffusion-loops for more exampl
 ---
 #### Updates
 
+**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).
+
 **7/11/2022**. Added pre-trained latent audio diffusion models [teticio/latent-audio-diffusion-256](https://huggingface.co/teticio/latent-audio-diffusion-256) and [teticio/latent-audio-diffusion-ddim-256](https://huggingface.co/teticio/latent-audio-diffusion-ddim-256). You can use the pre-trained VAE to train your own latent diffusion models on a different set of audio files.
 
 **22/10/2022**. Added DDIM encoder and ability to interpolate between audios in latent "noise" space. Mel spectrograms no longer have to be square (thanks to Tristan for this one), so you can set the vertical (frequency) and horizontal (time) resolutions independently.

audiodiffusion/__init__.py

@@ -1,62 +1,34 @@
-from math import acos, sin
-from typing import Iterable, Tuple, Union, List
+from typing import Iterable, Tuple, Union
 
 import torch
 import numpy as np
 from PIL import Image
 from tqdm.auto import tqdm
 from librosa.beat import beat_track
-from diffusers import (DiffusionPipeline, UNet2DConditionModel, DDIMScheduler,
-                       DDPMScheduler, AutoencoderKL)
-from diffusers.pipeline_utils import (AudioPipelineOutput, BaseOutput,
-                                      ImagePipelineOutput)
-
-from .mel import Mel
+#from diffusers import DiffusionPipeline
 
-VERSION = "1.2.7"
+VERSION = "1.3.0"
 
 
 class AudioDiffusion:
 
     def __init__(self,
                  model_id: str = "teticio/audio-diffusion-256",
-                 sample_rate: int = 22050,
-                 n_fft: int = 2048,
-                 hop_length: int = 512,
-                 top_db: int = 80,
                  cuda: bool = torch.cuda.is_available(),
                  progress_bar: Iterable = tqdm):
         """Class for generating audio using De-noising Diffusion Probabilistic Models.
 
         Args:
             model_id (String): name of model (local directory or Hugging Face Hub)
-            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
             cuda (bool): use CUDA?
             progress_bar (iterable): iterable callback for progress updates or None
         """
         self.model_id = model_id
-        pipeline = {
-            'LatentAudioDiffusionPipeline': LatentAudioDiffusionPipeline,
-            'AudioDiffusionPipeline': AudioDiffusionPipeline
-        }.get(
-            DiffusionPipeline.get_config_dict(self.model_id)['_class_name'],
-            AudioDiffusionPipeline)
-        self.pipe = pipeline.from_pretrained(self.model_id)
+        self.pipe = AudioDiffusionPipeline.from_pretrained(self.model_id)
         if cuda:
             self.pipe.to("cuda")
         self.progress_bar = progress_bar or (lambda _: _)
 
-        sample_size = self.pipe.get_input_dims()
-        self.mel = Mel(x_res=sample_size[1],
-                       y_res=sample_size[0],
-                       sample_rate=sample_rate,
-                       n_fft=n_fft,
-                       hop_length=hop_length,
-                       top_db=top_db)
-
     def generate_spectrogram_and_audio(
         self,
         steps: int = None,
@@ -79,8 +51,7 @@ class AudioDiffusion:
             (float, np.ndarray): sample rate and raw audio
         """
         images, (sample_rate,
-                 audios) = self.pipe(mel=self.mel,
-                                     batch_size=1,
+                 audios) = self.pipe(batch_size=1,
                                      steps=steps,
                                      generator=generator,
                                      step_generator=step_generator,
@@ -124,8 +95,7 @@ class AudioDiffusion:
         """
 
         images, (sample_rate,
-                 audios) = self.pipe(mel=self.mel,
-                                     batch_size=1,
+                 audios) = self.pipe(batch_size=1,
                                      audio_file=audio_file,
                                      raw_audio=raw_audio,
                                      slice=slice,
@@ -161,18 +131,274 @@ class AudioDiffusion:
         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
+
+
+warnings.filterwarnings("ignore")
+
+import numpy as np  # noqa: E402
+
+import librosa  # noqa: E402
+from PIL import Image  # noqa: E402
+
+
+class Mel(ConfigMixin):
+    """
+    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
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Dict[str, Any] = None,
+        subfolder: Optional[str] = None,
+        return_unused_kwargs=False,
+        **kwargs,
+    ):
+        r"""
+        Instantiate a Mel class from a pre-defined JSON configuration file inside a directory or Hub repo.
+
+        Parameters:
+            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
+                Can be either:
+
+                    - A string, the *model id* of a model repo on huggingface.co. Valid model ids should have an
+                      organization name, like `google/ddpm-celebahq-256`.
+                    - A path to a *directory* containing the mel configurations saved using [`~Mel.save_pretrained`],
+                      e.g., `./my_model_directory/`.
+            subfolder (`str`, *optional*):
+                In case the relevant files are located inside a subfolder of the model repo (either remote in
+                huggingface.co or downloaded locally), you can specify the folder name here.
+            return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+                Whether kwargs that are not consumed by the Python class should be returned or not.
+            cache_dir (`Union[str, os.PathLike]`, *optional*):
+                Path to a directory in which a downloaded pretrained model configuration should be cached if the
+                standard cache should not be used.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            resume_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
+                file exists.
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            output_loading_info(`bool`, *optional*, defaults to `False`):
+                Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
+            local_files_only(`bool`, *optional*, defaults to `False`):
+                Whether or not to only look at local files (i.e., do not try to download the model).
+            use_auth_token (`str` or *bool*, *optional*):
+                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+                when running `transformers-cli login` (stored in `~/.huggingface`).
+            revision (`str`, *optional*, defaults to `"main"`):
+                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+                identifier allowed by git.
+
+        <Tip>
+
+         It is required to be logged in (`huggingface-cli login`) when you want to use private or [gated
+         models](https://huggingface.co/docs/hub/models-gated#gated-models).
+
+        </Tip>
+
+        <Tip>
+
+        Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to
+        use this method in a firewalled environment.
+
+        </Tip>
+
+        """
+        config, kwargs = cls.load_config(
+            pretrained_model_name_or_path=pretrained_model_name_or_path,
+            subfolder=subfolder,
+            return_unused_kwargs=True,
+            **kwargs,
+        )
+        return cls.from_config(config, return_unused_kwargs=return_unused_kwargs, **kwargs)
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], push_to_hub: bool = False, **kwargs):
+        """
+        Save a mel configuration object to the directory `save_directory`, so that it can be re-loaded using the
+        [`~Mel.from_pretrained`] class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`):
+                Directory where the configuration JSON file will be saved (will be created if it does not exist).
+        """
+        self.save_config(save_directory=save_directory, push_to_hub=push_to_hub, **kwargs)
+
+#-----------------------------------------------------------------------------#
+
+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):
-    def __init__(self, unet: UNet2DConditionModel, scheduler: Union[DDIMScheduler, DDPMScheduler]):
+    """
+    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)
+        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)
+            `Tuple`: (height, width)
         """
-        input_module = self.vqvae if hasattr(self, "vqvae") else self.unet
+        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)
@@ -185,14 +411,13 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         """Returns default number of steps recommended for inference
 
         Returns:
-            int: number of steps
+            `int`: number of steps
         """
         return 50 if isinstance(self.scheduler, DDIMScheduler) else 1000
 
     @torch.no_grad()
     def __call__(
         self,
-        mel: Mel,
         batch_size: int = 1,
         audio_file: str = None,
         raw_audio: np.ndarray = None,
@@ -212,23 +437,22 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         """Generate random mel spectrogram from audio input and convert to audio.
 
         Args:
-            mel (Mel): instance of Mel class to perform image <-> audio
-            batch_size (int): number of samples to generate
-            audio_file (str): must be a file on disk due to Librosa limitation or
-            raw_audio (np.ndarray): audio as numpy array
-            slice (int): slice number of audio to convert
+            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
+            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
+            `List[PIL Image]`: mel spectrograms (`float`, `List[np.ndarray]`): sample rate and raw audios
         """
 
         steps = steps or self.get_default_steps()
@@ -238,7 +462,7 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         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()
-        mel.set_resolution(x_res=input_dims[1], y_res=input_dims[0])
+        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]),
@@ -249,15 +473,15 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         mask = None
 
         if audio_file is not None or raw_audio is not None:
-            mel.load_audio(audio_file, raw_audio)
-            input_image = mel.audio_slice_to_image(slice)
+            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 hasattr(self, "vqvae"):
+            if self.vqvae is not None:
                 input_images = self.vqvae.encode(torch.unsqueeze(input_images, 0)).latent_dist.sample(
                     generator=generator
                 )[0]
@@ -266,7 +490,9 @@ class AudioDiffusionPipeline(DiffusionPipeline):
             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] * mel.get_sample_rate() / mel.x_res / mel.hop_length
+            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:]))
@@ -289,7 +515,7 @@ class AudioDiffusionPipeline(DiffusionPipeline):
                 if mask_end > 0:
                     images[:, :, :, -mask_end:] = mask[:, step, :, -mask_end:]
 
-        if hasattr(self, "vqvae"):
+        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"]
@@ -303,9 +529,9 @@ class AudioDiffusionPipeline(DiffusionPipeline):
             else map(lambda _: Image.fromarray(_, mode="RGB").convert("L"), images)
         )
 
-        audios = list(map(lambda _: mel.image_to_audio(_), images))
+        audios = list(map(lambda _: self.mel.image_to_audio(_), images))
         if not return_dict:
-            return images, (mel.get_sample_rate(), audios)
+            return images, (self.mel.get_sample_rate(), audios)
 
         return BaseOutput(**AudioPipelineOutput(np.array(audios)[:, np.newaxis, :]), **ImagePipelineOutput(images))
 
@@ -314,11 +540,11 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         """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)
+            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)
+            `np.ndarray`: noise tensor of shape (batch_size, 1, height, width)
         """
 
         # Only works with DDIM as this method is deterministic
@@ -351,24 +577,22 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         """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
+            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
+            `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)
 
 
-class LatentAudioDiffusionPipeline(AudioDiffusionPipeline):
-    def __init__(
-        self, unet: UNet2DConditionModel, scheduler: Union[DDIMScheduler, DDPMScheduler], vqvae: AutoencoderKL
-    ):
-        super().__init__(unet=unet, scheduler=scheduler)
-        self.register_modules(vqvae=vqvae)
+import diffusers
 
-    def __call__(self, *args, **kwargs):
-        return super().__call__(*args, **kwargs)
+diffusers.Mel = Mel
+setattr(diffusers, Mel.__name__, Mel)
+diffusers.AudioDiffusionPipeline = AudioDiffusionPipeline
+setattr(diffusers, AudioDiffusionPipeline.__name__, AudioDiffusionPipeline)
+diffusers.pipeline_utils.LOADABLE_CLASSES['diffusers']['Mel'] = ["save_pretrained", "from_pretrained"]

audiodiffusion/mel.py

@@ -1,129 +0,0 @@
-import warnings
-
-
-warnings.filterwarnings("ignore")
-
-import numpy as np  # noqa: E402
-
-import librosa  # noqa: E402
-from PIL import Image  # noqa: E402
-
-
-class Mel:
-    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,
-    ):
-        """Class to convert audio to mel spectrograms and vice versa.
-
-        Args:
-            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
-        """
-        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

notebooks/audio_diffusion_pipeline.ipynb

@@ -89,7 +89,7 @@
     "\n",
     "#@markdown teticio/audio-diffusion-instrumental-hiphop-256 - trained on instrumental hiphop\n",
     "\n",
-    "model_id = \"teticio/audio-diffusion-256-new\"  #@param [\"teticio/audio-diffusion-256\", \"teticio/audio-diffusion-breaks-256\", \"audio-diffusion-instrumenal-hiphop-256\", \"teticio/audio-diffusion-ddim-256\"]"
+    "model_id = \"teticio/audio-diffusion-256\"  #@param [\"teticio/audio-diffusion-256\", \"teticio/audio-diffusion-breaks-256\", \"audio-diffusion-instrumenal-hiphop-256\", \"teticio/audio-diffusion-ddim-256\"]"
    ]
   },
   {
@@ -356,7 +356,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "audio_diffusion = DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256-new').to(device)\n",
+    "audio_diffusion = DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256').to(device)\n",
     "mel = audio_diffusion.mel\n",
     "sample_rate = mel.get_sample_rate()"
    ]
@@ -532,7 +532,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_id = \"teticio/latent-audio-diffusion-ddim-256-new\"  #@param [\"teticio/latent-audio-diffusion-256\", \"teticio/latent-audio-diffusion-ddim-256\"]"
+    "model_id = \"teticio/latent-audio-diffusion-ddim-256\"  #@param [\"teticio/latent-audio-diffusion-256\", \"teticio/latent-audio-diffusion-ddim-256\"]"
    ]
   },
   {

notebooks/test_mel.ipynb

@@ -25,18 +25,6 @@
     "    pass"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "21f27189",
-   "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,
@@ -46,7 +34,7 @@
    "source": [
     "from datasets import load_dataset\n",
     "from IPython.display import Audio\n",
-    "from audiodiffusion.mel import Mel"
+    "from audiodiffusion import Mel"
    ]
   },
   {

notebooks/test_model.ipynb

@@ -49,7 +49,6 @@
     "import numpy as np\n",
     "from datasets import load_dataset\n",
     "from IPython.display import Audio\n",
-    "from audiodiffusion.mel import Mel\n",
     "from audiodiffusion import AudioDiffusion"
    ]
   },
@@ -60,7 +59,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "mel = Mel()\n",
     "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
     "generator = torch.Generator(device=device)"
    ]
@@ -104,7 +102,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "audio_diffusion = AudioDiffusion(model_id=model_id)"
+    "audio_diffusion = AudioDiffusion(model_id=model_id)\n",
+    "mel = audio_diffusion.pipe.mel"
    ]
   },
   {
@@ -336,7 +335,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "audio_diffusion = AudioDiffusion(model_id='teticio/audio-diffusion-ddim-256')"
+    "audio_diffusion = AudioDiffusion(model_id='teticio/audio-diffusion-ddim-256')\n",
+    "mel = audio_diffusion.pipe.mel"
    ]
   },
   {
@@ -507,7 +507,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_id = \"teticio/latent-audio-diffusion-ddim-256\"  #@param [\"teticio/latent-audio-diffusion-256\", \"teticio/latent-audio-diffusion-ddim-256\"]"
+    "model_id = \"teticio/latent-audio-diffusion-ddim-256-new\"  #@param [\"teticio/latent-audio-diffusion-256\", \"teticio/latent-audio-diffusion-ddim-256\"]"
    ]
   },
   {
@@ -517,7 +517,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "audio_diffusion = AudioDiffusion(model_id=model_id)"
+    "audio_diffusion = AudioDiffusion(model_id=model_id)\n",
+    "mel = audio_diffusion.pipe.mel"
    ]
   },
   {
@@ -568,9 +569,10 @@
    "source": [
     "generator.manual_seed(seed)\n",
     "latents = torch.randn((1, audio_diffusion.pipe.unet.in_channels,\n",
-    "                     audio_diffusion.pipe.unet.sample_size[0],\n",
-    "                     audio_diffusion.pipe.unet.sample_size[1]),\n",
-    "                    generator=generator)\n",
+    "                       audio_diffusion.pipe.unet.sample_size[0],\n",
+    "                       audio_diffusion.pipe.unet.sample_size[1]),\n",
+    "                      device=device,\n",
+    "                      generator=generator)\n",
     "latents.shape"
    ]
   },
@@ -583,9 +585,10 @@
    "source": [
     "generator.manual_seed(seed2)\n",
     "latents2 = torch.randn((1, audio_diffusion.pipe.unet.in_channels,\n",
-    "                      audio_diffusion.pipe.unet.sample_size[0],\n",
-    "                      audio_diffusion.pipe.unet.sample_size[1]),\n",
-    "                     generator=generator)\n",
+    "                        audio_diffusion.pipe.unet.sample_size[0],\n",
+    "                        audio_diffusion.pipe.unet.sample_size[1]),\n",
+    "                       device=device,\n",
+    "                       generator=generator)\n",
     "latents2.shape"
    ]
   },

requirements.txt

@@ -1,7 +1,7 @@
 torch
 numpy
 Pillow
-diffusers>=0.4.1
+diffusers>=0.9.0
 librosa
 datasets
 gradio

scripts/audio_to_images.py

@@ -9,7 +9,7 @@ import pandas as pd
 from tqdm.auto import tqdm
 from datasets import Dataset, DatasetDict, Features, Image, Value
 
-from audiodiffusion.mel import Mel
+from audiodiffusion import Mel
 
 logging.basicConfig(level=logging.WARN)
 logger = logging.getLogger('audio_to_images')

scripts/train_unconditional.py

@@ -1,5 +1,9 @@
 # based on https://github.com/huggingface/diffusers/blob/main/examples/train_unconditional.py
 
+# TODO
+# Migrate to diffusers
+# from diffusers.hub_utils import Repository
+
 import argparse
 import os
 
@@ -9,8 +13,13 @@ 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 (DiffusionPipeline, DDPMScheduler, UNet2DModel,
-                       DDIMScheduler, AutoencoderKL)
+from diffusers import (
+    #AudioDiffusionPipeline,
+    DDPMScheduler,
+    UNet2DModel,
+    DDIMScheduler,
+    AutoencoderKL,
+)
 from diffusers.hub_utils import init_git_repo, push_to_hub
 from diffusers.optimization import get_scheduler
 from diffusers.training_utils import EMAModel
@@ -23,8 +32,8 @@ import numpy as np
 from tqdm.auto import tqdm
 from librosa.util import normalize
 
-from audiodiffusion.mel import Mel
-from audiodiffusion import LatentAudioDiffusionPipeline, AudioDiffusionPipeline
+#from diffusers import Mel, AudioDiffusionPipeline
+from audiodiffusion import Mel, AudioDiffusionPipeline
 
 logger = get_logger(__name__)
 
@@ -59,7 +68,7 @@ def main(args):
             split="train",
         )
     # Determine image resolution
-    resolution = dataset[0]['image'].height, dataset[0]['image'].width
+    resolution = dataset[0]["image"].height, dataset[0]["image"].width
 
     augmentations = Compose([
         ToTensor(),
@@ -67,9 +76,9 @@ def main(args):
     ])
 
     def transforms(examples):
-        if args.vae is not None and vqvae.config['in_channels'] == 3:
+        if args.vae is not None and vqvae.config["in_channels"] == 3:
             images = [
-                augmentations(image.convert('RGB'))
+                augmentations(image.convert("RGB"))
                 for image in examples["image"]
             ]
         else:
@@ -85,32 +94,27 @@ def main(args):
         try:
             vqvae = AutoencoderKL.from_pretrained(args.vae)
         except EnvironmentError:
-            vqvae = LatentAudioDiffusionPipeline.from_pretrained(
+            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 = {
-            'LatentAudioDiffusionPipeline': LatentAudioDiffusionPipeline,
-            'AudioDiffusionPipeline': AudioDiffusionPipeline
-        }.get(
-            DiffusionPipeline.get_config_dict(
-                args.from_pretrained)['_class_name'], AudioDiffusionPipeline)
-        pipeline = pipeline.from_pretrained(args.from_pretrained)
+        pipeline = AudioDiffusionPipeline.from_pretrained(args.from_pretrained)
+        mel = pipeline.mel
         model = pipeline.unet
-        if hasattr(pipeline, 'vqvae'):
+        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'],
+            if vqvae is None else vqvae.config["latent_channels"],
             out_channels=1
-            if vqvae is None else vqvae.config['latent_channels'],
+            if vqvae is None else vqvae.config["latent_channels"],
             layers_per_block=2,
             block_out_channels=(128, 128, 256, 256, 512, 512),
             down_block_types=(
@@ -171,11 +175,13 @@ def main(args):
         run = os.path.split(__file__)[-1].split(".")[0]
         accelerator.init_trackers(run)
 
-    mel = Mel(x_res=resolution[1],
-              y_res=resolution[0],
-              hop_length=args.hop_length,
-              sample_rate=args.sample_rate,
-              n_fft=args.n_fft)
+    mel = Mel(
+        x_res=resolution[1],
+        y_res=resolution[0],
+        hop_length=args.hop_length,
+        sample_rate=args.sample_rate,
+        n_fft=args.n_fft,
+    )
 
     global_step = 0
     for epoch in range(args.num_epochs):
@@ -256,20 +262,14 @@ def main(args):
             if (
                     epoch + 1
             ) % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
-                if vqvae is not None:
-                    pipeline = LatentAudioDiffusionPipeline(
-                        unet=accelerator.unwrap_model(
-                            ema_model.averaged_model if args.use_ema else model
-                        ),
-                        vqvae=vqvae,
-                        scheduler=noise_scheduler)
-                else:
-                    pipeline = AudioDiffusionPipeline(
-                        unet=accelerator.unwrap_model(
-                            ema_model.averaged_model if args.use_ema else model
-                        ),
-                        scheduler=noise_scheduler,
-                    )
+                pipeline = AudioDiffusionPipeline(
+                    vqvae=vqvae,
+                    unet=accelerator.unwrap_model(
+                        ema_model.averaged_model if args.use_ema else model
+                    ),
+                    mel=mel,
+                    scheduler=noise_scheduler,
+                )
 
                 # save the model
                 if args.push_to_hub:
@@ -287,12 +287,13 @@ def main(args):
                     pipeline.save_pretrained(output_dir)
 
             if (epoch + 1) % args.save_images_epochs == 0:
-                generator = torch.manual_seed(42)
+                generator = torch.Generator(
+                    device=clean_images.device).manual_seed(42)
                 # run pipeline in inference (sample random noise and denoise)
                 images, (sample_rate, audios) = pipeline(
-                    mel=mel,
                     generator=generator,
                     batch_size=args.eval_batch_size,
+                    return_dict=False
                 )
 
                 # denormalize the images and save to tensorboard
@@ -373,10 +374,12 @@ if __name__ == "__main__":
                         type=str,
                         default="ddpm",
                         help="ddpm or ddim")
-    parser.add_argument("--vae",
-                        type=str,
-                        default=None,
-                        help="pretrained VAE model for latent diffusion")
+    parser.add_argument(
+        "--vae",
+        type=str,
+        default=None,
+        help="pretrained VAE model for latent diffusion",
+    )
 
     args = parser.parse_args()
     env_local_rank = int(os.environ.get("LOCAL_RANK", -1))

scripts/train_vae.py

@@ -17,7 +17,8 @@ from datasets import load_from_disk, load_dataset
 from pytorch_lightning.callbacks import Callback, ModelCheckpoint
 from pytorch_lightning.utilities.distributed import rank_zero_only
 
-from audiodiffusion.mel import Mel
+#from diffusers import Mel
+from audiodiffusion import Mel
 from audiodiffusion.utils import convert_ldm_to_hf_vae
 
 

setup.cfg

@@ -15,6 +15,6 @@ install_requires =
     torch
     numpy
     Pillow
-    diffusers>=0.4.1
+    diffusers>=0.9.0
     librosa
     datasets