Removed choice and fixed to a single model.

app.py

@@ -8,7 +8,7 @@ from typing import Dict
 from pyharp import *
 
 
-DEMUX_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]
+#DEMUX_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]
 
 STEM_CHOICES = {
     "Vocals": 3,
@@ -19,10 +19,13 @@ STEM_CHOICES = {
 }
 
 
-models = dict(zip(DEMUX_MODELS, [pretrained.get_model(m) for m in DEMUX_MODELS]))
+#models = dict(zip(DEMUX_MODELS, [pretrained.get_model(m) for m in DEMUX_MODELS]))
+
+#for model in models.values():
+    #model.eval()
+
+model = pretrained.get_model('mdx_extra_q')
 
-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)
@@ -32,7 +35,7 @@ def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> Au
 
     with torch.no_grad():
         stems_batch = apply_model(
-            models[model_name],
+            model,
             waveform.unsqueeze(0),
             overlap=0.2,
             shifts=1,
@@ -56,13 +59,13 @@ def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> Au
 
 # Gradio Callback Function
 
-def process_fn_stem(audio_file_path: str, demucs_model: str, stem_choice: str):
+def process_fn_stem(audio_file_path: str, stem_choice: str):
     """
     PyHARP process function:
       - Separates the chosen stem using Demucs.
       - Saves the stem as a .wav file.
     """
-    stem_signal = separate_stem(audio_file_path, model_name=demucs_model, stem_choice=stem_choice)
+    stem_signal = separate_stem(audio_file_path, model_name='', stem_choice=stem_choice)
     stem_path = save_audio(stem_signal, f"{stem_choice.lower().replace(' ', '_')}.wav")
     return stem_path, LabelList(labels=[])
 
@@ -78,11 +81,11 @@ model_card = ModelCard(
 # Gradio UI
 with gr.Blocks() as demo:
 
-    dropdown_model = gr.Dropdown(
-        label="Demucs Model",
-        choices=DEMUX_MODELS,
-        value="mdx_extra_q"
-    )
+    #dropdown_model = gr.Dropdown(
+    #    label="Demucs Model",
+    #    choices=DEMUX_MODELS,
+    #    value="mdx_extra_q"
+    #)
 
     dropdown_stem = gr.Dropdown(
         label="Stem to Separate",
@@ -92,7 +95,7 @@ with gr.Blocks() as demo:
 
     app = build_endpoint(
         model_card=model_card,
-        components=[dropdown_model, dropdown_stem],
+        components=[dropdown_stem],
         process_fn=process_fn_stem
     )