Update app.py

app.py

@@ -1,121 +1,72 @@
-import streamlit as st
+import gradio as gr
 import os
 import tempfile
 import subprocess
 
-def separate_audio(audio_path):
+def separate_audio(audio_path, stem_count):
 
-    print(f"{audio_path=}")
     head, tail = os.path.split(audio_path)
-
     gradio_temp_path = head
     audio_filename = tail.split('.')[0]
-    print(f"{gradio_temp_path=}")
-    print(f"{audio_filename=}")
-    
-    command = f"spleeter separate -p spleeter:2stems {audio_path}"
-    command = command.split()
-    print(f"{command=}")
-
-    result = subprocess.run(command)
-    print(result)
-    
-    print("--------")
-    accompaniment_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"
-    vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
-    print(f"{accompaniment_path=}")
-    print(os.path.exists(accompaniment_path))
-    print(f"{vocals_path=}")
-    print(os.path.exists(vocals_path))
-
-    return vocals_path, accompaniment_path
 
-
-def separate_audio_by_stem(audio_path, stem_count):
-
-    print(f"{audio_path=}")
-    head, tail = os.path.split(audio_path)
-
-    gradio_temp_path = head
-    audio_filename = tail.split('.')[0]
-    print(f"{gradio_temp_path=}")
-    print(f"{audio_filename=}")
-    print(f"{stem_count=}")
-    
     command = f"spleeter separate -p spleeter:{stem_count}stems {audio_path}"
     command = command.split()
-    print(f"{command=}")
 
     result = subprocess.run(command)
-    print(result)
 
+    outputs = []
     if stem_count == 2:
         accompaniment_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/accompaniment.wav"
         vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
-
-        print(f"{accompaniment_path=} \t exists: {os.path.exists(accompaniment_path)}")
-        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
-        
-        return [
-            {'description': 'Accompaniment', 'path':accompaniment_path},
-            {'description': 'Vocals', 'path':vocals_path},
-        ]
-    
+        outputs.append({'description': 'Accompaniment', 'path': accompaniment_path})
+        outputs.append({'description': 'Vocals', 'path': vocals_path})
     elif stem_count == 4:
-        
         vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
         drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
         bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
         other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
-        
-        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
-        print(f"{drums_path=} \t exists: {os.path.exists(drums_path)}")
-        print(f"{bass_path=} \t exists: {os.path.exists(bass_path)}")
-        print(f"{other_path=} \t exists: {os.path.exists(other_path)}")
-
-        return [
-            {'description': 'Vocals', 'path':vocals_path},
-            {'description': 'Drums', 'path':drums_path},
-            {'description': 'Bass', 'path':bass_path},
-            {'description': 'Other', 'path':other_path},
-        ]
-        
+        outputs.extend([
+            {'description': 'Vocals', 'path': vocals_path},
+            {'description': 'Drums', 'path': drums_path},
+            {'description': 'Bass', 'path': bass_path},
+            {'description': 'Other', 'path': other_path},
+        ])
     elif stem_count == 5:
-
         piano_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/piano.wav"
         vocals_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/vocals.wav"
         drums_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/drums.wav"
         bass_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/bass.wav"
         other_path = f"{gradio_temp_path}/separated_audio/{audio_filename}/other.wav"
+        outputs.extend([
+            {'description': 'Vocals', 'path': vocals_path},
+            {'description': 'Piano', 'path': piano_path},
+            {'description': 'Drums', 'path': drums_path},
+            {'description': 'Bass', 'path': bass_path},
+            {'description': 'Other', 'path': other_path},
+        ])
+    return outputs
+
+def separate_audio_gradio(audio_file, stem_count):
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:        
+        temp_file.write(audio_file.read())
+        temp_file_path = temp_file.name
         
-        print(f"{piano_path=} \t exists: {os.path.exists(vocals_path)}")
-        print(f"{vocals_path=} \t exists: {os.path.exists(vocals_path)}")
-        print(f"{drums_path=} \t exists: {os.path.exists(drums_path)}")
-        print(f"{bass_path=} \t exists: {os.path.exists(bass_path)}")
-        print(f"{other_path=} \t exists: {os.path.exists(other_path)}")
-
-        return [
-            {'description': 'Vocals', 'path':vocals_path},
-            {'description': 'Piano', 'path':piano_path},
-            {'description': 'Drums', 'path':drums_path},
-            {'description': 'Bass', 'path':bass_path},
-            {'description': 'Other', 'path':other_path},
-        ]
+    separate_audios = separate_audio(temp_file_path, stem_count)
     
-uploaded_file = st.file_uploader("", type=["wav","mp3"])
-selected_stem_count = st.radio("Stem count", (2,4,5))
-
-if uploaded_file is not None:
-
-    if st.button("Submit"):
-        
-        with tempfile.NamedTemporaryFile(delete=False) as temp_file:        
-            temp_file.write(uploaded_file.read())
-            temp_file_path = temp_file.name
-            
-        separate_audios = separate_audio_by_stem(temp_file_path, selected_stem_count)
-        
-        st.write("Output Files:")
-        for audio in separate_audios:
-            st.write(audio['description'])
-            st.audio(audio['path'], format="audio/wav", start_time=0)
+    outputs = []
+    for audio in separate_audios:
+        outputs.append((audio['description'], gr.outputs.Audio(audio['path'])))
+    
+    return outputs
+
+input_audio = gr.inputs.File(label="Upload Audio File", type=["wav", "mp3"])
+input_stem_count = gr.inputs.Radio([2, 4, 5], label="Stem count")
+
+gr.Interface(
+    fn=separate_audio_gradio,
+    inputs=[input_audio, input_stem_count],
+    outputs="label",
+    title="Audio Separator",
+    description="Separate audio into vocals and accompaniment using Spleeter",
+    capture_session=True
+).launch()