text: import os
  import json
  
  import numpy as np
  import torch
  import torch.nn as nn
  from torch import Tensor
  
  # from transformer import Encoder, Decoder, PostNet
  # using conformer
  from torchtyping import TensorType
  
  from conformer import Encoder, Decoder
  from conformer.Models import ModelConfig
  from dataset import PreProcessConfig
  from transformer import PostNet
  from .modules import VariancePredictor, LengthRegulator, StatsJson
  from utils.tools import get_mask_from_lengths
  
  from typing import Optional, Tuple
  
  VarianceOutput = Tuple[Tensor, Optional[Tensor], Tensor]
  DecoderOutput = Tuple[Tensor, Tensor, Tensor, TensorType[bool], TensorType[bool], Tensor, Tensor]
  
  
  class VoiceVoxVariancePredictor(nn.Module):
      def __init__(
          self, preprocess_config: PreProcessConfig, model_config: ModelConfig, predict_energy: bool = True
      ):
          super(VoiceVoxVariancePredictor, self).__init__()
          self.predict_energy = predict_energy
  
          self.encoder = Encoder(preprocess_config, model_config)
          self.duration_predictor = VariancePredictor(model_config)
          self.pitch_predictor = VariancePredictor(model_config)
          self.energy_predictor = VariancePredictor(model_config)
  
      def forward(
          self,
          speakers: Tensor,
          texts: Tensor,
          src_lens: Tensor,
          max_src_len: Tensor,
          accents: Tensor,
      ) -> VarianceOutput:
          src_masks = get_mask_from_lengths(src_lens, max_src_len)
          output = self.encoder(texts, src_masks, accents, speakers)
  
          log_duration_prediction = self.duration_predictor(output, src_masks)
          pitch_prediction = self.pitch_predictor(output, src_masks)
  
          # energyなしでのFastSpeech2音声合成の品質検証のため
          if self.predict_energy:
              energy_prediction = self.energy_predictor(output, src_masks)
          else:
              energy_prediction = None
  
          return pitch_prediction, energy_prediction, log_duration_prediction
  
  
  class VoiceVoxDecoder(nn.Module):
      def __init__(
          self, preprocess_config: PreProcessConfig, model_config: ModelConfig
      ):
          super(VoiceVoxDecoder, self).__init__()
          self.encoder = Encoder(preprocess_config, model_config)
          self.length_regulator = LengthRegulator()
  
          pitch_quantization = model_config["variance_embedding"]["pitch_quantization"]
          energy_quantization = model_config["variance_embedding"]["energy_quantization"]
          n_bins = model_config["variance_embedding"]["n_bins"]
  
          assert pitch_quantization == "log"
          assert energy_quantization in ["linear", "log"]
          with open(
              os.path.join(preprocess_config["path"]["preprocessed_path"], "stats.json")
          ) as f:
              stats: StatsJson = json.load(f)
              pitch_data = stats["pitch"][:2]
              pitch_min, pitch_max = pitch_data[0], pitch_data[1]
              energy_data = stats["energy"][:2]
              energy_min, energy_max = energy_data[0], energy_data[1]
  
          self.pitch_bins = nn.Parameter(
              torch.exp(
                  torch.linspace(np.log(pitch_min), np.log(pitch_max), n_bins - 1)
              ),
              requires_grad=False,
          )
          if energy_quantization == "log":
              self.energy_bins = nn.Parameter(
                  torch.exp(
                      torch.linspace(np.log(energy_min), np.log(energy_max), n_bins - 1)
                  ),
                  requires_grad=False,
              )
          else:
              self.energy_bins = nn.Parameter(
                  torch.linspace(energy_min, energy_max, n_bins - 1),
                  requires_grad=False,
              )
  
          self.pitch_embedding = nn.Embedding(
              n_bins, model_config["conformer"]["encoder_hidden"]
          )
          self.energy_embedding = nn.Embedding(
              n_bins, model_config["conformer"]["encoder_hidden"]
          )
  
          self.decoder = Decoder(model_config)
          self.mel_linear = nn.Linear(
              model_config["conformer"]["decoder_hidden"],
              preprocess_config["preprocessing"]["mel"]["n_mel_channels"],
          )
          self.postnet = PostNet()
  
      def get_pitch_embedding(
          self, prediction: Tensor, target: Optional[Tensor], mask: TensorType[bool], control: float
      ) -> Tuple[Tensor, Tensor]:
          if target is not None:
              embedding = self.pitch_embedding(torch.bucketize(target, self.pitch_bins))
          else:
              prediction = prediction * control
              embedding = self.pitch_embedding(
                  torch.bucketize(prediction, self.pitch_bins)
              )
          return embedding
  
      def get_energy_embedding(
          self, prediction: Tensor, target: Optional[Tensor], mask: TensorType[bool], control: float
      ) -> Tuple[Tensor, Tensor]:
          if target is not None:
              embedding = self.energy_embedding(torch.bucketize(target, self.energy_bins))
          else:
              prediction = prediction * control
              embedding = self.energy_embedding(
                  torch.bucketize(prediction, self.energy_bins)
              )
          return embedding
  
      def forward(
          self,
          speakers: Tensor,
          texts: Tensor,
          src_lens: Tensor,
          max_src_len: Tensor,
          accents: Tensor,
          p_predictions: Tensor,
          e_predictions: Optional[Tensor],
          log_d_predictions: Tensor,
          mel_lens: Optional[Tensor] = None,
          max_mel_len: Optional[Tensor] = None,
          pitch_target: Optional[Tensor] = None,
          energy_target: Optional[Tensor] = None,
          duration_target: Optional[Tensor] = None,
          p_control: float = 1.0,
          e_control: float = 1.0,
          d_control: float = 1.0,
      ) -> DecoderOutput:
          src_masks = get_mask_from_lengths(src_lens, max_src_len)
          output = self.encoder(texts, src_masks, accents, speakers)
  
          pitch_embedding = self.get_pitch_embedding(
              p_predictions, pitch_target, src_masks, p_control
          )
          output += pitch_embedding
  
          if e_predictions is not None:
              energy_embedding = self.get_energy_embedding(
                  e_predictions, energy_target, src_masks, e_control
              )
              output += energy_embedding
  
          if duration_target is not None:
              output, mel_len = self.length_regulator(output, duration_target, max_mel_len)
              duration_rounded = duration_target
              mel_masks = (
                  get_mask_from_lengths(mel_lens, max_mel_len)
                  if mel_lens is not None
                  else None
              )
          else:
              duration_rounded = torch.clamp(
                  (torch.round(torch.exp(log_d_predictions) - 1) * d_control),
                  min=0,
              )
              output, mel_len = self.length_regulator(output, duration_rounded, max_mel_len)
              mel_masks = get_mask_from_lengths(mel_len)
  
          output, mel_masks = self.decoder(output, mel_masks)
          output = self.mel_linear(output)
  
          postnet_output = self.postnet(output) + output
  
          return output, postnet_output, duration_rounded, src_masks, mel_masks, src_lens, mel_lens
  