Update app.py

app.py

@@ -13,7 +13,7 @@ import time
 console = Console()
 
 @spaces.GPU
-def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, gr.HTML]:
+def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, List[str], gr.HTML]:
     log_messages = []
 
     def stream_log(message, style=""):
@@ -21,14 +21,14 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
         log_messages.append(formatted_message)
         return gr.HTML(f"<pre style='margin-bottom: 0;{style}'>{formatted_message}</pre>")
 
-    yield None, stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;")
+    yield None, None, stream_log("Initializing Demucs...", "color: #4CAF50; font-weight: bold;")
     time.sleep(1)  # Simulate initialization time
 
-    yield None, stream_log("Loading audio file...", "color: #2196F3;")
+    yield None, None, stream_log("Loading audio file...", "color: #2196F3;")
     time.sleep(0.5)  # Simulate loading time
     
     if audio_file is None:
-        yield None, stream_log("Error: No audio file provided", "color: #F44336;")
+        yield None, None, stream_log("Error: No audio file provided", "color: #F44336;")
         raise gr.Error("Please upload an audio file")
 
     # Use absolute paths
@@ -44,7 +44,7 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
         audio_file
     ]
 
-    yield None, stream_log("Preparing separation process...", "color: #FF9800;")
+    yield None, None, stream_log("Preparing separation process...", "color: #FF9800;")
     time.sleep(0.5)  # Simulate preparation time
 
     try:
@@ -68,85 +68,81 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
 
         if process.returncode != 0:
             error_output = process.stderr.read()
-            yield None, stream_log(f"Error: Separation failed", "color: #F44336;")
+            yield None, None, stream_log(f"Error: Separation failed", "color: #F44336;")
             raise gr.Error(f"Demucs separation failed. Check the logs for details.")
 
     except Exception as e:
-        yield None, stream_log(f"Unexpected error: {str(e)}", "color: #F44336;")
+        yield None, None, stream_log(f"Unexpected error: {str(e)}", "color: #F44336;")
         raise gr.Error(f"An unexpected error occurred: {str(e)}")
 
-    yield None, stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;")
+    yield None, None, stream_log("Separation completed successfully!", "color: #4CAF50; font-weight: bold;")
     time.sleep(0.5)  # Pause for effect
 
-    yield None, stream_log("Processing stems...", "color: #9C27B0;")
+    yield None, None, stream_log("Processing stems...", "color: #9C27B0;")
     time.sleep(0.5)  # Simulate processing time
 
     # Change the stem search directory using full path
     stem_search_dir = os.path.join(base_output_dir, model_name, os.path.splitext(os.path.basename(audio_file))[0])
-    yield None, stream_log(f"Searching for stems in: {stem_search_dir}")
+    yield None, None, stream_log(f"Searching for stems in: {stem_search_dir}")
 
     stems: Dict[str, str] = {}
     for stem in ["vocals", "drums", "bass", "other"]:
         stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
-        yield None, stream_log(f"Checking for {stem} stem at: {stem_path}")
+        yield None, None, stream_log(f"Checking for {stem} stem at: {stem_path}")
         if os.path.exists(stem_path):
             stems[stem] = stem_path
-            yield None, stream_log(f"Found {stem} stem")
+            yield None, None, stream_log(f"Found {stem} stem")
         else:
-            yield None, stream_log(f"Warning: {stem} stem not found")
+            yield None, None, stream_log(f"Warning: {stem} stem not found")
 
     if not stems:
-        yield None, stream_log("Error: No stems found. Checking alternative directory...")
+        yield None, None, stream_log("Error: No stems found. Checking alternative directory...")
         stem_search_dir = os.path.join(base_output_dir, model_name)
         for stem in ["vocals", "drums", "bass", "other"]:
             stem_path = os.path.join(stem_search_dir, f"{stem}.wav")
-            yield None, stream_log(f"Checking for {stem} stem at: {stem_path}")
+            yield None, None, stream_log(f"Checking for {stem} stem at: {stem_path}")
             if os.path.exists(stem_path):
                 stems[stem] = stem_path
