updated boxes

app.py

@@ -59,26 +59,27 @@
 # ).launch()
 
 
-
 import os
 import gradio as gr
 from scipy.io.wavfile import write
 import subprocess
 import torch
 
-from audio_separator import Separator
+from audio_separator import Separator  # Ensure this is correctly implemented
 
-def inference(audio):
+def inference(audio, vocals, bass, drums, other, piano, guitar, lead_vocals, backing_vocals):
     os.makedirs("out", exist_ok=True)
     audio_path = 'test.wav'
     write(audio_path, audio[0], audio[1])
-
-    # Check for CUDA availability
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    print(f"Using device: {device}")
+    if device=='cuda':
+        use_cuda=True
+        print(f"Using device: {device}")
+    else:
+        use_cuda=False
+        print(f"Using device: {device}")
 
     try:
-        # Using subprocess.run for better control
         command = f"python3 -m demucs.separate -n htdemucs_6s -d {device} {audio_path} -o out"
         process = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         print("Demucs script output:", process.stdout.decode())
@@ -86,9 +87,7 @@ def inference(audio):
         print("Error in Demucs script:", e.stderr.decode())
         return None
 
-    use_cuda = device == 'cuda'
     try:
-        # Separating the stems using your custom separator
         separator = Separator("./out/htdemucs_6s/test/vocals.wav", model_name='UVR_MDXNET_KARA_2', use_cuda=use_cuda, output_format='mp3')
         primary_stem_path, secondary_stem_path = separator.separate()
     except Exception as e:
@@ -96,42 +95,42 @@ def inference(audio):
         return None
 
     # Collecting all file paths
-    stem_names = ["vocals", "bass", "drums", "other", "piano", "guitar", "lead_vocals", "backing_vocals"]
-    files = {stem: f"./out/htdemucs_6s/test/{stem}.wav" for stem in stem_names[:6]}
-    files["lead_vocals"] = primary_stem_path
-    files["backing_vocals"] = secondary_stem_path
-
-    # Check if files exist and return dictionary
-    existing_files = {stem: file for stem, file in files.items() if os.path.isfile(file)}
-    if not existing_files:
-        print("No files were created.")
-        return None
-
-    return existing_files
-
-# Function to return selected audios along with their labels
-def get_selected_audios(stems, *selected_stems):
-    return [(stem, stems[stem]) for stem in selected_stems if stem in stems and stems[stem]]
-
-# Custom output component
-def custom_output_component(label, filepath):
-    if filepath:
-        return gr.Audio(filepath, label=label, visible=True)
-    else:
-        return gr.Label("", visible=False)
+    stem_files = {
+        "vocals": "./out/htdemucs_6s/test/vocals.wav",
+        "bass": "./out/htdemucs_6s/test/bass.wav",
+        "drums": "./out/htdemucs_6s/test/drums.wav",
+        "other": "./out/htdemucs_6s/test/other.wav",
+        "piano": "./out/htdemucs_6s/test/piano.wav",
+        "guitar": "./out/htdemucs_6s/test/guitar.wav",
+        "lead_vocals": primary_stem_path,
+        "backing_vocals": secondary_stem_path
+    }
+
+    # Filter out unchecked stems
+    selected_stems = {
+        "vocals": vocals,
+        "bass": bass,
+        "drums": drums,
+        "other": other,
+        "piano": piano,
+        "guitar": guitar,
+        "lead_vocals": lead_vocals,
+        "backing_vocals": backing_vocals
+    }
+
+    return [stem_files[stem] if selected_stems[stem] and os.path.isfile(stem_files[stem]) else None for stem in selected_stems]
+
+# Checkbox for each stem
+checkboxes = [gr.components.Checkbox(label=stem) for stem in ["Vocals", "Bass", "Drums", "Other", "Piano", "Guitar", "Lead Vocals", "Backing Vocals"]]
 
 # Gradio Interface
 title = "Source Separation Demo"
-description = "Music Source Separation in the Waveform Domain. To use it, simply upload your audio."
-
-audio_input = gr.Audio(type="numpy", label="Input")
-checkboxes = [gr.Checkbox(label=stem) for stem in ["vocals", "bass", "drums", "other", "piano", "guitar", "lead_vocals", "backing_vocals"]]
-outputs = [gr.Dynamic(custom_output_component, label=stem) for stem in ["vocals", "bass", "drums", "other", "piano", "guitar", "lead_vocals", "backing_vocals"]]
-
+description = "Music Source Separation in the Waveform Domain. To use it, simply upload your audio and select the stems you want to display."
 gr.Interface(
-    fn=lambda audio, *args: get_selected_audios(inference(audio), *args),
-    inputs=[audio_input, *checkboxes],
-    outputs=outputs,
+    inference, 
+    [gr.components.Audio(type="numpy", label="Input")] + checkboxes,
+    [gr.components.Audio(type="filepath", label=stem, optional=True) for stem in ["Vocals", "Bass", "Drums", "Other", "Piano", "Guitar", "Lead Vocals", "Backing Vocals"]],
     title=title,
-    description=description
+    description=description,
+    live=True  # Enable live updates
 ).launch()