repository·master·Indexed 23 days ago
https://github.com/sigsep/open-unmix-pytorchA PyTorch-based music source separation toolkit (version 1.3.0) for splitting audio into vocals, drums, bass, and other stems. It utilizes a three-layer bidirectional LSTM architecture and provides pre-trained models including umxl, umxhq, umx, and umxse. The toolkit includes a command-line interface for inference, support for torch.hub loading, and integration with museval for SISEC standard evaluation.
Open-Unmix is a deep neural network reference implementation for music source separation built with PyTorch (1.8+). It allows users to separate pop music into four distinct stems:
vocalsdrumsbassotherThe models are pre-trained on the MUSDB18 dataset and are designed for researchers, audio engineers, and artists.
The repository is organized into several key modules that handle different stages of the machine learning pipeline:
data.py: Contains various torch datasets used for training.train.py: Contains the logic required to initiate and run training.model.py: Contains the Open-Unmix PyTorch modules (the neural network architectures).test.py: Contains code for performing inference (predicting/unmixing) from audio files.eval.py: Contains code for objective evaluation using museval on the MUSDB18 dataset.utils.py: Provides utility tools such as audio loading and metadata loading.Open-Unmix uses multiple models to perform separation into multiple sources, where each model is trained for a specific target. Each source model is based on a three-layer bidirectional LSTM that predicts the magnitude spectrogram of a target source by applying a mask to the input mixture.
Key architectural stages include:
models.Separator class is a high-level abstraction used during inference. It orchestrates multiple Open-unmix spectrogram models (one for each target source) and combines their outputs using a multichannel generalized Wiener filter. This filtering is a differentiable, parameter-free version of norbert. Finally, it applies inverse STFTs using torchaudio to reconstruct the audio signals.You can load a custom trained Separator by providing a path to a directory instead of a model name to the --model argument.
When loading a directory (e.g., --model mymodel --targets vocals), the following file structure is expected:
mymodel/separator.jsonmymodel/vocals.pthmymodel/vocals.jsonIf the separator contains multiple targets (e.g., vocals and drums), it will generate output files for each. If the --residual option is used, an additional source will be produced containing the estimate of everything not included in the targets.
If you are implementing models that operate directly in the time domain (like WaveNet or WaveRNN), you must modify the training and inference logic:
Training: Instead of comparing spectrograms, compare the time-domain output y_hat directly with the target time-domain signal y using a loss function like MSE.
Inference: Skip the spectral Wiener filter step. Directly save the time-domain signal produced by the model. You can also calculate the residual by subtracting the estimate from the original audio.
# Training logic change
y_hat = unmix(x)
loss = criterion(y_hat, y)
# Inference logic change
est = unmix(audio_torch).cpu().detach().numpy()
estimates[target] = est[0].T
estimates['residual'] = audio - est[0].TThe primary way to separate audio files is using the umx command. By default, it separates a mixture file into four stems: vocals, drums, bass, and other.
Supported audio formats include any files readable by torchaudio (backends like soundfile or sox).
umx input_file.wavTo evaluate the performance of your separation against SiSEC standards, install the museval package and run the evaluation module provided by openunmix.
museval:pip install musevalpython -m openunmix.evaluate --outdir /path/to/musdb/estimates --evaldir /path/to/museval/resultspip install museval
python -m openunmix.evaluate --outdir /path/to/musdb/estimates --evaldir /path/to/museval/resultsYou can separate audio files (.wav, .flac, .ogg; .mp3 is not supported) using the umx command.
umxhq (default): Trained on MUSDB18-HQ (uncompressed, full bandwidth up to 22050 Hz).umx: Trained on regular MUSDB18 (compressed, bandwidth limited to 16 kHz).umxse: Speech enhancement model trained on the Voicebank+DEMAND corpus.To use the default umxhq model:
umx input_file.wav --model umxhqTo load a custom model trained by a user, provide the path to the model's root directory instead of a model name:
umx --model /path/to/model/root/directory input_file.wavNote: A model directory typically contains individual models for each target. If the directory contains vocals and drums, two output files will be generated. If the --residual-model option is selected, an additional source containing the estimate of all non-target instruments will be produced.
If automatic downloading via TorchHub fails, you can download the model weights (umx or umxhq) from Zenodo or Zenodo and use them locally.
umx-weights)..pth and .json files in a flat hierarchy within that directory:umx-weights/vocals-*.pth
umx-weights/drums-*.pth
umx-weights/bass-*.pth
umx-weights/other-*.pth
umx-weights/vocals.json
umx-weights/drums.json
umx-weights/bass.json
umx-weights/other.json
umx-weights/separator.jsonumx command pointing to your local directory using the --model flag.umx --model umx-weights --input test.wavTo evaluate your results against SISEC standards, install the museval package and run the evaluation module:
pip install museval
python -m openunmix.evaluate --outdir /path/to/musdb/estimates --evaldir /path/to/museval/resultsTo implement a new model architecture, you can build upon the provided spectrogram model template. A custom model should typically inherit from nn.Module and implement the forward pass.
If your model works with spectrograms, you will likely need to include a transformation step (like STFT) within the forward method to convert the time-domain input into the frequency domain.
from model import Spectrogram, STFT
class Model(nn.Module):
def __init__(
self,
n_fft=4096,
n_hop=1024,
nb_channels=2,
input_is_spectrogram=False,
sample_rate=44100.0,
):
"""
Input: (batch, channel, sample)
or (frame, batch, channels, frequency)
Output: (frame, batch, channels, frequency)
"""
super(OpenUnmix, self).__init__()
def forward(self, mix):
# transform to spectrogram on the fly
X = self.transform(mix)
nb_frames, nb_samples, nb_channels, nb_bins = x.data.shape
# transform X to estimate
# ....
return X