<div align="center">

<img src="assets/boost.svg" width="160px">

# boost

<i> General purpose audio quality enhancement model </i>

</div>

## Setup

```
cd boost
conda create -n boost-env python=3.10
conda activate boost-env
conda install conda-forge::ffmpeg
pip install -e .
pip install -U 'jsonargparse[signatures]>=4.27.7
```

## Codec data generation

To generate training data we perform codec cycling offline, storing audio to disk to use for training. 
This is achieved with `scripts/encode.py`, which requires `suno_utils`.

```
git clone git@github.com:suno-ai/glockenspiel.git
cd glockenspiel
pip install -e ./suno_utils
pip install -e ./descript-audio-codec
```

Run the script passing a directory of audio files and DAC checkpoints
```
CUDA_VISIBLE_DEVICES=0 python scripts/encode.py \
/app/suno/christian/reference-audio-wav \
/app/suno/christian/data/codec_audio/reference-audio-wav \
--max_frames 524288 \
--block_size 524288 \
--codec_paths checkpoints/dac/dac_2c_25x12.pt
```

```
CUDA_VISIBLE_DEVICES=0 python scripts/encode.py \
/app/suno/data/audio_2ch_48khz_lg/train/genius_hq \
/app/suno/christian/data/codec_audio/genius_hq \
--max_frames 524288 \
--block_size 524288 \
--codec_paths checkpoints/dac/dac_2c_25x12.pt
```

## Inference

```
CUDA_VISIBLE_DEVICES=1 python scripts/infer.py \
"/app/suno/christian/data/codec_audio/reference-audio-wav/02 Josie_s Song.dac_2c_25x12.wav" \
--use_gpu \
--num_steps 100 \
--batch_size 64 \
--ckpt_path /app/suno/christian/boost-logs/boost/61n3rmcc/checkpoints/last.ckpt 
```

## Training

Single node training
```
CUDA_VISIBLE_DEVICES=2 python scripts/main.py fit -c cfg/trainer.yaml -c cfg/models/model.yaml -c cfg/data/upsample_genius.yaml
```

HDemucs
```
CUDA_VISIBLE_DEVICES=3 python scripts/main.py fit -c cfg/trainer.yaml -c cfg/models/hdemucs.yaml -c cfg/data/upsample_genius.yaml
```

Cluster training
```
sbatch ./sruns/submit.sh
```