Refactor audio loading to use librosa for consistency and improved handling of stereo audio

utils/extract_conditions.py

@@ -1,15 +1,10 @@
-import torchaudio
 import numpy as np
-from scipy.signal import savgol_filter
-import librosa
 import torch
 import torchaudio
+import librosa
 import scipy.signal as signal
+from scipy.signal import savgol_filter
 from torchaudio import transforms as T
-import torch
-import torchaudio
-import librosa
-import numpy as np
 
 
 def compute_melody_v2(stereo_audio: torch.Tensor) -> np.ndarray:
@@ -23,7 +18,10 @@ def compute_melody_v2(stereo_audio: torch.Tensor) -> np.ndarray:
            每一列代表： [L1, R1, L2, R2, L3, R3, L4, R4]（按 frame 交錯），
            且每個值都 ∈ {1, 2, …, 128}，對應 CQT 的頻率 bin。
     """
-    audio, sr = torchaudio.load(stereo_audio)
+    audio_np, sr = librosa.load(stereo_audio, sr=None, mono=False)
+    if audio_np.ndim == 1:
+        audio_np = np.expand_dims(audio_np, 0)
+    audio = torch.as_tensor(audio_np, dtype=torch.float32)
     # 1. 先針對左、右聲道分別計算 CQT (128 bins)，回傳 cqt_db 形狀都是 (128, T_frames)
     cqt_left  = compute_music_represent(audio[0], sr)  # shape: (128, T_frames)
     cqt_right = compute_music_represent(audio[1], sr)  # shape: (128, T_frames)
@@ -95,7 +93,10 @@ def compute_melody(input_audio):
     sample_rate = 44100
 
     # Load audio file
-    wav, sr = torchaudio.load(input_audio)
+    wav_np, sr = librosa.load(input_audio, sr=None, mono=False)
+    if wav_np.ndim == 1:
+        wav_np = np.expand_dims(wav_np, 0)
+    wav = torch.as_tensor(wav_np, dtype=torch.float32)
     if sr != sample_rate:
         resample = torchaudio.transforms.Resample(orig_freq=sr, new_freq=sample_rate)
         wav = resample(wav)
@@ -127,7 +128,10 @@ def compute_dynamics(audio_file, hop_length=160, target_sample_rate=44100, cut=T
         dynamics_curve (numpy.ndarray): The computed dynamic values in dB.
     """
     # Load audio file
-    waveform, original_sample_rate = torchaudio.load(audio_file)
+    waveform_np, original_sample_rate = librosa.load(audio_file, sr=None, mono=False)
+    if waveform_np.ndim == 1:
+        waveform_np = np.expand_dims(waveform_np, 0)
+    waveform = torch.as_tensor(waveform_np, dtype=torch.float32)
     if original_sample_rate != target_sample_rate:
         resampler = torchaudio.transforms.Resample(orig_freq=original_sample_rate, new_freq=target_sample_rate)
         waveform = resampler(waveform)
@@ -173,9 +177,12 @@ def extract_melody_one_hot(audio_path,
         One-hot chromagram of the most prominent pitch class per frame.
     """
     # ---------------------------------------------------------
-    # 1. Load audio (Torchaudio => shape: (channels, samples))
+    # 1. Load audio (librosa => shape: (channels, samples))
     # ---------------------------------------------------------
-    audio, in_sr = torchaudio.load(audio_path)
+    audio_np, in_sr = librosa.load(audio_path, sr=None, mono=False)
+    if audio_np.ndim == 1:
+        audio_np = np.expand_dims(audio_np, 0)
+    audio = torch.as_tensor(audio_np, dtype=torch.float32)
 
     # Convert to mono by averaging channels: shape => (samples,)
     audio_mono = audio.mean(dim=0)

utils/stable_audio_dataset_utils.py

@@ -1,15 +1,20 @@
 import math
 import random
-import torch
-from torch import nn
 from typing import Tuple
-import torchaudio
+
+import librosa
+import numpy as np
+import torch
 import torch.nn.functional as F
+from torch import nn
 from torchaudio import transforms as T
 
 def load_audio_file(filename, target_sr=44100, target_samples=2097152):
     try:
-        audio, in_sr = torchaudio.load(filename)    
+        audio_np, in_sr = librosa.load(filename, sr=None, mono=False)
+        if audio_np.ndim == 1:
+            audio_np = np.expand_dims(audio_np, 0)
+        audio = torch.as_tensor(audio_np, dtype=torch.float32)
         # Resample if necessary
         if in_sr != target_sr:
             resampler = T.Resample(in_sr, target_sr)
@@ -126,4 +131,4 @@ class Stereo(nn.Module):
         elif signal_shape[0] > 2: #?, s -> 2,s
             signal = signal[:2, :]    
 
-    return signal
+    return signal