ZFTurbo/Music-Source-Separation-Training

Music Source Separation – Universal Training Code (MSST)

What it is – A Python library and set of scripts for training, evaluating, and running inference with a wide variety of deep‑learning models that separate a music track into its constituent sources (vocals, drums, bass, other). The repo bundles ready‑to‑use model implementations, configuration files, data handling utilities, and a small GUI, all aimed at making experimentation fast and reproducible.

Why it matters – Music source separation is a core task in audio signal processing and a popular benchmark for generative and discriminative audio models. By providing a single, extensible training framework that supports dozens of state‑of‑the‑art architectures (MDX23C, Demucs‑HT, RoFormer‑based models, Swin‑Upernet, SCNet, Conformer, etc.), the project lets researchers and hobbyists compare ideas without rewriting boilerplate code.

Key components

• Model zoo – Over a dozen architectures selectable via --model_type. Many are thin wrappers around existing open‑source projects (e.g., Lucidrains’ RoFormer, the Segmentation‑Models‑Pytorch library, the original Demucs‑HT code). The wrapper standardises the training loop and checkpoint handling.
• Config‑driven training – YAML files in configs/ describe dataset paths, hyper‑parameters, batch size, learning‑rate schedule, etc. Example configs for the MUSDB18 benchmark are provided.
• Training scripts – train.py for single‑GPU, train_ddp.py for multi‑GPU (PyTorch Distributed Data Parallel). Both accept LoRA fine‑tuning via the docs/LoRA.md guide.
• Inference utility – inference.py loads a checkpoint and processes a folder of wav files, writing separated stems to an output directory.
• Ensembling – ensemble.py can combine predictions from multiple models to improve SDR scores.
• GUI – A lightweight wxPython interface (gui-wx.py) for users who prefer a point‑and‑click workflow.
• Packaging – Installable via pip install msst, exposing a Python API, CLI, and documentation under docs/.

Typical workflow

1. Pick a model – e.g., mel_band_roformer.
2. Select a config – e.g., configs/config_mel_band_roformer_vocals.yaml (customise paths, batch size, etc.).
3. Launch training –

   python train.py \
       --model_type mel_band_roformer \
       --config_path configs/config_mel_band_roformer_vocals.yaml \
       --results_path results/ \
       --data_path datasets/train1 datasets/train2 \
       --valid_path datasets/musdb18hq/test \
       --num_workers 4 \
       --device_ids 0
   

4. Run inference –

   python inference.py \
       --model_type mel_band_roformer \
       --config_path configs/config_mel_band_roformer_vocals.yaml \
       --start_check_point results/last.ckpt \
       --input_folder input/wavs/ \
       --store_dir separation_results/
   

5. (Optional) Ensemble several checkpoints with ensemble.py to boost separation quality.

Hardware notes – Default batch sizes target a single NVIDIA A6000 (48 GB). Users with less VRAM should lower training.batch_size and/or increase training.gradient_accumulation_steps in the YAML config.

Pre‑trained models – The repository maintains a list (docs/pretrained_models.md) of publicly released checkpoints for many of the supported architectures, useful for quick inference or as a warm‑start for fine‑tuning.

Community & citation – The project is maintained by the team behind MVSep.com and builds on the KUIELab code for the SDX23 challenge. If you publish results using this code, cite the benchmark paper linked in the README.

---

Bottom line – MSST is a practical, well‑documented toolbox that abstracts away the repetitive parts of music source‑separation research, letting you focus on model ideas, data augmentation, or large‑scale experiments.