# 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: 6
  frames_per_clip: 100
  clips_per_utterance: 10
  padding_frames: 2
  dataloader_num_workers: 4
  dataloader_prefetch_factor: 4
model:
  # Generator Config
  generator_type: "ParallelWaveGANGenerator"
  generator_params:
    in_channels: 1        # Number of input channels.
    out_channels: 1       # Number of output channels.
    kernel_size: 3        # Kernel size of dilated convolution.
    layers: 30            # Number of residual block layers.
    stacks: 3             # Number of stacks i.e., dilation cycles.
    residual_channels: 64 # Number of channels in residual conv.
    gate_channels: 128    # Number of channels in gated conv.
    skip_channels: 64     # Number of channels in skip conv.
    aux_channels: 80      # Number of channels for auxiliary feature conv.
                          # Must be the same as num_mels.
    aux_context_window: 2 # Context window size for auxiliary feature.
                          # If set to 2, previous 2 and future 2 frames will be considered.
    dropout: 0.0          # Dropout rate. 0.0 means no dropout applied.
    use_weight_norm: true # Whether to use weight norm.
                          # If set to true, it will be applied to all of the conv layers.
    upsample_net: "ConvInUpsampleNetwork" # Upsampling network architecture.
    upsample_params:                      # Upsampling network parameters.
      upsample_scales: [4, 5, 3, 5]     # Upsampling scales. Prodcut of these must be the same as hop size.
  generator_grad_norm: 10         # Generator's gradient norm.
  generator_optimizer:
    lr: 0.0001             # Generator's learning rate.
    eps: 1.0e-6            # Generator's epsilon.
    weight_decay: 0.0      # Generator's weight decay coefficient.
  generator_scheduler_params:
    step_size: 200000      # Generator's scheduler step size.
    gamma: 0.5             # Generator's scheduler gamma.
                           # At each step size, lr will be multiplied by this parameter.

  # Discriminator Config
  discriminator_type: "ParallelWaveGANDiscriminator"
  discriminator_params:
    in_channels: 1        # Number of input channels.
    out_channels: 1       # Number of output channels.
    kernel_size: 3        # Number of output channels.
    layers: 10            # Number of conv layers.
    conv_channels: 64     # Number of chnn layers.
    bias: true            # Whether to use bias parameter in conv.
    use_weight_norm: true # Whether to use weight norm.
                          # If set to true, it will be applied to all of the conv layers.
    nonlinear_activation: "LeakyReLU" # Nonlinear function after each conv.
    nonlinear_activation_params:      # Nonlinear function parameters
        negative_slope: 0.2           # Alpha in LeakyReLU.
  discriminator_grad_norm: 1        # Discriminator's gradient norm.
  discriminator_optimizer:
    lr: 0.00005            # Discriminator's learning rate.
    eps: 1.0e-6            # Discriminator's epsilon.
    weight_decay: 0.0      # Discriminator's weight decay coefficient.
  discriminator_scheduler_params:
    step_size: 200000      # Discriminator's scheduler step size.
    gamma: 0.5             # Discriminator's scheduler gamma.
                           # At each step size, lr will be multiplied by this parameter.
  # 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: False
  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_feat_match: 25.0   # Loss balancing coefficient for feature matching loss.
  lambda_adv: 4.0          # Loss balancing coefficient for adversarial loss.

  n_iterations: 4000000
  discriminator_train_start_steps: 100000
