Update app.py

app.py

@@ -1,11 +1,14 @@
 import gradio as gr
 import torch
-import demucs.api
+# import demucs.api
 import os
 import spaces
 import subprocess
 from pydub import AudioSegment
 from typing import Tuple, Dict, List
+from demucs.apply import apply_model
+from demucs.separate import load_track, save_tracks
+from demucs.pretrained import get_model
 
 # check if cuda is available
 device: str = "cuda" if torch.cuda.is_available() else "cpu"
@@ -46,21 +49,28 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     """
 
     # initialize demucs separator
-    separator: demucs.api.Separator = demucs.api.Separator(model=model_name)
+    # separator: demucs.api.Separator = demucs.api.Separator(model=model_name)
+    separator = get_model(name=model_name)
 
     # separate audio file and capture log
     import io
     log_stream = io.StringIO()
-    origin, separated = separator.separate_audio_file(audio_file, progress=True, log_stream=log_stream)
+    # origin, separated = separator.separate_audio_file(audio_file, progress=True, log_stream=log_stream)
+    wav = load_track(audio_file, separator.samplerate, channels=separator.audio_channels)
+    ref = wav.mean(0)
+    wav = (wav - ref.view(1, -1)).to(device)
+    sources = apply_model(separator, wav, device=device, progress=True, log_stream=log_stream)
+    sources = sources * ref.view(1, -1) + ref.view(1, -1)
     separation_log = log_stream.getvalue()
 
     # get the output file paths
     output_dir: str = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file))[0])
     os.makedirs(output_dir, exist_ok=True)  # create output directory if it doesn't exist
     stems: Dict[str, str] = {}
-    for stem, source in separated.items():
+    for stem, source in zip(separator.sources, sources):
         stem_path: str = os.path.join(output_dir, f"{stem}.wav")
-        demucs.api.save_audio(source, stem_path, samplerate=separator.samplerate)
+        # demucs.api.save_audio(source, stem_path, samplerate=separator.samplerate)
+        save_tracks(source, stem_path, separator.samplerate)
         stems[stem] = stem_path
 
     # mix the selected stems