fish-speech/fish_speech/configs/hubert_vq_diffusion.yaml

defaults:
  - base
  - _self_

project: hubert_vq_diffusion

# Lightning Trainer
trainer:
  accelerator: gpu
  devices: 1
  strategy:
    _target_: lightning.pytorch.strategies.DDPStrategy
    static_graph: true
  gradient_clip_val: 1.0
  gradient_clip_algorithm: 'norm'
  precision: 16-mixed
  max_steps: 1_000_000
  val_check_interval: 1000

sample_rate: 44100
hop_length: 512
num_mels: 128
n_fft: 2048
win_length: 2048

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

val_dataset:
  _target_: fish_speech.datasets.vqgan.VQGANDataset
  filelist: data/vq_val_filelist.txt
  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: 0 #16
  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}

  text_encoder:
    _target_: fish_speech.models.vqgan.modules.encoders.TextEncoder
    in_channels: 1024
    out_channels: 128
    hidden_channels: 192
    hidden_channels_ffn: 768
    n_heads: 2
    n_layers: 4
    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: 1024
    vq_channels: 1024
    codebook_size: 2048
    downsample: 2
    kmeans_ckpt: results/hubert-vq-pretrain/kmeans.pt

  speaker_encoder:
    _target_: fish_speech.models.vqgan.modules.encoders.SpeakerEncoder
    in_channels: 128
    hidden_channels: 192
    out_channels: 512
    num_heads: 2
    num_layers: 4
    p_dropout: 0.1
  
  # denoiser:
  #   _target_: fish_speech.models.vq_diffusion.convnext_1d.ConvNext1DModel
  #   in_channels: 256
  #   out_channels: 128
  #   intermediate_dim: 512
  #   mlp_dim: 2048
  #   num_layers: 20
  #   dilation_cycle_length: 2
  #   time_embedding_type: "positional"

  denoiser:
    _target_: fish_speech.models.vq_diffusion.wavenet.WaveNet
    in_channels: 128
    out_channels: 128
    d_encoder: 128
    residual_channels: 512
    residual_layers: 20
    use_linear_bias: false
    dilation_cycle: 2

  # denoiser:
  #   _target_: fish_speech.models.vq_diffusion.unet1d.Unet1DDenoiser
  #   dim: 64
  #   dim_mults: [1, 2, 4]
  #   groups: 8
  #   pe_scale: 1000

  vocoder:
    _target_: fish_speech.models.vq_diffusion.adamos.ADaMoSHiFiGANV1

  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: 40
    f_max: 16000

  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