Model
GAN-based audio super-resolution. Takes any MP3 (unknown bitrate) and reconstructs a high-fidelity stereo FLAC.
Architecture: 1D U-Net generator (28.2M params) + multi-scale discriminator (25.1M params)
Training: Mixed-precision, gradient checkpointing | Peak VRAM: ~3 GB on RTX 4070
Prerequisites
- Python >= 3.10
- NVIDIA GPU with CUDA (tested on RTX 4070, 12 GB)
- FFmpeg (for data preparation)
- uv (Python package manager)
1# Install all dependencies from the project root
2uv sync
Quick Start
1. Prepare your data
Place your lossless FLAC files in root data/lossless/, then generate the lossy training pairs:
uv run python -m model.prepare_dataThis transcodes every FLAC to MP3 at 128, 256, and 320 kbps into root data/lossy/{128,256,320}/. Existing files are skipped on re-runs.
Custom paths:
uv run python -m model.prepare_data --input /path/to/flacs --output /path/to/lossy2. Train
uv run python -m model.train| Flag | Default | Description |
|---|
--epochs | 200 | Total training epochs |
--batch-size | 4 | Batch size (safe up to ~16 on 12 GB) |
--lr | 1e-4 | Learning rate for both G and D |
--resume | — | Path to a checkpoint dir to resume from |
--checkpoint-dir | model/checkpoints | Where to save checkpoints |
--data-lossless | data/lossless | Lossless FLAC directory |
--data-lossy | data/lossy | Lossy MP3 directory |
Training automatically:
- Mixes all three bitrates so the model generalizes to unknown compression
- Splits files 80/20 into train/val sets
- Saves the best checkpoint (lowest val spectral loss) to
checkpoints/best/
- Saves periodic checkpoints every 10 epochs
- Logs per-epoch metrics: G loss, D loss, spectral loss, feature matching loss, adversarial loss
Resume from a checkpoint:
uv run python -m model.train --resume model/checkpoints/epoch_503. Evaluate / Infer
Reconstruct a single MP3 to FLAC:
1uv run python -m model.evaluate \
2 --checkpoint model/checkpoints/best \
3 --input song.mp3 \
4 --output output/
Reconstruct an entire directory:
1uv run python -m model.evaluate \
2 --checkpoint model/checkpoints/best \
3 --input /path/to/mp3s/ \
4 --output /path/to/output/
Evaluate against reference FLACs (computes PESQ and SNR):
1uv run python -m model.evaluate \
2 --checkpoint model/checkpoints/best \
3 --input data/lossy/128/ \
4 --output output/ \
5 --reference data/lossless/
4. Export for HuggingFace
uv run python -m model.export --checkpoint model/checkpoints/bestProduces exported/model.safetensors and exported/config.json — only the generator weights (discriminator is not needed at inference).
Push to HuggingFace Hub:
huggingface-cli upload your-username/audio-super-resolution exported/
File Reference
| File | Purpose |
|---|
config.py | ModelConfig, TrainConfig, DataConfig dataclasses |
generator.py | 1D U-Net with skip connections and gradient checkpointing |
discriminator.py | Multi-scale discriminator (3 scales) |
losses.py | LSGAN + multi-scale spectral + feature matching losses |
dataset.py | Stereo audio pair dataset with random bitrate sampling |
train.py | Training loop with mixed precision and rich progress |
evaluate.py | Inference (MP3 to FLAC) and metrics (PESQ, SNR) |
export.py | Export generator weights as safetensors for HuggingFace |
prepare_data.py | FFmpeg batch transcoder (FLAC to MP3 at 128/256/320) |
utils.py | Audio I/O, normalization, checkpoint save/load |
Training Loss Components
| Loss | Weight | Purpose |
|---|
| Adversarial (LSGAN) | 1.0 | Perceptual realism |
| Multi-scale spectral | 100.0 | Frequency accuracy (STFT at 512, 1024, 2048) |
| Feature matching | 10.0 | Stability via discriminator feature alignment |
Data Layout
data/
├── lossless/ # Your original FLAC files
└── lossy/
├── 128/ # Generated: 128 kbps MP3
├── 256/ # Generated: 256 kbps MP3
└── 320/ # Generated: 320 kbps MP3
Files are paired by filename: data/lossless/track.flac matches data/lossy/128/track.mp3.
