# -*- coding: utf-8 -*-
# Copyright 2020 The FastSpeech Authors, The HuggingFace Inc. team and Minh Nguyen (@dathudeptrai)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tensorflow Model modules for FastSpeech."""

import numpy as np
import tensorflow as tf


def get_initializer(initializer_range=0.02):
    """Creates a `tf.initializers.truncated_normal` with the given range.

    Args:
        initializer_range: float, initializer range for stddev.

    Returns:
        TruncatedNormal initializer with stddev = `initializer_range`.

    """
    return tf.keras.initializers.TruncatedNormal(stddev=initializer_range)


def gelu(x):
    """Gaussian Error Linear unit."""
    cdf = 0.5 * (1.0 + tf.math.erf(x / tf.math.sqrt(2.0)))
    return x * cdf


def gelu_new(x):
    """Smoother gaussian Error Linear Unit."""
    cdf = 0.5 * (1.0 + tf.tanh((np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
    return x * cdf


def swish(x):
    """Swish activation function."""
    return x * tf.sigmoid(x)


def mish(x):
    return x * tf.math.tanh(tf.math.softplus(x))


ACT2FN = {
    "identity": tf.keras.layers.Activation("linear"),
    "tanh": tf.keras.layers.Activation("tanh"),
    "gelu": tf.keras.layers.Activation(gelu),
    "relu": tf.keras.activations.relu,
    "swish": tf.keras.layers.Activation(swish),
    "gelu_new": tf.keras.layers.Activation(gelu_new),
    "mish": tf.keras.layers.Activation(mish),
}


class TFFastSpeechEmbeddings(tf.keras.layers.Layer):
    """Construct charactor/phoneme/positional/speaker embeddings."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.vocab_size = config.vocab_size
        self.hidden_size = config.encoder_self_attention_params.hidden_size
        self.initializer_range = config.initializer_range
        self.config = config

        self.position_embeddings = tf.keras.layers.Embedding(
            config.max_position_embeddings + 1,
            self.hidden_size,
            weights=[self._sincos_embedding()],
            name="position_embeddings",
            trainable=False,
        )

        if config.n_speakers > 1:
            self.encoder_speaker_embeddings = tf.keras.layers.Embedding(
                config.n_speakers,
                self.hidden_size,
                embeddings_initializer=get_initializer(self.initializer_range),
                name="speaker_embeddings",
            )
            self.speaker_fc = tf.keras.layers.Dense(
                units=self.hidden_size, name="speaker_fc"
            )

    def build(self, input_shape):
        """Build shared charactor/phoneme embedding layers."""
        with tf.name_scope("charactor_embeddings"):
            self.charactor_embeddings = self.add_weight(
                "weight",
                shape=[self.vocab_size, self.hidden_size],
                initializer=get_initializer(self.initializer_range),
            )
        super().build(input_shape)

    def call(self, inputs, training=False):
        """Get charactor embeddings of inputs.

        Args:
            1. charactor, Tensor (int32) shape [batch_size, length].
            2. speaker_id, Tensor (int32) shape [batch_size]
        Returns:
            Tensor (float32) shape [batch_size, length, embedding_size].

        """
        return self._embedding(inputs, training=training)

    def _embedding(self, inputs, training=False):
        """Applies embedding based on inputs tensor."""
        input_ids, speaker_ids = inputs

        input_shape = tf.shape(input_ids)
        seq_length = input_shape[1]

        position_ids = tf.range(1, seq_length + 1, dtype=tf.int32)[tf.newaxis, :]

        # create embeddings
        inputs_embeds = tf.gather(self.charactor_embeddings, input_ids)
        position_embeddings = self.position_embeddings(position_ids)

        # sum embedding
        embeddings = inputs_embeds + position_embeddings
        if self.config.n_speakers > 1:
            speaker_embeddings = self.encoder_speaker_embeddings(speaker_ids)
            speaker_features = tf.math.softplus(self.speaker_fc(speaker_embeddings))
            # extended speaker embeddings
            extended_speaker_features = speaker_features[:, tf.newaxis, :]
            embeddings += extended_speaker_features

        return embeddings

    def _sincos_embedding(self):
        position_enc = np.array(
            [
                [
                    pos / np.power(10000, 2.0 * (i // 2) / self.hidden_size)
                    for i in range(self.hidden_size)
                ]
                for pos in range(self.config.max_position_embeddings + 1)
            ]
        )

        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])

        # pad embedding.
        position_enc[0] = 0.0

        return position_enc


class TFFastSpeechSelfAttention(tf.keras.layers.Layer):
    """Self attention module for fastspeech."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
        self.output_attentions = config.output_attentions
        self.num_attention_heads = config.num_attention_heads
        self.all_head_size = self.num_attention_heads * config.attention_head_size

        self.query = tf.keras.layers.Dense(
            self.all_head_size,
            kernel_initializer=get_initializer(config.initializer_range),
            name="query",
        )
        self.key = tf.keras.layers.Dense(
            self.all_head_size,
            kernel_initializer=get_initializer(config.initializer_range),
            name="key",
        )
        self.value = tf.keras.layers.Dense(
            self.all_head_size,
            kernel_initializer=get_initializer(config.initializer_range),
            name="value",
        )

        self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob)
        self.config = config

    def transpose_for_scores(self, x, batch_size):
        """Transpose to calculate attention scores."""
        x = tf.reshape(
            x,
            (batch_size, -1, self.num_attention_heads, self.config.attention_head_size),
        )
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        batch_size = tf.shape(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        attention_scores = tf.matmul(query_layer, key_layer, transpose_b=True)
        dk = tf.cast(tf.shape(key_layer)[-1], tf.float32)  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # extended_attention_masks for self attention encoder.
            extended_attention_mask = attention_mask[:, tf.newaxis, tf.newaxis, :]
            extended_attention_mask = tf.cast(extended_attention_mask, tf.float32)
            extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
            attention_scores = attention_scores + extended_attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = tf.nn.softmax(attention_scores, axis=-1)
        attention_probs = self.dropout(attention_probs, training=training)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(context_layer, (batch_size, -1, self.all_head_size))

        outputs = (
            (context_layer, attention_probs)
            if self.output_attentions
            else (context_layer,)
        )
        return outputs


class TFFastSpeechSelfOutput(tf.keras.layers.Layer):
    """Fastspeech output of self attention module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.dense = tf.keras.layers.Dense(
            config.hidden_size,
            kernel_initializer=get_initializer(config.initializer_range),
            name="dense",
        )
        self.LayerNorm = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="LayerNorm"
        )
        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)

    def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, input_tensor = inputs

        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class TFFastSpeechAttention(tf.keras.layers.Layer):
    """Fastspeech attention module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.self_attention = TFFastSpeechSelfAttention(config, name="self")
        self.dense_output = TFFastSpeechSelfOutput(config, name="output")

    def call(self, inputs, training=False):
        input_tensor, attention_mask = inputs

        self_outputs = self.self_attention(
            [input_tensor, attention_mask], training=training
        )
        attention_output = self.dense_output(
            [self_outputs[0], input_tensor], training=training
        )
        masked_attention_output = attention_output * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=tf.float32
        )
        outputs = (masked_attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs


class TFFastSpeechIntermediate(tf.keras.layers.Layer):
    """Intermediate representation module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.conv1d_1 = tf.keras.layers.Conv1D(
            config.intermediate_size,
            kernel_size=config.intermediate_kernel_size,
            kernel_initializer=get_initializer(config.initializer_range),
            padding="same",
            name="conv1d_1",
        )
        self.conv1d_2 = tf.keras.layers.Conv1D(
            config.hidden_size,
            kernel_size=config.intermediate_kernel_size,
            kernel_initializer=get_initializer(config.initializer_range),
            padding="same",
            name="conv1d_2",
        )
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def call(self, inputs):
        """Call logic."""
        hidden_states, attention_mask = inputs

        hidden_states = self.conv1d_1(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        hidden_states = self.conv1d_2(hidden_states)

        masked_hidden_states = hidden_states * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=tf.float32
        )
        return masked_hidden_states


class TFFastSpeechOutput(tf.keras.layers.Layer):
    """Output module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.LayerNorm = tf.keras.layers.LayerNormalization(
            epsilon=config.layer_norm_eps, name="LayerNorm"
        )
        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)

    def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, input_tensor = inputs

        hidden_states = self.dropout(hidden_states, training=training)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class TFFastSpeechLayer(tf.keras.layers.Layer):
    """Fastspeech module (FFT module on the paper)."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.attention = TFFastSpeechAttention(config, name="attention")
        self.intermediate = TFFastSpeechIntermediate(config, name="intermediate")
        self.bert_output = TFFastSpeechOutput(config, name="output")

    def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        attention_outputs = self.attention(
            [hidden_states, attention_mask], training=training
        )
        attention_output = attention_outputs[0]
        intermediate_output = self.intermediate(
            [attention_output, attention_mask], training=training
        )
        layer_output = self.bert_output(
            [intermediate_output, attention_output], training=training
        )
        masked_layer_output = layer_output * tf.cast(
            tf.expand_dims(attention_mask, 2), dtype=tf.float32
        )
        outputs = (masked_layer_output,) + attention_outputs[
            1:
        ]  # add attentions if we output them
        return outputs


class TFFastSpeechEncoder(tf.keras.layers.Layer):
    """Fast Speech encoder module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.layer = [
            TFFastSpeechLayer(config, name="layer_._{}".format(i))
            for i in range(config.num_hidden_layers)
        ]

    def call(self, inputs, training=False):
        """Call logic."""
        hidden_states, attention_mask = inputs

        all_hidden_states = ()
        all_attentions = ()
        for _, layer_module in enumerate(self.layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(
                [hidden_states, attention_mask], training=training
            )
            hidden_states = layer_outputs[0]

            if self.output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
        if self.output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # outputs, (hidden states), (attentions)


class TFFastSpeechDecoder(TFFastSpeechEncoder):
    """Fast Speech decoder module."""

    def __init__(self, config, **kwargs):
        self.is_compatible_encoder = kwargs.pop("is_compatible_encoder", True)

        super().__init__(config, **kwargs)
        self.config = config

        # create decoder positional embedding
        self.decoder_positional_embeddings = tf.keras.layers.Embedding(
            config.max_position_embeddings + 1,
            config.hidden_size,
            weights=[self._sincos_embedding()],
            name="position_embeddings",
            trainable=False,
        )

        if self.is_compatible_encoder is False:
            self.project_compatible_decoder = tf.keras.layers.Dense(
                units=config.hidden_size, name="project_compatible_decoder"
            )

        if config.n_speakers > 1:
            self.decoder_speaker_embeddings = tf.keras.layers.Embedding(
                config.n_speakers,
                config.hidden_size,
                embeddings_initializer=get_initializer(config.initializer_range),
                name="speaker_embeddings",
            )
            self.speaker_fc = tf.keras.layers.Dense(
                units=config.hidden_size, name="speaker_fc"
            )

    def call(self, inputs, training=False):
        hidden_states, speaker_ids, encoder_mask, decoder_pos = inputs

        if self.is_compatible_encoder is False:
            hidden_states = self.project_compatible_decoder(hidden_states)

        # calculate new hidden states.
        hidden_states += self.decoder_positional_embeddings(decoder_pos)

        if self.config.n_speakers > 1:
            speaker_embeddings = self.decoder_speaker_embeddings(speaker_ids)
            speaker_features = tf.math.softplus(self.speaker_fc(speaker_embeddings))
            # extended speaker embeddings
            extended_speaker_features = speaker_features[:, tf.newaxis, :]
            hidden_states += extended_speaker_features

        return super().call([hidden_states, encoder_mask], training=training)

    def _sincos_embedding(self):
        position_enc = np.array(
            [
                [
                    pos / np.power(10000, 2.0 * (i // 2) / self.config.hidden_size)
                    for i in range(self.config.hidden_size)
                ]
                for pos in range(self.config.max_position_embeddings + 1)
            ]
        )

        position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])
        position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])

        # pad embedding.
        position_enc[0] = 0.0

        return position_enc


class TFTacotronPostnet(tf.keras.layers.Layer):
    """Tacotron-2 postnet."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.conv_batch_norm = []
        for i in range(config.n_conv_postnet):
            conv = tf.keras.layers.Conv1D(
                filters=config.postnet_conv_filters
                if i < config.n_conv_postnet - 1
                else config.num_mels,
                kernel_size=config.postnet_conv_kernel_sizes,
                padding="same",
                name="conv_._{}".format(i),
            )
            batch_norm = tf.keras.layers.BatchNormalization(
                name="batch_norm_._{}".format(i)
            )
            self.conv_batch_norm.append((conv, batch_norm))
        self.dropout = tf.keras.layers.Dropout(
            rate=config.postnet_dropout_rate, name="dropout"
        )
        self.activation = [tf.nn.tanh] * (config.n_conv_postnet - 1) + [tf.identity]

    def call(self, inputs, training=False):
        """Call logic."""
        outputs, mask = inputs
        extended_mask = tf.cast(tf.expand_dims(mask, axis=2), tf.float32)
        for i, (conv, bn) in enumerate(self.conv_batch_norm):
            outputs = conv(outputs)
            outputs = bn(outputs)
            outputs = self.activation[i](outputs)
            outputs = self.dropout(outputs, training=training)
        return outputs * extended_mask


class TFFastSpeechDurationPredictor(tf.keras.layers.Layer):
    """FastSpeech duration predictor module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        super().__init__(**kwargs)
        self.conv_layers = []
        for i in range(config.num_duration_conv_layers):
            self.conv_layers.append(
                tf.keras.layers.Conv1D(
                    config.duration_predictor_filters,
                    config.duration_predictor_kernel_sizes,
                    padding="same",
                    name="conv_._{}".format(i),
                )
            )
            self.conv_layers.append(
                tf.keras.layers.LayerNormalization(
                    epsilon=config.layer_norm_eps, name="LayerNorm_._{}".format(i)
                )
            )
            self.conv_layers.append(tf.keras.layers.Activation(tf.nn.relu6))
            self.conv_layers.append(
                tf.keras.layers.Dropout(config.duration_predictor_dropout_probs)
            )
        self.conv_layers_sequence = tf.keras.Sequential(self.conv_layers)
        self.output_layer = tf.keras.layers.Dense(1)

    def call(self, inputs, training=False):
        """Call logic."""
        encoder_hidden_states, attention_mask = inputs
        attention_mask = tf.cast(tf.expand_dims(attention_mask, 2), tf.float32)

        # mask encoder hidden states
        masked_encoder_hidden_states = encoder_hidden_states * attention_mask

        # pass though first layer
        outputs = self.conv_layers_sequence(masked_encoder_hidden_states)
        outputs = self.output_layer(outputs)
        masked_outputs = outputs * attention_mask
        return tf.squeeze(tf.nn.relu6(masked_outputs), -1)  # make sure positive value.


class TFFastSpeechLengthRegulator(tf.keras.layers.Layer):
    """FastSpeech lengthregulator module."""

    def __init__(self, config, **kwargs):
        """Init variables."""
        self.enable_tflite_convertible = kwargs.pop("enable_tflite_convertible", False)
        super().__init__(**kwargs)
        self.config = config

    def call(self, inputs, training=False):
        """Call logic.
        Args:
            1. encoder_hidden_states, Tensor (float32) shape [batch_size, length, hidden_size]
            2. durations_gt, Tensor (float32/int32) shape [batch_size, length]
        """
        encoder_hidden_states, durations_gt = inputs
        outputs, encoder_masks = self._length_regulator(
            encoder_hidden_states, durations_gt
        )
        return outputs, encoder_masks

    def _length_regulator(self, encoder_hidden_states, durations_gt):
        """Length regulator logic."""
        sum_durations = tf.reduce_sum(durations_gt, axis=-1)  # [batch_size]
        max_durations = tf.reduce_max(sum_durations)

        input_shape = tf.shape(encoder_hidden_states)
        batch_size = input_shape[0]
        hidden_size = input_shape[-1]

        # initialize output hidden states and encoder masking.
        if self.enable_tflite_convertible:
            # There is only 1 batch in inference, so we don't have to use
            # `tf.While` op with 3-D output tensor.
            repeats = durations_gt[0]
            real_length = tf.reduce_sum(repeats)
            pad_size = max_durations - real_length
            # masks : [max_durations]
            masks = tf.sequence_mask([real_length], max_durations, dtype=tf.int32)
            repeat_encoder_hidden_states = tf.repeat(
                encoder_hidden_states[0], repeats=repeats, axis=0
            )
            repeat_encoder_hidden_states = tf.expand_dims(
                tf.pad(repeat_encoder_hidden_states, [[0, pad_size], [0, 0]]), 0
            )  # [1, max_durations, hidden_size]

            outputs = repeat_encoder_hidden_states
            encoder_masks = masks
        else:
            outputs = tf.zeros(shape=[0, max_durations, hidden_size], dtype=tf.float32)
            encoder_masks = tf.zeros(shape=[0, max_durations], dtype=tf.int32)

            def condition(
                i,
                batch_size,
                outputs,
                encoder_masks,
                encoder_hidden_states,
                durations_gt,
                max_durations,
            ):
                return tf.less(i, batch_size)

            def body(
                i,
                batch_size,
                outputs,
                encoder_masks,
                encoder_hidden_states,
                durations_gt,
                max_durations,
            ):
                repeats = durations_gt[i]
                real_length = tf.reduce_sum(repeats)
                pad_size = max_durations - real_length
                masks = tf.sequence_mask([real_length], max_durations, dtype=tf.int32)
                repeat_encoder_hidden_states = tf.repeat(
                    encoder_hidden_states[i], repeats=repeats, axis=0
                )
                repeat_encoder_hidden_states = tf.expand_dims(
                    tf.pad(repeat_encoder_hidden_states, [[0, pad_size], [0, 0]]), 0
                )  # [1, max_durations, hidden_size]
                outputs = tf.concat([outputs, repeat_encoder_hidden_states], axis=0)
                encoder_masks = tf.concat([encoder_masks, masks], axis=0)
                return [
                    i + 1,
                    batch_size,
                    outputs,
                    encoder_masks,
                    encoder_hidden_states,
                    durations_gt,
                    max_durations,
                ]

            # initialize iteration i.
            i = tf.constant(0, dtype=tf.int32)
            _, _, outputs, encoder_masks, _, _, _, = tf.while_loop(
                condition,
                body,
                [
                    i,
                    batch_size,
                    outputs,
                    encoder_masks,
                    encoder_hidden_states,
                    durations_gt,
                    max_durations,
                ],
                shape_invariants=[
                    i.get_shape(),
                    batch_size.get_shape(),
                    tf.TensorShape(
                        [
                            None,
                            None,
                            self.config.encoder_self_attention_params.hidden_size,
                        ]
                    ),
                    tf.TensorShape([None, None]),
                    encoder_hidden_states.get_shape(),
                    durations_gt.get_shape(),
                    max_durations.get_shape(),
                ],
            )

        return outputs, encoder_masks


class TFFastSpeech(tf.keras.Model):
    """TF Fastspeech module."""

    def __init__(self, config, **kwargs):
        """Init layers for fastspeech."""
        self.enable_tflite_convertible = kwargs.pop("enable_tflite_convertible", False)
        super().__init__(**kwargs)
        self.embeddings = TFFastSpeechEmbeddings(config, name="embeddings")
        self.encoder = TFFastSpeechEncoder(
            config.encoder_self_attention_params, name="encoder"
        )
        self.duration_predictor = TFFastSpeechDurationPredictor(
            config, name="duration_predictor"
        )
        self.length_regulator = TFFastSpeechLengthRegulator(
            config,
            enable_tflite_convertible=self.enable_tflite_convertible,
            name="length_regulator",
        )
        self.decoder = TFFastSpeechDecoder(
            config.decoder_self_attention_params,
            is_compatible_encoder=config.encoder_self_attention_params.hidden_size
            == config.decoder_self_attention_params.hidden_size,
            name="decoder",
        )
        self.mel_dense = tf.keras.layers.Dense(units=config.num_mels, name="mel_before")
        self.postnet = TFTacotronPostnet(config=config, name="postnet")

        self.setup_inference_fn()

    def _build(self):
        """Dummy input for building model."""
        # fake inputs
        input_ids = tf.convert_to_tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], tf.int32)
        attention_mask = tf.convert_to_tensor(
            [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], tf.int32
        )
        speaker_ids = tf.convert_to_tensor([0], tf.int32)
        duration_gts = tf.convert_to_tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], tf.int32)
        self(input_ids, attention_mask, speaker_ids, duration_gts)

    def call(
        self, input_ids, attention_mask, speaker_ids, duration_gts, training=False
    ):
        """Call logic."""
        embedding_output = self.embeddings([input_ids, speaker_ids], training=training)
        encoder_output = self.encoder(
            [embedding_output, attention_mask], training=training
        )
        last_encoder_hidden_states = encoder_output[0]

        # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
        # rather than just use last_hidden_states of encoder for duration_predictor.
        duration_outputs = self.duration_predictor(
            [last_encoder_hidden_states, attention_mask]
        )  # [batch_size, length]

        length_regulator_outputs, encoder_masks = self.length_regulator(
            [last_encoder_hidden_states, duration_gts], training=training
        )

        # create decoder positional embedding
        decoder_pos = tf.range(
            1, tf.shape(length_regulator_outputs)[1] + 1, dtype=tf.int32
        )
        masked_decoder_pos = tf.expand_dims(decoder_pos, 0) * encoder_masks

        decoder_output = self.decoder(
            [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
            training=training,
        )
        last_decoder_hidden_states = decoder_output[0]

        # here u can use sum or concat more than 1 hidden states layers from decoder.
        mel_before = self.mel_dense(last_decoder_hidden_states)
        mel_after = (
            self.postnet([mel_before, encoder_masks], training=training) + mel_before
        )

        outputs = (mel_before, mel_after, duration_outputs)
        return outputs

    def _inference(self, input_ids, attention_mask, speaker_ids, speed_ratios):
        """Call logic."""
        embedding_output = self.embeddings([input_ids, speaker_ids], training=False)
        encoder_output = self.encoder(
            [embedding_output, attention_mask], training=False
        )
        last_encoder_hidden_states = encoder_output[0]

        # duration predictor, here use last_encoder_hidden_states, u can use more hidden_states layers
        # rather than just use last_hidden_states of encoder for duration_predictor.
        duration_outputs = self.duration_predictor(
            [last_encoder_hidden_states, attention_mask]
        )  # [batch_size, length]
        duration_outputs = tf.math.exp(duration_outputs) - 1.0

        if speed_ratios is None:
            speed_ratios = tf.convert_to_tensor(np.array([1.0]), dtype=tf.float32)

        duration_outputs = tf.cast(
            tf.math.round(duration_outputs * speed_ratios), tf.int32
        )

        length_regulator_outputs, encoder_masks = self.length_regulator(
            [last_encoder_hidden_states, duration_outputs], training=False
        )

        # create decoder positional embedding
        decoder_pos = tf.range(
            1, tf.shape(length_regulator_outputs)[1] + 1, dtype=tf.int32
        )
        masked_decoder_pos = tf.expand_dims(decoder_pos, 0) * encoder_masks

        decoder_output = self.decoder(
            [length_regulator_outputs, speaker_ids, encoder_masks, masked_decoder_pos],
            training=False,
        )
        last_decoder_hidden_states = decoder_output[0]

        # here u can use sum or concat more than 1 hidden states layers from decoder.
        mel_before = self.mel_dense(last_decoder_hidden_states)
        mel_after = (
            self.postnet([mel_before, encoder_masks], training=False) + mel_before
        )

        outputs = (mel_before, mel_after, duration_outputs)
        return outputs

    def setup_inference_fn(self):
        self.inference = tf.function(
            self._inference,
            experimental_relax_shapes=True,
            input_signature=[
                tf.TensorSpec(shape=[None, None], dtype=tf.int32),
                tf.TensorSpec(shape=[None, None], dtype=tf.bool),
                tf.TensorSpec(shape=[None,], dtype=tf.int32),
                tf.TensorSpec(shape=[None,], dtype=tf.float32),
            ],
        )

        self.inference_tflite = tf.function(
            self._inference,
            experimental_relax_shapes=True,
            input_signature=[
                tf.TensorSpec(shape=[1, None], dtype=tf.int32),
                tf.TensorSpec(shape=[1, None], dtype=tf.bool),
                tf.TensorSpec(shape=[1,], dtype=tf.int32),
                tf.TensorSpec(shape=[1,], dtype=tf.float32),
            ],
        )

tensorflow.keras.layers.LayerNormalization tensorflow.convert_to_tensor tensorflow.concat numpy.sqrt tensorflow.zeros tensorflow.reduce_sum tensorflow.cast tensorflow.keras.Sequential tensorflow.pad numpy.sin tensorflow.gather tensorflow.math.softplus tensorflow.name_scope tensorflow.keras.initializers.TruncatedNormal tensorflow.matmul tensorflow.TensorShape tensorflow.shape tensorflow.less tensorflow.keras.layers.Dense numpy.power tensorflow.pow tensorflow.math.exp tensorflow.math.round numpy.array tensorflow.sequence_mask tensorflow.reduce_max tensorflow.nn.softmax tensorflow.keras.layers.Activation tensorflow.transpose tensorflow.range tensorflow.math.sqrt tensorflow.nn.relu6 tensorflow.constant tensorflow.reshape tensorflow.sigmoid numpy.cos tensorflow.expand_dims tensorflow.repeat tensorflow.keras.layers.Dropout tensorflow.TensorSpec
[(31, 'tensorflow.keras.initializers.TruncatedNormal', 'tf.keras.initializers.TruncatedNormal', ([], {'stddev': 'initializer_range'}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['"""linear"""'], {}), True, 'import tensorflow as tf\n'), (57, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['"""tanh"""'], {}), True, 'import tensorflow as tf\n'), (58, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['gelu'], {}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['swish'], {}), True, 'import tensorflow as tf\n'), (61, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['gelu_new'], {}), True, 'import tensorflow as tf\n'), (62, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['mish'], {}), True, 'import tensorflow as tf\n'), (48, 'tensorflow.sigmoid', 'tf.sigmoid', (['x'], {}), True, 'import tensorflow as tf\n'), (122, 'tensorflow.shape', 'tf.shape', (['input_ids'], {}), True, 'import tensorflow as tf\n'), (128, 'tensorflow.gather', 'tf.gather', (['self.charactor_embeddings', 'input_ids'], {}), True, 'import tensorflow as tf\n'), (153, 'numpy.sin', 'np.sin', (['position_enc[:, 0::2]'], {}), True, 'import numpy as np\n'), (154, 'numpy.cos', 'np.cos', (['position_enc[:, 1::2]'], {}), True, 'import numpy as np\n'), (193, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['config.attention_probs_dropout_prob'], {}), True, 'import tensorflow as tf\n'), (198, 'tensorflow.reshape', 'tf.reshape', (['x', '(batch_size, -1, self.num_attention_heads, self.config.attention_head_size)'], {}), True, 'import tensorflow as tf\n'), (202, 'tensorflow.transpose', 'tf.transpose', (['x'], {'perm': '[0, 2, 1, 3]'}), True, 'import tensorflow as tf\n'), (217, 'tensorflow.matmul', 'tf.matmul', (['query_layer', 'key_layer'], {'transpose_b': '(True)'}), True, 'import tensorflow as tf\n'), (229, 'tensorflow.nn.softmax', 'tf.nn.softmax', (['attention_scores'], {'axis': '(-1)'}), True, 'import tensorflow as tf\n'), (232, 'tensorflow.matmul', 'tf.matmul', (['attention_probs', 'value_layer'], {}), True, 'import tensorflow as tf\n'), (233, 'tensorflow.transpose', 'tf.transpose', (['context_layer'], {'perm': '[0, 2, 1, 3]'}), True, 'import tensorflow as tf\n'), (234, 'tensorflow.reshape', 'tf.reshape', (['context_layer', '(batch_size, -1, self.all_head_size)'], {}), True, 'import tensorflow as tf\n'), (255, 'tensorflow.keras.layers.LayerNormalization', 'tf.keras.layers.LayerNormalization', ([], {'epsilon': 'config.layer_norm_eps', 'name': '"""LayerNorm"""'}), True, 'import tensorflow as tf\n'), (258, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['config.hidden_dropout_prob'], {}), True, 'import tensorflow as tf\n'), (342, 'tensorflow.keras.layers.LayerNormalization', 'tf.keras.layers.LayerNormalization', ([], {'epsilon': 'config.layer_norm_eps', 'name': '"""LayerNorm"""'}), True, 'import tensorflow as tf\n'), (345, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['config.hidden_dropout_prob'], {}), True, 'import tensorflow as tf\n'), (495, 'numpy.sin', 'np.sin', (['position_enc[:, 0::2]'], {}), True, 'import numpy as np\n'), (496, 'numpy.cos', 'np.cos', (['position_enc[:, 1::2]'], {}), True, 'import numpy as np\n'), (524, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', ([], {'rate': 'config.postnet_dropout_rate', 'name': '"""dropout"""'}), True, 'import tensorflow as tf\n'), (566, 'tensorflow.keras.Sequential', 'tf.keras.Sequential', (['self.conv_layers'], {}), True, 'import tensorflow as tf\n'), (567, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', (['(1)'], {}), True, 'import tensorflow as tf\n'), (607, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['durations_gt'], {'axis': '(-1)'}), True, 'import tensorflow as tf\n'), (608, 'tensorflow.reduce_max', 'tf.reduce_max', (['sum_durations'], {}), True, 'import tensorflow as tf\n'), (610, 'tensorflow.shape', 'tf.shape', (['encoder_hidden_states'], {}), True, 'import tensorflow as tf\n'), (737, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'config.num_mels', 'name': '"""mel_before"""'}), True, 'import tensorflow as tf\n'), (745, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (746, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (749, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['[0]', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (750, 'tensorflow.convert_to_tensor', 'tf.convert_to_tensor', (['[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]', 'tf.int32'], {}), True, 'import tensorflow as tf\n'), (52, 'tensorflow.math.softplus', 'tf.math.softplus', (['x'], {}), True, 'import tensorflow as tf\n'), (92, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'self.hidden_size', 'name': '"""speaker_fc"""'}), True, 'import tensorflow as tf\n'), (98, 'tensorflow.name_scope', 'tf.name_scope', (['"""charactor_embeddings"""'], {}), True, 'import tensorflow as tf\n'), (125, 'tensorflow.range', 'tf.range', (['(1)', '(seq_length + 1)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (208, 'tensorflow.shape', 'tf.shape', (['hidden_states'], {}), True, 'import tensorflow as tf\n'), (219, 'tensorflow.math.sqrt', 'tf.math.sqrt', (['dk'], {}), True, 'import tensorflow as tf\n'), (224, 'tensorflow.cast', 'tf.cast', (['extended_attention_mask', 'tf.float32'], {}), True, 'import tensorflow as tf\n'), (451, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'config.hidden_size', 'name': '"""project_compatible_decoder"""'}), True, 'import tensorflow as tf\n'), (462, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'config.hidden_size', 'name': '"""speaker_fc"""'}), True, 'import tensorflow as tf\n'), (532, 'tensorflow.expand_dims', 'tf.expand_dims', (['mask'], {'axis': '(2)'}), True, 'import tensorflow as tf\n'), (572, 'tensorflow.expand_dims', 'tf.expand_dims', (['attention_mask', '(2)'], {}), True, 'import tensorflow as tf\n'), (581, 'tensorflow.nn.relu6', 'tf.nn.relu6', (['masked_outputs'], {}), True, 'import tensorflow as tf\n'), (619, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['repeats'], {}), True, 'import tensorflow as tf\n'), (622, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[real_length]', 'max_durations'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (623, 'tensorflow.repeat', 'tf.repeat', (['encoder_hidden_states[0]'], {'repeats': 'repeats', 'axis': '(0)'}), True, 'import tensorflow as tf\n'), (633, 'tensorflow.zeros', 'tf.zeros', ([], {'shape': '[0, max_durations, hidden_size]', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (634, 'tensorflow.zeros', 'tf.zeros', ([], {'shape': '[0, max_durations]', 'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (679, 'tensorflow.constant', 'tf.constant', (['(0)'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (777, 'tensorflow.expand_dims', 'tf.expand_dims', (['decoder_pos', '(0)'], {}), True, 'import tensorflow as tf\n'), (807, 'tensorflow.math.exp', 'tf.math.exp', (['duration_outputs'], {}), True, 'import tensorflow as tf\n'), (813, 'tensorflow.math.round', 'tf.math.round', (['(duration_outputs * speed_ratios)'], {}), True, 'import tensorflow as tf\n'), (824, 'tensorflow.expand_dims', 'tf.expand_dims', (['decoder_pos', '(0)'], {}), True, 'import tensorflow as tf\n'), (218, 'tensorflow.shape', 'tf.shape', (['key_layer'], {}), True, 'import tensorflow as tf\n'), (289, 'tensorflow.expand_dims', 'tf.expand_dims', (['attention_mask', '(2)'], {}), True, 'import tensorflow as tf\n'), (331, 'tensorflow.expand_dims', 'tf.expand_dims', (['attention_mask', '(2)'], {}), True, 'import tensorflow as tf\n'), (381, 'tensorflow.expand_dims', 'tf.expand_dims', (['attention_mask', '(2)'], {}), True, 'import tensorflow as tf\n'), (562, 'tensorflow.keras.layers.Activation', 'tf.keras.layers.Activation', (['tf.nn.relu6'], {}), True, 'import tensorflow as tf\n'), (564, 'tensorflow.keras.layers.Dropout', 'tf.keras.layers.Dropout', (['config.duration_predictor_dropout_probs'], {}), True, 'import tensorflow as tf\n'), (627, 'tensorflow.pad', 'tf.pad', (['repeat_encoder_hidden_states', '[[0, pad_size], [0, 0]]'], {}), True, 'import tensorflow as tf\n'), (645, 'tensorflow.less', 'tf.less', (['i', 'batch_size'], {}), True, 'import tensorflow as tf\n'), (657, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['repeats'], {}), True, 'import tensorflow as tf\n'), (659, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[real_length]', 'max_durations'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (660, 'tensorflow.repeat', 'tf.repeat', (['encoder_hidden_states[i]'], {'repeats': 'repeats', 'axis': '(0)'}), True, 'import tensorflow as tf\n'), (666, 'tensorflow.concat', 'tf.concat', (['[outputs, repeat_encoder_hidden_states]'], {'axis': '(0)'}), True, 'import tensorflow as tf\n'), (667, 'tensorflow.concat', 'tf.concat', (['[encoder_masks, masks]'], {'axis': '(0)'}), True, 'import tensorflow as tf\n'), (810, 'numpy.array', 'np.array', (['[1.0]'], {}), True, 'import numpy as np\n'), (36, 'tensorflow.math.sqrt', 'tf.math.sqrt', (['(2.0)'], {}), True, 'import tensorflow as tf\n'), (42, 'numpy.sqrt', 'np.sqrt', (['(2 / np.pi)'], {}), True, 'import numpy as np\n'), (664, 'tensorflow.pad', 'tf.pad', (['repeat_encoder_hidden_states', '[[0, pad_size], [0, 0]]'], {}), True, 'import tensorflow as tf\n'), (775, 'tensorflow.shape', 'tf.shape', (['length_regulator_outputs'], {}), True, 'import tensorflow as tf\n'), (822, 'tensorflow.shape', 'tf.shape', (['length_regulator_outputs'], {}), True, 'import tensorflow as tf\n'), (846, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None, None]', 'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (847, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None, None]', 'dtype': 'tf.bool'}), True, 'import tensorflow as tf\n'), (848, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None]', 'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (849, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[None]', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (857, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[1, None]', 'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (858, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[1, None]', 'dtype': 'tf.bool'}), True, 'import tensorflow as tf\n'), (859, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[1]', 'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (860, 'tensorflow.TensorSpec', 'tf.TensorSpec', ([], {'shape': '[1]', 'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (146, 'numpy.power', 'np.power', (['(10000)', '(2.0 * (i // 2) / self.hidden_size)'], {}), True, 'import numpy as np\n'), (488, 'numpy.power', 'np.power', (['(10000)', '(2.0 * (i // 2) / self.config.hidden_size)'], {}), True, 'import numpy as np\n'), (695, 'tensorflow.TensorShape', 'tf.TensorShape', (['[None, None, self.config.encoder_self_attention_params.hidden_size]'], {}), True, 'import tensorflow as tf\n'), (702, 'tensorflow.TensorShape', 'tf.TensorShape', (['[None, None]'], {}), True, 'import tensorflow as tf\n'), (42, 'tensorflow.pow', 'tf.pow', (['x', '(3)'], {}), True, 'import tensorflow as tf\n')]