model.yaml

# lightning.pytorch==2.6.1
model:
  class_path: kanade_tokenizer.pipeline.KanadeModel
  init_args:
    ssl_feature_extractor:
      class_path: kanade_tokenizer.module.ssl_extractor.SSLFeatureExtractor
      init_args:
        model_name: wavlm_base_plus
        output_layer: 9 # Use at most 9 layers
        sample_rate: 24000 # Consistent to the target sample rate for reconstruction

    local_encoder:
      class_path: kanade_tokenizer.module.transformer.Transformer
      init_args:
        dim: 768
        n_layers: 6
        n_heads: 12
        window_size: 125
        use_rope: true
        rope_theta: 10000.0
        max_seq_len: 512
        use_flash_attention: true

    local_quantizer:
      class_path: kanade_tokenizer.module.fsq.FiniteScalarQuantizer
      init_args:
        input_dim: 768 # Must match local encoder output dimension
        output_dim: 768 # Must match feature decoder input dimension
        levels: [8, 8, 8, 5, 5] # 12800

    feature_decoder:
      class_path: kanade_tokenizer.module.transformer.Transformer
      init_args:
        dim: 768
        n_layers: 6
        n_heads: 12
        window_size: 125
        use_rope: true
        rope_theta: 10000.0
        max_seq_len: 512
        use_flash_attention: true

    global_encoder:
      class_path: kanade_tokenizer.module.global_encoder.GlobalEncoder
      init_args:
        input_channels: 768 # WavLM base plus feature dimension
        output_channels: 128
        num_layers: 4
        dim: 384
        intermediate_dim: 1152

    mel_prenet:
      class_path: kanade_tokenizer.module.transformer.Transformer
      init_args:
        dim: 768
        output_dim: 512
        n_layers: 6
        n_heads: 12
        window_size: 65
        use_rope: true
        rope_theta: 10000.0
        max_seq_len: 512
        use_flash_attention: true

    mel_decoder:
      class_path: kanade_tokenizer.module.transformer.Transformer
      init_args:
        dim: 512
        output_dim: 100 # Number of mel frequency bins
        n_layers: 6
        n_heads: 8
        window_size: 65
        use_rope: true
        rope_theta: 10000.0
        max_seq_len: 512
        adanorm_condition_dim: 128 # Must match global encoder output dimension
        use_adaln_zero: true # Use AdaLNZero for conditioning
        use_flash_attention: true

    mel_postnet:
      class_path: kanade_tokenizer.module.postnet.PostNet
      init_args:
        input_channels: 100 # Number of mel frequency bins
        channels: 256
        kernel_size: 7
        num_layers: 4
        use_layer_norm: true

    config:
      # SSL Feature settings
      local_ssl_layers: [6, 9]
      global_ssl_layers: [1, 2, 3, 4]
      normalize_ssl_features: true

      # Down/up-sampling settings
      downsample_factor: 1
      mel_upsample_factor: 2
      use_conv_downsample: true
      mel_interpolation_mode: linear

      # Audio settings
      sample_rate: 24000
      n_fft: 1024
      hop_length: 256
      n_mels: 100
      padding: center