# Model setup
DAC.sample_rate: 48000
DAC.encoder_dim: 64
DAC.encoder_rates: [2, 3, 5, 8, 8]
DAC.latent_dim: 128
DAC.decoder_dim: 1536
DAC.decoder_rates: [8, 8, 5, 3, 2]

# Quantization
DAC.n_codebooks: 12
DAC.codebook_size: 2048
DAC.codebook_dim: 8
DAC.quantizer_dropout: 0.25

# Discriminator
Discriminator.sample_rate: 48000
Discriminator.rates: []
Discriminator.periods: [2, 3, 5, 7, 11]
Discriminator.fft_sizes: [2048, 1024, 512]
Discriminator.bands:
  - [0.0, 0.1]
  - [0.1, 0.25]
  - [0.25, 0.5]
  - [0.5, 0.75]
  - [0.75, 1.0]

# Optimization
AdamW.betas: [0.8, 0.99]
AdamW.lr: 0.0001
ExponentialLR.gamma: 0.999996

amp: false
val_batch_size: 100
device: cuda
num_iters: 400000
save_iters: [10000, 50000, 100000, 200000]
valid_freq: 1000
sample_freq: 10000
num_workers: 32
val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
seed: 0
lambdas:
  mel/loss: 15.0
  adv/feat_loss: 2.0
  adv/gen_loss: 1.0
  vq/commitment_loss: 0.25
  vq/codebook_loss: 1.0
  vq/orthogonal_loss: 10.0

wandb_log: true

VolumeNorm.db: [const, -16]

# Transforms
build_transform.preprocess:
  - Identity
build_transform.augment_prob: 0.0
build_transform.augment:
  - Identity
build_transform.postprocess:
  # - VolumeNorm
  - RescaleAudio
  - ShiftPhase

# Loss setup
MultiScaleSTFTLoss.window_lengths: [2048, 512]
MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
MelSpectrogramLoss.pow: 1.0
MelSpectrogramLoss.clamp_eps: 1.0e-5
MelSpectrogramLoss.mag_weight: 0.0

# Data
batch_size: 64
train/AudioDataset.duration: 1.0
train/AudioDataset.n_examples: 10000000

val/AudioDataset.duration: 1.0
val/build_transform.augment_prob: 1.0
val/AudioDataset.n_examples: 250

test/AudioDataset.duration: 1.0
test/build_transform.augment_prob: 1.0
test/AudioDataset.n_examples: 1000


AudioLoader.shuffle: true
AudioDataset.num_channels: 2
AudioDataset.without_replacement: true

train/build_dataset.folders:
  music_hq:
    - /app/suno/data/audio_2ch_48khz/music_tr.csv

val/build_dataset.folders:
  music_hq:
    - /app/suno/data/audio_2ch_48khz/music_val.csv

test/build_dataset.folders:
  music_hq:
    - /app/suno/data/audio_2ch_48khz/music_val.csv