Recombine audio

Add an option in the splitter to recombine stems at the end. Topic: recombine_audio
qinfrank · Jan 8, 2023 · a4784bb · a4784bb
1 parent 8e87c13
commit a4784bb
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 10 deletions.
diff --git a/riffusion/streamlit/pages/audio_to_audio.py b/riffusion/streamlit/pages/audio_to_audio.py
@@ -39,7 +39,7 @@ def render_audio_to_audio() -> None:
 
     audio_file = st.file_uploader(
         "Upload audio",
-        type=["mp3", "m4a", "ogg", "wav", "flac"],
+        type=["mp3", "m4a", "ogg", "wav", "flac", "webm"],
         label_visibility="collapsed",
     )
 

diff --git a/riffusion/streamlit/pages/split_audio.py b/riffusion/streamlit/pages/split_audio.py
@@ -1,5 +1,6 @@
 import io
 
+import pydub
 import streamlit as st
 
 from riffusion.audio_splitter import split_audio
@@ -37,16 +38,21 @@ def render_split_audio() -> None:
         label_visibility="collapsed",
     )
 
-    recombine = st.sidebar.checkbox(
-        "Recombine", value=False, help="Show recombined audio at the end for comparison"
+    stem_options = ["vocals", "drums", "bass", "guitar", "piano", "other"]
+    recombine = st.sidebar.multiselect(
+        "Recombine",
+        options=stem_options,
+        default=[],
+        help="Recombine these stems at the end",
     )
 
     if not audio_file:
         st.info("Upload audio to get started")
         return
 
     st.write("#### original")
-    st.audio(audio_file)
+    # TODO(hayk): This might be bogus, it can be other formats..
+    st.audio(audio_file, format="audio/mp3")
 
     if not st.button("Split", type="primary"):
         return
@@ -61,19 +67,22 @@ def render_split_audio() -> None:
         st.write(f"#### {name}")
         audio_bytes = io.BytesIO()
         stem.export(audio_bytes, format="mp3")
-        st.audio(audio_bytes)
+        st.audio(audio_bytes, format="audio/mp3")
 
     if recombine:
-        stems_list = list(stems.values())
-        recombined = stems_list[0]
-        for stem in stems_list[1:]:
-            recombined = recombined.overlay(stem)
+        recombined: pydub.AudioSegment = None
+        for name, stem in stems.items():
+            if name in recombine:
+                if recombined is None:
+                    recombined = stem
+                else:
+                    recombined = recombined.overlay(stem)
 
         # Display
         st.write("#### recombined")
         audio_bytes = io.BytesIO()
         recombined.export(audio_bytes, format="mp3")
-        st.audio(audio_bytes)
+        st.audio(audio_bytes, format="audio/mp3")
 
 
 if __name__ == "__main__":