feat: add separate examples, logs and improvements

app/helpers.py

@@ -1,5 +1,4 @@
 import json
-import logging
 import os
 import random
 from base64 import b64encode
@@ -8,7 +7,6 @@ from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
-import requests
 import streamlit as st
 from PIL import Image
 from pydub import AudioSegment
@@ -20,7 +18,7 @@ extensions = ["mp3", "wav", "ogg", "flac"]  # we will look for all those file ty
 
 
 def check_file_availability(url):
-    exit_status = os.system(f"wget --spider {url}")
+    exit_status = os.system(f"wget -o --spider {url}")
     return exit_status == 0
 
 
@@ -33,18 +31,6 @@ def url_is_valid(url):
         st.error("Extension not supported.")
         return False
     try:
-        r = requests.get(url)
-        r.raise_for_status()
-        return True
-    except requests.exceptions.HTTPError as err:
-        msg = (
-            "requests get failed with status code "
-            + str(err.response.status_code)
-            + " for url "
-            + url
-            + ". Try wget spider."
-        )
-        logging.error(msg)
         return check_file_availability(url)
     except Exception:
         st.error("URL is not valid.")
@@ -79,12 +65,19 @@ def plot_audio(_audio_segment: AudioSegment, *args, **kwargs) -> Image.Image:
 
 
 @st.cache_data(show_spinner=False)
-def load_list_of_songs():
-    return json.load(open("sample_songs.json"))
+def load_list_of_songs(path="sample_songs.json"):
+    if os.environ.get("PREPARE_SAMPLES"):
+        return json.load(open(path))
+    else:
+        st.error(
+            "No examples available. You need to set the environment variable `PREPARE_SAMPLES=true`"
+        )
 
 
 def get_random_song():
     sample_songs = load_list_of_songs()
+    if sample_songs is None:
+        return None, None
     name, url = random.choice(list(sample_songs.items()))
     return name, url
 

app/pages/Separate.py

@@ -1,21 +1,22 @@
 import os
 from pathlib import Path
+from typing import List
+from loguru import logger as log
 
 import streamlit as st
-from streamlit_option_menu import option_menu
-
-from service.demucs_runner import separator
+from footer import footer
+from header import header
 from helpers import (
     load_audio_segment,
+    load_list_of_songs,
     plot_audio,
     st_local_audio,
     url_is_valid,
 )
+from service.demucs_runner import separator
+from service.vocal_remover.runner import load_model, separate
+from streamlit_option_menu import option_menu
 
