Update app.py

app.py

@@ -2,15 +2,56 @@ import os
 import gradio as gr
 from audio_separator.separator import Separator
 
+# Define models
+models = {
+    "Roformer": {
+        'BS-Roformer-Viperx-1297.ckpt': 'model_bs_roformer_ep_317_sdr_12.9755.ckpt',
+        'BS-Roformer-Viperx-1296.ckpt': 'model_bs_roformer_ep_368_sdr_12.9628.ckpt',
+        'BS-Roformer-Viperx-1053.ckpt': 'model_bs_roformer_ep_937_sdr_10.5309.ckpt',
+        'Mel-Roformer-Viperx-1143.ckpt': 'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt'
+    },
+    "MDX23C": [
+        'MDX23C_D1581.ckpt',
+        'MDX23C-8KFFT-InstVoc_HQ.ckpt',
+        'MDX23C-8KFFT-InstVoc_HQ_2.ckpt',
+    ],
+    "MDXNet": [
+        'UVR-MDX-NET-Inst_full_292.onnx',
+        'UVR-MDX-NET_Inst_187_beta.onnx',
+        # Add remaining models as needed
+    ],
+    "VRArch": [
+        '1_HP-UVR.pth',
+        '2_HP-UVR.pth',
+        # Add remaining models as needed
+    ],
+    "Demucs": [
+        'htdemucs_ft.yaml',
+        'htdemucs.yaml',
+        # Add remaining models as needed
+    ]
+}
+
 # Function for separating the audio
-def separate_audio(audio):
+def separate_audio(audio, model_name):
     if audio is None:
         return None, None
 
-    # Step 1: Get the audio file path
-    input_file = audio
+    # Get the selected model path
+    model_path = None
+    for model_group in models.values():
+        if isinstance(model_group, dict):
+            model_path = model_group.get(model_name)
+            if model_path:
+                break
+        elif model_name in model_group:
+            model_path = model_name
+            break
+
+    if model_path is None:
+        return None, None
 
-    # Step 2: Set up output directory and models for separation
+    # Set up output directory and models for separation
     output_dir = "./output"
     os.makedirs(output_dir, exist_ok=True)
 
@@ -21,13 +62,12 @@ def separate_audio(audio):
     instrumental_path = os.path.join(output_dir, 'Instrumental.wav')
 
     with gr.Progress() as progress:
-        # Load model
         progress(0, "Loading model...")
-        separator.load_model(model_filename='model_bs_roformer_ep_317_sdr_12.9755.ckpt')
+        separator.load_model(model_filename=model_path)
         progress(20, "Model loaded. Starting separation...")
 
         # Step 3: Splitting track into Vocal and Instrumental
-        voc_inst = separator.separate(input_file)
+        voc_inst = separator.separate(audio)
 
         # Check if separation was successful
         if len(voc_inst) != 2:
@@ -49,6 +89,8 @@ with gr.Blocks(theme="NoCrypt/[email protected]") as demo:
     with gr.Row():
         with gr.Column():
             link_input = gr.Audio(label="Upload Audio File", type="filepath")
+            model_dropdown = gr.Dropdown(label="Select Model", choices=list(
+                models["Roformer"].keys()) + models["MDX23C"] + models["MDXNet"] + models["VRArch"] + models["Demucs"])
             separate_button = gr.Button("Separate Audio")
         
         with gr.Column():
@@ -58,7 +100,7 @@ with gr.Blocks(theme="NoCrypt/[email protected]") as demo:
     # Define button functionality
     separate_button.click(
         separate_audio,
-        inputs=[link_input],
+        inputs=[link_input, model_dropdown],
         outputs=[
             instrumental_output,
             vocals_output,