add ref

teticio · May 17, 2023 · a923a95 · a923a95
1 parent 110c8a4
commit a923a95
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/audiodiffusion/__init__.py b/audiodiffusion/__init__.py
@@ -9,7 +9,7 @@
 # from diffusers import AudioDiffusionPipeline
 from .pipeline_audio_diffusion import AudioDiffusionPipeline
 
-VERSION = "1.5.3"
+VERSION = "1.5.4"
 
 
 class AudioDiffusion:

diff --git a/audiodiffusion/mel.py b/audiodiffusion/mel.py
@@ -132,7 +132,7 @@ def get_sample_rate(self) -> int:
         """
         return self.sr
 
-    def audio_slice_to_image(self, slice: int) -> Image.Image:
+    def audio_slice_to_image(self, slice: int, ref=np.max) -> Image.Image:
         """Convert slice of audio to spectrogram.
 
         Args:
@@ -144,7 +144,7 @@ def audio_slice_to_image(self, slice: int) -> Image.Image:
         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)
+        log_S = librosa.power_to_db(S, ref=ref, 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