# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

import torch
from torch import nn

from .alignment import mas_width1 as mas
from .attribute_prediction_model import get_attribute_prediction_model
from .common import (
    AffineTransformationLayer,
    ConvAttention,
    Encoder,
    ExponentialClass,
    Invertible1x1Conv,
    Invertible1x1ConvLUS,
    LengthRegulator,
    LinearNorm,
    get_mask_from_lengths,
)
from .torch_env import device


class FlowStep(nn.Module):
    def __init__(
        self,
        n_mel_channels,
        n_context_dim,
        n_layers,
        affine_model="simple_conv",
        scaling_fn="exp",
        matrix_decomposition="",
        affine_activation="softplus",
        use_partial_padding=False,
        cache_inverse=False,
    ):
        super(FlowStep, self).__init__()
        if matrix_decomposition == "LUS":
            self.invtbl_conv = Invertible1x1ConvLUS(
                n_mel_channels, cache_inverse=cache_inverse
            )
        else:
            self.invtbl_conv = Invertible1x1Conv(
                n_mel_channels, cache_inverse=cache_inverse
            )

        self.affine_tfn = AffineTransformationLayer(
            n_mel_channels,
            n_context_dim,
            n_layers,
            affine_model=affine_model,
            scaling_fn=scaling_fn,
            affine_activation=affine_activation,
            use_partial_padding=use_partial_padding,
        )

    def enable_inverse_cache(self):
        self.invtbl_conv.cache_inverse = True

    def forward(self, z, context, inverse=False, seq_lens=None):
        if inverse:  # for inference z-> mel
            z = self.affine_tfn(z, context, inverse, seq_lens=seq_lens)
            z = self.invtbl_conv(z, inverse)
            return z
        else:  # training mel->z
            z, log_det_W = self.invtbl_conv(z)
            z, log_s = self.affine_tfn(z, context, seq_lens=seq_lens)
            return z, log_det_W, log_s


