yongyizang/MSR_UFormers

Music Restoration with UFormer

These checkpoints accompany the Music Source Restoration paper. (arxiv link coming soon)

A simple audio restoration toolkit based on UFormer, designed for music source restoration across different instrument groups. Supports long-file inference via overlap-add and multiple checkpoints for specialized restoration (e.g., guitars, vocals, synth).

Directory Structure

music-restoration/               # Root of the repository
├── model.py                     # UFormerConfig & UFormer definitions
├── inference.py                 # Inference entry-point (CLI, HF API, Gradio)
├── requirements.txt             # Dependencies: torch, soundfile, gradio, huggingface_hub
├── .gitattributes               # LFS tracking for checkpoints
├── README.md                    
└── checkpoints/                 # Pretrained weights
    ├── acoustic_guitar.pth
    ├── bass.pth
    ├── electric_guitar.pth
    ├── guitars.pth
    ├── keyboards.pth
    ├── orchestra.pth
    ├── rhythm_section.pth
    ├── synth.pth
    └── vocals.pth

Installation

pip install -r requirements.txt

Usage

Command-Line Interface

python inference.py \
  --input noisy.wav \
  --output restored.wav \
  --checkpoint guitars

Hugging Face Inference API

Send a JSON payload:

{
  "inputs": <raw WAV bytes>,
  "parameters": { "checkpoint": "vocals" }
}

Gradio Demo (Spaces)

python inference.py --serve

Open the provided local URL to upload audio and select a checkpoint.

Checkpoints

Available models for different instrument groups:

• acoustic_guitar
• bass
• electric_guitar
• guitars
• keyboards
• orchestra
• rhythm_section
• synth
• vocals

License

Apache 2.0 License. See LICENSE for details.