-                yield None, stream_log(f"Found {stem} stem")
+                yield None, None, stream_log(f"Found {stem} stem")
             else:
-                yield None, stream_log(f"Warning: {stem} stem not found")
+                yield None, None, stream_log(f"Warning: {stem} stem not found")
 
-    yield None, stream_log(f"All found stems: {list(stems.keys())}")
+    yield None, None, stream_log(f"All found stems: {list(stems.keys())}")
 
     selected_stems: List[str] = []
     for stem, selected in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]):
         if selected:
-            yield None, stream_log(f"{stem} is selected by user")
+            yield None, None, stream_log(f"{stem} is selected by user")
             if stem in stems:
                 selected_stems.append(stems[stem])
-                yield None, stream_log(f"Selected {stem} stem for mixing")
+                yield None, None, stream_log(f"Selected {stem} stem for mixing")
             else:
-                yield None, stream_log(f"Warning: {stem} was selected but not found")
+                yield None, None, stream_log(f"Warning: {stem} was selected but not found")
 
-    yield None, stream_log(f"Final selected stems: {selected_stems}")
+    yield None, None, stream_log(f"Final selected stems: {selected_stems}")
 
     if not selected_stems:
-        yield None, stream_log("Error: No stems selected for mixing", "color: #F44336;")
+        yield None, None, stream_log("Error: No stems selected for mixing", "color: #F44336;")
         raise gr.Error("Please select at least one stem to mix and ensure it was successfully separated.")
 
     output_file: str = os.path.join(output_dir, "mixed.wav")
-    yield None, stream_log("Mixing selected stems...", "color: #FF5722;")
+    yield None, None, stream_log("Mixing selected stems...", "color: #FF5722;")
     time.sleep(0.5)  # Simulate mixing time
     
-    if len(selected_stems) == 1:
-        mixed_audio = AudioSegment.from_wav(selected_stems[0])
-        mixed_audio.export(output_file, format="wav")
-    else:
-        mixed_audio: AudioSegment = AudioSegment.empty()
-        for stem_path in selected_stems:
-            mixed_audio += AudioSegment.from_wav(stem_path)
-        mixed_audio.export(output_file, format="wav")
+    mixed_audio: AudioSegment = AudioSegment.empty()
+    for stem_path in selected_stems:
+        mixed_audio += AudioSegment.from_wav(stem_path)
+    mixed_audio.export(output_file, format="wav")
 
     if mp3:
-        yield None, stream_log(f"Converting to MP3...", "color: #795548;")
+        yield None, None, stream_log(f"Converting to MP3...", "color: #795548;")
         time.sleep(0.5)  # Simulate conversion time
         mp3_output_file: str = os.path.splitext(output_file)[0] + ".mp3"
         mixed_audio.export(mp3_output_file, format="mp3", bitrate=str(mp3_bitrate) + "k")
         output_file = mp3_output_file
 
-    yield None, stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;")
-    yield output_file, gr.HTML(
+    yield None, None, stream_log("Process completed successfully!", "color: #4CAF50; font-weight: bold;")
+    yield output_file, list(stems.values()), gr.HTML(
         Panel.fit(
             Text("Separation and mixing completed successfully!", style="bold green"),
             title="Demucs Result",
@@ -178,13 +174,14 @@ with gr.Blocks() as iface:
             submit_btn = gr.Button("Process", variant="primary")
 
         with gr.Column(scale=1):
-            output_audio = gr.Audio(type="filepath", label="Processed Audio")
+            output_audio = gr.Audio(type="filepath", label="Processed Audio (Mixed)")
+            stems_output = gr.File(label="Individual Stems", file_count="multiple")
             separation_log = gr.HTML()
 
     submit_btn.click(
         fn=inference,
         inputs=[audio_input, model_dropdown, vocals_checkbox, drums_checkbox, bass_checkbox, other_checkbox, mp3_checkbox, mp3_bitrate],
-        outputs=[output_audio, separation_log]
+        outputs=[output_audio, stems_output, separation_log]
     )
     
     mp3_checkbox.change(