class_docstr: VITS module (generator + discriminator).
  
  This is a module of VITS described in `Conditional Variational Autoencoder
  with Adversarial Learning for End-to-End Text-to-Speech`_.
  
  .. _`Conditional Variational Autoencoder with Adversarial Learning for End-to-End
      Text-to-Speech`: https://arxiv.org/abs/2006.04558
docstr: Initialize VITS module.
  
  Args:
      idim (int): Input vocabrary size.
      odim (int): Acoustic feature dimension. The actual output channels will
          be 1 since VITS is the end-to-end text-to-wave model but for the
          compatibility odim is used to indicate the acoustic feature dimension.
      sampling_rate (int): Sampling rate, not used for the training but it will
          be referred in saving waveform during the inference.
      generator_type (str): Generator type.
      vocoder_generator_type (str): Type of vocoder generator to use in the model.
      generator_params (Dict[str, Any]): Parameter dict for generator.
      discriminator_type (str): Discriminator type.
      discriminator_params (Dict[str, Any]): Parameter dict for discriminator.
      generator_adv_loss_params (Dict[str, Any]): Parameter dict for generator
          adversarial loss.
      discriminator_adv_loss_params (Dict[str, Any]): Parameter dict for
          discriminator adversarial loss.
      feat_match_loss_params (Dict[str, Any]): Parameter dict for feat match loss.
      mel_loss_params (Dict[str, Any]): Parameter dict for mel loss.
      lambda_adv (float): Loss scaling coefficient for adversarial loss.
      lambda_mel (float): Loss scaling coefficient for mel spectrogram loss.
      lambda_feat_match (float): Loss scaling coefficient for feat match loss.
      lambda_dur (float): Loss scaling coefficient for duration loss.
      lambda_kl (float): Loss scaling coefficient for KL divergence loss.
      lambda_pitch (float): Loss scaling coefficient for pitch loss.
      lambda_phoneme (float): Loss scaling coefficient for phoneme loss.
      lambda_c_yin (float): Loss scaling coefficient for yin loss.
      cache_generator_outputs (bool): Whether to cache generator outputs.
