Update app.py

app.py

@@ -4,31 +4,33 @@ import demucs.separate
 import shlex
 import os
 import spaces
+import subprocess
 
 # Check if CUDA is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Define the inference function
 @spaces.GPU
-def inference(audio_file, model_name, two_stems, mp3, mp3_bitrate):
+def inference(audio_file, model_name, vocals, drums, bass, other, mp3, mp3_bitrate):
     """
-    Performs inference using Demucs.
+    Performs inference using Demucs and mixes the selected stems.
 
     Args:
         audio_file: The audio file to separate.
         model_name: The name of the Demucs model to use.
-        two_stems: The name of the stem to separate (for two-stems mode).
+        vocals: Whether to include vocals in the mix.
+        drums: Whether to include drums in the mix.
+        bass: Whether to include bass in the mix.
+        other: Whether to include other instruments in the mix.
         mp3: Whether to save the output as MP3.
         mp3_bitrate: The bitrate of the output MP3 file.
 
     Returns:
-        A list of separated audio files.
+        The path to the mixed audio file.
     """
 
-    # Construct the command line arguments
+    # Construct the command line arguments for Demucs
     cmd = f"demucs -n {model_name} --clip-mode clamp --shifts=1"
-    if two_stems:
-        cmd += f" --two-stems={two_stems}"
     if mp3:
         cmd += f" --mp3 --mp3-bitrate={mp3_bitrate}"
     cmd += f" {audio_file.name}"
@@ -38,9 +40,28 @@ def inference(audio_file, model_name, two_stems, mp3, mp3_bitrate):
 
     # Get the output file paths
     output_dir = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file.name))[0])
-    output_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))]
+    stems = {
+        "vocals": os.path.join(output_dir, "vocals.wav"),
+        "drums": os.path.join(output_dir, "drums.wav"),
+        "bass": os.path.join(output_dir, "bass.wav"),
+        "other": os.path.join(output_dir, "other.wav"),
+    }
 
-    return output_files
+    # Mix the selected stems
+    selected_stems = [stems[stem] for stem, include in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]) if include]
+    if not selected_stems:
+        raise gr.Error("Please select at least one stem to mix.")
+
+    output_file = os.path.join(output_dir, "mixed.wav")
+    if len(selected_stems) == 1:
+        # If only one stem is selected, just copy it
+        os.rename(selected_stems[0], output_file)
+    else:
+        # Otherwise, use ffmpeg to mix the stems
+        ffmpeg_cmd = ["ffmpeg", "-y", "-i"] + selected_stems + ["-filter_complex", "amix=inputs=" + str(len(selected_stems)) + ":duration=longest", output_file]
+        subprocess.run(ffmpeg_cmd, check=True)
+
+    return output_file
 
 # Define the Gradio interface
 iface = gr.Interface(
@@ -48,13 +69,16 @@ iface = gr.Interface(
     inputs=[
         gr.Audio(type="filepath"),
         gr.Dropdown(["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"], label="Model Name"),
-        gr.Dropdown(["vocals", "drums", "bass", "other"], label="Two Stems (Optional)"),  # Remove optional=True
+        gr.Checkbox(label="Vocals", value=True),
+        gr.Checkbox(label="Drums", value=True),
+        gr.Checkbox(label="Bass", value=True),
+        gr.Checkbox(label="Other", value=True),
         gr.Checkbox(label="Save as MP3"),
         gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False),
     ],
-    outputs=gr.Files(),
-    title="Demucs Music Source Separation",
-    description="Separate vocals, drums, bass, and other instruments from your music using Demucs.",
+    outputs=gr.Audio(type="filepath"),
+    title="Demucs Music Source Separation and Mixing",
+    description="Separate vocals, drums, bass, and other instruments from your music using Demucs and mix the selected stems.",
 )
 
 # Launch the Gradio interface