# -*- coding: utf-8 -*-
# Copyright 2020 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.
"""Train FastSpeech2."""

import sys
sys.path.append(".")
sys.path.append("/content/TensorFlowTTS/examples/")

from tensorflow_tts.optimizers import AdamWeightDecay
from tensorflow_tts.optimizers import WarmUp
from tensorflow_tts.models import TFFastSpeech2
from tensorflow_tts.configs import FastSpeech2Config
from tensorflow_tts.losses import GANCritic, GanLoss
from fastspeech2.fastspeech2_dataset import CharactorDurationF0EnergyMelDataset
from fastspeech.train_fastspeech import FastSpeechTrainer
from tqdm import tqdm
import tensorflow_tts
import yaml
import tensorflow as tf
import numpy as np
import argparse
import logging
import os

import torch.optim as optim


class FastSpeech2Trainer(FastSpeechTrainer):
    """FastSpeech2 Trainer class based on FastSpeechTrainer."""

    def __init__(
        self, config, steps=0, epochs=0, is_mixed_precision=False,
    ):
        """Initialize trainer.

        Args:
            steps (int): Initial global steps.
            epochs (int): Initial global epochs.
            config (dict): Config dict loaded from yaml format configuration file.
            is_mixed_precision (bool): Use mixed precision or not.

        """
        super(FastSpeech2Trainer, self).__init__(
            steps=steps,
            epochs=epochs,
            config=config,
            is_mixed_precision=is_mixed_precision,
        )
        # define metrics to aggregates data and use tf.summary logs them
        self.list_metrics_name = [
            "duration_loss",
            "f0_loss",
            "energy_loss",
            "mel_loss_before",
            "mel_loss_after",
            "gan_loss",
        ]
        self.init_train_eval_metrics(self.list_metrics_name)
        self.reset_states_train()
        self.reset_states_eval()

        self.config = config

        self.enc = GANCritic().to('cuda').double().train()
        self.opt_enc = optim.Adam(self.enc.parameters(), lr=1e-3, betas=(0.0, 0.9))





    def compile(self, model, optimizer):
        super().compile(model, optimizer)
        self.mse = tf.keras.losses.MeanSquaredError()
        self.mae = tf.keras.losses.MeanAbsoluteError()


    def _train_step(self, batch):
        """Train model one step."""
        charactor, duration, f0, energy, mel, bound, prom = batch
        self._one_step_fastspeech2(charactor, duration, f0, energy, mel, bound, prom)

        # update counts
        self.steps += 1
        self.tqdm.update(1)
        self._check_train_finish()

    @tf.function(
        experimental_relax_shapes=True,
        input_signature=[
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None, 80], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
        ],
    )
    def _one_step_fastspeech2(self, charactor, duration, f0, energy, mel, bound, prom):
        with tf.GradientTape() as tape:
            (
                mel_before,
                mel_after,
                duration_outputs,
                f0_outputs,
                energy_outputs,
            ) = self.model(
                charactor,
                attention_mask=tf.math.not_equal(charactor, 0),
                speaker_ids=tf.zeros(shape=[tf.shape(mel)[0]]),
                duration_gts=duration,
                f0_gts=f0,
                energy_gts=energy,
                bounds=bound,
                proms=prom,
                training=True,
            )
            log_duration = tf.math.log(tf.cast(tf.math.add(duration, 1), tf.float32))
            duration_loss = self.mse(log_duration, duration_outputs)
            f0_loss = self.mse(f0, f0_outputs)
            energy_loss = self.mse(energy, energy_outputs)
            mel_loss_before = self.mae(mel, mel_before)
            mel_loss_after = self.mae(mel, mel_after)
            tf.config.run_functions_eagerly(True)
            gan_loss = GanLoss(self.enc, self.opt_enc, mel, mel_after , 16)
            loss = (
                duration_loss + f0_loss + energy_loss + mel_loss_before + mel_loss_after - gan_loss
            )

            if self.is_mixed_precision:
                scaled_loss = self.optimizer.get_scaled_loss(loss)

        if self.is_mixed_precision:
            scaled_gradients = tape.gradient(
                scaled_loss, self.model.trainable_variables
            )
            gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
        else:
            gradients = tape.gradient(loss, self.model.trainable_variables)
        self.optimizer.apply_gradients(
            zip(gradients, self.model.trainable_variables), 5.0
        )

        # accumulate loss into metrics
        self.train_metrics["duration_loss"].update_state(duration_loss)
        self.train_metrics["f0_loss"].update_state(f0_loss)
        self.train_metrics["energy_loss"].update_state(energy_loss)
        self.train_metrics["mel_loss_before"].update_state(mel_loss_before)
        self.train_metrics["mel_loss_after"].update_state(mel_loss_after)
        self.train_metrics["gan_loss"].update_state(gan_loss)

    def _eval_epoch(self):
        """Evaluate model one epoch."""
        logging.info(f"(Steps: {self.steps}) Start evaluation.")

        # calculate loss for each batch
        for eval_steps_per_epoch, batch in enumerate(
            tqdm(self.eval_data_loader, desc="[eval]"), 1
        ):
            # eval one step
            charactor, duration, f0, energy, mel, bound, prom = batch
            self._eval_step(charactor, duration, f0, energy, mel, bound, prom)

            if eval_steps_per_epoch <= self.config["num_save_intermediate_results"]:
                # save intermedia
                self.generate_and_save_intermediate_result(batch)

        logging.info(
            f"(Steps: {self.steps}) Finished evaluation "
            f"({eval_steps_per_epoch} steps per epoch)."
        )

        # average loss
        for key in self.eval_metrics.keys():
            logging.info(
                f"(Steps: {self.steps}) eval_{key} = {self.eval_metrics[key].result():.4f}."
            )

        # record
        self._write_to_tensorboard(self.eval_metrics, stage="eval")

        # reset
        self.reset_states_eval()

    @tf.function(
        experimental_relax_shapes=True,
        input_signature=[
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None, 80], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
        ],
    )
    def _eval_step(self, charactor, duration, f0, energy, mel, bound, prom):
        """Evaluate model one step."""
        (
            mel_before,
            mel_after,
            duration_outputs,
            f0_outputs,
            energy_outputs,
        ) = self.model(
            charactor,
            attention_mask=tf.math.not_equal(charactor, 0),
            speaker_ids=tf.zeros(shape=[tf.shape(mel)[0]]),
            duration_gts=duration,
            f0_gts=f0,
            energy_gts=energy,
            bounds=bound,
            proms=prom,
            training=False,
        )
        log_duration = tf.math.log(tf.cast(tf.math.add(duration, 1), tf.float32))
        duration_loss = self.mse(log_duration, duration_outputs)
        f0_loss = self.mse(f0, f0_outputs)
        energy_loss = self.mse(energy, energy_outputs)
        mel_loss_before = self.mae(mel, mel_before)
        mel_loss_after = self.mae(mel, mel_after)
        gan_loss = self.ganloss(mel, mel_after)

        # accumulate loss into metrics
        self.eval_metrics["duration_loss"].update_state(duration_loss)
        self.eval_metrics["f0_loss"].update_state(f0_loss)
        self.eval_metrics["energy_loss"].update_state(energy_loss)
        self.eval_metrics["mel_loss_before"].update_state(mel_loss_before)
        self.eval_metrics["mel_loss_after"].update_state(mel_loss_after)
        self.eval_metrics["gan_loss"].update_state(gan_loss)

    def _check_log_interval(self):
        """Log to tensorboard."""
        if self.steps % self.config["log_interval_steps"] == 0:
            for metric_name in self.list_metrics_name:
                logging.info(
                    f"(Step: {self.steps}) train_{metric_name} = {self.train_metrics[metric_name].result():.4f}."
                )
            self._write_to_tensorboard(self.train_metrics, stage="train")

            # reset
            self.reset_states_train()

    @tf.function(
        experimental_relax_shapes=True,
        input_signature=[
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.int32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None, 80], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
            tf.TensorSpec([None, None], dtype=tf.float32),
        ],
    )
    def predict(self, charactor, duration, f0, energy, mel, bound, prom):
        """Predict."""
        mel_before, mel_after, _, _, _ = self.model(
            charactor,
            attention_mask=tf.math.not_equal(charactor, 0),
            speaker_ids=tf.zeros(shape=[tf.shape(mel)[0]]),
            duration_gts=duration,
            f0_gts=f0,
            energy_gts=energy,
            bounds=bound,
            proms=prom,
            training=False,
        )
        return mel_before, mel_after

    def generate_and_save_intermediate_result(self, batch):
        """Generate and save intermediate result."""
        import matplotlib.pyplot as plt

        # unpack input.
        charactor, duration, f0, energy, mel, bound, prom = batch

        # predict with tf.function.
        masked_mel_before, masked_mel_after = self.predict(
            charactor, duration, f0, energy, mel, bound, prom
        )

        # check directory
        dirname = os.path.join(self.config["outdir"], f"predictions/{self.steps}steps")
        if not os.path.exists(dirname):
            os.makedirs(dirname)

        for idx, (mel_gt, mel_pred_before, mel_pred_after) in enumerate(
            zip(mel, masked_mel_before, masked_mel_after), 1
        ):
            mel_gt = tf.reshape(mel_gt, (-1, 80)).numpy()  # [length, 80]
            mel_pred_before = tf.reshape(
                mel_pred_before, (-1, 80)
            ).numpy()  # [length, 80]
            mel_pred_after = tf.reshape(
                mel_pred_after, (-1, 80)
            ).numpy()  # [length, 80]

            # plit figure and save it
            figname = os.path.join(dirname, f"{idx}.png")
            fig = plt.figure(figsize=(10, 8))
            ax1 = fig.add_subplot(311)
            ax2 = fig.add_subplot(312)
            ax3 = fig.add_subplot(313)
            im = ax1.imshow(np.rot90(mel_gt), aspect="auto", interpolation="none")
            ax1.set_title("Target Mel-Spectrogram")
            fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax1)
            ax2.set_title("Predicted Mel-before-Spectrogram")
            im = ax2.imshow(
                np.rot90(mel_pred_before), aspect="auto", interpolation="none"
            )
            fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax2)
            ax3.set_title("Predicted Mel-after-Spectrogram")
            im = ax3.imshow(
                np.rot90(mel_pred_after), aspect="auto", interpolation="none"
            )
            fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax3)
            plt.tight_layout()
            plt.savefig(figname)
            plt.close()


