Duplicate from theadityamittal/music-separator-unet

Co-authored-by: Aditya Mittal <theadityamittal@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/*.pt filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,105 @@
+---
+tags:
+- audio
+- music-source-separation
+- u-net
+- pytorch
+license: mit
+datasets:
+- musdb18hq
+metrics:
+- SDR
+- SIR
+- SAR
+---
+
+# 🎸 Music-U-Net — 4-Stem Source Separator
+
+A PyTorch U-Net trained to split a full-band stereo **mixture** into  
+**drums · bass · other · vocals**.
+
+| Property              | Value |
+|-----------------------|-------|
+| Model type            | 2-D U-Net (6.2 M params) |
+| Input representation  | STFT magnitude (mono, 16 kHz) |
+| Output                | 4 magnitude masks (drums, bass, other, vocals) |
+| Training data         | 100 train + 50 test songs from **MUSDB-18 HQ** |
+| Checkpoint size       | ~24 MB (`state_dict`, FP32) |
+| License               | MIT |
+
+---
+
+## 🗂️ Contents  
+
+```
+checkpoints/unet\_best.pt   # model weights (state\_dict)
+config/default.yaml        # sample-rate, FFT size, etc.
+README.md                  # this card
+
+```
+
+---
+
+## 📝 Model Details  
+
+### Architecture  
+Classic symmetric U-Net over 2-D spectra:
+
+```
+Encoder:  \[C32]→\[C64]→\[C128]→\[C256]→\[C512]
+Decoder:  \[C256]←\[C128]←\[C64]←\[C32]
+```
+
+`ReLU` activations, batch-norm, skip-connections, 1×1 final conv to **4 channels**  
+(one per target stem) followed by soft masks --> multiplied by mixture magnitude.
+
+### Training  
+* **Loss**: L1( pred_mag·mix_phase , ref_mag·mix_phase )  
+* **Augment**: time/freq masking, Gaussian noise, ±3 dB gain  
+* **Optimizer**: Adam, LR 1e-4 → 1e-5 cosine decay, 50 epochs  
+* **Hardware**: single RTX 3090, 2 h total
+
+---
+
+## 📊 Evaluation (MUSDB-18 test, per-track average)
+
+| Metric | Mean | Std |
+|--------|------|-----|
+| **SDR** | **-0.14 dB** | 1.66 |
+| **SIR** | 3.93 dB | 1.86 |
+| **SAR** | 4.26 dB | 0.85 |
+
+*(baseline numbers; not state-of-the-art, but fast & lightweight)*
+
+---
+
+## 💻 Usage
+
+Try it live in the **Gradio Space** 👉 **[https://huggingface.co/spaces/theadityamittal/music-separator-space](https://huggingface.co/spaces/YOUR_USERNAME/music-separator-space)**
+
+---
+
+## ⚖ Limitations & Biases
+
+* Trained only on MUSDB-18 HQ → may fail on genres not represented (classical, EDM).
+* Uses mixture phase → audible bleeding & artifacts, negative SDR in some tracks.
+* No multi-channel or stem permutation handling.
+
+---
+
+## 📄 License
+
+Released under the MIT License.
+
+---
+
+## 🙏 Citation
+
+```bibtex
+@misc{music-unet-2025,
+  title   = {Music Source Separation with U-Net},
+  author  = {Your Name},
+  url     = {https://huggingface.co/YOUR_USERNAME/music-separator-unet},
+  year    = 2025
+}
+```

checkpoints/unet_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d751220f7215032954092f940c454bf992ffdb9a4186f7c94500e20c0248739
+size 31133505

config/default.yaml

@@ -0,0 +1,50 @@
+# config/default.yaml
+
+device: "mps"
+
+data:
+  raw_path: data/raw
+  splits: ["train", "test"]
+  processed_path: data/processed
+  sample_rate: 16000
+  n_fft: 1024
+  hop_length: 512
+  n_mels: 80
+  segment_length: 256
+
+  # for DataLoader
+  batch_size: 16
+  num_workers: 4
+
+  # list of all sources (including mixture)
+  sources: ["mixture", "drums", "bass", "other", "vocals"]
+
+model:
+  checkpoint_dir: models/checkpoints
+
+  # for UNet
+  chans: 32
+  num_pool_layers: 4
+
+training:
+  # for training loop
+  epochs: 50
+  lr: 1e-4
+  max_steps: null
+  log_interval: 50    # how many batches between progress logs
+
+augment:
+  # defaults for your SpectrogramTransforms
+  time_mask_param: 30
+  freq_mask_param: 15
+  time_warp_param: 40
+  stripe_time_width: 1
+  stripe_freq_width: 1
+  stripe_time_count: 2
+  stripe_freq_count: 2
+  noise_std: 0.01
+
+experiment:
+  # MLflow experiment metadata
+  name: default_experiment
+  run_name: run1