fish-speech/fish_speech/configs/vq_diffusion.yaml

defaults:
  - base
  - _self_

project: vq_diffusion

# Lightning Trainer
trainer:
  accelerator: gpu
  devices: 4
  strategy: ddp_find_unused_parameters_true
  gradient_clip_val: 1.0
  gradient_clip_algorithm: 'norm'
  precision: 16-mixed
  max_steps: 300_000
  val_check_interval: 5000

sample_rate: 24000
hop_length: 256
num_mels: 100
n_fft: 1024
win_length: 1024

# Dataset Configuration
train_dataset:
  _target_: fish_speech.datasets.vqgan.VQGANDataset
  filelist: data/filelist.split.train
  sample_rate: ${sample_rate}
  hop_length: ${hop_length}
  slice_frames: 512

val_dataset:
  _target_: fish_speech.datasets.vqgan.VQGANDataset
  filelist: data/filelist.split.valid
  sample_rate: ${sample_rate}
  hop_length: ${hop_length}

data:
  _target_: fish_speech.datasets.vqgan.VQGANDataModule
  train_dataset: ${train_dataset}
  val_dataset: ${val_dataset}
  num_workers: 8
  batch_size: 32
  val_batch_size: 4

# Model Configuration
model:
  _target_: fish_speech.models.vq_diffusion.lit_module.VQDiffusion
  sample_rate: ${sample_rate}
  hop_length: ${hop_length}
  speaker_use_feats: true

  downsample:
    _target_: fish_speech.models.vq_diffusion.lit_module.ConvDownSample
    dims: [128, 512, 128]
    kernel_sizes: [3, 3]
    strides: [2, 2]

  text_encoder:
    _target_: fish_speech.models.vqgan.modules.encoders.TextEncoder
    in_channels: 128
    out_channels: 128
    hidden_channels: 192
    hidden_channels_ffn: 768
    n_heads: 2
    n_layers: 6
    kernel_size: 1
    dropout: 0.1
    use_vae: false
    gin_channels: 512
    speaker_cond_layer: 0

  vq_encoder:
    _target_: fish_speech.models.vqgan.modules.encoders.VQEncoder
    in_channels: 128
    vq_channels: 128
    codebook_size: 4096
    downsample: 1

  speaker_encoder:
    _target_: fish_speech.models.vqgan.modules.encoders.SpeakerEncoder
    in_channels: 128
    hidden_channels: 192
    out_channels: 128
    num_heads: 2
    num_layers: 4
    p_dropout: 0.1
  
  denoiser:
    _target_: fish_speech.models.vq_diffusion.wavenet.WaveNet
    d_encoder: 128
    mel_channels: 100
    residual_channels: 512
    residual_layers: 20

  vocoder:
    _target_: fish_speech.models.vq_diffusion.bigvgan.BigVGAN

  mel_transform:
    _target_: fish_speech.models.vqgan.spectrogram.LogMelSpectrogram
    sample_rate: ${sample_rate}
    n_fft: ${n_fft}
    hop_length: ${hop_length}
    win_length: ${win_length}
    n_mels: ${num_mels}
    f_min: 0
    f_max: 12000

  feature_mel_transform:
    _target_: fish_speech.models.vqgan.spectrogram.LogMelSpectrogram
    sample_rate: 32000
    n_fft: 2048
    hop_length: 320
    win_length: 2048
    n_mels: 128

  optimizer:
    _target_: torch.optim.AdamW
    _partial_: true
    lr: 1e-4
    betas: [0.9, 0.999]
    eps: 1e-5

  lr_scheduler:
    _target_: torch.optim.lr_scheduler.LambdaLR
    _partial_: true
    lr_lambda:
      _target_: fish_speech.scheduler.get_cosine_schedule_with_warmup_lr_lambda
      _partial_: true
      num_warmup_steps: 0
      num_training_steps: ${trainer.max_steps}
      final_lr_ratio: 0.05

callbacks:
  grad_norm_monitor:
    sub_module: 
      - vq_encoder
      - text_encoder
      - speaker_encoder
      - denoiser