def main():
    tf.config.run_functions_eagerly(True)
    """Run training process."""
    parser = argparse.ArgumentParser(
        description="Train FastSpeech (See detail in tensorflow_tts/bin/train-fastspeech.py)"
    )
    parser.add_argument(
        "--train-dir",
        default=None,
        type=str,
        help="directory including training data. ",
    )
    parser.add_argument(
        "--dev-dir",
        default=None,
        type=str,
        help="directory including development data. ",
    )
    parser.add_argument(
        "--use-norm", default=1, type=int, help="usr norm-mels for train or raw."
    )
    parser.add_argument(
        "--f0-stat",
        default="./dump/stats_f0.npy",
        type=str,
        required=True,
        help="f0-stat path.",
    )
    parser.add_argument(
        "--energy-stat",
        default="./dump/stats_energy.npy",
        type=str,
        required=True,
        help="energy-stat path.",
    )
    parser.add_argument(
        "--outdir", type=str, required=True, help="directory to save checkpoints."
    )
    parser.add_argument(
        "--config", type=str, required=True, help="yaml format configuration file."
    )
    parser.add_argument(
        "--resume",
        default="",
        type=str,
        nargs="?",
        help='checkpoint file path to resume training. (default="")',
    )
    parser.add_argument(
        "--verbose",
        type=int,
        default=1,
        help="logging level. higher is more logging. (default=1)",
    )
    parser.add_argument(
        "--mixed_precision",
        default=0,
        type=int,
        help="using mixed precision for generator or not.",
    )
    args = parser.parse_args()

    # set mixed precision config
    if args.mixed_precision == 1:
        tf.config.optimizer.set_experimental_options({"auto_mixed_precision": True})

    args.mixed_precision = bool(args.mixed_precision)
    args.use_norm = bool(args.use_norm)

    # set logger
    if args.verbose > 1:
        logging.basicConfig(
            level=logging.DEBUG,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
    elif args.verbose > 0:
        logging.basicConfig(
            level=logging.INFO,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
    else:
        logging.basicConfig(
            level=logging.WARN,
            stream=sys.stdout,
            format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
        )
        logging.warning("Skip DEBUG/INFO messages")

    # check directory existence
    if not os.path.exists(args.outdir):
        os.makedirs(args.outdir)

    # check arguments
    if args.train_dir is None:
        raise ValueError("Please specify --train-dir")
    if args.dev_dir is None:
        raise ValueError("Please specify --valid-dir")

    # load and save config
    with open(args.config) as f:
        config = yaml.load(f, Loader=yaml.Loader)
    config.update(vars(args))
    config["version"] = tensorflow_tts.__version__
    with open(os.path.join(args.outdir, "config.yml"), "w") as f:
        yaml.dump(config, f, Dumper=yaml.Dumper)
    for key, value in config.items():
        logging.info(f"{key} = {value}")

    # get dataset
    if config["remove_short_samples"]:
        mel_length_threshold = config["mel_length_threshold"]
    else:
        mel_length_threshold = None

    if config["format"] == "npy":
        charactor_query = "*-ids.npy"
        mel_query = "*-raw-feats.npy" if args.use_norm is False else "*-norm-feats.npy"
        duration_query = "*-durations.npy"
        f0_query = "*-raw-f0.npy"
        energy_query = "*-raw-energy.npy"
    else:
        raise ValueError("Only npy are supported.")

    # define train/valid dataset
    train_dataset = CharactorDurationF0EnergyMelDataset(
        root_dir=args.train_dir,
        charactor_query=charactor_query,
        mel_query=mel_query,
        duration_query=duration_query,
        f0_query=f0_query,
        energy_query=energy_query,
        f0_stat=args.f0_stat,
        energy_stat=args.energy_stat,
        mel_length_threshold=mel_length_threshold,
        return_utt_id=False,
    ).create(
        is_shuffle=config["is_shuffle"],
        allow_cache=config["allow_cache"],
        batch_size=config["batch_size"],
    )

    #print(list(train_dataset)[0])

    valid_dataset = CharactorDurationF0EnergyMelDataset(
        root_dir=args.dev_dir,
        charactor_query=charactor_query,
        mel_query=mel_query,
        duration_query=duration_query,
        f0_query=f0_query,
        energy_query=energy_query,
        f0_stat=args.f0_stat,
        energy_stat=args.energy_stat,
        mel_length_threshold=mel_length_threshold,
        return_utt_id=False,
    ).create(
        is_shuffle=config["is_shuffle"],
        allow_cache=config["allow_cache"],
        batch_size=config["batch_size"],
    )

    fastspeech = TFFastSpeech2(config=FastSpeech2Config(**config["fastspeech_params"]))
    fastspeech._build()
    fastspeech.summary()

    # define trainer
    trainer = FastSpeech2Trainer(
        config=config, steps=0, epochs=0, is_mixed_precision=args.mixed_precision
    )

    # AdamW for fastspeech
    learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
        initial_learning_rate=config["optimizer_params"]["initial_learning_rate"],
        decay_steps=config["optimizer_params"]["decay_steps"],
        end_learning_rate=config["optimizer_params"]["end_learning_rate"],
    )

    learning_rate_fn = WarmUp(
        initial_learning_rate=config["optimizer_params"]["initial_learning_rate"],
        decay_schedule_fn=learning_rate_fn,
        warmup_steps=int(
            config["train_max_steps"] * config["optimizer_params"]["warmup_proportion"]
        ),
    )

    optimizer = AdamWeightDecay(
        learning_rate=learning_rate_fn,
        weight_decay_rate=config["optimizer_params"]["weight_decay"],
        beta_1=0.9,
        beta_2=0.98,
        epsilon=1e-6,
        exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
    )

    # compile trainer
    trainer.compile(model=fastspeech, optimizer=optimizer)

    # start training
    try:
        trainer.fit(
            train_dataset,
            valid_dataset,
            saved_path=os.path.join(config["outdir"], "checkpoints/"),
            resume=args.resume,
        )
    except KeyboardInterrupt:
        trainer.save_checkpoint()
        logging.info(f"Successfully saved checkpoint @ {trainer.steps}steps.")


if __name__ == "__main__":
    main()

numpy.rot90 tensorflow.math.add tensorflow.keras.optimizers.schedules.PolynomialDecay matplotlib.pyplot.tight_layout tensorflow.shape tensorflow.keras.losses.MeanSquaredError matplotlib.pyplot.figure tensorflow.reshape tensorflow.config.run_functions_eagerly matplotlib.pyplot.savefig tensorflow.GradientTape tensorflow.config.optimizer.set_experimental_options matplotlib.pyplot.close tensorflow.math.not_equal tensorflow.TensorSpec tensorflow.keras.losses.MeanAbsoluteError
[(18, 'sys.path.append', 'sys.path.append', (['"""."""'], {}), False, 'import sys\n'), (19, 'sys.path.append', 'sys.path.append', (['"""/content/TensorFlowTTS/examples/"""'], {}), False, 'import sys\n'), (336, 'tensorflow.config.run_functions_eagerly', 'tf.config.run_functions_eagerly', (['(True)'], {}), True, 'import tensorflow as tf\n'), (338, 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Train FastSpeech (See detail in tensorflow_tts/bin/train-fastspeech.py)"""'}), False, 'import argparse\n'), (507, 'tensorflow.keras.optimizers.schedules.PolynomialDecay', 'tf.keras.optimizers.schedules.PolynomialDecay', ([], {'initial_learning_rate': "config['optimizer_params']['initial_learning_rate']", 'decay_steps': "config['optimizer_params']['decay_steps']", 'end_learning_rate': "config['optimizer_params']['end_learning_rate']"}), True, 'import tensorflow as tf\n'), (521, 'tensorflow_tts.optimizers.AdamWeightDecay', 'AdamWeightDecay', ([], {'learning_rate': 'learning_rate_fn', 'weight_decay_rate': "config['optimizer_params']['weight_decay']", 'beta_1': '(0.9)', 'beta_2': '(0.98)', 'epsilon': '(1e-06)', 'exclude_from_weight_decay': "['LayerNorm', 'layer_norm', 'bias']"}), False, 'from tensorflow_tts.optimizers import AdamWeightDecay\n'), (85, 'tensorflow.keras.losses.MeanSquaredError', 'tf.keras.losses.MeanSquaredError', ([], {}), True, 'import tensorflow as tf\n'), (86, 'tensorflow.keras.losses.MeanAbsoluteError', 'tf.keras.losses.MeanAbsoluteError', ([], {}), True, 'import tensorflow as tf\n'), (166, 'logging.info', 'logging.info', (['f"""(Steps: {self.steps}) Start evaluation."""'], {}), False, 'import logging\n'), (180, 'logging.info', 'logging.info', (['f"""(Steps: {self.steps}) Finished evaluation ({eval_steps_per_epoch} steps per epoch)."""'], {}), False, 'import logging\n'), (296, 'os.path.join', 'os.path.join', (["self.config['outdir']", 'f"""predictions/{self.steps}steps"""'], {}), False, 'import os\n'), (399, 'tensorflow.config.optimizer.set_experimental_options', 'tf.config.optimizer.set_experimental_options', (["{'auto_mixed_precision': True}"], {}), True, 'import tensorflow as tf\n'), (406, 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.DEBUG', 'stream': 'sys.stdout', 'format': '"""%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"""'}), False, 'import logging\n'), (426, 'os.path.exists', 'os.path.exists', (['args.outdir'], {}), False, 'import os\n'), (427, 'os.makedirs', 'os.makedirs', (['args.outdir'], {}), False, 'import os\n'), (437, 'yaml.load', 'yaml.load', (['f'], {'Loader': 'yaml.Loader'}), False, 'import yaml\n'), (441, 'yaml.dump', 'yaml.dump', (['config', 'f'], {'Dumper': 'yaml.Dumper'}), False, 'import yaml\n'), (443, 'logging.info', 'logging.info', (['f"""{key} = {value}"""'], {}), False, 'import logging\n'), (112, 'tensorflow.GradientTape', 'tf.GradientTape', ([], {}), True, 'import tensorflow as tf\n'), (136, 'tensorflow.config.run_functions_eagerly', 'tf.config.run_functions_eagerly', (['(True)'], {}), True, 'import tensorflow as tf\n'), (137, 'tensorflow_tts.losses.GanLoss', 'GanLoss', (['self.enc', 'self.opt_enc', 'mel', 'mel_after', '(16)'], {}), False, 'from tensorflow_tts.losses import GANCritic, GanLoss\n'), (170, 'tqdm.tqdm', 'tqdm', (['self.eval_data_loader'], {'desc': '"""[eval]"""'}), False, 'from tqdm import tqdm\n'), (297, 'os.path.exists', 'os.path.exists', (['dirname'], {}), False, 'import os\n'), (298, 'os.makedirs', 'os.makedirs', (['dirname'], {}), False, 'import os\n'), (312, 'os.path.join', 'os.path.join', (['dirname', 'f"""{idx}.png"""'], {}), False, 'import os\n'), (313, 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(10, 8)'}), True, 'import matplotlib.pyplot as plt\n'), (330, 'matplotlib.pyplot.tight_layout', 'plt.tight_layout', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (331, 'matplotlib.pyplot.savefig', 'plt.savefig', (['figname'], {}), True, 'import matplotlib.pyplot as plt\n'), (332, 'matplotlib.pyplot.close', 'plt.close', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (412, 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.INFO', 'stream': 'sys.stdout', 'format': '"""%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"""'}), False, 'import logging\n'), (418, 'logging.basicConfig', 'logging.basicConfig', ([], {'level': 'logging.WARN', 'stream': 'sys.stdout', 'format': '"""%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"""'}), False, 'import logging\n'), (423, 'logging.warning', 'logging.warning', (['"""Skip DEBUG/INFO messages"""'], {}), False, 'import logging\n'), (440, 'os.path.join', 'os.path.join', (['args.outdir', '"""config.yml"""'], {}), False, 'import os\n'), (461, 'fastspeech2.fastspeech2_dataset.CharactorDurationF0EnergyMelDataset', 'CharactorDurationF0EnergyMelDataset', ([], {'root_dir': 'args.train_dir', 'charactor_query': 'charactor_query', 'mel_query': 'mel_query', 'duration_query': 'duration_query', 'f0_query': 'f0_query', 'energy_query': 'energy_query', 'f0_stat': 'args.f0_stat', 'energy_stat': 'args.energy_stat', 'mel_length_threshold': 'mel_length_threshold', 'return_utt_id': '(False)'}), False, 'from fastspeech2.fastspeech2_dataset import CharactorDurationF0EnergyMelDataset\n'), (480, 'fastspeech2.fastspeech2_dataset.CharactorDurationF0EnergyMelDataset', 'CharactorDurationF0EnergyMelDataset', ([], {'root_dir': 'args.dev_dir', 'charactor_query': 'charactor_query', 'mel_query': 'mel_query', 'duration_query': 'duration_query', 'f0_query': 'f0_query', 'energy_query': 'energy_query', 'f0_stat': 'args.f0_stat', 'energy_stat': 'args.energy_stat', 'mel_length_threshold': 'mel_length_threshold', 'return_utt_id': '(False)'}), False, 'from fastspeech2.fastspeech2_dataset import CharactorDurationF0EnergyMelDataset\n'), (497, 'tensorflow_tts.configs.FastSpeech2Config', 'FastSpeech2Config', ([], {}), False, 'from tensorflow_tts.configs import FastSpeech2Config\n'), (543, 'logging.info', 'logging.info', (['f"""Successfully saved checkpoint @ {trainer.steps}steps."""'], {}), False, 'import logging\n'), (102, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (103, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (104, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (105, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (106, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None, 80]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (107, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (108, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (219, 'tensorflow.math.not_equal', 'tf.math.not_equal', (['charactor', '(0)'], {}), True, 'import tensorflow as tf\n'), (228, 'tensorflow.math.add', 'tf.math.add', (['duration', '(1)'], {}), True, 'import tensorflow as tf\n'), (200, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (201, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (202, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (203, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (204, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None, 80]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (205, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (206, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (272, 'tensorflow.math.not_equal', 'tf.math.not_equal', (['charactor', '(0)'], {}), True, 'import tensorflow as tf\n'), (259, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (260, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.int32'}), True, 'import tensorflow as tf\n'), (261, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (262, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (263, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None, 80]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (264, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (265, 'tensorflow.TensorSpec', 'tf.TensorSpec', (['[None, None]'], {'dtype': 'tf.float32'}), True, 'import tensorflow as tf\n'), (317, 'numpy.rot90', 'np.rot90', (['mel_gt'], {}), True, 'import numpy as np\n'), (322, 'numpy.rot90', 'np.rot90', (['mel_pred_before'], {}), True, 'import numpy as np\n'), (327, 'numpy.rot90', 'np.rot90', (['mel_pred_after'], {}), True, 'import numpy as np\n'), (538, 'os.path.join', 'os.path.join', (["config['outdir']", '"""checkpoints/"""'], {}), False, 'import os\n'), (121, 'tensorflow.math.not_equal', 'tf.math.not_equal', (['charactor', '(0)'], {}), True, 'import tensorflow as tf\n'), (130, 'tensorflow.math.add', 'tf.math.add', (['duration', '(1)'], {}), True, 'import tensorflow as tf\n'), (303, 'tensorflow.reshape', 'tf.reshape', (['mel_gt', '(-1, 80)'], {}), True, 'import tensorflow as tf\n'), (304, 'tensorflow.reshape', 'tf.reshape', (['mel_pred_before', '(-1, 80)'], {}), True, 'import tensorflow as tf\n'), (307, 'tensorflow.reshape', 'tf.reshape', (['mel_pred_after', '(-1, 80)'], {}), True, 'import tensorflow as tf\n'), (76, 'tensorflow_tts.losses.GANCritic', 'GANCritic', ([], {}), False, 'from tensorflow_tts.losses import GANCritic, GanLoss\n'), (220, 'tensorflow.shape', 'tf.shape', (['mel'], {}), True, 'import tensorflow as tf\n'), (273, 'tensorflow.shape', 'tf.shape', (['mel'], {}), True, 'import tensorflow as tf\n'), (122, 'tensorflow.shape', 'tf.shape', (['mel'], {}), True, 'import tensorflow as tf\n')]