docstr_tok: ["Initialize","VITS","module",".","Args",":","idim","(","int",")",":","Input","vocabrary","size",".","odim","(","int",")",":","Acoustic","feature","dimension",".","The","actual","output","channels","will","be","1","since","VITS","is","the","end-to-end","text-to-wave","model","but","for","the","compatibility","odim","is","used","to","indicate","the","acoustic","feature","dimension",".","sampling_rate","(","int",")",":","Sampling","rate",",","not","used","for","the","training","but","it","will","be","referred","in","saving","waveform","during","the","inference",".","generator_type","(","str",")",":","Generator","type",".","vocoder_generator_type","(","str",")",":","Type","of","vocoder","generator","to","use","in","the","model",".","generator_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","generator",".","discriminator_type","(","str",")",":","Discriminator","type",".","discriminator_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","discriminator",".","generator_adv_loss_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","generator","adversarial","loss",".","discriminator_adv_loss_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","discriminator","adversarial","loss",".","feat_match_loss_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","feat","match","loss",".","mel_loss_params","(","Dict","[","str",",","Any","]",")",":","Parameter","dict","for","mel","loss",".","lambda_adv","(","float",")",":","Loss","scaling","coefficient","for","adversarial","loss",".","lambda_mel","(","float",")",":","Loss","scaling","coefficient","for","mel","spectrogram","loss",".","lambda_feat_match","(","float",")",":","Loss","scaling","coefficient","for","feat","match","loss",".","lambda_dur","(","float",")",":","Loss","scaling","coefficient","for","duration","loss",".","lambda_kl","(","float",")",":","Loss","scaling","coefficient","for","KL","divergence","loss",".","lambda_pitch","(","float",")",":","Loss","scaling","coefficient","for","pitch","loss",".","lambda_phoneme","(","float",")",":","Loss","scaling","coefficient","for","phoneme","loss",".","lambda_c_yin","(","float",")",":","Loss","scaling","coefficient","for","yin","loss",".","cache_generator_outputs","(","bool",")",":","Whether","to","cache","generator","outputs","."]
code: def __init__(
      self,
      # generator related
      idim: int,
      odim: int,
      sampling_rate: int = 22050,
      generator_type: str = "visinger",
      vocoder_generator_type: str = "hifigan",
      generator_params: Dict[str, Any] = {
          "hidden_channels": 192,
          "spks": None,
          "langs": None,
          "spk_embed_dim": None,
          "global_channels": -1,
          "segment_size": 32,
          "text_encoder_attention_heads": 2,
          "text_encoder_ffn_expand": 4,
          "text_encoder_blocks": 6,
          "text_encoder_positionwise_layer_type": "conv1d",
          "text_encoder_positionwise_conv_kernel_size": 1,
          "text_encoder_positional_encoding_layer_type": "rel_pos",
          "text_encoder_self_attention_layer_type": "rel_selfattn",
          "text_encoder_activation_type": "swish",
          "text_encoder_normalize_before": True,
          "text_encoder_dropout_rate": 0.1,
          "text_encoder_positional_dropout_rate": 0.0,
          "text_encoder_attention_dropout_rate": 0.0,
          "text_encoder_conformer_kernel_size": 7,
          "use_macaron_style_in_text_encoder": True,
          "use_conformer_conv_in_text_encoder": True,
          "decoder_kernel_size": 7,
          "decoder_channels": 512,
          "decoder_upsample_scales": [8, 8, 2, 2],
          "decoder_upsample_kernel_sizes": [16, 16, 4, 4],
          "decoder_resblock_kernel_sizes": [3, 7, 11],
          "decoder_resblock_dilations": [
              [1, 3, 5],
              [1, 3, 5],
              [1, 3, 5],
          ],
          "projection_filters": [0, 1, 1, 1],
          "projection_kernels": [0, 5, 7, 11],
          "use_weight_norm_in_decoder": True,
          "posterior_encoder_kernel_size": 5,
          "posterior_encoder_layers": 16,
          "posterior_encoder_stacks": 1,
          "posterior_encoder_base_dilation": 1,
          "posterior_encoder_dropout_rate": 0.0,
          "use_weight_norm_in_posterior_encoder": True,
          "flow_flows": 4,
          "flow_kernel_size": 5,
          "flow_base_dilation": 1,
          "flow_layers": 4,
          "flow_dropout_rate": 0.0,
          "use_weight_norm_in_flow": True,
          "use_only_mean_in_flow": True,
          "expand_f0_method": "repeat",
          "use_phoneme_predictor": False,
      },
      # discriminator related
      discriminator_type: str = "hifigan_multi_scale_multi_period_discriminator",
      discriminator_params: Dict[str, Any] = {
          "hifigan_multi_scale_multi_period_discriminator": {
              "scales": 1,
              "scale_downsample_pooling": "AvgPool1d",
              "scale_downsample_pooling_params": {
                  "kernel_size": 4,
                  "stride": 2,
                  "padding": 2,
              },
              "scale_discriminator_params": {
                  "in_channels": 1,
                  "out_channels": 1,
                  "kernel_sizes": [15, 41, 5, 3],
                  "channels": 128,
                  "max_downsample_channels": 1024,
                  "max_groups": 16,
                  "bias": True,
                  "downsample_scales": [2, 2, 4, 4, 1],
                  "nonlinear_activation": "LeakyReLU",
                  "nonlinear_activation_params": {
                      "negative_slope": 0.1
                  },
                  "use_weight_norm": True,
                  "use_spectral_norm": False,
              },
              "follow_official_norm": False,
              "periods": [2, 3, 5, 7, 11],
              "period_discriminator_params": {
                  "in_channels": 1,
                  "out_channels": 1,
                  "kernel_sizes": [5, 3],
                  "channels": 32,
                  "downsample_scales": [3, 3, 3, 3, 1],
                  "max_downsample_channels": 1024,
                  "bias": True,
                  "nonlinear_activation": "LeakyReLU",
                  "nonlinear_activation_params": {
                      "negative_slope": 0.1
                  },
                  "use_weight_norm": True,
                  "use_spectral_norm": False,
              },
          },
          "avocodo": {
              "combd": {
                  "combd_h_u": [
                      [16, 64, 256, 1024, 1024, 1024],
                      [16, 64, 256, 1024, 1024, 1024],
                      [16, 64, 256, 1024, 1024, 1024],
                  ],
                  "combd_d_k": [
                      [7, 11, 11, 11, 11, 5],
                      [11, 21, 21, 21, 21, 5],
                      [15, 41, 41, 41, 41, 5],
                  ],
                  "combd_d_s": [
                      [1, 1, 4, 4, 4, 1],
                      [1, 1, 4, 4, 4, 1],
                      [1, 1, 4, 4, 4, 1],
                  ],
                  "combd_d_d": [
                      [1, 1, 1, 1, 1, 1],
                      [1, 1, 1, 1, 1, 1],
                      [1, 1, 1, 1, 1, 1],
                  ],
                  "combd_d_g": [
                      [1, 4, 16, 64, 256, 1],
                      [1, 4, 16, 64, 256, 1],
                      [1, 4, 16, 64, 256, 1],
                  ],
                  "combd_d_p": [
                      [3, 5, 5, 5, 5, 2],
                      [5, 10, 10, 10, 10, 2],
                      [7, 20, 20, 20, 20, 2],
                  ],
                  "combd_op_f": [1, 1, 1],
                  "combd_op_k": [3, 3, 3],
                  "combd_op_g": [1, 1, 1],
              },
              "sbd": {
                  "use_sbd": True,
                  "sbd_filters": [
                      [64, 128, 256, 256, 256],
                      [64, 128, 256, 256, 256],
                      [64, 128, 256, 256, 256],
                      [32, 64, 128, 128, 128],
                  ],
                  "sbd_strides": [
                      [1, 1, 3, 3, 1],
                      [1, 1, 3, 3, 1],
                      [1, 1, 3, 3, 1],
                      [1, 1, 3, 3, 1],
                  ],
                  "sbd_kernel_sizes": [
                      [
                          [7, 7, 7],
                          [7, 7, 7],
                          [7, 7, 7],
                          [7, 7, 7],
                          [7, 7, 7],
                      ],
                      [
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                      ],
                      [
                          [3, 3, 3],
                          [3, 3, 3],
                          [3, 3, 3],
                          [3, 3, 3],
                          [3, 3, 3],
                      ],
                      [
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                          [5, 5, 5],
                      ],
                  ],
                  "sbd_dilations": [
                      [
                          [5, 7, 11],
                          [5, 7, 11],
                          [5, 7, 11],
                          [5, 7, 11],
                          [5, 7, 11],
                      ],
                      [
                          [3, 5, 7],
                          [3, 5, 7],
                          [3, 5, 7],
                          [3, 5, 7],
                          [3, 5, 7],
                      ],
                      [
                          [1, 2, 3],
                          [1, 2, 3],
                          [1, 2, 3],
                          [1, 2, 3],
                          [1, 2, 3],
                      ],
                      [
                          [1, 2, 3],
                          [1, 2, 3],
                          [1, 2, 3],
                          [2, 3, 5],
                          [2, 3, 5],
                      ],
                  ],
                  "sbd_band_ranges": [
                      [0, 6],
                      [0, 11],
                      [0, 16],
                      [0, 64],
                  ],
                  "sbd_transpose": [False, False, False, True],
                  "pqmf_config": {
                      "sbd": [16, 256, 0.03, 10.0],
                      "fsbd": [64, 256, 0.1, 9.0],
                  },
              },
              "pqmf_config": {
                  "lv1": [2, 256, 0.25, 10.0],
                  "lv2": [4, 192, 0.13, 10.0],
              },
          },
      },
      # loss related
      generator_adv_loss_params: Dict[str, Any] = {
          "average_by_discriminators": False,
          "loss_type": "mse",
      },
      discriminator_adv_loss_params: Dict[str, Any] = {
          "average_by_discriminators": False,
          "loss_type": "mse",
      },
      feat_match_loss_params: Dict[str, Any] = {
          "average_by_discriminators": False,
          "average_by_layers": False,
          "include_final_outputs": True,
      },
      mel_loss_params: Dict[str, Any] = {
          "fs": 22050,
          "n_fft": 1024,
          "hop_length": 256,
          "win_length": None,
          "window": "hann",
          "n_mels": 80,
          "fmin": 0,
          "fmax": None,
          "log_base": None,
      },
      lambda_adv: float = 1.0,
      lambda_mel: float = 45.0,
      lambda_feat_match: float = 2.0,
      lambda_dur: float = 0.1,
      lambda_kl: float = 1.0,
      lambda_pitch: float = 10.0,
      lambda_phoneme: float = 1.0,
      lambda_c_yin: float = 45.0,
      cache_generator_outputs: bool = True,
  ):
      """Initialize VITS module.
  
      Args:
          idim (int): Input vocabrary size.
          odim (int): Acoustic feature dimension. The actual output channels will
              be 1 since VITS is the end-to-end text-to-wave model but for the
              compatibility odim is used to indicate the acoustic feature dimension.
          sampling_rate (int): Sampling rate, not used for the training but it will
              be referred in saving waveform during the inference.
          generator_type (str): Generator type.
          vocoder_generator_type (str): Type of vocoder generator to use in the model.
          generator_params (Dict[str, Any]): Parameter dict for generator.
          discriminator_type (str): Discriminator type.
          discriminator_params (Dict[str, Any]): Parameter dict for discriminator.
          generator_adv_loss_params (Dict[str, Any]): Parameter dict for generator
              adversarial loss.
          discriminator_adv_loss_params (Dict[str, Any]): Parameter dict for
              discriminator adversarial loss.
          feat_match_loss_params (Dict[str, Any]): Parameter dict for feat match loss.
          mel_loss_params (Dict[str, Any]): Parameter dict for mel loss.
          lambda_adv (float): Loss scaling coefficient for adversarial loss.
          lambda_mel (float): Loss scaling coefficient for mel spectrogram loss.
          lambda_feat_match (float): Loss scaling coefficient for feat match loss.
          lambda_dur (float): Loss scaling coefficient for duration loss.
          lambda_kl (float): Loss scaling coefficient for KL divergence loss.
          lambda_pitch (float): Loss scaling coefficient for pitch loss.
          lambda_phoneme (float): Loss scaling coefficient for phoneme loss.
          lambda_c_yin (float): Loss scaling coefficient for yin loss.
          cache_generator_outputs (bool): Whether to cache generator outputs.
  
      """
      assert check_argument_types()
      super().__init__()
  
      # define modules
      generator_class = AVAILABLE_GENERATERS[generator_type]
      if "visinger" in generator_type or "pisinger" in generator_type:
          # NOTE(kan-bayashi): Update parameters for the compatibility.
          #   The idim and odim is automatically decided from input data,
          #   where idim represents #vocabularies and odim represents
          #   the input acoustic feature dimension.
          generator_params.update(vocabs=idim, aux_channels=odim)
      self.generator_type = generator_type
      self.use_flow = (
          True if generator_params["flow_flows"] > 0 else False
      )
      self.use_phoneme_predictor = generator_params[
          "use_phoneme_predictor"
      ]
      self.discriminator_type = discriminator_type
      if "avocodo" in discriminator_type:
          use_avocodo = True
          vocoder_generator_type = "avocodo"
      else:
          use_avocodo = False
      self.use_avocodo = use_avocodo
      self.vocoder_generator_type = vocoder_generator_type
      generator_params.update(generator_type=generator_type)
      generator_params.update(
          vocoder_generator_type=vocoder_generator_type
      )
      generator_params.update(fs=mel_loss_params["fs"])
      generator_params.update(hop_length=mel_loss_params["hop_length"])
      generator_params.update(win_length=mel_loss_params["win_length"])
      generator_params.update(n_fft=mel_loss_params["n_fft"])
      if vocoder_generator_type == "uhifigan" and use_avocodo:
          generator_params.update(use_avocodo=use_avocodo)
      self.generator = generator_class(
          **generator_params,
      )
  
      discriminator_class = AVAILABLE_DISCRIMINATORS[
          self.discriminator_type
      ]
      if use_avocodo:
          discriminator_params.update(
              projection_filters=generator_params["projection_filters"]
          )
          discriminator_params["sbd"].update(
              segment_size=generator_params["segment_size"]
              * mel_loss_params["hop_length"]
          )
      if "visinger2" in discriminator_type:
          discriminator_params["multi_freq_disc_params"].update(
              sample_rate=sampling_rate
          )
  
      self.discriminator = discriminator_class(
          **discriminator_params,
      )
      self.generator_adv_loss = GeneratorAdversarialLoss(
          **generator_adv_loss_params,
      )
      self.discriminator_adv_loss = DiscriminatorAdversarialLoss(
          **discriminator_adv_loss_params,
      )
      self.feat_match_loss = FeatureMatchLoss(
          **feat_match_loss_params,
      )
      self.mel_loss = MelSpectrogramLoss(
          **mel_loss_params,
      )
      if self.use_flow:
          self.kl_loss = KLDivergenceLoss()
      else:
          self.kl_loss = KLDivergenceLossWithoutFlow()
  
      self.ctc_loss = torch.nn.CTCLoss(idim - 1, reduction="mean")
      self.mse_loss = torch.nn.MSELoss()
  
      # coefficients
      self.lambda_adv = lambda_adv
      self.lambda_mel = lambda_mel
      self.lambda_kl = lambda_kl
      self.lambda_feat_match = lambda_feat_match
      self.lambda_dur = lambda_dur
      self.lambda_pitch = lambda_pitch
      self.lambda_phoneme = lambda_phoneme
      self.lambda_c_yin = lambda_c_yin
  
      # cache
      self.cache_generator_outputs = cache_generator_outputs
      self._cache = None
  
      # store sampling rate for saving wav file
      # (not used for the training)
      self.fs = sampling_rate
  
      # store parameters for test compatibility
      self.spks = self.generator.spks
      self.langs = self.generator.langs
      self.spk_embed_dim = self.generator.spk_embed_dim
  
