Abstracted model loading, added progress bar and increased threads.

app.py

@@ -19,11 +19,12 @@ STEM_CHOICES = {
 }
 
 
-def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
-    model = pretrained.get_model(model_name)
-    model.to('cuda' if torch.cuda.is_available() else 'cpu')
+models = dict(zip(DEMUX_MODELS, [pretrained.get_model(m) for m in DEMUX_MODELS]))
+
+for model in models.values():
     model.eval()
 
+def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
     waveform, sr = torchaudio.load(audio_file_path)
     is_mono = waveform.shape[0] == 1
     if is_mono:
@@ -31,11 +32,13 @@ def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> Au
 
     with torch.no_grad():
         stems_batch = apply_model(
-            model,
+            models[model_name],
             waveform.unsqueeze(0),
             overlap=0.2,
             shifts=1,
-            split=True
+            split=True,
+            progress=True,
+            num_workers=4
         )
 
     stems = stems_batch[0]
@@ -76,13 +79,13 @@ model_card = ModelCard(
 with gr.Blocks() as demo:
 
     dropdown_model = gr.Dropdown(
-        label="Select Demucs Model",
+        label="Demucs Model",
         choices=DEMUX_MODELS,
         value="mdx_extra_q"
     )
 
     dropdown_stem = gr.Dropdown(
-        label="Select Stem to Separate",
+        label="Stem to Separate",
         choices=list(STEM_CHOICES.keys()),
         value="Vocals"
     )