sync with latest version of diffusers

audiodiffusion/__init__.py

@@ -8,6 +8,8 @@ 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
 
@@ -83,7 +85,8 @@ class AudioDiffusion:
                                      generator=generator,
                                      step_generator=step_generator,
                                      eta=eta,
-                                     noise=noise)
+                                     noise=noise,
+                                     return_dict=False)
         return images[0], (sample_rate, audios[0])
 
     def generate_spectrogram_and_audio_from_audio(
@@ -133,7 +136,8 @@ class AudioDiffusion:
                                      mask_end_secs=mask_end_secs,
                                      step_generator=step_generator,
                                      eta=eta,
-                                     noise=noise)
+                                     noise=noise,
+                                     return_dict=False)
         return images[0], (sample_rate, audios[0])
 
     @staticmethod
@@ -158,9 +162,7 @@ class AudioDiffusion:
 
 
 class AudioDiffusionPipeline(DiffusionPipeline):
-
-    def __init__(self, unet: UNet2DConditionModel,
-                 scheduler: Union[DDIMScheduler, DDPMScheduler]):
+    def __init__(self, unet: UNet2DConditionModel, scheduler: Union[DDIMScheduler, DDPMScheduler]):
         super().__init__()
         self.register_modules(unet=unet, scheduler=scheduler)
 
@@ -170,11 +172,13 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         Returns:
             Tuple: (height, width)
         """
-        input_module = self.vqvae if hasattr(self, 'vqvae') else self.unet
+        input_module = self.vqvae if hasattr(self, "vqvae") 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
+            (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:
@@ -200,8 +204,11 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         mask_end_secs: float = 0,
         step_generator: torch.Generator = None,
         eta: float = 0,
-        noise: torch.Tensor = None
-    ) -> Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]]:
+        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:
@@ -218,10 +225,10 @@ class AudioDiffusionPipeline(DiffusionPipeline):
             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()
@@ -229,89 +236,78 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         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)
+            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])
         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)
+                (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:
             mel.load_audio(audio_file, raw_audio)
             input_image = mel.audio_slice_to_image(slice)
-            input_image = np.frombuffer(input_image.tobytes(),
-                                        dtype="uint8").reshape(
-                                            (input_image.height,
-                                             input_image.width))
-            input_image = ((input_image / 255) * 2 - 1)
-            input_images = torch.tensor(input_image[np.newaxis, :, :],
-                                        dtype=torch.float)
-
-            if hasattr(self, 'vqvae'):
-                input_images = self.vqvae.encode(
-                    torch.unsqueeze(input_images,
-                                    0).to(self.device)).latent_dist.sample(
-                                        generator=generator).cpu()[0]
+            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"):
+                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])
+                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] * mel.get_sample_rate() / mel.x_res / 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:]))
+            mask = self.scheduler.add_noise(input_images, noise, torch.tensor(self.scheduler.timesteps[start_step:]))
 
-        images = images.to(self.device)
-        for step, t in enumerate(
-                self.progress_bar(self.scheduler.timesteps[start_step:])):
-            model_output = self.unet(images, t)['sample']
+        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']
+                    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']
+                    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]
+                    images[:, :, :, :mask_start] = mask[:, step, :, :mask_start]
                 if mask_end > 0:
                     images[:, :, :, -mask_end:] = mask[:, step, :, -mask_end:]
 
-        if hasattr(self, 'vqvae'):
+        if hasattr(self, "vqvae"):
             # 0.18215 was scaling factor used in training to ensure unit variance
             images = 1 / 0.18215 * images
-            images = self.vqvae.decode(images)['sample']
+            images = 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))
+            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 _: mel.image_to_audio(_), images))
-        return images, (mel.get_sample_rate(), audios)
+        if not return_dict:
+            return images, (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:
@@ -328,35 +324,30 @@ class AudioDiffusionPipeline(DiffusionPipeline):
         # 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 = 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)
+        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)
+            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
+            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:
+    def slerp(x0: torch.Tensor, x1: torch.Tensor, alpha: float) -> torch.Tensor:
         """Spherical Linear intERPolation
 
         Args:
@@ -368,18 +359,14 @@ class AudioDiffusionPipeline(DiffusionPipeline):
             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)
