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
class_docstr_tok: ["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.
      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.
      cache_generator_outputs (bool): Whether to cache generator outputs.
      plot_pred_mos (bool): Whether to plot predicted MOS during the training.
      mos_pred_tool (str): MOS prediction tool name.
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",".","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",".","cache_generator_outputs","(","bool",")",":","Whether","to","cache","generator","outputs",".","plot_pred_mos","(","bool",")",":","Whether","to","plot","predicted","MOS","during","the","training",".","mos_pred_tool","(","str",")",":","MOS","prediction","tool","name","."]
code: def __init__(
      self,
      # generator related
      idim: int,
      odim: int,
      sampling_rate: int = 22050,
      generator_type: str = "vits_generator",
      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],
          ],
          "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,
          "stochastic_duration_predictor_kernel_size": 3,
          "stochastic_duration_predictor_dropout_rate": 0.5,
          "stochastic_duration_predictor_flows": 4,
          "stochastic_duration_predictor_dds_conv_layers": 3,
      },
      # discriminator related
      discriminator_type: str = "hifigan_multi_scale_multi_period_discriminator",
      discriminator_params: Dict[str, Any] = {
          "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,
          },
      },
      # 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 = 1.0,
      lambda_kl: float = 1.0,
      cache_generator_outputs: bool = True,
      plot_pred_mos: bool = False,
      mos_pred_tool: str = "utmos",
  ):
      """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.
          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.
          cache_generator_outputs (bool): Whether to cache generator outputs.
          plot_pred_mos (bool): Whether to plot predicted MOS during the training.
          mos_pred_tool (str): MOS prediction tool name.
  
      """
      assert check_argument_types()
      super().__init__()
  
      # define modules
      generator_class = AVAILABLE_GENERATERS[generator_type]
      if generator_type == "vits_generator":
          # 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 = generator_class(
          **generator_params,
      )
      discriminator_class = AVAILABLE_DISCRIMINATORS[discriminator_type]
      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,
      )
      self.kl_loss = KLDivergenceLoss()
  
      # 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
  
      # 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
  
      # plot pseudo mos during training
      self.plot_pred_mos = plot_pred_mos
      if plot_pred_mos:
          if mos_pred_tool == "utmos":
              # Load predictor for UTMOS22 (https://arxiv.org/abs/2204.02152)
              self.predictor = torch.hub.load(
                  "tarepan/SpeechMOS:v1.2.0", "utmos22_strong"
              )
          else:
              raise NotImplementedError(
                  f"Not supported mos_pred_tool: {mos_pred_tool}"
              )
  
imports: from contextlib import contextmanager
  from distutils.version import LooseVersion
  from typing import Any, Dict, Optional
  import torch
  from typeguard import check_argument_types
  from espnet2.gan_tts.abs_gan_tts import AbsGANTTS
  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.generator import VITSGenerator
  from espnet2.gan_tts.vits.loss import KLDivergenceLoss
  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: AbsGANTTS with example: obj = VITS(idim, odim, sampling_rate, 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, cache_generator_outputs, plot_pred_mos, mos_pred_tool)