fix imports

audiodiffusion/__init__.py

@@ -9,7 +9,7 @@ from tqdm.auto import tqdm
 # from diffusers import AudioDiffusionPipeline
 from .pipeline_audio_diffusion import AudioDiffusionPipeline
 
-VERSION = "1.4.0"
+VERSION = "1.4.1"
 
 
 class AudioDiffusion:

audiodiffusion/mel.py

@@ -23,8 +23,21 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 
 warnings.filterwarnings("ignore")
 
-import librosa  # noqa: E402
 import numpy as np  # noqa: E402
+
+
+try:
+    import librosa  # noqa: E402
+
+    _librosa_can_be_imported = True
+    _import_error = ""
+except Exception as e:
+    _librosa_can_be_imported = False
+    _import_error = (
+        f"Cannot import librosa because {e}. Make sure to correctly install librosa to be able to install it."
+    )
+
+
 from PIL import Image  # noqa: E402
 
 
@@ -61,6 +74,9 @@ class Mel(ConfigMixin, SchedulerMixin):
         self.set_resolution(x_res, y_res)
         self.audio = None
 
+        if not _librosa_can_be_imported:
+            raise ValueError(_import_error)
+
     def set_resolution(self, x_res: int, y_res: int):
         """Set resolution.
 
@@ -87,12 +103,7 @@ class Mel(ConfigMixin, SchedulerMixin):
 
         # 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.
@@ -131,11 +142,7 @@ class Mel(ConfigMixin, SchedulerMixin):
             `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,
+            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)
@@ -155,10 +162,6 @@ class Mel(ConfigMixin, SchedulerMixin):
         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

audiodiffusion/pipeline_audio_diffusion.py

@@ -21,13 +21,12 @@ from typing import List, Tuple, Union
 
 import numpy as np
 import torch
-from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, Mel, UNet2DConditionModel
+from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, 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