+        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):
+    def __init__(
+        self, unet: UNet2DConditionModel, scheduler: Union[DDIMScheduler, DDPMScheduler], vqvae: AutoencoderKL
+    ):
         super().__init__(unet=unet, scheduler=scheduler)
         self.register_modules(vqvae=vqvae)
 

audiodiffusion/mel.py

@@ -1,22 +1,25 @@
 import warnings
 
-warnings.filterwarnings('ignore')
 
-import librosa
-import numpy as np
-from PIL import Image
+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 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:
@@ -28,17 +31,26 @@ class Mel:
             top_db (int): loudest in decibels
             n_iter (int): number of iterations for Griffin Linn mel inversion
         """
-        self.x_res = x_res
-        self.y_res = y_res
+        self.hop_length = hop_length
         self.sr = sample_rate
         self.n_fft = n_fft
-        self.hop_length = hop_length
-        self.n_mels = self.y_res
-        self.slice_size = self.x_res * self.hop_length - 1
         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.
 
@@ -53,10 +65,7 @@ class Mel:
 
         # 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), ))
-            ])
+            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.
@@ -75,8 +84,7 @@ class Mel:
         Returns:
             np.ndarray: audio as numpy array
         """
-        return self.audio[self.slice_size * slice:self.slice_size *
-                          (slice + 1)]
+        return self.audio[self.slice_size * slice : self.slice_size * (slice + 1)]
 
     def get_sample_rate(self) -> int:
         """Get sample rate:
@@ -95,14 +103,11 @@ class Mel:
         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)
+        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)
+        bytedata = (((log_S + self.top_db) * 255 / self.top_db).clip(0, 255) + 0.5).astype(np.uint8)
         image = Image.fromarray(bytedata)
         return image
 
@@ -115,14 +120,10 @@ class Mel:
         Returns:
             audio (np.ndarray): raw audio
         """
-        bytedata = np.frombuffer(image.tobytes(), dtype="uint8").reshape(
-            (image.height, image.width))
+        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)
+            S, sr=self.sr, n_fft=self.n_fft, hop_length=self.hop_length, n_iter=self.n_iter
+        )
         return audio

notebooks/test_model.ipynb

@@ -61,7 +61,8 @@
    "outputs": [],
    "source": [
     "mel = Mel(x_res=256, y_res=256)\n",
-    "generator = torch.Generator()"
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "generator = torch.Generator(device=device)"
    ]
   },
   {
@@ -160,7 +161,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "seed = 16183389798189209330  #@param {type:\"integer\"}\n",
+    "seed = 2391504374279719  #@param {type:\"integer\"}\n",
     "generator.manual_seed(seed)\n",
     "image, (sample_rate, audio) = audio_diffusion.generate_spectrogram_and_audio(\n",
     "    generator=generator)\n",
@@ -270,7 +271,7 @@
     "overlap_samples = overlap_secs * mel.get_sample_rate()\n",
     "slice_size = mel.x_res * mel.hop_length\n",
     "stride = slice_size - overlap_samples\n",
-    "generator = torch.Generator()\n",
+    "generator = torch.Generator(device=device)\n",
     "seed = generator.seed()\n",
     "print(f'Seed = {seed}')\n",
     "track = np.array([])\n",
@@ -315,7 +316,7 @@
     "        raw_audio=raw_audio,\n",
     "        mask_start_secs=1,\n",
     "        mask_end_secs=1,\n",
-    "        step_generator=torch.Generator())\n",
+    "        step_generator=torch.Generator(device=device))\n",
     "display(Audio(audio, rate=sample_rate))\n",
     "display(Audio(audio2, rate=sample_rate))"
    ]
@@ -526,7 +527,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "seed = 6015487092443227811  #@param {type:\"integer\"}\n",
+    "seed = 3412253600050855  #@param {type:\"integer\"}\n",
     "generator.manual_seed(seed)\n",
     "image, (sample_rate, audio) = audio_diffusion.generate_spectrogram_and_audio(\n",
     "    generator=generator)\n",
@@ -541,7 +542,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "seed2 = 5623685468252603494  #@param {type:\"integer\"}\n",
+    "seed2 = 7016114633369557  #@param {type:\"integer\"}\n",
     "generator.manual_seed(seed2)\n",
     "image2, (sample_rate, audio2) = audio_diffusion.generate_spectrogram_and_audio(\n",
     "    generator=generator)\n",