imports: from contextlib import contextmanager
  from distutils.version import LooseVersion
  from typing import Any, Dict, Optional
  import torch
  from torch.nn import functional
  from typeguard import check_argument_types
  from espnet2.gan_svs.abs_gan_svs import AbsGANSVS
  from espnet2.gan_svs.avocodo.avocodo import SBD, AvocodoDiscriminator, AvocodoDiscriminatorPlus, CoMBD
  from espnet2.gan_svs.visinger2.visinger2_vocoder import VISinger2Discriminator
  from espnet2.gan_svs.vits.generator import VISingerGenerator
  from espnet2.gan_tts.hifigan import HiFiGANMultiPeriodDiscriminator, HiFiGANMultiScaleDiscriminator, HiFiGANMultiScaleMultiPeriodDiscriminator, HiFiGANPeriodDiscriminator, HiFiGANScaleDiscriminator
  from espnet2.gan_tts.hifigan.loss import DiscriminatorAdversarialLoss, FeatureMatchLoss, GeneratorAdversarialLoss, MelSpectrogramLoss
  from espnet2.gan_tts.utils import get_segments
  from espnet2.gan_tts.vits.loss import KLDivergenceLoss, KLDivergenceLossWithoutFlow
  from espnet2.torch_utils.device_funcs import force_gatherable
label_desc: Use image node_id 1 to create a new VITS object from inherited base classes: AbsGANSVS with example: obj = VITS(idim, odim, sampling_rate, generator_type, vocoder_generator_type, generator_params, discriminator_type, discriminator_params, generator_adv_loss_params, discriminator_adv_loss_params, feat_match_loss_params, mel_loss_params, lambda_adv, lambda_mel, lambda_feat_match, lambda_dur, lambda_kl, lambda_pitch, lambda_phoneme, lambda_c_yin, cache_generator_outputs)