shivbarca/Stem-Separator

1import gradio as gr2import subprocess3import os4import glob5import re6from pathlib import Path7 8def split_full_band(audio_filepath):9    """Step 1: Splits the full song into 4 stems and sets up the folders."""10    if not audio_filepath:11        return [], None, None, gr.update(visible=False)12        13    base_name = Path(audio_filepath).stem14    15    # 1. Create the folder structure16    root_dir = os.path.abspath(base_name)17    four_stems_dir = os.path.join(root_dir, f"{base_name} 4 stems")18    os.makedirs(four_stems_dir, exist_ok=True)19    20    # Clean up old flac files21    for f in glob.glob(os.path.join(four_stems_dir, "*.flac")):22        os.remove(f)23        24    # 2. Run htdemucs for the 4 stems25    command = [26        "audio-separator", 27        audio_filepath, 28        "--model_filename", "htdemucs_ft.yaml",29        "--output_dir", four_stems_dir,30        "--output_format", "FLAC"31    ]32    33    print(f"Starting Full Band Separation for '{base_name}'...")34    subprocess.run(command, check=True)35    36    # 3. Bulletproof Renaming Logic using Regex37    drums_filepath = ""38    for file_path in glob.glob(os.path.join(four_stems_dir, "*.flac")):39        filename = os.path.basename(file_path)40        41        # audio-separator always outputs files with the stem name in parentheses, e.g., "_(Vocals)_"42        match = re.search(r'_\((.*?)\)_', filename)43        if match:44            stem_type = match.group(1).lower()45            new_path = os.path.join(four_stems_dir, f"{base_name} {stem_type}.flac")46            os.rename(file_path, new_path)47            if stem_type == "drums":48                drums_filepath = new_path49                50    stems = glob.glob(os.path.join(four_stems_dir, "*.flac"))51    52    return stems, base_name, drums_filepath, gr.update(visible=True)53 54 55def split_drums_further(base_name, drums_filepath):56    """Step 2: Takes the isolated FLAC drums and splits them into WAV kit pieces."""57    if not drums_filepath or not os.path.exists(drums_filepath):58        return []59        60    four_stems_dir = os.path.join(os.path.abspath(base_name), f"{base_name} 4 stems")61    drum_stems_dir = os.path.join(four_stems_dir, f"{base_name} drum stems")62    os.makedirs(drum_stems_dir, exist_ok=True)63    64    for f in glob.glob(os.path.join(drum_stems_dir, "*.wav")):65        os.remove(f)66    67    # We pass the FLAC file in, but force WAV out for the final kit pieces68    command = [69        "audio-separator", 70        drums_filepath, 71        "--model_filename", "MDX23C-DrumSep-aufr33-jarredou.ckpt",72        "--output_dir", drum_stems_dir,73        "--output_format", "WAV"74    ]75    76    print(f"Starting Drum Stem Separation for '{base_name}'...")77    subprocess.run(command, check=True)78    79    # Bulletproof Renaming Logic for Drum Stems80    for file_path in glob.glob(os.path.join(drum_stems_dir, "*.wav")):81        filename = os.path.basename(file_path)82        83        match = re.search(r'_\((.*?)\)_', filename)84        if match:85            stem_type = match.group(1).lower()86            new_path = os.path.join(drum_stems_dir, f"{base_name} {stem_type}.wav")87            os.rename(file_path, new_path)88            89    drum_stems = glob.glob(os.path.join(drum_stems_dir, "*.wav"))90    return drum_stems91 92 93# --- Build the Dynamic User Interface ---94with gr.Blocks(theme=gr.themes.Soft()) as app:95    gr.Markdown("# 🎛️ The Full Band & Drum Separator")96    97    state_basename = gr.State()98    state_drumspath = gr.State()99    100    with gr.Row():101        audio_input = gr.File(102            file_types=[".wav", ".webm", ".flac", ".m4a"], 103            label="Upload Track (.wav, .webm, .flac, .m4a)"104        )105    106    full_band_btn = gr.Button("Step 1: Split Full Band (Vocals, Bass, Drums, Other)", variant="primary")107    output_4_stems = gr.File(label="4 Stems Output", file_count="multiple")108    109    with gr.Column(visible=False) as drum_split_ui:110        gr.Markdown("### 🥁 Drums Isolated! Break the kit down further:")111        drum_split_btn = gr.Button("Step 2: Split Drums into Kick, Snare, Hats, etc.", variant="secondary")112        output_drum_stems = gr.File(label="Drum Stems Output", file_count="multiple")113 114    full_band_btn.click(115        fn=split_full_band, 116        inputs=audio_input, 117        outputs=[output_4_stems, state_basename, state_drumspath, drum_split_ui]118    )119    120    drum_split_btn.click(121        fn=split_drums_further,122        inputs=[state_basename, state_drumspath],123        outputs=output_drum_stems124    )125 126if __name__ == "__main__":127    app.launch()