class RADTTS(torch.nn.Module):
    def __init__(
        self,
        n_speakers,
        n_speaker_dim,
        n_text,
        n_text_dim,
        n_flows,
        n_conv_layers_per_step,
        n_mel_channels,
        n_hidden,
        mel_encoder_n_hidden,
        dummy_speaker_embedding,
        n_early_size,
        n_early_every,
        n_group_size,
        affine_model,
        dur_model_config,
        f0_model_config,
        energy_model_config,
        v_model_config=None,
        include_modules="dec",
        scaling_fn="exp",
        matrix_decomposition="",
        learn_alignments=False,
        affine_activation="softplus",
        attn_use_CTC=True,
        use_speaker_emb_for_alignment=False,
        use_context_lstm=False,
        context_lstm_norm=None,
        text_encoder_lstm_norm=None,
        n_f0_dims=0,
        n_energy_avg_dims=0,
        context_lstm_w_f0_and_energy=True,
        use_first_order_features=False,
        unvoiced_bias_activation="",
        ap_pred_log_f0=False,
        **kwargs,
    ):
        super(RADTTS, self).__init__()
        assert n_early_size % 2 == 0
        self.do_mel_descaling = kwargs.get("do_mel_descaling", True)
        self.n_mel_channels = n_mel_channels
        self.n_f0_dims = n_f0_dims  # >= 1 to trains with f0
        self.n_energy_avg_dims = n_energy_avg_dims  # >= 1 trains with energy
        self.decoder_use_partial_padding = kwargs.get(
            "decoder_use_partial_padding", True
        )
        self.n_speaker_dim = n_speaker_dim
        assert self.n_speaker_dim % 2 == 0
        self.speaker_embedding = torch.nn.Embedding(n_speakers, self.n_speaker_dim)
        self.embedding = torch.nn.Embedding(n_text, n_text_dim)
        self.flows = torch.nn.ModuleList()
        self.encoder = Encoder(
            encoder_embedding_dim=n_text_dim,
            norm_fn=nn.InstanceNorm1d,
            lstm_norm_fn=text_encoder_lstm_norm,
        )
        self.dummy_speaker_embedding = dummy_speaker_embedding
        self.learn_alignments = learn_alignments
        self.affine_activation = affine_activation
        self.include_modules = include_modules
        self.attn_use_CTC = bool(attn_use_CTC)
        self.use_speaker_emb_for_alignment = use_speaker_emb_for_alignment
        self.use_context_lstm = bool(use_context_lstm)
        self.context_lstm_norm = context_lstm_norm
        self.context_lstm_w_f0_and_energy = context_lstm_w_f0_and_energy
        self.length_regulator = LengthRegulator()
        self.use_first_order_features = bool(use_first_order_features)
        self.decoder_use_unvoiced_bias = kwargs.get("decoder_use_unvoiced_bias", True)
        self.ap_pred_log_f0 = ap_pred_log_f0
        self.ap_use_unvoiced_bias = kwargs.get("ap_use_unvoiced_bias", True)
        self.attn_straight_through_estimator = kwargs.get(
            "attn_straight_through_estimator", False
        )
        if "atn" in include_modules or "dec" in include_modules:
            if self.learn_alignments:
                if self.use_speaker_emb_for_alignment:
                    self.attention = ConvAttention(
                        n_mel_channels, n_text_dim + self.n_speaker_dim
                    )
                else:
                    self.attention = ConvAttention(n_mel_channels, n_text_dim)

            self.n_flows = n_flows
            self.n_group_size = n_group_size

            n_flowstep_cond_dims = (
                self.n_speaker_dim
                + (n_text_dim + n_f0_dims + n_energy_avg_dims) * n_group_size
            )

            if self.use_context_lstm:
                n_in_context_lstm = self.n_speaker_dim + n_text_dim * n_group_size
                n_context_lstm_hidden = int(
                    (self.n_speaker_dim + n_text_dim * n_group_size) / 2
                )

                if self.context_lstm_w_f0_and_energy:
                    n_in_context_lstm = n_f0_dims + n_energy_avg_dims + n_text_dim
                    n_in_context_lstm *= n_group_size
                    n_in_context_lstm += self.n_speaker_dim

                    n_context_hidden = n_f0_dims + n_energy_avg_dims + n_text_dim
                    n_context_hidden = n_context_hidden * n_group_size / 2
                    n_context_hidden = self.n_speaker_dim + n_context_hidden
                    n_context_hidden = int(n_context_hidden)

                    n_flowstep_cond_dims = (
                        self.n_speaker_dim + n_text_dim * n_group_size
                    )

                self.context_lstm = torch.nn.LSTM(
                    input_size=n_in_context_lstm,
                    hidden_size=n_context_lstm_hidden,
                    num_layers=1,
                    batch_first=True,
                    bidirectional=True,
                )

                if context_lstm_norm is not None:
                    if "spectral" in context_lstm_norm:
                        print("Applying spectral norm to context encoder LSTM")
                        lstm_norm_fn_pntr = (
                            torch.nn.utils.parametrizations.spectral_norm
                        )
                    elif "weight" in context_lstm_norm:
                        print("Applying weight norm to context encoder LSTM")
                        lstm_norm_fn_pntr = torch.nn.utils.parametrizations.weight_norm

                    self.context_lstm = lstm_norm_fn_pntr(
                        self.context_lstm, "weight_hh_l0"
                    )
                    self.context_lstm = lstm_norm_fn_pntr(
                        self.context_lstm, "weight_hh_l0_reverse"
                    )

            if self.n_group_size > 1:
                self.unfold_params = {
                    "kernel_size": (n_group_size, 1),
                    "stride": n_group_size,
                    "padding": 0,
                    "dilation": 1,
                }
                self.unfold = nn.Unfold(**self.unfold_params)

            self.exit_steps = []
            self.n_early_size = n_early_size
            n_mel_channels = n_mel_channels * n_group_size

            for i in range(self.n_flows):
                if i > 0 and i % n_early_every == 0:  # early exitting
                    n_mel_channels -= self.n_early_size
                    self.exit_steps.append(i)

                self.flows.append(
                    FlowStep(
                        n_mel_channels,
                        n_flowstep_cond_dims,
                        n_conv_layers_per_step,
                        affine_model,
                        scaling_fn,
                        matrix_decomposition,
                        affine_activation=affine_activation,
                        use_partial_padding=self.decoder_use_partial_padding,
                    )
                )

        if "dpm" in include_modules:
            dur_model_config["hparams"]["n_speaker_dim"] = n_speaker_dim
            self.dur_pred_layer = get_attribute_prediction_model(dur_model_config)

        self.use_unvoiced_bias = False
        self.use_vpred_module = False
        self.ap_use_voiced_embeddings = kwargs.get("ap_use_voiced_embeddings", True)

        if self.decoder_use_unvoiced_bias or self.ap_use_unvoiced_bias:
            assert unvoiced_bias_activation in {"relu", "exp"}
            self.use_unvoiced_bias = True
            if unvoiced_bias_activation == "relu":
                unvbias_nonlin = nn.ReLU()
            elif unvoiced_bias_activation == "exp":
                unvbias_nonlin = ExponentialClass()
            else:
                exit(1)  # we won't reach here anyway due to the assertion
            self.unvoiced_bias_module = nn.Sequential(
                LinearNorm(n_text_dim, 1), unvbias_nonlin
            )

        # all situations in which the vpred module is necessary
        if (
            self.ap_use_voiced_embeddings
            or self.use_unvoiced_bias
            or "vpred" in include_modules
        ):
            self.use_vpred_module = True

        if self.use_vpred_module:
            v_model_config["hparams"]["n_speaker_dim"] = n_speaker_dim
            self.v_pred_module = get_attribute_prediction_model(v_model_config)
            # 4 embeddings, first two are scales, second two are biases
            if self.ap_use_voiced_embeddings:
                self.v_embeddings = torch.nn.Embedding(4, n_text_dim)

        if "apm" in include_modules:
            f0_model_config["hparams"]["n_speaker_dim"] = n_speaker_dim
            energy_model_config["hparams"]["n_speaker_dim"] = n_speaker_dim
            if self.use_first_order_features:
                f0_model_config["hparams"]["n_in_dim"] = 2
                energy_model_config["hparams"]["n_in_dim"] = 2
                if (
                    "spline_flow_params" in f0_model_config["hparams"]
                    and f0_model_config["hparams"]["spline_flow_params"] is not None
                ):
                    f0_model_config["hparams"]["spline_flow_params"][
                        "n_in_channels"
                    ] = 2
                if (
                    "spline_flow_params" in energy_model_config["hparams"]
                    and energy_model_config["hparams"]["spline_flow_params"] is not None
                ):
                    energy_model_config["hparams"]["spline_flow_params"][
                        "n_in_channels"
                    ] = 2
            else:
                if (
                    "spline_flow_params" in f0_model_config["hparams"]
                    and f0_model_config["hparams"]["spline_flow_params"] is not None
                ):
                    f0_model_config["hparams"]["spline_flow_params"][
                        "n_in_channels"
                    ] = f0_model_config["hparams"]["n_in_dim"]
                if (
                    "spline_flow_params" in energy_model_config["hparams"]
                    and energy_model_config["hparams"]["spline_flow_params"] is not None
                ):
                    energy_model_config["hparams"]["spline_flow_params"][
                        "n_in_channels"
                    ] = energy_model_config["hparams"]["n_in_dim"]

            self.f0_pred_module = get_attribute_prediction_model(f0_model_config)
            self.energy_pred_module = get_attribute_prediction_model(
                energy_model_config
            )

    def is_attribute_unconditional(self):
        """
        returns true if the decoder is conditioned on neither energy nor F0
        """
        return self.n_f0_dims == 0 and self.n_energy_avg_dims == 0

    def encode_speaker(self, spk_ids):
        spk_ids = spk_ids * 0 if self.dummy_speaker_embedding else spk_ids
        spk_vecs = self.speaker_embedding(spk_ids)
        return spk_vecs

    def encode_text(self, text, in_lens):
        # text_embeddings: b x len_text x n_text_dim
        text_embeddings = self.embedding(text).transpose(1, 2)
        # text_enc: b x n_text_dim x encoder_dim (512)
        if in_lens is None:
            text_enc = self.encoder.infer(text_embeddings).transpose(1, 2)
        else:
            text_enc = self.encoder(text_embeddings, in_lens).transpose(1, 2)

        return text_enc, text_embeddings

    def preprocess_context(
        self, context, speaker_vecs, out_lens=None, f0=None, energy_avg=None
    ):
        if self.n_group_size > 1:
            # unfolding zero-padded values
            context = self.unfold(context.unsqueeze(-1))
            if f0 is not None:
                f0 = self.unfold(f0[:, None, :, None])
            if energy_avg is not None:
                energy_avg = self.unfold(energy_avg[:, None, :, None])
        speaker_vecs = speaker_vecs[..., None].expand(-1, -1, context.shape[2])
        context_w_spkvec = torch.cat((context, speaker_vecs), 1)

        if self.use_context_lstm:
            if self.context_lstm_w_f0_and_energy:
                if f0 is not None:
                    context_w_spkvec = torch.cat((context_w_spkvec, f0), 1)

                if energy_avg is not None:
                    context_w_spkvec = torch.cat((context_w_spkvec, energy_avg), 1)

            unfolded_out_lens = (out_lens // self.n_group_size).long().cpu()
            unfolded_out_lens_packed = nn.utils.rnn.pack_padded_sequence(
                context_w_spkvec.transpose(1, 2),
                unfolded_out_lens,
                batch_first=True,
                enforce_sorted=False,
            )
            self.context_lstm.flatten_parameters()
            context_lstm_packed_output, _ = self.context_lstm(unfolded_out_lens_packed)
            context_lstm_padded_output, _ = nn.utils.rnn.pad_packed_sequence(
                context_lstm_packed_output, batch_first=True
            )
            context_w_spkvec = context_lstm_padded_output.transpose(1, 2)

        if not self.context_lstm_w_f0_and_energy:
            if f0 is not None:
                context_w_spkvec = torch.cat((context_w_spkvec, f0), 1)

            if energy_avg is not None:
                context_w_spkvec = torch.cat((context_w_spkvec, energy_avg), 1)

        return context_w_spkvec

    def enable_inverse_cache(self):
        for flow_step in self.flows:
            flow_step.enable_inverse_cache()

    def fold(self, mel):
        """Inverse of the self.unfold(mel.unsqueeze(-1)) operation used for the
        grouping or "squeeze" operation on input

        Args:
            mel: B x C x T tensor of temporal data
        """
        mel = nn.functional.fold(
            mel, output_size=(mel.shape[2] * self.n_group_size, 1), **self.unfold_params
        ).squeeze(-1)
        return mel

    def binarize_attention(self, attn, in_lens, out_lens):
        """For training purposes only. Binarizes attention with MAS. These will
        no longer recieve a gradient
        Args:
            attn: B x 1 x max_mel_len x max_text_len
        """
        b_size = attn.shape[0]
        with torch.no_grad():
            attn_cpu = attn.data.cpu().numpy()
            attn_out = torch.zeros_like(attn)
            for ind in range(b_size):
                hard_attn = mas(attn_cpu[ind, 0, : out_lens[ind], : in_lens[ind]])
                attn_out[ind, 0, : out_lens[ind], : in_lens[ind]] = torch.tensor(
                    hard_attn, device=attn.get_device()
                )
        return attn_out

    def get_first_order_features(self, feats, out_lens, dilation=1):
        """
        feats: b x max_length
        out_lens: b-dim
        """
        # add an extra column
        feats_extended_R = torch.cat(
            (feats, torch.zeros_like(feats[:, 0:dilation])), dim=1
        )
        feats_extended_L = torch.cat(
            (torch.zeros_like(feats[:, 0:dilation]), feats), dim=1
        )
        dfeats_R = feats_extended_R[:, dilation:] - feats
        dfeats_L = feats - feats_extended_L[:, 0:-dilation]

        return (dfeats_R + dfeats_L) * 0.5

    def apply_voice_mask_to_text(self, text_enc, voiced_mask):
        """
        text_enc: b x C x N
        voiced_mask: b x N
        """
        voiced_mask = voiced_mask.unsqueeze(1)
        voiced_embedding_s = self.v_embeddings.weight[0:1, :, None]
        unvoiced_embedding_s = self.v_embeddings.weight[1:2, :, None]
        voiced_embedding_b = self.v_embeddings.weight[2:3, :, None]
        unvoiced_embedding_b = self.v_embeddings.weight[3:4, :, None]
        scale = torch.sigmoid(
            voiced_embedding_s * voiced_mask + unvoiced_embedding_s * (1 - voiced_mask)
        )
        bias = 0.1 * torch.tanh(
            voiced_embedding_b * voiced_mask + unvoiced_embedding_b * (1 - voiced_mask)
        )
        return text_enc * scale + bias

    def forward(
        self,
        mel,
        speaker_ids,
        text,
        in_lens,
        out_lens,
        binarize_attention=False,
        attn_prior=None,
        f0=None,
        energy_avg=None,
        voiced_mask=None,
        p_voiced=None,
    ):
        speaker_vecs = self.encode_speaker(speaker_ids)
        text_enc, text_embeddings = self.encode_text(text, in_lens)

        log_s_list, log_det_W_list, z_mel = [], [], []
        attn = None
        attn_soft = None
        attn_hard = None
        if "atn" in self.include_modules or "dec" in self.include_modules:
            # make sure to do the alignments before folding
            attn_mask = get_mask_from_lengths(in_lens)[..., None] == 0

            text_embeddings_for_attn = text_embeddings
            if self.use_speaker_emb_for_alignment:
                speaker_vecs_expd = speaker_vecs[:, :, None].expand(
                    -1, -1, text_embeddings.shape[2]
                )
                text_embeddings_for_attn = torch.cat(
                    (text_embeddings_for_attn, speaker_vecs_expd.detach()), 1
                )

            # attn_mask shld be 1 for unsd t-steps in text_enc_w_spkvec tensor
            attn_soft, attn_logprob = self.attention(
                mel,
                text_embeddings_for_attn,
                out_lens,
                attn_mask,
                key_lens=in_lens,
                attn_prior=attn_prior,
            )

            if binarize_attention:
                attn = self.binarize_attention(attn_soft, in_lens, out_lens)
                attn_hard = attn
                if self.attn_straight_through_estimator:
                    attn_hard = attn_soft + (attn_hard - attn_soft).detach()
            else:
                attn = attn_soft

            context = torch.bmm(text_enc, attn.squeeze(1).transpose(1, 2))

        f0_bias = 0
        # unvoiced bias forward pass
        if self.use_unvoiced_bias:
            f0_bias = self.unvoiced_bias_module(context.permute(0, 2, 1))
            f0_bias = -f0_bias[..., 0]
            f0_bias = f0_bias * (~voiced_mask.bool()).float()

        # mel decoder forward pass
        if "dec" in self.include_modules:
            if self.n_group_size > 1:
                # might truncate some frames at the end, but that's ok
                # sometimes referred to as the "squeeeze" operation
                # invert this by calling self.fold(mel_or_z)
                mel = self.unfold(mel.unsqueeze(-1))
            z_out = []
            # where context is folded
            # mask f0 in case values are interpolated

            if f0 is None:
                f0_aug = None
            else:
                if self.decoder_use_unvoiced_bias:
                    f0_aug = f0 * voiced_mask + f0_bias
                else:
                    f0_aug = f0 * voiced_mask

            context_w_spkvec = self.preprocess_context(
                context, speaker_vecs, out_lens, f0_aug, energy_avg
            )

            log_s_list, log_det_W_list, z_out = [], [], []
            unfolded_seq_lens = out_lens // self.n_group_size
            for i, flow_step in enumerate(self.flows):
                if i in self.exit_steps:
                    z = mel[:, : self.n_early_size]
                    z_out.append(z)
                    mel = mel[:, self.n_early_size :]
                mel, log_det_W, log_s = flow_step(
                    mel, context_w_spkvec, seq_lens=unfolded_seq_lens
                )
                log_s_list.append(log_s)
                log_det_W_list.append(log_det_W)

            z_out.append(mel)
            z_mel = torch.cat(z_out, 1)

        # duration predictor forward pass
        duration_model_outputs = None
        if "dpm" in self.include_modules:
            if attn_hard is None:
                attn_hard = self.binarize_attention(attn_soft, in_lens, out_lens)

            # convert hard attention to durations
            attn_hard_reduced = attn_hard.sum(2)[:, 0, :]
            duration_model_outputs = self.dur_pred_layer(
                torch.detach(text_enc),
                torch.detach(speaker_vecs),
                torch.detach(attn_hard_reduced.float()),
                in_lens,
            )

        # f0, energy, vpred predictors forward pass
        f0_model_outputs = None
        energy_model_outputs = None
        vpred_model_outputs = None
        if "apm" in self.include_modules:
            if attn_hard is None:
                attn_hard = self.binarize_attention(attn_soft, in_lens, out_lens)

            # convert hard attention to durations
            if binarize_attention:
                text_enc_time_expanded = context.clone()
            else:
                text_enc_time_expanded = torch.bmm(
                    text_enc, attn_hard.squeeze(1).transpose(1, 2)
                )

            if self.use_vpred_module:
                # unvoiced bias requires  voiced mask prediction
                vpred_model_outputs = self.v_pred_module(
                    torch.detach(text_enc_time_expanded),
                    torch.detach(speaker_vecs),
                    torch.detach(voiced_mask),
                    out_lens,
                )

                # affine transform context using voiced mask
                if self.ap_use_voiced_embeddings:
                    text_enc_time_expanded = self.apply_voice_mask_to_text(
                        text_enc_time_expanded, voiced_mask
                    )

            # whether to use the unvoiced bias in the attribute predictor
            # circumvent in-place modification
            f0_target = f0.clone()
            if self.ap_use_unvoiced_bias:
                f0_target = torch.detach(f0_target * voiced_mask + f0_bias)
            else:
                f0_target = torch.detach(f0_target)

            # fit to log f0 in f0 predictor
            f0_target[voiced_mask.bool()] = torch.log(f0_target[voiced_mask.bool()])
            f0_target = f0_target / 6  # scale to ~ [0, 1] in log space
            energy_avg = energy_avg * 2 - 1  # scale to ~ [-1, 1]

            if self.use_first_order_features:
                df0 = self.get_first_order_features(f0_target, out_lens)
                denergy_avg = self.get_first_order_features(energy_avg, out_lens)

                f0_voiced = torch.cat((f0_target[:, None], df0[:, None]), dim=1)
                energy_avg = torch.cat(
                    (energy_avg[:, None], denergy_avg[:, None]), dim=1
                )

                f0_voiced = f0_voiced * 3  # scale to ~ 1 std
                energy_avg = energy_avg * 3  # scale to ~ 1 std
            else:
                f0_voiced = f0_target * 2  # scale to ~ 1 std
                energy_avg = energy_avg * 1.4  # scale to ~ 1 std

            f0_model_outputs = self.f0_pred_module(
                text_enc_time_expanded, torch.detach(speaker_vecs), f0_voiced, out_lens
            )

            energy_model_outputs = self.energy_pred_module(
                text_enc_time_expanded, torch.detach(speaker_vecs), energy_avg, out_lens
            )

        outputs = {
            "z_mel": z_mel,
            "log_det_W_list": log_det_W_list,
            "log_s_list": log_s_list,
            "duration_model_outputs": duration_model_outputs,
            "f0_model_outputs": f0_model_outputs,
            "energy_model_outputs": energy_model_outputs,
            "vpred_model_outputs": vpred_model_outputs,
            "attn_soft": attn_soft,
            "attn": attn,
            "text_embeddings": text_embeddings,
            "attn_logprob": attn_logprob,
        }

        return outputs

    def infer(
        self,
        speaker_id,
        text,
        sigma,
        sigma_dur=0.8,
        sigma_f0=0.8,
        sigma_energy=0.8,
        token_dur_scaling=1.0,
        token_duration_max=100,
        speaker_id_text=None,
        speaker_id_attributes=None,
        dur=None,
        f0=None,
        energy_avg=None,
        voiced_mask=None,
        f0_mean=0.0,
        f0_std=0.0,
        energy_mean=0.0,
        energy_std=0.0,
    ):
        batch_size = text.shape[0]
        n_tokens = text.shape[1]
        spk_vec = self.encode_speaker(speaker_id)
        spk_vec_text, spk_vec_attributes = spk_vec, spk_vec
        if speaker_id_text is not None:
            spk_vec_text = self.encode_speaker(speaker_id_text)
        if speaker_id_attributes is not None:
            spk_vec_attributes = self.encode_speaker(speaker_id_attributes)

        txt_enc, txt_emb = self.encode_text(text, None)

        if dur is None:
            # get token durations
            z_dur = (
                torch.randn(batch_size, 1, n_tokens, dtype=torch.float32, device=device)
                * sigma_dur
            )

            dur = self.dur_pred_layer.infer(z_dur, txt_enc, spk_vec_text)
            if dur.shape[-1] < txt_enc.shape[-1]:
                to_pad = txt_enc.shape[-1] - dur.shape[2]
                pad_fn = nn.ReplicationPad1d((0, to_pad))
                dur = pad_fn(dur)
            dur = dur[:, 0]
            dur = dur.clamp(0, token_duration_max)
            dur = dur * token_dur_scaling if token_dur_scaling > 0 else dur
            dur = (dur + 0.5).floor().int()

        out_lens = dur.sum(1).long().cpu() if dur.shape[0] != 1 else [dur.sum(1)]
        max_n_frames = max(out_lens)

        out_lens = torch.LongTensor(out_lens).to(txt_enc.device)

        # get attributes f0, energy, vpred, etc)
        txt_enc_time_expanded = self.length_regulator(
            txt_enc.transpose(1, 2), dur
        ).transpose(1, 2)

        if not self.is_attribute_unconditional():
            # if explicitly modeling attributes
            if voiced_mask is None:
                if self.use_vpred_module:
                    # get logits
                    voiced_mask = self.v_pred_module.infer(
                        None, txt_enc_time_expanded, spk_vec_attributes
                    )
                    voiced_mask = torch.sigmoid(voiced_mask[:, 0]) > 0.5
                    voiced_mask = voiced_mask.float()

            ap_txt_enc_time_expanded = txt_enc_time_expanded
            # voice mask augmentation only used for attribute prediction
            if self.ap_use_voiced_embeddings:
                ap_txt_enc_time_expanded = self.apply_voice_mask_to_text(
                    txt_enc_time_expanded, voiced_mask
                )

            f0_bias = 0
            # unvoiced bias forward pass
            if self.use_unvoiced_bias:
                f0_bias = self.unvoiced_bias_module(
                    txt_enc_time_expanded.permute(0, 2, 1)
                )
                f0_bias = -f0_bias[..., 0]
                f0_bias = f0_bias * (~voiced_mask.bool()).float()

            if f0 is None:
                n_f0_feature_channels = 2 if self.use_first_order_features else 1

                z_f0 = (
                    torch.randn(
                        batch_size,
                        n_f0_feature_channels,
                        max_n_frames,
                        dtype=torch.float32,
                    )
                    * sigma_f0
                ).to(device)

                f0 = self.infer_f0(
                    z_f0,
                    ap_txt_enc_time_expanded,
                    spk_vec_attributes,
                    voiced_mask,
                    out_lens,
                )[:, 0]

            if f0_mean > 0.0:
                vmask_bool = voiced_mask.bool()
                f0_mu, f0_sigma = f0[vmask_bool].mean(), f0[vmask_bool].std()
                f0[vmask_bool] = (f0[vmask_bool] - f0_mu) / f0_sigma
                f0_std = f0_std if f0_std > 0 else f0_sigma
                f0[vmask_bool] = f0[vmask_bool] * f0_std + f0_mean

            if energy_avg is None:
                n_energy_feature_channels = 2 if self.use_first_order_features else 1

                z_energy_avg = (
                    torch.randn(
                        batch_size,
                        n_energy_feature_channels,
                        max_n_frames,
                        dtype=torch.float32,
                        device=device,
                    )
                    * sigma_energy
                )

                energy_avg = self.infer_energy(
                    z_energy_avg, ap_txt_enc_time_expanded, spk_vec, out_lens
                )[:, 0]

            # replication pad, because ungrouping with different group sizes
            # may lead to mismatched lengths
            if energy_avg.shape[1] < out_lens[0]:
                to_pad = out_lens[0] - energy_avg.shape[1]
                pad_fn = nn.ReplicationPad1d((0, to_pad))
                f0 = pad_fn(f0[None])[0]
                energy_avg = pad_fn(energy_avg[None])[0]
            if f0.shape[1] < out_lens[0]:
                to_pad = out_lens[0] - f0.shape[1]
                pad_fn = nn.ReplicationPad1d((0, to_pad))
                f0 = pad_fn(f0[None])[0]

            if self.decoder_use_unvoiced_bias:
                context_w_spkvec = self.preprocess_context(
                    txt_enc_time_expanded,
                    spk_vec,
                    out_lens,
                    f0 * voiced_mask + f0_bias,
                    energy_avg,
                )
            else:
                context_w_spkvec = self.preprocess_context(
                    txt_enc_time_expanded,
                    spk_vec,
                    out_lens,
                    f0 * voiced_mask,
                    energy_avg,
                )
        else:
            context_w_spkvec = self.preprocess_context(
                txt_enc_time_expanded, spk_vec, out_lens, None, None
            )

        residual = torch.randn(
            batch_size,
            80 * self.n_group_size,
            max_n_frames // self.n_group_size,
            dtype=torch.float32,
        ).to(device)

        residual = residual * sigma

        # map from z sample to data
        exit_steps_stack = self.exit_steps.copy()
        mel = residual[:, len(exit_steps_stack) * self.n_early_size :]
        remaining_residual = residual[:, : len(exit_steps_stack) * self.n_early_size]
        unfolded_seq_lens = out_lens // self.n_group_size
        for i, flow_step in enumerate(reversed(self.flows)):
            curr_step = len(self.flows) - i - 1
            mel = flow_step(
                mel, context_w_spkvec, inverse=True, seq_lens=unfolded_seq_lens
            )
            if len(exit_steps_stack) > 0 and curr_step == exit_steps_stack[-1]:
                # concatenate the next chunk of z
                exit_steps_stack.pop()
                residual_to_add = remaining_residual[
                    :, len(exit_steps_stack) * self.n_early_size :
                ]
                remaining_residual = remaining_residual[
                    :, : len(exit_steps_stack) * self.n_early_size
                ]
                mel = torch.cat((residual_to_add, mel), 1)

        if self.n_group_size > 1:
            mel = self.fold(mel)
        if self.do_mel_descaling:
            mel = mel * 2 - 5.5

        return {
            "mel": mel,
            "dur": dur,
            "f0": f0,
            "energy_avg": energy_avg,
            "voiced_mask": voiced_mask,
        }

    def infer_f0(
        self, residual, txt_enc_time_expanded, spk_vec, voiced_mask=None, lens=None
    ):
        f0 = self.f0_pred_module.infer(residual, txt_enc_time_expanded, spk_vec, lens)

        if voiced_mask is not None and len(voiced_mask.shape) == 2:
            voiced_mask = voiced_mask[:, None]

        # constants
        if self.ap_pred_log_f0:
            if self.use_first_order_features:
                f0 = f0[:, 0:1, :] / 3
            else:
                f0 = f0 / 2
            f0 = f0 * 6
        else:
            f0 = f0 / 6
            f0 = f0 / 640

        if voiced_mask is None:
            voiced_mask = f0 > 0.0
        else:
            voiced_mask = voiced_mask.bool()

        # due to grouping, f0 might be 1 frame short
        voiced_mask = voiced_mask[:, :, : f0.shape[-1]]
        if self.ap_pred_log_f0:
            # if variable is set, decoder sees linear f0
            # mask = f0 > 0.0 if voiced_mask is None else voiced_mask.bool()
            f0[voiced_mask] = torch.exp(f0[voiced_mask])
        f0[~voiced_mask] = 0.0
        return f0

    def infer_energy(self, residual, txt_enc_time_expanded, spk_vec, lens):
        energy = self.energy_pred_module.infer(
            residual, txt_enc_time_expanded, spk_vec, lens
        )

        # magic constants
        if self.use_first_order_features:
            energy = energy / 3
        else:
            energy = energy / 1.4
        energy = (energy + 1) / 2
        return energy

    def remove_norms(self):
        """Removes spectral and weightnorms from model. Call before inference"""
        for name, module in self.named_modules():
            try:
                nn.utils.remove_spectral_norm(module, name="weight_hh_l0")
                print("Removed spectral norm from {}".format(name))
            except Exception as e:
                print(e)

            try:
                nn.utils.remove_spectral_norm(module, name="weight_hh_l0_reverse")
                print("Removed spectral norm from {}".format(name))
            except Exception as e:
                print(e)

            try:
                nn.utils.remove_weight_norm(module)
                print("Removed wnorm from {}".format(name))
            except Exception as e:
                print(e)
