class_docstr: JETS module (generator + discriminator).
  
  This is a module of JETS described in `JETS: Jointly Training FastSpeech2
  and HiFi-GAN for End to End Text to Speech'_.
  
  .. _`JETS: Jointly Training FastSpeech2 and HiFi-GAN for End to End Text to Speech`
      : https://arxiv.org/abs/2203.16852
class_docstr_tok: ["JETS","module","(","generator","+","discriminator",")",".","This","is","a","module","of","JETS","described","in","`","JETS",":","Jointly","Training","FastSpeech2","and","HiFi-GAN","for","End","to","End","Text","to","Speech","'","_",".","..","_","`","JETS",":","Jointly","Training","FastSpeech2","and","HiFi-GAN","for","End","to","End","Text","to","Speech","`",":","https",":","\/\/arxiv.org\/abs\/2203.16852"]
docstr: Initialize JETS module.
  
  Args:
      idim (int): Input vocabrary size.
      odim (int): Acoustic feature dimension. The actual output channels will
          be 1 since JETS 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_var (float): Loss scaling coefficient for variance loss.
      lambda_align (float): Loss scaling coefficient for alignment 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","JETS","module",".","Args",":","idim","(","int",")",":","Input","vocabrary","size",".","odim","(","int",")",":","Acoustic","feature","dimension",".","The","actual","output","channels","will","be","1","since","JETS","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_var","(","float",")",":","Loss","scaling","coefficient","for","variance","loss",".","lambda_align","(","float",")",":","Loss","scaling","coefficient","for","alignment","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 = "jets_generator",
      generator_params: Dict[str, Any] = {
          "adim": 256,
          "aheads": 2,
          "elayers": 4,
          "eunits": 1024,
          "dlayers": 4,
          "dunits": 1024,
          "positionwise_layer_type": "conv1d",
          "positionwise_conv_kernel_size": 1,
          "use_scaled_pos_enc": True,
          "use_batch_norm": True,
          "encoder_normalize_before": True,
          "decoder_normalize_before": True,
          "encoder_concat_after": False,
          "decoder_concat_after": False,
          "reduction_factor": 1,
          "encoder_type": "transformer",
          "decoder_type": "transformer",
          "transformer_enc_dropout_rate": 0.1,
          "transformer_enc_positional_dropout_rate": 0.1,
          "transformer_enc_attn_dropout_rate": 0.1,
          "transformer_dec_dropout_rate": 0.1,
          "transformer_dec_positional_dropout_rate": 0.1,
          "transformer_dec_attn_dropout_rate": 0.1,
          "conformer_rel_pos_type": "latest",
          "conformer_pos_enc_layer_type": "rel_pos",
          "conformer_self_attn_layer_type": "rel_selfattn",
          "conformer_activation_type": "swish",
          "use_macaron_style_in_conformer": True,
          "use_cnn_in_conformer": True,
          "zero_triu": False,
          "conformer_enc_kernel_size": 7,
          "conformer_dec_kernel_size": 31,
          "duration_predictor_layers": 2,
          "duration_predictor_chans": 384,
          "duration_predictor_kernel_size": 3,
          "duration_predictor_dropout_rate": 0.1,
          "energy_predictor_layers": 2,
          "energy_predictor_chans": 384,
          "energy_predictor_kernel_size": 3,
          "energy_predictor_dropout": 0.5,
          "energy_embed_kernel_size": 1,
          "energy_embed_dropout": 0.5,
          "stop_gradient_from_energy_predictor": False,
          "pitch_predictor_layers": 5,
          "pitch_predictor_chans": 384,
          "pitch_predictor_kernel_size": 5,
          "pitch_predictor_dropout": 0.5,
          "pitch_embed_kernel_size": 1,
          "pitch_embed_dropout": 0.5,
          "stop_gradient_from_pitch_predictor": True,
          "generator_out_channels": 1,
          "generator_channels": 512,
          "generator_global_channels": -1,
          "generator_kernel_size": 7,
          "generator_upsample_scales": [8, 8, 2, 2],
          "generator_upsample_kernel_sizes": [16, 16, 4, 4],
          "generator_resblock_kernel_sizes": [3, 7, 11],
          "generator_resblock_dilations": [
              [1, 3, 5],
              [1, 3, 5],
              [1, 3, 5],
          ],
          "generator_use_additional_convs": True,
          "generator_bias": True,
          "generator_nonlinear_activation": "LeakyReLU",
          "generator_nonlinear_activation_params": {
              "negative_slope": 0.1
          },
          "generator_use_weight_norm": True,
          "segment_size": 64,
          "spks": -1,
          "langs": -1,
          "spk_embed_dim": None,
          "spk_embed_integration_type": "add",
          "use_gst": False,
          "gst_tokens": 10,
          "gst_heads": 4,
          "gst_conv_layers": 6,
          "gst_conv_chans_list": [32, 32, 64, 64, 128, 128],
          "gst_conv_kernel_size": 3,
          "gst_conv_stride": 2,
          "gst_gru_layers": 1,
          "gst_gru_units": 128,
          "init_type": "xavier_uniform",
          "init_enc_alpha": 1.0,
          "init_dec_alpha": 1.0,
          "use_masking": False,
          "use_weighted_masking": False,
      },
      # 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_var: float = 1.0,
      lambda_align: float = 2.0,
      cache_generator_outputs: bool = True,
      plot_pred_mos: bool = False,
      mos_pred_tool: str = "utmos",
  ):
      """Initialize JETS module.
  
      Args:
          idim (int): Input vocabrary size.
          odim (int): Acoustic feature dimension. The actual output channels will
              be 1 since JETS 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_var (float): Loss scaling coefficient for variance loss.
          lambda_align (float): Loss scaling coefficient for alignment 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]
      generator_params.update(idim=idim, odim=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.var_loss = VarianceLoss()
      self.forwardsum_loss = ForwardSumLoss()
  
      # coefficients
      self.lambda_adv = lambda_adv
      self.lambda_mel = lambda_mel
      self.lambda_feat_match = lambda_feat_match
      self.lambda_var = lambda_var
      self.lambda_align = lambda_align
  
      # 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
      self.use_gst = getattr(self.generator, "use_gst", False)
  
      # 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}"
              )
  
file_path: espnet/espnet2/gan_tts/jets/jets.py
imports: 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.jets.generator import JETSGenerator
  from espnet2.gan_tts.jets.loss import ForwardSumLoss, VarianceLoss
  from espnet2.gan_tts.utils import get_segments
  from espnet2.torch_utils.device_funcs import force_gatherable
label_desc: Use image node_id 1 to create a new JETS object from inherited base classes: AbsGANTTS with example: obj = JETS(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_var, lambda_align, cache_generator_outputs, plot_pred_mos, mos_pred_tool)