Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| # import demucs.api | |
| import os | |
| import spaces | |
| import subprocess | |
| from pydub import AudioSegment | |
| from typing import Tuple, Dict, List | |
| from demucs.apply import apply_model | |
| from demucs.separate import load_track | |
| from demucs.pretrained import get_model | |
| from demucs.audio import save_audio # Import save_audio from demucs.audio | |
| # check if cuda is available | |
| device: str = "cuda" if torch.cuda.is_available() else "cpu" | |
| # check if sox is installed and install it if necessary | |
| try: | |
| subprocess.run(["sox", "--version"], check=True, capture_output=True) | |
| except FileNotFoundError: | |
| print("sox is not installed. trying to install it now...") | |
| try: | |
| subprocess.run(["apt-get", "update"], check=True) | |
| subprocess.run(["apt-get", "install", "-y", "sox"], check=True) | |
| print("sox has been installed.") | |
| except subprocess.CalledProcessError as e: | |
| print(f"error installing sox: {e}") | |
| print("please install sox manually or try adding the following repository to your sources list:") | |
| print("deb http://deb.debian.org/debian stretch main contrib non-free") | |
| exit(1) | |
| # define the inference function | |
| def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, str]: | |
| """ | |
| performs inference using demucs and mixes the selected stems. | |
| args: | |
| audio_file: the audio file to separate. | |
| model_name: the name of the demucs model to use. | |
| vocals: whether to include vocals in the mix. | |
| drums: whether to include drums in the mix. | |
| bass: whether to include bass in the mix. | |
| other: whether to include other instruments in the mix. | |
| mp3: whether to save the output as mp3. | |
| mp3_bitrate: the bitrate of the output mp3 file. | |
| returns: | |
| a tuple containing the path to the mixed audio file and the separation log. | |
| """ | |
| # initialize demucs separator | |
| # separator: demucs.api.Separator = demucs.api.Separator(model=model_name) | |
| separator = get_model(name=model_name) | |
| # separate audio file and capture log | |
| import io | |
| log_stream = io.StringIO() | |
| # origin, separated = separator.separate_audio_file(audio_file, progress=True, log_stream=log_stream) | |
| wav = load_track(audio_file, separator.samplerate, channels=separator.audio_channels) | |
| ref = wav.mean(0) | |
| wav = (wav - ref.view(1, -1)).to(device) | |
| sources = apply_model(separator, wav, device=device, progress=True, log_stream=log_stream) | |
| sources = sources * ref.view(1, -1) + ref.view(1, -1) | |
| separation_log = log_stream.getvalue() | |
| # get the output file paths | |
| output_dir: str = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file))[0]) | |
| os.makedirs(output_dir, exist_ok=True) # create output directory if it doesn't exist | |
| stems: Dict[str, str] = {} | |
| for stem, source in zip(separator.sources, sources): | |
| stem_path: str = os.path.join(output_dir, f"{stem}.wav") | |
| # demucs.api.save_audio(source, stem_path, samplerate=separator.samplerate) | |
| save_audio(source, stem_path, separator.samplerate) # Use save_audio | |
| stems[stem] = stem_path | |
| # mix the selected stems | |
| selected_stems: List[str] = [stems[stem] for stem, include in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]) if include] | |
| if not selected_stems: | |
| raise gr.Error("please select at least one stem to mix.") | |
| output_file: str = os.path.join(output_dir, "mixed.wav") | |
| if len(selected_stems) == 1: | |
| # if only one stem is selected, just copy it | |
| os.rename(selected_stems[0], output_file) | |
| else: | |
| # otherwise, use pydub to mix the stems | |
| 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") | |
| # automatically convert to mp3 if requested | |
| if mp3: | |
| 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 # update output_file to the mp3 file | |
| return output_file, separation_log | |
| # Define the Gradio interface | |
| with gr.Blocks() as iface: | |
| gr.Markdown("# Demucs Music Source Separation and Mixing") | |
| gr.Markdown("Separate vocals, drums, bass, and other instruments from your music using Demucs and mix the selected stems.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| audio_input = gr.Audio(type="filepath", label="Upload Audio File") | |
| model_dropdown = gr.Dropdown( | |
| ["htdemucs", "htdemucs_ft", "htdemucs_6s", "hdemucs_mmi", "mdx", "mdx_extra", "mdx_q", "mdx_extra_q"], | |
| label="Model Name", | |
| value="htdemucs_ft" | |
| ) | |
| with gr.Row(): | |
| vocals_checkbox = gr.Checkbox(label="Vocals", value=True) | |
| drums_checkbox = gr.Checkbox(label="Drums", value=True) | |
| with gr.Row(): | |
| bass_checkbox = gr.Checkbox(label="Bass", value=True) | |
| other_checkbox = gr.Checkbox(label="Other", value=True) | |
| mp3_checkbox = gr.Checkbox(label="Save as MP3", value=False) | |
| mp3_bitrate = gr.Slider(128, 320, step=32, label="MP3 Bitrate", visible=False) | |
| submit_btn = gr.Button("Process", variant="primary") | |
| with gr.Column(scale=1): | |
| output_audio = gr.Audio(type="filepath", label="Processed Audio") | |
| separation_log = gr.Textbox(label="Separation Log", lines=10) | |
| 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] | |
| ) | |
| # Make MP3 bitrate slider visible only when "Save as MP3" is checked | |
| mp3_checkbox.change( | |
| fn=lambda mp3: gr.update(visible=mp3), | |
| inputs=mp3_checkbox, | |
| outputs=mp3_bitrate | |
| ) | |
| # Launch the Gradio interface | |
| iface.launch() | |