<reponame>y-chan/FastSpeech2<filename>model/vv_fastspeech2.py
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

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
# from transformer import Encoder, Decoder, PostNet # using conformer # energyなしでのFastSpeech2音声合成の品質検証のため