MSST

MSST (Music Source Separation Training) is a PyTorch library for building and using music source separation models. It provides a common runtime for inference, validation, training, dataset preparation, preflight checks, and experiment tracking across many separation architectures.

The project is designed for both ready-to-use models and model development. A user can load a published config/checkpoint pair and separate audio; a researcher can train the same architecture on another dataset, evaluate it with consistent metrics, resume from a checkpoint, or run it on multiple devices. The Python API and command-line interface call the same packaged runtime, so notebooks, applications, and automated jobs use the same model and dataset implementations.

Getting started · Python API · CLI · Compatibility · Dataset layouts · Pretrained models

What is included

This repository contains:

• model implementations including RoFormer, MDX, Demucs, SCNet, BandIt, Apollo, Conformer, and related architectures;
• reusable inference for in-memory audio, individual files, and folders;
• training and fine-tuning with configurable datasets, losses, optimizers, validation, checkpointing, LoRA, and multi-device launchers;
• checkpoint validation and source-separation metrics;
• eight training dataset layouts, metadata caching, and audio augmentations;
• typed Python interfaces, a command-line interface, experiment records, tests, and package documentation.

An MSST operation is assembled from four main pieces:

The package contains executable Python code. Model configs, checkpoints, datasets, and generated results are external artifacts, so they can be updated, stored, and versioned independently. Published config/checkpoint pairs are listed in the pretrained model catalog.

Installation

MSST supports CPython 3.10, 3.11, 3.12, and 3.13.

pip install msst

Each architecture can have its own optional dependencies. Install the extra whose name matches the selected model_type. Add train or validation when those workflows are needed:

pip install "msst[bs_roformer]"
pip install "msst[validation,bs_roformer]"
pip install "msst[train,bs_roformer]"

To install dependencies for every supported model family:

pip install "msst[all-models]"

Most users should install only the model extras they need. all-models is a large environment and intentionally excludes dependencies that are restricted to a narrow platform. Those models use their own extra, so one unavailable architecture does not prevent the generally available model families from being installed. For example, install BSMamba2 separately with msst[bs_mamba2] on a supported system. Workflow extras can still be combined as msst[train,validation,all-models].

On a CUDA system, install a PyTorch build compatible with the target GPU before installing MSST. MSST does not bundle CUDA binaries.

Inference

Use msst.inference() to process an input folder. The model is loaded once, all supported audio files are separated, and the paths written to disk are returned:

import msst

written_files = msst.inference(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    checkpoint_path="checkpoints/model.ckpt",
    input_folder="audio/input",
    output_folder="results/separated",
    device_ids=0,
)

for path in written_files:
    print(path)

For repeated calls, use msst.Separator so the weights stay loaded:

import msst

with msst.Separator(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    checkpoint_path="checkpoints/model.ckpt",
    device_ids=0,
) as separator:
    stems = separator.separate_file("audio/song.wav")
    separator.separate_folder("audio/album", "results/album")

print(separator.instruments)
print(stems.keys())

Separator.separate() also accepts an in-memory NumPy waveform. Input and output format details are documented in the compatibility guide.

Training

msst.train() writes checkpoints, model metadata, and training history to results_path:

import msst

msst.train(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    checkpoint_path="checkpoints/initial.ckpt",
    data_path="datasets/train",
    valid_path="datasets/valid",
    results_path="results/training",
    device_ids=(0,),
    launcher="standard",
    dataset_type=1,
    num_workers=8,
)

The checkpoint is optional when training from scratch. The launcher selects the execution runtime:

msst train --launcher standard --help
msst train --launcher ddp --help
msst train --launcher accelerate --help

Dataset structures are described in the dataset layout guide.

Validation

msst.valid() evaluates a checkpoint against reference stems and returns a typed result containing aggregate and per-track metrics:

import msst

result = msst.valid(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    checkpoint_path="checkpoints/model.ckpt",
    valid_path="datasets/test",
    metrics=("sdr", "k_sdr"),
    device_ids=0,
)

print(result.averages)
print(result.per_track["sdr"])

Pass multiple device indices to run parallel validation.

Preflight checks

Use msst.check() before training to detect configuration, dataset, checkpoint, dependency, device, and memory problems early:

import msst

report = msst.check(
    mode="safe",
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    checkpoint_path="checkpoints/initial.ckpt",
    data_path="datasets/train",
    valid_path="datasets/valid",
    results_path="results/training",
    device_ids=(0,),
    dataset_type=1,
)

report.raise_for_errors()
for item in report.items:
    print(item.status, item.name, item.message)

Two modes are available:

• safe validates the planned run without executing model forward or backward.
• full includes one real forward, loss, backward, and in-memory optimizer step without writing a checkpoint. On CUDA it also reports peak allocated and reserved memory for that step.

The full check uses one representative batch; memory use can still vary with later batches and random augmentations.

Dataset metadata

Training metadata can be generated independently and reused by the training runtime:

import msst

metadata = msst.build_metadata(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    data_path="datasets/train",
    results_path="results/training",
    dataset_type=1,
    num_workers=8,
)

print(metadata.path)
print(metadata.track_count)
print(metadata.fingerprint)

The function writes both the metadata cache and a JSON manifest describing its inputs and fingerprint.

Experiments

Experiment tracking is optional. Create a portable experiment directory with Python or the command line:

import msst

experiment = msst.init_experiment("experiments/vocals")

msst init experiments/vocals

The resulting layout keeps inputs, outputs, logs, scripts, and run records separate:

vocals/
├── experiment.yaml
├── configs/
├── checkpoints/
├── jobs/
├── logs/
├── results/
└── runs/

Pass the experiment to inference, validation, or training:

msst.train(
    model_type="bs_roformer",
    config_path="configs/model.yaml",
    data_path="datasets/train",
    valid_path="datasets/valid",
    results_path="results/training",
    device_ids=(0,),
    experiment=experiment,
    run_name="baseline",
)

Each call creates a record under runs/ with its parameters, software environment, config snapshot, input hashes, status, results, and artifacts. Tracking does not change the paths supplied to the operation.

Command-line interface

The same workflows are available through the msst command:

msst --help
msst inference --help
msst train --help
msst valid --help
msst check --help
msst metadata --help
msst init --help

Model extras

The base installation supports model types that require no additional Python dependencies. Other architectures have an extra with the same name:

Extras can be combined in one installation command:

pip install "msst[train,validation,bs_roformer]"

See the model compatibility table for architecture-specific notes.

Python API

The package exports the main interfaces directly from msst:

import msst

msst.inference
msst.Separator
msst.train
msst.valid
msst.check
msst.build_metadata
msst.init_experiment
msst.load_experiment

Public functions accept str and pathlib.Path values and are annotated for type checkers. Package-specific exceptions inherit from msst.MSSTError. Use help(msst.inference) or the complete Python API reference for every parameter, return type, and error condition.

Additional documentation:

• Getting started
• Command-line interface
• Experiments and run tracking
• Augmentations
• LoRA training
• Package architecture

License and citation

MSST is distributed under the MIT License.

If MSST contributes to published work, cite:

@misc{solovyev2023benchmarks,
  title={Benchmarks and leaderboards for sound demixing tasks},
  author={Roman Solovyev and Alexander Stempkovskiy and Tatiana Habruseva},
  year={2023},
  eprint={2305.07489},
  archivePrefix={arXiv},
  primaryClass={cs.SD}
}

MSST is brought to you by MVSep.