-from service.vocal_remover.runner import separate, load_model
-
-from footer import footer
-from header import header
 
 label_sources = {
     "no_vocals.mp3": "🎶 Instrumental",
@@ -27,28 +28,104 @@ label_sources = {
     "other.mp3": "🎶 Other",
 }
 
-extensions = ["mp3", "wav", "ogg", "flac"]
+separation_mode_to_model = {
+    "Vocals & Instrumental (Faster)": ("vocal_remover", ["vocals.mp3", "no_vocals.mp3"]),
+    "Vocals & Instrumental (High Quality, Slower)": ("htdemucs", ["vocals.mp3", "no_vocals.mp3"]),
+    "Vocals, Drums, Bass & Other (Slower)": (
+        "htdemucs",
+        ["vocals.mp3", "drums.mp3", "bass.mp3", "other.mp3"],
+    ),
+    "Vocal, Drums, Bass, Guitar, Piano & Other (Slowest)": (
+        "htdemucs_6s",
+        ["vocals.mp3", "drums.mp3", "bass.mp3", "guitar.mp3", "piano.mp3", "other.mp3"],
+    ),
+}
 
+extensions = ["mp3", "wav", "ogg", "flac"]
 
 out_path = Path("/tmp")
 in_path = Path("/tmp")
 
 
+@st.cache_data(show_spinner=False)
+def get_sources(path, file_sources):
+    sources = {}
+    for file in file_sources:
+        fullpath = path / file
+        if fullpath.exists():
+            sources[file] = fullpath
+    return sources
+
+
 def reset_execution():
     st.session_state.executed = False
 
 
+def show_results(model_name: str, dir_name_output: str, file_sources: List):
+    sources = get_sources(out_path / Path(model_name) / dir_name_output, file_sources)
+    tab_sources = st.tabs([f"**{label_sources.get(k)}**" for k in sources.keys()])
+    for i, (file, pathname) in enumerate(sources.items()):
+        with tab_sources[i]:
+            cols = st.columns(2)
+            with cols[0]:
+                auseg = load_audio_segment(pathname, "mp3")
+                st.image(
+                    plot_audio(
+                        auseg,
+                        title="",
+                        file=file,
+                        model_name=model_name,
+                        dir_name_output=dir_name_output,
+                    ),
+                    use_column_width="always",
+                )
+            with cols[1]:
+                st_local_audio(pathname, key=f"output_{file}_{dir_name_output}")
+    log.info(f"Displaying results for {dir_name_output}")
+
+
 def body():
     filename = None
+    name_song = None
+    st.markdown(
+        """
+                <style>
+                    div[data-baseweb="tab-list"] {
+        align-items: center !important;
+        justify-content: center !important;
+        }
+        </style>""",
+        unsafe_allow_html=True,
+    )
+
     cols = st.columns([1, 3, 2, 1])
     with cols[1]:
-        with st.columns([1, 5, 1])[1]:
+        with st.columns([1, 8, 1])[1]:
             option = option_menu(
                 menu_title=None,
-                options=["Upload File", "From URL"],
-                icons=["cloud-upload-fill", "link-45deg"],
+                options=["Upload File", "From URL", "Examples"],
+                icons=["cloud-upload-fill", "link-45deg", "music-note-list"],
                 orientation="horizontal",
-                styles={"container": {"width": "100%", "margin": "0px", "padding": "0px"}},
+                styles={
+                    "container": {
+                        "width": "100%",
+                        "height": "3.5rem",
+                        "margin": "0px",
+                        "padding": "0px",
+                    },
+                    "icon": {"font-size": "1rem"},
+                    "nav-link": {
+                        "display": "flex",
+                        "height": "3rem",
+                        "justify-content": "center",
+                        "align-items": "center",
+                        "text-align": "center",
+                        "flex-direction": "column",
+                        "font-size": "1rem",
+                        "padding-left": "0px",
+                        "padding-right": "0px",
+                    },
+                },
                 key="option_separate",
             )
         if option == "Upload File":
@@ -64,18 +141,32 @@ def body():
                 filename = uploaded_file.name
                 st_local_audio(in_path / filename, key="input_upload_file")
 
-        elif option == "From URL":  # TODO: show examples
+        elif option == "From URL":
             url = st.text_input(
                 "Paste the URL of the audio file",
                 key="url_input",
                 help="Supported formats: mp3, wav, ogg, flac.",
             )
-            if url != "":
-                if url_is_valid(url):
-                    with st.spinner("Downloading audio..."):
-                        filename = url.split("/")[-1]
-                        os.system(f"wget -O {in_path / filename} {url}")
+            if url != "" and url_is_valid(url):
+                with st.spinner("Downloading audio..."):
+                    filename = url.split("/")[-1]
+                    os.system(f"wget -q -O {in_path / filename} {url}")
                 st_local_audio(in_path / filename, key="input_from_url")
+        elif option == "Examples":
+            samples_song = load_list_of_songs(path="separate_songs.json")
+            if samples_song is not None:
+                name_song = st.selectbox(
+                    label="Select a song",
+                    options=list(samples_song.keys()),
+                    format_func=lambda x: x.replace("_", " "),
+                    index=1,
+                    key="select_example",
+                )
+                if (Path("/tmp") / name_song).exists():
+                    st_local_audio(Path("/tmp") / name_song, key=f"input_from_sample_{name_song}")
+                else:
+                    name_song = None
+
     with cols[2]:
         separation_mode = st.selectbox(
             "Choose the separation mode",
@@ -92,6 +183,7 @@ def body():
             max_duration = 30
         else:
             max_duration = 15
+        model_name, file_sources = separation_mode_to_model[separation_mode]
 
     if filename is not None:
         song = load_audio_segment(in_path / filename, filename.split(".")[-1])
@@ -124,10 +216,10 @@ def body():
                     st.session_state.executed = False
 
                 if not st.session_state.executed:
+                    log.info(f"{option} - Separating {filename} with {separation_mode}...")
                     song.export(in_path / filename, format=filename.split(".")[-1])
                     with st.spinner("Separating source audio, it will take a while..."):
-                        if separation_mode == "Vocals & Instrumental (Faster)":
-                            model_name = "vocal_remover"
+                        if model_name == "vocal_remover":
                             model, device = load_model(pretrained_model="baseline.pth")
                             separate(
                                 input=in_path / filename,
@@ -137,13 +229,7 @@ def body():
                             )
                         else:
                             stem = None
-                            model_name = "htdemucs"
-                            if (
-                                separation_mode
-                                == "Vocal, Drums, Bass, Guitar, Piano & Other (Slowest)"
-                            ):
-                                model_name = "htdemucs_6s"
-                            elif separation_mode == "Vocals & Instrumental (High Quality, Slower)":
+                            if separation_mode == "Vocals & Instrumental (High Quality, Slower)":
                                 stem = "vocals"
 
                             separator(
@@ -162,39 +248,12 @@ def body():
                                 start_time=start_time,
                                 end_time=end_time,
                             )
-                last_dir = ".".join(filename.split(".")[:-1])
+                dir_name_output = ".".join(filename.split(".")[:-1])
                 filename = None
                 st.session_state.executed = True
-
-                def get_sources(path):
-                    sources = {}
-                    for file in [
-                        "no_vocals.mp3",
-                        "vocals.mp3",
-                        "drums.mp3",
-                        "bass.mp3",
-                        "guitar.mp3",
-                        "piano.mp3",
-                        "other.mp3",
-                    ]:
-                        fullpath = path / file
-                        if fullpath.exists():
-                            sources[file] = fullpath
-                    return sources
-
-                sources = get_sources(out_path / Path(model_name) / last_dir)
-                tab_sources = st.tabs([f"**{label_sources.get(k)}**" for k in sources.keys()])
-                for i, (file, pathname) in enumerate(sources.items()):
-                    with tab_sources[i]:
-                        cols = st.columns(2)
-                        with cols[0]:
-                            auseg = load_audio_segment(pathname, "mp3")
-                            st.image(
-                                plot_audio(auseg, title="", file=file),
-                                use_column_width="always",
-                            )
-                        with cols[1]:
-                            st_local_audio(pathname, key=f"output_{file}")
+                show_results(model_name, dir_name_output, file_sources)
+    elif name_song is not None and option == "Examples":
+        show_results(model_name, name_song, file_sources)
 
 
 if __name__ == "__main__":

app/style.py

@@ -124,6 +124,12 @@ CSS = (
         gap: 0rem;
     }
 
+    /* center the audio player in Separate page */
+    .css-keje6w.e1tzin5v1 {
+        display: flex;
+        justify-content: center;
+        align-items: center;
+    }
 
 </style>
 

requirements.in

@@ -14,3 +14,4 @@ resampy==0.4.2
 stqdm==0.0.5
 streamlit_option_menu==0.3.6
 htbuilder==0.6.1
+loguru==0.7.0

requirements.txt

@@ -38,7 +38,7 @@ contourpy==1.1.0
     # via matplotlib
 cycler==0.11.0
     # via matplotlib
-cython==0.29.35
+cython==0.29.36
     # via diffq
 decorator==5.1.1
     # via
@@ -91,14 +91,16 @@ kaleido==0.2.1
     # via -r requirements.in
 kiwisolver==1.4.4
     # via matplotlib
-lameenc==1.5.0
+lameenc==1.5.1
     # via demucs
-lazy-loader==0.2
+lazy-loader==0.3
     # via librosa
 librosa==0.10.0.post2
     # via -r requirements.in
 llvmlite==0.40.1
     # via numba
+loguru==0.7.0
+    # via -r requirements.in
 markdown-it-py==3.0.0
     # via rich
 markupsafe==2.1.3
@@ -152,7 +154,7 @@ pandas==1.5.3
     #   -r requirements.in
     #   altair
     #   streamlit
-pillow==9.5.0
+pillow==10.0.0
     # via
     #   matplotlib
     #   streamlit
@@ -271,7 +273,7 @@ tqdm==4.65.0
     #   stqdm
 treetable==0.2.5
     # via dora-search
-typing-extensions==4.7.0
+typing-extensions==4.7.1
     # via
     #   librosa
     #   rich

scripts/inference.py

@@ -1,7 +1,9 @@
 import argparse
+from pathlib import Path
 
 import warnings
 from app.service.vocal_remover.runner import load_model, separate
+from app.service.demucs_runner import separator
 
 warnings.simplefilter("ignore", UserWarning)
 warnings.simplefilter("ignore", FutureWarning)
@@ -14,16 +16,35 @@ def main():
     p.add_argument("--pretrained_model", "-P", type=str, default="baseline.pth")
     p.add_argument("--input", "-i", required=True)
     p.add_argument("--output_dir", "-o", type=str, default="")
+    p.add_argument("--only_no_vocals", "-n", action="store_true")
     args = p.parse_args()
 
+    input_file = args.input
+
     model, device = load_model(pretrained_model=args.pretrained_model)
     separate(
-        input=args.input,
+        input=input_file,
         model=model,
         device=device,
         output_dir=args.output_dir,
-        only_no_vocals=True,
+        only_no_vocals=args.only_no_vocals,
     )
+    if not args.only_no_vocals:
+        for stem, model_name in [("vocals", "htdemucs"), (None, "htdemucs"), (None, "htdemucs_6s")]:
+            separator(
+                tracks=[Path(input_file)],
+                out=Path(args.output_dir),
+                model=model_name,
+                shifts=1,
+                overlap=0.5,
+                stem=stem,
+                int24=False,
+                float32=False,
+                clip_mode="rescale",
+                mp3=True,
+                mp3_bitrate=320,
+                verbose=False,
+            )
 
 
 if __name__ == "__main__":

scripts/prepare_samples.sh

@@ -22,3 +22,21 @@ for name in $(echo "${json}" | jq -r 'keys[]'); do
     python inference.py --input /tmp/${name} --output /tmp
     echo "Done separating ${name}"
 done
+
+
+# Read JSON file into a variable
+json_separate=$(cat separate_songs.json)
+
+# Iterate through keys and values
+for name in $(echo "${json_separate}" | jq -r 'keys[]'); do
+    url=$(echo "${json_separate}" | jq -r --arg name "${name}" '.[$name]')
+    echo "Separating ${name} from ${url}"
+
+    # Download with pytube
+    yt-dlp ${url} -o "/tmp/${name}" --format "bestaudio/best"  --download-sections "*45-110"
+    mkdir -p "/tmp/vocal_remover"
+
+    # Run inference
+    python inference.py --input /tmp/${name} --output /tmp --only_no_vocals false
+    echo "Done separating ${name}"
+done

scripts/separate_songs.json

@@ -0,0 +1,8 @@
+{
+    "ABBA_-_Dancing_Queen": "https://www.youtube.com/watch?v=3qiMJt-JBb4",
+    "Queen_–_Bohemian_Rhapsody": "https://www.youtube.com/watch?v=yk3prd8GER4",
+    "Backstreet_Boys_-_I_Want_It_That_Way": "https://www.youtube.com/watch?v=qjlVAsvQLM8",
+    "The_Beatles_-_Let_It_Be": "https://www.youtube.com/watch?v=FIV73iG_e5I",
+    "Coldplay_-_Viva_La_Vida": "https://www.youtube.com/watch?v=a1EYnngNHIA",
+    "The_Cranberries_-_Zombie": "https://www.youtube.com/watch?v=8sM-rm4lFZg"
+}