---
license: mit
tags:
  - music-source-separation
  - diffusion
  - consistency-models
  - audio
---

# DiCoSe: Improving Music Source Separation with Diffusion and Consistency Refinement

Pre-trained checkpoints for **"Improving Music Source Separation with Diffusion and Consistency Refinement"**.

- Code: [github.com/Russell-Izadi-Bose/DiCoSe](https://github.com/Russell-Izadi-Bose/DiCoSe)
- Paper: [arXiv:2412.06965](https://arxiv.org/abs/2412.06965)
- Demo: [consistency-separation.github.io](https://consistency-separation.github.io/)

This repo hosts checkpoints for two experimental tracks described in the paper:

1. A custom **U-Net** separator trained on **Slakh2100**.
2. A **BS-RoFormer** separator (backbone from [Music-Source-Separation-Training](https://github.com/ZFTurbo/Music-Source-Separation-Training)) trained on **MUSDB18-HQ**.

For each track, three checkpoints are provided, corresponding to the three stages of the method: a Deterministic separator, a Diffusion refinement model trained on top of it, and a Consistency-Distilled (CD) model distilled from the diffusion model for fast (1-2 step) inference.

## Files

| File | Track | Stage | SDR (dB, avg across stems) |
|---|---|---|---|
| `Deterministic_model_unet/model.ckpt` | U-Net / Slakh2100 | Deterministic | 10.89 |
| `diffusion_model_unet/model.ckpt` | U-Net / Slakh2100 | Diffusion | 11.34 |
| `CD_unet/model.ckpt` | U-Net / Slakh2100 | Consistency-Distilled | 11.42 (T=1) → 11.95 (T=4) |
| `Deterministic_model_MSST_bs_roformer/model.ckpt` | BS-RoFormer / MUSDB18 | Deterministic | 9.84 |
| `diffusion_model_MSST_bs_roformer/model.ckpt` | BS-RoFormer / MUSDB18 | Diffusion | 10.34 |
| `CD_MSST_bs_roformer/model.ckpt` | BS-RoFormer / MUSDB18 | Consistency-Distilled | 10.41 (T=1) → 10.40 (T=2) |

SDR is the median-over-1s-chunks SDR (via `museval`), averaged across stems on the respective test set, as reported in the paper. The Consistency-Distilled (CD) checkpoints are a single model evaluated at different numbers of inference steps (T); more steps generally improve quality further.

## Usage

See the [GitHub repo](https://github.com/Russell-Izadi-Bose/DiCoSe) for the download script, environment setup, and eval configs that load these checkpoints. Training/eval code for the BS-RoFormer track is coming soon; checkpoints are published now for reference.

## Citation

```bibtex
@misc{karchkhadze2024improvingsourceextractiondiffusion,
  title={Improving Music Source Separation with Diffusion and Consistency Refinement},
  author={Tornike Karchkhadze and Mohammad Rasool Izadi and Shuo Zhang and Shlomo Dubnov},
  year={2024},
  eprint={2412.06965},
  archivePrefix={arXiv},
  primaryClass={cs.SD},
  url={https://arxiv.org/abs/2412.06965},
}
```

## License

MIT