Update app.py

app.py

@@ -9,7 +9,9 @@ import subprocess
 import torch
 from demucs import pretrained
 from demucs.apply import apply_model
-from demucs.audio import load_audio, save_audio
+import torchaudio
+import torch
+from pathlib import Path
 
 # Helper functions
 def audiosegment_to_array(audio):
@@ -72,14 +74,42 @@ def apply_treble_boost(audio, gain=10):
 # Vocal Isolation using Demucs
 def apply_vocal_isolation(audio_path):
     model = pretrained.get_model(name='htdemucs')
-    wav = load_audio(audio_path)
+    wav = load_track_local(audio_path, model.samplerate, channels=2)  # stereo
     ref = wav.mean(0)
     wav -= ref[:, None]
     sources = apply_model(model, wav[None])[0]
     wav += ref[:, None]
-    vocal_track = sources[3]  # index 3 = vocals
+
+    # Get vocals (index 3)
+    vocal_track = sources[3].cpu()
+
     out_path = os.path.join(tempfile.gettempdir(), "vocals.wav")
-    save_audio(vocal_track, out_path, samplerate=model.samplerate)
+    save_track(out_path, vocal_track, model.samplerate)
+    return out_path
+
+
+# Local copy of helper functions from demucs
+def load_track(track, sample_rate, mono=True):
+    wav, sr = torchaudio.load(str(track))
+    if sr != sample_rate:
+        wav = torchaudio.functional.resample(wav, sr, sample_rate)
+    if mono and wav.shape[0] == 2:
+        wav = wav.mean(0)
+    return wav
+
+
+def load_track_local(path, sample_rate, channels=2):
+    sig, rate = torchaudio.load(path)
+    if rate != sample_rate:
+        sig = torchaudio.functional.resample(sig, rate, sample_rate)
+    if channels == 1:
+        sig = sig.mean(0)
+    return sig
+
+
+def save_track(path, wav, sample_rate):
+    path = Path(path)
+    torchaudio.save(str(path), wav, sample_rate)
     return out_path
 
 # Apply selected effects in order