# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
dataset:
  batch_size: 64
  frames_per_clip: 64 # 16384
  clips_per_utterance: 10
  padding_frames: 0
  dataloader_num_workers: 4
  dataloader_prefetch_factor: 4
model:
  # Generator Config
  generator_type: "MelGANGenerator"
  generator_params:
    in_channels: 80               # Number of input channels.
    out_channels: 4               # Number of output channels.
    kernel_size: 7                # Kernel size of initial and final conv layers.
    channels: 384                 # Initial number of channels for conv layers.
    upsample_scales: [5, 5, 3] # List of Upsampling scales.
    stack_kernel_size: 3          # Kernel size of dilated conv layers in residual stack.
    stacks: 4                     # Number of stacks in a single residual stack module.
    use_weight_norm: True         # Whether to use weight normalization.
    use_causal_conv: False        # Whether to use causal convolution.
  generator_optimizer_type: "RAdam"
  generator_grad_norm: 10         # Generator's gradient norm.
  generator_optimizer:
    lr: 1.0e-3             # Generator's learning rate.
    eps: 1.0e-6            # Generator's epsilon.
    weight_decay: 0.0      # Generator's weight decay coefficient.
  generator_scheduler_type: "MultiStepLR" # Generator's scheduler type.
  generator_scheduler_params:
    gamma: 0.5                          # Generator's scheduler gamma.
    milestones:                         # At each milestone, lr will be multiplied by gamma.
      - 100000
      - 200000
      - 300000
      - 400000
      - 500000
      - 600000

  # Discriminator Config
  discriminator_type: ParallelWaveGANDiscriminator
  discriminator_params:
    in_channels: 1
    out_channels: 1
    scales: null
    layers: 10
    downsample_pooling: null
    downsample_pooling_params: null
    kernel_sizes: null
    kernel_size: 3
    channels: null
    max_downsample_channels: null
    downsample_scales: null
    nonlinear_activation: LeakyReLU
    nonlinear_activation_params:
      negative_slope: 0.2
    use_weight_norm: true
    bias: true
    conv_channels: 64
  discriminator_optimizer_type: RAdam
  discriminator_optimizer:
    lr: 5.0e-05
    eps: 1.0e-06
    weight_decay: 0.0
    amsgrad: null
  discriminator_grad_norm: 1
  discriminator_scheduler_type: StepLR
  discriminator_scheduler_params:
    gamma: 0.5
    step_size: 200000
    milestones: null
  # Loss
  stft_loss_params:
      fft_sizes: [1024, 2048, 512]  # List of FFT size for STFT-based loss.
      hop_sizes: [120, 240, 50]     # List of hop size for STFT-based loss
      win_lengths: [600, 1200, 240] # List of window length for STFT-based loss.
      window: "hann_window"         # Window function for STFT-based loss
  use_subband_stft_loss: true
  subband_stft_loss_params:
      fft_sizes: [384, 683, 171]  # List of FFT size for STFT-based loss.
      hop_sizes: [30, 60, 10]     # List of hop size for STFT-based loss
      win_lengths: [150, 300, 60] # List of window length for STFT-based loss.
      window: "hann_window"       # Window function for STFT-based loss
  use_feat_match_loss: false # Whether to use feature matching loss.
  lambda_adv: 2.5          # Loss balancing coefficient for adversarial loss.

  # Training
  n_iterations: 4000000
  discriminator_train_start_steps: 200000
