Inverse Drum Machine (IDM)

This repository contains the official implementation of the paper: "The Inverse Drum Machine: Source Separation Through Joint Transcription and Analysis-by-Synthesis" (paper link).

IDM does Drum Source Separation (DSS) using analysis-by-synthesis combined with deep learning. Unlike traditional supervised methods that require isolated stem recordings for training, IDM is trained to reconstruct full mixes using transcription annotations as supervision.

• IDM is trained using only drum mixtures and their corresponding transcriptions, eliminating the need for isolated stems.

• The model jointly trains for Automatic Drum Transcription (ADT) and One-shot drum Sample Synthesis (OSS) in an end-to-end manner.

• IDM is comparable to state-of-the-art supervised methods while using approximately 100 times fewer parameters.

• The modular architecture allows for optional external information, such as corrected transcriptions, to be provided at inference time.

*Evaluation code will be added very soon, for now only the training and inference code is available, as well as weights for the model described in the paper.

System Architecture

IDM is composed of three main modules:

Feature Extraction: A ConvNeXt encoder extracts frame-level features from the input Log-Mel Spectrograms. Synthesis Conditioning: This module transforms the features into synthesis parameters: transcription onsets, velocities, individual track gains, and a conditioning vector for the one-shot synthesizer. Decoder: A one-shot synthesizer and a sequencer work together to reconstruct the individual drum tracks. The one-shot synth is a Temporal Convolutional Network (TCN) conditioned on instrument type and mixture embedding via FiLM. The sequencer is implemented using FFT multiplication.

Installation

This project uses Poetry for dependency management. The following steps assume you are working in a Python 3.10+ environment (you can create one as such conda create --name idm python=3.10).

1. Install Poetry: If you don't have Poetry installed, follow the instructions on the official website. One way is to run the following command inside your python environment:

   curl -sSL https://install.python-poetry.org | python3 -

2. Clone the repository

3. Install dependencies:

   poetry install # Install main dependencies
   poetry install --with dev # Install development dependencies (wandb, jupyter, etc.)
   poetry install --with nmf # Install nmf dependencies (e.g. TorchNMF)
   poetry install --with dev,nmf # Install all optional dependencies

Usage

The demo.ipynb notebook provides a comprehensive guide on how to use the pre-trained models for drum separation.

Here is a basic example of how to load a model and process an audio file:

import torch
from idm.inference import load_model, separate

# Load the pre-trained model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model, name = load_model('idm-44-train-kits', device, log_dir="pretrained")

drum_loop_44_1khz = "path/to/your/drum_loop.wav"
output_dir = "output/separated"
# Separate the audio using Wiener filtering
stems = separate(drum_loop_44_1khz, model, output_dir=output_dir, masking='wiener', alpha=1.0)

Note: if the drum loop is too far from the training distribution, the separation quality will be bad. This could be due to the transcription module not detecting the onsets correctly. We encourage users to try manual overrides of the transcription at inference time for better results. If I have some time, I will try to improve the transcription module.

To override the transcription, you can pass a dictionary with the onsets to the separate function:

# Example of manual transcription override
print(model.train_classes) # manual transcription should match these classes
onsets = {
    'KD': [0.0, 0.5, 1.0],
    'SD': [0.25, 0.75],
    'HH_CHH': [0.125, 0.375, 0.625, 0.875]
}
stems = separate(drum_loop_44_1khz, model, save_to_disk=False, masking='wiener', alpha=1.0, onsets=onsets)

Training

To train the model, you can use the provided run.py script with the desired configuration files. The main configuration is configs/train.yaml, which can be customized as needed.

Results

For detailed results, audio examples, and comparisons with baseline models, please visit the accompanying project website: https://bernardo-torres.github.io/projects/inverse-drum-machine/

Citation

If you find this work useful in your research, please consider citing the following paper:

@article{torres2025inversedrummachine,
title={The Inverse Drum Machine: Source Separation Through Joint Transcription and Analysis-by-Synthesis},
author={Torres, Bernardo and Peeters, Geoffroy and Richard, Ga{\"e}l},
journal={IEEE Transactions on Audio, Speech and Language Processing},
year={2025},
doi={10.1109/TASLPRO.2025.3629286}
}

TODO

• Evaluation code
• Model weights
• Add baseline instructions
• Add StemGMD preprocessing instructions
• Callbacks for logging audio samples during training
• Train new transcription module