From 30480f2335233d6c259840ebec7a2563175ca85b Mon Sep 17 00:00:00 2001
From: Hayk Martiros <hayk.mart@gmail.com>
Date: Sun, 8 Jan 2023 20:12:42 +0000
Subject: [PATCH] Recombine audio

Topic: recombine_audio
---
 riffusion/streamlit/pages/audio_to_audio.py |  2 +-
 riffusion/streamlit/pages/split_audio.py    | 27 ++++++++++++++-------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/riffusion/streamlit/pages/audio_to_audio.py b/riffusion/streamlit/pages/audio_to_audio.py
index f9dffcf..4aa1fbc 100644
--- 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
index 55da427..dc5c235 100644
--- 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,8 +38,12 @@ 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:
@@ -46,7 +51,8 @@ def render_split_audio() -> None:
         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__":